Update test_openai_token_counter.py

chain-ml · Jun 7, 2024 · a1d3854 · a1d3854
1 parent 65ce415
commit a1d3854
Showing 1 changed file with 5 additions and 0 deletions.
diff --git a/tests/unit/llm/test_openai_token_counter.py b/tests/unit/llm/test_openai_token_counter.py
@@ -16,6 +16,11 @@ def test_token_counter_gpt_4_turbo(self):
         counter = OpenAITokenCounter.from_model(model)
         self.assertEqual(counter.token_limit, 128000)
 
+    def test_token_counter_gpt_4o(self):
+        model = "gpt-4o"
+        counter = OpenAITokenCounter.from_model(model)
+        self.assertEqual(counter.token_limit, 128000)
+
     def test_token_counter_exception(self):
         model = "gpt-4"
         counter = OpenAITokenCounter.from_model(model)