nuclia 4.9.2__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,
@@ -42,6 +44,8 @@ from nuclia.lib.nua_responses import (
     PushResponseV2,
     QueryInfo,
     RephraseModel,
+    RerankModel,
+    RerankResponse,
     RestrictedIDString,
     Sentence,
     Source,
@@ -77,6 +81,7 @@ PUSH_PROCESS = "/api/v2/processing/push"
 SCHEMA = "/api/v1/learning/configuration/schema"
 SCHEMA_KBID = "/api/v1/schema"
 CONFIG = "/api/v1/config"
+RERANK = "/api/v1/predict/rerank"
 
 ConvertType = TypeVar("ConvertType", bound=BaseModel)
 
@@ -125,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:
@@ -140,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]:
@@ -148,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
@@ -191,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,
@@ -209,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:
@@ -217,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:
@@ -232,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
@@ -249,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:
@@ -263,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:
@@ -285,6 +331,7 @@ class NuaClient:
             payload=body.model_dump(),
             output=SummarizedModel,
             timeout=timeout,
+            extra_headers=extra_headers,
         )
 
     def rephrase(
@@ -321,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,
         )
@@ -410,6 +459,20 @@ class NuaClient:
         activity_endpoint = f"{self.url}{STATUS_PROCESS}/{process_id}"
         return self._request("GET", activity_endpoint, ProcessRequestStatus)
 
+    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,
+            extra_headers=extra_headers,
+        )
+
 
 class AsyncNuaClient:
     def __init__(
@@ -445,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:
@@ -460,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]:
@@ -468,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
@@ -518,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,
@@ -537,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:
@@ -545,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(
@@ -564,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:
@@ -576,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
@@ -593,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:
@@ -608,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:
@@ -630,6 +732,7 @@ class AsyncNuaClient:
             payload=body.model_dump(),
             output=SummarizedModel,
             timeout=timeout,
+            extra_headers=extra_headers,
         )
 
     async def rephrase(
@@ -663,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(
@@ -792,3 +900,15 @@ class AsyncNuaClient:
         return await self._request(
             "GET", activity_endpoint, output=ProcessRequestStatus
         )
+
+    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,
+            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):
@@ -557,3 +563,16 @@ class QueryInfo(BaseModel):
     max_context: int
     entities: Optional[TokenSearch]
     sentence: Optional[SentenceSearch]
+
+
+class RerankModel(BaseModel):
+    question: str
+    user_id: str
+    context: dict[str, str] = {}
+
+
+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())