Async calls using new SDK by OpenAI

mock_test.py

@@ -1,5 +1,4 @@
 import asyncio
-
 from openai_multi_client import OpenAIMultiClient, Payload, OpenAIMultiOrderedClient
 
 
@@ -12,7 +11,7 @@ async def mock(payload: Payload):
         rand_fail = random.random()
         if rand_fail < 0.3:
             raise Exception("Mocked exception")
-        return {"response": f"mocked {payload.metadata['id']}"}
+        return {"response": f"mocked {payload.metadata.get('id')}"}
 
     if ordered:
         api = OpenAIMultiOrderedClient(custom_api=mock, max_retries=3, retry_multiplier=2)
@@ -35,9 +34,10 @@ def put_data():
         i += 1
         if response.failed:
             failed += 1
-            print(f"Failed {response.metadata['id']}: {i}/100")
+            print(f"Failed {response.metadata.get('id')}: {i}/100")
         else:
-            print(f"Got response {response.metadata['id']}: {i}/100")
+            print(f"Got response {response.metadata.get('id')}: {i}/100")
+
 
     print('*' * 20)
     print(f"Total failed: {failed}/100")

openai_multi_client/__init__.py

@@ -6,7 +6,9 @@
 
 from aioprocessing import AioJoinableQueue, AioQueue
 from tenacity import wait_random_exponential, stop_after_attempt, AsyncRetrying, RetryError
-import openai
+from openai import AsyncOpenAI
+
+aclient = AsyncOpenAI()
 
 logger = logging.getLogger(__name__)
 
@@ -76,17 +78,17 @@ async def _process_payload(self, payload: Payload) -> Payload:
         if self._mock_api:
             payload.response = await self._mock_api(payload)
         elif payload.endpoint == "completions":
-            payload.response = await openai.Completion.acreate(**payload.data)
+            payload.response = await aclient.completions.create(**payload.data)
         elif payload.endpoint == "chat.completions" or payload.endpoint == "chats":
-            payload.response = await openai.ChatCompletion.acreate(**payload.data)
+            payload.response = await aclient.chat.completions.create(**payload.data)
         elif payload.endpoint == "embeddings":
-            payload.response = await openai.Embedding.acreate(**payload.data)
+            payload.response = await aclient.embeddings.create(**payload.data)
         elif payload.endpoint == "edits":
-            payload.response = await openai.Edit.acreate(**payload.data)
+            payload.response = await aclient.edits.create(**payload.data)
         elif payload.endpoint == "images":
-            payload.response = await openai.Image.acreate(**payload.data)
+            payload.response = await aclient.images.generate(**payload.data)
         elif payload.endpoint == "fine-tunes":
-            payload.response = await openai.FineTune.acreate(**payload.data)
+            payload.response = await aclient.fine_tunes.create(**payload.data)
         else:
             raise ValueError(f"Unknown endpoint {payload.endpoint}")
         logger.debug(f"Processed {payload}")

real_test.py

@@ -37,7 +37,7 @@ def on_success(result: Payload):
         print(f"Failed {pid}")
     else:
         print(f"Got response for {pid}:",
-              result.response['choices'][0]['message']['content'])
+              result.response.choices[0].message.content)
 
 
 def test_callback():