feat: add support for studio async client (#129)

* feat: async support - stream, http, ai21 http

* fix: commit changes

* feat: studio resource, chat, chat completions, answer

* feat: beta, dataset, completion, custom model

* feat: embed, gec, improvements

* feat: paraphrase, segmentation, summarize, by segment

* feat: library

* feat: client

* refactor: sync and async http, ai21 http, ai21 client, resources

* test: update imports, create tests for async

* fix: base client

* fix: add pytest marker asyncio

* fix: add pytest asyncio to poetry

* fix: add delete to lib files, add examples, test examples

* fix: tests

* fix: fix stream, add stream tests, add readme

* fix: fix import on sm stub

* feat: async support - stream, http, ai21 http

* fix: commit changes

* feat: studio resource, chat, chat completions, answer

* feat: beta, dataset, completion, custom model

* feat: embed, gec, improvements

* feat: paraphrase, segmentation, summarize, by segment

* feat: library

* feat: client

* refactor: sync and async http, ai21 http, ai21 client, resources

* test: update imports, create tests for async

* fix: base client

* fix: add pytest marker asyncio

* fix: add pytest asyncio to poetry

* fix: add delete to lib files, add examples, test examples

* fix: tests

* fix: fix stream, add stream tests, add readme

* fix: fix import on sm stub

* fix: fix async http client, fix tests

* fix: remove commented out code

* fix: CR comments

* fix: fix failing test

* fix: fix failing test

* fix: fix failing test

* fix: fix library test

* fix: cr comments

Author: miri-bar

README.md

@@ -167,10 +167,10 @@ client = AI21Client(
     api_key='my_api_key',
 )
 
+system = "You're a support engineer in a SaaS company"
 messages = [
-    # Could be a dict or a ChatMessage object
-    ChatMessage(content="Hello, this is a readme", role="user"),
-    ChatMessage(content="You are correct, how can I help you?", role="assistant"),
+    ChatMessage(content=system, role="system"),
+    ChatMessage(content="Hello, I need help with a signup process.", role="user"),
 ]
 
 chat_completions = client.chat.completions.create(
@@ -179,6 +179,42 @@ chat_completions = client.chat.completions.create(
 )
 ```
 
+### Async Usage
+
+You can use the `AsyncAI21Client` to make asynchronous requests.
+There is no difference between the sync and the async client in terms of usage.
+
+```python
+import asyncio
+
+from ai21 import AsyncAI21Client
+from ai21.models.chat import ChatMessage
+
+system = "You're a support engineer in a SaaS company"
+messages = [
+    ChatMessage(content=system, role="system"),
+    ChatMessage(content="Hello, I need help with a signup process.", role="user"),
+]
+
+client = AsyncAI21Client(
+   # defaults to os.enviorn.get('AI21_API_KEY')
+    api_key='my_api_key',
+)
+
+
+async def main():
+    response = await client.chat.completions.create(
+        messages=messages,
+        model="jamba-instruct-preview",
+    )
+
+    print(response)
+
+
+asyncio.run(main())
+
+```
+
 A more detailed example can be found [here](examples/studio/chat/chat_completions.py).
 
 ## Older Models Support Usage
@@ -260,6 +296,33 @@ for chunk in response:
 
 ```
 
+### Async Streaming
+
+```python
+import asyncio
+
+from ai21 import AsyncAI21Client
+from ai21.models.chat import ChatMessage
+
+messages = [ChatMessage(content="What is the meaning of life?", role="user")]
+
+client = AsyncAI21Client()
+
+
+async def main():
+    response = await client.chat.completions.create(
+        messages=messages,
+        model="jamba-instruct-preview",
+        stream=True,
+    )
+    async for chunk in response:
+        print(chunk.choices[0].delta.content, end="")
+
+
+asyncio.run(main())
+
+```
+
 ---
 
 ## More Models

ai21/__init__.py

@@ -2,6 +2,7 @@
 
 from ai21.ai21_env_config import AI21EnvConfig
 from ai21.clients.studio.ai21_client import AI21Client
+from ai21.clients.studio.async_ai21_client import AsyncAI21Client
 from ai21.errors import (
     AI21APIError,
     APITimeoutError,
@@ -53,6 +54,7 @@ def __getattr__(name: str) -> Any:
 __all__ = [
     "AI21EnvConfig",
     "AI21Client",
+    "AsyncAI21Client",
     "AI21APIError",
     "APITimeoutError",
     "AI21Error",

ai21/ai21_http_client/__init__.py

@@ -0,0 +1,65 @@
+from typing import Optional, Dict, Any, BinaryIO
+
+import httpx
+
+from ai21.http_client.http_client import HttpClient
+from ai21.ai21_http_client.base_ai21_http_client import BaseAI21HTTPClient
+
+
+class AI21HTTPClient(BaseAI21HTTPClient):
+    def __init__(
+        self,
+        *,
+        api_key: Optional[str] = None,
+        requires_api_key: bool = True,
+        base_url: Optional[str] = None,
+        api_version: Optional[str] = None,
+        headers: Optional[Dict[str, Any]] = None,
+        timeout_sec: Optional[int] = None,
+        num_retries: Optional[int] = None,
+        via: Optional[str] = None,
+        http_client: Optional[HttpClient] = None,
+    ):
+        super().__init__(
+            api_key=api_key,
+            requires_api_key=requires_api_key,
+            base_url=base_url,
+            api_version=api_version,
+            headers=headers,
+            timeout_sec=timeout_sec,
+            num_retries=num_retries,
+            via=via,
+        )
+
+        headers = self._build_headers(passed_headers=headers)
+        self._http_client = self._init_http_client(http_client=http_client, headers=headers)
+
+    def _init_http_client(self, http_client: Optional[HttpClient], headers: Dict[str, Any]) -> HttpClient:
+        if http_client is None:
+            return HttpClient(
+                timeout_sec=self._timeout_sec,
+                num_retries=self._num_retries,
+                headers=headers,
+            )
+
+        http_client.add_headers(headers)
+
+        return http_client
+
+    def execute_http_request(
+        self,
+        method: str,
+        path: str,
+        params: Optional[Dict] = None,
+        body: Optional[Dict] = None,
+        stream: bool = False,
+        files: Optional[Dict[str, BinaryIO]] = None,
+    ) -> httpx.Response:
+        return self._http_client.execute_http_request(
+            method=method,
+            url=f"{self._base_url}{path}",
+            params=params or {},
+            files=files,
+            stream=stream,
+            body=body or {},
+        )

ai21/ai21_http_client/ai21_http_client.py

@@ -0,0 +1,65 @@
+from typing import Optional, Dict, Any, BinaryIO
+
+import httpx
+
+from ai21.http_client.async_http_client import AsyncHttpClient
+from ai21.ai21_http_client.base_ai21_http_client import BaseAI21HTTPClient
+
+
+class AsyncAI21HTTPClient(BaseAI21HTTPClient):
+    def __init__(
+        self,
+        *,
+        api_key: Optional[str] = None,
+        requires_api_key: bool = True,
+        base_url: Optional[str] = None,
+        api_version: Optional[str] = None,
+        headers: Optional[Dict[str, Any]] = None,
+        timeout_sec: Optional[int] = None,
+        num_retries: Optional[int] = None,
+        via: Optional[str] = None,
+        http_client: Optional[AsyncHttpClient] = None,
+    ):
+        super().__init__(
+            api_key=api_key,
+            requires_api_key=requires_api_key,
+            base_url=base_url,
+            api_version=api_version,
+            headers=headers,
+            timeout_sec=timeout_sec,
+            num_retries=num_retries,
+            via=via,
+        )
+
+        headers = self._build_headers(passed_headers=headers)
+        self._http_client = self._init_http_client(http_client=http_client, headers=headers)
+
+    def _init_http_client(self, http_client: Optional[AsyncHttpClient], headers: Dict[str, Any]) -> AsyncHttpClient:
+        if http_client is None:
+            return AsyncHttpClient(
+                timeout_sec=self._timeout_sec,
+                num_retries=self._num_retries,
+                headers=headers,
+            )
+
+        http_client.add_headers(headers)
+
+        return http_client
+
+    async def execute_http_request(
+        self,
+        method: str,
+        path: str,
+        params: Optional[Dict] = None,
+        body: Optional[Dict] = None,
+        stream: bool = False,
+        files: Optional[Dict[str, BinaryIO]] = None,
+    ) -> httpx.Response:
+        return await self._http_client.execute_http_request(
+            method=method,
+            url=f"{self._base_url}{path}",
+            params=params or {},
+            files=files,
+            stream=stream,
+            body=body or {},
+        )

ai21/ai21_http_client/async_ai21_http_client.py

@@ -1,14 +1,21 @@
 import platform
-from typing import Optional, Dict, Any, BinaryIO
+from typing import Optional, Dict, Any, BinaryIO, TypeVar, Union
+from abc import ABC, abstractmethod
 
 import httpx
 
 from ai21.errors import MissingApiKeyError
-from ai21.http_client import HttpClient
+from ai21.http_client.http_client import HttpClient
+from ai21.http_client.async_http_client import AsyncHttpClient
 from ai21.version import VERSION
 
 
-class AI21HTTPClient:
+_HttpClientT = TypeVar("_HttpClientT", bound=Union[HttpClient, AsyncHttpClient])
+
+
+class BaseAI21HTTPClient(ABC):
+    _http_client: Optional[_HttpClientT] = None
+
     def __init__(
         self,
         *,
@@ -34,9 +41,6 @@ def __init__(
         self._num_retries = num_retries
         self._via = via
 
-        headers = self._build_headers(passed_headers=headers)
-        self._http_client = self._init_http_client(http_client=http_client, headers=headers)
-
     def _build_headers(self, passed_headers: Optional[Dict[str, Any]]) -> Dict[str, Any]:
         headers = {
             "Content-Type": "application/json",
@@ -51,18 +55,6 @@ def _build_headers(self, passed_headers: Optional[Dict[str, Any]]) -> Dict[str,
 
         return headers
 
-    def _init_http_client(self, http_client: Optional[HttpClient], headers: Dict[str, Any]) -> HttpClient:
-        if http_client is None:
-            return HttpClient(
-                timeout_sec=self._timeout_sec,
-                num_retries=self._num_retries,
-                headers=headers,
-            )
-
-        http_client.add_headers(headers)
-
-        return http_client
-
     def _build_user_agent(self) -> str:
         user_agent = (
             f"AI21 studio SDK {VERSION} Python {platform.python_version()} Operating System {platform.platform()}"
@@ -73,6 +65,11 @@ def _build_user_agent(self) -> str:
 
         return user_agent
 
+    @abstractmethod
+    def _init_http_client(self, http_client: Optional[HttpClient], headers: Dict[str, Any]) -> _HttpClientT:
+        pass
+
+    @abstractmethod
     def execute_http_request(
         self,
         method: str,
@@ -82,11 +79,4 @@ def execute_http_request(
         stream: bool = False,
         files: Optional[Dict[str, BinaryIO]] = None,
     ) -> httpx.Response:
-        return self._http_client.execute_http_request(
-            method=method,
-            url=f"{self._base_url}{path}",
-            params=params or {},
-            files=files,
-            stream=stream,
-            body=body or {},
-        )
+        pass

ai21/ai21_http_client.py renamed to ai21/ai21_http_client/base_ai21_http_client.py

@@ -7,7 +7,7 @@
 
 from ai21.logger import logger
 from ai21.errors import AccessDenied, NotFound, APITimeoutError
-from ai21.http_client import handle_non_success_response
+from ai21.http_client.base_http_client import handle_non_success_response
 
 _ERROR_MSG_TEMPLATE = (
     r"Received client error \((.*?)\) from primary with message \"(.*?)\". "

ai21/clients/bedrock/bedrock_session.py

@@ -1,4 +1,4 @@
-from abc import ABC
+from abc import ABC, abstractmethod
 from typing import Any, Dict
 
 from ai21.models import AnswerResponse
@@ -7,6 +7,7 @@
 class Answer(ABC):
     _module_name = "answer"
 
+    @abstractmethod
     def create(
         self,
         context: str,