summaryrefslogtreecommitdiffstats
diff options
context:
space:
mode:
authorH Lohaus <hlohaus@users.noreply.github.com>2024-01-23 12:18:32 +0100
committerGitHub <noreply@github.com>2024-01-23 12:18:32 +0100
commit984c3e1f891b1242954413d7abf81d0a45aa1110 (patch)
tree1f1bee0d403b70866b2bf5f71cfc5a0fb2fe4efd
parentAdded CodeLlama and misc models from DeepInfra (#1506) (diff)
parentFix: Model sometimes not converted #1507 (diff)
downloadgpt4free-984c3e1f891b1242954413d7abf81d0a45aa1110.tar
gpt4free-984c3e1f891b1242954413d7abf81d0a45aa1110.tar.gz
gpt4free-984c3e1f891b1242954413d7abf81d0a45aa1110.tar.bz2
gpt4free-984c3e1f891b1242954413d7abf81d0a45aa1110.tar.lz
gpt4free-984c3e1f891b1242954413d7abf81d0a45aa1110.tar.xz
gpt4free-984c3e1f891b1242954413d7abf81d0a45aa1110.tar.zst
gpt4free-984c3e1f891b1242954413d7abf81d0a45aa1110.zip
Diffstat (limited to '')
-rw-r--r--etc/unittest/__main__.py1
-rw-r--r--etc/unittest/main.py8
-rw-r--r--etc/unittest/mocks.py10
-rw-r--r--etc/unittest/model.py27
-rw-r--r--g4f/__init__.py9
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: