Merge pull request #10 from pamelafox/motypes

Use more precise types

Author: pamelafox

.github/workflows/python.yaml

@@ -30,6 +30,6 @@ jobs:
           run: black . --check --verbose
         - name: Run unit tests
           run: |
-            python3 -m pytest -s -vv --cov --cov-fail-under=99
+            python3 -m pytest -s -vv --cov --cov-fail-under=98
         - name: Run type checks
           run: mypy .

CHANGELOG.md

@@ -2,6 +2,10 @@
 
 All notable changes to this project will be documented in this file.
 
+## [0.1.3] - May 2, 2024
+
+- Use openai type annotations for more precise type hints, and add a typing test.
+
 ## [0.1.2] - May 2, 2024
 
 - Add `py.typed` file so that mypy can find the type hints in this package.

README.md

@@ -32,10 +32,10 @@ Arguments:
 * `model` (`str`): The model name to use for token calculation, like gpt-3.5-turbo.
 * `system_prompt` (`str`): The initial system prompt message.
 * `tools` (`List[openai.types.chat.ChatCompletionToolParam]`): (Optional) The tools that will be used in the conversation. These won't be part of the final returned messages, but they will be used to calculate the token count.
-* `tool_choice` (`str | dict`): (Optional) The tool choice that will be used in the conversation. This won't be part of the final returned messages, but it will be used to calculate the token count.
+* `tool_choice` (`openai.types.chat.ChatCompletionNamedToolChoiceParam`): (Optional) The tool choice that will be used in the conversation. This won't be part of the final returned messages, but it will be used to calculate the token count.
 * `new_user_content` (`str | List[openai.types.chat.ChatCompletionContentPartParam]`): (Optional) The content of new user message to append.
-* `past_messages` (`list[dict]`): (Optional) The list of past messages in the conversation.
-* `few_shots` (`list[dict]`): (Optional) A few-shot list of messages to insert after the system prompt.
+* `past_messages` (`list[openai.types.chat.ChatCompletionMessageParam]`): (Optional) The list of past messages in the conversation.
+* `few_shots` (`list[openai.types.chat.ChatCompletionMessageParam]`): (Optional) A few-shot list of messages to insert after the system prompt.
 * `max_tokens` (`int`): (Optional) The maximum number of tokens allowed for the conversation.
 * `fallback_to_default` (`bool`): (Optional) Whether to fallback to default model/token limits if model is not found. Defaults to `False`.
 
@@ -83,7 +83,7 @@ Counts the number of tokens in a message.
 Arguments:
 
 * `model` (`str`): The model name to use for token calculation, like gpt-3.5-turbo.
-* `message` (`dict`): The message to count tokens for.
+* `message` (`openai.types.chat.ChatCompletionMessageParam`): The message to count tokens for.
 * `default_to_cl100k` (`bool`): Whether to default to the CL100k token limit if the model is not found.
 
 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.1.2"
+version = "0.1.3"
 authors = [{name = "Pamela Fox"}]
 requires-python = ">=3.9"
 readme = "README.md"

src/openai_messages_token_helper/message_builder.py

@@ -1,22 +1,26 @@
 import logging
 import unicodedata
+from collections.abc import Iterable
 from typing import Optional, Union
 
 from openai.types.chat import (
     ChatCompletionAssistantMessageParam,
     ChatCompletionContentPartParam,
     ChatCompletionMessageParam,
+    ChatCompletionNamedToolChoiceParam,
+    ChatCompletionRole,
     ChatCompletionSystemMessageParam,
+    ChatCompletionToolParam,
     ChatCompletionUserMessageParam,
 )
 
 from .model_helper import count_tokens_for_message, count_tokens_for_system_and_tools, get_token_limit
 
 
-def normalize_content(content: Union[str, list[ChatCompletionContentPartParam]]):
+def normalize_content(content: Union[str, Iterable[ChatCompletionContentPartParam]]):
     if isinstance(content, str):
         return unicodedata.normalize("NFC", content)
-    elif isinstance(content, list):
+    else:
         for part in content:
             if "image_url" not in part:
                 part["text"] = unicodedata.normalize("NFC", part["text"])
