Add async support

README.md

@@ -216,6 +216,32 @@ image_resp = openai.Image.create(prompt="two dogs playing chess, oil painting",
 
 ```
 
+## Async API
+
+Async support is available in the API by prepending `a` to a network-bound method:
+
+```python
+import openai
+openai.api_key = "sk-..."  # supply your API key however you choose
+
+async def create_completion():
+    completion_resp = await openai.Completion.acreate(prompt="This is a test", engine="davinci")
+
+```
+
+To make async requests more efficient, you can pass in your own
+``aiohttp.ClientSession``, but you must manage the client session closing at the end 
+of your program/event loop:
+
+```python
+import openai
+from aiohttp import ClientSession
+
+openai.aiosession.set(ClientSession())
+# At the end of your program, close the http session
+await openai.aiosession.get().close()
+```
+
 See the [usage guide](https://beta.openai.com/docs/guides/images) for more details.
 
 ## Requirements

openai/__init__.py

@@ -3,7 +3,8 @@
 # Originally forked from the MIT-licensed Stripe Python bindings.
 
 import os
-from typing import Optional
+from contextvars import ContextVar
+from typing import Optional, TYPE_CHECKING
 
 from openai.api_resources import (
     Answer,
@@ -24,6 +25,9 @@
 )
 from openai.error import APIError, InvalidRequestError, OpenAIError
 
+if TYPE_CHECKING:
+    from aiohttp import ClientSession
+
 api_key = os.environ.get("OPENAI_API_KEY")
 # Path of a file with an API key, whose contents can change. Supercedes
 # `api_key` if set.  The main use case is volume-mounted Kubernetes secrets,
@@ -44,6 +48,10 @@
 debug = False
 log = None  # Set to either 'debug' or 'info', controls console logging
 
+aiosession: ContextVar[Optional["ClientSession"]] = ContextVar(
+    "aiohttp-session", default=None
+)
+
 __all__ = [
     "APIError",
     "Answer",

openai/api_requestor.py

@@ -6,6 +6,7 @@
 from typing import Dict, Iterator, Optional, Tuple, Union, overload
 from urllib.parse import urlencode, urlsplit, urlunsplit
 
+import aiohttp
 import requests
 from typing_extensions import Literal
 
@@ -44,6 +45,19 @@ def _requests_proxies_arg(proxy) -> Optional[Dict[str, str]]:
         )
 
 
+def _aiohttp_proxies_arg(proxy) -> Optional[str]:
+    """Returns a value suitable for the 'proxies' argument to 'aiohttp.ClientSession.request."""
+    if proxy is None:
+        return None
+    elif isinstance(proxy, str):
+        return proxy
+    elif isinstance(proxy, dict):
+        return proxy["https"] if "https" in proxy else proxy["http"]
+    else:
+        raise ValueError(
+            "'openai.proxy' must be specified as either a string URL or a dict with string URL under the https and/or http keys."
+        )
+
 def _make_session() -> requests.Session:
     if not openai.verify_ssl_certs:
         warnings.warn("verify_ssl_certs is ignored; openai always verifies.")
@@ -72,6 +86,18 @@ def parse_stream(rbody):
             yield line
 
 
+async def parse_stream_async(rbody: aiohttp.StreamReader):
+    async for line in rbody:
+        if line:
+            if line == b"data: [DONE]":
+                continue
+            if hasattr(line, "decode"):
+                line = line.decode("utf-8")
+            if line.startswith("data: "):
+                line = line[len("data: ") :]
+            yield line
+
+
 class APIRequestor:
     def __init__(
         self,
@@ -181,6 +207,29 @@ def request(
         resp, got_stream = self._interpret_response(result, stream)
         return resp, got_stream, self.api_key
 
+    async def arequest(
+        self,
+        method,
+        url,
+        params=None,
+        headers=None,
+        files=None,
+        stream: bool = False,
+        request_id: Optional[str] = None,
+        request_timeout: Optional[Union[float, Tuple[float, float]]] = None,
+    ) -> Tuple[Union[OpenAIResponse, Iterator[OpenAIResponse]], bool, str]:
+        result = await self.arequest_raw(
+            method.lower(),
+            url,
+            params=params,
+            supplied_headers=headers,
+            files=files,
+            request_id=request_id,
+            request_timeout=request_timeout,
+        )
+        resp, got_stream = await self._interpret_async_response(result, stream)
+        return resp, got_stream, self.api_key
+
     def handle_error_response(self, rbody, rcode, resp, rheaders, stream_error=False):
         try:
             error_data = resp["error"]
@@ -310,18 +359,9 @@ def _validate_headers(
 
         return headers
 
-    def request_raw(
-        self,
-        method,
-        url,
-        *,
-        params=None,
-        supplied_headers: Dict[str, str] = None,
-        files=None,
-        stream: bool = False,
-        request_id: Optional[str] = None,
-        request_timeout: Optional[Union[float, Tuple[float, float]]] = None,
-    ) -> requests.Response:
+    def _prepare_request_raw(
+        self, url, supplied_headers, method, params, files, request_id: Optional[str],
+    ) -> Tuple[str, Dict[str, str], Optional[bytes]]:
         abs_url = "%s%s" % (self.api_base, url)
         headers = self._validate_headers(supplied_headers)
 
@@ -350,6 +390,24 @@ def request_raw(
         util.log_info("Request to OpenAI API", method=method, path=abs_url)
         util.log_debug("Post details", data=data, api_version=self.api_version)
 
+        return abs_url, headers, data
+
+    def request_raw(
+        self,
+        method,
+        url,
+        *,
+        params=None,
+        supplied_headers: Dict[str, str] = None,
+        files=None,
+        stream: bool = False,
+        request_id: Optional[str] = None,
+        request_timeout: Optional[Union[float, Tuple[float, float]]] = None,
+    ) -> requests.Response:
+        abs_url, headers, data = self._prepare_request_raw(
+            url, supplied_headers, method, params, files, request_id
+        )
+
         if not hasattr(_thread_context, "session"):
             _thread_context.session = _make_session()
         try:
@@ -380,6 +438,74 @@ def request_raw(
             )
         return result
 
+    async def arequest_raw(
+        self,
+        method,
+        url,
+        *,
+        params=None,
+        supplied_headers: Dict[str, str] = None,
+        files=None,
+        request_id: Optional[str] = None,
+        request_timeout: Optional[Union[float, Tuple[float, float]]] = None,
+    ) -> aiohttp.ClientResponse:
+        abs_url, headers, data = self._prepare_request_raw(
+            url, supplied_headers, method, params, files, request_id
+        )
+
+        timeout = aiohttp.ClientTimeout(
+            total=request_timeout if request_timeout else TIMEOUT_SECS
+        )
+        user_set_session = openai.aiosession.get()
+
+        if files:
+            data, content_type = requests.models.RequestEncodingMixin._encode_files(
+                files, data
+            )
+            headers["Content-Type"] = content_type
+
+        result = None
+        try:
+            if user_set_session:
+                result = await user_set_session.request(
+                    method,
+                    abs_url,
+                    headers=headers,
+                    data=data,
+                    proxy=_aiohttp_proxies_arg(openai.proxy),
+                    timeout=timeout,
+                )
+            else:
+                async with aiohttp.ClientSession() as session:
+                    result = await session.request(
+                        method,
+                        abs_url,
+                        headers=headers,
+                        data=data,
+                        proxy=_aiohttp_proxies_arg(openai.proxy),
+                        timeout=timeout,
+                    )
+            util.log_info(
+                "OpenAI API response",
+                path=abs_url,
+                response_code=result.status,
+                processing_ms=result.headers.get("OpenAI-Processing-Ms"),
+                request_id=result.headers.get("X-Request-Id"),
+            )
+            # Don't read the whole stream for debug logging unless necessary.
+            if openai.log == "debug":
+                util.log_debug(
+                    "API response body", body=result.content, headers=result.headers
+                )
+            return result
+        except aiohttp.ServerTimeoutError as e:
+            raise error.Timeout("Request timed out") from e
+        except aiohttp.ClientError as e:
+            raise error.APIConnectionError("Error communicating with OpenAI") from e
+        finally:
+            if result and not result.closed:
+                result.close()
+
     def _interpret_response(
         self, result: requests.Response, stream: bool
     ) -> Tuple[Union[OpenAIResponse, Iterator[OpenAIResponse]], bool]:
@@ -399,6 +525,29 @@ def _interpret_response(
                 False,
             )
 
+    async def _interpret_async_response(
+        self, result: aiohttp.ClientResponse, stream: bool
+    ) -> Tuple[Union[OpenAIResponse, Iterator[OpenAIResponse]], bool]:
+        """Returns the response(s) and a bool indicating whether it is a stream."""
+        if stream and "text/event-stream" in result.headers.get("Content-Type", ""):
+            return (
+                self._interpret_response_line(
+                    line, result.status, result.headers, stream=True
+                )
+                async for line in parse_stream_async(result.content)
+            ), True
+        else:
+            try:
+                await result.read()
+            except aiohttp.ClientError as e:
+                print(result.content, e)
+            return (
+                self._interpret_response_line(
+                    await result.read(), result.status, result.headers, stream=False
+                ),
+                False,
+            )
+
     def _interpret_response_line(
         self, rbody, rcode, rheaders, stream: bool
     ) -> OpenAIResponse:

openai/api_resources/abstract/api_resource.py

@@ -20,6 +20,13 @@ def retrieve(
         instance.refresh(request_id=request_id, request_timeout=request_timeout)
         return instance
 
+    @classmethod
+    def aretrieve(
+        cls, id, api_key=None, request_id=None, request_timeout=None, **params
+    ):
+        instance = cls(id, api_key, **params)
+        return instance.arefresh(request_id=request_id, request_timeout=request_timeout)
+
     def refresh(self, request_id=None, request_timeout=None):
         self.refresh_from(
             self.request(
@@ -31,6 +38,17 @@ def refresh(self, request_id=None, request_timeout=None):
         )
         return self
 
+    async def arefresh(self, request_id=None, request_timeout=None):
+        self.refresh_from(
+            await self.arequest(
+                "get",
+                self.instance_url(operation="refresh"),
+                request_id=request_id,
+                request_timeout=request_timeout,
+            )
+        )
+        return self
+
     @classmethod
     def class_url(cls):
         if cls == APIResource:
@@ -116,6 +134,31 @@ def _static_request(
             response, api_key, api_version, organization
         )
 
+    @classmethod
+    async def _astatic_request(
+        cls,
+        method_,
+        url_,
+        api_key=None,
+        api_base=None,
+        api_type=None,
+        request_id=None,
+        api_version=None,
+        organization=None,
+        **params,
+    ):
+        requestor = api_requestor.APIRequestor(
+            api_key,
+            api_version=api_version,
+            organization=organization,
+            api_base=api_base,
+            api_type=api_type,
+        )
+        response, _, api_key = await requestor.arequest(
+            method_, url_, params, request_id=request_id
+        )
+        return response
+
     @classmethod
     def _get_api_type_and_version(
         cls, api_type: Optional[str] = None, api_version: Optional[str] = None