nucliadb 2.46.1.post382__py3-none-any.whl → 6.2.1.post2777__py3-none-any.whl

nucliadb/search/search/fetch.py

@@ -17,26 +17,25 @@
 # 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/>.
 #
+import asyncio
 from contextvars import ContextVar
 from typing import Optional
 
-from nucliadb_protos.nodereader_pb2 import DocumentResult, ParagraphResult
-from nucliadb_protos.resources_pb2 import Paragraph
-
-from nucliadb.ingest.orm.resource import KB_REVERSE
+from nucliadb.common.ids import FIELD_TYPE_STR_TO_PB
+from nucliadb.common.maindb.utils import get_driver
 from nucliadb.ingest.orm.resource import Resource as ResourceORM
 from nucliadb.ingest.serialize import managed_serialize
-from nucliadb.middleware.transaction import get_read_only_transaction
 from nucliadb.search import SERVICE_NAME, logger
+from nucliadb.search.search import cache
 from nucliadb_models.common import FieldTypeName
 from nucliadb_models.resource import ExtractedDataTypeName, Resource
 from nucliadb_models.search import ResourceProperties
+from nucliadb_protos.nodereader_pb2 import DocumentResult, ParagraphResult
+from nucliadb_protos.resources_pb2 import Paragraph
+from nucliadb_utils import const
+from nucliadb_utils.utilities import has_feature
 
-from .cache import get_resource_from_cache
-
-rcache: ContextVar[Optional[dict[str, ResourceORM]]] = ContextVar(
-    "rcache", default=None
-)
+rcache: ContextVar[Optional[dict[str, ResourceORM]]] = ContextVar("rcache", default=None)
 
 
 async def fetch_resources(
@@ -46,20 +45,34 @@ async def fetch_resources(
     field_type_filter: list[FieldTypeName],
     extracted: list[ExtractedDataTypeName],
 ) -> dict[str, Resource]:
+    if ResourceProperties.EXTRACTED in show and has_feature(
+        const.Features.IGNORE_EXTRACTED_IN_SEARCH, context={"kbid": kbid}, default=False
+    ):
+        # Returning extracted metadata in search results is deprecated and this flag
+        # will be set to True for all KBs in the future.
+        show.remove(ResourceProperties.EXTRACTED)
+        extracted = []
+
     result = {}
-    txn = await get_read_only_transaction()
-    for resource in resources:
-        serialization = await managed_serialize(
-            txn,
-            kbid,
-            resource,
-            show,
-            field_type_filter=field_type_filter,
-            extracted=extracted,
-            service_name=SERVICE_NAME,
-        )
-        if serialization is not None:
-            result[resource] = serialization
+    async with get_driver().transaction(read_only=True) as txn:
+        tasks = []
+        for resource in resources:
+            tasks.append(
+                asyncio.create_task(
+                    managed_serialize(
+                        txn,
+                        kbid,
+                        resource,
+                        show,
+                        field_type_filter=field_type_filter,
+                        extracted=extracted,
+                        service_name=SERVICE_NAME,
+                    )
+                )
+            )
+        for resource, serialization in zip(resources, await asyncio.gather(*tasks)):
+            if serialization is not None:
+                result[resource] = serialization
     return result
 
 
@@ -67,7 +80,7 @@ async def get_paragraph_from_resource(
     orm_resource: ResourceORM, result: ParagraphResult
 ) -> Optional[Paragraph]:
     _, field_type, field = result.field.split("/")
-    field_type_int = KB_REVERSE[field_type]
+    field_type_int = FIELD_TYPE_STR_TO_PB[field_type]
     field_obj = await orm_resource.get_field(field, field_type_int, load=False)
     field_metadata = await field_obj.get_field_metadata()
     paragraph = None
@@ -81,7 +94,7 @@ async def get_paragraph_from_resource(
 
 
 async def get_labels_resource(result: DocumentResult, kbid: str) -> list[str]:
-    orm_resource = await get_resource_from_cache(kbid, result.uuid)
+    orm_resource = await cache.get_resource(kbid, result.uuid)
 
     if orm_resource is None:
         logger.error(f"{result.uuid} does not exist on DB")
@@ -97,7 +110,7 @@ async def get_labels_resource(result: DocumentResult, kbid: str) -> list[str]:
 
 
 async def get_labels_paragraph(result: ParagraphResult, kbid: str) -> list[str]:
-    orm_resource = await get_resource_from_cache(kbid, result.uuid)
+    orm_resource = await cache.get_resource(kbid, result.uuid)
 
     if orm_resource is None:
         logger.error(f"{result.uuid} does not exist on DB")
@@ -110,7 +123,7 @@ async def get_labels_paragraph(result: ParagraphResult, kbid: str) -> list[str]:
             labels.append(f"{classification.labelset}/{classification.label}")
 
     _, field_type, field = result.field.split("/")
-    field_type_int = KB_REVERSE[field_type]
+    field_type_int = FIELD_TYPE_STR_TO_PB[field_type]
     field_obj = await orm_resource.get_field(field, field_type_int, load=False)
     field_metadata = await field_obj.get_field_metadata()
     if field_metadata:
@@ -131,21 +144,15 @@ 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]]]:
-    orm_resource = await get_resource_from_cache(kbid, result.uuid)
+    orm_resource = await cache.get_resource(kbid, result.uuid)
 
     if orm_resource is None:
         logger.error(f"{result.uuid} does not exist on DB")
         return None
 
