nucliadb 6.9.1.post5192__py3-none-any.whl → 6.10.0.post5705__py3-none-any.whl

nucliadb/search/search/chat/rpc.py

@@ -0,0 +1,85 @@
+# 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 nucliadb.common.ids import FieldId
+from nucliadb.search.search.metrics import Metrics
+from nucliadb_models.augment import AugmentRequest, AugmentResponse
+from nucliadb_models.labels import KnowledgeBoxLabels
+from nucliadb_models.retrieval import RetrievalRequest, RetrievalResponse
+from nucliadb_models.search import FindRequest, Image, KnowledgeboxFindResults, NucliaDBClientType
+
+
+# TODO(decoupled-ask): replace this for a sdk.find call when moving /ask to RAO
+async def find(
+    kbid: str,
+    item: FindRequest,
+    x_ndb_client: NucliaDBClientType,
+    x_nucliadb_user: str,
+    x_forwarded_for: str,
+    # REVIEW(decoupled-ask): once in an SDK metrics, we'll lose track of metrics
+    metrics: Metrics,
+) -> tuple[KnowledgeboxFindResults, bool]:
+    from nucliadb.search.search.find import find
+
+    results, incomplete, _ = await find(
+        kbid, item, x_ndb_client, x_nucliadb_user, x_forwarded_for, metrics
+    )
+    return results, incomplete
+
+
+# TODO(decoupled-ask): replace this for a sdk.retrieve call when moving /ask to RAO
+async def retrieve(
+    kbid: str,
+    item: RetrievalRequest,
+    *,
+    x_ndb_client: NucliaDBClientType,
+    x_nucliadb_user: str,
+    x_forwarded_for: str,
+) -> RetrievalResponse:
+    from nucliadb.search.api.v1.retrieve import retrieve_endpoint
+
+    return await retrieve_endpoint(
+        kbid,
+        item,
+        x_ndb_client=x_ndb_client,
+        x_nucliadb_user=x_nucliadb_user,
+        x_forwarded_for=x_forwarded_for,
+    )
+
+
+# TODO(decoupled-ask): replace this for a sdk.augment call when moving /ask to RAO
+async def augment(kbid: str, item: AugmentRequest) -> AugmentResponse:
+    from nucliadb.search.api.v1.augment import augment_endpoint
+
+    return await augment_endpoint(kbid, item)
+
+
+# TODO(decoupled-ask): replace this for a sdk.labelsets call when moving /ask to RAO
+async def labelsets(kbid: str) -> KnowledgeBoxLabels:
+    from nucliadb.reader.api.v1.services import get_labelsets
+
+    return await get_labelsets(kbid)
+
+
+# TODO(decoupled-ask): replace this for a sdk.download call when moving /ask to RAO
+async def download_image(kbid: str, field_id: FieldId, path: str, *, mime_type: str) -> Image | None:
+    from nucliadb.search.search.hydrator.images import download_image
+
+    return await download_image(kbid, field_id, path, mime_type=mime_type)

nucliadb/search/search/fetch.py

@@ -19,7 +19,6 @@
 #
 import asyncio
 from contextvars import ContextVar
-from typing import Optional
 
 from nidx_protos.nodereader_pb2 import DocumentResult, ParagraphResult
 
@@ -36,7 +35,7 @@ from nucliadb_protos.resources_pb2 import Paragraph
 from nucliadb_utils import const
 from nucliadb_utils.utilities import has_feature
 
-rcache: ContextVar[Optional[dict[str, ResourceORM]]] = ContextVar("rcache", default=None)
+rcache: ContextVar[dict[str, ResourceORM] | None] = ContextVar("rcache", default=None)
 
 
 async def fetch_resources(
@@ -79,7 +78,7 @@ async def fetch_resources(
 
 async def get_paragraph_from_resource(
     orm_resource: ResourceORM, result: ParagraphResult
-) -> Optional[Paragraph]:
+) -> Paragraph | None:
     _, field_type, field = result.field.split("/")
     field_type_int = FIELD_TYPE_STR_TO_PB[field_type]
     field_obj = await orm_resource.get_field(field, field_type_int, load=False)
@@ -144,7 +143,7 @@ async def get_labels_paragraph(result: ParagraphResult, kbid: str) -> list[str]:
 
 async def get_seconds_paragraph(
     result: ParagraphResult, kbid: str
-) -> Optional[tuple[list[int], list[int]]]:
+) -> tuple[list[int], list[int]] | None:
     orm_resource = await cache.get_resource(kbid, result.uuid)
 
     if orm_resource is None:

