Implement typeahead search

requirements.txt

@@ -9,4 +9,5 @@
 opentelemetry-api==1.33.1
 opentelemetry-sdk==1.33.1
 
-cachetools==6.1.0
+cachetools==6.1.0
+blake3==1.0.5

src/marqo/base_model.py

@@ -25,4 +25,16 @@ class Config(StrictBaseModel.Config, ImmutableBaseModel.Config):
 
 
 class MarqoBaseModelV2(pydantic.BaseModel):
-    model_config = ConfigDict(validate_by_name=True, validate_assignment=True)
+    model_config = ConfigDict(validate_by_name=True, validate_assignment=True)
+
+
+class StrictBaseModelV2(MarqoBaseModelV2):
+    model_config = ConfigDict(extra='forbid')
+
+
+class ImmutableBaseModelV2(MarqoBaseModelV2):
+    model_config = ConfigDict(frozen=True)
+
+
+class ImmutableStrictBaseModelV2(StrictBaseModelV2, ImmutableBaseModelV2):
+    pass

src/marqo/config.py

@@ -8,6 +8,7 @@
 from marqo.core.inference.api import Inference, ModelManager
 from marqo.core.monitoring.monitoring import Monitoring
 from marqo.core.search.recommender import Recommender
+from marqo.core.typeahead.typeahead import Typeahead
 from marqo.logging import get_logger
 from marqo.tensor_search import enums
 from marqo.tensor_search import utils
@@ -47,6 +48,7 @@ def __init__(
         self.document = Document(vespa_client, self.index_management, self.inference)
         self.recommender = Recommender(vespa_client, self.index_management, self.inference)
         self.embed = Embed(vespa_client, self.index_management, self.inference)
+        self.typeahead = Typeahead(vespa_client, self.index_management)
 
         self.model_manager = model_manager
 

src/marqo/core/index_management/index_management.py

@@ -18,6 +18,7 @@
 from marqo.core.models.marqo_index import SemiStructuredMarqoIndex
 from marqo.core.models.marqo_index_request import MarqoIndexRequest
 from marqo.core.semi_structured_vespa_index.semi_structured_vespa_schema import SemiStructuredVespaSchema
+from marqo.core.typeahead.typeahead_vespa_schema import TypeaheadVespaSchema
 from marqo.core.vespa_index.vespa_schema import for_marqo_index_request as vespa_schema_factory
 from marqo.tensor_search.models.index_settings import IndexSettings
 from marqo.vespa.vespa_client import VespaClient
@@ -162,7 +163,8 @@ def batch_create_indexes(self, marqo_index_requests: List[MarqoIndexRequest]) ->
             OperationConflictError: If another index creation/deletion operation is
                 in progress and the lock cannot be acquired
         """
-        index_to_create: List[Tuple[str, MarqoIndex]] = []
+        index_to_create: List[Tuple[str, str, MarqoIndex]] = []
+
         for request in marqo_index_requests:
             # set the default prefixes if not provided
             if request.model.text_query_prefix is None:
@@ -171,13 +173,23 @@ def batch_create_indexes(self, marqo_index_requests: List[MarqoIndexRequest]) ->
                 request.model.text_chunk_prefix = request.model.get_default_text_chunk_prefix()
 
             schema, marqo_index = vespa_schema_factory(request).generate_schema()
-            index_to_create.append((schema, marqo_index))
-            logger.debug(f'Creating index {str(request.name)} with schema:\n{schema}')
+            logger.debug(f'Creating index {request.name} with schema:\n{schema}')
+
+            typeahead_schema, updated_marqo_index = TypeaheadVespaSchema(marqo_index).generate_schema()
+            logger.debug(
+                f'Creating typeahead schema for index {request.name} with schema: '
+                f'{updated_marqo_index.typeahead_schema_name}'
+            )
+
+            index_to_create.append((schema, typeahead_schema, updated_marqo_index))
 
         with self._vespa_deployment_lock():
-            self._get_vespa_application().batch_add_index_setting_and_schema(index_to_create)
+            vespa_app = self._get_vespa_application()
+
+            # Deploy schemas and index settings (this will deploy everything together)
+            vespa_app.batch_add_index_setting_and_schema(index_to_create)
 
-        return [index for _, index in index_to_create]
+        return [index for _, _, index in index_to_create]
 
     def delete_index_by_name(self, index_name: str) -> None:
         """
@@ -231,7 +243,7 @@ def is_subset(dict_a, dict_b):
 
         if (is_subset(marqo_index.tensor_field_map, existing_index.tensor_field_map) and
                 is_subset(marqo_index.field_map, existing_index.field_map) and
-                    is_subset(marqo_index.name_to_string_array_field_map, existing_index.name_to_string_array_field_map)):
+                is_subset(marqo_index.name_to_string_array_field_map, existing_index.name_to_string_array_field_map)):
             logger.debug(f'Another thread has updated the index {marqo_index.name} already.')
             return
 

