asyncio.py 2.7 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273
  1. import asyncio
  2. try:
  3. import nest_asyncio
  4. has_nest_asyncio = True
  5. except:
  6. has_nest_asyncio = False
  7. import unittest
  8. import g4f
  9. from g4f import ChatCompletion
  10. from g4f.client import Client
  11. from .mocks import ProviderMock, AsyncProviderMock, AsyncGeneratorProviderMock
  12. DEFAULT_MESSAGES = [{'role': 'user', 'content': 'Hello'}]
  13. class TestChatCompletion(unittest.TestCase):
  14. async def run_exception(self):
  15. return ChatCompletion.create(g4f.models.default, DEFAULT_MESSAGES, AsyncProviderMock)
  16. def test_exception(self):
  17. if has_nest_asyncio:
  18. self.skipTest('has nest_asyncio')
  19. self.assertRaises(g4f.errors.NestAsyncioError, asyncio.run, self.run_exception())
  20. def test_create(self):
  21. result = ChatCompletion.create(g4f.models.default, DEFAULT_MESSAGES, AsyncProviderMock)
  22. self.assertEqual("Mock", result)
  23. def test_create_generator(self):
  24. result = ChatCompletion.create(g4f.models.default, DEFAULT_MESSAGES, AsyncGeneratorProviderMock)
  25. self.assertEqual("Mock", result)
  26. def test_await_callback(self):
  27. client = Client(provider=AsyncGeneratorProviderMock)
  28. response = client.chat.completions.create(DEFAULT_MESSAGES, "", max_tokens=0)
  29. self.assertEqual("Mock", response.choices[0].message.content)
  30. class TestChatCompletionAsync(unittest.IsolatedAsyncioTestCase):
  31. async def test_base(self):
  32. result = await ChatCompletion.create_async(g4f.models.default, DEFAULT_MESSAGES, ProviderMock)
  33. self.assertEqual("Mock",result)
  34. async def test_async(self):
  35. result = await ChatCompletion.create_async(g4f.models.default, DEFAULT_MESSAGES, AsyncProviderMock)
  36. self.assertEqual("Mock",result)
  37. async def test_create_generator(self):
  38. result = await ChatCompletion.create_async(g4f.models.default, DEFAULT_MESSAGES, AsyncGeneratorProviderMock)
  39. self.assertEqual("Mock",result)
  40. class TestChatCompletionNestAsync(unittest.IsolatedAsyncioTestCase):
  41. def setUp(self) -> None:
  42. if not has_nest_asyncio:
  43. self.skipTest('"nest_asyncio" not installed')
  44. nest_asyncio.apply()
  45. async def test_create(self):
  46. result = await ChatCompletion.create_async(g4f.models.default, DEFAULT_MESSAGES, ProviderMock)
  47. self.assertEqual("Mock",result)
  48. async def test_nested(self):
  49. result = ChatCompletion.create(g4f.models.default, DEFAULT_MESSAGES, AsyncProviderMock)
  50. self.assertEqual("Mock",result)
  51. async def test_nested_generator(self):
  52. result = ChatCompletion.create(g4f.models.default, DEFAULT_MESSAGES, AsyncGeneratorProviderMock)
  53. self.assertEqual("Mock",result)
  54. if __name__ == '__main__':
  55. unittest.main()