Add async API for completion and completion_stream

.gitignore

@@ -1,6 +1,7 @@
 .env
 __pycache__
 .DS_Store
+.idea/
 **/.DS_Store
 
 anthropic.egg-info/

anthropic/api.py

@@ -1,13 +1,46 @@
-from typing import Dict, Iterator, Optional, Tuple, Union
+from typing import Dict, AsyncIterator, Iterator, Optional, Tuple, NamedTuple, Union
+import aiohttp
 import requests
 import requests.adapters
 import urllib.parse
 import json
+
+from aiohttp import ClientResponse
+
 from . import constants
 
+
 class ApiException(Exception):
     pass
 
+
+class Request(NamedTuple):
+    method: str
+    url: str
+    headers: Optional[Dict[str, str]]
+    data: bytes
+    stream: bool
+    timeout: Optional[Union[float, Tuple[float, float]]]
+
+
+def _process_request_error(
+    method: str, result: Union[requests.Response, ClientResponse]
+):
+    status_code = (
+        result.status if isinstance(result, ClientResponse) else result.status_code
+    )
+    if status_code != 200:
+        content = result.content.decode("utf-8")
+        try:
+            formatted_content = json.loads(content)
+        except json.decoder.JSONDecodeError:
+            formatted_content = content
+        raise ApiException(
+            f"{method} request failed with status code: {result.status_code}",
+            formatted_content,
+        )
+
+
 class Client:
     def __init__(
         self,
@@ -33,31 +66,28 @@ def _setup_session(self) -> requests.Session:
         )
         return self._session
 
-    def _request_raw(
+    def _request_params(
         self,
+        headers: Optional[Dict[str, str]],
         method: str,
-        path: str,
         params: dict,
-        headers: Optional[Dict[str, str]] = None,
-        request_timeout: Optional[Union[float, Tuple[float, float]]] = None,
-    ) -> requests.Response:
+        path: str,
+        request_timeout: Optional[Union[float, Tuple[float, float]]],
+    ) -> Request:
         method = method.lower()
-
         abs_url = urllib.parse.urljoin(self.api_url, path)
         final_headers = {
             "Accept": "application/json",
             "Client": constants.ANTHROPIC_CLIENT_VERSION,
             "X-API-Key": self.api_key,
             **(headers or {}),
         }
-
         if params.get("disable_checks"):
             del params["disable_checks"]
         else:
             # NOTE: disabling_checks can lead to very poor sampling quality from our API.
             # _Please_ read the docs on "Claude instructions when using the API" before disabling this
             _validate_prompt(params["prompt"])
-
         data = None
         if params:
             if method in {"get"}:
