coding-alt/AutoGPT

1import unittest2 3import tests.context4from autogpt.token_counter import count_message_tokens, count_string_tokens5 6 7class TestTokenCounter(unittest.TestCase):8    def test_count_message_tokens(self):9        messages = [10            {"role": "user", "content": "Hello"},11            {"role": "assistant", "content": "Hi there!"},12        ]13        self.assertEqual(count_message_tokens(messages), 17)14 15    def test_count_message_tokens_with_name(self):16        messages = [17            {"role": "user", "content": "Hello", "name": "John"},18            {"role": "assistant", "content": "Hi there!"},19        ]20        self.assertEqual(count_message_tokens(messages), 17)21 22    def test_count_message_tokens_empty_input(self):23        self.assertEqual(count_message_tokens([]), 3)24 25    def test_count_message_tokens_invalid_model(self):26        messages = [27            {"role": "user", "content": "Hello"},28            {"role": "assistant", "content": "Hi there!"},29        ]30        with self.assertRaises(KeyError):31            count_message_tokens(messages, model="invalid_model")32 33    def test_count_message_tokens_gpt_4(self):34        messages = [35            {"role": "user", "content": "Hello"},36            {"role": "assistant", "content": "Hi there!"},37        ]38        self.assertEqual(count_message_tokens(messages, model="gpt-4-0314"), 15)39 40    def test_count_string_tokens(self):41        string = "Hello, world!"42        self.assertEqual(43            count_string_tokens(string, model_name="gpt-3.5-turbo-0301"), 444        )45 46    def test_count_string_tokens_empty_input(self):47        self.assertEqual(count_string_tokens("", model_name="gpt-3.5-turbo-0301"), 0)48 49    def test_count_message_tokens_invalid_model(self):50        messages = [51            {"role": "user", "content": "Hello"},52            {"role": "assistant", "content": "Hi there!"},53        ]54        with self.assertRaises(NotImplementedError):55            count_message_tokens(messages, model="invalid_model")56 57    def test_count_string_tokens_gpt_4(self):58        string = "Hello, world!"59        self.assertEqual(count_string_tokens(string, model_name="gpt-4-0314"), 4)60 61 62if __name__ == "__main__":63    unittest.main()64