Merge pull request #14 from pamelafox/fewshottols

Support tools in few shots

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=98
+            python3 -m pytest -s -vv --cov --cov-fail-under=97
         - name: Run type checks
           run: mypy .

CHANGELOG.md

@@ -2,6 +2,14 @@
 
 All notable changes to this project will be documented in this file.
 
+## [0.1.7] - Aug 3, 2024
+
+- Fix bug where you couldn't pass in example tool calls in `few_shots` to `build_messages`.
+
+## [0.1.6] - Aug 2, 2024
+
+- Fix bug where you couldn't pass in `tools` and `default_to_cl100k` to True with a non-OpenAI model.
+
 ## [0.1.5] - June 4, 2024
 
 - Remove spurious `print` call when counting tokens for function calling.

CONTRIBUTING.md

@@ -17,10 +17,16 @@ python3 -m pytest
 
 ## Publishing
 
-Publish to PyPi:
+1. Update the CHANGELOG with description of changes
 
-```shell
-export FLIT_USERNAME=__token__
-export FLIT_PASSWORD=<your-pypi-token>
-flit publish
-```
+2. Update the version number in pyproject.toml
+
+3. Push the changes to the main branch
+
+4. Publish to PyPi:
+
+    ```shell
+    export FLIT_USERNAME=__token__
+    export FLIT_PASSWORD=<your-pypi-token>
+    flit publish
+    ```

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

src/openai_messages_token_helper/message_builder.py

@@ -7,17 +7,21 @@
     ChatCompletionAssistantMessageParam,
     ChatCompletionContentPartParam,
     ChatCompletionMessageParam,
+    ChatCompletionMessageToolCallParam,
     ChatCompletionNamedToolChoiceParam,
     ChatCompletionRole,
     ChatCompletionSystemMessageParam,
+    ChatCompletionToolMessageParam,
     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, Iterable[ChatCompletionContentPartParam]]):
+def normalize_content(content: Union[str, Iterable[ChatCompletionContentPartParam], None]):
+    if content is None:
+        return None
     if isinstance(content, str):
         return unicodedata.normalize("NFC", content)
     else:
@@ -48,7 +52,12 @@ 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
+        self,
+        role: ChatCompletionRole,
+        content: Union[str, Iterable[ChatCompletionContentPartParam], None],
+        index: int = 0,
+        tool_calls: Optional[Iterable[ChatCompletionMessageToolCallParam]] = None,
+        tool_call_id: Optional[str] = None,
     ):
         """
         Inserts a message into the conversation at the specified index,
@@ -63,8 +72,14 @@ def insert_message(
             message = ChatCompletionUserMessageParam(role="user", content=normalize_content(content))
         elif role == "assistant" and isinstance(content, str):
             message = ChatCompletionAssistantMessageParam(role="assistant", content=normalize_content(content))
+        elif role == "assistant" and tool_calls is not None:
+            message = ChatCompletionAssistantMessageParam(role="assistant", tool_calls=tool_calls)
+        elif role == "tool" and tool_call_id is not None:
+            message = ChatCompletionToolMessageParam(
+                role="tool", tool_call_id=tool_call_id, content=normalize_content(content)
+            )
         else:
-            raise ValueError(f"Invalid role: {role}")
+            raise ValueError("Invalid message for builder")
         self.messages.insert(index, message)
 
 
@@ -102,9 +117,17 @@ def build_messages(
     message_builder = _MessageBuilder(system_prompt)
 
     for shot in reversed(few_shots):
-        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"])
+        if shot["role"] is None or (shot.get("content") is None and shot.get("tool_calls") is None):
+            raise ValueError("Few-shot messages must have role and either content or tool_calls")
+        tool_call_id = shot.get("tool_call_id")
+        if tool_call_id is not None and not isinstance(tool_call_id, str):
+            raise ValueError("tool_call_id must be a string value")
+        tool_calls = shot.get("tool_calls")
+        if tool_calls is not None and not isinstance(tool_calls, Iterable):
+            raise ValueError("tool_calls must be a list of tool calls")
+        message_builder.insert_message(
+            shot["role"], shot.get("content"), tool_calls=tool_calls, tool_call_id=tool_call_id
+        )
 
     append_index = len(few_shots)
 

tests/test_messagebuilder.py

@@ -200,6 +200,61 @@ def test_messagebuilder_system_fewshots():
     assert messages[5]["content"] == user_message_pm["message"]["content"]
 
 
+def test_messagebuilder_system_fewshotstools():
+    messages = build_messages(
+        model="gpt-35-turbo",
+        system_prompt=system_message_short["message"]["content"],
+        new_user_content=user_message_pm["message"]["content"],
+        past_messages=[],
+        few_shots=[
+            {"role": "user", "content": "good options for climbing gear that can be used outside?"},
+            {
+                "role": "assistant",
+                "tool_calls": [
+                    {
+                        "id": "call_abc123",
+                        "type": "function",
+                        "function": {
+                            "arguments": '{"search_query":"climbing gear outside"}',
+                            "name": "search_database",
+                        },
+                    }
+                ],
+            },
+            {
+                "role": "tool",
+                "tool_call_id": "call_abc123",
+                "content": "Search results for climbing gear that can be used outside: ...",
+            },
+            {"role": "user", "content": "are there any shoes less than $50?"},
+            {
+                "role": "assistant",
+                "tool_calls": [
+                    {
+                        "id": "call_abc456",
+                        "type": "function",
+                        "function": {
+                            "arguments": '{"search_query":"shoes","price_filter":{"comparison_operator":"<","value":50}}',
+                            "name": "search_database",
+                        },
+                    }
+                ],
+            },
+            {"role": "tool", "tool_call_id": "call_abc456", "content": "Search results for shoes cheaper than 50: ..."},
+        ],
+    )
+    # Make sure messages are in the right order
+    assert messages[0]["role"] == "system"
+    assert messages[1]["role"] == "user"
+    assert messages[2]["role"] == "assistant"
+    assert messages[3]["role"] == "tool"
+    assert messages[4]["role"] == "user"
+    assert messages[5]["role"] == "assistant"
+    assert messages[6]["role"] == "tool"
+    assert messages[7]["role"] == "user"
+    assert messages[7]["content"] == user_message_pm["message"]["content"]
+
+
 def test_messagebuilder_system_tools():
     """Tests that the system message token count is considered."""
     messages = build_messages(