nucliadb 6.7.2.post4911__py3-none-any.whl → 6.7.2.post4917__py3-none-any.whl

nucliadb/search/predict.py

@@ -447,6 +447,10 @@ class DummyPredictEngine(PredictEngine):
         self.cluster_url = "http://localhost:8000"
         self.public_url = "http://localhost:8000"
         self.calls = []
+        self.ndjson_reasoning = [
+            b'{"chunk": {"type": "reasoning", "text": "dummy "}}\n',
+            b'{"chunk": {"type": "reasoning", "text": "reasoning"}}\n',
+        ]
         self.ndjson_answer = [
             b'{"chunk": {"type": "text", "text": "valid "}}\n',
             b'{"chunk": {"type": "text", "text": "answer "}}\n',
@@ -482,8 +486,11 @@ class DummyPredictEngine(PredictEngine):
         self.calls.append(("chat_query_ndjson", item))
 
         async def generate():
-            for item in self.ndjson_answer:
-                yield GenerativeChunk.model_validate_json(item)
+            if item.reasoning is not False:
+                for chunk in self.ndjson_reasoning:
+                    yield GenerativeChunk.model_validate_json(chunk)
+            for chunk in self.ndjson_answer:
+                yield GenerativeChunk.model_validate_json(chunk)
 
         return (DUMMY_LEARNING_ID, DUMMY_LEARNING_MODEL, generate())
 

nucliadb/search/search/chat/ask.py

@@ -20,7 +20,7 @@
 import dataclasses
 import functools
 import json
-from typing import AsyncGenerator, Optional, cast
+from typing import AsyncGenerator, Optional, Union, cast
 
 from nuclia_models.common.consumption import Consumption
 from nuclia_models.predict.generative_responses import (
@@ -28,6 +28,7 @@ from nuclia_models.predict.generative_responses import (
     GenerativeChunk,
     JSONGenerativeResponse,
     MetaGenerativeResponse,
+    ReasoningGenerativeResponse,
     StatusGenerativeResponse,
     TextGenerativeResponse,
 )
@@ -102,6 +103,7 @@ from nucliadb_models.search import (
     PromptContext,
     PromptContextOrder,
     RagStrategyName,
+    ReasoningAskResponseItem,
     Relations,
     RelationsAskResponseItem,
     RetrievalAskResponseItem,
@@ -167,6 +169,7 @@ class AskResult:
 
         # Computed from the predict chat answer stream
         self._answer_text = ""
+        self._reasoning_text: Optional[str] = None
         self._object: Optional[JSONGenerativeResponse] = None
         self._status: Optional[StatusGenerativeResponse] = None
         self._citations: Optional[CitationsGenerativeResponse] = None
@@ -220,12 +223,23 @@ class AskResult:
     async def _stream(self) -> AsyncGenerator[AskResponseItemType, None]:
         # First, stream out the predict answer
         first_chunk_yielded = False
+        first_reasoning_chunk_yielded = False
         with self.metrics.time("stream_predict_answer"):
             async for answer_chunk in self._stream_predict_answer_text():
-                yield AnswerAskResponseItem(text=answer_chunk)
-                if not first_chunk_yielded:
-                    self.metrics.record_first_chunk_yielded()
-                    first_chunk_yielded = True
+                if isinstance(answer_chunk, TextGenerativeResponse):
+                    yield AnswerAskResponseItem(text=answer_chunk.text)
+                    if not first_chunk_yielded:
+                        self.metrics.record_first_chunk_yielded()
+                        first_chunk_yielded = True
+                elif isinstance(answer_chunk, ReasoningGenerativeResponse):
+                    yield ReasoningAskResponseItem(text=answer_chunk.text)
+                    if not first_reasoning_chunk_yielded:
+                        self.metrics.record_first_reasoning_chunk_yielded()
+                        first_reasoning_chunk_yielded = True
+                else:
+                    # This is a trick so mypy generates an error if this branch can be reached,
+                    # that is, if we are missing some ifs
+                    _a: int = "a"
 
         if self._object is not None:
             yield JSONAskResponseItem(object=self._object.object)
@@ -274,8 +288,10 @@ class AskResult:
             audit_answer = json.dumps(self._object.object).encode("utf-8")
         self.auditor.audit(
             text_answer=audit_answer,
+            text_reasoning=self._reasoning_text,
             generative_answer_time=self.metrics["stream_predict_answer"],
             generative_answer_first_chunk_time=self.metrics.get_first_chunk_time() or 0,
+            generative_reasoning_first_chunk_time=self.metrics.get_first_reasoning_chunk_time(),
             rephrase_time=self.metrics.get("rephrase"),
             status_code=self.status_code,
         )
@@ -384,6 +400,7 @@ class AskResult:
 
         response = SyncAskResponse(
             answer=self._answer_text,
+            reasoning=self._reasoning_text,
             answer_json=answer_json,
             status=self.status_code.prettify(),
             relations=self._relations,
@@ -420,7 +437,9 @@ class AskResult:
                 )
         return self._relations
 
-    async def _stream_predict_answer_text(self) -> AsyncGenerator[str, None]:
+    async def _stream_predict_answer_text(
+        self,
+    ) -> AsyncGenerator[Union[TextGenerativeResponse, ReasoningGenerativeResponse], None]:
         """
         Reads the stream of the generative model, yielding the answer text but also parsing
         other items like status codes, citations and miscellaneous metadata.
@@ -435,7 +454,13 @@ class AskResult:
             item = generative_chunk.chunk
             if isinstance(item, TextGenerativeResponse):
                 self._answer_text += item.text
-                yield item.text
+                yield item
+            elif isinstance(item, ReasoningGenerativeResponse):
+                if self._reasoning_text is None:
+                    self._reasoning_text = item.text
+                else:
+                    self._reasoning_text += item.text
+                yield item
             elif isinstance(item, JSONGenerativeResponse):
                 self._object = item
             elif isinstance(item, StatusGenerativeResponse):
@@ -559,11 +584,13 @@ async def ask(
             origin=origin,
             generative_answer_time=0,
             generative_answer_first_chunk_time=0,
+            generative_reasoning_first_chunk_time=None,
             rephrase_time=metrics.get("rephrase"),
             user_query=user_query,
             rephrased_query=rephrased_query,
             retrieval_rephrase_query=err.main_query.rephrased_query if err.main_query else None,
             text_answer=b"",
+            text_reasoning=None,
             status_code=AnswerStatusCode.NO_RETRIEVAL_DATA,
             chat_history=chat_history,
             query_context={},
@@ -625,6 +652,7 @@ async def ask(
         json_schema=ask_request.answer_json_schema,
         rerank_context=False,
         top_k=ask_request.top_k,
+        reasoning=ask_request.reasoning,
     )
 
     nuclia_learning_id = None

nucliadb/search/search/chat/query.py

@@ -308,11 +308,13 @@ def maybe_audit_chat(
     origin: str,
     generative_answer_time: float,
     generative_answer_first_chunk_time: float,
+    generative_reasoning_first_chunk_time: Optional[float],
     rephrase_time: Optional[float],
     user_query: str,
     rephrased_query: Optional[str],
     retrieval_rephrase_query: Optional[str],
     text_answer: bytes,
+    text_reasoning: Optional[str],
     status_code: AnswerStatusCode,
     chat_history: list[ChatContextMessage],
     query_context: PromptContext,
@@ -344,12 +346,14 @@ def maybe_audit_chat(
         question=user_query,
         generative_answer_time=generative_answer_time,
         generative_answer_first_chunk_time=generative_answer_first_chunk_time,
+        generative_reasoning_first_chunk_time=generative_reasoning_first_chunk_time,
         rephrase_time=rephrase_time,
         rephrased_question=rephrased_query,
         retrieval_rephrased_question=retrieval_rephrase_query,
         chat_context=chat_history_context,
         retrieved_context=chat_retrieved_context,
         answer=audit_answer,
+        reasoning=text_reasoning,
         learning_id=learning_id,
         status_code=int(status_code.value),
         model=model,
@@ -401,8 +405,10 @@ class ChatAuditor:
     def audit(
         self,
         text_answer: bytes,
+        text_reasoning: Optional[str],
         generative_answer_time: float,
         generative_answer_first_chunk_time: float,
+        generative_reasoning_first_chunk_time: Optional[float],
         rephrase_time: Optional[float],
         status_code: AnswerStatusCode,
     ):
@@ -416,8 +422,10 @@ class ChatAuditor:
             retrieval_rephrase_query=self.retrieval_rephrased_query,
             generative_answer_time=generative_answer_time,
             generative_answer_first_chunk_time=generative_answer_first_chunk_time,
+            generative_reasoning_first_chunk_time=generative_reasoning_first_chunk_time,
             rephrase_time=rephrase_time,
             text_answer=text_answer,
+            text_reasoning=text_reasoning,
             status_code=status_code,
             chat_history=self.chat_history,
             query_context=self.query_context,

nucliadb/search/search/metrics.py

@@ -49,6 +49,10 @@ buckets = [
 ]
 
 generative_first_chunk_histogram = metrics.Histogram(
+    name="generative_reasoning_first_chunk",
+    buckets=buckets,
+)
+reasoning_first_chunk_histogram = metrics.Histogram(
     name="generative_first_chunk",
     buckets=buckets,
 )
@@ -107,12 +111,24 @@ class AskMetrics(Metrics):
         super().__init__(id="ask")
         self.global_start = time.monotonic()
         self.first_chunk_yielded_at: Optional[float] = None
+        self.first_reasoning_chunk_yielded_at: Optional[float] = None
 
     def record_first_chunk_yielded(self):
         self.first_chunk_yielded_at = time.monotonic()
         generative_first_chunk_histogram.observe(self.first_chunk_yielded_at - self.global_start)
 
+    def record_first_reasoning_chunk_yielded(self):
+        self.first_reasoning_chunk_yielded_at = time.monotonic()
+        reasoning_first_chunk_histogram.observe(
+            self.first_reasoning_chunk_yielded_at - self.global_start
+        )
+
     def get_first_chunk_time(self) -> Optional[float]:
         if self.first_chunk_yielded_at is None:
             return None
         return self.first_chunk_yielded_at - self.global_start
+
+    def get_first_reasoning_chunk_time(self) -> Optional[float]:
+        if self.first_reasoning_chunk_yielded_at is None:
+            return None
+        return self.first_reasoning_chunk_yielded_at - self.global_start

nucliadb/search/search/predict_proxy.py

@@ -28,6 +28,7 @@ from multidict import CIMultiDictProxy
 from nuclia_models.predict.generative_responses import (
     GenerativeChunk,
     JSONGenerativeResponse,
+    ReasoningGenerativeResponse,
     StatusGenerativeResponse,
     TextGenerativeResponse,
 )
@@ -140,8 +141,10 @@ async def predict_proxy(
                 client_type=client_type,
                 origin=origin,
                 text_answer=content,
+                text_reasoning=None,
                 generative_answer_time=metrics[PREDICT_ANSWER_METRIC],
                 generative_answer_first_chunk_time=None,
+                generative_reasoning_first_chunk_time=None,
                 status_code=AnswerStatusCode(str(llm_status_code)),
             )
 
@@ -173,23 +176,33 @@ async def chat_streaming_generator(
     is_json: bool,
 ):
     first = True
+    first_reasoning = True
     status_code = AnswerStatusCode.ERROR.value
     text_answer = ""
+    text_reasoning = ""
     json_object = None
     metrics = AskMetrics()
     with metrics.time(PREDICT_ANSWER_METRIC):
         async for chunk in predict_response.content:
-            if first:
-                metrics.record_first_chunk_yielded()
-                first = False
-
             yield chunk
 
             if is_json:
                 try:
                     parsed_chunk = GenerativeChunk.model_validate_json(chunk).chunk
+                    if first and isinstance(
+                        parsed_chunk,
+                        (TextGenerativeResponse, JSONGenerativeResponse, StatusGenerativeResponse),
+                    ):
+                        metrics.record_first_chunk_yielded()
+                        first = False
+
                     if isinstance(parsed_chunk, TextGenerativeResponse):
                         text_answer += parsed_chunk.text
+                    elif isinstance(parsed_chunk, ReasoningGenerativeResponse):
+                        if first_reasoning:
+                            metrics.record_first_reasoning_chunk_yielded()
+                            first_reasoning = False
+                        text_reasoning += parsed_chunk.text
                     elif isinstance(parsed_chunk, JSONGenerativeResponse):
                         json_object = parsed_chunk.object
                     elif isinstance(parsed_chunk, StatusGenerativeResponse):
@@ -218,8 +231,10 @@ async def chat_streaming_generator(
         client_type=client_type,
         origin=origin,
         text_answer=text_answer.encode() if json_object is None else json.dumps(json_object).encode(),
+        text_reasoning=text_reasoning if text_reasoning else None,
         generative_answer_time=metrics[PREDICT_ANSWER_METRIC],
         generative_answer_first_chunk_time=metrics.get_first_chunk_time(),
+        generative_reasoning_first_chunk_time=metrics.get_first_reasoning_chunk_time(),
         status_code=AnswerStatusCode(status_code),
     )
 
@@ -232,8 +247,10 @@ def audit_predict_proxy_endpoint(
     client_type: NucliaDBClientType,
     origin: str,
     text_answer: bytes,
+    text_reasoning: Optional[str],
     generative_answer_time: float,
     generative_answer_first_chunk_time: Optional[float],
+    generative_reasoning_first_chunk_time: Optional[float],
     status_code: AnswerStatusCode,
 ):
     maybe_audit_chat(
@@ -250,8 +267,10 @@ def audit_predict_proxy_endpoint(
         query_context_order={},
         model=headers.get(NUCLIA_LEARNING_MODEL_HEADER),
         text_answer=text_answer,
+        text_reasoning=text_reasoning,
         generative_answer_time=generative_answer_time,
         generative_answer_first_chunk_time=generative_answer_first_chunk_time or 0,
+        generative_reasoning_first_chunk_time=generative_reasoning_first_chunk_time,
         rephrase_time=None,
         status_code=status_code,
     )

{nucliadb-6.7.2.post4911.dist-info → nucliadb-6.7.2.post4917.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: nucliadb
-Version: 6.7.2.post4911
+Version: 6.7.2.post4917
 Summary: NucliaDB
 Author-email: Nuclia <nucliadb@nuclia.com>
 License-Expression: AGPL-3.0-or-later
@@ -19,13 +19,13 @@ Classifier: Programming Language :: Python :: 3.12
 Classifier: Programming Language :: Python :: 3 :: Only
 Requires-Python: <4,>=3.9
 Description-Content-Type: text/markdown
-Requires-Dist: nucliadb-telemetry[all]>=6.7.2.post4911
-Requires-Dist: nucliadb-utils[cache,fastapi,storages]>=6.7.2.post4911
-Requires-Dist: nucliadb-protos>=6.7.2.post4911
-Requires-Dist: nucliadb-models>=6.7.2.post4911
-Requires-Dist: nidx-protos>=6.7.2.post4911
+Requires-Dist: nucliadb-telemetry[all]>=6.7.2.post4917
+Requires-Dist: nucliadb-utils[cache,fastapi,storages]>=6.7.2.post4917
+Requires-Dist: nucliadb-protos>=6.7.2.post4917
+Requires-Dist: nucliadb-models>=6.7.2.post4917
+Requires-Dist: nidx-protos>=6.7.2.post4917
 Requires-Dist: nucliadb-admin-assets>=1.0.0.post1224
-Requires-Dist: nuclia-models>=0.46.0
+Requires-Dist: nuclia-models>=0.47.0
 Requires-Dist: uvicorn[standard]
 Requires-Dist: argdantic
 Requires-Dist: aiohttp>=3.11.11

{nucliadb-6.7.2.post4911.dist-info → nucliadb-6.7.2.post4917.dist-info}/RECORD

@@ -217,7 +217,7 @@ nucliadb/search/__init__.py,sha256=tnypbqcH4nBHbGpkINudhKgdLKpwXQCvDtPchUlsyY4,1
 nucliadb/search/app.py,sha256=-WEX1AZRA8R_9aeOo9ovOTwjXW_7VfwWN7N2ccSoqXg,3387
 nucliadb/search/lifecycle.py,sha256=hiylV-lxsAWkqTCulXBg0EIfMQdejSr8Zar0L_GLFT8,2218
 nucliadb/search/openapi.py,sha256=t3Wo_4baTrfPftg2BHsyLWNZ1MYn7ZRdW7ht-wFOgRs,1016
-nucliadb/search/predict.py,sha256=PdcJz--hK884GTTZ4m_QldR6dfNuI0vHpUWumJ2NGoA,23207
+nucliadb/search/predict.py,sha256=7AtA5K0YBydjUwP0mcMox66CZMI250_QBIOm6XVXxJs,23559
 nucliadb/search/predict_models.py,sha256=ozuQZGWAK7v8W6UOk3xXQ_zW7YUNtRy_4l5LE5BuT_A,8172
 nucliadb/search/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 nucliadb/search/run.py,sha256=aFb-CXRi_C8YMpP_ivNj8KW1BYhADj88y8K9Lr_nUPI,1402
@@ -257,9 +257,9 @@ nucliadb/search/search/graph_strategy.py,sha256=VNa-XLQnJometUTbDJumhHf_LUHphHBo
 nucliadb/search/search/hydrator.py,sha256=CL8XYq5QMSL0hvOND7XnHcJrM-BeVRvpwN5lL55yAD0,6963
 nucliadb/search/search/ingestion_agents.py,sha256=IK6yOPEF9rST_uoqspdVdPk0pldjDhc8Fz1eRovnySI,3311
 nucliadb/search/search/merge.py,sha256=XiRBsxhYPshPV7lZXD-9E259KZOPIf4I2tKosY0lPo4,22470
-nucliadb/search/search/metrics.py,sha256=3I6IN0qDSmqIvUaWJmT3rt-Jyjs6LcvnKI8ZqCiuJPY,3501
+nucliadb/search/search/metrics.py,sha256=yodhoyn-smFdS7rKUn_XXNNXT93WqVdZj9F3TCbbQTI,4160
 nucliadb/search/search/paragraphs.py,sha256=pNAEiYqJGGUVcEf7xf-PFMVqz0PX4Qb-WNG-_zPGN2o,7799
-nucliadb/search/search/predict_proxy.py,sha256=Df8F5K-oS4TIXJc_y8UDViJTo7st5L0kMgxYPFZ39Vk,8806
+nucliadb/search/search/predict_proxy.py,sha256=bU18TKE7GzSt8Hqr38jE-2Lz_GPJRTFTGsFvMdPZWIM,9872
 nucliadb/search/search/query.py,sha256=lYCesbUv-B7IyVFQoCCurcxl_Azc5nq3jtVQJ9tk1Ao,11552
 nucliadb/search/search/rank_fusion.py,sha256=xZtXhbmKb_56gs73u6KkFm2efvTATOSMmpOV2wrAIqE,9613
 nucliadb/search/search/rerankers.py,sha256=2LNC0I28EXriffMuBlOYzjQq0vCTjpCxaK29f852n3s,7473
@@ -267,11 +267,11 @@ nucliadb/search/search/shards.py,sha256=mc5DK-MoCv9AFhlXlOFHbPvetcyNDzTFOJ5rimK8
 nucliadb/search/search/summarize.py,sha256=3lLdwsM28W505bKvmK7JLXmz7kcjd8Hp70LQs391ofY,5087
 nucliadb/search/search/utils.py,sha256=ajRIXfdTF67dBVahQCXW-rSv6gJpUMPt3QhJrWqArTQ,2175
 nucliadb/search/search/chat/__init__.py,sha256=cp15ZcFnHvpcu_5-aK2A4uUyvuZVV_MJn4bIXMa20ks,835
-nucliadb/search/search/chat/ask.py,sha256=c8gvGkspW5KUbDhVzcUfd7eF7SMfqzwKfCPidJniRLI,39417
+nucliadb/search/search/chat/ask.py,sha256=Akq5nCl-0P9Hcil2hzuOEIlhEA9IxJtlQo4dz54BG3s,40884
 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=Bk69WVki5XIzXFYO6o1uQw1feHtuMwfEx_A1TT3piR0,54136
-nucliadb/search/search/chat/query.py,sha256=AhOPMf68p2BRjKz7CdkcUIDMANtxr00oGt42iKUUjAw,16698
+nucliadb/search/search/chat/query.py,sha256=qWrwVEX_GrDV7LFRyC21BURtl-WsRt8BkIocY8njkKM,17147
 nucliadb/search/search/query_parser/__init__.py,sha256=cp15ZcFnHvpcu_5-aK2A4uUyvuZVV_MJn4bIXMa20ks,835
 nucliadb/search/search/query_parser/exceptions.py,sha256=sVl9gRNzhE-s480LBBVkiXzNRbKhYRQN5F3it5tNNp8,939
 nucliadb/search/search/query_parser/fetcher.py,sha256=0Eg_7x9BaAQ1AuTK6NXQMUoGFAXIZiMRurR32tydeNM,17198
@@ -379,8 +379,8 @@ nucliadb/writer/tus/local.py,sha256=7jYa_w9b-N90jWgN2sQKkNcomqn6JMVBOVeDOVYJHto,
 nucliadb/writer/tus/s3.py,sha256=vu1BGg4VqJ_x2P1u2BxqPKlSfw5orT_a3R-Ln5oPUpU,8483
 nucliadb/writer/tus/storage.py,sha256=ToqwjoYnjI4oIcwzkhha_MPxi-k4Jk3Lt55zRwaC1SM,2903
 nucliadb/writer/tus/utils.py,sha256=MSdVbRsRSZVdkaum69_0wku7X3p5wlZf4nr6E0GMKbw,2556
-nucliadb-6.7.2.post4911.dist-info/METADATA,sha256=sWlsMPT1JH0UwpkUlpsNhj5RYdZyyd77MzrsmFh7fcc,4158
-nucliadb-6.7.2.post4911.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-nucliadb-6.7.2.post4911.dist-info/entry_points.txt,sha256=XqGfgFDuY3zXQc8ewXM2TRVjTModIq851zOsgrmaXx4,1268
-nucliadb-6.7.2.post4911.dist-info/top_level.txt,sha256=hwYhTVnX7jkQ9gJCkVrbqEG1M4lT2F_iPQND1fCzF80,20
-nucliadb-6.7.2.post4911.dist-info/RECORD,,
+nucliadb-6.7.2.post4917.dist-info/METADATA,sha256=J1t2FzIphJvmbhKEK2lJG0Xy-MC85goxjoTLD4lod4Q,4158
+nucliadb-6.7.2.post4917.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+nucliadb-6.7.2.post4917.dist-info/entry_points.txt,sha256=XqGfgFDuY3zXQc8ewXM2TRVjTModIq851zOsgrmaXx4,1268
+nucliadb-6.7.2.post4917.dist-info/top_level.txt,sha256=hwYhTVnX7jkQ9gJCkVrbqEG1M4lT2F_iPQND1fCzF80,20
+nucliadb-6.7.2.post4917.dist-info/RECORD,,