nucliadb 4.0.0.post542__py3-none-any.whl → 6.2.1.post2798__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,10 +59,22 @@ 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
@@ -80,6 +82,15 @@ TitleScore = str
 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)
 
@@ -97,7 +108,7 @@ async def get_sort_value(
         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
 
@@ -118,8 +129,7 @@ async def get_sort_value(
 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,
@@ -148,15 +158,9 @@ async def merge_documents_results(
                 raw_resource_list.append((result, sort_value))
         total += document_response.total
 
-    skip = page * count
-    end = skip + count
-    length = len(raw_resource_list)
-
-    if length > end:
-        next_page = True
-
-    # We need to cut first and then sort, otherwise pagination will be wrong if the order is DESC
-    raw_resource_list = raw_resource_list[min(skip, length) : min(end, length)]
+    # 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))
 
     result_resource_list: list[ResourceResult] = []
@@ -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,12 +340,11 @@ 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,
-):
+) -> Paragraphs:
     raw_paragraph_list: list[tuple[ParagraphResult, SortValue]] = []
     facets: dict[str, Any] = {}
     query = None
@@ -380,76 +374,68 @@ async def merge_paragraph_results(
 
     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)
 

nucliadb/search/search/metrics.py

@@ -17,10 +17,81 @@
 # 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 contextlib
+import time
+from typing import Optional
+
 from nucliadb_telemetry import metrics
 
 merge_observer = metrics.Observer("merge_results", labels={"type": ""})
 node_features = metrics.Counter("nucliadb_node_features", labels={"type": ""})
-query_parse_dependency_observer = metrics.Observer(
-    "query_parse_dependency", labels={"type": ""}
+query_parse_dependency_observer = metrics.Observer("query_parse_dependency", labels={"type": ""})
+
+buckets = [
+    0.005,
+    0.01,
+    0.025,
+    0.05,
+    0.075,
+    0.1,
+    0.25,
+    0.5,
+    0.75,
+    1.0,
+    2.5,
+    5.0,
+    7.5,
+    10.0,
+    30.0,
+    60.0,
+    metrics.INF,
+]
+
+generative_first_chunk_histogram = metrics.Histogram(
+    name="generative_first_chunk",
+    buckets=buckets,
+)
+rag_histogram = metrics.Histogram(
+    name="rag",
+    labels={"step": ""},
+    buckets=buckets,
 )
+
+
+class RAGMetrics:
+    def __init__(self):
+        self.global_start = time.monotonic()
+        self._start_times: dict[str, float] = {}
+        self._end_times: dict[str, float] = {}
+        self.first_chunk_yielded_at: Optional[float] = None
+
+    @contextlib.contextmanager
+    def time(self, step: str):
+        self._start(step)
+        try:
+            yield
+        finally:
+            self._end(step)
+
+    def steps(self) -> dict[str, float]:
+        return {step: self.elapsed(step) for step in self._end_times.keys()}
+
+    def elapsed(self, step: str) -> float:
+        return self._end_times[step] - self._start_times[step]
+
+    def record_first_chunk_yielded(self):
+        self.first_chunk_yielded_at = time.monotonic()
+        generative_first_chunk_histogram.observe(self.first_chunk_yielded_at - self.global_start)
+
+    def get_first_chunk_time(self) -> Optional[float]:
+        if self.first_chunk_yielded_at is None:
+            return None
+        return self.first_chunk_yielded_at - self.global_start
+
+    def _start(self, step: str):
+        self._start_times[step] = time.monotonic()
+
+    def _end(self, step: str):
+        self._end_times[step] = time.monotonic()
+        elapsed = self.elapsed(step)
+        rag_histogram.observe(elapsed, labels={"step": step})