nuclia 4.9.3__py3-none-any.whl → 4.9.4__py3-none-any.whl

nuclia/lib/kb.py

@@ -324,13 +324,18 @@ class NucliaDBClient(BaseNucliaDBClient):
         handle_http_sync_errors(response)
         return int(response.headers.get("Upload-Offset"))
 
-    def summarize(self, request: SummarizeRequest, timeout: int = 1000):
+    def summarize(
+        self,
+        request: SummarizeRequest,
+        extra_headers: Optional[dict[str, str]] = None,
+        timeout: int = 1000,
+    ):
         if self.url is None or self.writer_session is None:
             raise Exception("KB not configured")
         url = f"{self.url}{SUMMARIZE_URL}"
         assert self.reader_session
         response = self.reader_session.post(
-            url, json=request.model_dump(), timeout=timeout
+            url, json=request.model_dump(), headers=extra_headers, timeout=timeout
         )
         handle_http_sync_errors(response)
         return response
@@ -569,12 +574,21 @@ class AsyncNucliaDBClient(BaseNucliaDBClient):
         await handle_http_async_errors(response)
         return response
 
-    async def ask(self, request: AskRequest, timeout: int = 1000):
+    async def ask(
+        self,
+        request: AskRequest,
+        extra_headers: Optional[dict[str, str]] = None,
+        timeout: int = 1000,
+    ):
         if self.url is None or self.reader_session is None:
             raise Exception("KB not configured")
         url = f"{self.url}{ASK_URL}"
         req = self.reader_session.build_request(
-            "POST", url, json=request.model_dump(), timeout=timeout
+            "POST",
+            url,
+            json=request.model_dump(),
+            headers=extra_headers,
+            timeout=timeout,
         )
         response = await self.reader_session.send(req, stream=True)
         await handle_http_async_errors(response)
@@ -681,13 +695,18 @@ class AsyncNucliaDBClient(BaseNucliaDBClient):
         await handle_http_async_errors(response)
         return int(response.headers.get("Upload-Offset"))
 
-    async def summarize(self, request: SummarizeRequest, timeout: int = 1000):
+    async def summarize(
+        self,
+        request: SummarizeRequest,
+        extra_headers: Optional[dict[str, str]] = None,
+        timeout: int = 1000,
+    ):
         if self.url is None or self.writer_session is None:
             raise Exception("KB not configured")
         url = f"{self.url}{SUMMARIZE_URL}"
         assert self.reader_session
         response = await self.reader_session.post(
-            url, json=request.model_dump(), timeout=timeout
+            url, json=request.model_dump(), headers=extra_headers, timeout=timeout
         )
         await handle_http_async_errors(response)
         return response

nuclia/lib/nua.py

@@ -18,6 +18,7 @@ from pydantic import BaseModel
 
 from nuclia import REGIONAL
 from nuclia.exceptions import NuaAPIException
+from nuclia_models.common.consumption import ConsumptionGenerative
 from nuclia_models.predict.generative_responses import (
     GenerativeChunk,
     GenerativeFullResponse,
@@ -28,6 +29,7 @@ from nuclia_models.predict.generative_responses import (
     StatusGenerativeResponse,
     ToolsGenerativeResponse,
 )
+from nuclia_models.common.consumption import Consumption
 from nuclia.lib.nua_responses import (
     ChatModel,
     ChatResponse,
@@ -128,9 +130,12 @@ class NuaClient:
         url: str,
         output: Type[ConvertType],
         payload: Optional[dict[Any, Any]] = None,
+        extra_headers: Optional[dict[str, str]] = None,
         timeout: int = 60,
     ) -> ConvertType:
-        resp = self.client.request(method, url, json=payload, timeout=timeout)
+        resp = self.client.request(
+            method, url, json=payload, timeout=timeout, headers=extra_headers
+        )
         if resp.status_code != 200:
             raise NuaAPIException(code=resp.status_code, detail=resp.content.decode())
         try:
@@ -143,6 +148,7 @@ class NuaClient:
         self,
         method: str,
         url: str,
+        extra_headers: Optional[dict[str, str]] = None,
         payload: Optional[dict[Any, Any]] = None,
         timeout: int = 60,
     ) -> Iterator[GenerativeChunk]:
@@ -151,8 +157,9 @@ class NuaClient:
             url,
             json=payload,
             timeout=timeout,
+            headers=extra_headers,
         ) as response:
-            if response.headers.get("content-type") == "application/x-ndjson":
+            if response.headers.get("transfer-encoding") == "chunked":
                 for json_body in response.iter_lines():
                     try:
                         yield GenerativeChunk.model_validate_json(json_body)  # type: ignore
@@ -194,17 +201,31 @@ class NuaClient:
             endpoint = f"{self.url}{CONFIG}/{kbid}"
         return self._request("GET", endpoint, output=StoredLearningConfiguration)
 
-    def sentence_predict(self, text: str, model: Optional[str] = None) -> Sentence:
+    def sentence_predict(
+        self,
+        text: str,
+        model: Optional[str] = None,
+        extra_headers: Optional[dict[str, str]] = None,
+    ) -> Sentence:
         endpoint = f"{self.url}{SENTENCE_PREDICT}?text={text}"
         if model:
             endpoint += f"&model={model}"
-        return self._request("GET", endpoint, output=Sentence)
+        return self._request(
+            "GET", endpoint, output=Sentence, extra_headers=extra_headers
+        )
 
-    def tokens_predict(self, text: str, model: Optional[str] = None) -> Tokens:
+    def tokens_predict(
+        self,
+        text: str,
+        model: Optional[str] = None,
+        extra_headers: Optional[dict[str, str]] = None,
+    ) -> Tokens:
         endpoint = f"{self.url}{TOKENS_PREDICT}?text={text}"
         if model:
             endpoint += f"&model={model}"
-        return self._request("GET", endpoint, output=Tokens)
+        return self._request(
+            "GET", endpoint, output=Tokens, extra_headers=extra_headers
+        )
 
     def query_predict(
         self,
@@ -212,6 +233,7 @@ class NuaClient:
         semantic_model: Optional[str] = None,
         token_model: Optional[str] = None,
         generative_model: Optional[str] = None,
+        extra_headers: Optional[dict[str, str]] = None,
     ) -> QueryInfo:
         endpoint = f"{self.url}{QUERY_PREDICT}?text={text}"
         if semantic_model:
@@ -220,10 +242,16 @@ class NuaClient:
             endpoint += f"&token_model={token_model}"
         if generative_model:
             endpoint += f"&generative_model={generative_model}"
-        return self._request("GET", endpoint, output=QueryInfo)
+        return self._request(
+            "GET", endpoint, output=QueryInfo, extra_headers=extra_headers
+        )
 
     def generate(
-        self, body: ChatModel, model: Optional[str] = None, timeout: int = 300
+        self,
+        body: ChatModel,
+        model: Optional[str] = None,
+        extra_headers: Optional[dict[str, str]] = None,
+        timeout: int = 300,
     ) -> GenerativeFullResponse:
         endpoint = f"{self.url}{CHAT_PREDICT}"
         if model:
@@ -235,6 +263,7 @@ class NuaClient:
             endpoint,
             payload=body.model_dump(),
             timeout=timeout,