nucliadb/search/search/filters.py

@@ -18,7 +18,7 @@
 # along with this program. If not, see <http://www.gnu.org/licenses/>.
 #
 from collections.abc import Iterator
-from typing import Any, Optional, Union
+from typing import Any
 
 from nucliadb.common.exceptions import InvalidQueryError
 from nucliadb_models.labels import translate_alias_to_system_label
@@ -108,7 +108,7 @@ def split_labels_by_type(
 
 def is_paragraph_labelset_kind(labelset_id: str, classification_labels: knowledgebox_pb2.Labels) -> bool:
     try:
-        labelset: Optional[knowledgebox_pb2.LabelSet] = classification_labels.labelset.get(labelset_id)
+        labelset: knowledgebox_pb2.LabelSet | None = classification_labels.labelset.get(labelset_id)
         if labelset is None:
             return False
         return knowledgebox_pb2.LabelSet.LabelSetKind.PARAGRAPHS in labelset.kind
@@ -117,7 +117,7 @@ def is_paragraph_labelset_kind(labelset_id: str, classification_labels: knowledg
         return False
 
 
-def flatten_filter_literals(filters: Union[list[str], dict[str, Any]]) -> list[str]:
+def flatten_filter_literals(filters: list[str] | dict[str, Any]) -> list[str]:
     if isinstance(filters, list):
         return filters
     else:
@@ -130,20 +130,17 @@ def iter_filter_expression_literals(expression: dict[str, Any]) -> Iterator[str]
         return
 
     if "not" in expression:
-        for label in iter_filter_expression_literals(expression["not"]):
-            yield label
+        yield from iter_filter_expression_literals(expression["not"])
         return
 
     if "and" in expression:
         for and_term in expression["and"]:
-            for label in iter_filter_expression_literals(and_term):
-                yield label
+            yield from iter_filter_expression_literals(and_term)
         return
 
     if "or" in expression:
         for or_term in expression["or"]:
-            for label in iter_filter_expression_literals(or_term):
-                yield label
+            yield from iter_filter_expression_literals(or_term)
         return
 
 
@@ -151,7 +148,7 @@ def has_classification_label_filters(filters: list[str]) -> bool:
     return any(label.startswith(CLASSIFICATION_LABEL_PREFIX) for label in filters)
 
 
-def convert_to_node_filters(filters: Union[list[str], list[Filter]]) -> dict[str, Any]:
+def convert_to_node_filters(filters: list[str] | list[Filter]) -> dict[str, Any]:
     if len(filters) == 0:
         return {}
 
@@ -161,7 +158,7 @@ def convert_to_node_filters(filters: Union[list[str], list[Filter]]) -> dict[str
     return {"and": [convert_filter_to_node_schema(fltr) for fltr in filters]}
 
 
-def convert_filter_to_node_schema(fltr: Union[str, Filter]) -> dict[str, Any]:
+def convert_filter_to_node_schema(fltr: str | Filter) -> dict[str, Any]:
     if isinstance(fltr, str):
         return {"literal": fltr}
 

nucliadb/search/search/find.py

@@ -23,7 +23,6 @@ from time import time
 from nucliadb.common.external_index_providers.base import ExternalIndexManager
 from nucliadb.common.external_index_providers.manager import get_external_index_manager
 from nucliadb.common.models_utils import to_proto
-from nucliadb.search.requesters.utils import Method, nidx_query
 from nucliadb.search.search.find_merge import (
     build_find_response,
     compose_find_resources,
@@ -38,14 +37,16 @@ from nucliadb.search.search.metrics import (
 )
 from nucliadb.search.search.query_parser.models import ParsedQuery
 from nucliadb.search.search.query_parser.parsers import parse_find
-from nucliadb.search.search.query_parser.parsers.unit_retrieval import legacy_convert_retrieval_to_proto
-from nucliadb.search.search.rank_fusion import (
-    get_rank_fusion,
+from nucliadb.search.search.query_parser.parsers.unit_retrieval import (
+    convert_retrieval_to_proto,
+    get_rephrased_query,
+    is_incomplete,
 )
 from nucliadb.search.search.rerankers import (
     RerankingOptions,
     get_reranker,
 )
+from nucliadb.search.search.retrieval import text_block_search
 from nucliadb.search.settings import settings
 from nucliadb_models.search import (
     FindRequest,
@@ -68,18 +69,16 @@ async def find(
 ) -> tuple[KnowledgeboxFindResults, bool, ParsedQuery]:
     external_index_manager = await get_external_index_manager(kbid=kbid)
     if external_index_manager is not None:
-        return await _external_index_retrieval(
+        return await _external_index_find(
             kbid,
             item,
             external_index_manager,
         )
     else:
-        return await _index_node_retrieval(
-            kbid, item, x_ndb_client, x_nucliadb_user, x_forwarded_for, metrics
-        )
+        return await _ndb_index_find(kbid, item, x_ndb_client, x_nucliadb_user, x_forwarded_for, metrics)
 
 
-async def _index_node_retrieval(
+async def _ndb_index_find(
     kbid: str,
     item: FindRequest,
     x_ndb_client: NucliaDBClientType,
@@ -95,32 +94,37 @@ async def _index_node_retrieval(
         assert parsed.retrieval.rank_fusion is not None and parsed.retrieval.reranker is not None, (
             "find parser must provide rank fusion and reranker algorithms"
         )
-        rank_fusion = get_rank_fusion(parsed.retrieval.rank_fusion)
         reranker = get_reranker(parsed.retrieval.reranker)
-        (
-            pb_query,
-            incomplete_results,
-            autofilters,
-            rephrased_query,
-        ) = await legacy_convert_retrieval_to_proto(parsed)
+        incomplete_results = is_incomplete(parsed.retrieval)
+        rephrased_query = get_rephrased_query(parsed)
 
     with metrics.time("index_search"):
-        results, queried_shards = await nidx_query(kbid, Method.SEARCH, pb_query)
+        text_blocks, pb_query, pb_response, queried_shards = await text_block_search(
+            kbid, parsed.retrieval
+        )
 
     # Rank fusion merge, cut, hydrate and rerank
     with metrics.time("results_merge"):
-        search_results = await build_find_response(
-            results,
-            retrieval=parsed.retrieval,
-            kbid=kbid,
-            query=pb_query.body,
-            rephrased_query=rephrased_query,
+        resource_hydration_options = ResourceHydrationOptions(
             show=item.show,
             extracted=item.extracted,
             field_type_filter=item.field_type_filter,
+        )
+        text_block_hydration_options = TextBlockHydrationOptions(
             highlight=item.highlight,
-            rank_fusion_algorithm=rank_fusion,
+            ematches=pb_response.paragraph.ematches,  # type: ignore
+        )
+        search_results = await build_find_response(
+            pb_response,
+            text_blocks,
+            pb_response.graph,
+            retrieval=parsed.retrieval,
+            kbid=kbid,
+            query=item.query,
+            rephrased_query=rephrased_query,
             reranker=reranker,
+            resource_hydration_options=resource_hydration_options,
+            text_block_hydration_options=text_block_hydration_options,
         )
 
     search_time = time() - start_time
@@ -137,7 +141,6 @@ async def _index_node_retrieval(
         )
 
     search_results.shards = queried_shards
-    search_results.autofilters = autofilters
 
     ndb_time = metrics["index_search"] + metrics["results_merge"]
     if metrics["index_search"] > settings.slow_node_query_log_threshold:
@@ -168,7 +171,7 @@ async def _index_node_retrieval(
     return search_results, incomplete_results, parsed
 
 
-async def _external_index_retrieval(
+async def _external_index_find(
     kbid: str,
     item: FindRequest,
     external_index_manager: ExternalIndexManager,
@@ -180,12 +183,12 @@ async def _external_index_retrieval(
     parsed = await parse_find(kbid, item)
     assert parsed.retrieval.reranker is not None, "find parser must provide a reranking algorithm"
     reranker = get_reranker(parsed.retrieval.reranker)
-    search_request, incomplete_results, _, rephrased_query = await legacy_convert_retrieval_to_proto(
-        parsed
-    )
+    incomplete_results = is_incomplete(parsed.retrieval)
+    rephrased_query = get_rephrased_query(parsed)
+    search_request = convert_retrieval_to_proto(parsed.retrieval)
 
     # Query index
-    query_results = await external_index_manager.query(search_request)  # noqa
+    query_results = await external_index_manager.query(search_request)
 
     # Hydrate and rerank results
     text_blocks, resources, best_matches = await hydrate_and_rerank(
@@ -220,7 +223,6 @@ async def _external_index_retrieval(
         page_number=0,
         page_size=item.top_k,
         relations=None,  # Not implemented for external indexes yet
-        autofilters=[],  # Not implemented for external indexes yet
         min_score=results_min_score,
         best_matches=best_matches,
         # These are not used for external indexes