Fix: Model sometimes not converted #1507

This commit is contained in:
Heiner Lohaus 2024-01-23 12:17:16 +01:00
parent b6feec9e7d
commit e4f743881c
5 changed files with 45 additions and 10 deletions

View File

@ -2,5 +2,6 @@ import unittest
from .asyncio import *
from .backend import *
from .main import *
from .model import *
unittest.main()

View File

@ -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()
self.assertEqual(get_last_provider(), ProviderMock)

View File

@ -22,4 +22,12 @@ class AsyncGeneratorProviderMock(AsyncGeneratorProvider):
async def create_async_generator(
model, messages, stream, **kwargs
):
yield "Mock"
yield "Mock"
class ModelProviderMock(AbstractProvider):
working = True
def create_completion(
model, messages, stream, **kwargs
):
yield model

27
etc/unittest/model.py Normal file
View File

@ -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)

View File

@ -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: