From 9a957a89ace8589a9e2b68793220872d2f7f548b Mon Sep 17 00:00:00 2001
From: Oleksandr Kozachuk <ddeus.lp@mailnull.com>
Date: Sat, 30 Mar 2024 14:17:38 +0100
Subject: [PATCH] Switch to current version of OpenAI.

---
 chatmastermind/ais/openai.py | 32 ++++++++++-------------
 tests/test_ais_openai.py     | 50 ++++++++++++++----------------------
 2 files changed, 33 insertions(+), 49 deletions(-)

diff --git a/chatmastermind/ais/openai.py b/chatmastermind/ais/openai.py
index 7e28032..21f14f2 100644
--- a/chatmastermind/ais/openai.py
+++ b/chatmastermind/ais/openai.py
@@ -47,12 +47,12 @@ class OpenAIAnswer:
             self.finished = True
         if not self.finished:
             found_choice = False
-            for choice in chunk['choices']:
-                if not choice['finish_reason']:
-                    self.streams[choice['index']].data.append(choice['delta']['content'])
-                    self.tokens.completion += len(self.encoding.encode(choice['delta']['content']))
+            for choice in chunk.choices:
+                if not choice.finish_reason:
+                    self.streams[choice.index].data.append(choice.delta.content)
+                    self.tokens.completion += len(self.encoding.encode(choice.delta.content))
                     self.tokens.total = self.tokens.prompt + self.tokens.completion
-                if choice['index'] == self.idx:
+                if choice.index == self.idx:
                     found_choice = True
             if not found_choice:
                 return False
@@ -68,7 +68,10 @@ class OpenAI(AI):
         self.ID = config.ID
         self.name = config.name
         self.config = config
-        openai.api_key = self.config.api_key
+        self.client = openai.OpenAI(api_key=self.config.api_key)
+
+    def _completions(self, *args, **kw):  # type: ignore
+        return self.client.chat.completions.create(*args, **kw)
 
     def request(self,
                 question: Message,
@@ -83,7 +86,7 @@ class OpenAI(AI):
         self.encoding = tiktoken.encoding_for_model(self.config.model)
         oai_chat, prompt_tokens = self.openai_chat(chat, self.config.system, question)
         tokens: Tokens = Tokens(prompt_tokens, 0, prompt_tokens)
-        response = openai.ChatCompletion.create(
+        response = self._completions(
             model=self.config.model,
             messages=oai_chat,
             temperature=self.config.temperature,
@@ -114,9 +117,8 @@ class OpenAI(AI):
         Return all models supported by this AI.
         """
         ret = []
-        for engine in sorted(openai.Engine.list()['data'], key=lambda x: x['id']):
-            if engine['ready']:
-                ret.append(engine['id'])
+        for engine in sorted(self.client.models.list().data, key=lambda x: x.id):
+            ret.append(engine.id)
         ret.sort()
         return ret
 
@@ -124,14 +126,8 @@ class OpenAI(AI):
         """
         Print all models supported by the current AI.
         """
-        not_ready = []
-        for engine in sorted(openai.Engine.list()['data'], key=lambda x: x['id']):
-            if engine['ready']:
-                print(engine['id'])
-            else:
-                not_ready.append(engine['id'])
-        if len(not_ready) > 0:
-            print('\nNot ready: ' + ', '.join(not_ready))
+        for model in self.models():
+            print(model)
 
     def openai_chat(self, chat: Chat, system: str,
                     question: Optional[Message] = None) -> tuple[ChatType, int]:
diff --git a/tests/test_ais_openai.py b/tests/test_ais_openai.py
index eab84e6..d1e1a4c 100644
--- a/tests/test_ais_openai.py
+++ b/tests/test_ais_openai.py
@@ -9,43 +9,31 @@ from chatmastermind.configuration import OpenAIConfig
 
 class OpenAITest(unittest.TestCase):
 
-    @mock.patch('openai.ChatCompletion.create')
+    @mock.patch('chatmastermind.ais.openai.OpenAI._completions')
     def test_request(self, mock_create: mock.MagicMock) -> None:
         # Create a test instance of OpenAI
         config = OpenAIConfig()
         openai = OpenAI(config)
 
         # Set up the mock response from openai.ChatCompletion.create
-        mock_chunk1 = {
-            'choices': [
-                {
-                    'index': 0,
-                    'delta': {
-                        'content': 'Answer 1'
-                    },
-                    'finish_reason': None
-                },
-                {
-                    'index': 1,
-                    'delta': {
-                        'content': 'Answer 2'
-                    },
-                    'finish_reason': None
-                }
-            ],
-        }
-        mock_chunk2 = {
-            'choices': [
-                {
-                    'index': 0,
-                    'finish_reason': 'stop'
-                },
-                {
-                    'index': 1,
-                    'finish_reason': 'stop'
-                }
-            ],
-        }
+        class mock_obj:
+            pass
+        mock_chunk1 = mock_obj()
+        mock_chunk1.choices = [mock_obj(), mock_obj()]  # type: ignore
+        mock_chunk1.choices[0].index = 0  # type: ignore
+        mock_chunk1.choices[0].delta = mock_obj()  # type: ignore
+        mock_chunk1.choices[0].delta.content = 'Answer 1'  # type: ignore
+        mock_chunk1.choices[0].finish_reason = None  # type: ignore
+        mock_chunk1.choices[1].index = 1  # type: ignore
+        mock_chunk1.choices[1].delta = mock_obj()  # type: ignore
+        mock_chunk1.choices[1].delta.content = 'Answer 2'  # type: ignore
+        mock_chunk1.choices[1].finish_reason = None  # type: ignore
+        mock_chunk2 = mock_obj()
+        mock_chunk2.choices = [mock_obj(), mock_obj()]  # type: ignore
+        mock_chunk2.choices[0].index = 0  # type: ignore
+        mock_chunk2.choices[0].finish_reason = 'stop'  # type: ignore
+        mock_chunk2.choices[1].index = 1  # type: ignore
+        mock_chunk2.choices[1].finish_reason = 'stop'  # type: ignore
         mock_create.return_value = iter([mock_chunk1, mock_chunk2])
 
         # Create test data