Expanded estimate_tokens test to compare against tiktoken

Author: jamesbraza

packages/lmi/tests/test_embeddings.py

@@ -4,6 +4,7 @@
 
 import litellm
 import pytest
+import tiktoken
 from litellm.caching import Cache, InMemoryCache
 from pytest_subtests import SubTests
 
@@ -20,17 +21,29 @@
 from lmi.utils import VCR_DEFAULT_MATCH_ON, encode_image_as_url
 
 
-def test_estimate_tokens(stub_png_image: bytes) -> None:
-    # Test text-only
-    text_only = "Hello world"
-    assert estimate_tokens(text_only) == 2.75
+def test_estimate_tokens(subtests: SubTests, stub_png_image: bytes) -> None:
+    with subtests.test(msg="text only"):
+        text_only = "Hello world"
+        text_only_estimated_token_count = estimate_tokens(text_only)
+        assert text_only_estimated_token_count == 2.75, (
+            "Expected a reasonable token estimate"
+        )
+        text_only_actual_token_count = len(
+            tiktoken.get_encoding("cl100k_base").encode(text_only)
+        )
+        assert text_only_estimated_token_count == pytest.approx(
+            text_only_actual_token_count, abs=1
+        ), "Estimation should be within one token of what tiktoken"
 
     # Test multimodal (text + image)
-    multimodal = [
-        "What is in this image?",
-        encode_image_as_url(image_type="png", image_data=stub_png_image),
-    ]
-    assert estimate_tokens(multimodal) == 90.5
+    with subtests.test(msg="multimodal"):  # Text + image
+        multimodal = [
+            "What is in this image?",
+            encode_image_as_url(image_type="png", image_data=stub_png_image),
+        ]
+        assert estimate_tokens(multimodal) == 90.5, (
+            "Expected a reasonable token estimate"
+        )
 
 
 class TestLiteLLMEmbeddingModel: