nucliadb 6.5.0.post4413__py3-none-any.whl → 6.5.0.post4415__py3-none-any.whl

nucliadb/search/api/v1/predict_proxy.py

@@ -20,7 +20,7 @@
 import json
 from typing import Union
 
-from fastapi import Request
+from fastapi import Header, Request
 from fastapi.responses import Response, StreamingResponse
 from fastapi_versioning import version
 
@@ -30,10 +30,17 @@ from nucliadb.search import predict
 from nucliadb.search.api.v1.router import KB_PREFIX, api
 from nucliadb.search.search.predict_proxy import PredictProxiedEndpoints, predict_proxy
 from nucliadb_models.resource import NucliaDBRoles
+from nucliadb_models.search import NucliaDBClientType
 from nucliadb_utils.authentication import requires
 from nucliadb_utils.exceptions import LimitsExceededError
 
-DESCRIPTION = "Convenience endpoint that proxies requests to the Predict API. It adds the Knowledge Box configuration settings as headers to the predict API request. Refer to the Predict API documentation for more details about the request and response models: https://docs.nuclia.dev/docs/nua-api#tag/Predict"  # noqa: E501
+DESCRIPTION = (
+    "Convenience endpoint that proxies requests to the Predict API."
+    " It adds the Knowledge Box configuration settings as headers to"
+    " the predict API request. Refer to the Predict API documentation"
+    " for more details about the request and response models:"
+    " https://docs.nuclia.dev/docs/nua-api#tag/Predict"
+)
 
 
 @api.get(
@@ -58,20 +65,28 @@ async def predict_proxy_endpoint(
     request: Request,
     kbid: str,
     endpoint: PredictProxiedEndpoints,
+    x_nucliadb_user: str = Header(""),
+    x_ndb_client: NucliaDBClientType = Header(NucliaDBClientType.API),
+    x_forwarded_for: str = Header(""),
 ) -> Union[Response, StreamingResponse, HTTPClientError]:
     try:
         payload = await request.json()
     except json.JSONDecodeError:
         payload = None
     try:
-        return await predict_proxy(
+        response = await predict_proxy(
             kbid,
             endpoint,
             request.method,
             params=request.query_params,
             json=payload,
             headers=dict(request.headers),
+            user_id=x_nucliadb_user,
+            client_type=x_ndb_client,
+            origin=x_forwarded_for,
         )
+
+        return response
     except KnowledgeBoxNotFound:
         return HTTPClientError(status_code=404, detail="Knowledge box not found")
     except LimitsExceededError as exc:

nucliadb/search/search/predict_proxy.py

@@ -17,15 +17,34 @@
 # You should have received a copy of the GNU Affero General Public License
 # along with this program. If not, see <http://www.gnu.org/licenses/>.
 #
+import json
 from enum import Enum
 from typing import Any, Optional, Union
 
+import aiohttp
 from fastapi.datastructures import QueryParams
 from fastapi.responses import Response, StreamingResponse
+from multidict import CIMultiDictProxy
+from nuclia_models.predict.generative_responses import (
+    GenerativeChunk,
+    JSONGenerativeResponse,
+    StatusGenerativeResponse,
+    TextGenerativeResponse,
+)
+from pydantic import ValidationError
 
 from nucliadb.common import datamanagers
-from nucliadb.search.predict import PredictEngine
+from nucliadb.search import logger
+from nucliadb.search.predict import (
+    NUCLIA_LEARNING_ID_HEADER,
+    NUCLIA_LEARNING_MODEL_HEADER,
+    AnswerStatusCode,
+    PredictEngine,
+)
+from nucliadb.search.search.chat.query import maybe_audit_chat
+from nucliadb.search.search.metrics import AskMetrics
 from nucliadb.search.utilities import get_predict
+from nucliadb_models.search import NucliaDBClientType
 
 
 class PredictProxiedEndpoints(str, Enum):
@@ -46,12 +65,17 @@ ALLOWED_HEADERS = [
     "Accept",  # To allow 'application/x-ndjson' on the /chat endpoint
 ]
 
+PREDICT_ANSWER_METRIC = "predict_answer_proxy_metric"
+
 
 async def predict_proxy(
     kbid: str,
     endpoint: PredictProxiedEndpoints,
     method: str,
     params: QueryParams,
+    user_id: str,
+    client_type: NucliaDBClientType,
+    origin: str,
     json: Optional[Any] = None,
     headers: dict[str, str] = {},
 ) -> Union[Response, StreamingResponse]:
@@ -71,22 +95,61 @@ async def predict_proxy(
         headers={**user_headers, **predict_headers},
     )
 
-    # Proxy the response back to the client
     status_code = predict_response.status
     media_type = predict_response.headers.get("Content-Type")
     response: Union[Response, StreamingResponse]
+    user_query = json.get("question") if json is not None else ""
     if predict_response.headers.get("Transfer-Encoding") == "chunked":
+        if endpoint == PredictProxiedEndpoints.CHAT:
+            streaming_generator = chat_streaming_generator(
+                predict_response=predict_response,
+                kbid=kbid,
+                user_id=user_id,
+                client_type=client_type,
+                origin=origin,
+                user_query=user_query,
+                is_json="json" in (media_type or ""),
+            )
+        else:
+            streaming_generator = predict_response.content.iter_any()
+
         response = StreamingResponse(
-            content=predict_response.content.iter_any(),
+            content=streaming_generator,
             status_code=status_code,
             media_type=media_type,
         )
     else:
+        metrics = AskMetrics()
+        with metrics.time(PREDICT_ANSWER_METRIC):
+            content = await predict_response.read()
+
+        if endpoint == PredictProxiedEndpoints.CHAT:
+            try:
+                llm_status_code = int(content[-1:].decode())  # Decode just the last char
+                if llm_status_code != 0:
+                    llm_status_code = -llm_status_code
+            except ValueError:
+                llm_status_code = -1
+
+            audit_predict_proxy_endpoint(
+                predict_response.headers,
+                kbid=kbid,
+                user_id=user_id,
+                user_query=user_query,
+                client_type=client_type,
+                origin=origin,
+                text_answer=content,
+                generative_answer_time=metrics[PREDICT_ANSWER_METRIC],
+                generative_answer_first_chunk_time=None,
+                status_code=AnswerStatusCode(str(llm_status_code)),
+            )
+
         response = Response(
-            content=await predict_response.read(),
+            content=content,
             status_code=status_code,
             media_type=media_type,
         )
+
     nuclia_learning_id = predict_response.headers.get("NUCLIA-LEARNING-ID")
     if nuclia_learning_id:
         response.headers["NUCLIA-LEARNING-ID"] = nuclia_learning_id
@@ -97,3 +160,94 @@ async def predict_proxy(
 async def exists_kb(kbid: str) -> bool:
     async with datamanagers.with_ro_transaction() as txn:
         return await datamanagers.kb.exists_kb(txn, kbid=kbid)
+
+
+async def chat_streaming_generator(
+    predict_response: aiohttp.ClientResponse,
+    kbid: str,
+    user_id: str,
+    client_type: NucliaDBClientType,
+    origin: str,
+    user_query: str,
+    is_json: bool,
+):
+    stream = predict_response.content.iter_any()
+    first = True
+    status_code = AnswerStatusCode.ERROR.value
+    text_answer = ""
+    json_object = None
+    metrics = AskMetrics()
+    with metrics.time(PREDICT_ANSWER_METRIC):
+        async for chunk in stream:
+            if first:
+                metrics.record_first_chunk_yielded()
+                first = False
+
+            yield chunk
+
+            if is_json:
+                try:
+                    parsed_chunk = GenerativeChunk.model_validate(chunk)
+                    if isinstance(parsed_chunk, TextGenerativeResponse):
+                        text_answer += parsed_chunk.text
+                    elif isinstance(parsed_chunk, JSONGenerativeResponse):
+                        json_object = parsed_chunk.object
+                    elif isinstance(parsed_chunk, StatusGenerativeResponse):
+                        status_code = parsed_chunk.code
+                except ValidationError:
+                    logger.warning(
+                        f"Unexpected item in predict answer stream: {chunk.decode()}",
+                        extra={"kbid": kbid},
+                    )
+            else:
+                text_answer += chunk.decode()
+
+    if is_json is False and chunk:  # Ensure chunk is not empty before decoding
+        # If response is text the status_code comes at the last chunk of data
+        status_code = chunk.decode()
+
+    audit_predict_proxy_endpoint(
+        headers=predict_response.headers,
+        kbid=kbid,
+        user_id=user_id,
+        user_query=user_query,
+        client_type=client_type,
+        origin=origin,
+        text_answer=text_answer.encode() if json_object is None else json.dumps(json_object).encode(),
+        generative_answer_time=metrics[PREDICT_ANSWER_METRIC],
+        generative_answer_first_chunk_time=metrics.get_first_chunk_time(),
+        status_code=AnswerStatusCode(status_code),
+    )
+
+
+def audit_predict_proxy_endpoint(
+    headers: CIMultiDictProxy,
+    kbid: str,
+    user_id: str,
+    user_query: str,
+    client_type: NucliaDBClientType,
+    origin: str,
+    text_answer: bytes,
+    generative_answer_time: float,
+    generative_answer_first_chunk_time: Optional[float],
+    status_code: AnswerStatusCode,
+):
+    maybe_audit_chat(
+        kbid=kbid,
+        user_id=user_id,
+        client_type=client_type,
+        origin=origin,
+        user_query=user_query,
+        rephrased_query=None,
+        retrieval_rephrase_query=None,
+        chat_history=[],
+        learning_id=headers.get(NUCLIA_LEARNING_ID_HEADER),
+        query_context={},
+        query_context_order={},
+        model=headers.get(NUCLIA_LEARNING_MODEL_HEADER),
+        text_answer=text_answer,
+        generative_answer_time=generative_answer_time,
+        generative_answer_first_chunk_time=generative_answer_first_chunk_time or 0,
+        rephrase_time=None,
+        status_code=status_code,
+    )

{nucliadb-6.5.0.post4413.dist-info → nucliadb-6.5.0.post4415.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: nucliadb
-Version: 6.5.0.post4413
+Version: 6.5.0.post4415
 Summary: NucliaDB
 Author-email: Nuclia <nucliadb@nuclia.com>
 License-Expression: AGPL-3.0-or-later
@@ -19,11 +19,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.5.0.post4413
-Requires-Dist: nucliadb-utils[cache,fastapi,storages]>=6.5.0.post4413
-Requires-Dist: nucliadb-protos>=6.5.0.post4413
-Requires-Dist: nucliadb-models>=6.5.0.post4413
-Requires-Dist: nidx-protos>=6.5.0.post4413
+Requires-Dist: nucliadb-telemetry[all]>=6.5.0.post4415
+Requires-Dist: nucliadb-utils[cache,fastapi,storages]>=6.5.0.post4415
+Requires-Dist: nucliadb-protos>=6.5.0.post4415
+Requires-Dist: nucliadb-models>=6.5.0.post4415
+Requires-Dist: nidx-protos>=6.5.0.post4415
 Requires-Dist: nucliadb-admin-assets>=1.0.0.post1224
 Requires-Dist: nuclia-models>=0.24.2
 Requires-Dist: uvicorn[standard]

{nucliadb-6.5.0.post4413.dist-info → nucliadb-6.5.0.post4415.dist-info}/RECORD

@@ -221,7 +221,7 @@ nucliadb/search/api/v1/feedback.py,sha256=kNLc4dHz2SXHzV0PwC1WiRAwY88fDptPcP-kO0
 nucliadb/search/api/v1/find.py,sha256=iMjyq4y0JOMC_x1B8kUfVdkCoc9G9Ark58kPLLY4HDw,10824
 nucliadb/search/api/v1/graph.py,sha256=gthqxCOn9biE6D6s93jRGLglk0ono8U7OyS390kWiI8,4178
 nucliadb/search/api/v1/knowledgebox.py,sha256=e9xeLPUqnQTx33i4A8xuV93ENvtJGrpjPlLRbGJtAI8,8415
-nucliadb/search/api/v1/predict_proxy.py,sha256=Q03ZTvWp7Sq0x71t5Br4LHxTiYsRd6-GCb4YuKqhynM,3131
+nucliadb/search/api/v1/predict_proxy.py,sha256=TnXKAqf_Go-9QVi6L5z4cXjnuNRe7XLJjF5QH_uwA1I,3504
 nucliadb/search/api/v1/router.py,sha256=mtT07rBZcVfpa49doaw9b1tj3sdi3qLH0gn9Io6NYM0,988
 nucliadb/search/api/v1/search.py,sha256=eqlrvRE7IlMpunNwD1RJwt6RgMV01sIDJLgxxE7CFcE,12297
 nucliadb/search/api/v1/suggest.py,sha256=gaJE60r8-z6TVO05mQRKBITwXn2_ofM3B4-OtpOgZEk,6343
@@ -250,7 +250,7 @@ nucliadb/search/search/merge.py,sha256=XiRBsxhYPshPV7lZXD-9E259KZOPIf4I2tKosY0lP
 nucliadb/search/search/metrics.py,sha256=3I6IN0qDSmqIvUaWJmT3rt-Jyjs6LcvnKI8ZqCiuJPY,3501
 nucliadb/search/search/paragraphs.py,sha256=pNAEiYqJGGUVcEf7xf-PFMVqz0PX4Qb-WNG-_zPGN2o,7799
 nucliadb/search/search/pgcatalog.py,sha256=QtgArjoM-dW_B1oO0aXqp5au7GlLG8jAct9jevUHatw,10997
-nucliadb/search/search/predict_proxy.py,sha256=JwgBeEg1j4LnCjPCvTUrnmOd9LceJAt3iAu4m9cmJBo,3390
+nucliadb/search/search/predict_proxy.py,sha256=2451FR2DJLedA4eV8KbMbY4la5WZCgjBUHO_iQVaA8I,8634
 nucliadb/search/search/query.py,sha256=0qIQdt548L3jtKOyKo06aGJ73SLBxAW3N38_Hc1M3Uw,11528
 nucliadb/search/search/rank_fusion.py,sha256=xZtXhbmKb_56gs73u6KkFm2efvTATOSMmpOV2wrAIqE,9613
 nucliadb/search/search/rerankers.py,sha256=E2J1QdKAojqbhHM3KAyaOXKf6tJyETUxKs4tf_BEyqk,7472
@@ -370,8 +370,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.5.0.post4413.dist-info/METADATA,sha256=B1pDmizFK3zLtNJQUCZAOb4jLeQgnyw3vcryDfJO5y4,4152
-nucliadb-6.5.0.post4413.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-nucliadb-6.5.0.post4413.dist-info/entry_points.txt,sha256=XqGfgFDuY3zXQc8ewXM2TRVjTModIq851zOsgrmaXx4,1268
-nucliadb-6.5.0.post4413.dist-info/top_level.txt,sha256=hwYhTVnX7jkQ9gJCkVrbqEG1M4lT2F_iPQND1fCzF80,20
-nucliadb-6.5.0.post4413.dist-info/RECORD,,
+nucliadb-6.5.0.post4415.dist-info/METADATA,sha256=lSbgbK7GgTJXjm5IIzYAeFt7-W8LdME3CKarZjmQwbg,4152
+nucliadb-6.5.0.post4415.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+nucliadb-6.5.0.post4415.dist-info/entry_points.txt,sha256=XqGfgFDuY3zXQc8ewXM2TRVjTModIq851zOsgrmaXx4,1268
+nucliadb-6.5.0.post4415.dist-info/top_level.txt,sha256=hwYhTVnX7jkQ9gJCkVrbqEG1M4lT2F_iPQND1fCzF80,20
+nucliadb-6.5.0.post4415.dist-info/RECORD,,