kaizenkodo/ChatMastermind
2
0
Fork 0
Code Issues 2 Pull Requests Packages Projects Releases Wiki Activity

Switch to current version of OpenAI.

Browse Source
This commit is contained in:
Oleksandr Kozachuk 2024-03-30 14:17:38 +01:00
parent 5d1bb1f9e4
commit 9a957a89ac
2 changed files with 33 additions and 49 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

32
chatmastermind/ais/openai.py
View File

@ -47,12 +47,12 @@ class OpenAIAnswer:
self.finished = True self.finished = True
if not self.finished: if not self.finished:
found_choice = False found_choice = False
for choice in chunk['choices']: for choice in chunk.choices:
if not choice['finish_reason']: if not choice.finish_reason:
self.streams[choice['index']].data.append(choice['delta']['content']) self.streams[choice.index].data.append(choice.delta.content)
self.tokens.completion += len(self.encoding.encode(choice['delta']['content'])) self.tokens.completion += len(self.encoding.encode(choice.delta.content))
self.tokens.total = self.tokens.prompt + self.tokens.completion self.tokens.total = self.tokens.prompt + self.tokens.completion
if choice['index'] == self.idx: if choice.index == self.idx:
found_choice = True found_choice = True
if not found_choice: if not found_choice:
return False return False
@ -68,7 +68,10 @@ class OpenAI(AI):
self.ID = config.ID self.ID = config.ID
self.name = config.name self.name = config.name
self.config = config 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, def request(self,
question: Message, question: Message,
@ -83,7 +86,7 @@ class OpenAI(AI):
self.encoding = tiktoken.encoding_for_model(self.config.model) self.encoding = tiktoken.encoding_for_model(self.config.model)
oai_chat, prompt_tokens = self.openai_chat(chat, self.config.system, question) oai_chat, prompt_tokens = self.openai_chat(chat, self.config.system, question)
tokens: Tokens = Tokens(prompt_tokens, 0, prompt_tokens) tokens: Tokens = Tokens(prompt_tokens, 0, prompt_tokens)
response = openai.ChatCompletion.create( response = self._completions(
model=self.config.model, model=self.config.model,
messages=oai_chat, messages=oai_chat,
temperature=self.config.temperature, temperature=self.config.temperature,
@ -114,9 +117,8 @@ class OpenAI(AI):
Return all models supported by this AI. Return all models supported by this AI.
""" """
ret = [] ret = []
for engine in sorted(openai.Engine.list()['data'], key=lambda x: x['id']): for engine in sorted(self.client.models.list().data, key=lambda x: x.id):
if engine['ready']: ret.append(engine.id)
ret.append(engine['id'])
ret.sort() ret.sort()
return ret return ret
@ -124,14 +126,8 @@ class OpenAI(AI):
""" """
Print all models supported by the current AI. Print all models supported by the current AI.
""" """
not_ready = [] for model in self.models():
for engine in sorted(openai.Engine.list()['data'], key=lambda x: x['id']): print(model)
if engine['ready']:
print(engine['id'])
else:
not_ready.append(engine['id'])
if len(not_ready) > 0:
print('\nNot ready: ' + ', '.join(not_ready))
def openai_chat(self, chat: Chat, system: str, def openai_chat(self, chat: Chat, system: str,
question: Optional[Message] = None) -> tuple[ChatType, int]: question: Optional[Message] = None) -> tuple[ChatType, int]:

50
tests/test_ais_openai.py
View File

@ -9,43 +9,31 @@ from chatmastermind.configuration import OpenAIConfig
class OpenAITest(unittest.TestCase): 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: def test_request(self, mock_create: mock.MagicMock) -> None:
# Create a test instance of OpenAI # Create a test instance of OpenAI
config = OpenAIConfig() config = OpenAIConfig()
openai = OpenAI(config) openai = OpenAI(config)
# Set up the mock response from openai.ChatCompletion.create # Set up the mock response from openai.ChatCompletion.create
mock_chunk1 = { class mock_obj:
'choices': [ pass
{ mock_chunk1 = mock_obj()
'index': 0, mock_chunk1.choices = [mock_obj(), mock_obj()] # type: ignore
'delta': { mock_chunk1.choices[0].index = 0 # type: ignore
'content': 'Answer 1' mock_chunk1.choices[0].delta = mock_obj() # type: ignore
}, mock_chunk1.choices[0].delta.content = 'Answer 1' # type: ignore
'finish_reason': None 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
'index': 1, mock_chunk1.choices[1].delta.content = 'Answer 2' # type: ignore
'delta': { mock_chunk1.choices[1].finish_reason = None # type: ignore
'content': 'Answer 2' mock_chunk2 = mock_obj()
}, mock_chunk2.choices = [mock_obj(), mock_obj()] # type: ignore
'finish_reason': None 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_chunk2 = {
'choices': [
{
'index': 0,
'finish_reason': 'stop'
},
{
'index': 1,
'finish_reason': 'stop'
}
],
}
mock_create.return_value = iter([mock_chunk1, mock_chunk2]) mock_create.return_value = iter([mock_chunk1, mock_chunk2])
# Create test data # Create test data