nucliadb 6.2.0.post2628__py3-none-any.whl → 6.2.0.post2646__py3-none-any.whl

nucliadb/search/predict.py

@@ -21,12 +21,13 @@ import json
 import os
 import random
 from enum import Enum
-from typing import Any, AsyncIterator, Dict, Literal, Optional, Union
+from typing import Any, AsyncIterator, Optional
 from unittest.mock import AsyncMock, Mock
 
 import aiohttp
 import backoff
-from pydantic import BaseModel, Field, ValidationError
+from nuclia_models.predict.generative_responses import GenerativeChunk
+from pydantic import ValidationError
 
 from nucliadb.common import datamanagers
 from nucliadb.search import logger
@@ -129,53 +130,6 @@ class AnswerStatusCode(str, Enum):
         }[self]
 
 
-class TextGenerativeResponse(BaseModel):
-    type: Literal["text"] = "text"
-    text: str
-
-
-class JSONGenerativeResponse(BaseModel):
-    type: Literal["object"] = "object"
-    object: Dict[str, Any]
-
-
-class MetaGenerativeResponse(BaseModel):
-    type: Literal["meta"] = "meta"
-    input_tokens: int
-    output_tokens: int
-    timings: dict[str, float]
-
-
-class CitationsGenerativeResponse(BaseModel):
-    type: Literal["citations"] = "citations"
-    citations: dict[str, Any]
-
-
-class RerankGenerativeResponse(BaseModel):
-    type: Literal["rerank"] = "rerank"
-    context_scores: Dict[str, float]
-
-
-class StatusGenerativeResponse(BaseModel):
-    type: Literal["status"] = "status"
-    code: str
-    details: Optional[str] = None
-
-
-GenerativeResponse = Union[
-    TextGenerativeResponse,
-    JSONGenerativeResponse,
-    MetaGenerativeResponse,
-    CitationsGenerativeResponse,
-    RerankGenerativeResponse,
-    StatusGenerativeResponse,
-]
-
-
-class GenerativeChunk(BaseModel):
-    chunk: GenerativeResponse = Field(..., discriminator="type")
-
-
 async def start_predict_engine():
     if nuclia_settings.dummy_predict:
         predict_util = DummyPredictEngine()

nucliadb/search/search/chat/ask.py

@@ -22,6 +22,14 @@ import functools
 import json
 from typing import AsyncGenerator, Optional, cast
 
+from nuclia_models.predict.generative_responses import (
+    CitationsGenerativeResponse,
+    GenerativeChunk,
+    JSONGenerativeResponse,
+    MetaGenerativeResponse,
+    StatusGenerativeResponse,
+    TextGenerativeResponse,
+)
 from pydantic_core import ValidationError
 
 from nucliadb.common.datamanagers.exceptions import KnowledgeBoxNotFound
