nucliadb 6.4.0.post4265__py3-none-any.whl → 6.4.0.post4271__py3-none-any.whl

nucliadb/search/predict.py

@@ -22,6 +22,7 @@ import json
 import logging
 import os
 import random
+from dataclasses import dataclass
 from enum import Enum
 from typing import Any, AsyncGenerator, Optional
 from unittest.mock import AsyncMock, Mock
@@ -108,7 +109,7 @@ RERANK = "/rerank"
 
 NUCLIA_LEARNING_ID_HEADER = "NUCLIA-LEARNING-ID"
 NUCLIA_LEARNING_MODEL_HEADER = "NUCLIA-LEARNING-MODEL"
-
+NUCLIA_LEARNING_CHAT_HISTORY_HEADER = "NUCLIA-LEARNING-CHAT-HISTORY"
 
 predict_observer = metrics.Observer(
     "predict_engine",
@@ -139,6 +140,12 @@ class AnswerStatusCode(str, Enum):
         }[self]
 
 
+@dataclass
+class RephraseResponse:
+    rephrased_query: str
+    use_chat_history: Optional[bool]
+
+
 async def start_predict_engine():
     if nuclia_settings.dummy_predict:
         predict_util = DummyPredictEngine()
@@ -267,7 +274,7 @@ class PredictEngine:
         return await func(**request_args)
 
     @predict_observer.wrap({"type": "rephrase"})
-    async def rephrase_query(self, kbid: str, item: RephraseModel) -> str:
+    async def rephrase_query(self, kbid: str, item: RephraseModel) -> RephraseResponse:
         try:
             self.check_nua_key_is_configured_for_onprem()
         except NUAKeyMissingError:
@@ -477,9 +484,9 @@ class DummyPredictEngine(PredictEngine):
         response.headers = {NUCLIA_LEARNING_ID_HEADER: DUMMY_LEARNING_ID}
         return response
 
-    async def rephrase_query(self, kbid: str, item: RephraseModel) -> str:
+    async def rephrase_query(self, kbid: str, item: RephraseModel) -> RephraseResponse:
         self.calls.append(("rephrase_query", item))
