nucliadb 6.2.1.post2949__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
@@ -121,12 +121,14 @@ class AnswerStatusCode(str, Enum):
     SUCCESS = "0"
     ERROR = "-1"
     NO_CONTEXT = "-2"
+    NO_RETRIEVAL_DATA = "-3"
 
     def prettify(self) -> str:
         return {
             AnswerStatusCode.SUCCESS: "success",
             AnswerStatusCode.ERROR: "error",
             AnswerStatusCode.NO_CONTEXT: "no_context",
+            AnswerStatusCode.NO_RETRIEVAL_DATA: "no_retrieval_data",
         }[self]
 
 
@@ -266,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
@@ -442,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():
@@ -553,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

@@ -49,6 +49,7 @@ from nucliadb.search.search.chat.query import (
     ChatAuditor,
     get_find_results,
     get_relations_results,
+    maybe_audit_chat,
     rephrase_query,
     sorted_prompt_context_list,
     tokens_to_chars,
@@ -128,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,
@@ -395,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):
@@ -433,14 +437,14 @@ class NotEnoughContextAskResult(AskResult):
         """
         yield self._ndjson_encode(RetrievalAskResponseItem(results=self.main_results))
         yield self._ndjson_encode(AnswerAskResponseItem(text=NOT_ENOUGH_CONTEXT_ANSWER))
-        status = AnswerStatusCode.NO_CONTEXT
+        status = AnswerStatusCode.NO_RETRIEVAL_DATA
         yield self._ndjson_encode(StatusAskResponseItem(code=status.value, status=status.prettify()))
 
     async def json(self) -> str:
         return SyncAskResponse(
             answer=NOT_ENOUGH_CONTEXT_ANSWER,
             retrieval_results=self.main_results,
-            status=AnswerStatusCode.NO_CONTEXT,
+            status=AnswerStatusCode.NO_RETRIEVAL_DATA.prettify(),
         ).model_dump_json()
 
 
@@ -487,6 +491,31 @@ async def ask(
             resource=resource,
         )
     except NoRetrievalResultsError as err:
+        try:
+            rephrase_time = metrics.elapsed("rephrase")
+        except KeyError:
+            # Not all ask requests have a rephrase step
+            rephrase_time = None
+
+        maybe_audit_chat(
+            kbid=kbid,
+            user_id=user_id,
+            client_type=client_type,
+            origin=origin,
+            generative_answer_time=0,
+            generative_answer_first_chunk_time=0,
+            rephrase_time=rephrase_time,
+            user_query=user_query,
+            rephrased_query=rephrased_query,
+            text_answer=b"",
+            status_code=AnswerStatusCode.NO_RETRIEVAL_DATA,
+            chat_history=chat_history,
+            query_context={},
+            query_context_order={},
+            learning_id=None,
+            model=ask_request.generative_model,
+        )
+
         # If a retrieval was attempted but no results were found,
         # early return the ask endpoint without querying the generative model
         return NotEnoughContextAskResult(
@@ -536,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,
@@ -564,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

@@ -285,8 +285,8 @@ def maybe_audit_chat(
     chat_history: list[ChatContextMessage],
     query_context: PromptContext,
     query_context_order: PromptContextOrder,
-    learning_id: str,
-    model: str,
+    learning_id: Optional[str],
+    model: Optional[str],
 ):
     audit = get_audit()
     if audit is None:
@@ -324,7 +324,7 @@ def maybe_audit_chat(
 
 
 def parse_audit_answer(raw_text_answer: bytes, status_code: AnswerStatusCode) -> Optional[str]:
-    if status_code == AnswerStatusCode.NO_CONTEXT:
+    if status_code == AnswerStatusCode.NO_CONTEXT or status_code == AnswerStatusCode.NO_RETRIEVAL_DATA:
         # We don't want to audit "Not enough context to answer this." and instead set a None.
         return None
     return raw_text_answer.decode()
@@ -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.post2949.dist-info → nucliadb-6.2.1.post2971.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.2
 Name: nucliadb
-Version: 6.2.1.post2949
+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.post2949
-Requires-Dist: nucliadb-utils[cache,fastapi,storages]>=6.2.1.post2949
-Requires-Dist: nucliadb-protos>=6.2.1.post2949
-Requires-Dist: nucliadb-models>=6.2.1.post2949
+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.post2949.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=AB8E5epRR_aUsfKYfWw96WsrJUAdalZE4uolkrTor8Q,20799
+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=4PHueXt4dMF6NHuBYII--reCslbAVCzAE4LgWFAiKdY,35143
+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=4cmTxnqnvQGPDKTdaNJL4Au8aop1reHLXzkGNhf4NWg,15345
+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.post2949.dist-info/METADATA,sha256=N3Z4fQUusKc5JJ5ybVTWHvLHQH7ZK7Yb8I2oNor1vT8,4689
-nucliadb-6.2.1.post2949.dist-info/WHEEL,sha256=In9FTNxeP60KnTkGw7wk6mJPYd_dQSjEZmXdBdMCI-8,91
-nucliadb-6.2.1.post2949.dist-info/entry_points.txt,sha256=XqGfgFDuY3zXQc8ewXM2TRVjTModIq851zOsgrmaXx4,1268
-nucliadb-6.2.1.post2949.dist-info/top_level.txt,sha256=hwYhTVnX7jkQ9gJCkVrbqEG1M4lT2F_iPQND1fCzF80,20
-nucliadb-6.2.1.post2949.dist-info/zip-safe,sha256=AbpHGcgLb-kRsJGnwFEktk7uzpZOCcBY74-YBdrKVGs,1
-nucliadb-6.2.1.post2949.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,,