-    paragraph = await get_paragraph_from_resource(
-        orm_resource=orm_resource, result=result
-    )
+    paragraph = await get_paragraph_from_resource(orm_resource=orm_resource, result=result)
 
-    if (
-        paragraph is not None
-        and len(paragraph.end_seconds) > 0
-        and paragraph.end_seconds[0] > 0
-    ):
+    if paragraph is not None and len(paragraph.end_seconds) > 0 and paragraph.end_seconds[0] > 0:
         return (list(paragraph.start_seconds), list(paragraph.end_seconds))
 
     return None

nucliadb/search/search/filters.py

@@ -59,9 +59,7 @@ def translate_label(literal: str) -> str:
     if len(literal) == 0:
         raise InvalidQueryError("filters", "Invalid empty label")
     if literal[0] != "/":
-        raise InvalidQueryError(
-            "filters", f"Invalid label. It must start with a `/`: {literal}"
-        )
+        raise InvalidQueryError("filters", f"Invalid label. It must start with a `/`: {literal}")
     return translate_alias_to_system_label(literal)
 
 
@@ -109,13 +107,9 @@ def split_labels_by_type(
     return field_labels, paragraph_labels
 
 
-def is_paragraph_labelset_kind(
-    labelset_id: str, classification_labels: knowledgebox_pb2.Labels
-) -> bool:
+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: Optional[knowledgebox_pb2.LabelSet] = classification_labels.labelset.get(labelset_id)
         if labelset is None:
             return False
         return knowledgebox_pb2.LabelSet.LabelSetKind.PARAGRAPHS in labelset.kind
@@ -124,32 +118,32 @@ def is_paragraph_labelset_kind(
         return False
 
 
-def flat_filter_labels(filters: Union[list[str], dict[str, Any]]) -> list[str]:
+def flatten_filter_literals(filters: Union[list[str], dict[str, Any]]) -> list[str]:
     if isinstance(filters, list):
         return filters
     else:
-        return list(iter_filter_labels_expression(filters))
+        return list(iter_filter_expression_literals(filters))
 
 
-def iter_filter_labels_expression(expression: dict[str, Any]) -> Iterator[str]:
+def iter_filter_expression_literals(expression: dict[str, Any]) -> Iterator[str]:
     if "literal" in expression:
         yield expression["literal"]
         return
 
     if "not" in expression:
-        for label in iter_filter_labels_expression(expression["not"]):
+        for label in iter_filter_expression_literals(expression["not"]):
             yield label
         return
 
     if "and" in expression:
         for and_term in expression["and"]:
-            for label in iter_filter_labels_expression(and_term):
+            for label in iter_filter_expression_literals(and_term):
                 yield label
         return
 
     if "or" in expression:
         for or_term in expression["or"]:
-            for label in iter_filter_labels_expression(or_term):
+            for label in iter_filter_expression_literals(or_term):
                 yield label
         return
 

nucliadb/search/search/find.py

@@ -18,13 +18,39 @@
 # along with this program. If not, see <http://www.gnu.org/licenses/>.
 #
 import logging
+from dataclasses import dataclass
 from time import time
 from typing import Optional
 
+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, debug_nodes_info, node_query
-from nucliadb.search.search.find_merge import find_merge_results
+from nucliadb.search.search.find_merge import (
+    build_find_response,
+    compose_find_resources,
+    hydrate_and_rerank,
+)
+from nucliadb.search.search.hydrator import (
+    ResourceHydrationOptions,
+    TextBlockHydrationOptions,
+)
+from nucliadb.search.search.metrics import (
+    RAGMetrics,
+)
 from nucliadb.search.search.query import QueryParser
+from nucliadb.search.search.query_parser.parser import parse_find
+from nucliadb.search.search.rank_fusion import (
+    RankFusionAlgorithm,
+    get_rank_fusion,
+)
+from nucliadb.search.search.rerankers import (
+    Reranker,
+    RerankingOptions,
+    get_reranker,
+)
 from nucliadb.search.search.utils import (
+    filter_hidden_resources,
     min_score_from_payload,
     should_disable_vector_search,
 )
@@ -32,6 +58,7 @@ from nucliadb.search.settings import settings
 from nucliadb_models.search import (
     FindRequest,
     KnowledgeboxFindResults,
+    MinScore,
     NucliaDBClientType,
     SearchOptions,
 )
@@ -47,73 +74,76 @@ async def find(
     x_nucliadb_user: str,
     x_forwarded_for: str,
     generative_model: Optional[str] = None,
+    metrics: RAGMetrics = RAGMetrics(),
 ) -> tuple[KnowledgeboxFindResults, bool, QueryParser]:
-    audit = get_audit()
-    start_time = time()
+    external_index_manager = await get_external_index_manager(kbid=kbid)
+    if external_index_manager is not None:
+        return await _external_index_retrieval(
+            kbid,
+            item,
+            external_index_manager,
+            generative_model,
+        )
+    else:
+        return await _index_node_retrieval(
+            kbid, item, x_ndb_client, x_nucliadb_user, x_forwarded_for, generative_model, metrics
+        )
 
-    item.min_score = min_score_from_payload(item.min_score)
 
-    if SearchOptions.VECTOR in item.features:
-        if should_disable_vector_search(item):
-            item.features.remove(SearchOptions.VECTOR)
+async def _index_node_retrieval(
+    kbid: str,
+    item: FindRequest,
+    x_ndb_client: NucliaDBClientType,
+    x_nucliadb_user: str,
+    x_forwarded_for: str,
+    generative_model: Optional[str] = None,
+    metrics: RAGMetrics = RAGMetrics(),
+) -> tuple[KnowledgeboxFindResults, bool, QueryParser]:
+    audit = get_audit()
+    start_time = time()
 
-    query_parser = QueryParser(
-        kbid=kbid,
-        features=item.features,
-        query=item.query,
-        filters=item.filters,
-        faceted=None,
-        sort=None,
-        page_number=item.page_number,
-        page_size=item.page_size,
-        min_score=item.min_score,
-        range_creation_start=item.range_creation_start,
-        range_creation_end=item.range_creation_end,
-        range_modification_start=item.range_modification_start,
-        range_modification_end=item.range_modification_end,
-        fields=item.fields,
-        user_vector=item.vector,
-        vectorset=item.vectorset,
-        with_duplicates=item.with_duplicates,
-        with_synonyms=item.with_synonyms,
-        autofilter=item.autofilter,
-        key_filters=item.resource_filters,
-        security=item.security,
-        generative_model=generative_model,
-        rephrase=item.rephrase,
-    )
-    pb_query, incomplete_results, autofilters = await query_parser.parse()
-    results, query_incomplete_results, queried_nodes = await node_query(
-        kbid, Method.SEARCH, pb_query, target_shard_replicas=item.shards
+    query_parser, rank_fusion, reranker = await query_parser_from_find_request(
+        kbid, item, generative_model=generative_model
     )
+    with metrics.time("query_parse"):
+        pb_query, incomplete_results, autofilters = await query_parser.parse()
+
+    with metrics.time("node_query"):
+        results, query_incomplete_results, queried_nodes = await node_query(
+            kbid, Method.SEARCH, pb_query, target_shard_replicas=item.shards
+        )
     incomplete_results = incomplete_results or query_incomplete_results
 
-    # We need to merge
-    search_results = await find_merge_results(
-        results,
-        count=item.page_size,
-        page=item.page_number,
-        kbid=kbid,
-        show=item.show,
-        field_type_filter=item.field_type_filter,
-        extracted=item.extracted,
-        requested_relations=pb_query.relation_subgraph,
-        min_score_bm25=query_parser.min_score.bm25,
-        min_score_semantic=query_parser.min_score.semantic,
-        highlight=item.highlight,
-    )
+    # Rank fusion merge, cut, hydrate and rerank
+    with metrics.time("results_merge"):
+        search_results = await build_find_response(
+            results,
+            kbid=kbid,
+            query=pb_query.body,
+            relation_subgraph_query=pb_query.relations.subgraph,
+            min_score_bm25=pb_query.min_score_bm25,
+            min_score_semantic=pb_query.min_score_semantic,
+            top_k=item.top_k,
+            show=item.show,
+            extracted=item.extracted,
+            field_type_filter=item.field_type_filter,
+            highlight=item.highlight,
+            rank_fusion_algorithm=rank_fusion,
+            reranker=reranker,
+        )
 
     search_time = time() - start_time
     if audit is not None:
-        await audit.search(
+        audit.search(
             kbid,
             x_nucliadb_user,
-            x_ndb_client.to_proto(),
+            to_proto.client_type(x_ndb_client),
             x_forwarded_for,
             pb_query,
             search_time,
             len(search_results.resources),
         )
+
     if item.debug:
         search_results.nodes = debug_nodes_info(queried_nodes)
 
@@ -121,17 +151,147 @@ async def find(
     search_results.shards = queried_shards
     search_results.autofilters = autofilters
 
-    if search_time > settings.slow_find_log_threshold:
+    if metrics.elapsed("node_query") > settings.slow_node_query_log_threshold:
         logger.warning(
-            "Slow query",
+            "Slow node query",
+            extra={
+                "kbid": kbid,
+                "user": x_nucliadb_user,
+                "client": x_ndb_client,
+                "query": item.model_dump_json(),
+                "time": search_time,
+                "nodes": debug_nodes_info(queried_nodes),
+                "durations": metrics.steps(),
+            },
+        )
+    elif search_time > settings.slow_find_log_threshold:
+        logger.info(
+            "Slow find query",
             extra={
                 "kbid": kbid,
                 "user": x_nucliadb_user,
                 "client": x_ndb_client,
-                "query": item.json(),
+                "query": item.model_dump_json(),
                 "time": search_time,
                 "nodes": debug_nodes_info(queried_nodes),
+                "durations": metrics.steps(),
             },
         )
 
     return search_results, incomplete_results, query_parser
+
+
+async def _external_index_retrieval(
+    kbid: str,
+    item: FindRequest,
+    external_index_manager: ExternalIndexManager,
+    generative_model: Optional[str] = None,
+) -> tuple[KnowledgeboxFindResults, bool, QueryParser]:
+    """
+    Parse the query, query the external index, and hydrate the results.
+    """
+    # Parse query
+    query_parser, _, reranker = await query_parser_from_find_request(
+        kbid, item, generative_model=generative_model
+    )
+    search_request, incomplete_results, _ = await query_parser.parse()
+
+    # Query index
+    query_results = await external_index_manager.query(search_request)  # noqa
+
+    # Hydrate and rerank results
+    text_blocks, resources, best_matches = await hydrate_and_rerank(
+        query_results.iter_matching_text_blocks(),
+        kbid,
+        resource_hydration_options=ResourceHydrationOptions(
+            show=item.show,
+            extracted=item.extracted,
+            field_type_filter=item.field_type_filter,
+        ),
+        text_block_hydration_options=TextBlockHydrationOptions(),
+        reranker=reranker,
+        reranking_options=RerankingOptions(
+            kbid=kbid,
+            query=search_request.body,
+        ),
+        top_k=query_parser.top_k,
+    )
+    find_resources = compose_find_resources(text_blocks, resources)
+
+    results_min_score = MinScore(
+        bm25=0,
+        semantic=query_parser.min_score.semantic,
+    )
+    retrieval_results = KnowledgeboxFindResults(
+        resources=find_resources,
+        query=item.query,
+        total=0,
+        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
+        shards=None,
+        nodes=None,
+    )
+
+    return retrieval_results, incomplete_results, query_parser
+
+
+@dataclass
+class ScoredParagraph:
+    id: str
+    score: float
+
+
+async def query_parser_from_find_request(
+    kbid: str, item: FindRequest, *, generative_model: Optional[str] = None
+) -> tuple[QueryParser, RankFusionAlgorithm, Reranker]:
+    item.min_score = min_score_from_payload(item.min_score)
+
+    if SearchOptions.SEMANTIC in item.features:
+        if should_disable_vector_search(item):
+            item.features.remove(SearchOptions.SEMANTIC)
+
+    hidden = await filter_hidden_resources(kbid, item.show_hidden)
+
+    # XXX this is becoming the new /find query parsing, this should be moved to
+    # a cleaner abstraction
+
+    parsed = parse_find(item)
+
+    rank_fusion = get_rank_fusion(parsed.rank_fusion)
+    reranker = get_reranker(parsed.reranker)
+
+    query_parser = QueryParser(
+        kbid=kbid,
+        features=item.features,
+        query=item.query,
+        label_filters=item.filters,
+        keyword_filters=item.keyword_filters,
+        faceted=None,
+        sort=None,
+        top_k=item.top_k,
+        min_score=item.min_score,
+        range_creation_start=item.range_creation_start,
+        range_creation_end=item.range_creation_end,
+        range_modification_start=item.range_modification_start,
+        range_modification_end=item.range_modification_end,
+        fields=item.fields,
+        user_vector=item.vector,
+        vectorset=item.vectorset,
+        with_duplicates=item.with_duplicates,
+        with_synonyms=item.with_synonyms,
+        autofilter=item.autofilter,
+        key_filters=item.resource_filters,
+        security=item.security,
+        generative_model=generative_model,
+        rephrase=item.rephrase,
+        rephrase_prompt=item.rephrase_prompt,
+        hidden=hidden,
+        rank_fusion=rank_fusion,
+        reranker=reranker,
+    )
+    return (query_parser, rank_fusion, reranker)