nucliadb 6.5.0.post4408__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/train/api/utils.py

@@ -24,7 +24,7 @@ from typing import Optional
 from nucliadb.train.utils import get_shard_manager
 
 
-async def get_kb_partitions(kbid: str, prefix: Optional[str] = None):
+async def get_kb_partitions(kbid: str, prefix: Optional[str] = None) -> list[str]:
     shard_manager = get_shard_manager()
     shards = await shard_manager.get_shards_by_kbid_inner(kbid=kbid)
     valid_shards = []

nucliadb/train/api/v1/shards.py

@@ -26,6 +26,7 @@ from fastapi import HTTPException, Request
 from fastapi.responses import StreamingResponse
 from fastapi_versioning import version
 
+from nucliadb.common.cluster.exceptions import ShardNotFound
 from nucliadb.train.api.utils import get_kb_partitions
 from nucliadb.train.api.v1.router import KB_PREFIX, api
 from nucliadb.train.generator import generate_train_data
@@ -49,7 +50,10 @@ async def object_get_response(
     kbid: str,
     shard: str,
 ) -> StreamingResponse:
-    partitions = await get_kb_partitions(kbid, shard)
+    try:
+        partitions = await get_kb_partitions(kbid, prefix=shard)
+    except ShardNotFound:
+        raise HTTPException(status_code=404, detail=f"No shards found for kb")
     if shard not in partitions:
         raise HTTPException(status_code=404, detail=f"Partition {shard} not found")
     trainset, filter_expression = await get_trainset(request)

nucliadb/train/api/v1/trainset.py

@@ -20,9 +20,10 @@
 
 from typing import Optional
 
-from fastapi import Request
+from fastapi import HTTPException, Request
 from fastapi_versioning import version
 
+from nucliadb.common.cluster.exceptions import ShardNotFound
 from nucliadb.train.api.utils import get_kb_partitions
 from nucliadb.train.api.v1.router import KB_PREFIX, api
 from nucliadb_models.resource import NucliaDBRoles
@@ -57,5 +58,8 @@ async def get_partitions_prefix(request: Request, kbid: str, prefix: str) -> Tra
 
 
 async def get_partitions(kbid: str, prefix: Optional[str] = None) -> TrainSetPartitions:
-    all_keys = await get_kb_partitions(kbid, prefix)
+    try:
+        all_keys = await get_kb_partitions(kbid, prefix)
+    except ShardNotFound:
+        raise HTTPException(status_code=404, detail=f"No shards found for kb")
     return TrainSetPartitions(partitions=all_keys)

nucliadb/train/generator.py

@@ -17,12 +17,14 @@
 # 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/>.
 #
-
 from typing import AsyncIterator, Callable, Optional
 
 from fastapi import HTTPException
+from grpc import StatusCode
+from grpc.aio import AioRpcError
 
 from nucliadb.common.cache import resource_cache
+from nucliadb.train import logger
 from nucliadb.train.generators.field_classifier import (
     field_classification_batch_generator,
 )
