94 lines
3.3 KiB
Python
94 lines
3.3 KiB
Python
import unittest
|
|
from unittest import mock
|
|
from chatmastermind.ais.openai import OpenAI
|
|
from chatmastermind.message import Message, Question, Answer
|
|
from chatmastermind.chat import Chat
|
|
from chatmastermind.ai import AIResponse, Tokens
|
|
from chatmastermind.configuration import OpenAIConfig
|
|
|
|
|
|
class OpenAITest(unittest.TestCase):
|
|
|
|
@mock.patch('openai.ChatCompletion.create')
|
|
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'
|
|
}
|
|
],
|
|
}
|
|
mock_create.return_value = iter([mock_chunk1, mock_chunk2])
|
|
|
|
# Create test data
|
|
question = Message(Question('Question'))
|
|
chat = Chat([
|
|
Message(Question('Question 1'), answer=Answer('Answer 1')),
|
|
Message(Question('Question 2'), answer=Answer('Answer 2')),
|
|
# add message without an answer -> expect to be skipped
|
|
Message(Question('Question 3'))
|
|
])
|
|
|
|
# Make the request
|
|
response = openai.request(question, chat, num_answers=2)
|
|
|
|
# Assert the AIResponse
|
|
self.assertIsInstance(response, AIResponse)
|
|
self.assertEqual(len(response.messages), 2)
|
|
self.assertEqual(response.messages[0].answer, 'Answer 1')
|
|
self.assertEqual(response.messages[1].answer, 'Answer 2')
|
|
self.assertIsNotNone(response.tokens)
|
|
self.assertIsInstance(response.tokens, Tokens)
|
|
assert response.tokens
|
|
self.assertEqual(response.tokens.prompt, 53)
|
|
self.assertEqual(response.tokens.completion, 6)
|
|
self.assertEqual(response.tokens.total, 59)
|
|
|
|
# Assert the mock call to openai.ChatCompletion.create
|
|
mock_create.assert_called_once_with(
|
|
model=f'{config.model}',
|
|
messages=[
|
|
{'role': 'system', 'content': f'{config.system}'},
|
|
{'role': 'user', 'content': 'Question 1'},
|
|
{'role': 'assistant', 'content': 'Answer 1'},
|
|
{'role': 'user', 'content': 'Question 2'},
|
|
{'role': 'assistant', 'content': 'Answer 2'},
|
|
{'role': 'user', 'content': 'Question'}
|
|
],
|
|
temperature=config.temperature,
|
|
max_tokens=config.max_tokens,
|
|
top_p=config.top_p,
|
|
n=2,
|
|
stream=True,
|
|
frequency_penalty=config.frequency_penalty,
|
|
presence_penalty=config.presence_penalty
|
|
)
|