From c76ed1dad645b35e815ab1f8b425bd6c054acac5 Mon Sep 17 00:00:00 2001 From: Josephasafg Date: Wed, 8 May 2024 11:59:05 +0300 Subject: [PATCH 01/12] feat: Added httpx to pyproject --- poetry.lock | 93 ++++++++++++++++++++++++++++++++++++++++++++++++-- pyproject.toml | 1 + 2 files changed, 92 insertions(+), 2 deletions(-) diff --git a/poetry.lock b/poetry.lock index 5b4e0c38..4dda8269 100644 --- a/poetry.lock +++ b/poetry.lock @@ -1,4 +1,4 @@ -# This file is automatically @generated by Poetry 1.7.0 and should not be changed by hand. +# This file is automatically @generated by Poetry 1.8.2 and should not be changed by hand. [[package]] name = "ai21-tokenizer" @@ -29,6 +29,28 @@ files = [ [package.dependencies] typing-extensions = {version = ">=4.0.0", markers = "python_version < \"3.9\""} +[[package]] +name = "anyio" +version = "4.3.0" +description = "High level compatibility layer for multiple asynchronous event loop implementations" +optional = false +python-versions = ">=3.8" +files = [ + {file = "anyio-4.3.0-py3-none-any.whl", hash = "sha256:048e05d0f6caeed70d731f3db756d35dcc1f35747c8c403364a8332c630441b8"}, + {file = "anyio-4.3.0.tar.gz", hash = "sha256:f75253795a87df48568485fd18cdd2a3fa5c4f7c5be8e5e36637733fce06fed6"}, +] + +[package.dependencies] +exceptiongroup = {version = ">=1.0.2", markers = "python_version < \"3.11\""} +idna = ">=2.8" +sniffio = ">=1.1" +typing-extensions = {version = ">=4.1", markers = "python_version < \"3.11\""} + +[package.extras] +doc = ["Sphinx (>=7)", "packaging", "sphinx-autodoc-typehints (>=1.2.0)", "sphinx-rtd-theme"] +test = ["anyio[trio]", "coverage[toml] (>=7)", "exceptiongroup (>=1.2.0)", "hypothesis (>=4.0)", "psutil (>=5.9)", "pytest (>=7.0)", "pytest-mock (>=3.6.1)", "trustme", "uvloop (>=0.17)"] +trio = ["trio (>=0.23)"] + [[package]] name = "black" version = "24.3.0" @@ -392,6 +414,62 @@ gitdb = ">=4.0.1,<5" [package.extras] test = ["black", "coverage[toml]", "ddt (>=1.1.1,!=1.4.3)", "mock", "mypy", "pre-commit", "pytest (>=7.3.1)", "pytest-cov", "pytest-instafail", "pytest-mock", "pytest-sugar", "sumtypes"] +[[package]] +name = "h11" +version = "0.14.0" +description = "A pure-Python, bring-your-own-I/O implementation of HTTP/1.1" +optional = false +python-versions = ">=3.7" +files = [ + {file = "h11-0.14.0-py3-none-any.whl", hash = "sha256:e3fe4ac4b851c468cc8363d500db52c2ead036020723024a109d37346efaa761"}, + {file = "h11-0.14.0.tar.gz", hash = "sha256:8f19fbbe99e72420ff35c00b27a34cb9937e902a8b810e2c88300c6f0a3b699d"}, +] + +[[package]] +name = "httpcore" +version = "1.0.5" +description = "A minimal low-level HTTP client." +optional = false +python-versions = ">=3.8" +files = [ + {file = "httpcore-1.0.5-py3-none-any.whl", hash = "sha256:421f18bac248b25d310f3cacd198d55b8e6125c107797b609ff9b7a6ba7991b5"}, + {file = "httpcore-1.0.5.tar.gz", hash = "sha256:34a38e2f9291467ee3b44e89dd52615370e152954ba21721378a87b2960f7a61"}, +] + +[package.dependencies] +certifi = "*" +h11 = ">=0.13,<0.15" + +[package.extras] +asyncio = ["anyio (>=4.0,<5.0)"] +http2 = ["h2 (>=3,<5)"] +socks = ["socksio (==1.*)"] +trio = ["trio (>=0.22.0,<0.26.0)"] + +[[package]] +name = "httpx" +version = "0.27.0" +description = "The next generation HTTP client." +optional = false +python-versions = ">=3.8" +files = [ + {file = "httpx-0.27.0-py3-none-any.whl", hash = "sha256:71d5465162c13681bff01ad59b2cc68dd838ea1f10e51574bac27103f00c91a5"}, + {file = "httpx-0.27.0.tar.gz", hash = "sha256:a0cb88a46f32dc874e04ee956e4c2764aba2aa228f650b06788ba6bda2962ab5"}, +] + +[package.dependencies] +anyio = "*" +certifi = "*" +httpcore = "==1.*" +idna = "*" +sniffio = "*" + +[package.extras] +brotli = ["brotli", "brotlicffi"] +cli = ["click (==8.*)", "pygments (==2.*)", "rich (>=10,<14)"] +http2 = ["h2 (>=3,<5)"] +socks = ["socksio (==1.*)"] + [[package]] name = "huggingface-hub" version = "0.22.2" @@ -1375,6 +1453,17 @@ files = [ {file = "smmap-5.0.1.tar.gz", hash = "sha256:dceeb6c0028fdb6734471eb07c0cd2aae706ccaecab45965ee83f11c8d3b1f62"}, ] +[[package]] +name = "sniffio" +version = "1.3.1" +description = "Sniff out which async library your code is running under" +optional = false +python-versions = ">=3.7" +files = [ + {file = "sniffio-1.3.1-py3-none-any.whl", hash = "sha256:2f6da418d1f1e0fddd844478f41680e794e6051915791a034ff65e5f100525a2"}, + {file = "sniffio-1.3.1.tar.gz", hash = "sha256:f4324edc670a0f49750a81b895f35c3adb843cca46f0530f79fc1babb23789dc"}, +] + [[package]] name = "tokenizers" version = "0.15.2" @@ -1624,4 +1713,4 @@ aws = ["boto3"] [metadata] lock-version = "2.0" python-versions = "^3.8" -content-hash = "915b491d362897a01bb63a9fcc57b19915319a8f6d24e9157444520bed4c5b87" +content-hash = "5acd8eae6220934a23cb3ed8b03845d0afbc592aa94729758571eb2936814a26" diff --git a/pyproject.toml b/pyproject.toml index 9546e1e8..a887070f 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -61,6 +61,7 @@ ai21-tokenizer = "^0.9.0" boto3 = { version = "^1.28.82", optional = true } dataclasses-json = "^0.6.3" typing-extensions = "^4.9.0" +httpx = "^0.27.0" [tool.poetry.group.dev.dependencies] From 9fff865aa834bb3d230f5086add6f9868e91ff97 Mon Sep 17 00:00:00 2001 From: Josephasafg Date: Wed, 8 May 2024 13:24:34 +0300 Subject: [PATCH 02/12] feat: httpx instead of requests --- ai21/http_client.py | 50 +++++++++++++++------------------------------ 1 file changed, 16 insertions(+), 34 deletions(-) diff --git a/ai21/http_client.py b/ai21/http_client.py index 229c5156..40aae793 100644 --- a/ai21/http_client.py +++ b/ai21/http_client.py @@ -1,8 +1,8 @@ import json from typing import Optional, Dict, Any, BinaryIO -import requests -from requests.adapters import HTTPAdapter, Retry, RetryError +import httpx +from httpx import ConnectError from ai21.errors import ( BadRequest, @@ -39,25 +39,16 @@ def handle_non_success_response(status_code: int, response_text: str): raise AI21APIError(status_code, details=response_text) -def requests_retry_session(session, retries=0): - retry = Retry( - total=retries, - read=retries, - connect=retries, - backoff_factor=RETRY_BACK_OFF_FACTOR, - status_forcelist=RETRY_ERROR_CODES, - allowed_methods=frozenset(RETRY_METHOD_WHITELIST), +def requests_retry_session(retries: int) -> httpx.HTTPTransport: + return httpx.HTTPTransport( + retries=retries, ) - adapter = HTTPAdapter(max_retries=retry) - session.mount("https://", adapter) - session.mount("http://", adapter) - return session class HttpClient: def __init__( self, - session: Optional[requests.Session] = None, + session: Optional[httpx.Client] = None, timeout_sec: int = None, num_retries: int = None, headers: Dict = None, @@ -66,7 +57,7 @@ def __init__( self._num_retries = num_retries or DEFAULT_NUM_RETRIES self._headers = headers or {} self._apply_retry_policy = self._num_retries > 0 - self._session = self._init_session(session) + self._client = self._init_client(session) def execute_http_request( self, @@ -82,7 +73,7 @@ def execute_http_request( try: if method == "GET": - response = self._session.request( + response = self._client.request( method=method, url=url, headers=headers, @@ -98,7 +89,7 @@ def execute_http_request( headers.pop( "Content-Type" ) # multipart/form-data 'Content-Type' is being added when passing rb files and payload - response = self._session.request( + response = self._client.request( method=method, url=url, headers=headers, @@ -107,34 +98,25 @@ def execute_http_request( timeout=timeout, ) else: - response = self._session.request(method=method, url=url, headers=headers, data=data, timeout=timeout) - except ConnectionError as connection_error: + response = self._client.request(method=method, url=url, headers=headers, data=data, timeout=timeout) + except ConnectError as connection_error: logger.error(f"Calling {method} {url} failed with ConnectionError: {connection_error}") raise connection_error - except RetryError as retry_error: - logger.error( - f"Calling {method} {url} failed with RetryError after {self._num_retries} attempts: {retry_error}" - ) - raise retry_error except Exception as exception: logger.error(f"Calling {method} {url} failed with Exception: {exception}") raise exception - if response.status_code != 200: + if response.status_code != httpx.codes.OK: logger.error(f"Calling {method} {url} failed with a non-200 response code: {response.status_code}") handle_non_success_response(response.status_code, response.text) return response.json() - def _init_session(self, session: Optional[requests.Session]) -> requests.Session: - if session is not None: - return session + def _init_client(self, client: Optional[httpx.Client]) -> httpx.Client: + if client is not None: + return client - return ( - requests_retry_session(requests.Session(), retries=self._num_retries) - if self._apply_retry_policy - else requests.Session() - ) + return requests_retry_session(retries=self._num_retries) if self._apply_retry_policy else httpx.Client() def add_headers(self, headers: Dict[str, Any]) -> None: self._headers.update(headers) From 803186005c10b71029f650e17c5bd33609fa2cb0 Mon Sep 17 00:00:00 2001 From: Josephasafg Date: Wed, 8 May 2024 13:26:20 +0300 Subject: [PATCH 03/12] feat: Removed requests --- poetry.lock | 2 +- pyproject.toml | 1 - 2 files changed, 1 insertion(+), 2 deletions(-) diff --git a/poetry.lock b/poetry.lock index 4dda8269..c072902b 100644 --- a/poetry.lock +++ b/poetry.lock @@ -1713,4 +1713,4 @@ aws = ["boto3"] [metadata] lock-version = "2.0" python-versions = "^3.8" -content-hash = "5acd8eae6220934a23cb3ed8b03845d0afbc592aa94729758571eb2936814a26" +content-hash = "f3e6fa54edc225744d4751857e4cc96a74de0b673c9d31b645ee8903ed129790" diff --git a/pyproject.toml b/pyproject.toml index a887070f..55e8056f 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -56,7 +56,6 @@ packages = [ [tool.poetry.dependencies] python = "^3.8" -requests = "^2.31.0" ai21-tokenizer = "^0.9.0" boto3 = { version = "^1.28.82", optional = true } dataclasses-json = "^0.6.3" From d67d97c16ecc97e5cac25d13dfb19d74139d660d Mon Sep 17 00:00:00 2001 From: Josephasafg Date: Wed, 8 May 2024 13:56:32 +0300 Subject: [PATCH 04/12] fix: not given --- .../studio/resources/studio_library.py | 77 +++++++++++-------- .../clients/studio/conftest.py | 2 +- 2 files changed, 44 insertions(+), 35 deletions(-) diff --git a/ai21/clients/studio/resources/studio_library.py b/ai21/clients/studio/resources/studio_library.py index 3288c7c7..8b6e6622 100644 --- a/ai21/clients/studio/resources/studio_library.py +++ b/ai21/clients/studio/resources/studio_library.py @@ -1,8 +1,11 @@ +from __future__ import annotations from typing import Optional, List from ai21.ai21_http_client import AI21HTTPClient from ai21.clients.studio.resources.studio_resource import StudioResource from ai21.models import FileResponse, LibraryAnswerResponse, LibrarySearchResponse +from ai21.types import NotGiven, NOT_GIVEN +from ai21.utils.typing import remove_not_given class StudioLibrary(StudioResource): @@ -22,14 +25,14 @@ def create( self, file_path: str, *, - path: Optional[str] = None, - labels: Optional[List[str]] = None, - public_url: Optional[str] = None, + path: Optional[str] | NotGiven = NOT_GIVEN, + labels: Optional[List[str]] | NotGiven = NOT_GIVEN, + public_url: Optional[str] | NotGiven = NOT_GIVEN, **kwargs, ) -> str: url = f"{self._client.get_base_url()}/{self._module_name}" files = {"file": open(file_path, "rb")} - body = {"path": path, "labels": labels, "publicUrl": public_url, **kwargs} + body = remove_not_given({"path": path, "labels": labels, "publicUrl": public_url, **kwargs}) raw_response = self._post(url=url, files=files, body=body) @@ -44,12 +47,12 @@ def get(self, file_id: str) -> FileResponse: def list( self, *, - offset: Optional[int] = None, - limit: Optional[int] = None, + offset: Optional[int] | NotGiven = NOT_GIVEN, + limit: Optional[int] | NotGiven = NOT_GIVEN, **kwargs, ) -> List[FileResponse]: url = f"{self._client.get_base_url()}/{self._module_name}" - params = {"offset": offset, "limit": limit} + params = remove_not_given({"offset": offset, "limit": limit}) raw_response = self._get(url=url, params=params) return [FileResponse.from_dict(file) for file in raw_response] @@ -58,16 +61,18 @@ def update( self, file_id: str, *, - public_url: Optional[str] = None, - labels: Optional[List[str]] = None, + public_url: Optional[str] | NotGiven = NOT_GIVEN, + labels: Optional[List[str]] | NotGiven = NOT_GIVEN, **kwargs, ) -> None: url = f"{self._client.get_base_url()}/{self._module_name}/{file_id}" - body = { - "publicUrl": public_url, - "labels": labels, - **kwargs, - } + body = remove_not_given( + { + "publicUrl": public_url, + "labels": labels, + **kwargs, + } + ) self._put(url=url, body=body) def delete(self, file_id: str) -> None: @@ -82,19 +87,21 @@ def create( self, query: str, *, - path: Optional[str] = None, - field_ids: Optional[List[str]] = None, - max_segments: Optional[int] = None, + path: Optional[str] | NotGiven = NOT_GIVEN, + field_ids: Optional[List[str]] | NotGiven = NOT_GIVEN, + max_segments: Optional[int] | NotGiven = NOT_GIVEN, **kwargs, ) -> LibrarySearchResponse: url = f"{self._client.get_base_url()}/{self._module_name}" - body = { - "query": query, - "path": path, - "fieldIds": field_ids, - "maxSegments": max_segments, - **kwargs, - } + body = remove_not_given( + { + "query": query, + "path": path, + "fieldIds": field_ids, + "maxSegments": max_segments, + **kwargs, + } + ) raw_response = self._post(url=url, body=body) return LibrarySearchResponse.from_dict(raw_response) @@ -106,18 +113,20 @@ def create( self, question: str, *, - path: Optional[str] = None, - field_ids: Optional[List[str]] = None, - labels: Optional[List[str]] = None, + path: Optional[str] | NotGiven = NOT_GIVEN, + field_ids: Optional[List[str]] | NotGiven = NOT_GIVEN, + labels: Optional[List[str]] | NotGiven = NOT_GIVEN, **kwargs, ) -> LibraryAnswerResponse: url = f"{self._client.get_base_url()}/{self._module_name}" - body = { - "question": question, - "path": path, - "fieldIds": field_ids, - "labels": labels, - **kwargs, - } + body = remove_not_given( + { + "question": question, + "path": path, + "fieldIds": field_ids, + "labels": labels, + **kwargs, + } + ) raw_response = self._post(url=url, body=body) return LibraryAnswerResponse.from_dict(raw_response) diff --git a/tests/integration_tests/clients/studio/conftest.py b/tests/integration_tests/clients/studio/conftest.py index 19da28b3..c1ea9400 100644 --- a/tests/integration_tests/clients/studio/conftest.py +++ b/tests/integration_tests/clients/studio/conftest.py @@ -42,7 +42,7 @@ def file_in_library(): # Delete any file that might be in the library due to failed tests files = client.library.files.list() for file in files: - _delete_uploaded_file(file.file_id) + _delete_uploaded_file(client, file.file_id) file_id = client.library.files.create(file_path=LIBRARY_FILE_TO_UPLOAD, labels=DEFAULT_LABELS) _wait_for_file_to_process(client, file_id) From e7b306cdcd2f053bfe063600ef38fb443ea11a14 Mon Sep 17 00:00:00 2001 From: Josephasafg Date: Wed, 8 May 2024 14:01:40 +0300 Subject: [PATCH 05/12] fix: setup --- setup.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/setup.py b/setup.py index 2ba17979..e65e5028 100755 --- a/setup.py +++ b/setup.py @@ -21,6 +21,6 @@ packages=find_packages(exclude=["tests", "tests.*"]), keywords=["python", "sdk", "ai", "ai21", "jurassic", "ai21-python", "llm"], install_requires=[ - "requests", + "httpx", ], ) From 020cefb2245662d1c0ba788dc2bb2dad1d2fe68d Mon Sep 17 00:00:00 2001 From: Josephasafg Date: Wed, 8 May 2024 15:35:33 +0300 Subject: [PATCH 06/12] feat: Added tenacity for retry --- ai21/http_client.py | 87 +++++++++++++++++++++++++++------------------ poetry.lock | 17 ++++++++- pyproject.toml | 1 + 3 files changed, 69 insertions(+), 36 deletions(-) diff --git a/ai21/http_client.py b/ai21/http_client.py index 40aae793..4f9fe81e 100644 --- a/ai21/http_client.py +++ b/ai21/http_client.py @@ -3,6 +3,7 @@ import httpx from httpx import ConnectError +from tenacity import retry, retry_if_result, stop_after_attempt, wait_fixed from ai21.errors import ( BadRequest, @@ -17,7 +18,6 @@ DEFAULT_TIMEOUT_SEC = 300 DEFAULT_NUM_RETRIES = 0 -RETRY_BACK_OFF_FACTOR = 0.5 TIME_BETWEEN_RETRIES = 1000 RETRY_ERROR_CODES = (408, 429, 500, 503) RETRY_METHOD_WHITELIST = ["GET", "POST", "PUT"] @@ -39,7 +39,7 @@ def handle_non_success_response(status_code: int, response_text: str): raise AI21APIError(status_code, details=response_text) -def requests_retry_session(retries: int) -> httpx.HTTPTransport: +def _requests_retry_session(retries: int) -> httpx.HTTPTransport: return httpx.HTTPTransport( retries=retries, ) @@ -59,6 +59,17 @@ def __init__( self._apply_retry_policy = self._num_retries > 0 self._client = self._init_client(session) + # Since we can't use the retry decorator on a method of a class as we can't access class attributes, + # we have to wrap the method in a function + self._request = retry( + wait=wait_fixed(TIME_BETWEEN_RETRIES), + retry=retry_if_result(self._should_retry), + stop=stop_after_attempt(self._num_retries), + )(self._request) + + def _should_retry(self, response: httpx.Response) -> bool: + return response.status_code in RETRY_ERROR_CODES and response.request.method in RETRY_METHOD_WHITELIST + def execute_http_request( self, method: str, @@ -66,39 +77,8 @@ def execute_http_request( params: Optional[Dict] = None, files: Optional[Dict[str, BinaryIO]] = None, ): - timeout = self._timeout_sec - headers = self._headers - data = json.dumps(params).encode() - logger.debug(f"Calling {method} {url} {headers} {data}") - try: - if method == "GET": - response = self._client.request( - method=method, - url=url, - headers=headers, - timeout=timeout, - params=params, - ) - elif files is not None: - if method != "POST": - raise ValueError( - f"execute_http_request supports only POST for files upload, but {method} was supplied instead" - ) - if "Content-Type" in headers: - headers.pop( - "Content-Type" - ) # multipart/form-data 'Content-Type' is being added when passing rb files and payload - response = self._client.request( - method=method, - url=url, - headers=headers, - data=params, - files=files, - timeout=timeout, - ) - else: - response = self._client.request(method=method, url=url, headers=headers, data=data, timeout=timeout) + response = self._request(files=files, method=method, params=params, url=url) except ConnectError as connection_error: logger.error(f"Calling {method} {url} failed with ConnectionError: {connection_error}") raise connection_error @@ -112,11 +92,48 @@ def execute_http_request( return response.json() + def _request( + self, files: Optional[Dict[str, BinaryIO]], method: str, params: Optional[Dict], url: str + ) -> httpx.Response: + timeout = self._timeout_sec + headers = self._headers + data = json.dumps(params).encode() + logger.debug(f"Calling {method} {url} {headers} {data}") + + if method == "GET": + return self._client.request( + method=method, + url=url, + headers=headers, + timeout=timeout, + params=params, + ) + + if files is not None: + if method != "POST": + raise ValueError( + f"execute_http_request supports only POST for files upload, but {method} was supplied instead" + ) + if "Content-Type" in headers: + headers.pop( + "Content-Type" + ) # multipart/form-data 'Content-Type' is being added when passing rb files and payload + return self._client.request( + method=method, + url=url, + headers=headers, + data=params, + files=files, + timeout=timeout, + ) + + return self._client.request(method=method, url=url, headers=headers, data=data, timeout=timeout) + def _init_client(self, client: Optional[httpx.Client]) -> httpx.Client: if client is not None: return client - return requests_retry_session(retries=self._num_retries) if self._apply_retry_policy else httpx.Client() + return _requests_retry_session(retries=self._num_retries) if self._apply_retry_policy else httpx.Client() def add_headers(self, headers: Dict[str, Any]) -> None: self._headers.update(headers) diff --git a/poetry.lock b/poetry.lock index c072902b..9409cfcc 100644 --- a/poetry.lock +++ b/poetry.lock @@ -1464,6 +1464,21 @@ files = [ {file = "sniffio-1.3.1.tar.gz", hash = "sha256:f4324edc670a0f49750a81b895f35c3adb843cca46f0530f79fc1babb23789dc"}, ] +[[package]] +name = "tenacity" +version = "8.3.0" +description = "Retry code until it succeeds" +optional = false +python-versions = ">=3.8" +files = [ + {file = "tenacity-8.3.0-py3-none-any.whl", hash = "sha256:3649f6443dbc0d9b01b9d8020a9c4ec7a1ff5f6f3c6c8a036ef371f573fe9185"}, + {file = "tenacity-8.3.0.tar.gz", hash = "sha256:953d4e6ad24357bceffbc9707bc74349aca9d245f68eb65419cf0c249a1949a2"}, +] + +[package.extras] +doc = ["reno", "sphinx"] +test = ["pytest", "tornado (>=4.5)", "typeguard"] + [[package]] name = "tokenizers" version = "0.15.2" @@ -1713,4 +1728,4 @@ aws = ["boto3"] [metadata] lock-version = "2.0" python-versions = "^3.8" -content-hash = "f3e6fa54edc225744d4751857e4cc96a74de0b673c9d31b645ee8903ed129790" +content-hash = "c6c474d713d3660255aade619131c42d2c57ba74d992b7e20be02275601a5b48" diff --git a/pyproject.toml b/pyproject.toml index 55e8056f..473cdac0 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -61,6 +61,7 @@ boto3 = { version = "^1.28.82", optional = true } dataclasses-json = "^0.6.3" typing-extensions = "^4.9.0" httpx = "^0.27.0" +tenacity = "^8.3.0" [tool.poetry.group.dev.dependencies] From 2721928565bc3aa9c33988e17032599eac82c5b3 Mon Sep 17 00:00:00 2001 From: Josephasafg Date: Wed, 8 May 2024 15:44:10 +0300 Subject: [PATCH 07/12] fix: conftest --- tests/unittests/conftest.py | 6 +++--- tests/unittests/test_ai21_http_client.py | 18 +++++++++--------- 2 files changed, 12 insertions(+), 12 deletions(-) diff --git a/tests/unittests/conftest.py b/tests/unittests/conftest.py index 02e4d467..68c4efb5 100644 --- a/tests/unittests/conftest.py +++ b/tests/unittests/conftest.py @@ -1,5 +1,5 @@ import pytest -import requests +import httpx @pytest.fixture @@ -8,5 +8,5 @@ def dummy_api_host() -> str: @pytest.fixture -def mock_requests_session(mocker) -> requests.Session: - return mocker.Mock(spec=requests.Session) +def mock_httpx_client(mocker) -> httpx.Client: + return mocker.Mock(spec=httpx.Client) diff --git a/tests/unittests/test_ai21_http_client.py b/tests/unittests/test_ai21_http_client.py index 86cde2ee..863225d1 100644 --- a/tests/unittests/test_ai21_http_client.py +++ b/tests/unittests/test_ai21_http_client.py @@ -1,8 +1,8 @@ import platform from typing import Optional +import httpx import pytest -import requests from ai21.ai21_http_client import AI21HTTPClient from ai21.http_client import HttpClient @@ -94,12 +94,12 @@ def test__execute_http_request__( params, headers, dummy_api_host: str, - mock_requests_session: requests.Session, + mock_httpx_client: httpx.Client, ): response_json = {"test_key": "test_value"} - mock_requests_session.request.return_value = MockResponse(response_json, 200) + mock_httpx_client.request.return_value = MockResponse(response_json, 200) - http_client = HttpClient(session=mock_requests_session) + http_client = HttpClient(session=mock_httpx_client) client = AI21HTTPClient(http_client=http_client, api_key=_DUMMY_API_KEY, api_host=dummy_api_host, api_version="v1") response = client.execute_http_request(**params) @@ -107,7 +107,7 @@ def test__execute_http_request__( if "files" in params: # We split it because when calling requests with "files", "params" is turned into "data" - mock_requests_session.request.assert_called_once_with( + mock_httpx_client.request.assert_called_once_with( timeout=300, headers=headers, files=params["files"], @@ -116,18 +116,18 @@ def test__execute_http_request__( method=params["method"], ) else: - mock_requests_session.request.assert_called_once_with(timeout=300, headers=headers, **params) + mock_httpx_client.request.assert_called_once_with(timeout=300, headers=headers, **params) def test__execute_http_request__when_files_with_put_method__should_raise_value_error( dummy_api_host: str, - mock_requests_session: requests.Session, + mock_httpx_client: httpx.Client, ): response_json = {"test_key": "test_value"} - http_client = HttpClient(session=mock_requests_session) + http_client = HttpClient(session=mock_httpx_client) client = AI21HTTPClient(http_client=http_client, api_key=_DUMMY_API_KEY, api_host=dummy_api_host, api_version="v1") - mock_requests_session.request.return_value = MockResponse(response_json, 200) + mock_httpx_client.request.return_value = MockResponse(response_json, 200) with pytest.raises(ValueError): params = {"method": "PUT", "url": "test_url", "params": {"foo": "bar"}, "files": {"file": "test_file"}} client.execute_http_request(**params) From c7ac6e766dea09ba2d5b534dd6803226ef6bb00d Mon Sep 17 00:00:00 2001 From: Josephasafg Date: Wed, 8 May 2024 16:33:31 +0300 Subject: [PATCH 08/12] test: Added tests --- ai21/http_client.py | 16 ++++++-- tests/unittests/test_ai21_http_client.py | 4 +- tests/unittests/test_http_client.py | 40 +++++++++++++++++++ tests/unittests/tokenizers/__init__.py | 0 .../test_ai21_tokenizer.py | 0 5 files changed, 54 insertions(+), 6 deletions(-) create mode 100644 tests/unittests/test_http_client.py delete mode 100644 tests/unittests/tokenizers/__init__.py rename tests/unittests/{tokenizers => tokenizerss}/test_ai21_tokenizer.py (100%) diff --git a/ai21/http_client.py b/ai21/http_client.py index 4f9fe81e..da154061 100644 --- a/ai21/http_client.py +++ b/ai21/http_client.py @@ -3,7 +3,7 @@ import httpx from httpx import ConnectError -from tenacity import retry, retry_if_result, stop_after_attempt, wait_fixed +from tenacity import retry, retry_if_result, stop_after_attempt, wait_fixed, RetryError from ai21.errors import ( BadRequest, @@ -18,7 +18,7 @@ DEFAULT_TIMEOUT_SEC = 300 DEFAULT_NUM_RETRIES = 0 -TIME_BETWEEN_RETRIES = 1000 +TIME_BETWEEN_RETRIES = 1 RETRY_ERROR_CODES = (408, 429, 500, 503) RETRY_METHOD_WHITELIST = ["GET", "POST", "PUT"] @@ -48,7 +48,7 @@ def _requests_retry_session(retries: int) -> httpx.HTTPTransport: class HttpClient: def __init__( self, - session: Optional[httpx.Client] = None, + client: Optional[httpx.Client] = None, timeout_sec: int = None, num_retries: int = None, headers: Dict = None, @@ -57,7 +57,7 @@ def __init__( self._num_retries = num_retries or DEFAULT_NUM_RETRIES self._headers = headers or {} self._apply_retry_policy = self._num_retries > 0 - self._client = self._init_client(session) + self._client = self._init_client(client) # Since we can't use the retry decorator on a method of a class as we can't access class attributes, # we have to wrap the method in a function @@ -79,6 +79,14 @@ def execute_http_request( ): try: response = self._request(files=files, method=method, params=params, url=url) + except RetryError as retry_error: + last_attempt = retry_error.last_attempt + + if last_attempt.failed: + raise last_attempt.exception() + else: + response = last_attempt.result() + except ConnectError as connection_error: logger.error(f"Calling {method} {url} failed with ConnectionError: {connection_error}") raise connection_error diff --git a/tests/unittests/test_ai21_http_client.py b/tests/unittests/test_ai21_http_client.py index 863225d1..1d730f1a 100644 --- a/tests/unittests/test_ai21_http_client.py +++ b/tests/unittests/test_ai21_http_client.py @@ -99,7 +99,7 @@ def test__execute_http_request__( response_json = {"test_key": "test_value"} mock_httpx_client.request.return_value = MockResponse(response_json, 200) - http_client = HttpClient(session=mock_httpx_client) + http_client = HttpClient(client=mock_httpx_client) client = AI21HTTPClient(http_client=http_client, api_key=_DUMMY_API_KEY, api_host=dummy_api_host, api_version="v1") response = client.execute_http_request(**params) @@ -124,7 +124,7 @@ def test__execute_http_request__when_files_with_put_method__should_raise_value_e mock_httpx_client: httpx.Client, ): response_json = {"test_key": "test_value"} - http_client = HttpClient(session=mock_httpx_client) + http_client = HttpClient(client=mock_httpx_client) client = AI21HTTPClient(http_client=http_client, api_key=_DUMMY_API_KEY, api_host=dummy_api_host, api_version="v1") mock_httpx_client.request.return_value = MockResponse(response_json, 200) diff --git a/tests/unittests/test_http_client.py b/tests/unittests/test_http_client.py new file mode 100644 index 00000000..9358440b --- /dev/null +++ b/tests/unittests/test_http_client.py @@ -0,0 +1,40 @@ +import pytest +from unittest.mock import Mock +from urllib.request import Request + +import httpx + +from ai21 import TooManyRequestsError +from ai21.http_client import HttpClient + +_METHOD = "GET" +_URL = "http://test_url" + + +def test__execute_http_request__when_retry_error_code_once__should_retry_and_succeed(mock_httpx_client: Mock) -> None: + request = Request(method=_METHOD, url=_URL) + retries = 3 + mock_httpx_client.request.side_effect = [ + httpx.Response(status_code=429, request=request), + httpx.Response(status_code=200, request=request, json={"test_key": "test_value"}), + ] + + client = HttpClient(client=mock_httpx_client, num_retries=retries) + client.execute_http_request(method=_METHOD, url=_URL) + assert mock_httpx_client.request.call_count == retries - 1 + + +def test__execute_http_request__when_retry_error__should_retry_and_stop(mock_httpx_client: Mock) -> None: + request = Request(method=_METHOD, url=_URL) + retries = 3 + mock_httpx_client.request.side_effect = [ + httpx.Response(status_code=429, request=request), + httpx.Response(status_code=429, request=request), + httpx.Response(status_code=429, request=request), + ] + + client = HttpClient(client=mock_httpx_client, num_retries=retries) + with pytest.raises(TooManyRequestsError): + client.execute_http_request(method=_METHOD, url=_URL) + + assert mock_httpx_client.request.call_count == retries diff --git a/tests/unittests/tokenizers/__init__.py b/tests/unittests/tokenizers/__init__.py deleted file mode 100644 index e69de29b..00000000 diff --git a/tests/unittests/tokenizers/test_ai21_tokenizer.py b/tests/unittests/tokenizerss/test_ai21_tokenizer.py similarity index 100% rename from tests/unittests/tokenizers/test_ai21_tokenizer.py rename to tests/unittests/tokenizerss/test_ai21_tokenizer.py From 2f3599a7456793d22c4d6ed0961a306609b3f5bd Mon Sep 17 00:00:00 2001 From: Josephasafg Date: Wed, 8 May 2024 16:37:48 +0300 Subject: [PATCH 09/12] fix: Rename --- tests/unittests/tokenizers/__init__.py | 0 .../unittests/{tokenizerss => tokenizers}/test_ai21_tokenizer.py | 0 2 files changed, 0 insertions(+), 0 deletions(-) create mode 100644 tests/unittests/tokenizers/__init__.py rename tests/unittests/{tokenizerss => tokenizers}/test_ai21_tokenizer.py (100%) diff --git a/tests/unittests/tokenizers/__init__.py b/tests/unittests/tokenizers/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/tests/unittests/tokenizerss/test_ai21_tokenizer.py b/tests/unittests/tokenizers/test_ai21_tokenizer.py similarity index 100% rename from tests/unittests/tokenizerss/test_ai21_tokenizer.py rename to tests/unittests/tokenizers/test_ai21_tokenizer.py From 3401934daba4a8d2ad7d99a11ed25fe87db26b78 Mon Sep 17 00:00:00 2001 From: Josephasafg Date: Thu, 9 May 2024 09:06:55 +0300 Subject: [PATCH 10/12] fix: Modified test --- tests/unittests/test_http_client.py | 13 ++++++------- 1 file changed, 6 insertions(+), 7 deletions(-) diff --git a/tests/unittests/test_http_client.py b/tests/unittests/test_http_client.py index 9358440b..73347f09 100644 --- a/tests/unittests/test_http_client.py +++ b/tests/unittests/test_http_client.py @@ -4,8 +4,8 @@ import httpx -from ai21 import TooManyRequestsError -from ai21.http_client import HttpClient +from ai21.errors import ServiceUnavailable +from ai21.http_client import HttpClient, RETRY_ERROR_CODES _METHOD = "GET" _URL = "http://test_url" @@ -26,15 +26,14 @@ def test__execute_http_request__when_retry_error_code_once__should_retry_and_suc def test__execute_http_request__when_retry_error__should_retry_and_stop(mock_httpx_client: Mock) -> None: request = Request(method=_METHOD, url=_URL) - retries = 3 + retries = len(RETRY_ERROR_CODES) + mock_httpx_client.request.side_effect = [ - httpx.Response(status_code=429, request=request), - httpx.Response(status_code=429, request=request), - httpx.Response(status_code=429, request=request), + httpx.Response(status_code=status_code, request=request) for status_code in RETRY_ERROR_CODES ] client = HttpClient(client=mock_httpx_client, num_retries=retries) - with pytest.raises(TooManyRequestsError): + with pytest.raises(ServiceUnavailable): client.execute_http_request(method=_METHOD, url=_URL) assert mock_httpx_client.request.call_count == retries From 8c505fcbccb4a1ba0759dc4913994c8a6abf2e01 Mon Sep 17 00:00:00 2001 From: Josephasafg Date: Mon, 13 May 2024 10:02:30 +0300 Subject: [PATCH 11/12] fix: CR --- ai21/http_client.py | 26 ++++++++++++-------------- 1 file changed, 12 insertions(+), 14 deletions(-) diff --git a/ai21/http_client.py b/ai21/http_client.py index da154061..4ef88145 100644 --- a/ai21/http_client.py +++ b/ai21/http_client.py @@ -1,9 +1,8 @@ -import json from typing import Optional, Dict, Any, BinaryIO import httpx from httpx import ConnectError -from tenacity import retry, retry_if_result, stop_after_attempt, wait_fixed, RetryError +from tenacity import retry, retry_if_result, stop_after_attempt, wait_exponential, RetryError from ai21.errors import ( BadRequest, @@ -19,6 +18,7 @@ DEFAULT_TIMEOUT_SEC = 300 DEFAULT_NUM_RETRIES = 0 TIME_BETWEEN_RETRIES = 1 +RETRY_BACK_OFF_FACTOR = 0.5 RETRY_ERROR_CODES = (408, 429, 500, 503) RETRY_METHOD_WHITELIST = ["GET", "POST", "PUT"] @@ -62,7 +62,7 @@ def __init__( # Since we can't use the retry decorator on a method of a class as we can't access class attributes, # we have to wrap the method in a function self._request = retry( - wait=wait_fixed(TIME_BETWEEN_RETRIES), + 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) @@ -105,8 +105,7 @@ def _request( ) -> httpx.Response: timeout = self._timeout_sec headers = self._headers - data = json.dumps(params).encode() - logger.debug(f"Calling {method} {url} {headers} {data}") + logger.debug(f"Calling {method} {url} {headers} {params}") if method == "GET": return self._client.request( @@ -126,16 +125,15 @@ def _request( headers.pop( "Content-Type" ) # multipart/form-data 'Content-Type' is being added when passing rb files and payload - return self._client.request( - method=method, - url=url, - headers=headers, - data=params, - files=files, - timeout=timeout, - ) - return self._client.request(method=method, url=url, headers=headers, data=data, timeout=timeout) + return self._client.request( + method=method, + url=url, + headers=headers, + json=params, + timeout=timeout, + files=files, + ) def _init_client(self, client: Optional[httpx.Client]) -> httpx.Client: if client is not None: From c5b988efb9475d5aaa0a3dda4b88d51e110a911b Mon Sep 17 00:00:00 2001 From: Josephasafg Date: Mon, 13 May 2024 10:46:07 +0300 Subject: [PATCH 12/12] fix: request --- ai21/http_client.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/ai21/http_client.py b/ai21/http_client.py index 4ef88145..e199cb1b 100644 --- a/ai21/http_client.py +++ b/ai21/http_client.py @@ -1,3 +1,4 @@ +import json from typing import Optional, Dict, Any, BinaryIO import httpx @@ -125,12 +126,15 @@ def _request( headers.pop( "Content-Type" ) # multipart/form-data 'Content-Type' is being added when passing rb files and payload + data = params + else: + data = json.dumps(params).encode() if params else None return self._client.request( method=method, url=url, headers=headers, - json=params, + data=data, timeout=timeout, files=files, )