nucliadb 6.4.0.post4227__py3-none-any.whl → 6.4.0.post4241__py3-none-any.whl

nucliadb/search/api/v1/predict_proxy.py

@@ -21,7 +21,7 @@ import json
 from typing import Union
 
 from fastapi import Request
-from fastapi.responses import JSONResponse, StreamingResponse
+from fastapi.responses import Response, StreamingResponse
 from fastapi_versioning import version
 
 from nucliadb.common.datamanagers.exceptions import KnowledgeBoxNotFound
@@ -58,14 +58,19 @@ async def predict_proxy_endpoint(
     request: Request,
     kbid: str,
     endpoint: PredictProxiedEndpoints,
-) -> Union[JSONResponse, StreamingResponse, HTTPClientError]:
+) -> Union[Response, StreamingResponse, HTTPClientError]:
     try:
         payload = await request.json()
     except json.JSONDecodeError:
         payload = None
     try:
         return await predict_proxy(
-            kbid, endpoint, request.method, params=request.query_params, json=payload
+            kbid,
+            endpoint,
+            request.method,
+            params=request.query_params,
+            json=payload,
+            headers=dict(request.headers),
         )
     except KnowledgeBoxNotFound:
         return HTTPClientError(status_code=404, detail="Knowledge box not found")

nucliadb/search/predict.py

@@ -262,7 +262,7 @@ class PredictEngine:
         jitter=backoff.random_jitter,
         max_tries=MAX_TRIES,
     )
-    async def make_request(self, method: str, **request_args):
+    async def make_request(self, method: str, **request_args) -> aiohttp.ClientResponse:
         func = getattr(self.session, method.lower())
         return await func(**request_args)
 
@@ -311,8 +311,8 @@ class PredictEngine:
             timeout=None,
         )
         await self.check_response(kbid, resp, expected_status=200)
-        ident = resp.headers.get(NUCLIA_LEARNING_ID_HEADER)
-        model = resp.headers.get(NUCLIA_LEARNING_MODEL_HEADER)
+        ident = resp.headers.get(NUCLIA_LEARNING_ID_HEADER) or "unknown"
+        model = resp.headers.get(NUCLIA_LEARNING_MODEL_HEADER) or "unknown"
         return ident, model, get_chat_ndjson_generator(resp)
 
     @predict_observer.wrap({"type": "query"})
@@ -471,7 +471,9 @@ class DummyPredictEngine(PredictEngine):
 
     async def make_request(self, method: str, **request_args):
         response = Mock(status=200)
-        response.json = AsyncMock(return_value={"foo": "bar"})
+        json_data = {"foo": "bar"}
+        response.json = AsyncMock(return_value=json_data)
+        response.read = AsyncMock(return_value=json.dumps(json_data).encode("utf-8"))
         response.headers = {NUCLIA_LEARNING_ID_HEADER: DUMMY_LEARNING_ID}
         return response
 

nucliadb/search/search/predict_proxy.py

@@ -21,7 +21,7 @@ from enum import Enum
 from typing import Any, Optional, Union
 
 from fastapi.datastructures import QueryParams
-from fastapi.responses import JSONResponse, StreamingResponse
+from fastapi.responses import Response, StreamingResponse
 
 from nucliadb.common import datamanagers
 from nucliadb.search.predict import PredictEngine
@@ -42,20 +42,25 @@ class PredictProxiedEndpoints(str, Enum):
     REMI = "remi"
 
 
+ALLOWED_HEADERS = [
+    "Accept",  # To allow 'application/x-ndjson' on the /chat endpoint
+]
+
+
 async def predict_proxy(
     kbid: str,
     endpoint: PredictProxiedEndpoints,
     method: str,
     params: QueryParams,
     json: Optional[Any] = None,
-) -> Union[JSONResponse, StreamingResponse]:
+    headers: dict[str, str] = {},
+) -> Union[Response, StreamingResponse]:
     if not await exists_kb(kbid=kbid):
         raise datamanagers.exceptions.KnowledgeBoxNotFound()
 
     predict: PredictEngine = get_predict()
