From e4f743881c915ba250b859d01ef78c7b90a96bf5 Mon Sep 17 00:00:00 2001
From: Heiner Lohaus <hlohaus@users.noreply.github.com>
Date: Tue, 23 Jan 2024 12:17:16 +0100
Subject: Fix: Model sometimes not converted #1507

---
 etc/unittest/__main__.py |  1 +
 etc/unittest/main.py     |  8 +++-----
 etc/unittest/mocks.py    | 10 +++++++++-
 etc/unittest/model.py    | 27 +++++++++++++++++++++++++++
 g4f/__init__.py          |  9 +++++----
 5 files changed, 45 insertions(+), 10 deletions(-)
 create mode 100644 etc/unittest/model.py

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:
-- 
cgit v1.2.3