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

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/sdk/task.py

@@ -6,9 +6,9 @@ from nuclia.sdk.auth import NucliaAuth, AsyncNucliaAuth
 from nuclia_models.worker.tasks import (
     ApplyOptions,
     TaskStartKB,
-    TaskResponse,
     TaskList,
     TaskName,
+    TaskResponse,
     PARAMETERS_TYPING,
     PublicTaskSet,
     TASKS,
@@ -64,7 +64,7 @@ class NucliaTask:
         return TaskResponse.model_validate(response.json())
 
     @kb
-    def delete(self, *args, task_id: str, **kwargs):
+    def delete(self, *args, task_id: str, cleanup: bool = False, **kwargs):
         """
         Delete task
 
@@ -72,7 +72,7 @@ class NucliaTask:
         """
         ndb: NucliaDBClient = kwargs["ndb"]
         try:
-            _ = ndb.delete_task(task_id=task_id)
+            _ = ndb.delete_task(task_id=task_id, cleanup=cleanup)
         except InvalidPayload:
             pass
 
@@ -158,7 +158,7 @@ class AsyncNucliaTask:
         return TaskResponse.model_validate(response.json())
 
     @kb
-    async def delete(self, *args, task_id: str, **kwargs):
+    async def delete(self, *args, task_id: str, cleanup: bool = False, **kwargs):
         """
         Delete task
 
@@ -166,7 +166,7 @@ class AsyncNucliaTask:
         """
         ndb: AsyncNucliaDBClient = kwargs["ndb"]
         try:
-            _ = await ndb.delete_task(task_id=task_id)
+            _ = await ndb.delete_task(task_id=task_id, cleanup=cleanup)
         except InvalidPayload:
             pass
 

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