dearwolf/aider
1
0
Fork 0
mirror of https://github.com/Aider-AI/aider.git synced 2025-05-30 01:04:59 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 3

test openrouter model properties

Browse source
This commit is contained in:
Joshua Vial 2023-08-23 21:37:22 +12:00
parent d1e0a196b5
commit 668a0500ff
2 changed files with 27 additions and 7 deletions
Download patch file Download diff file Expand all files Collapse all files

6
aider/models/openrouter.py
View file

@ -28,9 +28,9 @@ class OpenRouterModel(Model):
found = next((details for details in cached_model_details if details.get('id') == name), None)
if found:
self.max_context_tokens = int(found.context_length)
self.prompt_price = float(found.get('pricing').get('prompt')) * 1000
self.completion_price = float(found.get('pricing').get('completion')) * 1000
self.max_context_tokens = int(found.get('context_length'))
self.prompt_price = round(float(found.get('pricing').get('prompt')) * 1000,6)
self.completion_price = round(float(found.get('pricing').get('completion')) * 1000,6)
else:
raise ValueError(f'invalid openrouter model: {name}')

28
tests/test_models.py
View file

@ -1,6 +1,7 @@
import unittest
from unittest.mock import patch
from aider.models import Model
from aider.models import Model, OpenRouterModel
class TestModels(unittest.TestCase):
@ -23,12 +24,31 @@ class TestModels(unittest.TestCase):
model = Model.create("gpt-4-32k-2123")
self.assertEqual(model.max_context_tokens, 32 * 1024)
def test_openrouter_models(self):
@patch('openai.Model.list')
def test_openrouter_model_properties(self, mock_model_list):
import openai
openai.api_base = 'https://openrouter.ai/api/v1'
model = Model.create("gpt-3.5-turbo")
self.assertEqual(model.name, 'openai/gpt-3.5-turbo')
mock_model_list.return_value = {
'data': [
{
'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']})()
model = OpenRouterModel("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)
if __name__ == "__main__":
unittest.main()