From bcb89dcf4ac12882717306dba556036fe9ef901a Mon Sep 17 00:00:00 2001
From: ishaan-jaff <ishaanjaffer0324@gmail.com>
Date: Mon, 11 Sep 2023 17:44:50 -0700
Subject: [PATCH] add test cases for token trimmer

---
 litellm/tests/test_utils.py | 58 +++++++++++++++++++++++++++++++++++++
 1 file changed, 58 insertions(+)
 create mode 100644 litellm/tests/test_utils.py

diff --git a/litellm/tests/test_utils.py b/litellm/tests/test_utils.py
new file mode 100644
index 000000000..f59e647be
--- /dev/null
+++ b/litellm/tests/test_utils.py
@@ -0,0 +1,58 @@
+import sys, os
+import traceback
+from dotenv import load_dotenv
+
+load_dotenv()
+import os
+
+sys.path.insert(
+    0, os.path.abspath("../..")
+)  # Adds the parent directory to the system path
+import pytest
+import litellm
+from litellm.utils import safe_messages, get_token_count
+
+# Assuming your trim_messages, shorten_message_to_fit_limit, and get_token_count functions are all in a module named 'message_utils'
+
+# Test 1: Check trimming of normal message
+def test_basic_trimming():
+    messages = [{"role": "user", "content": "This is a long message that definitely exceeds the token limit."}]
+    trimmed_messages = safe_messages(messages, model="claude-2", max_tokens=8)
+    print("trimmed messages")
+    print(trimmed_messages)
+    print(get_token_count(messages=trimmed_messages, model="claude-2"))
+    assert (get_token_count(messages=trimmed_messages, model="claude-2")) <= 8
+# test_basic_trimming()
+
+def test_multiple_messages_trimming():
+    messages = [
+        {"role": "user", "content": "This is a long message that will exceed the token limit."},
+        {"role": "user", "content": "This is another long message that will also exceed the limit."}
+    ]
+    trimmed_messages = safe_messages(messages=messages, model="gpt-3.5-turbo", max_tokens=20)
+    print("Trimmed messages")
+    print(trimmed_messages)
+    print(get_token_count(messages=trimmed_messages, model="gpt-3.5-turbo"))
+    assert(get_token_count(messages=trimmed_messages, model="gpt-3.5-turbo")) <= 20
+# test_multiple_messages_trimming()
+
+def test_multiple_messages_no_trimming():
+    messages = [
+        {"role": "user", "content": "This is a long message that will exceed the token limit."},
+        {"role": "user", "content": "This is another long message that will also exceed the limit."}
+    ]
+    trimmed_messages = safe_messages(messages=messages, model="gpt-3.5-turbo", max_tokens=100)
+    print("Trimmed messages")
+    print(trimmed_messages)
+    assert(messages==trimmed_messages)
+
+# test_multiple_messages_no_trimming()
+
+
+def test_large_trimming():
+    messages = [{"role": "user", "content": "This is a singlelongwordthatexceedsthelimit."}, {"role": "user", "content": "This is a singlelongwordthatexceedsthelimit."},{"role": "user", "content": "This is a singlelongwordthatexceedsthelimit."},{"role": "user", "content": "This is a singlelongwordthatexceedsthelimit."},{"role": "user", "content": "This is a singlelongwordthatexceedsthelimit."}]
+    trimmed_messages = safe_messages(messages, max_tokens=20, model="random")
+    print("trimmed messages")
+    print(trimmed_messages)
+    assert(get_token_count(messages=trimmed_messages, model="random")) <= 20
+# test_large_trimming()
\ No newline at end of file