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

nucliadb/common/http_clients/processing.py

@@ -25,6 +25,7 @@ import aiohttp
 import jwt
 import pydantic
 
+from nucliadb_utils.helpers import MessageProgressUpdater
 from nucliadb_utils.settings import nuclia_settings
 
 from .utils import check_status
@@ -51,6 +52,16 @@ def get_processing_api_url() -> str:
         return nuclia_settings.nuclia_processing_cluster_url + "/api/v1/internal/processing"
 
 
+def get_processing_api_v2_url() -> str:
+    if nuclia_settings.nuclia_service_account:
+        return (
+            nuclia_settings.nuclia_public_url.format(zone=nuclia_settings.nuclia_zone)
+            + "/api/v2/processing"
+        )
+    else:
+        return nuclia_settings.nuclia_processing_cluster_url + "/api/v2/internal/processing"
+
+
 class PullResponse(pydantic.BaseModel):
     status: str
     payload: Optional[str] = None
@@ -150,11 +161,38 @@ class StatsResponse(pydantic.BaseModel):
     scheduled: int
 
 
+class PullRequestV2(pydantic.BaseModel):
+    timeout: float = 5
+    limit: int = 1
+    ack: list[str] = []
+
+
+class InProgressRequest(pydantic.BaseModel):
+    ack: list[str] = []
+
+
+class PulledMessage(pydantic.BaseModel):
+    payload: bytes
+    headers: dict[str, str]
+    ack_token: str
+    seq: int
+
+
+class PullResponseV2(pydantic.BaseModel):
+    messages: list[PulledMessage]
+    ttl: float
+    pending: int
+
+
+JSON_HEADERS = {"Content-Type": "application/json"}
+
+
 class ProcessingHTTPClient:
     def __init__(self):
         self.session = aiohttp.ClientSession()
         self.base_url = get_processing_api_url()
-        self.headers = {}
+        self.base_url_v2 = get_processing_api_v2_url()
+        self.headers: dict[str, str] = {}
         if nuclia_settings.nuclia_service_account is not None:
             self.headers["X-STF-NUAKEY"] = f"Bearer {nuclia_settings.nuclia_service_account}"
 
@@ -193,6 +231,31 @@ class ProcessingHTTPClient:
             data = PullPosition.model_validate_json(resp_text)
             return data.cursor
 