+            extra_headers=extra_headers,
         ):
             if isinstance(chunk.chunk, TextGenerativeResponse):
                 result.answer += chunk.chunk.text
@@ -252,10 +281,19 @@ class NuaClient:
                 result.code = chunk.chunk.code
             elif isinstance(chunk.chunk, ToolsGenerativeResponse):
                 result.tools = chunk.chunk.tools
+            elif isinstance(chunk.chunk, ConsumptionGenerative):
+                result.consumption = Consumption(
+                    normalized_tokens=chunk.chunk.normalized_tokens,
+                    customer_key_tokens=chunk.chunk.customer_key_tokens,
+                )
         return result
 
     def generate_stream(
-        self, body: ChatModel, model: Optional[str] = None, timeout: int = 300
+        self,
+        body: ChatModel,
+        model: Optional[str] = None,
+        extra_headers: Optional[dict[str, str]] = None,
+        timeout: int = 300,
     ) -> Iterator[GenerativeChunk]:
         endpoint = f"{self.url}{CHAT_PREDICT}"
         if model:
@@ -266,11 +304,16 @@ class NuaClient:
             endpoint,
             payload=body.model_dump(),
             timeout=timeout,
+            extra_headers=extra_headers,
         ):
             yield gr
 
     def summarize(
-        self, documents: dict[str, str], model: Optional[str] = None, timeout: int = 300
+        self,
+        documents: dict[str, str],
+        model: Optional[str] = None,
+        extra_headers: Optional[dict[str, str]] = None,
+        timeout: int = 300,
     ) -> SummarizedModel:
         endpoint = f"{self.url}{SUMMARIZE_PREDICT}"
         if model:
@@ -288,6 +331,7 @@ class NuaClient:
             payload=body.model_dump(),
             output=SummarizedModel,
             timeout=timeout,
+            extra_headers=extra_headers,
         )
 
     def rephrase(
@@ -324,11 +368,13 @@ class NuaClient:
     def remi(
         self,
         request: RemiRequest,
+        extra_headers: Optional[dict[str, str]] = None,
     ) -> RemiResponse:
         endpoint = f"{self.url}{REMI_PREDICT}"
         return self._request(
             "POST",
             endpoint,
+            extra_headers=extra_headers,
             payload=request.model_dump(),
             output=RemiResponse,
         )
@@ -413,10 +459,18 @@ class NuaClient:
         activity_endpoint = f"{self.url}{STATUS_PROCESS}/{process_id}"
         return self._request("GET", activity_endpoint, ProcessRequestStatus)
 
-    def rerank(self, model: RerankModel) -> RerankResponse:
+    def rerank(
+        self,
+        model: RerankModel,
+        extra_headers: Optional[dict[str, str]] = None,
+    ) -> RerankResponse:
         endpoint = f"{self.url}{RERANK}"
         return self._request(
-            "POST", endpoint, payload=model.model_dump(), output=RerankResponse
+            "POST",
+            endpoint,
+            payload=model.model_dump(),
+            output=RerankResponse,
+            extra_headers=extra_headers,
         )
 
 
@@ -454,9 +508,12 @@ class AsyncNuaClient:
         url: str,
         output: Type[ConvertType],
         payload: Optional[dict[Any, Any]] = None,
+        extra_headers: Optional[dict[str, str]] = None,
         timeout: int = 60,
     ) -> ConvertType:
-        resp = await self.client.request(method, url, json=payload, timeout=timeout)
+        resp = await self.client.request(
+            method, url, json=payload, timeout=timeout, headers=extra_headers
+        )
         if resp.status_code != 200:
             raise NuaAPIException(code=resp.status_code, detail=resp.content.decode())
         try:
@@ -469,6 +526,7 @@ class AsyncNuaClient:
         self,
         method: str,
         url: str,
+        extra_headers: Optional[dict[str, str]] = None,
         payload: Optional[dict[Any, Any]] = None,
         timeout: int = 60,
     ) -> AsyncIterator[GenerativeChunk]:
@@ -477,8 +535,9 @@ class AsyncNuaClient:
             url,
             json=payload,
             timeout=timeout,
+            headers=extra_headers,
         ) as response:
-            if response.headers.get("content-type") == "application/x-ndjson":
+            if response.headers.get("transfer-encoding") == "chunked":
                 async for json_body in response.aiter_lines():
                     try:
                         yield GenerativeChunk.model_validate_json(json_body)  # type: ignore
@@ -527,18 +586,30 @@ class AsyncNuaClient:
         return await self._request("GET", endpoint, output=StoredLearningConfiguration)
 
     async def sentence_predict(
-        self, text: str, model: Optional[str] = None
+        self,
+        text: str,
+        model: Optional[str] = None,
+        extra_headers: Optional[dict[str, str]] = None,
     ) -> Sentence:
         endpoint = f"{self.url}{SENTENCE_PREDICT}?text={text}"
         if model:
             endpoint += f"&model={model}"
-        return await self._request("GET", endpoint, output=Sentence)
+        return await self._request(
+            "GET", endpoint, output=Sentence, extra_headers=extra_headers
+        )
 
-    async def tokens_predict(self, text: str, model: Optional[str] = None) -> Tokens:
+    async def tokens_predict(
+        self,
+        text: str,
+        model: Optional[str] = None,
+        extra_headers: Optional[dict[str, str]] = None,
+    ) -> Tokens:
         endpoint = f"{self.url}{TOKENS_PREDICT}?text={text}"
         if model:
             endpoint += f"&model={model}"
-        return await self._request("GET", endpoint, output=Tokens)
+        return await self._request(
+            "GET", endpoint, output=Tokens, extra_headers=extra_headers
+        )
 
     async def query_predict(
         self,
@@ -546,6 +617,7 @@ class AsyncNuaClient:
         semantic_model: Optional[str] = None,
         token_model: Optional[str] = None,
         generative_model: Optional[str] = None,
+        extra_headers: Optional[dict[str, str]] = None,
     ) -> QueryInfo:
         endpoint = f"{self.url}{QUERY_PREDICT}?text={text}"
         if semantic_model:
@@ -554,7 +626,9 @@ class AsyncNuaClient:
             endpoint += f"&token_model={token_model}"
         if generative_model:
             endpoint += f"&generative_model={generative_model}"
-        return await self._request("GET", endpoint, output=QueryInfo)
+        return await self._request(
+            "GET", endpoint, output=QueryInfo, extra_headers=extra_headers
+        )
 
     @deprecated(version="2.1.0", reason="You should use generate function")
     async def generate_predict(
@@ -573,7 +647,11 @@ class AsyncNuaClient:
         )
 
     async def generate(
-        self, body: ChatModel, model: Optional[str] = None, timeout: int = 300
+        self,
+        body: ChatModel,
+        model: Optional[str] = None,
+        extra_headers: Optional[dict[str, str]] = None,
+        timeout: int = 300,
     ) -> GenerativeFullResponse:
         endpoint = f"{self.url}{CHAT_PREDICT}"
         if model:
@@ -585,6 +663,7 @@ class AsyncNuaClient:
             endpoint,
             payload=body.model_dump(),
             timeout=timeout,