@@ -36,14 +40,19 @@ class _MessageBuilder:
     """
 
     def __init__(self, system_content: str):
-        self.messages: list[ChatCompletionMessageParam] = [
-            ChatCompletionSystemMessageParam(role="system", content=normalize_content(system_content))
-        ]
+        self.system_message = ChatCompletionSystemMessageParam(role="system", content=normalize_content(system_content))
+        self.messages: list[ChatCompletionMessageParam] = []
 
-    def insert_message(self, role: str, content: Union[str, list[ChatCompletionContentPartParam]], index: int = 1):
+    @property
+    def all_messages(self) -> list[ChatCompletionMessageParam]:
+        return [self.system_message] + self.messages
+
+    def insert_message(
+        self, role: ChatCompletionRole, content: Union[str, Iterable[ChatCompletionContentPartParam]], index: int = 0
+    ):
         """
         Inserts a message into the conversation at the specified index,
-        or at index 1 (after system message) if no index is specified.
+        or at index 0 if no index is specified.
         Args:
             role (str): The role of the message sender (either "user", "system", or "assistant").
             content (str | List[ChatCompletionContentPartParam]): The content of the message.
@@ -63,11 +72,11 @@ def build_messages(
     model: str,
     system_prompt: str,
     *,
-    tools: Optional[list[dict[str, dict]]] = None,
-    tool_choice: Optional[Union[str, dict]] = None,
+    tools: Optional[list[ChatCompletionToolParam]] = None,
+    tool_choice: Optional[ChatCompletionNamedToolChoiceParam] = None,
     new_user_content: Union[str, list[ChatCompletionContentPartParam], None] = None,  # list is for GPT4v usage
-    past_messages: list[dict[str, str]] = [],  # *not* including system prompt
-    few_shots=[],  # will always be inserted after system prompt
+    past_messages: list[ChatCompletionMessageParam] = [],  # *not* including system prompt
+    few_shots: list[ChatCompletionMessageParam] = [],  # will always be inserted after system prompt
     max_tokens: Optional[int] = None,
     fallback_to_default: bool = False,
 ) -> list[ChatCompletionMessageParam]:
