Support AsyncParallelFunctionCall in message_to_X_message (#406)

* Add AsyncParallelFunctionCall to message_to_X_message * Add tests
jackmpcollins · Jan 27, 2025 · 5f5119f · 5f5119f
1 parent 20b36b9
commit 5f5119f
Show file tree

Hide file tree

Showing 4 changed files with 81 additions and 3 deletions.
diff --git a/src/magentic/chat_model/anthropic_chat_model.py b/src/magentic/chat_model/anthropic_chat_model.py
@@ -33,7 +33,12 @@
     StreamParser,
     StreamState,
 )
-from magentic.function_call import FunctionCall, ParallelFunctionCall, _create_unique_id
+from magentic.function_call import (
+    AsyncParallelFunctionCall,
+    FunctionCall,
+    ParallelFunctionCall,
+    _create_unique_id,
+)
 from magentic.streaming import AsyncStreamedStr, StreamedStr
 from magentic.vision import UserImageMessage
 
@@ -205,6 +210,15 @@ def _(message: AssistantMessage[Any]) -> MessageParam:
 
 @async_message_to_anthropic_message.register(AssistantMessage)
 async def _(message: AssistantMessage[Any]) -> MessageParam:
+    if isinstance(message.content, AsyncParallelFunctionCall):
+        return {
+            "role": AnthropicMessageRole.ASSISTANT.value,
+            "content": [
+                _function_call_to_tool_call_block(function_call)
+                async for function_call in message.content
+            ],
+        }
+
     if isinstance(message.content, AsyncStreamedResponse):
         content_blocks: list[TextBlockParam | ToolUseBlockParam] = []
         async for item in message.content:

diff --git a/src/magentic/chat_model/openai_chat_model.py b/src/magentic/chat_model/openai_chat_model.py
@@ -45,7 +45,12 @@
     StreamParser,
     StreamState,
 )
-from magentic.function_call import FunctionCall, ParallelFunctionCall, _create_unique_id
+from magentic.function_call import (
+    AsyncParallelFunctionCall,
+    FunctionCall,
+    ParallelFunctionCall,
+    _create_unique_id,
+)
 from magentic.streaming import AsyncStreamedStr, StreamedStr
 from magentic.vision import UserImageMessage
 
@@ -200,6 +205,15 @@ def _(message: AssistantMessage[Any]) -> ChatCompletionMessageParam:
 
 @async_message_to_openai_message.register(AssistantMessage)
 async def _(message: AssistantMessage[Any]) -> ChatCompletionMessageParam:
+    if isinstance(message.content, AsyncParallelFunctionCall):
+        return {
+            "role": OpenaiMessageRole.ASSISTANT.value,
+            "tool_calls": [
+                _function_call_to_tool_call_block(function_call)
+                async for function_call in message.content
+            ],
+        }
+
     if isinstance(message.content, AsyncStreamedResponse):
         content: list[str] = []
         function_calls: list[FunctionCall[Any]] = []

diff --git a/tests/chat_model/test_anthropic_chat_model.py b/tests/chat_model/test_anthropic_chat_model.py
@@ -128,6 +128,30 @@ def test_message_to_anthropic_message(message, expected_anthropic_message):
 
 async_message_to_anthropic_message_test_cases = [
     *message_to_anthropic_message_test_cases,
+    (
+        AssistantMessage(
+            AsyncParallelFunctionCall(
+                async_iter([FunctionCall(plus, 1, 2), FunctionCall(plus, 3, 4)])
+            )
+        ),
+        {
+            "role": "assistant",
+            "content": [
+                {
+                    "type": "tool_use",
+                    "id": ANY,
+                    "name": "plus",
+                    "input": {"a": 1, "b": 2},
+                },
+                {
+                    "type": "tool_use",
+                    "id": ANY,
+                    "name": "plus",
+                    "input": {"a": 3, "b": 4},
+                },
+            ],
+        },
+    ),
     (
         AssistantMessage(
             AsyncStreamedResponse(

diff --git a/tests/chat_model/test_openai_chat_model.py b/tests/chat_model/test_openai_chat_model.py
@@ -27,7 +27,11 @@
     async_message_to_openai_message,
     message_to_openai_message,
 )
-from magentic.function_call import FunctionCall, ParallelFunctionCall
+from magentic.function_call import (
+    AsyncParallelFunctionCall,
+    FunctionCall,
+    ParallelFunctionCall,
+)
 from magentic.streaming import AsyncStreamedStr, StreamedStr, async_iter
 
 
@@ -137,6 +141,28 @@ def test_message_to_openai_message(message, expected_openai_message):
 
 async_message_to_openai_message_test_cases = [
     *message_to_openai_message_test_cases,
+    (
+        AssistantMessage(
+            AsyncParallelFunctionCall(
+                async_iter([FunctionCall(plus, 1, 2), FunctionCall(plus, 3, 4)])
+            )
+        ),
+        {
+            "role": "assistant",
+            "tool_calls": [
+                {
+                    "id": ANY,
+                    "type": "function",
+                    "function": {"name": "plus", "arguments": '{"a":1,"b":2}'},
+                },
+                {
+                    "id": ANY,
+                    "type": "function",
+                    "function": {"name": "plus", "arguments": '{"a":3,"b":4}'},
+                },
+            ],
+        },
+    ),
     (
         AssistantMessage(
             AsyncStreamedResponse(