init

AlgoriAgent/tests/token_utils_test.py

@@ -0,0 +1,52 @@
+# -*- coding: utf-8 -*-
+""" Unit test for token_utils."""
+import unittest
+from agentscope.utils.token_utils import get_openai_max_length
+from agentscope.utils.token_utils import count_openai_token
+
+
+class TokenUtilsTest(unittest.TestCase):
+    """Token utils test."""
+
+    def test_get_openai_max_length(self) -> None:
+        """Test the function get_openai_max_length."""
+        self.assertEqual(get_openai_max_length("gpt-4"), 8192)
+        self.assertEqual(get_openai_max_length("gpt-3.5-turbo"), 4096)
+        with self.assertRaises(KeyError):
+            get_openai_max_length("non-existing-model")
+
+    def test_count_openai_token(self) -> None:
+        """Test the function count_openai_token."""
+        test_model_1 = "text-davinci-003"
+        test_content_str = "This is a test string."
+        self.assertEqual(
+            count_openai_token(test_content_str, test_model_1),
+            6,
+        )
+
+        test_content_list = [
+            {
+                "role": "system",
+                "content": "This is a system content.",
+            },
+            {
+                "role": "user",
+                "name": "example_user",
+                "content": "This is a user content.",
+            },
+        ]
+
+        test_model_2 = "gpt-4"
+        self.assertEqual(
+            count_openai_token(test_content_list, test_model_2),
+            26,
+        )
+
+        # Test unsupported model
+        unsupported_model = "unsupported-model"
+        with self.assertRaises(NotImplementedError):
+            count_openai_token(test_content_str, unsupported_model)
+
+
+if __name__ == "__main__":
+    unittest.main()