diff options
| -rw-r--r-- | etc/unittest/__main__.py | 1 | ||||
| -rw-r--r-- | etc/unittest/main.py | 8 | ||||
| -rw-r--r-- | etc/unittest/mocks.py | 10 | ||||
| -rw-r--r-- | etc/unittest/model.py | 27 | ||||
| -rw-r--r-- | g4f/__init__.py | 9 |
5 files changed, 45 insertions, 10 deletions
diff --git a/etc/unittest/__main__.py b/etc/unittest/__main__.py index 243c56b2..a133343e 100644 --- a/etc/unittest/__main__.py +++ b/etc/unittest/__main__.py @@ -2,5 +2,6 @@ import unittest from .asyncio import * from .backend import * from .main import * +from .model import * unittest.main()
\ No newline at end of file diff --git a/etc/unittest/main.py b/etc/unittest/main.py index 9b9ca011..f5eb5138 100644 --- a/etc/unittest/main.py +++ b/etc/unittest/main.py @@ -1,4 +1,3 @@ -from .include import DEFAULT_MESSAGES import unittest import asyncio import g4f @@ -6,6 +5,8 @@ from g4f import ChatCompletion, get_last_provider from g4f.Provider import RetryProvider from .mocks import ProviderMock +DEFAULT_MESSAGES = [{'role': 'user', 'content': 'Hello'}] + class NoTestChatCompletion(unittest.TestCase): def no_test_create_default(self): @@ -31,7 +32,4 @@ class TestGetLastProvider(unittest.TestCase): def test_get_last_provider_async(self): coroutine = ChatCompletion.create_async(g4f.models.default, DEFAULT_MESSAGES, ProviderMock) asyncio.run(coroutine) - self.assertEqual(get_last_provider(), ProviderMock) - -if __name__ == '__main__': - unittest.main()
\ No newline at end of file + self.assertEqual(get_last_provider(), ProviderMock)
\ No newline at end of file diff --git a/etc/unittest/mocks.py b/etc/unittest/mocks.py index a9505997..885bdaee 100644 --- a/etc/unittest/mocks.py +++ b/etc/unittest/mocks.py @@ -22,4 +22,12 @@ class AsyncGeneratorProviderMock(AsyncGeneratorProvider): async def create_async_generator( model, messages, stream, **kwargs ): - yield "Mock"
\ No newline at end of file + yield "Mock" + +class ModelProviderMock(AbstractProvider): + working = True + + def create_completion( + model, messages, stream, **kwargs + ): + yield model
\ No newline at end of file diff --git a/etc/unittest/model.py b/etc/unittest/model.py new file mode 100644 index 00000000..0d318e7a --- /dev/null +++ b/etc/unittest/model.py @@ -0,0 +1,27 @@ +import unittest +import g4f +from g4f import ChatCompletion +from .mocks import ModelProviderMock + +DEFAULT_MESSAGES = [{'role': 'user', 'content': 'Hello'}] + +test_model = g4f.models.Model( + name = "test/test_model", + base_provider = "", + best_provider = ModelProviderMock +) +g4f.models.ModelUtils.convert["test_model"] = test_model + +class TestPassModel(unittest.TestCase): + + def test_model_instance(self): + response = ChatCompletion.create(test_model, DEFAULT_MESSAGES) + self.assertEqual(test_model.name, response) + + def test_model_name(self): + response = ChatCompletion.create("test_model", DEFAULT_MESSAGES) + self.assertEqual(test_model.name, response) + + def test_model_pass(self): + response = ChatCompletion.create("test/test_model", DEFAULT_MESSAGES, ModelProviderMock) + self.assertEqual(test_model.name, response)
\ No newline at end of file diff --git a/g4f/__init__.py b/g4f/__init__.py index 2b0e5b46..996acf6c 100644 --- a/g4f/__init__.py +++ b/g4f/__init__.py @@ -45,12 +45,13 @@ def get_model_and_provider(model : Union[Model, str], else: raise ProviderNotFoundError(f'Provider not found: {provider}') + if isinstance(model, str): + if model in ModelUtils.convert: + model = ModelUtils.convert[model] + if not provider: if isinstance(model, str): - if model in ModelUtils.convert: - model = ModelUtils.convert[model] - else: - raise ModelNotFoundError(f'Model not found: {model}') + raise ModelNotFoundError(f'Model not found: {model}') provider = model.best_provider if not provider: |