src/marqo/core/index_management/vespa_application_package.py

@@ -693,15 +693,20 @@ def rollback(self, marqo_version: str) -> None:
         self._store.save_file(new_backup.to_zip_stream().read(), self._BACKUP_FILE)
         self._deploy()
 
-    def batch_add_index_setting_and_schema(self, indexes: List[Tuple[str, MarqoIndex]]) -> None:
-        for schema, index in indexes:
+    def batch_add_index_setting_and_schema(self, indexes: List[Tuple[str, str, MarqoIndex]]) -> None:
+        for schema, typeahead_schema, index in indexes:
             if self.has_index(index.name):
                 raise IndexExistsError(f"Index {index.name} already exists")
 
+            # Add index settings and schema
             self._index_setting_store.save_index_setting(index)
             self._store.save_file(schema, 'schemas', f'{index.schema_name}.sd')
             self._service_xml.add_schema(index.schema_name)
 
+            # Add typeahead schema if provided
+            self._store.save_file(typeahead_schema, 'schemas', f'{index.typeahead_schema_name}.sd')
+            self._service_xml.add_schema(index.typeahead_schema_name)
+
         self._persist_index_settings()
         self._store.save_file(self._service_xml.to_xml(), self._SERVICES_XML_FILE)
         self._deploy()
@@ -714,6 +719,9 @@ def batch_delete_index_setting_and_schema(self, index_names: List[str]) -> None:
             self._index_setting_store.delete_index_setting(index.name)
             self._store.remove_file('schemas', f'{index.schema_name}.sd')
             self._service_xml.remove_schema(index.schema_name)
+            if index.typeahead_schema_name is not None:
+                self._store.remove_file('schemas', f'{index.typeahead_schema_name}.sd')
+                self._service_xml.remove_schema(index.typeahead_schema_name)
 
         self._add_schema_removal_override()
         self._persist_index_settings()

src/marqo/core/models/marqo_index.py

