diff options
Diffstat (limited to '')
| -rw-r--r-- | etc/unittest/__main__.py | 1 | ||||
| -rw-r--r-- | etc/unittest/async_client.py | 56 | ||||
| -rw-r--r-- | etc/unittest/integration.py | 10 |
3 files changed, 66 insertions, 1 deletions
diff --git a/etc/unittest/__main__.py b/etc/unittest/__main__.py index 3a459dba..351c2bb3 100644 --- a/etc/unittest/__main__.py +++ b/etc/unittest/__main__.py @@ -4,6 +4,7 @@ from .backend import * from .main import * from .model import * from .client import * +from .async_client import * from .include import * from .integration import * diff --git a/etc/unittest/async_client.py b/etc/unittest/async_client.py new file mode 100644 index 00000000..a49b90ed --- /dev/null +++ b/etc/unittest/async_client.py @@ -0,0 +1,56 @@ +import unittest + +from g4f.client import AsyncClient, ChatCompletion, ChatCompletionChunk +from .mocks import AsyncGeneratorProviderMock, ModelProviderMock, YieldProviderMock + +DEFAULT_MESSAGES = [{'role': 'user', 'content': 'Hello'}] + +class AsyncTestPassModel(unittest.IsolatedAsyncioTestCase): + + async def test_response(self): + client = AsyncClient(provider=AsyncGeneratorProviderMock) + response = await client.chat.completions.create(DEFAULT_MESSAGES, "") + self.assertIsInstance(response, ChatCompletion) + self.assertEqual("Mock", response.choices[0].message.content) + + async def test_pass_model(self): + client = AsyncClient(provider=ModelProviderMock) + response = await client.chat.completions.create(DEFAULT_MESSAGES, "Hello") + self.assertIsInstance(response, ChatCompletion) + self.assertEqual("Hello", response.choices[0].message.content) + + async def test_max_tokens(self): + client = AsyncClient(provider=YieldProviderMock) + messages = [{'role': 'user', 'content': chunk} for chunk in ["How ", "are ", "you", "?"]] + response = await client.chat.completions.create(messages, "Hello", max_tokens=1) + self.assertIsInstance(response, ChatCompletion) + self.assertEqual("How ", response.choices[0].message.content) + response = await client.chat.completions.create(messages, "Hello", max_tokens=2) + self.assertIsInstance(response, ChatCompletion) + self.assertEqual("How are ", response.choices[0].message.content) + + async def test_max_stream(self): + client = AsyncClient(provider=YieldProviderMock) + messages = [{'role': 'user', 'content': chunk} for chunk in ["How ", "are ", "you", "?"]] + response = client.chat.completions.create(messages, "Hello", stream=True) + async for chunk in response: + self.assertIsInstance(chunk, ChatCompletionChunk) + if chunk.choices[0].delta.content is not None: + 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 = [chunk async for chunk in response] + self.assertEqual(len(response), 3) + for chunk in response: + if chunk.choices[0].delta.content is not None: + self.assertEqual(chunk.choices[0].delta.content, "You ") + + async def test_stop(self): + client = AsyncClient(provider=YieldProviderMock) + messages = [{'role': 'user', 'content': chunk} for chunk in ["How ", "are ", "you", "?"]] + response = await 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()
\ No newline at end of file diff --git a/etc/unittest/integration.py b/etc/unittest/integration.py index f69d4770..d87a7f6b 100644 --- a/etc/unittest/integration.py +++ b/etc/unittest/integration.py @@ -8,7 +8,7 @@ except ImportError: has_nest_asyncio = False from g4f.client import Client, ChatCompletion -from g4f.Provider import Bing, OpenaiChat +from g4f.Provider import Bing, OpenaiChat, DuckDuckGo DEFAULT_MESSAGES = [{"role": "system", "content": 'Response in json, Example: {"success: true"}'}, {"role": "user", "content": "Say success true in json"}] @@ -19,11 +19,19 @@ class TestProviderIntegration(unittest.TestCase): self.skipTest("nest_asyncio is not installed") def test_bing(self): + self.skipTest("Not stable") client = Client(provider=Bing) response = client.chat.completions.create(DEFAULT_MESSAGES, "", response_format={"type": "json_object"}) self.assertIsInstance(response, ChatCompletion) self.assertIn("success", json.loads(response.choices[0].message.content)) + def test_duckduckgo(self): + self.skipTest("Not working") + client = Client(provider=DuckDuckGo) + response = client.chat.completions.create(DEFAULT_MESSAGES, "", response_format={"type": "json_object"}) + self.assertIsInstance(response, ChatCompletion) + self.assertIn("success", json.loads(response.choices[0].message.content)) + def test_openai(self): client = Client(provider=OpenaiChat) response = client.chat.completions.create(DEFAULT_MESSAGES, "", response_format={"type": "json_object"}) |