+            extra_headers=extra_headers,
         ):
             if isinstance(chunk.chunk, TextGenerativeResponse):
                 result.answer += chunk.chunk.text
@@ -602,11 +681,20 @@ class AsyncNuaClient:
                 result.code = chunk.chunk.code
             elif isinstance(chunk.chunk, ToolsGenerativeResponse):
                 result.tools = chunk.chunk.tools
+            elif isinstance(chunk.chunk, ConsumptionGenerative):
+                result.consumption = Consumption(
+                    normalized_tokens=chunk.chunk.normalized_tokens,
+                    customer_key_tokens=chunk.chunk.customer_key_tokens,
+                )
 
         return result
 
     async def generate_stream(
-        self, body: ChatModel, model: Optional[str] = None, timeout: int = 300
+        self,
+        body: ChatModel,
+        model: Optional[str] = None,
+        extra_headers: Optional[dict[str, str]] = None,
+        timeout: int = 300,
     ) -> AsyncIterator[GenerativeChunk]:
         endpoint = f"{self.url}{CHAT_PREDICT}"
         if model:
@@ -617,11 +705,16 @@ class AsyncNuaClient:
             endpoint,
             payload=body.model_dump(),
             timeout=timeout,
+            extra_headers=extra_headers,
         ):
             yield gr
 
     async def summarize(
-        self, documents: dict[str, str], model: Optional[str] = None, timeout: int = 300
+        self,
+        documents: dict[str, str],
+        model: Optional[str] = None,
+        extra_headers: Optional[dict[str, str]] = None,
+        timeout: int = 300,
     ) -> SummarizedModel:
         endpoint = f"{self.url}{SUMMARIZE_PREDICT}"
         if model:
@@ -639,6 +732,7 @@ class AsyncNuaClient:
             payload=body.model_dump(),
             output=SummarizedModel,
             timeout=timeout,
+            extra_headers=extra_headers,
         )
 
     async def rephrase(
@@ -672,13 +766,18 @@ class AsyncNuaClient:
             output=RephraseModel,
         )
 
-    async def remi(self, request: RemiRequest) -> RemiResponse:
+    async def remi(
+        self,
+        request: RemiRequest,
+        extra_headers: Optional[dict[str, str]] = None,
+    ) -> RemiResponse:
         endpoint = f"{self.url}{REMI_PREDICT}"
         return await self._request(
             "POST",
             endpoint,
             payload=request.model_dump(),
             output=RemiResponse,
+            extra_headers=extra_headers,
         )
 
     async def generate_retrieval(
@@ -802,8 +901,14 @@ class AsyncNuaClient:
             "GET", activity_endpoint, output=ProcessRequestStatus
         )
 
-    async def rerank(self, model: RerankModel) -> RerankResponse:
+    async def rerank(
+        self, model: RerankModel, extra_headers: Optional[dict[str, str]] = None
+    ) -> RerankResponse:
         endpoint = f"{self.url}{RERANK}"
         return await self._request(
-            "POST", endpoint, payload=model.model_dump(), output=RerankResponse
+            "POST",
+            endpoint,
+            payload=model.model_dump(),
+            output=RerankResponse,
+            extra_headers=extra_headers,
         )

nuclia/lib/nua_responses.py

@@ -5,6 +5,7 @@ from typing import Any, Dict, List, Optional, Union, cast
 import pydantic
 from pydantic import BaseModel, Field, RootModel, model_validator
 from typing_extensions import Annotated, Self
+from nuclia_models.common.consumption import Consumption
 
 
 class GenerativeOption(BaseModel):
@@ -51,6 +52,7 @@ class ConfigSchema(BaseModel):
 class Sentence(BaseModel):
     data: List[float]
     time: float
+    consumption: Optional[Consumption] = None
 
 
 class Author(str, Enum):
@@ -137,6 +139,7 @@ class Token(BaseModel):
 class Tokens(BaseModel):
     tokens: List[Token]
     time: float
+    consumption: Optional[Consumption] = None
 
 
 class SummarizeResource(BaseModel):
@@ -155,6 +158,7 @@ class SummarizedResource(BaseModel):
 class SummarizedModel(BaseModel):
     resources: Dict[str, SummarizedResource]
     summary: str = ""
+    consumption: Optional[Consumption] = None
 
 
 class RephraseModel(RootModel[str]):
@@ -535,6 +539,7 @@ class StoredLearningConfiguration(BaseModel):
 class SentenceSearch(BaseModel):
     data: List[float] = []
     time: float
+    consumption: Optional[Consumption] = None
 
 
 class Ner(BaseModel):
@@ -547,6 +552,7 @@ class Ner(BaseModel):
 class TokenSearch(BaseModel):
     tokens: List[Ner] = []
     time: float
+    consumption: Optional[Consumption] = None
 
 
 class QueryInfo(BaseModel):
@@ -569,3 +575,4 @@ class RerankResponse(BaseModel):
     context_scores: dict[str, float] = Field(
         description="Scores for each context given by the reranker"
     )
+    consumption: Optional[Consumption] = None

nuclia/sdk/kb.py

@@ -297,9 +297,15 @@ class NucliaKB:
         )
 
     @kb
-    def summarize(self, *, resources: List[str], **kwargs):
+    def summarize(
+        self, *, resources: List[str], show_consumption: bool = False, **kwargs
+    ):
         ndb: NucliaDBClient = kwargs["ndb"]
-        return ndb.ndb.summarize(kbid=ndb.kbid, resources=resources)
+        return ndb.ndb.summarize(
+            kbid=ndb.kbid,
+            resources=resources,
+            headers={"X-Show-Consumption": str(show_consumption).lower()},
+        )
 
     @kb
     def notifications(self, **kwargs):
@@ -715,6 +721,7 @@ class AsyncNucliaKB:
         resources: List[str],
         generative_model: Optional[str] = None,
         summary_kind: Optional[str] = None,
+        show_consumption: bool = False,
         timeout: int = 1000,
         **kwargs,
     ) -> SummarizedModel:
@@ -725,6 +732,7 @@ class AsyncNucliaKB:
                 generative_model=generative_model,
                 summary_kind=SummaryKind(summary_kind),
             ),
+            extra_headers={"X-Show-Consumption": str(show_consumption).lower()},
             timeout=timeout,
         )
         return SummarizedModel.model_validate(resp.json())

nuclia/sdk/predict.py

@@ -52,9 +52,19 @@ class NucliaPredict:
         nc.del_config_predict(kbid)
 
     @nua
-    def sentence(self, text: str, model: Optional[str] = None, **kwargs) -> Sentence:
+    def sentence(
+        self,
+        text: str,
+        model: Optional[str] = None,
+        show_consumption: bool = False,
+        **kwargs,
+    ) -> Sentence:
         nc: NuaClient = kwargs["nc"]