@@ -287,6 +287,7 @@ class MarqoIndex(ImmutableBaseModel, ABC):
     """
     name: str
     schema_name: str
+    typeahead_schema_name: Optional[str] = None
     type: IndexType  # We need this so that we can deserialize the correct subclass
     model: Model
     normalize_embeddings: bool

src/marqo/core/models/typeahead.py

@@ -0,0 +1,129 @@
+"""Pydantic models for typeahead API requests and responses."""
+
+from typing import List, Optional, Dict
+
+from pydantic import Field, field_validator
+
+from marqo.base_model import ImmutableStrictBaseModelV2
+from marqo.core.exceptions import InvalidArgumentError
+from marqo.tensor_search.enums import EnvVars
+from marqo.tensor_search.utils import read_env_vars_and_defaults_ints
+
+
+class TypeaheadRequest(ImmutableStrictBaseModelV2):
+    """Request model for typeahead suggestions."""
+
+    q: str = Field(..., description="Partial user search input")
+    limit: int = Field(default=10, ge=0, description="Maximum number of suggestions to return")
+    fuzzy_edit_distance: int = Field(
+        default=2,
+        ge=0,
+        alias="fuzzyEditDistance",
+        description="Maximum edit distance for fuzzy matching"
+    )
+    min_fuzzy_match_length: int = Field(
+        default=3,
+        ge=0,
+        alias="minFuzzyMatchLength",
+        description="Minimum length to switch to fuzzy matching"
+    )
+    popularity_weight: Optional[float] = Field(
+        default=None,
+        alias="popularityWeight",
+        description="Weight for popularity score in ranking"
+    )
+    bm25_weight: Optional[float] = Field(
+        default=None,
+        alias="bm25Weight",
+        description="Weight for BM25 score in ranking"
+    )
+
+    @field_validator('q')
+    def validate_q(cls, v: str) -> str:
+        if not v or not v.strip():
+            raise ValueError("q is required")
+        return v.strip()
+
+
+class TypeaheadSuggestion(ImmutableStrictBaseModelV2):
+    """Individual suggestion in typeahead response."""
+
+    suggestion: str = Field(..., description="The suggested query text")
+    score: float = Field(..., alias="_score", description="Relevance score for the suggestion")
+    metadata: Optional[dict] = Field(default=None, description="Additional metadata")
+
+
+class TypeaheadResponse(ImmutableStrictBaseModelV2):
+    """Response model for typeahead suggestions."""
+
+    suggestions: List[TypeaheadSuggestion] = Field(..., description="List of suggestions")
+    processing_time_ms: Optional[float] = Field(
+        default=None,
+        alias="processingTimeMs",
+        description="Processing time in milliseconds"
+    )
+
+
+class TypeaheadAddQueryRequest(ImmutableStrictBaseModelV2):
+    query: str = Field(..., description="User search query")
+    # Please note that popularity is not mandatory. This is to support multiple popularity values in metadata for future
+    popularity: float = Field(default=0.0, description="Popularity score")
+    metadata: Dict[str, float] = Field(default_factory=dict, description="Additional metadata")
+
+    @field_validator('query')
+    def validate_q(cls, v: str) -> str:
+        if not v or not v.strip():
+            raise ValueError("query is required")
+        return v.strip()
+
+
+class TypeaheadIndexRequest(ImmutableStrictBaseModelV2):
+    queries: List[TypeaheadAddQueryRequest]
+
+    @field_validator('queries')
+    def validate_queries_batch_size(cls, queries):
+        query_count = len(queries)
+        max_queries = read_env_vars_and_defaults_ints(EnvVars.MARQO_MAX_DOCUMENTS_BATCH_SIZE)
+        if query_count == 0:
+            raise InvalidArgumentError("Received empty index queries request")
+        elif query_count > max_queries:
+            raise InvalidArgumentError(
+                f"Number of queries in index request ({query_count}) exceeds limit of {max_queries}. "
+                f"Please break up your request into smaller batches."
+            )
+        return queries
+
+
+class TypeaheadIndexError(ImmutableStrictBaseModelV2):
+    query: Optional[str] = None
+    message: str
+    code: int = 400
+
+
+class TypeaheadIndexResponse(ImmutableStrictBaseModelV2):
+    indexed: int = Field(..., description="Indexed queries")
+    errors: List[TypeaheadIndexError] = Field(default_factory=list, description="Index Errors")
+    processing_time_ms: float = Field(
+        alias="processingTimeMs",
+        description="Processing time in milliseconds"
+    )
+
+
+class TypeaheadStatsResponse(ImmutableStrictBaseModelV2):
+    indexed_queries: int = Field(
+        alias="indexedQueries",
+        description="Number of indexed queries"
+    )
+
+
+class TypeaheadQuery(ImmutableStrictBaseModelV2):
+    """Represents a query from the typeahead schema."""
+    query: str = Field(..., description="The query string")
+    popularity: float = Field(..., description="Popularity score")
+    metadata: Dict[str, float] = Field(..., description="Additional metadata")
+    last_updated_at: Optional[int] = Field(None, alias="lastUpdatedAt", description="Last updated timestamp")
+
+
+class TypeaheadGetQueriesResponse(ImmutableStrictBaseModelV2):
+    """Response model for getting typeahead queries."""
+    queries: List[TypeaheadQuery] = Field(..., description="List of retrieved queries")

src/marqo/core/search/hybrid_search.py

@@ -16,7 +16,6 @@
 from marqo.core.vespa_index.vespa_index import for_marqo_index as vespa_index_factory
 from marqo.core.structured_vespa_index.common import RANK_PROFILE_HYBRID_CUSTOM_SEARCHER
 from marqo.core.models.interpolation_method import InterpolationMethod
-from marqo.tensor_search import index_meta_cache
 from marqo.tensor_search import utils
 from marqo.tensor_search.enums import (
     SearchMethod
@@ -28,7 +27,6 @@
 from marqo.tensor_search.telemetry import RequestMetricsStore
 from marqo.tensor_search.tensor_search import run_vectorise_pipeline, gather_documents_from_response, logger
 from marqo.vespa.exceptions import VespaStatusError
-import semver
 from marqo.tensor_search.models.sort_by_model import SortByModel
 from marqo.tensor_search.models.relevance_cutoff_model import RelevanceCutoffModel
 

src/marqo/core/typeahead/text_normalization.py

@@ -0,0 +1,35 @@
+import unicodedata
+from typing import List
+
+
+def normalize_text(text: str) -> str:
+    """
+    Normalize text by removing accents and converting to lowercase.
+
+    Args:
+        text: Input text to normalize
+
+    Returns:
+        Normalized text with accents removed and lowercased
+    """
+    if not text:
+        return ""
+
+    # Normalize to NFKD form and remove accents
+    normalized = unicodedata.normalize('NFKD', text)
+    # Filter out combining characters (accents)
+    without_accents = ''.join(c for c in normalized if not unicodedata.combining(c))
+    # Convert to lowercase
+    return without_accents.lower()
+
+
+def generate_prefixes(text: str) -> List[str]:
+    result = []
+    prefix = ""
+    for ch in text:
+        if ch.isspace():
+            prefix = ""   # reset when hitting whitespace
+        else:
+            prefix += ch
+            result.append(prefix)
+    return result