+    async def in_progress(self, ack_token: str):
+        url = self.base_url_v2 + "/pull/in_progress"
+        request = InProgressRequest(ack=[ack_token])
+        async with self.session.post(
+            url, headers=self.headers | JSON_HEADERS, data=request.model_dump_json()
+        ) as resp:
+            resp_text = await resp.text()
+            check_status(resp, resp_text)
+
+    async def pull_v2(
+        self, ack_tokens: list[str], limit: int = 1, timeout: float = 5
+    ) -> Optional[PullResponseV2]:
+        url = self.base_url_v2 + "/pull"
+        request = PullRequestV2(limit=limit, timeout=timeout, ack=ack_tokens)
+        async with self.session.post(
+            url, headers=self.headers | JSON_HEADERS, data=request.model_dump_json()
+        ) as resp:
+            resp_text = await resp.text()
+            check_status(resp, resp_text)
+
+            if resp.status == 204:
+                return None
+            else:
+                return PullResponseV2.model_validate_json(resp_text)
+
     async def requests(
         self,
         cursor: Optional[str] = None,
@@ -225,3 +288,19 @@ class ProcessingHTTPClient:
             resp_text = await resp.text()
             check_status(resp, resp_text)
             return StatsResponse.model_validate_json(resp_text)
+
+
+class ProcessingPullMessageProgressUpdater(MessageProgressUpdater):
+    """
+    Context manager to send progress updates to NATS.
+
+    This should allow lower ack_wait time settings without causing
+    messages to be redelivered.
+    """
+
+    def __init__(self, client: ProcessingHTTPClient, msg: PulledMessage, timeout: float):
+        async def update_msg() -> bool:
+            await client.in_progress(msg.ack_token)
+            return False
+
+        super().__init__(str(msg.seq), update_msg, timeout)

nucliadb/ingest/app.py

@@ -32,7 +32,7 @@ from nucliadb.ingest.consumer import service as consumer_service
 from nucliadb.ingest.partitions import assign_partitions
 from nucliadb.ingest.processing import start_processing_engine, stop_processing_engine
 from nucliadb.ingest.service import start_grpc
-from nucliadb.ingest.settings import settings
+from nucliadb.ingest.settings import ProcessingPullMode, settings
 from nucliadb.ingest.utils import start_ingest as start_ingest_utility
 from nucliadb.ingest.utils import stop_ingest as stop_ingest_utility
 from nucliadb_telemetry import errors
@@ -101,7 +101,12 @@ async def initialize_grpc():  # pragma: no cover
 
 async def initialize_pull_workers() -> list[Callable[[], Awaitable[None]]]:
     finalizers = await initialize_grpc()
-    pull_workers = await consumer_service.start_pull_workers(SERVICE_NAME)
+    if settings.processing_pull_mode == ProcessingPullMode.V1:
+        pull_workers = await consumer_service.start_pull_workers(SERVICE_NAME)
+    elif settings.processing_pull_mode == ProcessingPullMode.V2:
+        pull_workers = [await consumer_service.start_ingest_processed_consumer_v2(SERVICE_NAME)]
+    else:
+        raise Exception("Processing pull workers not enabled and it is required")
 
     return pull_workers + finalizers
 
@@ -113,7 +118,11 @@ async def main_consumer():  # pragma: no cover
     grpc_health_finalizer = await health.start_grpc_health_service(settings.grpc_port)
 
     # pull workers could be pulled out into it's own deployment
-    pull_workers = await consumer_service.start_pull_workers(SERVICE_NAME)
+    if settings.processing_pull_mode == ProcessingPullMode.V1:
+        pull_workers = await consumer_service.start_pull_workers(SERVICE_NAME)
+    else:
+        # In v2, pull workers run inside the ingest consumer
+        pull_workers = []
     ingest_consumers = await consumer_service.start_ingest_consumers(SERVICE_NAME)
 
     await run_until_exit(
@@ -134,7 +143,16 @@ async def main_ingest_processed_consumer():  # pragma: no cover
     await start_processing_engine()
     metrics_server = await serve_metrics()
     grpc_health_finalizer = await health.start_grpc_health_service(settings.grpc_port)
-    consumer = await consumer_service.start_ingest_processed_consumer(SERVICE_NAME)
+
+    if settings.processing_pull_mode == ProcessingPullMode.V1:
+        consumer = await consumer_service.start_ingest_processed_consumer(SERVICE_NAME)
+    elif settings.processing_pull_mode == ProcessingPullMode.V2:
+        consumer = await consumer_service.start_ingest_processed_consumer_v2(SERVICE_NAME)
+    else:
+        # Off
+        async def fake_consumer(): ...
+
+        consumer = fake_consumer
 
     await run_until_exit(
         [grpc_health_finalizer, consumer, metrics_server.shutdown, stop_processing_engine] + finalizers

nucliadb/ingest/consumer/consumer.py

@@ -39,7 +39,7 @@ from nucliadb_protos.writer_pb2 import BrokerMessage, BrokerMessageBlobReference
 from nucliadb_telemetry import context, errors, metrics
 from nucliadb_utils import const
 from nucliadb_utils.cache.pubsub import PubSubDriver
-from nucliadb_utils.nats import MessageProgressUpdater, NatsConnectionManager
+from nucliadb_utils.nats import NatsConnectionManager, NatsMessageProgressUpdater
 from nucliadb_utils.settings import nats_consumer_settings
 from nucliadb_utils.storages.storage import Storage
 
@@ -181,7 +181,7 @@ class IngestConsumer:
         start = time.monotonic()
 
         async with (
-            MessageProgressUpdater(msg, nats_consumer_settings.nats_ack_wait * 0.66),
+            NatsMessageProgressUpdater(msg, nats_consumer_settings.nats_ack_wait * 0.66),
             self.lock,
         ):
             try:

nucliadb/ingest/consumer/pull.py

@@ -19,21 +19,34 @@
 #
 import asyncio
 import base64
+from contextlib import contextmanager
 from datetime import datetime, timezone
 from typing import Optional
 
 from aiohttp.client_exceptions import ClientConnectorError
+from opentelemetry import trace
+from opentelemetry.context import Context
+from opentelemetry.propagate import extract
+from opentelemetry.trace import (
+    Link,
+)
 
 from nucliadb.common import datamanagers
 from nucliadb.common.back_pressure.materializer import BackPressureMaterializer
 from nucliadb.common.back_pressure.utils import BackPressureException
-from nucliadb.common.http_clients.processing import ProcessingHTTPClient, get_nua_api_id
+from nucliadb.common.http_clients.processing import (
+    ProcessingHTTPClient,
+    ProcessingPullMessageProgressUpdater,
+    get_nua_api_id,
+)
 from nucliadb.common.maindb.driver import Driver
-from nucliadb.ingest import logger, logger_activity
+from nucliadb.ingest import SERVICE_NAME, logger, logger_activity
 from nucliadb.ingest.orm.exceptions import ReallyStopPulling
 from nucliadb.ingest.orm.processor import Processor
 from nucliadb_protos.writer_pb2 import BrokerMessage, BrokerMessageBlobReference
 from nucliadb_telemetry import errors
+from nucliadb_telemetry.metrics import Gauge
+from nucliadb_telemetry.utils import get_telemetry
 from nucliadb_utils import const
 from nucliadb_utils.cache.pubsub import PubSubDriver
 from nucliadb_utils.settings import nuclia_settings
@@ -41,6 +54,8 @@ from nucliadb_utils.storages.storage import Storage
 from nucliadb_utils.transaction import MaxTransactionSizeExceededError
 from nucliadb_utils.utilities import get_storage, get_transaction_utility
 
+processing_pending_messages = Gauge("nucliadb_processing_pending_messages")
+
 
 class PullWorker:
     """
@@ -234,3 +249,152 @@ class PullWorker:
                 except Exception:
                     logger.exception("Unhandled error pulling messages from processing")
                     await asyncio.sleep(self.pull_time_error_backoff)
+
+
+@contextmanager
+def run_in_span(headers: dict[str, str]):
+    # Create a span for handling this message
+    tracer_provider = get_telemetry(SERVICE_NAME)
+    if tracer_provider is None:
+        yield
+        return
+
+    tracer = tracer_provider.get_tracer(__name__)
+    our_span = tracer.start_span("handle_processing_pull")
+
+    # Try to retrieve processing context to link to it
+    witness = Context()
+    processor_context = extract(headers, context=witness)
+    if processor_context != witness:
+        # We successfully extracted a context, we link from the processor span to ours for ease of navigation
+        with tracer.start_as_current_span(
+            f"Pulled from proxy", links=[Link(our_span.get_span_context())], context=processor_context
+        ):
+            # And link from our span back to the processor span
+            our_span.add_link(trace.get_current_span().get_span_context())
+
+    # Go back to our context
+    trace.set_span_in_context(our_span)
+    with trace.use_span(our_span, end_on_exit=True):
+        yield
+
+
+class PullV2Worker:
+    """
+    The pull worker is responsible for pulling messages from the pull processing
+    http endpoint and processing them
+
+    The processing pull endpoint is also described as the "processing proxy" at times.
+    """
+
+    def __init__(
+        self,
+        driver: Driver,
+        storage: Storage,
+        pull_time_error_backoff: int,
+        pubsub: Optional[PubSubDriver] = None,
+        pull_time_empty_backoff: float = 5.0,
+        pull_api_timeout: int = 60,
+    ):
+        self.pull_time_error_backoff = pull_time_error_backoff
+        self.pull_time_empty_backoff = pull_time_empty_backoff
+        self.pull_api_timeout = pull_api_timeout
+
+        self.processor = Processor(driver, storage, pubsub, "-1")
+
+    async def handle_message(self, seq: int, payload: bytes) -> None:
+        pb = BrokerMessage()
+        data = base64.b64decode(payload)
+        pb.ParseFromString(data)
+
+        logger.debug(f"Resource: {pb.uuid} KB: {pb.kbid} ProcessingID: {pb.processing_id}")
+
+        await self.processor.process(
+            pb,
+            seq,
+            transaction_check=False,
+        )
+
+    async def loop(self):
+        """
+        Run this forever
+        """
+        while True:
+            try:
+                await self._loop()
+            except ReallyStopPulling:
+                logger.info("Exiting...")
+                break
+            except Exception as e:
+                errors.capture_exception(e)
+                logger.exception("Exception on worker", exc_info=e)
+                await asyncio.sleep(10)
+
+    async def _loop(self):
+        headers = {}
+        data = None
+        if nuclia_settings.nuclia_service_account is not None:
+            headers["X-STF-NUAKEY"] = f"Bearer {nuclia_settings.nuclia_service_account}"
+            # parse jwt sub to get pull type id
+            try:
+                get_nua_api_id()
+            except Exception as exc:
+                logger.exception("Could not read NUA API Key. Can not start pull worker")
+                raise ReallyStopPulling() from exc
+
+        ack_tokens = []
+        async with ProcessingHTTPClient() as processing_http_client:
+            while True:
+                try:
+                    # The code is only really prepared to pull 1 message at a time. If changing this, review MessageProgressUpdate usage
+                    pull = await processing_http_client.pull_v2(ack_tokens=ack_tokens, limit=1)
+                    ack_tokens.clear()
+                    if pull is None:
+                        processing_pending_messages.set(0)
+                        logger_activity.debug(f"No messages waiting in processing pull")
+                        await asyncio.sleep(self.pull_time_empty_backoff)
+                        continue
+
+                    logger.info("Message received from proxy", extra={"seq": [pull.messages[0].seq]})
+                    processing_pending_messages.set(pull.pending)
+                    try:
+                        for message in pull.messages:
+                            async with ProcessingPullMessageProgressUpdater(
+                                processing_http_client, message, pull.ttl * 0.66
+                            ):
+                                with run_in_span(message.headers):
+                                    await self.handle_message(message.seq, message.payload)
+                                    ack_tokens.append(message.ack_token)
+                    except Exception as e:
+                        errors.capture_exception(e)
+                        logger.exception("Error while pulling and processing message/s")
+                        raise e
+
+                except (
+                    asyncio.exceptions.CancelledError,
+                    RuntimeError,
+                    KeyboardInterrupt,
+                    SystemExit,
+                ):
+                    if ack_tokens:
+                        await processing_http_client.pull_v2(ack_tokens=ack_tokens, limit=0)
+                    logger.info(f"Pull task was canceled, exiting")
+                    raise ReallyStopPulling()
+
+                except ClientConnectorError:
+                    logger.error(
+                        f"Could not connect to processing engine, \
+                         {processing_http_client.base_url} verify your internet connection"
+                    )
+                    await asyncio.sleep(self.pull_time_error_backoff)
+
+                except MaxTransactionSizeExceededError as e:
+                    if data is not None:
+                        payload_length = 0
+                        if data.payload:
+                            payload_length = len(base64.b64decode(data.payload))
+                        logger.error(f"Message too big for transaction: {payload_length}")
+                    raise e
+                except Exception:
+                    logger.exception("Unhandled error pulling messages from processing")
+                    await asyncio.sleep(self.pull_time_error_backoff)

nucliadb/ingest/consumer/service.py

@@ -28,7 +28,7 @@ from nucliadb.common.back_pressure.utils import is_back_pressure_enabled
 from nucliadb.common.maindb.utils import setup_driver
 from nucliadb.ingest import SERVICE_NAME, logger
 from nucliadb.ingest.consumer.consumer import IngestConsumer, IngestProcessedConsumer
-from nucliadb.ingest.consumer.pull import PullWorker
+from nucliadb.ingest.consumer.pull import PullV2Worker, PullWorker
 from nucliadb.ingest.settings import settings
 from nucliadb_utils.exceptions import ConfigurationError
 from nucliadb_utils.settings import indexing_settings, transaction_settings
@@ -177,6 +177,32 @@ async def start_ingest_processed_consumer(
     return nats_connection_manager.finalize
 
 
+async def start_ingest_processed_consumer_v2(
+    service_name: Optional[str] = None,
+) -> Callable[[], Awaitable[None]]:
+    """
+    This is not meant to be deployed with a stateful set like the other consumers.
+
+    We are not maintaining transactionability based on the nats sequence id from this
+    consumer and we will start off by not separating writes by partition AND
+    allowing NATS to manage the queue group for us.
+    """
+    driver = await setup_driver()
+    pubsub = await get_pubsub()
+    storage = await get_storage(service_name=service_name or SERVICE_NAME)
+
+    consumer = PullV2Worker(
+        driver=driver,
+        storage=storage,
+        pubsub=pubsub,
+        pull_time_error_backoff=settings.pull_time_error_backoff,
+        pull_api_timeout=settings.pull_api_timeout,
+    )
+    task = asyncio.create_task(consumer.loop())
+    task.add_done_callback(_handle_task_result)
+    return partial(_exit_tasks, [task])
+
+
 async def start_auditor() -> Callable[[], Awaitable[None]]:
     audit = get_audit()
     assert audit is not None

nucliadb/ingest/settings.py

@@ -63,6 +63,13 @@ class DriverSettings(BaseSettings):
     )
 
 
+# For use during migration from pull v1 to pull v2
+class ProcessingPullMode(Enum):
+    OFF = "off"
+    V1 = "v1"
+    V2 = "v2"
+
+
 class Settings(DriverSettings):
     grpc_port: int = 8030
 
@@ -85,5 +92,7 @@ class Settings(DriverSettings):
 
     max_concurrent_ingest_processing: int = 5
 
+    processing_pull_mode: ProcessingPullMode = ProcessingPullMode.V1
+
 
 settings = Settings()

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/tasks/consumer.py

@@ -30,7 +30,7 @@ from nucliadb.tasks.logger import logger
 from nucliadb.tasks.models import Callback, MsgType
 from nucliadb.tasks.utils import NatsConsumer, NatsStream, create_nats_stream_if_not_exists
 from nucliadb_telemetry import errors
-from nucliadb_utils.nats import MessageProgressUpdater
+from nucliadb_utils.nats import NatsMessageProgressUpdater
 from nucliadb_utils.settings import nats_consumer_settings
 
 BEFORE_NAK_SLEEP_SECONDS = 2
@@ -124,7 +124,7 @@ class NatsTaskConsumer(Generic[MsgType]):
             f"Message received: subject:{subject}, seqid: {seqid}, reply: {reply}",
             extra={"consumer_name": self.name},
         )
-        async with MessageProgressUpdater(msg, nats_consumer_settings.nats_ack_wait * 0.66):
+        async with NatsMessageProgressUpdater(msg, nats_consumer_settings.nats_ack_wait * 0.66):
             try:
                 task_msg = self.msg_type.model_validate_json(msg.data)
             except pydantic.ValidationError as e:

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: nucliadb
-Version: 6.4.0.post4271
+Version: 6.4.0.post4279
 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.post4279
+Requires-Dist: nucliadb-utils[cache,fastapi,storages]>=6.4.0.post4279
+Requires-Dist: nucliadb-protos>=6.4.0.post4279
+Requires-Dist: nucliadb-models>=6.4.0.post4279
+Requires-Dist: nidx-protos>=6.4.0.post4279
 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.post4279.dist-info}/RECORD

@@ -101,7 +101,7 @@ nucliadb/common/external_index_providers/settings.py,sha256=EGHnIkwxqe6aypwKegXT
 nucliadb/common/http_clients/__init__.py,sha256=cp15ZcFnHvpcu_5-aK2A4uUyvuZVV_MJn4bIXMa20ks,835
 nucliadb/common/http_clients/auth.py,sha256=srfpgAbs2wmqA9u_l-HxsV4YoO77Tse4y3gm3q2YvYM,2112
 nucliadb/common/http_clients/exceptions.py,sha256=47Y8OjkaGV_F18G07FpJhOzgWKUIexhlILyuVtICz8s,1100
-nucliadb/common/http_clients/processing.py,sha256=Gj6smPfyQv4mQdWFSG9Zda2cLfQ-Dhad2LVeE9gh7oQ,7127
+nucliadb/common/http_clients/processing.py,sha256=VzxzFArNsHWGmFoX0c5OrQB3vFW841aeyuP5NgzPQGo,9581
 nucliadb/common/http_clients/pypi.py,sha256=VHIUjwJEJVntVUo_FRoXIo8sLmluy7sa9-iXSITcrMY,1540
 nucliadb/common/http_clients/utils.py,sha256=yGUkHNS41abHiBoHqo_Mg3QSqGsS7rUtbfGftbEC57U,1529
 nucliadb/common/maindb/__init__.py,sha256=cp15ZcFnHvpcu_5-aK2A4uUyvuZVV_MJn4bIXMa20ks,835
@@ -122,20 +122,20 @@ nucliadb/export_import/models.py,sha256=dbjScNkiMRv4X3Ktudy1JRliD25bfoDTy3JmEZgQ
 nucliadb/export_import/tasks.py,sha256=DWbdqY97ffoyfipelGXz3Jqz1iam6JCjQSh367Fc3NA,2947
 nucliadb/export_import/utils.py,sha256=8XOVMYXXw8b4ikojG7RjQ4tKN3Xu7nfu2yCUOqD50sk,23216
 nucliadb/ingest/__init__.py,sha256=fsw3C38VP50km3R-nHL775LNGPpJ4JxqXJ2Ib1f5SqE,1011
-nucliadb/ingest/app.py,sha256=BKmjpdBEskHcRIHwOnI_jG4gFGs6dV0KKVH9MLJeA48,7546
+nucliadb/ingest/app.py,sha256=Eympy8nbz09VDNPF28MuIeKMb7wgB9cTSOObS8uvL0o,8372
 nucliadb/ingest/partitions.py,sha256=2NIhMYbNT0TNBL6bX1UMSi7vxFGICstCKEqsB0TXHOE,2410
 nucliadb/ingest/processing.py,sha256=QmkHq-BU4vub7JRWe9VHvQ2DcAmT6-CzgFXuZxXhcBU,20953
 nucliadb/ingest/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 nucliadb/ingest/serialize.py,sha256=-TIjibJTbMqAowzRvyrG3R209vKqBZqXpdrQL9Dq4lo,16135
-nucliadb/ingest/settings.py,sha256=0B-wQNa8FLqtNcQgRzh-fuIuGptM816XHcbH1NQKfmE,3050
+nucliadb/ingest/settings.py,sha256=inB5SpkSI6sRd-ftlJIHFH6XlbuiSaRdL-F2WGyseUw,3249
 nucliadb/ingest/utils.py,sha256=l1myURu3r8oA11dx3GpHw-gNTUc1AFX8xdPm9Lgl2rA,2275
 nucliadb/ingest/consumer/__init__.py,sha256=cp15ZcFnHvpcu_5-aK2A4uUyvuZVV_MJn4bIXMa20ks,835
 nucliadb/ingest/consumer/auditing.py,sha256=xK21DIa_ZAiOJVVbnkmT4jgCRGshNGyPyxsqhE6kROE,7204
-nucliadb/ingest/consumer/consumer.py,sha256=OgS1fr5Yo55u-XbC6zypTH1aJ562Y1vZHnPDlJJpCXQ,13703
+nucliadb/ingest/consumer/consumer.py,sha256=GfdlrNlnt7PWYyk75xtyzn2SHZse7475U4U9q_9jKr0,13711
 nucliadb/ingest/consumer/materializer.py,sha256=tgD_rDI2twQzcz8kKNiW_L4YIth16IGh9mUfD5wiSD4,3858
 nucliadb/ingest/consumer/metrics.py,sha256=ji1l_4cKiHJthQd8YNem1ft4iMbw9KThmVvJmLcv3Xg,1075
-nucliadb/ingest/consumer/pull.py,sha256=vv1AyN0EhVgbgnZyT0D_1_IB4hWy7jPd4lAWPAOHGNc,10374
-nucliadb/ingest/consumer/service.py,sha256=GhuqlK-9Lvhzd8kBox8wOlKlJgM3W_gssKoWSfVVdoI,7897
+nucliadb/ingest/consumer/pull.py,sha256=gfdyQ8IMFA_bpGnEpmRB9qmOJywBEwxn7pGYaueZizU,16874
+nucliadb/ingest/consumer/service.py,sha256=WXBN8dY7MlmYWxqQHIbIO7w_SdVJRY1RuHAWlQUXf8o,8852
 nucliadb/ingest/consumer/shard_creator.py,sha256=w0smEu01FU_2cjZnsfBRNqT_Ntho11X17zTMST-vKbc,4359
 nucliadb/ingest/consumer/utils.py,sha256=jpX8D4lKzuPCpArQLZeX_Zczq3pfen_zAf8sPJfOEZU,2642
 nucliadb/ingest/fields/__init__.py,sha256=cp15ZcFnHvpcu_5-aK2A4uUyvuZVV_MJn4bIXMa20ks,835
@@ -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
@@ -290,7 +290,7 @@ nucliadb/standalone/static/favicon.ico,sha256=96pKGp6Sx457JkTfjy1dtApMhkitixfU6i
 nucliadb/standalone/static/index.html,sha256=PEZfuEQFYnYACAL1ceN8xC0im8lBrUx838RkE8tbvgA,3833
 nucliadb/standalone/static/logo.svg,sha256=-wQqSvPGTdlKjUP6pHE6kiq005pgYjDzp9nPl0X71Mk,2639
 nucliadb/tasks/__init__.py,sha256=oFJ3A8HD7w11mBu-IixYE_KxA7juMGlYQb7YD_y6WPM,975
-nucliadb/tasks/consumer.py,sha256=4CWfBdXVr2a25n7seldbQ0PaK0FcxJZuWgosU6aODS8,6956
+nucliadb/tasks/consumer.py,sha256=E7_9bY5o7BVlioWX9yO9yimDJaKeuj-P-tiNCJcaRz8,6964
 nucliadb/tasks/logger.py,sha256=C7keOEO_mjLVp5VbqAZ2QXfqVB2Hot7NgBlUP_SDSMw,924
 nucliadb/tasks/models.py,sha256=qrZKi5DNDQ07waMsp5L4_Fi7WRs57YiO-kmXlrBzEAA,1168
 nucliadb/tasks/producer.py,sha256=UnpJAzhj_GElsCoO5G6T4m6MshsgOaqR2tVzJmEta64,2625
@@ -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.post4279.dist-info/METADATA,sha256=ISm2mlidMfyHGlEXRXBJcbnPe52rP58sBjj5NLRFf68,4223
+nucliadb-6.4.0.post4279.dist-info/WHEEL,sha256=DnLRTWE75wApRYVsjgc6wsVswC54sMSJhAEd4xhDpBk,91
+nucliadb-6.4.0.post4279.dist-info/entry_points.txt,sha256=XqGfgFDuY3zXQc8ewXM2TRVjTModIq851zOsgrmaXx4,1268
+nucliadb-6.4.0.post4279.dist-info/top_level.txt,sha256=hwYhTVnX7jkQ9gJCkVrbqEG1M4lT2F_iPQND1fCzF80,20
+nucliadb-6.4.0.post4279.dist-info/RECORD,,