-        return nc.sentence_predict(text, model)
+        return nc.sentence_predict(
+            text,
+            model,
+            extra_headers={"X-Show-Consumption": str(show_consumption).lower()},
+        )
 
     @nua
     def query(
@@ -63,19 +73,25 @@ class NucliaPredict:
         semantic_model: Optional[str] = None,
         token_model: Optional[str] = None,
         generative_model: Optional[str] = None,
+        show_consumption: bool = False,
         **kwargs,
     ) -> QueryInfo:
         nc: NuaClient = kwargs["nc"]
         return nc.query_predict(
-            text,
+            text=text,
             semantic_model=semantic_model,
             token_model=token_model,
             generative_model=generative_model,
+            extra_headers={"X-Show-Consumption": str(show_consumption).lower()},
         )
 
     @nua
     def generate(
-        self, text: Union[str, ChatModel], model: Optional[str] = None, **kwargs
+        self,
+        text: Union[str, ChatModel],
+        model: Optional[str] = None,
+        show_consumption: bool = False,
+        **kwargs,
     ) -> GenerativeFullResponse:
         nc: NuaClient = kwargs["nc"]
         if isinstance(text, str):
@@ -88,11 +104,19 @@ class NucliaPredict:
         else:
             body = text
 
-        return nc.generate(body, model)
+        return nc.generate(
+            body=body,
+            model=model,
+            extra_headers={"X-Show-Consumption": str(show_consumption).lower()},
+        )
 
     @nua
     def generate_stream(
-        self, text: Union[str, ChatModel], model: Optional[str] = None, **kwargs
+        self,
+        text: Union[str, ChatModel],
+        model: Optional[str] = None,
+        show_consumption: bool = False,
+        **kwargs,
     ) -> Iterator[GenerativeChunk]:
         nc: NuaClient = kwargs["nc"]
         if isinstance(text, str):
@@ -105,20 +129,42 @@ class NucliaPredict:
         else:
             body = text
 
-        for chunk in nc.generate_stream(body, model):
+        for chunk in nc.generate_stream(
+            body=body,
+            model=model,
+            extra_headers={"X-Show-Consumption": str(show_consumption).lower()},
+        ):
             yield chunk
 
     @nua
-    def tokens(self, text: str, model: Optional[str] = None, **kwargs) -> Tokens:
+    def tokens(
+        self,
+        text: str,
+        model: Optional[str] = None,
+        show_consumption: bool = False,
+        **kwargs,
+    ) -> Tokens:
         nc: NuaClient = kwargs["nc"]
-        return nc.tokens_predict(text, model)
+        return nc.tokens_predict(
+            text,
+            model,
+            extra_headers={"X-Show-Consumption": str(show_consumption).lower()},
+        )
 
     @nua
     def summarize(
-        self, texts: dict[str, str], model: Optional[str] = None, **kwargs
+        self,
+        texts: dict[str, str],
+        model: Optional[str] = None,
+        show_consumption: bool = False,
+        **kwargs,
     ) -> SummarizedModel:
         nc: NuaClient = kwargs["nc"]
