fix failing tests

This commit is contained in:
Joshua Vial 2023-12-11 22:37:23 +13:00
parent 9ceaf97f08
commit 90d5071709
2 changed files with 16 additions and 13 deletions

View file

@ -12,6 +12,7 @@ def is_image_file(file_name):
:param file_name: The name of the file to check. :param file_name: The name of the file to check.
:return: True if the file is an image, False otherwise. :return: True if the file is an image, False otherwise.
""" """
file_name = str(file_name) # Convert file_name to string
return any(file_name.endswith(ext) for ext in IMAGE_EXTENSIONS) return any(file_name.endswith(ext) for ext in IMAGE_EXTENSIONS)
@ -50,4 +51,6 @@ def is_gpt4_with_openai_base_url(model_name, client):
:param client: The OpenAI client instance. :param client: The OpenAI client instance.
:return: True if conditions are met, False otherwise. :return: True if conditions are met, False otherwise.
""" """
if client is None or not hasattr(client, 'base_url'):
return False
return model_name.startswith("gpt-4") and "api.openai.com" in client.base_url.host return model_name.startswith("gpt-4") and "api.openai.com" in client.base_url.host

View file

@ -29,19 +29,19 @@ class TestModels(unittest.TestCase):
def test_openrouter_model_properties(self): def test_openrouter_model_properties(self):
client = MagicMock() client = MagicMock()
client.models.list.return_value = { class ModelData:
"data": [ def __init__(self, id, object, context_length, pricing):
{ self.id = id
"id": "openai/gpt-4", self.object = object
"object": "model", self.context_length = context_length
"context_length": "8192", self.pricing = pricing
"pricing": {"prompt": "0.00006", "completion": "0.00012"},
} model_data = ModelData("openai/gpt-4", "model", "8192", {"prompt": "0.00006", "completion": "0.00012"})
] class ModelList:
} def __init__(self, data):
client.models.list.return_value = type( self.data = data
"", (), {"data": client.models.list.return_value["data"]}
)() client.models.list.return_value = ModelList([model_data])
model = OpenRouterModel(client, "gpt-4") model = OpenRouterModel(client, "gpt-4")
self.assertEqual(model.name, "openai/gpt-4") self.assertEqual(model.name, "openai/gpt-4")