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

nuclia/sdk/predict.py

@@ -13,6 +13,8 @@ from nuclia.lib.nua_responses import (
     ConfigSchema,
     LearningConfigurationCreation,
     QueryInfo,
+    RerankModel,
+    RerankResponse,
     Sentence,
     StoredLearningConfiguration,
     SummarizedModel,
@@ -50,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(
@@ -61,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):
@@ -86,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):
@@ -103,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(
@@ -133,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(
@@ -143,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
 
@@ -160,7 +222,26 @@ 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, show_consumption: bool = False, **kwargs
+    ) -> RerankResponse:
+        """
+        Perform a reranking of the results based on the question and context provided.
+
+        :param request: RerankModel
+        :return: RerankResponse
+        """
+        nc: NuaClient = kwargs["nc"]
+        return nc.rerank(
+            request,
+            extra_headers={"X-Show-Consumption": str(show_consumption).lower()},
+        )
 
 
 class AsyncNucliaPredict:
@@ -195,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):
@@ -214,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):
@@ -231,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(
@@ -246,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(
@@ -285,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
@@ -298,4 +426,23 @@ 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, show_consumption: bool = False, **kwargs
+    ) -> RerankResponse:
+        """
+        Perform a reranking of the results based on the question and context provided.
+
+        :param request: RerankModel
+        :return: RerankResponse
+        """
+        nc: AsyncNuaClient = kwargs["nc"]
+        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