import unittest from g4f.client import Client, ChatCompletion, ChatCompletionChunk from .mocks import AsyncGeneratorProviderMock, ModelProviderMock, YieldProviderMock DEFAULT_MESSAGES = [{'role': 'user', 'content': 'Hello'}] class TestPassModel(unittest.TestCase): def test_response(self): client = Client(provider=AsyncGeneratorProviderMock) response = client.chat.completions.create(DEFAULT_MESSAGES, "") self.assertIsInstance(response, ChatCompletion) self.assertEqual("Mock", response.choices[0].message.content) def test_pass_model(self): client = Client(provider=ModelProviderMock) response = client.chat.completions.create(DEFAULT_MESSAGES, "Hello") self.assertIsInstance(response, ChatCompletion) self.assertEqual("Hello", response.choices[0].message.content) def test_max_tokens(self): client = Client(provider=YieldProviderMock) messages = [{'role': 'user', 'content': chunk} for chunk in ["How ", "are ", "you", "?"]] response = client.chat.completions.create(messages, "Hello", max_tokens=1) self.assertIsInstance(response, ChatCompletion) self.assertEqual("How ", response.choices[0].message.content) response = client.chat.completions.create(messages, "Hello", max_tokens=2) self.assertIsInstance(response, ChatCompletion) self.assertEqual("How are ", response.choices[0].message.content) def test_max_stream(self): client = Client(provider=YieldProviderMock) messages = [{'role': 'user', 'content': chunk} for chunk in ["How ", "are ", "you", "?"]] response = client.chat.completions.create(messages, "Hello", stream=True) for chunk in response: self.assertIsInstance(chunk, ChatCompletionChunk) self.assertIsInstance(chunk.choices[0].delta.content, str) messages = [{'role': 'user', 'content': chunk} for chunk in ["You ", "You ", "Other", "?"]] response = client.chat.completions.create(messages, "Hello", stream=True, max_tokens=2) response = list(response) self.assertEqual(len(response), 2) for chunk in response: self.assertEqual(chunk.choices[0].delta.content, "You ") def test_stop(self): client = Client(provider=YieldProviderMock) messages = [{'role': 'user', 'content': chunk} for chunk in ["How ", "are ", "you", "?"]] response = client.chat.completions.create(messages, "Hello", stop=["and"]) self.assertIsInstance(response, ChatCompletion) self.assertEqual("How are you?", response.choices[0].message.content) if __name__ == '__main__': unittest.main()