@@ -29,13 +37,7 @@ from nucliadb.models.responses import HTTPClientError
 from nucliadb.search import logger, predict
 from nucliadb.search.predict import (
     AnswerStatusCode,
-    CitationsGenerativeResponse,
-    GenerativeChunk,
-    JSONGenerativeResponse,
-    MetaGenerativeResponse,
     RephraseMissingContextError,
-    StatusGenerativeResponse,
-    TextGenerativeResponse,
 )
 from nucliadb.search.search.chat.exceptions import (
     AnswerJsonSchemaTooLong,
@@ -277,6 +279,8 @@ class AskResult:
                 tokens=AskTokens(
                     input=self._metadata.input_tokens,
                     output=self._metadata.output_tokens,
+                    input_nuclia=self._metadata.input_nuclia_tokens,
+                    output_nuclia=self._metadata.output_nuclia_tokens,
                 ),
                 timings=AskTimings(
                     generative_first_chunk=self._metadata.timings.get("generative_first_chunk"),
@@ -317,6 +321,8 @@ class AskResult:
                 tokens=AskTokens(
                     input=self._metadata.input_tokens,
                     output=self._metadata.output_tokens,
+                    input_nuclia=self._metadata.input_nuclia_tokens,
+                    output_nuclia=self._metadata.output_nuclia_tokens,
                 ),
                 timings=AskTimings(
                     generative_first_chunk=self._metadata.timings.get("generative_first_chunk"),

{nucliadb-6.2.0.post2628.dist-info → nucliadb-6.2.0.post2646.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: nucliadb
-Version: 6.2.0.post2628
+Version: 6.2.0.post2646
 Home-page: https://docs.nuclia.dev/docs/management/nucliadb/intro
 Author: NucliaDB Community
 Author-email: nucliadb@nuclia.com
@@ -22,12 +22,13 @@ Classifier: Programming Language :: Python :: 3.12
 Classifier: Programming Language :: Python :: 3 :: Only
 Requires-Python: >=3.9, <4
 Description-Content-Type: text/markdown
-Requires-Dist: nucliadb-telemetry[all]>=6.2.0.post2628
-Requires-Dist: nucliadb-utils[cache,fastapi,storages]>=6.2.0.post2628
-Requires-Dist: nucliadb-protos>=6.2.0.post2628
-Requires-Dist: nucliadb-models>=6.2.0.post2628
+Requires-Dist: nucliadb-telemetry[all]>=6.2.0.post2646
+Requires-Dist: nucliadb-utils[cache,fastapi,storages]>=6.2.0.post2646
+Requires-Dist: nucliadb-protos>=6.2.0.post2646
+Requires-Dist: nucliadb-models>=6.2.0.post2646
 Requires-Dist: nucliadb-admin-assets>=1.0.0.post1224
 Requires-Dist: nucliadb-node-binding>=2.26.0
+Requires-Dist: nuclia-models>=0.24.2
 Requires-Dist: uvicorn
 Requires-Dist: argdantic
 Requires-Dist: aiohttp>=3.9.4

{nucliadb-6.2.0.post2628.dist-info → nucliadb-6.2.0.post2646.dist-info}/RECORD

@@ -184,7 +184,7 @@ nucliadb/search/__init__.py,sha256=tnypbqcH4nBHbGpkINudhKgdLKpwXQCvDtPchUlsyY4,1
 nucliadb/search/app.py,sha256=6UV7rO0f3w5bNFXLdQM8bwUwXayMGnM4hF6GGv7WPv4,4260
 nucliadb/search/lifecycle.py,sha256=DW8v4WUi4rZqc7xTOi3rE67W7877WG7fH9oTZbolHdE,2099
 nucliadb/search/openapi.py,sha256=t3Wo_4baTrfPftg2BHsyLWNZ1MYn7ZRdW7ht-wFOgRs,1016
-nucliadb/search/predict.py,sha256=R0LtFV6WSWtFJjd0rrPLBSPSUrC4cMedK-pqWzalUUU,21808
+nucliadb/search/predict.py,sha256=AB8E5epRR_aUsfKYfWw96WsrJUAdalZE4uolkrTor8Q,20799
 nucliadb/search/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 nucliadb/search/run.py,sha256=aFb-CXRi_C8YMpP_ivNj8KW1BYhADj88y8K9Lr_nUPI,1402
 nucliadb/search/settings.py,sha256=vem3EcyYlTPSim0kEK-xe-erF4BZg0CT_LAb8ZRQAE8,1684
@@ -228,7 +228,7 @@ nucliadb/search/search/shards.py,sha256=mM2aCHWhl_gwkCENXDShPukS-_qnB5tFS3UAJuzM
 nucliadb/search/search/summarize.py,sha256=ksmYPubEQvAQgfPdZHfzB_rR19B2ci4IYZ6jLdHxZo8,4996
 nucliadb/search/search/utils.py,sha256=iF2tbBA56gRMJH1TlE2hMrqeXqjoeOPt4KgRdp2m9Ek,3313
 nucliadb/search/search/chat/__init__.py,sha256=cp15ZcFnHvpcu_5-aK2A4uUyvuZVV_MJn4bIXMa20ks,835
-nucliadb/search/search/chat/ask.py,sha256=7yUPEMluZ553O4FdcghyQI3Hw042P3QL06T0AwDctJI,33799
+nucliadb/search/search/chat/ask.py,sha256=_emhBlraFGsScWOmEbc8MbYi-aES9XZ1NnhSXR_79qQ,34138
 nucliadb/search/search/chat/exceptions.py,sha256=Siy4GXW2L7oPhIR86H3WHBhE9lkV4A4YaAszuGGUf54,1356
 nucliadb/search/search/chat/images.py,sha256=PA8VWxT5_HUGfW1ULhKTK46UBsVyINtWWqEM1ulzX1E,3095
 nucliadb/search/search/chat/prompt.py,sha256=TIzjI_882hJ--KLKCY8rJomtJ_CMJ-MHYtHqivgG8Lk,46819
@@ -332,9 +332,9 @@ nucliadb/writer/tus/local.py,sha256=7jYa_w9b-N90jWgN2sQKkNcomqn6JMVBOVeDOVYJHto,
 nucliadb/writer/tus/s3.py,sha256=vF0NkFTXiXhXq3bCVXXVV-ED38ECVoUeeYViP8uMqcU,8357
 nucliadb/writer/tus/storage.py,sha256=ToqwjoYnjI4oIcwzkhha_MPxi-k4Jk3Lt55zRwaC1SM,2903
 nucliadb/writer/tus/utils.py,sha256=MSdVbRsRSZVdkaum69_0wku7X3p5wlZf4nr6E0GMKbw,2556
-nucliadb-6.2.0.post2628.dist-info/METADATA,sha256=RP8ybjNxSveiosTS71zdP1M6A7V5x1TQCJDMZ61jFxQ,4390
-nucliadb-6.2.0.post2628.dist-info/WHEEL,sha256=PZUExdf71Ui_so67QXpySuHtCi3-J3wvF4ORK6k_S8U,91
-nucliadb-6.2.0.post2628.dist-info/entry_points.txt,sha256=XqGfgFDuY3zXQc8ewXM2TRVjTModIq851zOsgrmaXx4,1268
-nucliadb-6.2.0.post2628.dist-info/top_level.txt,sha256=hwYhTVnX7jkQ9gJCkVrbqEG1M4lT2F_iPQND1fCzF80,20
-nucliadb-6.2.0.post2628.dist-info/zip-safe,sha256=AbpHGcgLb-kRsJGnwFEktk7uzpZOCcBY74-YBdrKVGs,1
-nucliadb-6.2.0.post2628.dist-info/RECORD,,
+nucliadb-6.2.0.post2646.dist-info/METADATA,sha256=iPtmKTExj1QBmR76v6ok474FV7HIBJBBn3JgxeCNl6U,4427
+nucliadb-6.2.0.post2646.dist-info/WHEEL,sha256=PZUExdf71Ui_so67QXpySuHtCi3-J3wvF4ORK6k_S8U,91
+nucliadb-6.2.0.post2646.dist-info/entry_points.txt,sha256=XqGfgFDuY3zXQc8ewXM2TRVjTModIq851zOsgrmaXx4,1268
+nucliadb-6.2.0.post2646.dist-info/top_level.txt,sha256=hwYhTVnX7jkQ9gJCkVrbqEG1M4lT2F_iPQND1fCzF80,20
+nucliadb-6.2.0.post2646.dist-info/zip-safe,sha256=AbpHGcgLb-kRsJGnwFEktk7uzpZOCcBY74-YBdrKVGs,1
+nucliadb-6.2.0.post2646.dist-info/RECORD,,