@@ -78,11 +87,11 @@ def build_messages(
     Args:
         model (str): The model name to use for token calculation, like gpt-3.5-turbo.
         system_prompt (str): The initial system prompt message.
-        tools (list[dict]): A list of tools to include in the conversation.
-        tool_choice (str | dict): The tool to use in the conversation.
+        tools (list[ChatCompletionToolParam]): A list of tools to include in the conversation.
+        tool_choice (ChatCompletionNamedToolChoiceParam): The tool to use in the conversation.
         new_user_content (str | List[ChatCompletionContentPartParam]): Content of new user message to append.
-        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.
+        past_messages (list[ChatCompletionMessageParam]): The list of past messages in the conversation.
+        few_shots (list[ChatCompletionMessageParam]): 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.
     """
@@ -93,17 +102,19 @@ def build_messages(
     message_builder = _MessageBuilder(system_prompt)
 
     for shot in reversed(few_shots):
-        message_builder.insert_message(shot.get("role"), shot.get("content"))
+        if shot["role"] is None or shot["content"] is None:
+            raise ValueError("Few-shot messages must have both role and content")
+        message_builder.insert_message(shot["role"], shot["content"])
 
-    append_index = len(few_shots) + 1
+    append_index = len(few_shots)
 
     if new_user_content:
         message_builder.insert_message("user", new_user_content, index=append_index)
 
     total_token_count = count_tokens_for_system_and_tools(
-        model, message_builder.messages[0], tools, tool_choice, default_to_cl100k=fallback_to_default
+        model, message_builder.system_message, tools, tool_choice, default_to_cl100k=fallback_to_default
     )
-    for existing_message in message_builder.messages[1:]:
+    for existing_message in message_builder.messages:
         total_token_count += count_tokens_for_message(model, existing_message, default_to_cl100k=fallback_to_default)
 
     newest_to_oldest = list(reversed(past_messages))
@@ -112,6 +123,9 @@ def build_messages(
         if (total_token_count + potential_message_count) > max_tokens:
             logging.info("Reached max tokens of %d, history will be truncated", max_tokens)
             break
+
+        if message["role"] is None or message["content"] is None:
+            raise ValueError("Few-shot messages must have both role and content")
         message_builder.insert_message(message["role"], message["content"], index=append_index)
         total_token_count += potential_message_count
-    return message_builder.messages
+    return message_builder.all_messages

src/openai_messages_token_helper/model_helper.py

@@ -1,9 +1,14 @@
 from __future__ import annotations
 
 import logging
-from collections.abc import Mapping
 
 import tiktoken
+from openai.types.chat import (
+    ChatCompletionMessageParam,
+    ChatCompletionNamedToolChoiceParam,
+    ChatCompletionSystemMessageParam,
+    ChatCompletionToolParam,
+)
 
 from .function_format import format_function_definitions
 from .images_helper import count_tokens_for_image
@@ -69,7 +74,7 @@ def encoding_for_model(model: str, default_to_cl100k=False) -> tiktoken.Encoding
             raise
 
 
-def count_tokens_for_message(model: str, message: Mapping[str, object], default_to_cl100k=False) -> int:
+def count_tokens_for_message(model: str, message: ChatCompletionMessageParam, default_to_cl100k=False) -> int:
     """
     Calculate the number of tokens required to encode a message. Based off cookbook:
     https://github.com/openai/openai-cookbook/blob/main/examples/How_to_count_tokens_with_tiktoken.ipynb
@@ -113,9 +118,9 @@ def count_tokens_for_message(model: str, message: Mapping[str, object], default_
 
 def count_tokens_for_system_and_tools(
     model: str,
-    system_message: Mapping[str, object] | None = None,
-    tools: list[dict[str, dict]] | None = None,
-    tool_choice: str | dict | None = None,
+    system_message: ChatCompletionSystemMessageParam | None = None,
+    tools: list[ChatCompletionToolParam] | None = None,
+    tool_choice: ChatCompletionNamedToolChoiceParam | None = None,
     default_to_cl100k: bool = False,
 ) -> int:
     """

tests/test_messagebuilder.py

@@ -1,4 +1,7 @@
+import typing
+
 import pytest
+from openai.types.chat import ChatCompletionMessageParam, ChatCompletionNamedToolChoiceParam, ChatCompletionToolParam
 from openai_messages_token_helper import build_messages, count_tokens_for_message
 
 from .functions import search_sources_toolchoice_auto
@@ -213,3 +216,48 @@ def test_messagebuilder_system_tools():
         max_tokens=90,
     )
     assert messages == [search_sources_toolchoice_auto["system_message"], user_message_pm["message"]]
+
+
+def test_messagebuilder_typing() -> None:
+    tools: list[ChatCompletionToolParam] = [
+        {
+            "type": "function",
+            "function": {
+                "name": "search_sources",
+                "description": "Retrieve sources from the Azure AI Search index",
+                "parameters": {
+                    "type": "object",
+                    "properties": {
+                        "search_query": {
+                            "type": "string",
+                            "description": "Query string to retrieve documents from azure search eg: 'Health care plan'",
+                        }
+                    },
+                    "required": ["search_query"],
+                },
+            },
+        }
+    ]
+    tool_choice: ChatCompletionNamedToolChoiceParam = {
+        "type": "function",
+        "function": {"name": "search_sources"},
+    }
+
+    past_messages: list[ChatCompletionMessageParam] = [
+        {"role": "user", "content": "What are my health plans?"},
+        {"role": "assistant", "content": "Here are some tools you can use to search for sources."},
+    ]
+
+    messages = build_messages(
+        model="gpt-35-turbo",
+        system_prompt="Here are some tools you can use to search for sources.",
+        tools=tools,
+        tool_choice=tool_choice,
+        past_messages=past_messages,
+        new_user_content="What are my health plans?",
+        max_tokens=90,
+    )
+
+    assert isinstance(messages, list)
+    if hasattr(typing, "assert_type"):
+        typing.assert_type(messages[0], ChatCompletionMessageParam)