nucliadb 6.3.7.post4081__py3-none-any.whl → 6.3.7.post4114__py3-none-any.whl

nucliadb/search/search/query_parser/models.py

@@ -21,10 +21,7 @@ from dataclasses import dataclass
 from datetime import datetime
 from typing import Literal, Optional, Union
 
-from pydantic import (
-    BaseModel,
-    Field,
-)
+from pydantic import BaseModel, ConfigDict, Field
 
 from nucliadb.search.search.query_parser.fetcher import Fetcher
 from nucliadb_models import search as search_models
@@ -35,8 +32,7 @@ from nucliadb_protos import nodereader_pb2, utils_pb2
 # query
 
 
-@dataclass
-class _TextQuery:
+class _TextQuery(BaseModel):
     query: str
     is_synonyms_query: bool
     min_score: float
@@ -48,24 +44,23 @@ FulltextQuery = _TextQuery
 KeywordQuery = _TextQuery
 
 
-@dataclass
-class SemanticQuery:
+class SemanticQuery(BaseModel):
     query: Optional[list[float]]
     vectorset: str
     min_score: float
 
 
-@dataclass
-class RelationQuery:
-    detected_entities: list[utils_pb2.RelationNode]
+class RelationQuery(BaseModel):
+    model_config = ConfigDict(arbitrary_types_allowed=True)
+
+    entry_points: list[utils_pb2.RelationNode]
     # list[subtype]
     deleted_entity_groups: list[str]
     # subtype -> list[entity]
     deleted_entities: dict[str, list[str]]
 
 
-@dataclass
-class Query:
+class Query(BaseModel):
     fulltext: Optional[FulltextQuery] = None
     keyword: Optional[KeywordQuery] = None
     semantic: Optional[SemanticQuery] = None
@@ -75,8 +70,9 @@ class Query:
 # filters
 
 
