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

nucliadb/search/search/merge.py

@@ -22,19 +22,10 @@ import datetime
 import math
 from typing import Any, Optional, Set, Union
 
-from nucliadb_protos.nodereader_pb2 import (
-    DocumentResult,
-    DocumentScored,
-    DocumentSearchResponse,
-    EntitiesSubgraphRequest,
-    ParagraphResult,
-    ParagraphSearchResponse,
-    RelationSearchResponse,
-    SearchResponse,
-    SuggestResponse,
-    VectorSearchResponse,
-)
-
+from nucliadb.common.ids import FieldId, ParagraphId
+from nucliadb.common.models_utils.from_proto import RelationTypePbMap
+from nucliadb.search.search import cache
+from nucliadb.search.search.cut import cut_page
 from nucliadb.search.search.fetch import (
     fetch_resources,
     get_labels_paragraph,
@@ -43,11 +34,11 @@ from nucliadb.search.search.fetch import (
 )
 from nucliadb_models.common import FieldTypeName
 from nucliadb_models.labels import translate_system_to_alias_label
-from nucliadb_models.metadata import RelationTypePbMap
 from nucliadb_models.resource import ExtractedDataTypeName
 from nucliadb_models.search import (
     DirectionalRelation,
     EntitySubgraph,
+    EntityType,
     KnowledgeboxSearchResults,
     KnowledgeboxSuggestResults,
     MinScore,
@@ -56,7 +47,6 @@ from nucliadb_models.search import (
     RelatedEntities,
     RelatedEntity,
     RelationDirection,
-    RelationNodeTypeMap,
     Relations,
     ResourceProperties,
     ResourceResult,
@@ -69,38 +59,59 @@ from nucliadb_models.search import (
     SortOrder,
     TextPosition,
 )
+from nucliadb_protos.nodereader_pb2 import (
+    DocumentResult,
+    DocumentScored,
+    DocumentSearchResponse,
+    EntitiesSubgraphRequest,
+    ParagraphResult,
+    ParagraphSearchResponse,
+    RelationSearchResponse,
+    SearchResponse,
+    SuggestResponse,
+    VectorSearchResponse,
+)
+from nucliadb_protos.utils_pb2 import RelationNode
 
-from .cache import get_resource_cache, get_resource_from_cache
 from .metrics import merge_observer
-from .paragraphs import ExtractedTextCache, get_paragraph_text, get_text_sentence
+from .paragraphs import get_paragraph_text, get_text_sentence
 
 Bm25Score = tuple[float, float]
 TimestampScore = datetime.datetime
 TitleScore = str
-Score = Union[Bm25Score, TimestampScore, TitleScore]
+SortValue = Union[Bm25Score, TimestampScore, TitleScore]
+
+
+def relation_node_type_to_entity_type(node_type: RelationNode.NodeType.ValueType) -> EntityType:
+    return {
+        RelationNode.NodeType.ENTITY: EntityType.ENTITY,
+        RelationNode.NodeType.LABEL: EntityType.LABEL,
+        RelationNode.NodeType.RESOURCE: EntityType.RESOURCE,
+        RelationNode.NodeType.USER: EntityType.USER,
+    }[node_type]
 
 
 def sort_results_by_score(results: Union[list[ParagraphResult], list[DocumentResult]]):
     results.sort(key=lambda x: (x.score.bm25, x.score.booster), reverse=True)
 
 
-async def text_score(
+async def get_sort_value(
     item: Union[DocumentResult, ParagraphResult],
     sort_field: SortField,
     kbid: str,
-) -> Optional[Score]:
+) -> Optional[SortValue]:
     """Returns the score for given `item` and `sort_field`. If the resource is being
     deleted, it might appear on search results but not in maindb. In this
     specific case, return None.
-
     """
     if sort_field == SortField.SCORE:
         return (item.score.bm25, item.score.booster)
 
     score: Any = None
-    resource = await get_resource_from_cache(kbid, item.uuid)
+    resource = await cache.get_resource(kbid, item.uuid)
     if resource is None:
         return score
+
     basic = await resource.get_basic()
     if basic is None:
         return score
@@ -118,13 +129,12 @@ async def text_score(
 async def merge_documents_results(
     document_responses: list[DocumentSearchResponse],
     resources: list[str],
-    count: int,
-    page: int,
+    top_k: int,
     kbid: str,
     sort: SortOptions,
     min_score: float,
 ) -> Resources:
-    raw_resource_list: list[tuple[DocumentResult, Score]] = []
+    raw_resource_list: list[tuple[DocumentResult, SortValue]] = []
     facets: dict[str, Any] = {}
     query = None
     total = 0
@@ -143,24 +153,18 @@ async def merge_documents_results(
         if document_response.next_page:
             next_page = True
         for result in document_response.results:
-            score = await text_score(result, sort.field, kbid)
-            if score is not None:
-                raw_resource_list.append((result, score))
+            sort_value = await get_sort_value(result, sort.field, kbid)
+            if sort_value is not None:
+                raw_resource_list.append((result, sort_value))
         total += document_response.total
 
+    # We need to cut first and then sort, otherwise the page will be wrong if the order is DESC
+    raw_resource_list, has_more = cut_page(raw_resource_list, top_k)
+    next_page = next_page or has_more
     raw_resource_list.sort(key=lambda x: x[1], reverse=(sort.order == SortOrder.DESC))
 
-    skip = page * count
-    end = skip + count
-    length = len(raw_resource_list)
-
-    if length > end:
-        next_page = True
-
     result_resource_list: list[ResourceResult] = []
-    for result, _ in raw_resource_list[min(skip, length) : min(end, length)]:
-        # /f/file
-
+    for result, _ in raw_resource_list:
         labels = await get_labels_resource(result, kbid)
         _, field_type, field = result.field.split("/")
 
@@ -181,8 +185,8 @@ async def merge_documents_results(
         results=result_resource_list,
         query=query,
         total=total,
-        page_number=page,
-        page_size=count,
+        page_number=0,  # Bw/c with pagination
+        page_size=top_k,
         next_page=next_page,
         min_score=min_score,
     )
@@ -207,65 +211,58 @@ async def merge_suggest_paragraph_results(
     if len(suggest_responses) > 1:
         sort_results_by_score(raw_paragraph_list)
 
-    rcache = get_resource_cache(clear=True)
-    etcache = ExtractedTextCache()
-    try:
-        result_paragraph_list: list[Paragraph] = []
-        for result in raw_paragraph_list[:10]:
-            _, field_type, field = result.field.split("/")
-            text = await get_paragraph_text(
-                kbid=kbid,
-                rid=result.uuid,
-                field=result.field,
-                start=result.start,
-                end=result.end,
-                split=result.split,
-                highlight=highlight,
-                ematches=ematches,  # type: ignore
-                matches=result.matches,  # type: ignore
-                extracted_text_cache=etcache,
-            )
-            labels = await get_labels_paragraph(result, kbid)
-            new_paragraph = Paragraph(
-                score=result.score.bm25,
-                rid=result.uuid,
-                field_type=field_type,
-                field=field,
-                text=text,
-                labels=labels,
-                position=TextPosition(
-                    index=result.metadata.position.index,
-                    start=result.metadata.position.start,
-                    end=result.metadata.position.end,
-                    page_number=result.metadata.position.page_number,
+    result_paragraph_list: list[Paragraph] = []
+    for result in raw_paragraph_list[:10]:
+        _, field_type, field = result.field.split("/")
+        text = await get_paragraph_text(
+            kbid=kbid,
+            paragraph_id=ParagraphId(
+                field_id=FieldId(
+                    rid=result.uuid,
+                    type=field_type,
+                    key=field,
+                    subfield_id=result.split,
                 ),
-            )
-            if len(result.metadata.position.start_seconds) or len(
-                result.metadata.position.end_seconds
-            ):
-                new_paragraph.start_seconds = list(
-                    result.metadata.position.start_seconds
-                )
-                new_paragraph.end_seconds = list(result.metadata.position.end_seconds)
-            else:
-                # TODO: Remove once we are sure all data has been migrated!
-                seconds_positions = await get_seconds_paragraph(result, kbid)
-                if seconds_positions is not None:
-                    new_paragraph.start_seconds = seconds_positions[0]
-                    new_paragraph.end_seconds = seconds_positions[1]
-            result_paragraph_list.append(new_paragraph)
-        return Paragraphs(results=result_paragraph_list, query=query, min_score=0)
-    finally:
-        etcache.clear()
-        rcache.clear()
+                paragraph_start=result.start,
+                paragraph_end=result.end,
+            ),
+            highlight=highlight,
+            ematches=ematches,  # type: ignore
+            matches=result.matches,  # type: ignore
+        )
+        labels = await get_labels_paragraph(result, kbid)
+        new_paragraph = Paragraph(
+            score=result.score.bm25,
+            rid=result.uuid,
+            field_type=field_type,
+            field=field,
+            text=text,
+            labels=labels,
+            position=TextPosition(
+                index=result.metadata.position.index,
+                start=result.metadata.position.start,
+                end=result.metadata.position.end,
+                page_number=result.metadata.position.page_number,
+            ),
+        )
+        if len(result.metadata.position.start_seconds) or len(result.metadata.position.end_seconds):
+            new_paragraph.start_seconds = list(result.metadata.position.start_seconds)
+            new_paragraph.end_seconds = list(result.metadata.position.end_seconds)
+        else:
+            # TODO: Remove once we are sure all data has been migrated!
+            seconds_positions = await get_seconds_paragraph(result, kbid)
+            if seconds_positions is not None:
+                new_paragraph.start_seconds = seconds_positions[0]
+                new_paragraph.end_seconds = seconds_positions[1]
+        result_paragraph_list.append(new_paragraph)
+    return Paragraphs(results=result_paragraph_list, query=query, min_score=0)
 
 
 async def merge_vectors_results(
     vector_responses: list[VectorSearchResponse],
     resources: list[str],
     kbid: str,
-    count: int,
-    page: int,
+    top_k: int,
     min_score: Optional[float] = None,
 ):
     facets: dict[str, Any] = {}
@@ -282,12 +279,10 @@ async def merge_vectors_results(
     if len(vector_responses) > 1:
         raw_vectors_list.sort(key=lambda x: x.score, reverse=True)
 
-    skip = page * count
-    end_element = skip + count
-    length = len(raw_vectors_list)
+    raw_vectors_list, _ = cut_page(raw_vectors_list, top_k)
 
     result_sentence_list: list[Sentence] = []
-    for result in raw_vectors_list[min(skip, length) : min(end_element, length)]:
+    for result in raw_vectors_list:
         id_count = result.doc_id.id.count("/")
         if id_count == 4:
             rid, field_type, field, index, position = result.doc_id.id.split("/")
@@ -335,8 +330,8 @@ async def merge_vectors_results(
     return Sentences(
         results=result_sentence_list,
         facets=facets,
-        page_number=page,
-        page_size=count,
+        page_number=0,  # Bw/c with pagination
+        page_size=top_k,
         min_score=round(min_score or 0, ndigits=3),
     )
 
@@ -345,13 +340,12 @@ async def merge_paragraph_results(
     paragraph_responses: list[ParagraphSearchResponse],
     resources: list[str],
     kbid: str,
-    count: int,
-    page: int,
+    top_k: int,
     highlight: bool,
     sort: SortOptions,
     min_score: float,
-):
-    raw_paragraph_list: list[tuple[ParagraphResult, Score]] = []
+) -> Paragraphs:
+    raw_paragraph_list: list[tuple[ParagraphResult, SortValue]] = []
     facets: dict[str, Any] = {}
     query = None
     next_page = False
@@ -373,83 +367,75 @@ async def merge_paragraph_results(
         if paragraph_response.next_page:
             next_page = True
         for result in paragraph_response.results:
-            score = await text_score(result, sort.field, kbid)
+            score = await get_sort_value(result, sort.field, kbid)
             if score is not None:
                 raw_paragraph_list.append((result, score))
         total += paragraph_response.total
 
     raw_paragraph_list.sort(key=lambda x: x[1], reverse=(sort.order == SortOrder.DESC))
 
-    skip = page * count
-    end = skip + count
-    length = len(raw_paragraph_list)
-
-    if length > end:
-        next_page = True
+    raw_paragraph_list, has_more = cut_page(raw_paragraph_list, top_k)
+    next_page = next_page or has_more
 
     result_paragraph_list: list[Paragraph] = []
-    etcache = ExtractedTextCache()
-    try:
-        for result, _ in raw_paragraph_list[min(skip, length) : min(end, length)]:
-            _, field_type, field = result.field.split("/")
-            text = await get_paragraph_text(
-                kbid=kbid,
-                rid=result.uuid,
-                field=result.field,
-                start=result.start,
-                end=result.end,
-                split=result.split,
-                highlight=highlight,
-                ematches=ematches,
-                matches=result.matches,  # type: ignore
-                extracted_text_cache=etcache,
-            )
-            labels = await get_labels_paragraph(result, kbid)
-            fuzzy_result = len(result.matches) > 0
-            new_paragraph = Paragraph(
-                score=result.score.bm25,
-                rid=result.uuid,
-                field_type=field_type,
-                field=field,
-                text=text,
-                labels=labels,
-                position=TextPosition(
-                    index=result.metadata.position.index,
-                    start=result.metadata.position.start,
-                    end=result.metadata.position.end,
-                    page_number=result.metadata.position.page_number,
+    for result, _ in raw_paragraph_list:
+        _, field_type, field = result.field.split("/")
+        text = await get_paragraph_text(
+            kbid=kbid,
+            paragraph_id=ParagraphId(
+                field_id=FieldId(
+                    rid=result.uuid,
+                    type=field_type,
+                    key=field,
+                    subfield_id=result.split,
                 ),
-                fuzzy_result=fuzzy_result,
-            )
-            if len(result.metadata.position.start_seconds) or len(
-                result.metadata.position.end_seconds
-            ):
-                new_paragraph.start_seconds = list(
-                    result.metadata.position.start_seconds
-                )
-                new_paragraph.end_seconds = list(result.metadata.position.end_seconds)
-            else:
-                # TODO: Remove once we are sure all data has been migrated!
-                seconds_positions = await get_seconds_paragraph(result, kbid)
-                if seconds_positions is not None:
-                    new_paragraph.start_seconds = seconds_positions[0]
-                    new_paragraph.end_seconds = seconds_positions[1]
-
-            result_paragraph_list.append(new_paragraph)
-            if new_paragraph.rid not in resources:
-                resources.append(new_paragraph.rid)
-        return Paragraphs(
-            results=result_paragraph_list,
-            facets=facets,
-            query=query,
-            total=total,
-            page_number=page,
-            page_size=count,
-            next_page=next_page,
-            min_score=min_score,
+                paragraph_start=result.start,
+                paragraph_end=result.end,
+            ),
+            highlight=highlight,
+            ematches=ematches,
+            matches=result.matches,  # type: ignore
         )
-    finally:
-        etcache.clear()
+        labels = await get_labels_paragraph(result, kbid)
+        fuzzy_result = len(result.matches) > 0
+        new_paragraph = Paragraph(
+            score=result.score.bm25,
+            rid=result.uuid,
+            field_type=field_type,
+            field=field,
+            text=text,
+            labels=labels,
+            position=TextPosition(
+                index=result.metadata.position.index,
+                start=result.metadata.position.start,
+                end=result.metadata.position.end,
+                page_number=result.metadata.position.page_number,
+            ),
+            fuzzy_result=fuzzy_result,
+        )
+        if len(result.metadata.position.start_seconds) or len(result.metadata.position.end_seconds):
+            new_paragraph.start_seconds = list(result.metadata.position.start_seconds)
+            new_paragraph.end_seconds = list(result.metadata.position.end_seconds)
+        else:
+            # TODO: Remove once we are sure all data has been migrated!
+            seconds_positions = await get_seconds_paragraph(result, kbid)
+            if seconds_positions is not None:
+                new_paragraph.start_seconds = seconds_positions[0]
+                new_paragraph.end_seconds = seconds_positions[1]
+
+        result_paragraph_list.append(new_paragraph)
+        if new_paragraph.rid not in resources:
+            resources.append(new_paragraph.rid)
+    return Paragraphs(
+        results=result_paragraph_list,
+        facets=facets,
+        query=query,
+        total=total,
+        page_number=0,  # Bw/c with pagination
+        page_size=top_k,
+        next_page=next_page,
+        min_score=min_score,
+    )
 
 
 @merge_observer.wrap({"type": "merge_relations"})
@@ -458,9 +444,7 @@ async def merge_relations_results(
     query: EntitiesSubgraphRequest,
 ) -> Relations:
     loop = asyncio.get_event_loop()
-    return await loop.run_in_executor(
-        None, _merge_relations_results, relations_responses, query
-    )
+    return await loop.run_in_executor(None, _merge_relations_results, relations_responses, query)
 
 
 def _merge_relations_results(
@@ -483,7 +467,7 @@ def _merge_relations_results(
                 relations.entities[origin.value].related_to.append(
                     DirectionalRelation(
                         entity=destination.value,
-                        entity_type=RelationNodeTypeMap[destination.ntype],
+                        entity_type=relation_node_type_to_entity_type(destination.ntype),
                         relation=relation_type,
                         relation_label=relation_label,
                         direction=RelationDirection.OUT,
@@ -493,7 +477,7 @@ def _merge_relations_results(
                 relations.entities[destination.value].related_to.append(
                     DirectionalRelation(
                         entity=origin.value,
-                        entity_type=RelationNodeTypeMap[origin.ntype],
+                        entity_type=relation_node_type_to_entity_type(origin.ntype),
                         relation=relation_type,
                         relation_label=relation_label,
                         direction=RelationDirection.IN,
@@ -506,8 +490,7 @@ def _merge_relations_results(
 @merge_observer.wrap({"type": "merge"})
 async def merge_results(
     search_responses: list[SearchResponse],
-    count: int,
-    page: int,
+    top_k: int,
     kbid: str,
     show: list[ResourceProperties],
     field_type_filter: list[FieldTypeName],
@@ -530,77 +513,59 @@ async def merge_results(
 
     api_results = KnowledgeboxSearchResults()
 
-    rcache = get_resource_cache(clear=True)
-    try:
-        resources: list[str] = list()
-        api_results.fulltext = await merge_documents_results(
-            documents, resources, count, page, kbid, sort, min_score=min_score.bm25
-        )
+    resources: list[str] = list()
+    api_results.fulltext = await merge_documents_results(
+        documents, resources, top_k, kbid, sort, min_score=min_score.bm25
+    )
 
-        api_results.paragraphs = await merge_paragraph_results(
-            paragraphs,
-            resources,
-            kbid,
-            count,
-            page,
-            highlight,
-            sort,
-            min_score=min_score.bm25,
-        )
+    api_results.paragraphs = await merge_paragraph_results(
+        paragraphs,
+        resources,
+        kbid,
+        top_k,
+        highlight,
+        sort,
+        min_score=min_score.bm25,
+    )
 
-        api_results.sentences = await merge_vectors_results(
-            vectors, resources, kbid, count, page, min_score=min_score.semantic
-        )
+    api_results.sentences = await merge_vectors_results(
+        vectors, resources, kbid, top_k, min_score=min_score.semantic
+    )
 
-        api_results.relations = await merge_relations_results(
-            relations, requested_relations
-        )
+    api_results.relations = await merge_relations_results(relations, requested_relations)
 
-        api_results.resources = await fetch_resources(
-            resources, kbid, show, field_type_filter, extracted
-        )
-        return api_results
-    finally:
-        rcache.clear()
+    api_results.resources = await fetch_resources(resources, kbid, show, field_type_filter, extracted)
+    return api_results
 
 
 async def merge_paragraphs_results(
-    paragraph_responses: list[ParagraphSearchResponse],
-    count: int,
-    page: int,
+    responses: list[SearchResponse],
+    top_k: int,
     kbid: str,
-    show: list[ResourceProperties],
-    field_type_filter: list[FieldTypeName],
-    extracted: list[ExtractedDataTypeName],
     highlight_split: bool,
     min_score: float,
 ) -> ResourceSearchResults:
     paragraphs = []
-    for result in paragraph_responses:
-        paragraphs.append(result)
+    for result in responses:
+        paragraphs.append(result.paragraph)
 
     api_results = ResourceSearchResults()
 
-    rcache = get_resource_cache(clear=True)
-    try:
-        resources: list[str] = list()
-        api_results.paragraphs = await merge_paragraph_results(
-            paragraphs,
-            resources,
-            kbid,
-            count,
-            page,
-            highlight=highlight_split,
-            sort=SortOptions(
-                field=SortField.SCORE,
-                order=SortOrder.DESC,
-                limit=None,
-            ),
-            min_score=min_score,
-        )
-        return api_results
-    finally:
-        rcache.clear()
+    resources: list[str] = list()
+    api_results.paragraphs = await merge_paragraph_results(
+        paragraphs,
+        resources,
+        kbid,
+        top_k,
+        highlight=highlight_split,
+        sort=SortOptions(
+            field=SortField.SCORE,
+            order=SortOrder.DESC,
+            limit=None,
+        ),
+        min_score=min_score,
+    )
+    return api_results
 
 
 async def merge_suggest_entities_results(
@@ -609,8 +574,7 @@ async def merge_suggest_entities_results(
     unique_entities: Set[RelatedEntity] = set()
     for response in suggest_responses:
         response_entities = (
-            RelatedEntity(family=e.subtype, value=e.value)
-            for e in response.entity_results.nodes
+            RelatedEntity(family=e.subtype, value=e.value) for e in response.entity_results.nodes
         )
         unique_entities.update(response_entities)