-        return nc.summarize(texts, model)
+        return nc.summarize(
+            documents=texts,
+            model=model,
+            extra_headers={"X-Show-Consumption": str(show_consumption).lower()},
+        )
 
     @nua
     def rephrase(
@@ -135,7 +181,12 @@ class NucliaPredict:
 
     @nua
     def rag(
-        self, question: str, context: list[str], model: Optional[str] = None, **kwargs
+        self,
+        question: str,
+        context: list[str],
+        model: Optional[str] = None,
+        show_consumption: bool = False,
+        **kwargs,
     ) -> GenerativeFullResponse:
         nc: NuaClient = kwargs["nc"]
         body = ChatModel(
@@ -145,10 +196,19 @@ class NucliaPredict:
             query_context=context,
         )
 
-        return nc.generate(body, model)
+        return nc.generate(
+            body,
+            model,
+            extra_headers={"X-Show-Consumption": str(show_consumption).lower()},
+        )
 
     @nua
-    def remi(self, request: Optional[RemiRequest] = None, **kwargs) -> RemiResponse:
+    def remi(
+        self,
+        request: Optional[RemiRequest] = None,
+        show_consumption: bool = False,
+        **kwargs,
+    ) -> RemiResponse:
         """
         Perform a REMi evaluation over a RAG experience
 
@@ -162,10 +222,15 @@ class NucliaPredict:
         if request is None:
             request = RemiRequest(**kwargs)
         nc: NuaClient = kwargs["nc"]
-        return nc.remi(request)
+        return nc.remi(
+            request=request,
+            extra_headers={"X-Show-Consumption": str(show_consumption).lower()},
+        )
 
     @nua
-    def rerank(self, request: RerankModel, **kwargs) -> RerankResponse:
+    def rerank(
+        self, request: RerankModel, show_consumption: bool = False, **kwargs
+    ) -> RerankResponse:
         """
         Perform a reranking of the results based on the question and context provided.
 
@@ -173,7 +238,10 @@ class NucliaPredict:
         :return: RerankResponse
         """
         nc: NuaClient = kwargs["nc"]
-        return nc.rerank(request)
+        return nc.rerank(
+            request,
+            extra_headers={"X-Show-Consumption": str(show_consumption).lower()},
+        )
 
 
 class AsyncNucliaPredict:
@@ -208,14 +276,26 @@ class AsyncNucliaPredict:
 
     @nua
     async def sentence(
-        self, text: str, model: Optional[str] = None, **kwargs
+        self,
+        text: str,
+        model: Optional[str] = None,
+        show_consumption: bool = False,
+        **kwargs,
     ) -> Sentence:
         nc: AsyncNuaClient = kwargs["nc"]
-        return await nc.sentence_predict(text, model)
+        return await nc.sentence_predict(
+            text,
+            model,
+            extra_headers={"X-Show-Consumption": str(show_consumption).lower()},
+        )
 
     @nua
     async def generate(
-        self, text: Union[str, ChatModel], model: Optional[str] = None, **kwargs
+        self,
+        text: Union[str, ChatModel],
+        model: Optional[str] = None,
+        show_consumption: bool = False,
+        **kwargs,
     ) -> GenerativeFullResponse:
         nc: AsyncNuaClient = kwargs["nc"]
         if isinstance(text, str):
@@ -227,11 +307,19 @@ class AsyncNucliaPredict:
             )
         else:
             body = text
-        return await nc.generate(body, model)
+        return await nc.generate(
+            body=body,
+            model=model,
+            extra_headers={"X-Show-Consumption": str(show_consumption).lower()},
+        )
 
     @nua
     async def generate_stream(
-        self, text: Union[str, ChatModel], model: Optional[str] = None, **kwargs
+        self,
+        text: Union[str, ChatModel],
+        model: Optional[str] = None,
+        show_consumption: bool = False,
+        **kwargs,
     ) -> AsyncIterator[GenerativeChunk]:
         nc: AsyncNuaClient = kwargs["nc"]
         if isinstance(text, str):
@@ -244,13 +332,27 @@ class AsyncNucliaPredict:
         else:
             body = text
 
-        async for chunk in nc.generate_stream(body, model):
+        async for chunk in nc.generate_stream(
+            body=body,
+            model=model,
+            extra_headers={"X-Show-Consumption": str(show_consumption).lower()},
+        ):
             yield chunk
 
     @nua
-    async def tokens(self, text: str, model: Optional[str] = None, **kwargs) -> Tokens:
+    async def tokens(
+        self,
+        text: str,
+        model: Optional[str] = None,
+        show_consumption: bool = False,
+        **kwargs,
+    ) -> Tokens:
         nc: AsyncNuaClient = kwargs["nc"]
-        return await nc.tokens_predict(text, model)
+        return await nc.tokens_predict(
+            text,
+            model,
+            extra_headers={"X-Show-Consumption": str(show_consumption).lower()},
+        )
 
     @nua
     async def query(
@@ -259,22 +361,32 @@ class AsyncNucliaPredict:
         semantic_model: Optional[str] = None,
         token_model: Optional[str] = None,
         generative_model: Optional[str] = None,
+        show_consumption: bool = False,
         **kwargs,
     ) -> QueryInfo:
         nc: AsyncNuaClient = kwargs["nc"]
         return await nc.query_predict(
-            text,
+            text=text,
             semantic_model=semantic_model,
             token_model=token_model,
             generative_model=generative_model,
+            extra_headers={"X-Show-Consumption": str(show_consumption).lower()},
         )
 
     @nua
     async def summarize(
-        self, texts: dict[str, str], model: Optional[str] = None, **kwargs
+        self,
+        texts: dict[str, str],
+        model: Optional[str] = None,
+        show_consumption: bool = False,
+        **kwargs,
     ) -> SummarizedModel:
         nc: AsyncNuaClient = kwargs["nc"]
-        return await nc.summarize(texts, model)
+        return await nc.summarize(
+            documents=texts,
+            model=model,
+            extra_headers={"X-Show-Consumption": str(show_consumption).lower()},
+        )
 
     @nua
     async def rephrase(
@@ -298,7 +410,10 @@ class AsyncNucliaPredict:
 
     @nua
     async def remi(
-        self, request: Optional[RemiRequest] = None, **kwargs
+        self,
+        request: Optional[RemiRequest] = None,
+        show_consumption: bool = False,
+        **kwargs,
     ) -> RemiResponse:
         """
         Perform a REMi evaluation over a RAG experience
@@ -311,10 +426,15 @@ class AsyncNucliaPredict:
             request = RemiRequest(**kwargs)
 
         nc: AsyncNuaClient = kwargs["nc"]
-        return await nc.remi(request)
+        return await nc.remi(
+            request=request,
+            extra_headers={"X-Show-Consumption": str(show_consumption).lower()},
+        )
 
     @nua
-    async def rerank(self, request: RerankModel, **kwargs) -> RerankResponse:
+    async def rerank(
+        self, request: RerankModel, show_consumption: bool = False, **kwargs
+    ) -> RerankResponse:
         """
         Perform a reranking of the results based on the question and context provided.
 
@@ -322,4 +442,7 @@ class AsyncNucliaPredict:
         :return: RerankResponse
         """
         nc: AsyncNuaClient = kwargs["nc"]
-        return await nc.rerank(request)
+        return await nc.rerank(
+            request,
+            extra_headers={"X-Show-Consumption": str(show_consumption).lower()},
+        )

nuclia/sdk/search.py

@@ -30,6 +30,7 @@ from nuclia.lib.kb import AsyncNucliaDBClient, NucliaDBClient
 from nuclia.sdk.logger import logger
 from nuclia.sdk.auth import AsyncNucliaAuth, NucliaAuth
 from nuclia.sdk.resource import RagImagesStrategiesParse, RagStrategiesParse
+from nuclia_models.common.consumption import Consumption, TokensDetail
 
 
 @dataclass
@@ -49,6 +50,7 @@ class AskAnswer:
     relations: Optional[Relations]
     predict_request: Optional[ChatModel]
     error_details: Optional[str]
+    consumption: Optional[Consumption]
 
     def __str__(self):
         if self.answer:
@@ -184,8 +186,9 @@ class NucliaSearch:
         filters: Optional[Union[List[str], List[Filter]]] = None,
         rag_strategies: Optional[list[RagStrategies]] = None,
         rag_images_strategies: Optional[list[RagImagesStrategies]] = None,
+        show_consumption: bool = False,
         **kwargs,
-    ):
+    ) -> AskAnswer:
         """
         Answer a question.
 
@@ -217,7 +220,11 @@ class NucliaSearch:
         else:
             raise ValueError("Invalid query type. Must be str, dict or AskRequest.")
 
-        ask_response: SyncAskResponse = ndb.ndb.ask(kbid=ndb.kbid, content=req)
+        ask_response: SyncAskResponse = ndb.ndb.ask(
+            kbid=ndb.kbid,
+            content=req,
+            headers={"X-Show-Consumption": str(show_consumption).lower()},
+        )
 
         result = AskAnswer(
             answer=ask_response.answer.encode(),
@@ -239,6 +246,7 @@ class NucliaSearch:
             else None,
             relations=ask_response.relations,
             prompt_context=ask_response.prompt_context,
+            consumption=ask_response.consumption,
         )
 
         if ask_response.prompt_context:
@@ -257,8 +265,9 @@ class NucliaSearch:
         schema: Union[str, Dict[str, Any]],
         query: Union[str, dict, AskRequest, None] = None,
         filters: Optional[Union[List[str], List[Filter]]] = None,
+        show_consumption: bool = False,
         **kwargs,
-    ):
+    ) -> Optional[AskAnswer]:
         """
         Answer a question.
 
@@ -272,10 +281,10 @@ class NucliaSearch:
                         schema_json = json.load(json_file_handler)
                     except Exception:
                         logger.exception("File format is not JSON")
-                        return
+                        return None
             else:
                 logger.exception("File not found")
-                return
+                return None
         else:
             schema_json = schema
 
@@ -303,7 +312,11 @@ class NucliaSearch:
                 req.filters = filters
         else:
             raise ValueError("Invalid query type. Must be str, dict or AskRequest.")
-        ask_response: SyncAskResponse = ndb.ndb.ask(kbid=ndb.kbid, content=req)
+        ask_response: SyncAskResponse = ndb.ndb.ask(
+            kbid=ndb.kbid,
+            content=req,
+            headers={"X-Show-Consumption": str(show_consumption).lower()},
+        )
 
         result = AskAnswer(
             answer=ask_response.answer.encode(),
@@ -325,6 +338,7 @@ class NucliaSearch:
             else None,
             relations=ask_response.relations,
             prompt_context=ask_response.prompt_context,
+            consumption=ask_response.consumption,
         )
         if ask_response.metadata is not None:
             if ask_response.metadata.timings is not None:
@@ -483,9 +497,10 @@ class AsyncNucliaSearch:
         *,
         query: Union[str, dict, AskRequest],
         filters: Optional[List[str]] = None,
+        show_consumption: bool = False,
         timeout: int = 100,
         **kwargs,
