nucliadb 6.7.2.post4874__py3-none-any.whl → 6.10.0.post5705__py3-none-any.whl

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

@@ -17,7 +17,8 @@
 # You should have received a copy of the GNU Affero General Public License
 # along with this program. If not, see <http://www.gnu.org/licenses/>.
 #
-from typing import Optional
+
+from typing_extensions import assert_never
 
 from nucliadb.search.search.query_parser.fetcher import Fetcher
 from nucliadb.search.search.query_parser.models import (
@@ -26,7 +27,7 @@ from nucliadb.search.search.query_parser.models import (
 from nucliadb_models.search import AskRequest, MaxTokens
 
 
-async def parse_ask(kbid: str, item: AskRequest, *, fetcher: Optional[Fetcher] = None) -> Generation:
+async def parse_ask(kbid: str, item: AskRequest, *, fetcher: Fetcher | None = None) -> Generation:
     fetcher = fetcher or fetcher_for_ask(kbid, item)
     parser = _AskParser(kbid, item, fetcher)
     return await parser.parse()
@@ -63,10 +64,8 @@ class _AskParser:
             )
         elif isinstance(self.item.max_tokens, MaxTokens):
             max_tokens = self.item.max_tokens
-        else:  # pragma: nocover
-            # This is a trick so mypy generates an error if this branch can be reached,
-            # that is, if we are missing some ifs
-            _a: int = "a"
+        else:  # pragma: no cover
+            assert_never(self.item.max_tokens)
 
         max_context_tokens = await self.fetcher.get_max_context_tokens(max_tokens)
         max_answer_tokens = self.fetcher.get_max_answer_tokens(max_tokens)

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

@@ -18,14 +18,13 @@
 # along with this program. If not, see <http://www.gnu.org/licenses/>.
 #
 
+from typing_extensions import assert_never
+
 from nucliadb.common import datamanagers
+from nucliadb.common.catalog.interface import CatalogExpression, CatalogQuery
 from nucliadb.common.exceptions import InvalidQueryError
-from nucliadb.common.filter_expression import FacetFilterTypes, facet_from_filter
+from nucliadb.common.filter_expression import FacetFilter, facet_from_filter
 from nucliadb.search.search.filters import translate_label
-from nucliadb.search.search.query_parser.models import (
-    CatalogExpression,
-    CatalogQuery,
-)
 from nucliadb_models import search as search_models
 from nucliadb_models.filters import (
     And,
@@ -81,7 +80,6 @@ async def parse_catalog(kbid: str, item: search_models.CatalogRequest) -> Catalo
         sort = SortOptions(
             field=SortField.CREATED,
             order=SortOrder.DESC,
-            limit=None,
         )
 
     if isinstance(item.query, search_models.CatalogQuery):
@@ -185,18 +183,16 @@ async def parse_filter_expression(expr: ResourceFilterExpression, kbid: str) ->
             if rid is None:
                 raise InvalidQueryError("slug", f"Cannot find slug {expr.slug}")
             cat.resource_id = rid
-        else:  # pragma: nocover
+        else:  # pragma: no cover
             # Cannot happen due to model validation
             raise ValueError("Resource needs id or slug")
     elif isinstance(expr, DateCreated):
         cat.date = CatalogExpression.Date(field="created_at", since=expr.since, until=expr.until)
     elif isinstance(expr, DateModified):
         cat.date = CatalogExpression.Date(field="modified_at", since=expr.since, until=expr.until)
-    elif isinstance(expr, FacetFilterTypes):
+    elif isinstance(expr, FacetFilter):
         cat.facet = facet_from_filter(expr)
     else:
-        # This is a trick so mypy generates an error if this branch can be reached,
-        # that is, if we are missing some ifs
-        _a: int = "a"
+        assert_never(expr)
 
     return cat

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

@@ -19,9 +19,7 @@
 #
 import re
 import string
-from typing import Optional, Union
 
-from nucliadb.common.exceptions import InvalidQueryError
 from nucliadb.search import logger
 from nucliadb.search.search.query_parser.fetcher import Fetcher
 from nucliadb.search.search.query_parser.models import (
@@ -32,15 +30,20 @@ from nucliadb_models import search as search_models
 
 DEFAULT_GENERIC_SEMANTIC_THRESHOLD = 0.7
 
-# -* is an invalid query in tantivy and it won't return results but if you add some whitespaces
-# between - and *, it will actually trigger a tantivy bug and panic
-INVALID_QUERY = re.compile(r"- +\*")
 
+def validate_query_syntax(query: str) -> str:
+    """Filter some queries that panic tantivy, better than returning the 500"""
 
-def validate_query_syntax(query: str):
-    # Filter some queries that panic tantivy, better than returning the 500
+    # -* is an invalid query in tantivy and it won't return results but if you add some whitespaces
+    # between - and *, it will actually trigger a tantivy bug and panic
+    INVALID_QUERY = re.compile(r"- *\*+")
     if INVALID_QUERY.search(query):
-        raise InvalidQueryError("query", "Invalid query syntax")
+        # remove the * and extra spaces, as it's probably what doesn't have
+        # meaning in both cases: -* and - *
+        fixed = re.sub(INVALID_QUERY, "- ", query)
+        query = fixed
+
+    return query
 
 
 def is_empty_query(request: search_models.BaseSearchRequest) -> bool:
@@ -85,6 +88,7 @@ async def parse_keyword_query(
     fetcher: Fetcher,
 ) -> KeywordQuery:
     query = item.query
+
     # If there was a rephrase with image, we should use the rephrased query for keyword search
     rephrased_query = await fetcher.get_rephrased_query()
     if item.query_image is not None and rephrased_query is not None:
@@ -98,6 +102,10 @@ async def parse_keyword_query(
             query = synonyms_query
             is_synonyms_query = True
 
+    # after all query transformations, pass a validator that can fix some
+    # queries that trigger a panic on the index
+    query = validate_query_syntax(query)
+
     min_score = parse_keyword_min_score(item.min_score)
 
     return KeywordQuery(
@@ -108,7 +116,7 @@ async def parse_keyword_query(
 
 
 async def parse_semantic_query(
-    item: Union[search_models.SearchRequest, search_models.FindRequest],
+    item: search_models.SearchRequest | search_models.FindRequest,
     *,
     fetcher: Fetcher,
 ) -> SemanticQuery:
@@ -121,7 +129,7 @@ async def parse_semantic_query(
 
 
 def parse_keyword_min_score(
-    min_score: Optional[Union[float, search_models.MinScore]],
+    min_score: float | search_models.MinScore | None,
 ) -> float:
     # Keep backward compatibility with the deprecated min_score payload
     # parameter being a float (specifying semantic)
@@ -132,7 +140,7 @@ def parse_keyword_min_score(
 
 
 async def parse_semantic_min_score(
-    min_score: Optional[Union[float, search_models.MinScore]],
+    min_score: float | search_models.MinScore | None,
     *,
     fetcher: Fetcher,
 ):
@@ -161,7 +169,7 @@ async def query_with_synonyms(
     query: str,
     *,
     fetcher: Fetcher,
-) -> Optional[str]:
+) -> str | None:
     """
     Replace the terms in the query with an expression that will make it match with the configured synonyms.
     We're using the Tantivy's query language here: https://docs.rs/tantivy/latest/tantivy/query/struct.QueryParser.html
@@ -183,7 +191,7 @@ async def query_with_synonyms(
     variants: dict[str, str] = {}
     for term, term_synonyms in synonyms.terms.items():
         if len(term_synonyms.synonyms) > 0:
-            variants[term] = "({})".format(" OR ".join([term] + list(term_synonyms.synonyms)))
+            variants[term] = "({})".format(" OR ".join([term, *list(term_synonyms.synonyms)]))
 
     # Split the query into terms
     query_terms = query.split()

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

@@ -18,7 +18,6 @@
 # along with this program. If not, see <http://www.gnu.org/licenses/>.
 #
 
-from typing import Optional
 
 from nidx_protos import nodereader_pb2
 from pydantic import ValidationError
@@ -27,7 +26,6 @@ from nucliadb.common.exceptions import InvalidQueryError
 from nucliadb.common.filter_expression import parse_expression
 from nucliadb.common.models_utils.from_proto import RelationNodeTypeMap
 from nucliadb.search.search.metrics import query_parser_observer
-from nucliadb.search.search.query import expand_entities
 from nucliadb.search.search.query_parser.exceptions import InternalParserError
 from nucliadb.search.search.query_parser.fetcher import Fetcher
 from nucliadb.search.search.query_parser.models import (
@@ -57,7 +55,6 @@ from .common import (
     parse_semantic_query,
     parse_top_k,
     should_disable_vector_search,
-    validate_query_syntax,
 )
 
 
@@ -66,7 +63,7 @@ async def parse_find(
     kbid: str,
     item: FindRequest,
     *,
-    fetcher: Optional[Fetcher] = None,
+    fetcher: Fetcher | None = None,
 ) -> ParsedQuery:
     fetcher = fetcher or fetcher_for_find(kbid, item)
     parser = _FindParser(kbid, item, fetcher)
@@ -94,8 +91,8 @@ class _FindParser:
         self.fetcher = fetcher
 
         # cached data while parsing
-        self._query: Optional[Query] = None
-        self._top_k: Optional[int] = None
+        self._query: Query | None = None
+        self._top_k: int | None = None
 
     async def parse(self) -> UnitRetrieval:
         self._validate_request()
@@ -123,11 +120,11 @@ class _FindParser:
         try:
             rank_fusion = self._parse_rank_fusion()
         except ValidationError as exc:
-            raise InternalParserError(f"Parsing error in rank fusion: {str(exc)}") from exc
+            raise InternalParserError(f"Parsing error in rank fusion: {exc!s}") from exc
         try:
             reranker = self._parse_reranker()
         except ValidationError as exc:
-            raise InternalParserError(f"Parsing error in reranker: {str(exc)}") from exc
+            raise InternalParserError(f"Parsing error in reranker: {exc!s}") from exc
 
         # Adjust retrieval windows. Our current implementation assume:
         # `top_k <= reranker.window <= rank_fusion.window`
@@ -146,8 +143,6 @@ class _FindParser:
         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
@@ -173,15 +168,8 @@ class _FindParser:
     async def _parse_relation_query(self) -> RelationQuery:
         detected_entities = await self._get_detected_entities()
 
-        deleted_entity_groups = await self.fetcher.get_deleted_entity_groups()
-
-        meta_cache = await self.fetcher.get_entities_meta_cache()
-        deleted_entities = meta_cache.deleted_entities
-
         return RelationQuery(
-            entry_points=detected_entities,
-            deleted_entity_groups=deleted_entity_groups,
-            deleted_entities=deleted_entities,
+            entry_points=detected_entities, deleted_entity_groups=[], deleted_entities={}
         )
 
     async def _parse_graph_query(self) -> GraphQuery:
@@ -208,9 +196,6 @@ class _FindParser:
         else:
             detected_entities = await self.fetcher.get_detected_entities()
 
-        meta_cache = await self.fetcher.get_entities_meta_cache()
-        detected_entities = expand_entities(meta_cache, detected_entities)
-
         return detected_entities
 
     async def _parse_filters(self) -> Filters:
@@ -256,17 +241,9 @@ class _FindParser:
             else:
                 filter_operator = nodereader_pb2.FilterOperator.AND
 
-        autofilter = None
-        if self.item.autofilter:
-            if self._query.relation is not None:
-                autofilter = self._query.relation.entry_points
-            else:
-                autofilter = await self._get_detected_entities()
-
         hidden = await filter_hidden_resources(self.kbid, self.item.show_hidden)
 
         return Filters(
-            autofilter=autofilter,
             facets=[],
             field_expression=field_expr,
             paragraph_expression=paragraph_expr,

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

@@ -18,9 +18,9 @@
 # along with this program. If not, see <http://www.gnu.org/licenses/>.
 #
 
-from typing import Optional, Union
 
 from nidx_protos import nodereader_pb2
+from typing_extensions import assert_never
 
 from nucliadb.common.filter_expression import add_and_expression, parse_expression
 from nucliadb.common.models_utils.from_proto import RelationNodeTypeMap, RelationTypeMap
@@ -56,11 +56,11 @@ async def parse_graph_relation_search(
     return pb
 
 
-AnyGraphRequest = Union[
-    graph_requests.GraphSearchRequest,
-    graph_requests.GraphNodesSearchRequest,
-    graph_requests.GraphRelationsSearchRequest,
-]
+AnyGraphRequest = (
+    graph_requests.GraphSearchRequest
+    | graph_requests.GraphNodesSearchRequest
+    | graph_requests.GraphRelationsSearchRequest
+)
 
 
 async def _parse_common(kbid: str, item: AnyGraphRequest) -> nodereader_pb2.GraphSearchRequest:
@@ -78,7 +78,7 @@ async def _parse_common(kbid: str, item: AnyGraphRequest) -> nodereader_pb2.Grap
     return pb
 
 
-async def _parse_filters(kbid: str, item: AnyGraphRequest) -> Optional[nodereader_pb2.FilterExpression]:
+async def _parse_filters(kbid: str, item: AnyGraphRequest) -> nodereader_pb2.FilterExpression | None:
     filter_expr = nodereader_pb2.FilterExpression()
     if item.filter_expression:
         if item.filter_expression.field:
@@ -100,7 +100,7 @@ async def _parse_filters(kbid: str, item: AnyGraphRequest) -> Optional[nodereade
         return None
 
 
-def _parse_security(kbid: str, item: AnyGraphRequest) -> Optional[utils_pb2.Security]:
+def _parse_security(kbid: str, item: AnyGraphRequest) -> utils_pb2.Security | None:
     if item.security is not None and len(item.security.groups) > 0:
         security_pb = utils_pb2.Security()
         for group_id in item.security.groups:
@@ -153,10 +153,8 @@ def parse_path_query(expr: graph_requests.GraphPathQuery) -> nodereader_pb2.Grap
     elif isinstance(expr, graph_requests.Generated):
         _set_generated_to_pb(expr, pb)
 
-    else:  # pragma: nocover
-        # This is a trick so mypy generates an error if this branch can be reached,
-        # that is, if we are missing some ifs
-        _a: int = "a"
+    else:  # pragma: no cover
+        assert_never(expr)
 
     return pb
 
@@ -182,10 +180,8 @@ def _parse_node_query(expr: graph_requests.GraphNodesQuery) -> nodereader_pb2.Gr
     elif isinstance(expr, graph_requests.Generated):
         _set_generated_to_pb(expr, pb)
 
-    else:  # pragma: nocover
-        # This is a trick so mypy generates an error if this branch can be reached,
-        # that is, if we are missing some ifs
-        _a: int = "a"
+    else:  # pragma: no cover
+        assert_never(expr)
 
     return pb
 
@@ -212,10 +208,8 @@ def _parse_relation_query(
     elif isinstance(expr, graph_requests.Generated):
         _set_generated_to_pb(expr, pb)
 
-    else:  # pragma: nocover
-        # This is a trick so mypy generates an error if this branch can be reached,
-        # that is, if we are missing some ifs
-        _a: int = "a"
+    else:  # pragma: no cover
+        assert_never(expr)
 
     return pb
 
@@ -230,10 +224,8 @@ def _set_node_to_pb(node: graph_requests.GraphNode, pb: nodereader_pb2.GraphQuer
             pb.fuzzy.kind = nodereader_pb2.GraphQuery.Node.MatchLocation.PREFIX
             pb.fuzzy.distance = 1
 
-        else:  # pragma: nocover
-            # This is a trick so mypy generates an error if this branch can be reached,
-            # that is, if we are missing some ifs
-            _a: int = "a"
+        else:  # pragma: no cover
+            assert_never(node.match)
 
     if node.type is not None:
         pb.node_type = RelationNodeTypeMap[node.type]
@@ -263,7 +255,5 @@ def _set_generated_to_pb(generated: graph_requests.Generated, pb: nodereader_pb2
 
         pb.facet.facet = facet
 
-    else:  # pragma: nocover
-        # This is a trick so mypy generates an error if this branch can be reached,
-        # that is, if we are missing some ifs
-        _a: int = "a"
+    else:  # pragma: no cover
+        assert_never(generated.by)

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

@@ -0,0 +1,207 @@
+# Copyright (C) 2021 Bosutech XXI S.L.
+#
+# nucliadb is offered under the AGPL v3.0 and as commercial software.
+# For commercial licensing, contact us at info@nuclia.com.
+#
+# AGPL:
+# This program is free software: you can redistribute it and/or modify
+# it under the terms of the GNU Affero General Public License as
+# published by the Free Software Foundation, either version 3 of the
+# License, or (at your option) any later version.
+#
+# This program is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
+# GNU Affero General Public License for more details.
+#
+# You should have received a copy of the GNU Affero General Public License
+# along with this program. If not, see <http://www.gnu.org/licenses/>.
+#
+from nidx_protos import nodereader_pb2
+from pydantic import ValidationError
+
+from nucliadb.common.exceptions import InvalidQueryError
+from nucliadb.common.filter_expression import parse_expression
+from nucliadb.search.search.metrics import query_parser_observer
+from nucliadb.search.search.query_parser.exceptions import InternalParserError
+from nucliadb.search.search.query_parser.fetcher import Fetcher
+from nucliadb.search.search.query_parser.models import (
+    Filters,
+    GraphQuery,
+    KeywordQuery,
+    Query,
+    RankFusion,
+    ReciprocalRankFusion,
+    SemanticQuery,
+    UnitRetrieval,
+)
+from nucliadb.search.search.query_parser.parsers.common import query_with_synonyms, validate_query_syntax
+from nucliadb.search.search.utils import filter_hidden_resources
+from nucliadb_models import search as search_models
+from nucliadb_models.filters import FilterExpression
+from nucliadb_models.retrieval import RetrievalRequest
+from nucliadb_models.search import MAX_RANK_FUSION_WINDOW
+
+
+@query_parser_observer.wrap({"type": "parse_retrieve"})
+async def parse_retrieve(kbid: str, item: RetrievalRequest) -> UnitRetrieval:
+    fetcher = Fetcher(
+        kbid=kbid,
+        query=item.query.keyword.query if item.query.keyword else "",
+        user_vector=item.query.semantic.query if item.query.semantic else None,
+        vectorset=item.query.semantic.vectorset if item.query.semantic else None,
+        # Retrieve doesn't use images for now
+        query_image=None,
+        # Retrieve doesn't do rephrasing
+        rephrase=False,
+        rephrase_prompt=None,
+        generative_model=None,
+    )
+    parser = _RetrievalParser(kbid, item, fetcher)
+    retrieval = await parser.parse()
+    return retrieval
+
+
+class _RetrievalParser:
+    def __init__(self, kbid: str, item: RetrievalRequest, fetcher: Fetcher):
+        self.kbid = kbid
+        self.item = item
+        self.fetcher = fetcher
+
+    async def parse(self) -> UnitRetrieval:
+        top_k = self.item.top_k
+        query = await self._parse_query()
+        filters = await self._parse_filters()
+        try:
+            rank_fusion = self._parse_rank_fusion()
+        except ValidationError as exc:
+            raise InternalParserError(f"Parsing error in rank fusion: {exc!s}") from exc
+
+        # ensure top_k and rank_fusion are coherent
+        if top_k > rank_fusion.window:
+            raise InvalidQueryError(
+                "rank_fusion.window", "Rank fusion window must be greater or equal to top_k"
+            )
+
+        retrieval = UnitRetrieval(
+            query=query,
+            top_k=top_k,
+            filters=filters,
+            rank_fusion=rank_fusion,
+            reranker=None,
+        )
+        return retrieval
+
+    async def _parse_query(self) -> Query:
+        keyword = None
+        if self.item.query.keyword is not None:
+            keyword_query, is_synonyms_query = await self._parse_keyword_query()
+            keyword = KeywordQuery(
+                query=keyword_query,
+                is_synonyms_query=is_synonyms_query,
+                min_score=self.item.query.keyword.min_score,
+            )
+
+        semantic = None
+        if self.item.query.semantic is not None:
+            vectorset, query_vector = await self._parse_semantic_query()
+            semantic = SemanticQuery(
+                query=query_vector,
+                vectorset=vectorset,
+                min_score=self.item.query.semantic.min_score,
+            )
+
+        graph = None
+        if self.item.query.graph is not None:
+            graph = GraphQuery(query=self.item.query.graph.query)
+
+        return Query(keyword=keyword, semantic=semantic, graph=graph)
+
+    async def _parse_keyword_query(self) -> tuple[str, bool]:
+        assert self.item.query.keyword is not None
+        keyword_query = self.item.query.keyword.query
+        is_synonyms_query = False
+        if self.item.query.keyword.with_synonyms:
+            synonyms_query = await query_with_synonyms(keyword_query, fetcher=self.fetcher)
+            if synonyms_query is not None:
+                keyword_query = synonyms_query
+                is_synonyms_query = True
+
+        # after all query transformations, pass a validator that can fix some
+        # queries that trigger a panic on the index
+        keyword_query = validate_query_syntax(keyword_query)
+        return keyword_query, is_synonyms_query
+
+    async def _parse_semantic_query(self) -> tuple[str, list[float]]:
+        # Make sure the vectorset exists in the KB
+        assert self.item.query.semantic is not None
+        vectorset = self.item.query.semantic.vectorset
+        await self.fetcher.validate_vectorset(self.kbid, vectorset)
+
+        # Calculate the matryoshka dimension if applicable
+        user_vector = self.item.query.semantic.query
+        matryoshka_dimension = await self.fetcher.get_matryoshka_dimension_cached(self.kbid, vectorset)
+        if matryoshka_dimension is not None:
+            if len(user_vector) < matryoshka_dimension:
+                raise InvalidQueryError(
+                    "vector",
+                    f"Invalid vector length, please check valid embedding size for {vectorset} model",
+                )
+
+            # KB using a matryoshka embeddings model, cut the query vector
+            # accordingly
+            query_vector = user_vector[:matryoshka_dimension]
+        return vectorset, query_vector
+
+    async def _parse_filters(self) -> Filters:
+        filters = Filters()
+        if self.item.filters is None:
+            return filters
+
+        if self.item.filters.filter_expression is not None:
+            if self.item.filters.filter_expression.field is not None:
+                filters.field_expression = await parse_expression(
+                    self.item.filters.filter_expression.field,
+                    self.kbid,
+                )
+            if self.item.filters.filter_expression.paragraph is not None:
+                filters.paragraph_expression = await parse_expression(
+                    self.item.filters.filter_expression.paragraph,
+                    self.kbid,
+                )
+            if self.item.filters.filter_expression.operator == FilterExpression.Operator.OR:
+                filter_operator = nodereader_pb2.FilterOperator.OR
+            else:
+                filter_operator = nodereader_pb2.FilterOperator.AND
+            filters.filter_expression_operator = filter_operator
+
+        filters.hidden = await filter_hidden_resources(self.kbid, self.item.filters.show_hidden)
+        filters.security = self.item.filters.security
+        filters.with_duplicates = self.item.filters.with_duplicates
+
+        return filters
+
+    def _parse_rank_fusion(self) -> RankFusion:
+        rank_fusion: RankFusion
+
+        top_k = self.item.top_k
+        window = min(top_k, MAX_RANK_FUSION_WINDOW)
+
+        if isinstance(self.item.rank_fusion, search_models.RankFusionName):
+            if self.item.rank_fusion == search_models.RankFusionName.RECIPROCAL_RANK_FUSION:
+                rank_fusion = ReciprocalRankFusion(window=window)
+            else:
+                raise InternalParserError(f"Unknown rank fusion algorithm: {self.item.rank_fusion}")
+
+        elif isinstance(self.item.rank_fusion, search_models.ReciprocalRankFusion):
+            user_window = self.item.rank_fusion.window
+            rank_fusion = ReciprocalRankFusion(
+                k=self.item.rank_fusion.k,
+                boosting=self.item.rank_fusion.boosting,
+                window=min(max(user_window or 0, top_k), 500),
+            )
+
+        else:
+            raise InternalParserError(f"Unknown rank fusion {self.item.rank_fusion}")
+
+        return rank_fusion