From e5ec4a92fe5f1a555be82081b0ef7b390b717c18 Mon Sep 17 00:00:00 2001
From: ishaan-jaff <ishaanjaffer0324@gmail.com>
Date: Mon, 13 Nov 2023 07:42:08 -0800
Subject: [PATCH] (test) add test token counter

---
 litellm/tests/test_utils.py | 36 +++++++++++++++++++++++++++++++++++-
 1 file changed, 35 insertions(+), 1 deletion(-)

diff --git a/litellm/tests/test_utils.py b/litellm/tests/test_utils.py
index 22a4af93c3f..739464a03a0 100644
--- a/litellm/tests/test_utils.py
+++ b/litellm/tests/test_utils.py
@@ -10,7 +10,7 @@ sys.path.insert(
 )  # Adds the parent directory to the system path
 import pytest
 import litellm
-from litellm.utils import trim_messages, get_token_count, get_valid_models, check_valid_key, validate_environment, function_to_dict
+from litellm.utils import trim_messages, get_token_count, get_valid_models, check_valid_key, validate_environment, function_to_dict, token_counter
 
 # Assuming your trim_messages, shorten_message_to_fit_limit, and get_token_count functions are all in a module named 'message_utils'
 
@@ -179,3 +179,37 @@ def test_function_to_dict():
     print("passed")
 # test_function_to_dict()
 
+
+def test_token_counter():
+    messages = [
+        {
+            "role": "user",
+            "content": "hi how are you what time is it"
+        }
+    ]
+    tokens = token_counter(
+        model = "gpt-3.5-turbo",
+        messages=messages
+    )
+    print(tokens)
+
+    tokens = token_counter(
+        model = "claude-2",
+        messages=messages
+    )
+    print("claude-2")
+    print(tokens)
+
+    tokens = token_counter(
+        model = "command-nightly",
+        messages=messages
+    )
+    print("command-nightly")
+    print(tokens)
+
+# test_token_counter()
+
+
+
+
+