-    ):
+    ) -> AskAnswer:
         """
         Answer a question.
 
@@ -509,7 +524,11 @@ class AsyncNucliaSearch:
             req = query
         else:
             raise ValueError("Invalid query type. Must be str, dict or AskRequest.")
-        ask_stream_response = await ndb.ask(req, timeout=timeout)
+        ask_stream_response = await ndb.ask(
+            req,
+            timeout=timeout,
+            extra_headers={"X-Show-Consumption": str(show_consumption).lower()},
+        )
         result = AskAnswer(
             answer=b"",
             learning_id=ask_stream_response.headers.get("NUCLIA-LEARNING-ID", ""),
@@ -526,6 +545,7 @@ class AsyncNucliaSearch:
             predict_request=None,
             relations=None,
             prompt_context=None,
+            consumption=None,
         )
         async for line in ask_stream_response.aiter_lines():
             try:
@@ -548,6 +568,19 @@ class AsyncNucliaSearch:
                     result.timings = ask_response_item.timings.model_dump()
                 if ask_response_item.tokens:
                     result.tokens = ask_response_item.tokens.model_dump()
+            elif ask_response_item.type == "consumption":
+                result.consumption = Consumption(
+                    normalized_tokens=TokensDetail(
+                        input=ask_response_item.normalized_tokens.input,
+                        output=ask_response_item.normalized_tokens.output,
+                        image=ask_response_item.normalized_tokens.image,
+                    ),
+                    customer_key_tokens=TokensDetail(
+                        input=ask_response_item.customer_key_tokens.input,
+                        output=ask_response_item.customer_key_tokens.output,
+                        image=ask_response_item.customer_key_tokens.image,
+                    ),
+                )
             elif ask_response_item.type == "status":
                 result.status = ask_response_item.status
             elif ask_response_item.type == "prequeries":
@@ -569,6 +602,7 @@ class AsyncNucliaSearch:
         *,
         query: Union[str, dict, AskRequest],
         filters: Optional[List[str]] = None,
+        show_consumption: bool = False,
         timeout: int = 100,
         **kwargs,
     ) -> AsyncIterator[AskResponseItem]:
@@ -593,7 +627,11 @@ class AsyncNucliaSearch:
             req = query
         else:
             raise ValueError("Invalid query type. Must be str, dict or AskRequest.")
-        ask_stream_response = await ndb.ask(req, timeout=timeout)
+        ask_stream_response = await ndb.ask(
+            req,
+            timeout=timeout,
+            extra_headers={"X-Show-Consumption": str(show_consumption).lower()},
+        )
         async for line in ask_stream_response.aiter_lines():
             try:
                 ask_response_item = AskResponseItem.model_validate_json(line)
@@ -609,9 +647,10 @@ class AsyncNucliaSearch:
         query: Union[str, dict, AskRequest],
         schema: Dict[str, Any],
         filters: Optional[List[str]] = None,
+        show_consumption: bool = False,
         timeout: int = 100,
         **kwargs,
-    ):
+    ) -> AskAnswer:
         """
         Answer a question.
 
@@ -635,7 +674,11 @@ class AsyncNucliaSearch:
             req = query
         else:
             raise ValueError("Invalid query type. Must be str, dict or AskRequest.")
-        ask_stream_response = await ndb.ask(req, timeout=timeout)
+        ask_stream_response = await ndb.ask(
+            req,
+            timeout=timeout,
+            extra_headers={"X-Show-Consumption": str(show_consumption).lower()},
+        )
         result = AskAnswer(
             answer=b"",
             learning_id=ask_stream_response.headers.get("NUCLIA-LEARNING-ID", ""),
@@ -652,6 +695,7 @@ class AsyncNucliaSearch:
             predict_request=None,
             relations=None,
             prompt_context=None,
+            consumption=None,
         )
         async for line in ask_stream_response.aiter_lines():
             try:
@@ -674,6 +718,19 @@ class AsyncNucliaSearch:
                     result.timings = ask_response_item.timings.model_dump()
                 if ask_response_item.tokens:
                     result.tokens = ask_response_item.tokens.model_dump()
+            elif ask_response_item.type == "consumption":
+                result.consumption = Consumption(
+                    normalized_tokens=TokensDetail(
+                        input=ask_response_item.normalized_tokens.input,
+                        output=ask_response_item.normalized_tokens.output,
+                        image=ask_response_item.normalized_tokens.image,
+                    ),
+                    customer_key_tokens=TokensDetail(
+                        input=ask_response_item.customer_key_tokens.input,
+                        output=ask_response_item.customer_key_tokens.output,
+                        image=ask_response_item.customer_key_tokens.image,
+                    ),
+                )
             elif ask_response_item.type == "status":
                 result.status = ask_response_item.status
             elif ask_response_item.type == "prequeries":

nuclia/tests/test_kb/test_search.py

@@ -132,7 +132,11 @@ def test_ask_json(testing_config):
 async def test_ask_json_async(testing_config):
     search = AsyncNucliaSearch()
     results = await search.ask_json(
-        query="Who is hedy Lamarr?", filters=["/icon/application/pdf"], schema=SCHEMA
+        query="Who is hedy Lamarr?",
+        filters=["/icon/application/pdf"],
+        schema=SCHEMA,
+        show_consumption=True,
     )
 
     assert "TECHNOLOGY" in results.object["document_type"]
+    assert results.consumption is not None

nuclia/tests/test_nua/test_predict.py

@@ -1,4 +1,7 @@
-from nuclia_models.predict.generative_responses import TextGenerativeResponse
+from nuclia_models.predict.generative_responses import (
+    TextGenerativeResponse,
+    ConsumptionGenerative,
+)
 
 from nuclia.lib.nua_responses import ChatModel, RerankModel, UserPrompt
 from nuclia.sdk.predict import AsyncNucliaPredict, NucliaPredict
@@ -8,9 +11,12 @@ from nuclia_models.predict.remi import RemiRequest
 
 def test_predict(testing_config):
     np = NucliaPredict()
-    embed = np.sentence(text="This is my text", model="multilingual-2024-05-06")
+    embed = np.sentence(
+        text="This is my text", model="multilingual-2024-05-06", show_consumption=True
+    )
     assert embed.time > 0
     assert len(embed.data) == 1024
+    assert embed.consumption is not None
 
 
 def test_predict_query(testing_config):
@@ -20,11 +26,14 @@ def test_predict_query(testing_config):
         semantic_model="multilingual-2024-05-06",
         token_model="multilingual",
         generative_model="chatgpt-azure-4o-mini",
+        show_consumption=True,
     )
     assert query.language == "en"
     assert query.visual_llm is True
     assert query.entities and query.entities.tokens[0].text == "Ramon"
     assert query.sentence and len(query.sentence.data) == 1024
+    assert query.entities.consumption is not None
+    assert query.sentence.consumption is not None
 
 
 def test_rag(testing_config):
@@ -36,14 +45,34 @@ def test_rag(testing_config):
             "Eudald Camprubí is CEO at the same company as Ramon Navarro",
         ],
         model="chatgpt-azure-4o-mini",
+        show_consumption=True,
     )
     assert "Eudald" in generated.answer
+    assert generated.consumption is not None
 
 
 def test_generative(testing_config):
     np = NucliaPredict()
     generated = np.generate(text="How much is 2 + 2?", model="chatgpt-azure-4o-mini")
     assert "4" in generated.answer
+    assert generated.consumption is None
+
+
+@pytest.mark.asyncio
+async def test_generative_with_consumption(testing_config):
+    np = NucliaPredict()
+    generated = np.generate(
+        text="How much is 2 + 2?", model="chatgpt-azure-4o-mini", show_consumption=True
+    )
+    assert "4" in generated.answer
+    assert generated.consumption is not None
+
+    anp = AsyncNucliaPredict()
+    async_generated = await anp.generate(
+        text="How much is 2 + 2?", model="chatgpt-azure-4o-mini", show_consumption=True
+    )
+    assert "4" in async_generated.answer
+    assert async_generated.consumption is not None
 
 
 @pytest.mark.asyncio