@@ -93,7 +95,16 @@ async def generate_train_data(
     # This cache size is an arbitrary number, once we have a metric in place and
     # we analyze memory consumption, we can adjust it with more knoweldge
     with resource_cache(size=settings.resource_cache_size):
-        async for item in batch_generator(kbid, trainset, shard_replica_id, filter_expression):
-            payload = item.SerializeToString()
-            yield len(payload).to_bytes(4, byteorder="big", signed=False)
-            yield payload
+        try:
+            async for item in batch_generator(kbid, trainset, shard_replica_id, filter_expression):
+                payload = item.SerializeToString()
+                yield len(payload).to_bytes(4, byteorder="big", signed=False)
+                yield payload
+        except AioRpcError as exc:
+            if exc.code() == StatusCode.NOT_FOUND:
+                logger.warning(
+                    f"Shard not found in nidx. Halting the stream",
+                    extra={"kbid": kbid, "shard": shard, "shard_replica_id": shard_replica_id},
+                )
+                return
+            raise

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: nucliadb
-Version: 6.5.0.post4408
+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.post4408
-Requires-Dist: nucliadb-utils[cache,fastapi,storages]>=6.5.0.post4408
-Requires-Dist: nucliadb-protos>=6.5.0.post4408
-Requires-Dist: nucliadb-models>=6.5.0.post4408
-Requires-Dist: nidx-protos>=6.5.0.post4408
+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.post4408.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
@@ -303,7 +303,7 @@ nucliadb/tests/config.py,sha256=JN_Jhgj-fwM9_8IeO9pwxr6C1PiwRDrXxm67Y38rU30,2080
 nucliadb/tests/vectors.py,sha256=CcNKx-E8LPpyvRyljbmb-Tn_wST9Juw2CBoogWrKiTk,62843
 nucliadb/train/__init__.py,sha256=NVwe5yULoHXb80itIJT8YJYEz2xbiOPQ7_OMys6XJw8,1301
 nucliadb/train/app.py,sha256=z6xlGVVVaJmZZmLPIVTgkjD-wIz5b0NYlXAQp7hBHYw,2652
-nucliadb/train/generator.py,sha256=fwFYal7VsV0EP7J_g3IOJ-WLpjwqrVo0gEP7vxIlxGs,4152
+nucliadb/train/generator.py,sha256=H8JLkQ23QQVo4CAdg1ZZh_cncPz7COEfaDu1l-h-0hM,4616
 nucliadb/train/lifecycle.py,sha256=3HadM4GRsYb2m-v4jtdr9C-KBEBx8GlrJDArPYi3SWQ,1960
 nucliadb/train/models.py,sha256=BmgmMjDsu_1Ih5JDAqo6whhume90q0ASJcDP9dkMQm8,1198
 nucliadb/train/nodes.py,sha256=6QD9ZnAacPyFCzs1jxJzsL3CkVNyADvsfnPjjrvn-NU,5610
@@ -317,11 +317,11 @@ nucliadb/train/upload.py,sha256=fTjH1KEL-0ogf3LV0T6ODO0QdPGwdZShSUtFUCAcUlA,3256
 nucliadb/train/uploader.py,sha256=xdLGz1ToDue9Q_M8A-_KYkO-V6fWKYOZQ6IGM4FuwWA,6424
 nucliadb/train/utils.py,sha256=OxQ8No19nxOxUhwNYDICNt8n2H-gowkAgu2Vt3Hatzk,3163
 nucliadb/train/api/__init__.py,sha256=cp15ZcFnHvpcu_5-aK2A4uUyvuZVV_MJn4bIXMa20ks,835
-nucliadb/train/api/utils.py,sha256=1E5h1dS3dLppXD4k6qH6jmaY0WYa_ChaNHiTW9y336Q,1307
+nucliadb/train/api/utils.py,sha256=7verIqWRzDVZUsPrc2OQMGSKcd_01aG2lheEc0z5rP8,1320
 nucliadb/train/api/v1/__init__.py,sha256=P4vCIv93r_Cq1WFDDNjy_Wg7zBkzx0S4euXwfPy1LA4,928
 nucliadb/train/api/v1/router.py,sha256=ukdxn5q1oMar6NSPobgJczWsSxLCHw6DYKlb3zwCiSo,910
-nucliadb/train/api/v1/shards.py,sha256=olgjQFKeYSSCbpYhglCRq6Q9WsM9hxOZgk6mOBb4g0o,3367
-nucliadb/train/api/v1/trainset.py,sha256=kpnpDgiMWr1FKHZJgwH7hue5kzilA8-i9X0YHlNeHuU,2113
+nucliadb/train/api/v1/shards.py,sha256=0TLDWdrMwLAD3lb0AmdvpJF_gALxK-dw0ns_lfGSOzY,3553
+nucliadb/train/api/v1/trainset.py,sha256=1Iep1Ze0mmRREHsP3mpJ8cC6sU-_avlqh5ItTRQnxP8,2307
 nucliadb/train/generators/__init__.py,sha256=cp15ZcFnHvpcu_5-aK2A4uUyvuZVV_MJn4bIXMa20ks,835
 nucliadb/train/generators/field_classifier.py,sha256=UcA5snqLNjIHw0VBzXo9ZtSua6o7wBU3tV9_d5qWpRA,3542
 nucliadb/train/generators/field_streaming.py,sha256=p0xu39D5gaSQc-LagKwpgsVaxm2ULTkWZDPi-Ad1lHc,8378
@@ -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.post4408.dist-info/METADATA,sha256=5PRi3_KM32DTF2_BSWmfndESWte8q1aH9fPh0SYJ6aE,4152
-nucliadb-6.5.0.post4408.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-nucliadb-6.5.0.post4408.dist-info/entry_points.txt,sha256=XqGfgFDuY3zXQc8ewXM2TRVjTModIq851zOsgrmaXx4,1268
-nucliadb-6.5.0.post4408.dist-info/top_level.txt,sha256=hwYhTVnX7jkQ9gJCkVrbqEG1M4lT2F_iPQND1fCzF80,20
-nucliadb-6.5.0.post4408.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,,