nucliadb 6.4.0.post4271__py3-none-any.whl → 6.4.0.post4276__py3-none-any.whl

nucliadb/search/search/chat/ask.py

@@ -507,17 +507,18 @@ async def ask(
             logger.info("Failed to rephrase ask query, using original")
 
     try:
-        retrieval_results = await retrieval_step(
-            kbid=kbid,
-            # Prefer the rephrased query for retrieval if available
-            main_query=rephrased_query or user_query,
-            ask_request=ask_request,
-            client_type=client_type,
-            user_id=user_id,
-            origin=origin,
-            metrics=metrics,
-            resource=resource,
-        )
+        with metrics.time("retrieval"):
+            retrieval_results = await retrieval_step(
+                kbid=kbid,
+                # Prefer the rephrased query for retrieval if available
+                main_query=rephrased_query or user_query,
+                ask_request=ask_request,
+                client_type=client_type,
+                user_id=user_id,
+                origin=origin,
+                metrics=metrics,
+                resource=resource,
+            )
     except NoRetrievalResultsError as err:
         maybe_audit_chat(
             kbid=kbid,
@@ -562,6 +563,7 @@ async def ask(
             image_strategies=ask_request.rag_images_strategies,
             max_context_characters=tokens_to_chars(generation.max_context_tokens),
             visual_llm=generation.use_visual_llm,
+            metrics=metrics.child_span("context_building"),
         )
         (
             prompt_context,
@@ -747,45 +749,44 @@ async def retrieval_in_kb(
 ) -> RetrievalResults:
     prequeries = parse_prequeries(ask_request)
     graph_strategy = parse_graph_strategy(ask_request)
-    with metrics.time("retrieval"):
-        main_results, prequeries_results, parsed_query = await get_find_results(
+    main_results, prequeries_results, parsed_query = await get_find_results(
+        kbid=kbid,
+        query=main_query,
+        item=ask_request,
+        ndb_client=client_type,
+        user=user_id,
+        origin=origin,
+        metrics=metrics.child_span("hybrid_retrieval"),
+        prequeries_strategy=prequeries,
+    )
+
+    if graph_strategy is not None:
+        assert parsed_query.retrieval.reranker is not None, (
+            "find parser must provide a reranking algorithm"
+        )
+        reranker = get_reranker(parsed_query.retrieval.reranker)
+        graph_results, graph_request = await get_graph_results(
             kbid=kbid,
             query=main_query,
             item=ask_request,
             ndb_client=client_type,
             user=user_id,
             origin=origin,
-            metrics=metrics,
-            prequeries_strategy=prequeries,
+            graph_strategy=graph_strategy,
+            metrics=metrics.child_span("graph_retrieval"),
+            text_block_reranker=reranker,
         )
 
-        if graph_strategy is not None:
-            assert parsed_query.retrieval.reranker is not None, (
-                "find parser must provide a reranking algorithm"
-            )
-            reranker = get_reranker(parsed_query.retrieval.reranker)
-            graph_results, graph_request = await get_graph_results(
-                kbid=kbid,
-                query=main_query,
-                item=ask_request,
-                ndb_client=client_type,
-                user=user_id,
-                origin=origin,
-                graph_strategy=graph_strategy,
-                metrics=metrics,
-                text_block_reranker=reranker,
-            )
-
-            if prequeries_results is None:
-                prequeries_results = []
+        if prequeries_results is None:
+            prequeries_results = []
 
-            prequery = PreQuery(id="graph", request=graph_request, weight=graph_strategy.weight)
-            prequeries_results.append((prequery, graph_results))
+        prequery = PreQuery(id="graph", request=graph_request, weight=graph_strategy.weight)
+        prequeries_results.append((prequery, graph_results))
 
-        if len(main_results.resources) == 0 and all(
-            len(prequery_result.resources) == 0 for (_, prequery_result) in prequeries_results or []
-        ):
-            raise NoRetrievalResultsError(main_results, prequeries_results)
+    if len(main_results.resources) == 0 and all(
+        len(prequery_result.resources) == 0 for (_, prequery_result) in prequeries_results or []
+    ):
+        raise NoRetrievalResultsError(main_results, prequeries_results)
 
     main_query_weight = prequeries.main_query_weight if prequeries is not None else 1.0
     best_matches = compute_best_matches(
@@ -836,21 +837,20 @@ async def retrieval_in_resource(
                 )
             add_resource_filter(prequery.request, [resource])
 
-    with metrics.time("retrieval"):
-        main_results, prequeries_results, parsed_query = await get_find_results(
-            kbid=kbid,
-            query=main_query,
-            item=ask_request,
-            ndb_client=client_type,
-            user=user_id,
-            origin=origin,
-            metrics=metrics,
-            prequeries_strategy=prequeries,
-        )
-        if len(main_results.resources) == 0 and all(
-            len(prequery_result.resources) == 0 for (_, prequery_result) in prequeries_results or []
-        ):
-            raise NoRetrievalResultsError(main_results, prequeries_results)
+    main_results, prequeries_results, parsed_query = await get_find_results(
+        kbid=kbid,
+        query=main_query,
+        item=ask_request,
+        ndb_client=client_type,
+        user=user_id,
+        origin=origin,
+        metrics=metrics.child_span("hybrid_retrieval"),
+        prequeries_strategy=prequeries,
+    )
+    if len(main_results.resources) == 0 and all(
+        len(prequery_result.resources) == 0 for (_, prequery_result) in prequeries_results or []
+    ):
+        raise NoRetrievalResultsError(main_results, prequeries_results)
     main_query_weight = prequeries.main_query_weight if prequeries is not None else 1.0
     best_matches = compute_best_matches(
         main_results=main_results,

nucliadb/search/search/chat/prompt.py

@@ -41,6 +41,7 @@ from nucliadb.search.search.chat.images import (
     get_paragraph_image,
 )
 from nucliadb.search.search.hydrator import hydrate_field_text, hydrate_resource_text
+from nucliadb.search.search.metrics import Metrics
 from nucliadb.search.search.paragraphs import get_paragraph_text
 from nucliadb_models.labels import translate_alias_to_system_label
 from nucliadb_models.metadata import Extra, Origin
@@ -244,6 +245,7 @@ async def full_resource_prompt_context(
     ordered_paragraphs: list[FindParagraph],
     resource: Optional[str],
     strategy: FullResourceStrategy,
+    metrics: Metrics,
 ) -> None:
     """
     Algorithm steps:
@@ -298,6 +300,8 @@ async def full_resource_prompt_context(
             context[field.full()] = extracted_text
             added_fields.add(field.full())
 
+    metrics.set("full_resource_ops", len(added_fields))
+
     if strategy.include_remaining_text_blocks:
         for paragraph in ordered_paragraphs:
             pid = cast(ParagraphId, parse_text_block_id(paragraph.id))
@@ -309,6 +313,7 @@ async def extend_prompt_context_with_metadata(
     context: CappedPromptContext,
     kbid: str,
     strategy: MetadataExtensionStrategy,
+    metrics: Metrics,
 ) -> None:
     text_block_ids: list[TextBlockId] = []
     for text_block_id in context.text_block_ids():
@@ -321,18 +326,25 @@ async def extend_prompt_context_with_metadata(
     if len(text_block_ids) == 0:  # pragma: no cover
         return
 
+    ops = 0
     if MetadataExtensionType.ORIGIN in strategy.types:
+        ops += 1
         await extend_prompt_context_with_origin_metadata(context, kbid, text_block_ids)
 
     if MetadataExtensionType.CLASSIFICATION_LABELS in strategy.types:
+        ops += 1
         await extend_prompt_context_with_classification_labels(context, kbid, text_block_ids)
 
     if MetadataExtensionType.NERS in strategy.types:
+        ops += 1
         await extend_prompt_context_with_ner(context, kbid, text_block_ids)
 
     if MetadataExtensionType.EXTRA_METADATA in strategy.types:
+        ops += 1
         await extend_prompt_context_with_extra_metadata(context, kbid, text_block_ids)
 
+    metrics.set("metadata_extension_ops", ops * len(text_block_ids))
+
 
 def parse_text_block_id(text_block_id: str) -> TextBlockId:
     try:
@@ -464,6 +476,7 @@ async def field_extension_prompt_context(
     kbid: str,
     ordered_paragraphs: list[FindParagraph],
     strategy: FieldExtensionStrategy,
+    metrics: Metrics,
 ) -> None:
     """
     Algorithm steps:
@@ -493,6 +506,8 @@ async def field_extension_prompt_context(
     tasks = [hydrate_field_text(kbid, fid) for fid in extend_field_ids]
     field_extracted_texts = await run_concurrently(tasks)
 
+    metrics.set("field_extension_ops", len(field_extracted_texts))
+
     for result in field_extracted_texts:
         if result is None:  # pragma: no cover
             continue
@@ -619,6 +634,7 @@ async def neighbouring_paragraphs_prompt_context(
     kbid: str,
     ordered_text_blocks: list[FindParagraph],
     strategy: NeighbouringParagraphsStrategy,
+    metrics: Metrics,
 ) -> None:
     """
     This function will get the paragraph texts and then craft a context with the neighbouring paragraphs of the
@@ -658,6 +674,9 @@ async def neighbouring_paragraphs_prompt_context(
         return
 
     results: list[tuple[ParagraphId, str]] = await asyncio.gather(*paragraph_ops)
+
+    metrics.set("neighbouring_paragraphs_ops", len(results))
+
     # Add the paragraph texts to the context
     for pid, text in results:
         if text != "":
@@ -670,8 +689,10 @@ async def conversation_prompt_context(
     ordered_paragraphs: list[FindParagraph],
     conversational_strategy: ConversationalStrategy,
     visual_llm: bool,
+    metrics: Metrics,
 ):
     analyzed_fields: List[str] = []
+    ops = 0
     async with get_driver().transaction(read_only=True) as txn:
         storage = await get_storage()
         kb = KnowledgeBoxORM(txn, storage, kbid)
@@ -701,6 +722,7 @@ async def conversation_prompt_context(
 
                 attachments: List[resources_pb2.FieldRef] = []
                 if conversational_strategy.full:
+                    ops += 5
                     extracted_text = await field_obj.get_extracted_text()
                     for current_page in range(1, cmetadata.pages + 1):
                         conv = await field_obj.db_get_value(current_page)
@@ -749,6 +771,7 @@ async def conversation_prompt_context(
                             break
 
                     for message in messages:
+                        ops += 1
                         text = message.content.text.strip()
                         pid = f"{rid}/{field_type}/{field_id}/{message.ident}/0-{len(message.content.text) + 1}"
                         context[pid] = text
@@ -757,6 +780,7 @@ async def conversation_prompt_context(
                 if conversational_strategy.attachments_text:
                     # add on the context the images if vlm enabled
                     for attachment in attachments:
+                        ops += 1
                         field: File = await resource.get_field(
                             attachment.field_id, attachment.field_type, load=True
                         )  # type: ignore
@@ -767,6 +791,7 @@ async def conversation_prompt_context(
 
                 if conversational_strategy.attachments_images and visual_llm:
                     for attachment in attachments:
+                        ops += 1
                         file_field: File = await resource.get_field(
                             attachment.field_id, attachment.field_type, load=True
                         )  # type: ignore
@@ -776,6 +801,7 @@ async def conversation_prompt_context(
                             context.images[pid] = image
 
                 analyzed_fields.append(field_unique_id)
+    metrics.set("conversation_ops", ops)
 
 
 async def hierarchy_prompt_context(
@@ -783,6 +809,7 @@ async def hierarchy_prompt_context(
     kbid: str,
     ordered_paragraphs: list[FindParagraph],
     strategy: HierarchyResourceStrategy,
+    metrics: Metrics,
 ) -> None:
     """
     This function will get the paragraph texts (possibly with extra characters, if extra_characters > 0) and then
@@ -842,6 +869,8 @@ async def hierarchy_prompt_context(
         else:
             resources[rid].paragraphs.append((paragraph, extended_paragraph_text))
 
+    metrics.set("hierarchy_ops", len(resources))
+
     # Modify the first paragraph of each resource to include the title and summary of the resource, as well as the
     # extended paragraph text of all the paragraphs in the resource.
     for values in resources.values():
@@ -886,6 +915,7 @@ class PromptContextBuilder:
         image_strategies: Optional[Sequence[ImageRagStrategy]] = None,
         max_context_characters: Optional[int] = None,
         visual_llm: bool = False,
+        metrics: Metrics = Metrics("prompt_context_builder"),
     ):
         self.kbid = kbid
         self.ordered_paragraphs = ordered_paragraphs
@@ -896,6 +926,7 @@ class PromptContextBuilder:
         self.image_strategies = image_strategies
         self.max_context_characters = max_context_characters
         self.visual_llm = visual_llm
+        self.metrics = metrics
 
     def prepend_user_context(self, context: CappedPromptContext):
         # Chat extra context passed by the user is the most important, therefore
@@ -920,6 +951,7 @@ class PromptContextBuilder:
         return context, context_order, context_images
 
     async def _build_context_images(self, context: CappedPromptContext) -> None:
+        ops = 0
         if self.image_strategies is None or len(self.image_strategies) == 0:
             # Nothing to do
             return
@@ -958,6 +990,7 @@ class PromptContextBuilder:
                 if page_image_id not in context.images:
                     image = await get_page_image(self.kbid, pid, paragraph_page_number)
                     if image is not None:
+                        ops += 1
                         context.images[page_image_id] = image
                         page_images_added += 1
                     else:
@@ -977,6 +1010,7 @@ class PromptContextBuilder:
             ):
                 pimage = await get_paragraph_image(self.kbid, pid, paragraph.reference)
                 if pimage is not None:
+                    ops += 1
                     context.images[paragraph.id] = pimage
                 else:
                     logger.warning(
@@ -987,6 +1021,7 @@ class PromptContextBuilder:
                             "reference": paragraph.reference,
                         },
                     )
+        self.metrics.set("image_ops", ops)
 
     async def _build_context(self, context: CappedPromptContext) -> None:
         if self.strategies is None or len(self.strategies) == 0:
@@ -1038,17 +1073,17 @@ class PromptContextBuilder:
                 self.ordered_paragraphs,
                 self.resource,
                 full_resource,
+                self.metrics,
             )
             if metadata_extension:
-                await extend_prompt_context_with_metadata(context, self.kbid, metadata_extension)
+                await extend_prompt_context_with_metadata(
+                    context, self.kbid, metadata_extension, self.metrics
+                )
             return
 
         if hierarchy:
             await hierarchy_prompt_context(
-                context,
-                self.kbid,
-                self.ordered_paragraphs,
-                hierarchy,
+                context, self.kbid, self.ordered_paragraphs, hierarchy, self.metrics
             )
         if neighbouring_paragraphs:
             await neighbouring_paragraphs_prompt_context(
@@ -1056,6 +1091,7 @@ class PromptContextBuilder:
                 self.kbid,
                 self.ordered_paragraphs,
                 neighbouring_paragraphs,
+                self.metrics,
             )
         if field_extension:
             await field_extension_prompt_context(
@@ -1063,6 +1099,7 @@ class PromptContextBuilder:
                 self.kbid,
                 self.ordered_paragraphs,
                 field_extension,
+                self.metrics,
             )
         if conversational_strategy:
             await conversation_prompt_context(
@@ -1071,9 +1108,12 @@ class PromptContextBuilder:
                 self.ordered_paragraphs,
                 conversational_strategy,
                 self.visual_llm,
+                self.metrics,
             )
         if metadata_extension:
-            await extend_prompt_context_with_metadata(context, self.kbid, metadata_extension)
+            await extend_prompt_context_with_metadata(
+                context, self.kbid, metadata_extension, self.metrics
+            )
 
 
 def get_paragraph_page_number(paragraph: FindParagraph) -> Optional[int]:

{nucliadb-6.4.0.post4271.dist-info → nucliadb-6.4.0.post4276.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: nucliadb
-Version: 6.4.0.post4271
+Version: 6.4.0.post4276
 Summary: NucliaDB
 Author-email: Nuclia <nucliadb@nuclia.com>
 License: AGPL
@@ -20,11 +20,11 @@ Classifier: Programming Language :: Python :: 3.12
 Classifier: Programming Language :: Python :: 3 :: Only
 Requires-Python: <4,>=3.9
 Description-Content-Type: text/markdown
-Requires-Dist: nucliadb-telemetry[all]>=6.4.0.post4271
-Requires-Dist: nucliadb-utils[cache,fastapi,storages]>=6.4.0.post4271
-Requires-Dist: nucliadb-protos>=6.4.0.post4271
-Requires-Dist: nucliadb-models>=6.4.0.post4271
-Requires-Dist: nidx-protos>=6.4.0.post4271
+Requires-Dist: nucliadb-telemetry[all]>=6.4.0.post4276
+Requires-Dist: nucliadb-utils[cache,fastapi,storages]>=6.4.0.post4276
+Requires-Dist: nucliadb-protos>=6.4.0.post4276
+Requires-Dist: nucliadb-models>=6.4.0.post4276
+Requires-Dist: nidx-protos>=6.4.0.post4276
 Requires-Dist: nucliadb-admin-assets>=1.0.0.post1224
 Requires-Dist: nuclia-models>=0.24.2
 Requires-Dist: uvicorn[standard]

{nucliadb-6.4.0.post4271.dist-info → nucliadb-6.4.0.post4276.dist-info}/RECORD

@@ -255,10 +255,10 @@ nucliadb/search/search/shards.py,sha256=mc5DK-MoCv9AFhlXlOFHbPvetcyNDzTFOJ5rimK8
 nucliadb/search/search/summarize.py,sha256=ksmYPubEQvAQgfPdZHfzB_rR19B2ci4IYZ6jLdHxZo8,4996
 nucliadb/search/search/utils.py,sha256=ajRIXfdTF67dBVahQCXW-rSv6gJpUMPt3QhJrWqArTQ,2175
 nucliadb/search/search/chat/__init__.py,sha256=cp15ZcFnHvpcu_5-aK2A4uUyvuZVV_MJn4bIXMa20ks,835
-nucliadb/search/search/chat/ask.py,sha256=4cpeWC7Q4NI30vrOeq22N0Vw0PQS0Ko9b47R5ISlgUw,37833
+nucliadb/search/search/chat/ask.py,sha256=aaNj0MeAbx9dyeKpQJdm3VsHMq9OmcCESxahbgSxvCk,37805
 nucliadb/search/search/chat/exceptions.py,sha256=Siy4GXW2L7oPhIR86H3WHBhE9lkV4A4YaAszuGGUf54,1356
 nucliadb/search/search/chat/images.py,sha256=PA8VWxT5_HUGfW1ULhKTK46UBsVyINtWWqEM1ulzX1E,3095
-nucliadb/search/search/chat/prompt.py,sha256=Jnja-Ss7skgnnDY8BymVfdeYsFPnIQFL8tEvcRXTKUE,47356
+nucliadb/search/search/chat/prompt.py,sha256=e8C7_MPr6Cn3nJHA4hWpeW3629KVI1ZUQA_wZf9Kiu4,48503
 nucliadb/search/search/chat/query.py,sha256=6v6twBUTWfUUzklVV6xqJSYPkAshnIrBH9wbTcjQvkI,17063
 nucliadb/search/search/query_parser/__init__.py,sha256=cp15ZcFnHvpcu_5-aK2A4uUyvuZVV_MJn4bIXMa20ks,835
 nucliadb/search/search/query_parser/exceptions.py,sha256=szAOXUZ27oNY-OSa9t2hQ5HHkQQC0EX1FZz_LluJHJE,1224
@@ -368,8 +368,8 @@ nucliadb/writer/tus/local.py,sha256=7jYa_w9b-N90jWgN2sQKkNcomqn6JMVBOVeDOVYJHto,
 nucliadb/writer/tus/s3.py,sha256=vF0NkFTXiXhXq3bCVXXVV-ED38ECVoUeeYViP8uMqcU,8357
 nucliadb/writer/tus/storage.py,sha256=ToqwjoYnjI4oIcwzkhha_MPxi-k4Jk3Lt55zRwaC1SM,2903
 nucliadb/writer/tus/utils.py,sha256=MSdVbRsRSZVdkaum69_0wku7X3p5wlZf4nr6E0GMKbw,2556
-nucliadb-6.4.0.post4271.dist-info/METADATA,sha256=DKjN0N70XgUrzoOshpgsNU9xMWIweEdjtQgziYbveU0,4223
-nucliadb-6.4.0.post4271.dist-info/WHEEL,sha256=DnLRTWE75wApRYVsjgc6wsVswC54sMSJhAEd4xhDpBk,91
-nucliadb-6.4.0.post4271.dist-info/entry_points.txt,sha256=XqGfgFDuY3zXQc8ewXM2TRVjTModIq851zOsgrmaXx4,1268
-nucliadb-6.4.0.post4271.dist-info/top_level.txt,sha256=hwYhTVnX7jkQ9gJCkVrbqEG1M4lT2F_iPQND1fCzF80,20
-nucliadb-6.4.0.post4271.dist-info/RECORD,,
+nucliadb-6.4.0.post4276.dist-info/METADATA,sha256=tOuAZanwYJGX1BhhJIuqi8rnoevYyz3fey7vfyQ_sJ4,4223
+nucliadb-6.4.0.post4276.dist-info/WHEEL,sha256=DnLRTWE75wApRYVsjgc6wsVswC54sMSJhAEd4xhDpBk,91
+nucliadb-6.4.0.post4276.dist-info/entry_points.txt,sha256=XqGfgFDuY3zXQc8ewXM2TRVjTModIq851zOsgrmaXx4,1268
+nucliadb-6.4.0.post4276.dist-info/top_level.txt,sha256=hwYhTVnX7jkQ9gJCkVrbqEG1M4lT2F_iPQND1fCzF80,20
+nucliadb-6.4.0.post4276.dist-info/RECORD,,