Merge pull request #13 from pamelafox/supporttoolsfallback

Support fallback when calculating tokens for non-openAI models

Author: pamelafox

.github/workflows/python.yaml

@@ -25,7 +25,7 @@ jobs:
             python -m pip install --upgrade pip
             python3 -m pip install -e '.[dev]'
         - name: Lint with ruff
-          run: ruff .
+          run: ruff check .
         - name: Check formatting with black
           run: black . --check --verbose
         - name: Run unit tests

src/openai_messages_token_helper/model_helper.py

@@ -144,7 +144,6 @@ def count_tokens_for_system_and_tools(
     if system_message:
         tokens += count_tokens_for_message(model, system_message, default_to_cl100k)
     if tools:
-        encoding = tiktoken.encoding_for_model(model)
         tokens += len(encoding.encode(format_function_definitions(tools)))
         tokens += 9  # Additional tokens for function definition of tools
     # If there's a system message and tools are present, subtract four tokens

tests/test_modelhelper.py

@@ -1,7 +1,7 @@
 import pytest
 from openai_messages_token_helper import count_tokens_for_message, count_tokens_for_system_and_tools, get_token_limit
 
-from .functions import FUNCTION_COUNTS
+from .functions import FUNCTION_COUNTS, search_sources_toolchoice_auto
 from .messages import system_message, system_message_with_name, text_and_image_message, user_message
 
 
@@ -105,3 +105,17 @@ def test_count_tokens_for_system_and_tools(function_count_pair):
     assert (
         diff >= 0 and diff <= 3
     ), f"Expected {expected_tokens} tokens, got {counted_tokens}. Counted tokens is only allowed to be off by 3 in the over-counting direction."
+
+
+def test_count_tokens_for_system_and_tools_fallback(caplog):
+    function_count_pair = search_sources_toolchoice_auto
+    with caplog.at_level("WARNING"):
+        counted_tokens = count_tokens_for_system_and_tools(
+            "llama-3.1",
+            function_count_pair["system_message"],
+            function_count_pair["tools"],
+            function_count_pair["tool_choice"],
+            default_to_cl100k=True,
+        )
+        assert counted_tokens == function_count_pair["count"]
+        assert "Model llama-3.1 not found, defaulting to CL100k encoding" in caplog.text