nucliadb 6.2.1.post2954__py3-none-any.whl → 6.2.1.post2971__py3-none-any.whl

nucliadb/search/predict.py

@@ -21,7 +21,7 @@ import json
 import os
 import random
 from enum import Enum
-from typing import Any, AsyncIterator, Optional
+from typing import Any, AsyncGenerator, Optional
 from unittest.mock import AsyncMock, Mock
 
 import aiohttp
@@ -268,7 +268,7 @@ class PredictEngine:
     @predict_observer.wrap({"type": "chat_ndjson"})
     async def chat_query_ndjson(
         self, kbid: str, item: ChatModel
-    ) -> tuple[str, str, AsyncIterator[GenerativeChunk]]:
+    ) -> tuple[str, str, AsyncGenerator[GenerativeChunk, None]]:
         """
         Chat query using the new stream format
         Format specs: https://github.com/ndjson/ndjson-spec
@@ -444,7 +444,7 @@ class DummyPredictEngine(PredictEngine):
 
     async def chat_query_ndjson(
         self, kbid: str, item: ChatModel
-    ) -> tuple[str, str, AsyncIterator[GenerativeChunk]]:
+    ) -> tuple[str, str, AsyncGenerator[GenerativeChunk, None]]:
         self.calls.append(("chat_query_ndjson", item))
 
         async def generate():
@@ -555,7 +555,7 @@ def get_answer_generator(response: aiohttp.ClientResponse):
 
 def get_chat_ndjson_generator(
     response: aiohttp.ClientResponse,
-) -> AsyncIterator[GenerativeChunk]:
+) -> AsyncGenerator[GenerativeChunk, None]:
     async def _parse_generative_chunks(gen):
         async for chunk in gen:
             try:

nucliadb/search/search/chat/ask.py

@@ -129,7 +129,7 @@ class AskResult:
         main_results: KnowledgeboxFindResults,
         prequeries_results: Optional[list[PreQueryResult]],
         nuclia_learning_id: Optional[str],
-        predict_answer_stream: AsyncGenerator[GenerativeChunk, None],
+        predict_answer_stream: Optional[AsyncGenerator[GenerativeChunk, None]],
         prompt_context: PromptContext,
         prompt_context_order: PromptContextOrder,
         auditor: ChatAuditor,
@@ -396,6 +396,9 @@ class AskResult:
         This method does not assume any order in the stream of items, but it assumes that at least
         the answer text is streamed in order.
         """
+        if self.predict_answer_stream is None:
+            # In some cases, clients may want to skip the answer generation step
+            return
         async for generative_chunk in self.predict_answer_stream:
             item = generative_chunk.chunk
             if isinstance(item, TextGenerativeResponse):
@@ -562,14 +565,18 @@ async def ask(
         rerank_context=False,
         top_k=ask_request.top_k,
     )
-    with metrics.time("stream_start"):
-        predict = get_predict()
-        (
-            nuclia_learning_id,
-            nuclia_learning_model,
-            predict_answer_stream,
-        ) = await predict.chat_query_ndjson(kbid, chat_model)
-        debug_chat_model = chat_model
+
+    nuclia_learning_id = None
+    nuclia_learning_model = None
+    predict_answer_stream = None
+    if ask_request.generate_answer:
+        with metrics.time("stream_start"):
+            predict = get_predict()
+            (
+                nuclia_learning_id,
+                nuclia_learning_model,
+                predict_answer_stream,
+            ) = await predict.chat_query_ndjson(kbid, chat_model)
 
     auditor = ChatAuditor(
         kbid=kbid,
@@ -590,13 +597,13 @@ async def ask(
         main_results=retrieval_results.main_query,
         prequeries_results=retrieval_results.prequeries,
         nuclia_learning_id=nuclia_learning_id,
-        predict_answer_stream=predict_answer_stream,  # type: ignore
+        predict_answer_stream=predict_answer_stream,
         prompt_context=prompt_context,
         prompt_context_order=prompt_context_order,
         auditor=auditor,
         metrics=metrics,
         best_matches=retrieval_results.best_matches,
-        debug_chat_model=debug_chat_model,
+        debug_chat_model=chat_model,
     )
 
 

nucliadb/search/search/chat/query.py

@@ -349,7 +349,7 @@ class ChatAuditor:
         learning_id: Optional[str],
         query_context: PromptContext,
         query_context_order: PromptContextOrder,
-        model: str,
+        model: Optional[str],
     ):
         self.kbid = kbid
         self.user_id = user_id

