Merge pull request #15 from pamelafox/toolchoice

Fix tool choice type

Author: pamelafox

CHANGELOG.md

@@ -2,6 +2,10 @@
 
 All notable changes to this project will be documented in this file.
 
+## [0.1.8] - Aug 3, 2024
+
+- Fix the type for the tool_choice param to be inclusive of "auto" and other options.
+
 ## [0.1.7] - Aug 3, 2024
 
 - Fix bug where you couldn't pass in example tool calls in `few_shots` to `build_messages`.

README.md

@@ -32,7 +32,7 @@ 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` (`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.
+* `tool_choice` (`openai.types.chat.ChatCompletionToolChoiceOptionParam`): (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[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.

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.7"
+version = "0.1.8"
 authors = [{name = "Pamela Fox"}]
 requires-python = ">=3.9"
 readme = "README.md"

src/openai_messages_token_helper/message_builder.py

@@ -8,9 +8,9 @@
     ChatCompletionContentPartParam,
     ChatCompletionMessageParam,
     ChatCompletionMessageToolCallParam,
-    ChatCompletionNamedToolChoiceParam,
     ChatCompletionRole,
     ChatCompletionSystemMessageParam,
+    ChatCompletionToolChoiceOptionParam,
     ChatCompletionToolMessageParam,
     ChatCompletionToolParam,
     ChatCompletionUserMessageParam,
@@ -88,7 +88,7 @@ def build_messages(
     system_prompt: str,
     *,
     tools: Optional[list[ChatCompletionToolParam]] = None,
-    tool_choice: Optional[ChatCompletionNamedToolChoiceParam] = None,
+    tool_choice: Optional[ChatCompletionToolChoiceOptionParam] = None,
     new_user_content: Union[str, list[ChatCompletionContentPartParam], None] = None,  # list is for GPT4v usage
     past_messages: list[ChatCompletionMessageParam] = [],  # *not* including system prompt
     few_shots: list[ChatCompletionMessageParam] = [],  # will always be inserted after system prompt
@@ -103,7 +103,7 @@ def build_messages(
         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[ChatCompletionToolParam]): A list of tools to include in the conversation.
-        tool_choice (ChatCompletionNamedToolChoiceParam): The tool to use in the conversation.
+        tool_choice (ChatCompletionToolChoiceOptionParam): The tool to use in the conversation.
         new_user_content (str | List[ChatCompletionContentPartParam]): Content of new user message to append.
         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.

src/openai_messages_token_helper/model_helper.py

@@ -5,8 +5,8 @@
 import tiktoken
 from openai.types.chat import (
     ChatCompletionMessageParam,
-    ChatCompletionNamedToolChoiceParam,
     ChatCompletionSystemMessageParam,
+    ChatCompletionToolChoiceOptionParam,
     ChatCompletionToolParam,
 )
 
@@ -121,7 +121,7 @@ def count_tokens_for_system_and_tools(
     model: str,
     system_message: ChatCompletionSystemMessageParam | None = None,
     tools: list[ChatCompletionToolParam] | None = None,
-    tool_choice: ChatCompletionNamedToolChoiceParam | None = None,
+    tool_choice: ChatCompletionToolChoiceOptionParam | None = None,
     default_to_cl100k: bool = False,
 ) -> int:
     """

tests/test_messagebuilder.py

@@ -1,7 +1,11 @@
 import typing
 
 import pytest
-from openai.types.chat import ChatCompletionMessageParam, ChatCompletionNamedToolChoiceParam, ChatCompletionToolParam
+from openai.types.chat import (
+    ChatCompletionMessageParam,
+    ChatCompletionToolChoiceOptionParam,
+    ChatCompletionToolParam,
+)
 from openai_messages_token_helper import build_messages, count_tokens_for_message
 
 from .functions import search_sources_toolchoice_auto
@@ -293,7 +297,7 @@ def test_messagebuilder_typing() -> None:
             },
         }
     ]
-    tool_choice: ChatCompletionNamedToolChoiceParam = {
+    tool_choice: ChatCompletionToolChoiceOptionParam = {
         "type": "function",
         "function": {"name": "search_sources"},
     }
@@ -316,3 +320,17 @@ def test_messagebuilder_typing() -> None:
     assert isinstance(messages, list)
     if hasattr(typing, "assert_type"):
         typing.assert_type(messages[0], ChatCompletionMessageParam)
+
+    messages = build_messages(
+        model="gpt-35-turbo",
+        system_prompt="Here are some tools you can use to search for sources.",
+        tools=tools,
+        tool_choice="auto",
+        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)