nucliadb 6.6.1.post4573__py3-none-any.whl → 6.6.1.post4584__py3-none-any.whl

nucliadb/search/api/v1/ask.py

@@ -59,6 +59,7 @@ async def ask_knowledgebox_endpoint(
     kbid: str,
     item: AskRequest,
     x_ndb_client: NucliaDBClientType = Header(NucliaDBClientType.API),
+    x_show_consumption: bool = Header(default=False),
     x_nucliadb_user: str = Header(""),
     x_forwarded_for: str = Header(""),
     x_synchronous: bool = Header(
@@ -97,7 +98,13 @@ async def ask_knowledgebox_endpoint(
             return HTTPClientError(status_code=422, detail=detail)
 
     return await create_ask_response(
-        kbid, item, x_nucliadb_user, x_ndb_client, x_forwarded_for, x_synchronous
+        kbid=kbid,
+        ask_request=item,
+        user_id=x_nucliadb_user,
+        client_type=x_ndb_client,
+        origin=x_forwarded_for,
+        x_synchronous=x_synchronous,
+        extra_predict_headers={"X-Show-Consumption": str(x_show_consumption).lower()},
     )
 
 
@@ -110,6 +117,7 @@ async def create_ask_response(
     origin: str,
     x_synchronous: bool,
     resource: Optional[str] = None,
+    extra_predict_headers: Optional[dict[str, str]] = None,
 ) -> Response:
     maybe_log_request_payload(kbid, "/ask", ask_request)
     ask_request.max_tokens = parse_max_tokens(ask_request.max_tokens)
@@ -122,6 +130,7 @@ async def create_ask_response(
                 client_type=client_type,
                 origin=origin,
                 resource=resource,
+                extra_predict_headers=extra_predict_headers,
             )
         except AnswerJsonSchemaTooLong as err:
             return HTTPClientError(status_code=400, detail=str(err))

nucliadb/search/api/v1/resource/ask.py

@@ -48,6 +48,7 @@ async def resource_ask_endpoint_by_uuid(
     kbid: str,
     rid: str,
     item: AskRequest,
+    x_show_consumption: bool = Header(default=False),
     x_ndb_client: NucliaDBClientType = Header(NucliaDBClientType.API),
     x_nucliadb_user: str = Header(""),
     x_forwarded_for: str = Header(""),
@@ -58,13 +59,14 @@ async def resource_ask_endpoint_by_uuid(
     ),
 ) -> Union[StreamingResponse, HTTPClientError, Response]:
     return await create_ask_response(
-        kbid,
-        item,
-        x_nucliadb_user,
-        x_ndb_client,
-        x_forwarded_for,
-        x_synchronous,
+        kbid=kbid,
+        ask_request=item,
+        user_id=x_nucliadb_user,
+        client_type=x_ndb_client,
+        origin=x_forwarded_for,
+        x_synchronous=x_synchronous,
         resource=rid,
+        extra_predict_headers={"X-Show-Consumption": str(x_show_consumption).lower()},
     )
 
 
@@ -83,6 +85,7 @@ async def resource_ask_endpoint_by_slug(
     kbid: str,
     slug: str,
     item: AskRequest,
+    x_show_consumption: bool = Header(default=False),
     x_ndb_client: NucliaDBClientType = Header(NucliaDBClientType.API),
     x_nucliadb_user: str = Header(""),
     x_forwarded_for: str = Header(""),
@@ -96,11 +99,12 @@ async def resource_ask_endpoint_by_slug(
     if resource_id is None:
         return HTTPClientError(status_code=404, detail="Resource not found")
     return await create_ask_response(
-        kbid,
-        item,
-        x_nucliadb_user,
-        x_ndb_client,
-        x_forwarded_for,
-        x_synchronous,
+        kbid=kbid,
+        ask_request=item,
+        user_id=x_nucliadb_user,
+        client_type=x_ndb_client,
+        origin=x_forwarded_for,
+        x_synchronous=x_synchronous,
         resource=resource_id,
+        extra_predict_headers={"X-Show-Consumption": str(x_show_consumption).lower()},
     )

nucliadb/search/api/v1/summarize.py

@@ -19,7 +19,7 @@
 #
 from typing import Union
 
-from fastapi import Request
+from fastapi import Header, Request
 from fastapi_versioning import version
 
 from nucliadb.common.datamanagers.exceptions import KnowledgeBoxNotFound
@@ -47,9 +47,14 @@ async def summarize_endpoint(
     request: Request,
     kbid: str,
     item: SummarizeRequest,
+    x_show_consumption: bool = Header(default=False),
 ) -> Union[SummarizedResponse, HTTPClientError]:
     try:
-        return await summarize(kbid, item)
+        return await summarize(
+            kbid=kbid,
+            request=item,
+            extra_predict_headers={"X-Show-Consumption": str(x_show_consumption).lower()},
+        )
     except KnowledgeBoxNotFound:
         return HTTPClientError(status_code=404, detail="Knowledge box not found")
     except NoResourcesToSummarize:

nucliadb/search/predict.py

@@ -293,7 +293,7 @@ class PredictEngine:
 
     @predict_observer.wrap({"type": "chat_ndjson"})
     async def chat_query_ndjson(
-        self, kbid: str, item: ChatModel
+        self, kbid: str, item: ChatModel, extra_headers: Optional[dict[str, str]] = None
     ) -> tuple[str, str, AsyncGenerator[GenerativeChunk, None]]:
         """
         Chat query using the new stream format
@@ -314,7 +314,7 @@ class PredictEngine:
             "POST",
             url=self.get_predict_url(CHAT, kbid),
             json=item.model_dump(),
-            headers=headers,
+            headers={**headers, **(extra_headers or {})},
             timeout=None,
         )
         await self.check_response(kbid, resp, expected_status=200)
@@ -396,7 +396,9 @@ class PredictEngine:
         return convert_relations(data)
 
     @predict_observer.wrap({"type": "summarize"})
-    async def summarize(self, kbid: str, item: SummarizeModel) -> SummarizedResponse:
+    async def summarize(
+        self, kbid: str, item: SummarizeModel, extra_headers: Optional[dict[str, str]] = None
+    ) -> SummarizedResponse:
         try:
             self.check_nua_key_is_configured_for_onprem()
         except NUAKeyMissingError:
@@ -407,7 +409,7 @@ class PredictEngine:
             "POST",
             url=self.get_predict_url(SUMMARIZE, kbid),
             json=item.model_dump(),
-            headers=self.get_predict_headers(kbid),
+            headers={**self.get_predict_headers(kbid), **(extra_headers or {})},
             timeout=None,
         )
         await self.check_response(kbid, resp, expected_status=200)
@@ -489,7 +491,7 @@ class DummyPredictEngine(PredictEngine):
         return RephraseResponse(rephrased_query=DUMMY_REPHRASE_QUERY, use_chat_history=None)
 
     async def chat_query_ndjson(
-        self, kbid: str, item: ChatModel
+        self, kbid: str, item: ChatModel, extra_headers: Optional[dict[str, str]] = None
     ) -> tuple[str, str, AsyncGenerator[GenerativeChunk, None]]:
         self.calls.append(("chat_query_ndjson", item))
 
@@ -559,7 +561,9 @@ class DummyPredictEngine(PredictEngine):
         else:
             return DUMMY_RELATION_NODE
 
-    async def summarize(self, kbid: str, item: SummarizeModel) -> SummarizedResponse:
+    async def summarize(
+        self, kbid: str, item: SummarizeModel, extra_headers: Optional[dict[str, str]] = None
+    ) -> SummarizedResponse:
         self.calls.append(("summarize", (kbid, item)))
         response = SummarizedResponse(
             summary="global summary",

nucliadb/search/search/chat/ask.py

@@ -22,6 +22,7 @@ import functools
 import json
 from typing import AsyncGenerator, Optional, cast
 
+from nuclia_models.common.consumption import Consumption
 from nuclia_models.predict.generative_responses import (
     CitationsGenerativeResponse,
     GenerativeChunk,
@@ -83,6 +84,7 @@ from nucliadb_models.search import (
     ChatModel,
     ChatOptions,
     CitationsAskResponseItem,
+    ConsumptionResponseItem,
     DebugAskResponseItem,
     ErrorAskResponseItem,
     FindOptions,
@@ -106,6 +108,7 @@ from nucliadb_models.search import (
     StatusAskResponseItem,
     SyncAskMetadata,
     SyncAskResponse,
+    TokensDetail,
     UserPrompt,
     parse_custom_prompt,
     parse_rephrase_prompt,
@@ -169,6 +172,7 @@ class AskResult:
         self._citations: Optional[CitationsGenerativeResponse] = None
         self._metadata: Optional[MetaGenerativeResponse] = None
         self._relations: Optional[Relations] = None
+        self._consumption: Optional[Consumption] = None
 
     @property
     def status_code(self) -> AnswerStatusCode:
@@ -299,6 +303,20 @@ class AskResult:
                 ),
             )
 
+        if self._consumption is not None:
+            yield ConsumptionResponseItem(
+                normalized_tokens=TokensDetail(
+                    input=self._consumption.normalized_tokens.input,
+                    output=self._consumption.normalized_tokens.output,
+                    image=self._consumption.normalized_tokens.image,
+                ),
+                customer_key_tokens=TokensDetail(
+                    input=self._consumption.customer_key_tokens.input,
+                    output=self._consumption.customer_key_tokens.output,
+                    image=self._consumption.customer_key_tokens.image,
+                ),
+            )
+
         # Stream out the relations results
         should_query_relations = (
             self.ask_request_with_relations and self.status_code == AnswerStatusCode.SUCCESS
@@ -341,6 +359,7 @@ class AskResult:
                     generative_total=self._metadata.timings.get("generative"),
                 ),
             )
+
         citations = {}
         if self._citations is not None:
             citations = self._citations.citations
@@ -373,6 +392,7 @@ class AskResult:
             prequeries=prequeries_results,
             citations=citations,
             metadata=metadata,
+            consumption=self._consumption,
             learning_id=self.nuclia_learning_id or "",
             augmented_context=self.augmented_context,
         )
@@ -424,6 +444,8 @@ class AskResult:
                 self._citations = item
             elif isinstance(item, MetaGenerativeResponse):
                 self._metadata = item
+            elif isinstance(item, Consumption):
+                self._consumption = item
             else:
                 logger.warning(
                     f"Unexpected item in predict answer stream: {item}",
@@ -486,6 +508,7 @@ async def ask(
     client_type: NucliaDBClientType,
     origin: str,
     resource: Optional[str] = None,
+    extra_predict_headers: Optional[dict[str, str]] = None,
 ) -> AskResult:
     metrics = AskMetrics()
     chat_history = ask_request.chat_history or []
@@ -613,7 +636,9 @@ async def ask(
                 nuclia_learning_id,
                 nuclia_learning_model,
                 predict_answer_stream,
-            ) = await predict.chat_query_ndjson(kbid, chat_model)
+            ) = await predict.chat_query_ndjson(
+                kbid=kbid, item=chat_model, extra_headers=extra_predict_headers
+            )
 
     auditor = ChatAuditor(
         kbid=kbid,

nucliadb/search/search/predict_proxy.py

@@ -63,6 +63,7 @@ class PredictProxiedEndpoints(str, Enum):
 
 ALLOWED_HEADERS = [
     "Accept",  # To allow 'application/x-ndjson' on the /chat endpoint
+    "X-show-consumption",  # To show token consumption in the response
 ]
 
 PREDICT_ANSWER_METRIC = "predict_answer_proxy_metric"

nucliadb/search/search/summarize.py

@@ -45,7 +45,9 @@ class NoResourcesToSummarize(Exception):
     pass
 
 
-async def summarize(kbid: str, request: SummarizeRequest) -> SummarizedResponse:
+async def summarize(
+    kbid: str, request: SummarizeRequest, extra_predict_headers: Optional[dict[str, str]]
+) -> SummarizedResponse:
     predict_request = SummarizeModel()
     predict_request.generative_model = request.generative_model
     predict_request.user_prompt = request.user_prompt
@@ -62,7 +64,7 @@ async def summarize(kbid: str, request: SummarizeRequest) -> SummarizedResponse:
         raise NoResourcesToSummarize()
 
     predict = get_predict()
-    return await predict.summarize(kbid, predict_request)
+    return await predict.summarize(kbid=kbid, item=predict_request, extra_headers=extra_predict_headers)
 
 
 async def get_extracted_texts(kbid: str, resource_uuids_or_slugs: list[str]) -> ExtractedTexts:

{nucliadb-6.6.1.post4573.dist-info → nucliadb-6.6.1.post4584.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: nucliadb
-Version: 6.6.1.post4573
+Version: 6.6.1.post4584
 Summary: NucliaDB
 Author-email: Nuclia <nucliadb@nuclia.com>
 License-Expression: AGPL-3.0-or-later
@@ -19,13 +19,13 @@ 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.6.1.post4573
-Requires-Dist: nucliadb-utils[cache,fastapi,storages]>=6.6.1.post4573
-Requires-Dist: nucliadb-protos>=6.6.1.post4573
-Requires-Dist: nucliadb-models>=6.6.1.post4573
-Requires-Dist: nidx-protos>=6.6.1.post4573
+Requires-Dist: nucliadb-telemetry[all]>=6.6.1.post4584
+Requires-Dist: nucliadb-utils[cache,fastapi,storages]>=6.6.1.post4584
+Requires-Dist: nucliadb-protos>=6.6.1.post4584
+Requires-Dist: nucliadb-models>=6.6.1.post4584
+Requires-Dist: nidx-protos>=6.6.1.post4584
 Requires-Dist: nucliadb-admin-assets>=1.0.0.post1224
-Requires-Dist: nuclia-models>=0.24.2
+Requires-Dist: nuclia-models>=0.43.0
 Requires-Dist: uvicorn[standard]
 Requires-Dist: argdantic
 Requires-Dist: aiohttp>=3.11.11

{nucliadb-6.6.1.post4573.dist-info → nucliadb-6.6.1.post4584.dist-info}/RECORD

@@ -212,7 +212,7 @@ nucliadb/search/__init__.py,sha256=tnypbqcH4nBHbGpkINudhKgdLKpwXQCvDtPchUlsyY4,1
 nucliadb/search/app.py,sha256=-WEX1AZRA8R_9aeOo9ovOTwjXW_7VfwWN7N2ccSoqXg,3387
 nucliadb/search/lifecycle.py,sha256=hiylV-lxsAWkqTCulXBg0EIfMQdejSr8Zar0L_GLFT8,2218
 nucliadb/search/openapi.py,sha256=t3Wo_4baTrfPftg2BHsyLWNZ1MYn7ZRdW7ht-wFOgRs,1016
-nucliadb/search/predict.py,sha256=__0qwIU2CIRYRTYsbG9zZEjXXrxNe8puZWYJIyOT6dg,23492
+nucliadb/search/predict.py,sha256=xZtZaydg1pzXOSEDg0xyWNbbgA4zMQ59gbHi0wNuAxk,23770
 nucliadb/search/predict_models.py,sha256=pm4ykuWH9bTXxj5RlI2F6pmXSXOVt64WL_sRlc2u6Tk,6144
 nucliadb/search/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 nucliadb/search/run.py,sha256=aFb-CXRi_C8YMpP_ivNj8KW1BYhADj88y8K9Lr_nUPI,1402
@@ -220,7 +220,7 @@ nucliadb/search/settings.py,sha256=vem3EcyYlTPSim0kEK-xe-erF4BZg0CT_LAb8ZRQAE8,1
 nucliadb/search/utilities.py,sha256=9SsRDw0rJVXVoLBfF7rBb6q080h-thZc7u8uRcTiBeY,1037
 nucliadb/search/api/__init__.py,sha256=cp15ZcFnHvpcu_5-aK2A4uUyvuZVV_MJn4bIXMa20ks,835
 nucliadb/search/api/v1/__init__.py,sha256=DH16OYnw9jQ38OpKlmdXeoq2j40ZPXZRtGvClKOkMhw,1239
-nucliadb/search/api/v1/ask.py,sha256=b4tz33HNsfT5DXv_2DMc_jirnFsHuobreWkbAKkzj5o,5337
+nucliadb/search/api/v1/ask.py,sha256=hZUnk1opZuXp1IwTiingSatlUefg2CZ9r_Z9sUwZMaU,5698
 nucliadb/search/api/v1/catalog.py,sha256=5ZY3d8sVia1traUxVS0Q4aQJmgcOuXzbxis_uY4ulE4,8077
 nucliadb/search/api/v1/feedback.py,sha256=kNLc4dHz2SXHzV0PwC1WiRAwY88fDptPcP-kO0q-FrQ,2620
 nucliadb/search/api/v1/find.py,sha256=j6mxEyxjlLnZSqCT_N2LmOJlytsm1vkY4KFFmJRrtP8,10904
@@ -230,10 +230,10 @@ nucliadb/search/api/v1/predict_proxy.py,sha256=TnXKAqf_Go-9QVi6L5z4cXjnuNRe7XLJj
 nucliadb/search/api/v1/router.py,sha256=mtT07rBZcVfpa49doaw9b1tj3sdi3qLH0gn9Io6NYM0,988
 nucliadb/search/api/v1/search.py,sha256=eqlrvRE7IlMpunNwD1RJwt6RgMV01sIDJLgxxE7CFcE,12297
 nucliadb/search/api/v1/suggest.py,sha256=gaJE60r8-z6TVO05mQRKBITwXn2_ofM3B4-OtpOgZEk,6343
-nucliadb/search/api/v1/summarize.py,sha256=VAHJvE6V3xUgEBfqNKhgoxmDqCvh30RnrEIBVhMcNLU,2499
+nucliadb/search/api/v1/summarize.py,sha256=eJzgFJWUO80STx3lHc_0h9RZVaBCWF196nZUecfmqbE,2700
 nucliadb/search/api/v1/utils.py,sha256=5Ve-frn7LAE2jqAgB85F8RSeqxDlyA08--gS-AdOLS4,1434
 nucliadb/search/api/v1/resource/__init__.py,sha256=cp15ZcFnHvpcu_5-aK2A4uUyvuZVV_MJn4bIXMa20ks,835
-nucliadb/search/api/v1/resource/ask.py,sha256=nsVzBSanSSlf0Ody6LSTjdEy75Vg283_YhbkAtWEjh8,3637
+nucliadb/search/api/v1/resource/ask.py,sha256=PlOXa17lnmj3KA9bARNfDqvnx7Pe9OTnwz-OwgGTUjU,4035
 nucliadb/search/api/v1/resource/ingestion_agents.py,sha256=AZ5_cH1jbf7d5wh_gz6EHLEKAzEOMrQZwEZAu1Q_3FE,4846
 nucliadb/search/api/v1/resource/search.py,sha256=PZR7fs5oYD0RKqKoD38NZMAnOJzBv35NB2YOr2xy1ck,4923
 nucliadb/search/api/v1/resource/utils.py,sha256=-NjZqAQtFEXKpIh8ui5S26ItnJ5rzmmG0BHxGSS9QPw,1141
@@ -255,15 +255,15 @@ nucliadb/search/search/merge.py,sha256=XiRBsxhYPshPV7lZXD-9E259KZOPIf4I2tKosY0lP
 nucliadb/search/search/metrics.py,sha256=3I6IN0qDSmqIvUaWJmT3rt-Jyjs6LcvnKI8ZqCiuJPY,3501
 nucliadb/search/search/paragraphs.py,sha256=pNAEiYqJGGUVcEf7xf-PFMVqz0PX4Qb-WNG-_zPGN2o,7799
 nucliadb/search/search/pgcatalog.py,sha256=_AiyW6it66UX6BsZbM3-230IQhiEG4utoKYboviyOFI,16799
-nucliadb/search/search/predict_proxy.py,sha256=jJj9LrxBOFo5AMn_EB7MNo1hqiTN-PkeBYINIpVKzl8,8735
+nucliadb/search/search/predict_proxy.py,sha256=Df8F5K-oS4TIXJc_y8UDViJTo7st5L0kMgxYPFZ39Vk,8806
 nucliadb/search/search/query.py,sha256=0qIQdt548L3jtKOyKo06aGJ73SLBxAW3N38_Hc1M3Uw,11528
 nucliadb/search/search/rank_fusion.py,sha256=xZtXhbmKb_56gs73u6KkFm2efvTATOSMmpOV2wrAIqE,9613
 nucliadb/search/search/rerankers.py,sha256=E2J1QdKAojqbhHM3KAyaOXKf6tJyETUxKs4tf_BEyqk,7472
 nucliadb/search/search/shards.py,sha256=mc5DK-MoCv9AFhlXlOFHbPvetcyNDzTFOJ5rimK8PC8,2636
-nucliadb/search/search/summarize.py,sha256=ksmYPubEQvAQgfPdZHfzB_rR19B2ci4IYZ6jLdHxZo8,4996
+nucliadb/search/search/summarize.py,sha256=S4-mUS8d-rvHFcsr8Pa8N5NTxU6ZTxLFZTMKTTOOpr4,5098
 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=0sgfiCbNaCZrTvYaRGtf5xL6VnzRgzofINiEP4IvhWs,38278
+nucliadb/search/search/chat/ask.py,sha256=vJ3TSdr-cT_xh43UnoYugqxnHv_-LFSCYoU7o0NnI1M,39368
 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=QwHULUDqe_pS2HZvQH1vzqpYEHQG_-UagXCNtLLtJEI,52997
@@ -375,8 +375,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.6.1.post4573.dist-info/METADATA,sha256=Qs_hupFPWx7a2jlWX2rRG9RNDF9avNHSxN079rwtP5E,4158
-nucliadb-6.6.1.post4573.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-nucliadb-6.6.1.post4573.dist-info/entry_points.txt,sha256=XqGfgFDuY3zXQc8ewXM2TRVjTModIq851zOsgrmaXx4,1268
-nucliadb-6.6.1.post4573.dist-info/top_level.txt,sha256=hwYhTVnX7jkQ9gJCkVrbqEG1M4lT2F_iPQND1fCzF80,20
-nucliadb-6.6.1.post4573.dist-info/RECORD,,
+nucliadb-6.6.1.post4584.dist-info/METADATA,sha256=1YXr0ZLW8EuBsP9B8fzdG8R__ktC5ANI378Dw50dRjM,4158
+nucliadb-6.6.1.post4584.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+nucliadb-6.6.1.post4584.dist-info/entry_points.txt,sha256=XqGfgFDuY3zXQc8ewXM2TRVjTModIq851zOsgrmaXx4,1268
+nucliadb-6.6.1.post4584.dist-info/top_level.txt,sha256=hwYhTVnX7jkQ9gJCkVrbqEG1M4lT2F_iPQND1fCzF80,20
+nucliadb-6.6.1.post4584.dist-info/RECORD,,