From f95548909b8d38b9daab6f4b53a85d83906ea00d Mon Sep 17 00:00:00 2001
From: rafayaar <rafaymahmood1999@gmail.com>
Date: Wed, 10 Jan 2024 17:21:43 +0500
Subject: [PATCH 1/2] Code Updated according to new AsyncOpenAI

---
 mock_test.py                    |  9 ++++++---
 openai_multi_client/__init__.py | 16 +++++++++-------
 real_test.py                    |  5 ++++-
 3 files changed, 19 insertions(+), 11 deletions(-)

diff --git a/mock_test.py b/mock_test.py
index 70430bf..51f741f 100644
--- a/mock_test.py
+++ b/mock_test.py
@@ -1,5 +1,7 @@
 import asyncio
+from dotenv import load_dotenv
 
+load_dotenv()
 from openai_multi_client import OpenAIMultiClient, Payload, OpenAIMultiOrderedClient
 
 
@@ -12,7 +14,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 +37,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")
diff --git a/openai_multi_client/__init__.py b/openai_multi_client/__init__.py
index 1dd3675..e6f5c97 100644
--- a/openai_multi_client/__init__.py
+++ b/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}")
diff --git a/real_test.py b/real_test.py
index da09303..dbdd8ac 100644
--- a/real_test.py
+++ b/real_test.py
@@ -1,4 +1,7 @@
 import os
+from dotenv import load_dotenv
+
+load_dotenv()
 from openai_multi_client import OpenAIMultiClient, Payload
 
 
@@ -37,7 +40,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():

From b22dffacd7ac667d4d6c7ff39d3e4cf5edf56e15 Mon Sep 17 00:00:00 2001
From: rafayaar <rafaymahmood1999@gmail.com>
Date: Wed, 10 Jan 2024 17:29:14 +0500
Subject: [PATCH 2/2] Removing extra code

---
 mock_test.py | 3 ---
 real_test.py | 3 ---
 2 files changed, 6 deletions(-)

diff --git a/mock_test.py b/mock_test.py
index 51f741f..5911d86 100644
--- a/mock_test.py
+++ b/mock_test.py
@@ -1,7 +1,4 @@
 import asyncio
-from dotenv import load_dotenv
-
-load_dotenv()
 from openai_multi_client import OpenAIMultiClient, Payload, OpenAIMultiOrderedClient
 
 
diff --git a/real_test.py b/real_test.py
index dbdd8ac..335de48 100644
--- a/real_test.py
+++ b/real_test.py
@@ -1,7 +1,4 @@
 import os
-from dotenv import load_dotenv
-
-load_dotenv()
 from openai_multi_client import OpenAIMultiClient, Payload