From 229709304aa49a241accc44692bc2cebafd57fb2 Mon Sep 17 00:00:00 2001 From: Miri Bar Date: Sun, 25 Aug 2024 15:53:17 +0300 Subject: [PATCH 1/6] fix: allow positional argument in EmbedResult for backward compatibility --- ai21/models/responses/embed_response.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/ai21/models/responses/embed_response.py b/ai21/models/responses/embed_response.py index 7ded40f6..e0cdf403 100644 --- a/ai21/models/responses/embed_response.py +++ b/ai21/models/responses/embed_response.py @@ -6,6 +6,9 @@ class EmbedResult(AI21BaseModel): embedding: List[float] + def __init__(self, embedding: List[float], **kwargs): + super().__init__(embedding=embedding, **kwargs) + class EmbedResponse(AI21BaseModel): id: str From e86edab3acc56b6f0d946e68e898d49e2f31da98 Mon Sep 17 00:00:00 2001 From: Miri Bar Date: Sun, 25 Aug 2024 20:27:17 +0300 Subject: [PATCH 2/6] chore: remove kwargs --- ai21/models/responses/embed_response.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/ai21/models/responses/embed_response.py b/ai21/models/responses/embed_response.py index e0cdf403..6cb10921 100644 --- a/ai21/models/responses/embed_response.py +++ b/ai21/models/responses/embed_response.py @@ -6,8 +6,8 @@ class EmbedResult(AI21BaseModel): embedding: List[float] - def __init__(self, embedding: List[float], **kwargs): - super().__init__(embedding=embedding, **kwargs) + def __init__(self, embedding: List[float]): + super().__init__(embedding=embedding) class EmbedResponse(AI21BaseModel): From 80f37a8f848b11e006c44fbd8346c3c3da3369e8 Mon Sep 17 00:00:00 2001 From: Miri Bar Date: Sun, 25 Aug 2024 20:46:43 +0300 Subject: [PATCH 3/6] fix: handle non 200 response for streaming --- ai21/http_client/async_http_client.py | 7 ++++++- ai21/http_client/base_http_client.py | 3 +++ ai21/http_client/http_client.py | 7 ++++++- tests/unittests/test_http_client.py | 27 ++++++++++++++++++++++++++- 4 files changed, 41 insertions(+), 3 deletions(-) diff --git a/ai21/http_client/async_http_client.py b/ai21/http_client/async_http_client.py index 13fd536d..a005e333 100644 --- a/ai21/http_client/async_http_client.py +++ b/ai21/http_client/async_http_client.py @@ -103,7 +103,12 @@ async def execute_http_request( logger.error( f"Calling {method} {self._base_url} failed with a non-200 response code: {response.status_code}" ) - handle_non_success_response(response.status_code, response.text) + + if stream: + details = self._extract_streaming_error_details(response) + handle_non_success_response(response.status_code, details) + else: + handle_non_success_response(response.status_code, response.text) return response diff --git a/ai21/http_client/base_http_client.py b/ai21/http_client/base_http_client.py index 669b82dc..8339715c 100644 --- a/ai21/http_client/base_http_client.py +++ b/ai21/http_client/base_http_client.py @@ -171,3 +171,6 @@ def _prepare_url(self, options: RequestOptions) -> str: return f"{options.url}{options.path}" return options.url + + def _extract_streaming_error_details(self, response: httpx.Response) -> str: + return response.read().decode("utf-8") diff --git a/ai21/http_client/http_client.py b/ai21/http_client/http_client.py index bacc10bb..bed3d994 100644 --- a/ai21/http_client/http_client.py +++ b/ai21/http_client/http_client.py @@ -102,7 +102,12 @@ def execute_http_request( f"Calling {method} {self._base_url} failed with a non-200 " f"response code: {response.status_code} headers: {response.headers}" ) - handle_non_success_response(response.status_code, response.text) + + if stream: + details = self._extract_streaming_error_details(response) + handle_non_success_response(response.status_code, details) + else: + handle_non_success_response(response.status_code, response.text) return response diff --git a/tests/unittests/test_http_client.py b/tests/unittests/test_http_client.py index 9b68fbcc..02c44542 100644 --- a/tests/unittests/test_http_client.py +++ b/tests/unittests/test_http_client.py @@ -4,7 +4,7 @@ import httpx -from ai21.errors import ServiceUnavailable +from ai21.errors import ServiceUnavailable, Unauthorized from ai21.http_client.base_http_client import RETRY_ERROR_CODES from ai21.http_client.http_client import AI21HTTPClient from ai21.http_client.async_http_client import AsyncAI21HTTPClient @@ -42,6 +42,17 @@ def test__execute_http_request__when_retry_error__should_retry_and_stop(mock_htt assert mock_httpx_client.send.call_count == retries +def test__execute_http_request__when_streaming__should_handle_non_200_response_code(mock_httpx_client: Mock) -> None: + error_details = "test_error" + request = Request(method=_METHOD, url=_URL) + response = httpx.Response(status_code=401, request=request, text=error_details) + mock_httpx_client.send.return_value = response + + client = AI21HTTPClient(client=mock_httpx_client, base_url=_URL, api_key=_API_KEY) + with pytest.raises(Unauthorized, match=error_details): + client.execute_http_request(method=_METHOD, stream=True) + + @pytest.mark.asyncio async def test__execute_async_http_request__when_retry_error_code_once__should_retry_and_succeed( mock_httpx_async_client: Mock, @@ -74,3 +85,17 @@ async def test__execute_async_http_request__when_retry_error__should_retry_and_s await client.execute_http_request(method=_METHOD) assert mock_httpx_async_client.send.call_count == retries + + +@pytest.mark.asyncio +async def test__execute_async_http_request__when_streaming__should_handle_non_200_response_code( + mock_httpx_async_client: Mock, +) -> None: + error_details = "test_error" + request = Request(method=_METHOD, url=_URL) + response = httpx.Response(status_code=401, request=request, text=error_details) + mock_httpx_async_client.send.return_value = response + + client = AsyncAI21HTTPClient(client=mock_httpx_async_client, base_url=_URL, api_key=_API_KEY) + with pytest.raises(Unauthorized, match=error_details): + await client.execute_http_request(method=_METHOD, stream=True) From 7279c2cf4fc01eabe3912c524182030a7c350938 Mon Sep 17 00:00:00 2001 From: Miri Bar Date: Mon, 26 Aug 2024 09:47:29 +0300 Subject: [PATCH 4/6] chore: add try-except --- ai21/http_client/base_http_client.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/ai21/http_client/base_http_client.py b/ai21/http_client/base_http_client.py index 8339715c..a9fd329c 100644 --- a/ai21/http_client/base_http_client.py +++ b/ai21/http_client/base_http_client.py @@ -173,4 +173,7 @@ def _prepare_url(self, options: RequestOptions) -> str: return options.url def _extract_streaming_error_details(self, response: httpx.Response) -> str: - return response.read().decode("utf-8") + try: + return response.read().decode("utf-8") + except Exception: + return "could not extract streaming error details" From 5e6189c062f90e03b1e78345f823a98ea04862f6 Mon Sep 17 00:00:00 2001 From: Paz Shalev Date: Mon, 26 Aug 2024 10:41:27 +0300 Subject: [PATCH 5/6] fix: fix _request redefining of a member by a method, by renaming the method --- ai21/http_client/async_http_client.py | 4 ++-- ai21/http_client/base_http_client.py | 2 +- ai21/http_client/http_client.py | 6 +++--- 3 files changed, 6 insertions(+), 6 deletions(-) diff --git a/ai21/http_client/async_http_client.py b/ai21/http_client/async_http_client.py index a005e333..cc3f49ec 100644 --- a/ai21/http_client/async_http_client.py +++ b/ai21/http_client/async_http_client.py @@ -55,7 +55,7 @@ def __init__( wait=wait_exponential(multiplier=RETRY_BACK_OFF_FACTOR, min=TIME_BETWEEN_RETRIES), retry=retry_if_result(self._should_retry), stop=stop_after_attempt(self._num_retries), - )(self._request) + )(self._run_request) self._streaming_decoder = _SSEDecoder() async def execute_http_request( @@ -112,7 +112,7 @@ async def execute_http_request( return response - async def _request(self, options: RequestOptions) -> httpx.Response: + async def _run_request(self, options: RequestOptions) -> httpx.Response: request = self._build_request(options) _logger.debug(f"Calling {request.method} {request.url} {request.headers}, {options.body}") diff --git a/ai21/http_client/base_http_client.py b/ai21/http_client/base_http_client.py index a9fd329c..4660aae7 100644 --- a/ai21/http_client/base_http_client.py +++ b/ai21/http_client/base_http_client.py @@ -118,7 +118,7 @@ def execute_http_request( pass @abstractmethod - def _request( + def _run_request( self, options: RequestOptions, ) -> httpx.Response: diff --git a/ai21/http_client/http_client.py b/ai21/http_client/http_client.py index bed3d994..e3cf61de 100644 --- a/ai21/http_client/http_client.py +++ b/ai21/http_client/http_client.py @@ -54,7 +54,7 @@ def __init__( wait=wait_exponential(multiplier=RETRY_BACK_OFF_FACTOR, min=TIME_BETWEEN_RETRIES), retry=retry_if_result(self._should_retry), stop=stop_after_attempt(self._num_retries), - )(self._request) + )(self._run_request) self._streaming_decoder = _SSEDecoder() def execute_http_request( @@ -81,7 +81,7 @@ def execute_http_request( timeout=self._timeout_sec, url=self._base_url, ) - response = self._request(options=options) + response = self._run_request(options=options) except RetryError as retry_error: last_attempt = retry_error.last_attempt @@ -111,7 +111,7 @@ def execute_http_request( return response - def _request(self, options: RequestOptions) -> httpx.Response: + def _run_request(self, options: RequestOptions) -> httpx.Response: request = self._build_request(options) _logger.debug(f"Calling {request.method} {request.url} {request.headers}, {options.body}") From e55185cd543fbc5f694bb9db7847da8fb72156ac Mon Sep 17 00:00:00 2001 From: Paz Shalev Date: Mon, 26 Aug 2024 10:55:38 +0300 Subject: [PATCH 6/6] fix: fix execute_http_request --- ai21/http_client/http_client.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/ai21/http_client/http_client.py b/ai21/http_client/http_client.py index e3cf61de..9aa4a8e4 100644 --- a/ai21/http_client/http_client.py +++ b/ai21/http_client/http_client.py @@ -81,7 +81,7 @@ def execute_http_request( timeout=self._timeout_sec, url=self._base_url, ) - response = self._run_request(options=options) + response = self._request(options=options) except RetryError as retry_error: last_attempt = retry_error.last_attempt