@@ -70,31 +100,100 @@ def _request_raw(
                 final_headers["Content-Type"] = "application/json"
             else:
                 raise ValueError(f"Unrecognized method: {method}")
-
         # If we're requesting a stream from the server, let's tell requests to expect the same
         stream = params.get("stream", None)
-        result = self._session.request(
+        return Request(
             method,
             abs_url,
-            headers=final_headers,
-            data=data,
-            stream=stream,
-            timeout=request_timeout
-            if request_timeout
-            else self.default_request_timeout,
+            final_headers,
+            data,
+            stream,
+            request_timeout or self.default_request_timeout,
         )
-        if result.status_code != 200:
-            content = result.content.decode("utf-8")
-            try:
-                formatted_content = json.loads(content)
-            except json.decoder.JSONDecodeError:
-                formatted_content = content
-            raise ApiException(
-                f'{method} request failed with status code: {result.status_code}',
-                formatted_content
-            )
+
+    def _request_raw(
+        self,
+        method: str,
+        path: str,
+        params: dict,
+        headers: Optional[Dict[str, str]] = None,
+        request_timeout: Optional[Union[float, Tuple[float, float]]] = None,
+    ) -> requests.Response:
+        request = self._request_params(headers, method, params, path, request_timeout)
+        result = self._session.request(
+            request.method,
+            request.url,
+            headers=request.headers,
+            data=request.data,
+            stream=request.stream,
+            timeout=request.timeout,
+        )
+
+        _process_request_error(method, result)
         return result
 
+    async def _arequest_as_json(
+        self,
+        method: str,
+        path: str,
+        params: dict,
+        headers: Optional[Dict[str, str]] = None,
+        request_timeout: Optional[Union[float, Tuple[float, float]]] = None,
+    ) -> dict:
+        request = self._request_params(headers, method, params, path, request_timeout)
+        async with aiohttp.ClientSession() as session:
+            async with session.request(
+                request.method,
+                request.url,
+                headers=request.headers,
+                data=request.data,
+                timeout=request.timeout,
+            ) as result:
+                _process_request_error(method, result)
+                content = await result.text()
+                json_body = json.loads(content)
+                return json_body
+
+    async def _arequest_as_stream(
+        self,
+        method: str,
+        path: str,
+        params: dict,
+        headers: Optional[Dict[str, str]] = None,
+        request_timeout: Optional[Union[float, Tuple[float, float]]] = None,
+    ) -> AsyncIterator[dict]:
+        request = self._request_params(headers, method, params, path, request_timeout)
+        awaiting_ping_data = False
+        async with aiohttp.ClientSession() as session:
+            async with session.request(
+                request.method,
+                request.url,
+                headers=request.headers,
+                data=request.data,
+                timeout=request.timeout,
+            ) as result:
+                _process_request_error(method, result)
+                async for line in result.content:
+                    line = line.strip()
+                    if not line:
+                        continue
+                    if line == b"event: ping":
+                        awaiting_ping_data = True
+                        continue
+                    if awaiting_ping_data:
+                        awaiting_ping_data = False
+                        continue
+
+                    if line == b"data: [DONE]":
+                        continue
+
+                    line = line.decode("utf-8")
+
+                    prefix = "data: "
+                    if line.startswith(prefix):
+                        line = line[len(prefix) :]
+                    yield json.loads(line)
+
     def _request_as_json(self, *args, **kwargs) -> dict:
         result = self._request_raw(*args, **kwargs)
         content = result.content.decode("utf-8")
@@ -143,6 +242,22 @@ def completion(self, **kwargs) -> dict:
             params=kwargs,
         )
 
+    async def acompletion(self, **kwargs) -> dict:
+        return await self._arequest_as_json(
+            "post",
+            "/v1/complete",
+            params=kwargs,
+        )
+
+    async def acompletion_stream(self, **kwargs) -> AsyncIterator[dict]:
+        new_kwargs = {"stream": True, **kwargs}
+        return self._arequest_as_stream(
+            "post",
+            "/v1/complete",
+            params=new_kwargs,
+        )
+
+
 def _validate_prompt(prompt: str) -> None:
     if not prompt.startswith(constants.HUMAN_PROMPT):
         raise ValueError(

examples/basic_async.py

@@ -0,0 +1,18 @@
+import asyncio
+import os
+import anthropic
+
+
+async def main(max_tokens_to_sample: int = 100):
+    c = anthropic.Client(os.environ["ANTHROPIC_API_KEY"])
+    resp = await c.acompletion(
+        prompt=f"{anthropic.HUMAN_PROMPT} How many toes do dogs have?{anthropic.AI_PROMPT}",
+        stop_sequences=[anthropic.HUMAN_PROMPT],
+        model="claude-v1",
+        max_tokens_to_sample=max_tokens_to_sample,
+    )
+    print(resp)
+
+
+if __name__ == "__main__":
+    asyncio.run(main())

examples/basic_async_stream.py

@@ -0,0 +1,21 @@
+import anthropic
+import asyncio
+import os
+
+
+async def main(max_tokens_to_sample: int = 200):
+    c = anthropic.Client(os.environ["ANTHROPIC_API_KEY"])
+
+    response = await c.acompletion_stream(
+        prompt=f"{anthropic.HUMAN_PROMPT} How many toes do dogs have?{anthropic.AI_PROMPT}",
+        stop_sequences=[anthropic.HUMAN_PROMPT],
+        max_tokens_to_sample=max_tokens_to_sample,
+        model="claude-v1",
+        stream=True,
+    )
+    async for data in response:
+        print(data)
+
+
+if __name__ == "__main__":
+    asyncio.run(main())

pyproject.toml

@@ -25,6 +25,7 @@ dependencies = [
     'importlib-metadata; python_version<"3.8"',
     "requests",
     "tokenizers",
+    "aiohttp",
 ]
 
 [project.urls]