@@ -70,13 +99,18 @@ def test_stream_generative(testing_config):
 @pytest.mark.asyncio
 async def test_async_stream_generative(testing_config):
     np = AsyncNucliaPredict()
+    consumption_found = False
+    found = False
     async for stream in np.generate_stream(
-        text="How much is 2 + 2?", model="chatgpt-azure-4o-mini"
+        text="How much is 2 + 2?", model="chatgpt-azure-4o-mini", show_consumption=True
     ):
         if isinstance(stream.chunk, TextGenerativeResponse) and stream.chunk.text:
             if "4" in stream.chunk.text:
                 found = True
+        elif isinstance(stream.chunk, ConsumptionGenerative):
+            consumption_found = True
     assert found
+    assert consumption_found
 
 
 SCHEMA = {
@@ -148,6 +182,8 @@ def test_nua_remi(testing_config):
     assert results.context_relevance[1] < 2
     assert results.groundedness[1] < 2
 
+    assert results.consumption is None
+
 
 @pytest.mark.asyncio
 async def test_nua_async_remi(testing_config):
@@ -161,7 +197,8 @@ async def test_nua_async_remi(testing_config):
                 "Paris is the capital of France.",
                 "Berlin is the capital of Germany.",
             ],
-        )
+        ),
+        show_consumption=True,
     )
     assert results.answer_relevance.score >= 4
 
@@ -171,6 +208,8 @@ async def test_nua_async_remi(testing_config):
     assert results.context_relevance[1] < 2
     assert results.groundedness[1] < 2
 
+    assert results.consumption is not None
+
 
 def test_nua_rerank(testing_config):
     np = NucliaPredict()
@@ -185,3 +224,37 @@ def test_nua_rerank(testing_config):
         )
     )
     assert results.context_scores["1"] > results.context_scores["2"]
+    assert results.consumption is None
+
+
+@pytest.mark.asyncio
+async def test_nua_rerank_with_consumption(testing_config):
+    np = NucliaPredict()
+    results = np.rerank(
+        RerankModel(
+            user_id="Nuclia PY CLI",
+            question="What is the capital of France?",
+            context={
+                "1": "Paris is the capital of France.",
+                "2": "Berlin is the capital of Germany.",
+            },
+        ),
+        show_consumption=True,
+    )
+    assert results.context_scores["1"] > results.context_scores["2"]
+    assert results.consumption is not None
+
+    anp = AsyncNucliaPredict()
+    async_results = await anp.rerank(
+        RerankModel(
+            user_id="Nuclia PY CLI",
+            question="What is the capital of France?",
+            context={
+                "1": "Paris is the capital of France.",
+                "2": "Berlin is the capital of Germany.",
+            },
+        ),
+        show_consumption=True,
+    )
+    assert async_results.context_scores["1"] > async_results.context_scores["2"]
+    assert async_results.consumption is not None