{nucliadb-6.2.1.post2954.dist-info → nucliadb-6.2.1.post2971.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.2
 Name: nucliadb
-Version: 6.2.1.post2954
+Version: 6.2.1.post2971
 Home-page: https://docs.nuclia.dev/docs/management/nucliadb/intro
 Author: NucliaDB Community
 Author-email: nucliadb@nuclia.com
@@ -22,10 +22,10 @@ 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.1.post2954
-Requires-Dist: nucliadb-utils[cache,fastapi,storages]>=6.2.1.post2954
-Requires-Dist: nucliadb-protos>=6.2.1.post2954
-Requires-Dist: nucliadb-models>=6.2.1.post2954
+Requires-Dist: nucliadb-telemetry[all]>=6.2.1.post2971
+Requires-Dist: nucliadb-utils[cache,fastapi,storages]>=6.2.1.post2971
+Requires-Dist: nucliadb-protos>=6.2.1.post2971
+Requires-Dist: nucliadb-models>=6.2.1.post2971
 Requires-Dist: nucliadb-admin-assets>=1.0.0.post1224
 Requires-Dist: nucliadb-node-binding>=2.26.0
 Requires-Dist: nuclia-models>=0.24.2

{nucliadb-6.2.1.post2954.dist-info → nucliadb-6.2.1.post2971.dist-info}/RECORD

@@ -189,7 +189,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=Jan9hMsiIE5WEcyjHMA79Gz8-2pMbHbb8Kb3SgoXw8M,20897
+nucliadb/search/predict.py,sha256=EWOiWVUX9U_TE19Cl6bpCr6Mjs7hjuvCcG26C7e6KnQ,20919
 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
@@ -234,11 +234,11 @@ 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=sopRYv2yyNUr-rSILpRm4rWI6exBQoGTaONhW1VBiVY,36008
+nucliadb/search/search/chat/ask.py,sha256=tE1Q5V58oLMCo-T9s0N6Kko-1RWn1e4kHfbbPBsD2uU,36266
 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=r2JTiRWH3YHPdeRAG5w6gD0g0fWVxdTjYIR86qAVa7k,47106
-nucliadb/search/search/chat/query.py,sha256=EfZVJ_Nxo2Jwrbl4ICUL61srKv22PCPt1tZDXMZZQq4,15418
+nucliadb/search/search/chat/query.py,sha256=y7W5VuKl1XiZuNsxZIcxxHcFXSG6It2W5CoftZ-ekAc,15428
 nucliadb/search/search/query_parser/__init__.py,sha256=cp15ZcFnHvpcu_5-aK2A4uUyvuZVV_MJn4bIXMa20ks,835
 nucliadb/search/search/query_parser/exceptions.py,sha256=tuzl7ZyvVsRz6u0_3zMe60vx39nd3pi641prs-5nC0E,872
 nucliadb/search/search/query_parser/models.py,sha256=-VlCDXUCgOroAZw1Leqhj2VMgRv_CD2w40PXXOBLaUM,2332
@@ -340,9 +340,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.1.post2954.dist-info/METADATA,sha256=uqKCYU8MpuEtw2VE0oX1sX-WQTRD-ACMAf1jRBOO7P0,4689
-nucliadb-6.2.1.post2954.dist-info/WHEEL,sha256=In9FTNxeP60KnTkGw7wk6mJPYd_dQSjEZmXdBdMCI-8,91
-nucliadb-6.2.1.post2954.dist-info/entry_points.txt,sha256=XqGfgFDuY3zXQc8ewXM2TRVjTModIq851zOsgrmaXx4,1268
-nucliadb-6.2.1.post2954.dist-info/top_level.txt,sha256=hwYhTVnX7jkQ9gJCkVrbqEG1M4lT2F_iPQND1fCzF80,20
-nucliadb-6.2.1.post2954.dist-info/zip-safe,sha256=AbpHGcgLb-kRsJGnwFEktk7uzpZOCcBY74-YBdrKVGs,1
-nucliadb-6.2.1.post2954.dist-info/RECORD,,
+nucliadb-6.2.1.post2971.dist-info/METADATA,sha256=Ftcrf80Q8tp56jdibzPDOdTJq2jG0aEGuC-bMD-4418,4689
+nucliadb-6.2.1.post2971.dist-info/WHEEL,sha256=In9FTNxeP60KnTkGw7wk6mJPYd_dQSjEZmXdBdMCI-8,91
+nucliadb-6.2.1.post2971.dist-info/entry_points.txt,sha256=XqGfgFDuY3zXQc8ewXM2TRVjTModIq851zOsgrmaXx4,1268
+nucliadb-6.2.1.post2971.dist-info/top_level.txt,sha256=hwYhTVnX7jkQ9gJCkVrbqEG1M4lT2F_iPQND1fCzF80,20
+nucliadb-6.2.1.post2971.dist-info/zip-safe,sha256=AbpHGcgLb-kRsJGnwFEktk7uzpZOCcBY74-YBdrKVGs,1
+nucliadb-6.2.1.post2971.dist-info/RECORD,,