Add arg to message_builder

Author: pamelafox

CHANGELOG.md

@@ -2,6 +2,11 @@
 
 All notable changes to this project will be documented in this file.
 
+## [0.0.6] - April 24, 2024
+
+- Add keyword argument `fallback_to_default` to `build_messages` function to allow for defaulting to the CL100k token encoder and minimum GPT token limit if the model is not found.
+- Fixed usage of `past_messages` argument of `build_messages` to not skip the last past message. (New user message should *not* be passed in)
+
 ## [0.0.5] - April 24, 2024
 
 - Add keyword argument `default_to_cl100k` to `count_tokens_for_message` function to allow for defaulting to the CL100k token limit if the model is not found.

README.md

@@ -35,6 +35,7 @@ Arguments:
 * `past_messages` (`list[dict]`): The list of past messages in the conversation.
 * `few_shots` (`list[dict]`): A few-shot list of messages to insert after the system prompt.
 * `max_tokens` (`int`): The maximum number of tokens allowed for the conversation.
+* `fallback_to_default` (`bool`): Whether to fallback to default model/token limits if model is not found. Defaults to `False`.
 
 Returns:
 

pyproject.toml

@@ -1,7 +1,7 @@
 [project]
 name = "openai-messages-token-helper"
 description = "A helper library for estimating tokens used by messages sent through OpenAI Chat Completions API."
-version = "0.0.5"
+version = "0.0.6"
 authors = [{name = "Pamela Fox"}]
 requires-python = ">=3.9"
 readme = "README.md"

src/openai_messages_token_helper/message_builder.py

@@ -1,6 +1,5 @@
 import logging
 import unicodedata
-from collections.abc import Mapping
 from typing import Optional, Union
 
 from openai.types.chat import (
@@ -52,9 +51,6 @@ def insert_message(self, role: str, content: Union[str, list[ChatCompletionConte
             raise ValueError(f"Invalid role: {role}")
         self.messages.insert(index, message)
 
-    def count_tokens_for_message(self, message: Mapping[str, object]):
-        return count_tokens_for_message(self.model, message)
-
     def normalize_content(self, content: Union[str, list[ChatCompletionContentPartParam]]):
         if isinstance(content, str):
             return unicodedata.normalize("NFC", content)
@@ -72,6 +68,7 @@ def build_messages(
     past_messages: list[dict[str, str]] = [],  # *not* including system prompt
     few_shots=[],  # will always be inserted after system prompt
     max_tokens: Optional[int] = None,
+    fallback_to_default: bool = False,
 ) -> list[ChatCompletionMessageParam]:
     """
     Build a list of messages for a chat conversation, given the system prompt, new user message,
@@ -84,10 +81,11 @@ def build_messages(
         past_messages (list[dict]): The list of past messages in the conversation.
         few_shots (list[dict]): A few-shot list of messages to insert after the system prompt.
         max_tokens (int): The maximum number of tokens allowed for the conversation.
+        fallback_to_default (bool): Whether to fallback to default model if the model is not found.
     """
     message_builder = MessageBuilder(system_prompt, model)
     if max_tokens is None:
-        max_tokens = get_token_limit(model)
+        max_tokens = get_token_limit(model, default_to_minimum=fallback_to_default)
 
     for shot in reversed(few_shots):
         message_builder.insert_message(shot.get("role"), shot.get("content"))
@@ -99,11 +97,11 @@ def build_messages(
 
     total_token_count = 0
     for existing_message in message_builder.messages:
-        total_token_count += message_builder.count_tokens_for_message(existing_message)
+        total_token_count += count_tokens_for_message(model, existing_message, default_to_cl100k=fallback_to_default)
 
-    newest_to_oldest = list(reversed(past_messages[:-1]))
+    newest_to_oldest = list(reversed(past_messages))
     for message in newest_to_oldest:
-        potential_message_count = message_builder.count_tokens_for_message(message)
+        potential_message_count = count_tokens_for_message(model, message, default_to_cl100k=fallback_to_default)
         if (total_token_count + potential_message_count) > max_tokens:
             logging.info("Reached max tokens of %d, history will be truncated", max_tokens)
             break

tests/test_messagebuilder.py

@@ -1,3 +1,4 @@
+import pytest
 from openai_messages_token_helper import build_messages, count_tokens_for_message
 
 from .messages import system_message_short, system_message_unicode, user_message, user_message_unicode
@@ -33,3 +34,24 @@ def test_messagebuilder_unicode_append():
     assert messages == [system_message_unicode["message"], user_message_unicode["message"]]
     assert count_tokens_for_message("gpt-35-turbo", messages[0]) == system_message_unicode["count"]
     assert count_tokens_for_message("gpt-35-turbo", messages[1]) == user_message_unicode["count"]
+
+
+def test_messagebuilder_model_error():
+    model = "phi-3"
+    with pytest.raises(ValueError, match="Called with unknown model name: phi-3"):
+        build_messages(
+            model, system_message_short["message"]["content"], new_user_message=user_message["message"]["content"]
+        )
+
+
+def test_messagebuilder_model_fallback():
+    model = "phi-3"
+    messages = build_messages(
+        model,
+        system_message_short["message"]["content"],
+        new_user_message=user_message["message"]["content"],
+        fallback_to_default=True,
+    )
+    assert messages == [system_message_short["message"], user_message["message"]]
+    assert count_tokens_for_message(model, messages[0], default_to_cl100k=True) == system_message_short["count"]
+    assert count_tokens_for_message(model, messages[1], default_to_cl100k=True) == user_message["count"]