Skip to content

Commit c72ca6f

Browse files
leoguillaumeleoguillaume
andauthored
hotfix: tei params (#309)
Co-authored-by: leoguillaume <leo.guillaume@modernisation.gouv.fr>
1 parent c4f0e80 commit c72ca6f

1 file changed

Lines changed: 8 additions & 2 deletions

File tree

app/clients/model/_teimodelclient.py

Lines changed: 8 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -87,7 +87,13 @@ def _format_request(self, json: Optional[dict] = None, files: Optional[dict] = N
8787

8888
return url, headers, json, files, data
8989

90-
def _format_response(self, json: dict, response: httpx.Response, additional_data: Dict[str, Any] = {}) -> httpx.Response:
90+
def _format_response(
91+
self,
92+
json: dict,
93+
response: httpx.Response,
94+
additional_data: Dict[str, Any] = {},
95+
request_latency: float = 0.0,
96+
) -> httpx.Response:
9197
"""
9298
Format a response from a client model and add usage data and model ID to the response. This method can be overridden by a subclass to add additional headers or parameters.
9399
@@ -105,7 +111,7 @@ def _format_response(self, json: dict, response: httpx.Response, additional_data
105111
data = response.json()
106112
if isinstance(data, list): # for TEI reranking
107113
data = {"data": data}
108-
data.update(self._get_additional_data(json=json, data=data, stream=False))
114+
data.update(self._get_additional_data(json=json, data=data, stream=False, request_latency=request_latency))
109115
data.update(additional_data)
110116
response = httpx.Response(status_code=response.status_code, content=dumps(data))
111117

0 commit comments

Comments
 (0)