-
-    # Add KB configuration headers
-    headers = predict.get_predict_headers(kbid)
+    predict_headers = predict.get_predict_headers(kbid)
+    user_headers = {k: v for k, v in headers.items() if k.capitalize() in ALLOWED_HEADERS}
 
     # Proxy the request to predict API
     predict_response = await predict.make_request(
@@ -63,22 +68,24 @@ async def predict_proxy(
         url=predict.get_predict_url(endpoint, kbid),
         json=json,
         params=params,
-        headers=headers,
+        headers={**user_headers, **predict_headers},
     )
 
     # Proxy the response back to the client
     status_code = predict_response.status
-    response: Union[JSONResponse, StreamingResponse]
+    media_type = predict_response.headers.get("Content-Type")
+    response: Union[Response, StreamingResponse]
     if predict_response.headers.get("Transfer-Encoding") == "chunked":
         response = StreamingResponse(
             content=predict_response.content.iter_any(),
             status_code=status_code,
-            media_type=predict_response.headers.get("Content-Type"),
+            media_type=media_type,
         )
     else:
-        response = JSONResponse(
-            content=await predict_response.json(),
+        response = Response(
+            content=await predict_response.read(),
             status_code=status_code,
+            media_type=media_type,
         )
     nuclia_learning_id = predict_response.headers.get("NUCLIA-LEARNING-ID")
     if nuclia_learning_id:

{nucliadb-6.4.0.post4227.dist-info → nucliadb-6.4.0.post4241.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: nucliadb
-Version: 6.4.0.post4227
+Version: 6.4.0.post4241
 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.post4227
-Requires-Dist: nucliadb-utils[cache,fastapi,storages]>=6.4.0.post4227
-Requires-Dist: nucliadb-protos>=6.4.0.post4227
-Requires-Dist: nucliadb-models>=6.4.0.post4227
-Requires-Dist: nidx-protos>=6.4.0.post4227
+Requires-Dist: nucliadb-telemetry[all]>=6.4.0.post4241
+Requires-Dist: nucliadb-utils[cache,fastapi,storages]>=6.4.0.post4241
+Requires-Dist: nucliadb-protos>=6.4.0.post4241
+Requires-Dist: nucliadb-models>=6.4.0.post4241
+Requires-Dist: nidx-protos>=6.4.0.post4241
 Requires-Dist: nucliadb-admin-assets>=1.0.0.post1224
 Requires-Dist: nuclia-models>=0.24.2
 Requires-Dist: uvicorn[standard]

{nucliadb-6.4.0.post4227.dist-info → nucliadb-6.4.0.post4241.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=VJr5Itx8FE7CZIGYcP-fRgd2YGxAnP9Qj9NxiwWiwcc,22819
+nucliadb/search/predict.py,sha256=BYkKL2-3-MNT8JnE7y7XTEMKMnynUm2y4VJZP1jRjdQ,22987
 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
@@ -218,7 +218,7 @@ nucliadb/search/api/v1/feedback.py,sha256=kNLc4dHz2SXHzV0PwC1WiRAwY88fDptPcP-kO0
 nucliadb/search/api/v1/find.py,sha256=JFbGDRFBHBTApYR1qHp9RngbE_QDb96fXORMdjcN6lg,10654
 nucliadb/search/api/v1/graph.py,sha256=ItVpzJbqfDLjoIo2fTb2mKGCM1Z34sx7CBb3gNmj6IQ,4274
 nucliadb/search/api/v1/knowledgebox.py,sha256=e9xeLPUqnQTx33i4A8xuV93ENvtJGrpjPlLRbGJtAI8,8415
-nucliadb/search/api/v1/predict_proxy.py,sha256=QrGzo0hKjtmyGZ6pjlJHYAh4hxwVUIOTcVcerRCw7eE,3047
+nucliadb/search/api/v1/predict_proxy.py,sha256=Q03ZTvWp7Sq0x71t5Br4LHxTiYsRd6-GCb4YuKqhynM,3131
 nucliadb/search/api/v1/router.py,sha256=mtT07rBZcVfpa49doaw9b1tj3sdi3qLH0gn9Io6NYM0,988
 nucliadb/search/api/v1/search.py,sha256=Or-mUvmBAyh0Y55NqTYNXe_BWR0lLLaTSL2ChjJaE2M,12402
 nucliadb/search/api/v1/suggest.py,sha256=Em7ApddZNHMHjL_ZfXmUIVUk504f58J96JlxJXnIxaM,6438
@@ -247,7 +247,7 @@ nucliadb/search/search/merge.py,sha256=Abg9YblQJvH2jDvXVT45MNxaIpNa7TTpsiUSJqb3N
 nucliadb/search/search/metrics.py,sha256=HJVQPLOIwLuc733G4keqEgx1-Dcg97hyWGKZQTojiSE,2973
 nucliadb/search/search/paragraphs.py,sha256=pNAEiYqJGGUVcEf7xf-PFMVqz0PX4Qb-WNG-_zPGN2o,7799
 nucliadb/search/search/pgcatalog.py,sha256=s_J98fsX_RuFXwpejpkGqG-tD9ELuzz4YQ6U3ew5h2g,9313
-nucliadb/search/search/predict_proxy.py,sha256=IFI3v_ODz2_UU1XZnyaD391fE7-2C0npSmj_HmDvzS4,3123
+nucliadb/search/search/predict_proxy.py,sha256=JwgBeEg1j4LnCjPCvTUrnmOd9LceJAt3iAu4m9cmJBo,3390
 nucliadb/search/search/query.py,sha256=-gvKsyGmKYpsoEVzKkq3HJUMcs_3LD3TYUueOcJsTec,11511
 nucliadb/search/search/rank_fusion.py,sha256=xZtXhbmKb_56gs73u6KkFm2efvTATOSMmpOV2wrAIqE,9613
 nucliadb/search/search/rerankers.py,sha256=PvhExUb8zZYghiFHRgGotw6h6bU--Rft09wE8arvtAw,7424
@@ -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.post4227.dist-info/METADATA,sha256=eZVBfu03hI3bpE7r2yLzoiPcUtqH82BsKHO51-dFfN0,4223
-nucliadb-6.4.0.post4227.dist-info/WHEEL,sha256=0CuiUZ_p9E4cD6NyLD6UG80LBXYyiSYZOKDm5lp32xk,91
-nucliadb-6.4.0.post4227.dist-info/entry_points.txt,sha256=XqGfgFDuY3zXQc8ewXM2TRVjTModIq851zOsgrmaXx4,1268
-nucliadb-6.4.0.post4227.dist-info/top_level.txt,sha256=hwYhTVnX7jkQ9gJCkVrbqEG1M4lT2F_iPQND1fCzF80,20
-nucliadb-6.4.0.post4227.dist-info/RECORD,,
+nucliadb-6.4.0.post4241.dist-info/METADATA,sha256=twJEtMMEkGzMqaXb2hvjgRRgPmvl4ESDTJuUw3ObLuk,4223
+nucliadb-6.4.0.post4241.dist-info/WHEEL,sha256=0CuiUZ_p9E4cD6NyLD6UG80LBXYyiSYZOKDm5lp32xk,91
+nucliadb-6.4.0.post4241.dist-info/entry_points.txt,sha256=XqGfgFDuY3zXQc8ewXM2TRVjTModIq851zOsgrmaXx4,1268
+nucliadb-6.4.0.post4241.dist-info/top_level.txt,sha256=hwYhTVnX7jkQ9gJCkVrbqEG1M4lT2F_iPQND1fCzF80,20
+nucliadb-6.4.0.post4241.dist-info/RECORD,,