{nuclia-4.9.3.dist-info → nuclia-4.9.4.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: nuclia
-Version: 4.9.3
+Version: 4.9.4
 Summary: Nuclia Python SDK
 Author-email: Nuclia <info@nuclia.com>
 License-Expression: MIT
@@ -24,9 +24,9 @@ Requires-Dist: requests
 Requires-Dist: httpx
 Requires-Dist: httpcore>=1.0.0
 Requires-Dist: prompt_toolkit
-Requires-Dist: nucliadb_sdk<7,>=6.5
-Requires-Dist: nucliadb_models<7,>=6.5
-Requires-Dist: nuclia-models>=0.41.1
+Requires-Dist: nucliadb_sdk<7,>=6.6.1
+Requires-Dist: nucliadb_models<7,>=6.6.1
+Requires-Dist: nuclia-models>=0.45.0
 Requires-Dist: tqdm
 Requires-Dist: aiofiles
 Requires-Dist: backoff

{nuclia-4.9.3.dist-info → nuclia-4.9.4.dist-info}/RECORD

@@ -9,11 +9,11 @@ nuclia/cli/run.py,sha256=B1hP0upSbSCqqT89WAwsd93ZxkAoF6ajVyLOdYmo8fU,1560
 nuclia/cli/utils.py,sha256=iZ3P8juBdAGvaRUd2BGz7bpUXNDHdPrC5p876yyZ2Cs,1223
 nuclia/lib/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 nuclia/lib/conversations.py,sha256=M6qhL9NPEKroYF767S-Q2XWokRrjX02kpYTzRvZKwUE,149
-nuclia/lib/kb.py,sha256=vSLfmV6HqPvWJwPVw4iIzkDnm0M96_ImTH3QMZdkd_I,30985
+nuclia/lib/kb.py,sha256=hbWg-0pX1OCt9ezUXYcOfyWHahgmZLUQe2nlSycHetU,31364
 nuclia/lib/models.py,sha256=ekEQrVIFU3aFvt60yQh-zpWkGNORBMSc7c5Hd_VzPzI,1564
-nuclia/lib/nua.py,sha256=jVUV8I50iJpUQD6yn4V7wNDDSDP8oEWrRxnmrs33hnI,27591
+nuclia/lib/nua.py,sha256=qP6Gv9ck117kdqNuKVl_qHodrDE8CS8xdh_bdzeJMiQ,30631
 nuclia/lib/nua_chat.py,sha256=ApL1Y1FWvAVUt-Y9a_8TUSJIhg8-UmBSy8TlDPn6tD8,3874
-nuclia/lib/nua_responses.py,sha256=FyCedSKHVGn4w380obV9B3D1JaqVyNTqda4-OCs4A9A,13637
+nuclia/lib/nua_responses.py,sha256=KHAjoW3fu2QT7u4F4WfL4DQIM8lCE5W_14vuRtCmtZg,13970
 nuclia/lib/utils.py,sha256=9l6DxBk-11WqUhXRn99cqeuVTUOJXj-1S6zckal7wOk,6312
 nuclia/sdk/__init__.py,sha256=-nAw8i53XBdmbfTa1FJZ0FNRMNakimDVpD6W4OdES-c,1374
 nuclia/sdk/accounts.py,sha256=7XQ3K9_jlSuk2Cez868FtazZ05xSGab6h3Mt1qMMwIE,647
@@ -22,17 +22,17 @@ nuclia/sdk/auth.py,sha256=o9CiWt3P_UHQW_L5Jq_9IOB-KpSCXFsVTouQcZp8BJM,25072
 nuclia/sdk/backup.py,sha256=adbPcNEbHGZW698o028toXKfDkDrmk5QRIDSiN6SPys,6529
 nuclia/sdk/export_import.py,sha256=y5cTOxhILwRPIvR2Ya12bk-ReGbeDzA3C9TPxgnOHD4,9756
 nuclia/sdk/extract_strategy.py,sha256=NZBLLThdLyQYw8z1mT9iRhFjkE5sQP86-3QhsTiyV9o,2540
-nuclia/sdk/kb.py,sha256=2-H9FOvPgsG-ZYNUA8D4FYAJhX3K8m2VwbwSQy1JV7c,27044
+nuclia/sdk/kb.py,sha256=rY1yVx_9S01twMXgrkpKSNpPChTupwYEw7EAIOwzLYA,27321
 nuclia/sdk/kbs.py,sha256=nXEvg5ddZYdDS8Kie7TrN-s1meU9ecYLf9FlT5xr-ro,9131
 nuclia/sdk/logger.py,sha256=UHB81eS6IGmLrsofKxLh8cmF2AsaTj_HXP0tGqMr_HM,57
 nuclia/sdk/logs.py,sha256=3jfORpo8fzZiXFFSbGY0o3Bre1ZgJaKQCXgxP1keNHw,9614
 nuclia/sdk/nua.py,sha256=6t0m0Sx-UhqNU2Hx9v6vTwy0m3a30K4T0KmP9G43MzY,293
 nuclia/sdk/nucliadb.py,sha256=bOESIppPgY7IrNqrYY7T3ESoxwttbOSTm5zj1xUS1jI,1288
-nuclia/sdk/predict.py,sha256=Gom-BlISXawThNq-f7fb1te5tY4catlIdmO-pMJyqkk,9815
+nuclia/sdk/predict.py,sha256=ZQMtoeQBXaPuaJRxOmO23nnKkVKRpNCMi0JgXWSJ9lM,12836
 nuclia/sdk/process.py,sha256=WuNnqaWprp-EABWDC_z7O2woesGIlYWnDUKozh7Ibr4,2241
 nuclia/sdk/remi.py,sha256=BEb3O9R2jOFlOda4vjFucKKGO1c2eTkqYZdFlIy3Zmo,4357
 nuclia/sdk/resource.py,sha256=0lSvD4e1FpN5iM9W295dOKLJ8hsXfIe8HKdo0HsVg20,13976
-nuclia/sdk/search.py,sha256=1mLJzDO-W3ObReDL1xK8zgjIxkdpuIhyf87-L4bISKQ,25180
+nuclia/sdk/search.py,sha256=tHpx0gwVwSepa3dn7AX2iWs7JO0pQkjJk7v8TEoL7Gg,27742
 nuclia/sdk/task.py,sha256=UawH-7IneRIGVOiLdDQ2vDBwe5eI51UXRbLRRVeu3C4,6095
 nuclia/sdk/upload.py,sha256=ZBzYROF3yP-77HcaR06OBsFjJAbTOCvF-nlxaqQZsT4,22720
 nuclia/sdk/zones.py,sha256=1ARWrTsTuzj8zguanpX3OaIw-3Qq_ULS_g4GG2mHxOA,342
@@ -49,7 +49,7 @@ nuclia/tests/test_kb/test_labels.py,sha256=IUdTq4mzv0OrOkwBWWy4UwKGKyJybtoHrgvXr
 nuclia/tests/test_kb/test_logs.py,sha256=Z9ELtiiU9NniITJzeWt92GCcERKYy9Nwc_fUVPboRU0,3121
 nuclia/tests/test_kb/test_remi.py,sha256=OX5N-MHbgcwpLg6fBjrAK_KhqkMspJo_VKQHCBCayZ8,2080
 nuclia/tests/test_kb/test_resource.py,sha256=05Xgmg5fwcPW2PZKnUSSjr6MPXp5w8XDgx8plfNcR68,1102
-nuclia/tests/test_kb/test_search.py,sha256=bsQhfB6-NYFwY3gqkrVJJnru153UgZqEhV22ho4VeWM,4660
+nuclia/tests/test_kb/test_search.py,sha256=8v3u-VcczBCTc-rHlWVM5mUtLJKNBZSM0V4v2I1uE3E,4751
 nuclia/tests/test_kb/test_tasks.py,sha256=tfJl1js2o0_dKyXdOxiwdj929kbzMkfl5XjO8HVMInQ,3456
 nuclia/tests/test_manage/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 nuclia/tests/test_manage/test_account.py,sha256=u9NhRK8gJLS7BEY618aGoYoV2rgDLHZUeSsWWYkDhF8,289
@@ -57,15 +57,15 @@ nuclia/tests/test_manage/test_auth.py,sha256=I5ho9rKhrzCiP57cDVcs2zrXyy1uWlKxzfE
 nuclia/tests/test_manage/test_kb.py,sha256=gqBMxmIuPUKC6kP2V_IapS9u72QR99V8CDQlJOa8sHU,1959
 nuclia/tests/test_nua/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 nuclia/tests/test_nua/test_agent.py,sha256=iPA8lVw7CyONS-0fVG7rDzv8T-LnUlNPMNynTnP-2ic,334
-nuclia/tests/test_nua/test_predict.py,sha256=8by69GgXuOZKAgksjSjgmbNr98jCXanF1HOo29oNuMg,5798
+nuclia/tests/test_nua/test_predict.py,sha256=EKgpvHZHA4pf3ENE17j0B7CKewqHeJryR7D8JSTGJMo,8154
 nuclia/tests/test_nucliadb/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 nuclia/tests/test_nucliadb/test_crud.py,sha256=GuY76HRvt2DFaNgioKm5n0Aco1HnG7zzV_zKom5N8xc,1173
 nuclia/tests/unit/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 nuclia/tests/unit/test_export_import.py,sha256=xo_wVbjUnNlVV65ZGH7LtZ38qy39EkJp2hjOuTHC1nU,980
 nuclia/tests/unit/test_nua_responses.py,sha256=t_hIdVztTi27RWvpfTJUYcCL0lpKdZFegZIwLdaPNh8,319
-nuclia-4.9.3.dist-info/licenses/LICENSE,sha256=Ops2LTti_HJtpmWcanuUTdTY3vKDR1myJ0gmGBKC0FA,1063
-nuclia-4.9.3.dist-info/METADATA,sha256=rxQGLEDjPHayY_NPuXkmXpfsp4_BfoBC2ZjLlOjhyoM,2337
-nuclia-4.9.3.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-nuclia-4.9.3.dist-info/entry_points.txt,sha256=iZHOyXPNS54r3eQmdi5So20xO1gudI9K2oP4sQsCJRw,46
-nuclia-4.9.3.dist-info/top_level.txt,sha256=cqn_EitXOoXOSUvZnd4q6QGrhm04pg8tLAZtem-Zfdo,7
-nuclia-4.9.3.dist-info/RECORD,,
+nuclia-4.9.4.dist-info/licenses/LICENSE,sha256=Ops2LTti_HJtpmWcanuUTdTY3vKDR1myJ0gmGBKC0FA,1063
+nuclia-4.9.4.dist-info/METADATA,sha256=2fk95GmG2Qa-GGAOPxWXWgA-gAGv-vEIHMRftKOsgA0,2341
+nuclia-4.9.4.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+nuclia-4.9.4.dist-info/entry_points.txt,sha256=iZHOyXPNS54r3eQmdi5So20xO1gudI9K2oP4sQsCJRw,46
+nuclia-4.9.4.dist-info/top_level.txt,sha256=cqn_EitXOoXOSUvZnd4q6QGrhm04pg8tLAZtem-Zfdo,7
+nuclia-4.9.4.dist-info/RECORD,,