Few shots for tools

Author: pamelafox

src/openai_messages_token_helper/message_builder.py

@@ -10,6 +10,7 @@
     ChatCompletionNamedToolChoiceParam,
     ChatCompletionRole,
     ChatCompletionSystemMessageParam,
+    ChatCompletionToolMessageParam,
     ChatCompletionToolParam,
     ChatCompletionUserMessageParam,
 )
@@ -48,7 +49,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]],
+        index: int = 0,
+        tool_calls: Optional[list[ChatCompletionToolParam]] = None,
+        tool_call_id: Optional[str] = None,
     ):
         """
         Inserts a message into the conversation at the specified index,
@@ -63,8 +69,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 +114,11 @@ 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")
+        message_builder.insert_message(
+            shot["role"], shot.get("content"), tool_calls=shot.get("tool_calls"), tool_call_id=shot.get("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(