-        return DUMMY_REPHRASE_QUERY
+        return RephraseResponse(rephrased_query=DUMMY_REPHRASE_QUERY, use_chat_history=None)
 
     async def chat_query_ndjson(
         self, kbid: str, item: ChatModel
@@ -624,7 +631,7 @@ def get_chat_ndjson_generator(
 
 async def _parse_rephrase_response(
     resp: aiohttp.ClientResponse,
-) -> str:
+) -> RephraseResponse:
     """
     Predict api is returning a json payload that is a string with the following format:
     <rephrased_query><status_code>
@@ -632,12 +639,15 @@ async def _parse_rephrase_response(
     it will raise an exception if the status code is not 0
     """
     content = await resp.json()
+
     if content.endswith("0"):
-        return content[:-1]
+        content = content[:-1]
     elif content.endswith("-1"):
         raise RephraseError(content[:-2])
     elif content.endswith("-2"):
         raise RephraseMissingContextError(content[:-2])
-    else:
-        # bw compatibility
-        return content
+
+    use_chat_history = None
+    if NUCLIA_LEARNING_CHAT_HISTORY_HEADER in resp.headers:
+        use_chat_history = resp.headers[NUCLIA_LEARNING_CHAT_HISTORY_HEADER] == "true"
+    return RephraseResponse(rephrased_query=content, use_chat_history=use_chat_history)

nucliadb/search/search/chat/ask.py

@@ -488,14 +488,21 @@ async def ask(
     if len(chat_history) > 0 or len(user_context) > 0:
         try:
             with metrics.time("rephrase"):
-                rephrased_query = await rephrase_query(
+                rephrase_response = await rephrase_query(
                     kbid,
                     chat_history=chat_history,
                     query=user_query,
                     user_id=user_id,
                     user_context=user_context,
                     generative_model=ask_request.generative_model,
+                    chat_history_relevance_threshold=ask_request.chat_history_relevance_threshold,
                 )
+                rephrased_query = rephrase_response.rephrased_query
+                if rephrase_response.use_chat_history is False:
+                    # Ignored if the question is not relevant enough with the chat history
+                    logger.info("Chat history was ignored for this request")
+                    chat_history = []
+
         except RephraseMissingContextError:
             logger.info("Failed to rephrase ask query, using original")
 

nucliadb/search/search/chat/query.py

@@ -27,7 +27,7 @@ from nidx_protos.nodereader_pb2 import (
 
 from nucliadb.common.models_utils import to_proto
 from nucliadb.search import logger
-from nucliadb.search.predict import AnswerStatusCode
+from nucliadb.search.predict import AnswerStatusCode, RephraseResponse
 from nucliadb.search.requesters.utils import Method, node_query
 from nucliadb.search.search.chat.exceptions import NoRetrievalResultsError
 from nucliadb.search.search.exceptions import IncompleteFindResultsError
@@ -71,7 +71,8 @@ async def rephrase_query(
     user_id: str,
     user_context: list[str],
     generative_model: Optional[str] = None,
-) -> str:
+    chat_history_relevance_threshold: Optional[float] = None,
+) -> RephraseResponse:
     predict = get_predict()
     req = RephraseModel(
         question=query,
@@ -79,6 +80,7 @@ async def rephrase_query(
         user_id=user_id,
         user_context=user_context,
         generative_model=generative_model,
+        chat_history_relevance_threshold=chat_history_relevance_threshold,
     )
     return await predict.rephrase_query(kbid, req)
 

{nucliadb-6.4.0.post4265.dist-info → nucliadb-6.4.0.post4271.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: nucliadb
-Version: 6.4.0.post4265
+Version: 6.4.0.post4271
 Summary: NucliaDB
 Author-email: Nuclia <nucliadb@nuclia.com>
 License: AGPL
@@ -20,11 +20,11 @@ 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.4.0.post4265
-Requires-Dist: nucliadb-utils[cache,fastapi,storages]>=6.4.0.post4265
-Requires-Dist: nucliadb-protos>=6.4.0.post4265
-Requires-Dist: nucliadb-models>=6.4.0.post4265
-Requires-Dist: nidx-protos>=6.4.0.post4265
+Requires-Dist: nucliadb-telemetry[all]>=6.4.0.post4271
+Requires-Dist: nucliadb-utils[cache,fastapi,storages]>=6.4.0.post4271
+Requires-Dist: nucliadb-protos>=6.4.0.post4271
+Requires-Dist: nucliadb-models>=6.4.0.post4271
+Requires-Dist: nidx-protos>=6.4.0.post4271
 Requires-Dist: nucliadb-admin-assets>=1.0.0.post1224
 Requires-Dist: nuclia-models>=0.24.2
 Requires-Dist: uvicorn[standard]

{nucliadb-6.4.0.post4265.dist-info → nucliadb-6.4.0.post4271.dist-info}/RECORD

@@ -204,7 +204,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=BYkKL2-3-MNT8JnE7y7XTEMKMnynUm2y4VJZP1jRjdQ,22987
+nucliadb/search/predict.py,sha256=__0qwIU2CIRYRTYsbG9zZEjXXrxNe8puZWYJIyOT6dg,23492
 nucliadb/search/predict_models.py,sha256=ZAe0dneUsPmV9uBar57cCFADCGOrYDsJHuqKlA5zWag,5937
 nucliadb/search/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 nucliadb/search/run.py,sha256=aFb-CXRi_C8YMpP_ivNj8KW1BYhADj88y8K9Lr_nUPI,1402
@@ -255,11 +255,11 @@ nucliadb/search/search/shards.py,sha256=mc5DK-MoCv9AFhlXlOFHbPvetcyNDzTFOJ5rimK8
 nucliadb/search/search/summarize.py,sha256=ksmYPubEQvAQgfPdZHfzB_rR19B2ci4IYZ6jLdHxZo8,4996
 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=tkU431_tZpFQ8Au9RpGrHO78D0vCKE-VaBT5BNTL0pA,37393
+nucliadb/search/search/chat/ask.py,sha256=4cpeWC7Q4NI30vrOeq22N0Vw0PQS0Ko9b47R5ISlgUw,37833
 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=Jnja-Ss7skgnnDY8BymVfdeYsFPnIQFL8tEvcRXTKUE,47356
-nucliadb/search/search/chat/query.py,sha256=yK7J8vwJT7gKeBPn21uYqQw0tEJJ_HRtGLSt2bnJEuQ,16895
+nucliadb/search/search/chat/query.py,sha256=6v6twBUTWfUUzklVV6xqJSYPkAshnIrBH9wbTcjQvkI,17063
 nucliadb/search/search/query_parser/__init__.py,sha256=cp15ZcFnHvpcu_5-aK2A4uUyvuZVV_MJn4bIXMa20ks,835
 nucliadb/search/search/query_parser/exceptions.py,sha256=szAOXUZ27oNY-OSa9t2hQ5HHkQQC0EX1FZz_LluJHJE,1224
 nucliadb/search/search/query_parser/fetcher.py,sha256=SkvBRDfSKmuz-QygNKLAU4AhZhhDo1dnOZmt1zA28RA,16851
@@ -368,8 +368,8 @@ 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.4.0.post4265.dist-info/METADATA,sha256=DdUvK3SLFsfNpySAdkMKFq9721wSRZ-R0n7fFHGEpxQ,4223
-nucliadb-6.4.0.post4265.dist-info/WHEEL,sha256=DnLRTWE75wApRYVsjgc6wsVswC54sMSJhAEd4xhDpBk,91
-nucliadb-6.4.0.post4265.dist-info/entry_points.txt,sha256=XqGfgFDuY3zXQc8ewXM2TRVjTModIq851zOsgrmaXx4,1268
-nucliadb-6.4.0.post4265.dist-info/top_level.txt,sha256=hwYhTVnX7jkQ9gJCkVrbqEG1M4lT2F_iPQND1fCzF80,20
-nucliadb-6.4.0.post4265.dist-info/RECORD,,
+nucliadb-6.4.0.post4271.dist-info/METADATA,sha256=DKjN0N70XgUrzoOshpgsNU9xMWIweEdjtQgziYbveU0,4223
+nucliadb-6.4.0.post4271.dist-info/WHEEL,sha256=DnLRTWE75wApRYVsjgc6wsVswC54sMSJhAEd4xhDpBk,91
+nucliadb-6.4.0.post4271.dist-info/entry_points.txt,sha256=XqGfgFDuY3zXQc8ewXM2TRVjTModIq851zOsgrmaXx4,1268
+nucliadb-6.4.0.post4271.dist-info/top_level.txt,sha256=hwYhTVnX7jkQ9gJCkVrbqEG1M4lT2F_iPQND1fCzF80,20
+nucliadb-6.4.0.post4271.dist-info/RECORD,,