From d00dca034d1fd48346fc2c735eac924df1aa4ab4 Mon Sep 17 00:00:00 2001 From: Paul Gauthier Date: Wed, 21 Jun 2023 14:13:34 -0700 Subject: [PATCH] More Model test cases --- tests/test_models.py | 24 +++++++++++++++++++++--- 1 file changed, 21 insertions(+), 3 deletions(-) diff --git a/tests/test_models.py b/tests/test_models.py index 84d04f6d7..af2a6f8d7 100644 --- a/tests/test_models.py +++ b/tests/test_models.py @@ -1,10 +1,28 @@ import unittest + from aider.models import Model + class TestModels(unittest.TestCase): def test_max_context_tokens(self): - model = Model('gpt-3.5') - self.assertEqual(model.max_context_tokens, 4*1024) + model = Model("gpt-3.5-turbo") + self.assertEqual(model.max_context_tokens, 4 * 1024) -if __name__ == '__main__': + model = Model("gpt-3.5-turbo-16k") + self.assertEqual(model.max_context_tokens, 16 * 1024) + + model = Model("gpt-4") + self.assertEqual(model.max_context_tokens, 8 * 1024) + + model = Model("gpt-4-32k") + self.assertEqual(model.max_context_tokens, 32 * 1024) + + model = Model("gpt-4-0101") + self.assertEqual(model.max_context_tokens, 8 * 1024) + + model = Model("gpt-4-32k-2123") + self.assertEqual(model.max_context_tokens, 32 * 1024) + + +if __name__ == "__main__": unittest.main()