This commit is contained in:
Paul Gauthier 2023-09-01 15:53:19 -07:00
parent 30a3cc0847
commit 1334392418
8 changed files with 49 additions and 41 deletions

View file

@ -24,33 +24,33 @@ class TestModels(unittest.TestCase):
model = Model.create("gpt-4-32k-2123")
self.assertEqual(model.max_context_tokens, 32 * 1024)
@patch('openai.Model.list')
@patch("openai.Model.list")
def test_openrouter_model_properties(self, mock_model_list):
import openai
old_base = openai.api_base
openai.api_base = 'https://openrouter.ai/api/v1'
openai.api_base = "https://openrouter.ai/api/v1"
mock_model_list.return_value = {
'data': [
"data": [
{
'id': 'openai/gpt-4',
'object': 'model',
'context_length': '8192',
'pricing': {
'prompt': '0.00006',
'completion': '0.00012'
}
"id": "openai/gpt-4",
"object": "model",
"context_length": "8192",
"pricing": {"prompt": "0.00006", "completion": "0.00012"},
}
]
}
mock_model_list.return_value = type('', (), {'data': mock_model_list.return_value['data']})()
mock_model_list.return_value = type(
"", (), {"data": mock_model_list.return_value["data"]}
)()
model = OpenRouterModel("gpt-4")
self.assertEqual(model.name, 'openai/gpt-4')
self.assertEqual(model.name, "openai/gpt-4")
self.assertEqual(model.max_context_tokens, 8192)
self.assertEqual(model.prompt_price, 0.06)
self.assertEqual(model.completion_price, 0.12)
openai.api_base = old_base
if __name__ == "__main__":
unittest.main()