-@dataclass
-class Filters:
+class Filters(BaseModel):
+    model_config = ConfigDict(arbitrary_types_allowed=True)
+
     field_expression: Optional[nodereader_pb2.FilterExpression] = None
     paragraph_expression: Optional[nodereader_pb2.FilterExpression] = None
     filter_expression_operator: nodereader_pb2.FilterOperator.ValueType = (
@@ -125,30 +121,29 @@ Reranker = Union[NoopReranker, PredictReranker]
 # retrieval and generation operations
 
 
-@dataclass
-class UnitRetrieval:
+class UnitRetrieval(BaseModel):
     query: Query
     top_k: int
-    filters: Filters
-    # TODO: rank fusion depends on the response building, not the retrieval
-    rank_fusion: RankFusion
-    # TODO: reranking fusion depends on the response building, not the retrieval
-    reranker: Reranker
+    filters: Filters = Field(default_factory=Filters)
+    rank_fusion: Optional[RankFusion] = None
+    reranker: Optional[Reranker] = None
 
 
-@dataclass
-class Generation:
+# TODO: augmentation things: hydration...
+
+
+class Generation(BaseModel):
     use_visual_llm: bool
     max_context_tokens: int
     max_answer_tokens: Optional[int]
 
 
-@dataclass
-class ParsedQuery:
+class ParsedQuery(BaseModel):
+    model_config = ConfigDict(arbitrary_types_allowed=True)
+
     fetcher: Fetcher
     retrieval: UnitRetrieval
     generation: Optional[Generation] = None
-    # TODO: add merge, rank fusion, rerank...
 
 
 ### Catalog

nucliadb/search/search/query_parser/parsers/common.py

@@ -28,7 +28,6 @@ from nucliadb.search.search.query_parser.models import (
     KeywordQuery,
     SemanticQuery,
 )
-from nucliadb.search.search.utils import should_disable_vector_search
 from nucliadb_models import search as search_models
 
 DEFAULT_GENERIC_SEMANTIC_THRESHOLD = 0.7
@@ -38,28 +37,40 @@ DEFAULT_GENERIC_SEMANTIC_THRESHOLD = 0.7
 INVALID_QUERY = re.compile(r"- +\*")
 
 
-def validate_base_request(item: search_models.BaseSearchRequest):
+def validate_query_syntax(query: str):
     # Filter some queries that panic tantivy, better than returning the 500
-    if INVALID_QUERY.search(item.query):
+    if INVALID_QUERY.search(query):
         raise InvalidQueryError("query", "Invalid query syntax")
 
-    # synonyms are not compatible with vector/graph search
-    if (
-        item.with_synonyms
-        and item.query
-        and (
-            search_models.SearchOptions.SEMANTIC in item.features
-            or search_models.SearchOptions.RELATIONS in item.features
-        )
-    ):
-        raise InvalidQueryError(
-            "synonyms",
-            "Search with custom synonyms is only supported on paragraph and document search",
-        )
-
-    if search_models.SearchOptions.SEMANTIC in item.features:
-        if should_disable_vector_search(item):
-            item.features.remove(search_models.SearchOptions.SEMANTIC)
+
+def is_empty_query(request: search_models.BaseSearchRequest) -> bool:
+    return len(request.query) == 0
+
+
+def has_user_vectors(request: search_models.BaseSearchRequest) -> bool:
+    return request.vector is not None and len(request.vector) > 0
+
+
+def is_exact_match_only_query(request: search_models.BaseSearchRequest) -> bool:
+    """
+    '"something"' -> True
+    'foo "something" else' -> False
+    """
+    query = request.query.strip()
+    return len(query) > 0 and query.startswith('"') and query.endswith('"')
+
+
+def should_disable_vector_search(request: search_models.BaseSearchRequest) -> bool:
+    if has_user_vectors(request):
+        return False
+
+    if is_exact_match_only_query(request):
+        return True
+
+    if is_empty_query(request):
+        return True
+
+    return False
 
 
 def parse_top_k(item: search_models.BaseSearchRequest) -> int:
@@ -92,7 +103,7 @@ async def parse_keyword_query(
 
 
 async def parse_semantic_query(
-    item: search_models.BaseSearchRequest,
+    item: Union[search_models.SearchRequest, search_models.FindRequest],
     *,
     fetcher: Fetcher,
 ) -> SemanticQuery:

nucliadb/search/search/query_parser/parsers/find.py

@@ -53,7 +53,8 @@ from .common import (
     parse_keyword_query,
     parse_semantic_query,
     parse_top_k,
-    validate_base_request,
+    should_disable_vector_search,
+    validate_query_syntax,
 )
 
 
@@ -93,7 +94,7 @@ class _FindParser:
         self._top_k: Optional[int] = None
 
     async def parse(self) -> UnitRetrieval:
-        validate_base_request(self.item)
+        self._validate_request()
 
         self._top_k = parse_top_k(self.item)
 
@@ -101,13 +102,13 @@ class _FindParser:
 
         self._query = Query()
 
-        if search_models.SearchOptions.KEYWORD in self.item.features:
+        if search_models.FindOptions.KEYWORD in self.item.features:
             self._query.keyword = await parse_keyword_query(self.item, fetcher=self.fetcher)
 
-        if search_models.SearchOptions.SEMANTIC in self.item.features:
+        if search_models.FindOptions.SEMANTIC in self.item.features:
             self._query.semantic = await parse_semantic_query(self.item, fetcher=self.fetcher)
 
-        if search_models.SearchOptions.RELATIONS in self.item.features:
+        if search_models.FindOptions.RELATIONS in self.item.features:
             self._query.relation = await self._parse_relation_query()
 
         # TODO: graph search
@@ -130,13 +131,35 @@ class _FindParser:
         if isinstance(reranker, PredictReranker):
             rank_fusion.window = max(rank_fusion.window, reranker.window)
 
-        return UnitRetrieval(
+        retrieval = UnitRetrieval(
             query=self._query,
             top_k=self._top_k,
             filters=filters,
             rank_fusion=rank_fusion,
             reranker=reranker,
         )
+        return retrieval
+
+    def _validate_request(self):
+        validate_query_syntax(self.item.query)
+
+        # synonyms are not compatible with vector/graph search
+        if (
+            self.item.with_synonyms
+            and self.item.query
+            and (
+                search_models.FindOptions.SEMANTIC in self.item.features
+                or search_models.FindOptions.RELATIONS in self.item.features
+            )
+        ):
+            raise InvalidQueryError(
+                "synonyms",
+                "Search with custom synonyms is only supported on paragraph and document search",
+            )
+
+        if search_models.FindOptions.SEMANTIC in self.item.features:
+            if should_disable_vector_search(self.item):
+                self.item.features.remove(search_models.FindOptions.SEMANTIC)
 
     async def _parse_relation_query(self) -> RelationQuery:
         detected_entities = await self._get_detected_entities()
@@ -147,7 +170,7 @@ class _FindParser:
         deleted_entities = meta_cache.deleted_entities
 
         return RelationQuery(
-            detected_entities=detected_entities,
+            entry_points=detected_entities,
             deleted_entity_groups=deleted_entity_groups,
             deleted_entities=deleted_entities,
         )
@@ -220,7 +243,7 @@ class _FindParser:
         autofilter = None
         if self.item.autofilter:
             if self._query.relation is not None:
-                autofilter = self._query.relation.detected_entities
+                autofilter = self._query.relation.entry_points
             else:
                 autofilter = await self._get_detected_entities()
 

nucliadb/search/search/query_parser/parsers/search.py

@@ -26,10 +26,8 @@ from nucliadb.search.search.query_parser.fetcher import Fetcher
 from nucliadb.search.search.query_parser.filter_expression import parse_expression
 from nucliadb.search.search.query_parser.models import (
     Filters,
-    NoopReranker,
     ParsedQuery,
     Query,
-    RankFusion,
     RelationQuery,
     UnitRetrieval,
     _TextQuery,
@@ -46,7 +44,13 @@ from nucliadb_models.search import (
 )
 from nucliadb_protos import nodereader_pb2, utils_pb2
 
-from .common import parse_keyword_query, parse_semantic_query, parse_top_k, validate_base_request
+from .common import (
+    parse_keyword_query,
+    parse_semantic_query,
+    parse_top_k,
+    should_disable_vector_search,
+    validate_query_syntax,
+)
 
 INDEX_SORTABLE_FIELDS = [
     SortField.CREATED,
@@ -87,7 +91,7 @@ class _SearchParser:
         self._top_k: Optional[int] = None
 
     async def parse(self) -> UnitRetrieval:
-        validate_base_request(self.item)
+        self._validate_request()
 
         self._top_k = parse_top_k(self.item)
 
@@ -113,14 +117,33 @@ class _SearchParser:
 
         filters = await self._parse_filters()
 
-        return UnitRetrieval(
+        retrieval = UnitRetrieval(
             query=self._query,
             top_k=self._top_k,
             filters=filters,
-            # TODO: this should be in a post retrieval step
-            rank_fusion=RankFusion(window=self._top_k),
-            reranker=NoopReranker(),
         )
+        return retrieval
+
+    def _validate_request(self):
+        validate_query_syntax(self.item.query)
+
+        # synonyms are not compatible with vector/graph search
+        if (
+            self.item.with_synonyms
+            and self.item.query
+            and (
+                search_models.SearchOptions.SEMANTIC in self.item.features
+                or search_models.SearchOptions.RELATIONS in self.item.features
+            )
+        ):
+            raise InvalidQueryError(
+                "synonyms",
+                "Search with custom synonyms is only supported on paragraph and document search",
+            )
+
+        if search_models.SearchOptions.SEMANTIC in self.item.features:
+            if should_disable_vector_search(self.item):
+                self.item.features.remove(search_models.SearchOptions.SEMANTIC)
 
     async def _parse_text_query(self) -> _TextQuery:
         assert self._top_k is not None, "top_k must be parsed before text query"
@@ -140,7 +163,7 @@ class _SearchParser:
         meta_cache = await self.fetcher.get_entities_meta_cache()
         deleted_entities = meta_cache.deleted_entities
         return RelationQuery(
-            detected_entities=detected_entities,
+            entry_points=detected_entities,
             deleted_entity_groups=deleted_entity_groups,
             deleted_entities=deleted_entities,
         )
@@ -231,7 +254,7 @@ class _SearchParser:
         autofilter = None
         if self.item.autofilter:
             if self._query.relation is not None:
-                autofilter = self._query.relation.detected_entities
+                autofilter = self._query.relation.entry_points
             else:
                 autofilter = await self._get_detected_entities()