Add parallel tool call support to OpenAI integration and AI driver

wrmsr · wrmsr · commit dcd54526aa70 · 2026-05-26T13:46:20.000-07:00
Implemented grouping of AI messages to generate a single assistant message with multiple tool calls, added `parallel_tool_calls` flag in OpenAI request construction, and refactored tool execution to batch parallel results. Updated CLI formatting, added necessary imports, and introduced comprehensive tests for request message generation, request flag behavior, and streaming with parallel tools. Adjusted driver logic to handle parallel tool messages without sequential loops.
diff --git a/omlish/text/docwrap/cli.py b/omlish/text/docwrap/cli.py
@@ -75,6 +75,7 @@ def _main(argv: ta.Sequence[str] | None = None) -> None:
 
     parser.add_argument('-s', '--start-line', type=int)
     parser.add_argument('-e', '--end-line', type=int)
+
     parser.add_argument('-i', '--in-place', action='store_true')
 
     args = parser.parse_args(argv)
diff --git a/ommlds/minichain/backends/openai/format.py b/ommlds/minichain/backends/openai/format.py
@@ -1,3 +1,4 @@
+import itertools
 import typing as ta
 
 from omlish import cached
@@ -40,48 +41,65 @@
 def build_oai_request_msgs(mc_chat: Chat) -> ta.Sequence[pt.ChatCompletionMessage]:
     oai_msgs: list[pt.ChatCompletionMessage] = []
 
-    for mc_msg in mc_chat:
-        if isinstance(mc_msg, SystemMessage):
-            oai_msgs.append(pt.SystemChatCompletionMessage(
-                content=check.isinstance(mc_msg.c, str),
-            ))
-
-        elif isinstance(mc_msg, AiMessage):
-            oai_msgs.append(pt.AssistantChatCompletionMessage(
-                content=check.isinstance(mc_msg.c, (str, None)),
-            ))
-
-        elif isinstance(mc_msg, ToolUseMessage):
-            oai_msgs.append(pt.AssistantChatCompletionMessage(
-                tool_calls=[pt.AssistantChatCompletionMessage.ToolCall(
-                    id=check.not_none(mc_msg.tu.id),
-                    function=pt.AssistantChatCompletionMessage.ToolCall.Function(
-                        arguments=check.not_none(mc_msg.tu.raw_args),
-                        name=mc_msg.tu.name,
-                    ),
-                )],
-            ))
-
-        elif isinstance(mc_msg, UserMessage):
-            oai_msgs.append(pt.UserChatCompletionMessage(
-                content=render_content_str(mc_msg.c),
-            ))
-
-        elif isinstance(mc_msg, ToolUseResultMessage):
-            tc: str
-            if isinstance(mc_msg.tur.c, str):
-                tc = mc_msg.tur.c
-            elif isinstance(mc_msg.tur.c, JsonContent):
-                tc = json.dumps_compact(mc_msg.tur.c)
-            else:
-                raise TypeError(mc_msg.tur.c)
-            oai_msgs.append(pt.ToolChatCompletionMessage(
-                tool_call_id=check.not_none(mc_msg.tur.id),
-                content=tc,
-            ))
+    for _, g in itertools.groupby(mc_chat, lambda mc_m: isinstance(mc_m, AnyAiMessage)):
+        mc_msgs = list(g)
+
+        if isinstance(mc_msgs[0], AnyAiMessage):
+            tups: list[tuple[AiMessage | None, list[ToolUseMessage]]] = []
+            for mc_msg in mc_msgs:
+                if isinstance(mc_msg, AiMessage):
+                    tups.append((mc_msg, []))
+
+                elif isinstance(mc_msg, ToolUseMessage):
+                    if not tups:
+                        tups.append((None, []))
+                    tups[-1][1].append(mc_msg)
+
+                else:
+                    raise TypeError(mc_msg)
+
+            for mc_ai_msg, mc_tu_msgs in tups:
+                oai_msgs.append(pt.AssistantChatCompletionMessage(
+                    content=check.isinstance(mc_ai_msg.c, (str, None)) if mc_ai_msg is not None else None,
+                    tool_calls=[
+                        pt.AssistantChatCompletionMessage.ToolCall(
+                            id=check.not_none(mc_tu_msg.tu.id),
+                            function=pt.AssistantChatCompletionMessage.ToolCall.Function(
+                                arguments=check.not_none(mc_tu_msg.tu.raw_args),
+                                name=mc_tu_msg.tu.name,
+                            ),
+                        )
+                        for mc_tu_msg in mc_tu_msgs
+                    ] if mc_tu_msgs else None,
+                ))
 
         else:
-            raise TypeError(mc_msg)
+            for mc_msg in mc_msgs:
+                if isinstance(mc_msg, SystemMessage):
+                    oai_msgs.append(pt.SystemChatCompletionMessage(
+                        content=check.isinstance(mc_msg.c, str),
+                    ))
+
+                elif isinstance(mc_msg, UserMessage):
+                    oai_msgs.append(pt.UserChatCompletionMessage(
+                        content=render_content_str(mc_msg.c),
+                    ))
+
+                elif isinstance(mc_msg, ToolUseResultMessage):
+                    tc: str
+                    if isinstance(mc_msg.tur.c, str):
+                        tc = mc_msg.tur.c
+                    elif isinstance(mc_msg.tur.c, JsonContent):
+                        tc = json.dumps_compact(mc_msg.tur.c)
+                    else:
+                        raise TypeError(mc_msg.tur.c)
+                    oai_msgs.append(pt.ToolChatCompletionMessage(
+                        tool_call_id=check.not_none(mc_msg.tur.id),
+                        content=tc,
+                    ))
+
+                else:
+                    raise TypeError(mc_msg)
 
     return oai_msgs
 
@@ -228,6 +246,7 @@ def oai_request(self) -> pt.ChatCompletionRequest:
             messages=build_oai_request_msgs(self._chat),
             top_p=1,
             tools=tools or None,
+            parallel_tool_calls=True if (tools and not (self._mandatory_kwargs or {}).get('stream')) else None,
             frequency_penalty=0.0,
             presence_penalty=0.0,
             **po.kwargs,
diff --git a/ommlds/minichain/backends/openai/tests/test_chat.py b/ommlds/minichain/backends/openai/tests/test_chat.py
@@ -8,6 +8,7 @@
 from omlish.http import all as http
 from omlish.secrets.tests.harness import HarnessSecrets
 
+from .....backends.openai import protocol as pt
 from ....chat.choices.services import ChatChoicesRequest
 from ....chat.choices.services import ChatChoicesService
 from ....chat.messages import Message
@@ -22,9 +23,87 @@
 from ....tools.types import ToolDtype
 from ....tools.types import ToolParam
 from ....tools.types import ToolSpec
+from ....tools.types import ToolUse
 from ....tools.types import ToolUseResult
 from ....wrappers.updateoptions import UpdateOptionsService
 from ..chat import OpenaiChatChoicesService
+from ..format import OpenaiChatRequestHandler
+from ..format import build_oai_request_msgs
+
+
+def test_openai_parallel_tool_call_request_messages():
+    chat: list[Message] = [
+        UserMessage('Use both tools.'),
+        ToolUseMessage(ToolUse(
+            id='call_1',
+            name='first_tool',
+            args={'value': 1},
+            raw_args='{"value":1}',
+        )),
+        ToolUseMessage(ToolUse(
+            id='call_2',
+            name='second_tool',
+            args={'value': 2},
+            raw_args='{"value":2}',
+        )),
+        ToolUseResultMessage(ToolUseResult(
+            id='call_1',
+            name='first_tool',
+            c='first result',
+        )),
+        ToolUseResultMessage(ToolUseResult(
+            id='call_2',
+            name='second_tool',
+            c='second result',
+        )),
+    ]
+
+    oai_msgs = list(build_oai_request_msgs(chat))
+
+    assert len(oai_msgs) == 4
+    assistant_msg = check.isinstance(oai_msgs[1], pt.AssistantChatCompletionMessage)
+    assert assistant_msg.content is None
+    tool_calls = list(check.not_none(assistant_msg.tool_calls))
+    assert [tc.id for tc in tool_calls] == ['call_1', 'call_2']
+    assert [tc.function.name for tc in tool_calls] == ['first_tool', 'second_tool']
+    assert [tc.function.arguments for tc in tool_calls] == ['{"value":1}', '{"value":2}']
+
+    tool_result_msgs = [
+        check.isinstance(oai_msg, pt.ToolChatCompletionMessage)
+        for oai_msg in oai_msgs[2:]
+    ]
+    assert [m.tool_call_id for m in tool_result_msgs] == ['call_1', 'call_2']
+
+
+def test_openai_parallel_tool_calls_enabled_with_tools():
+    tool_spec = ToolSpec(
+        'get_weather',
+        params=[
+            ToolParam(
+                'location',
+                type=ToolDtype.of(str),
+            ),
+        ],
+    )
+
+    req = OpenaiChatRequestHandler(
+        [UserMessage('What is the weather in Seattle?')],
+        Tool(tool_spec),
+        model='gpt-test',
+    ).oai_request()
+
+    assert req.tools is not None
+    assert req.parallel_tool_calls is True
+
+
+def test_openai_parallel_tool_calls_omitted_without_tools():
+    req = OpenaiChatRequestHandler(
+        [UserMessage('Hi!')],
+        model='gpt-test',
+    ).oai_request()
+
+    assert req.tools is None
+    assert req.parallel_tool_calls is None
 
 
 @pytest.mark.online
diff --git a/ommlds/minichain/backends/openai/tests/test_stream.py b/ommlds/minichain/backends/openai/tests/test_stream.py
@@ -85,7 +85,42 @@ async def test_openai_stream_tools(harness):
         [
             SystemMessage("You are a helpful agent. Use any tools available to you to answer the user's questions."),
             UserMessage('What is the weather in Seattle?'),
-            UserMessage(''),
+        ],
+        [
+            Tool(tool_spec),
+        ],
+    )
+
+    async with (await llm.invoke(foo_req)).v as it:
+        async for o in it:
+            print(o)
+        print(it.outputs)
+
+
+@pytest.mark.asyncs('asyncio')
+@pytest.mark.online
+async def test_openai_stream_parallel_tools(harness):
+    tool_spec = ToolSpec(
+        'get_weather',
+        params=[
+            ToolParam(
+                'location',
+                type=ToolDtype.of(str),
+                desc='The location to get the weather for.',
+            ),
+        ],
+        desc='Gets the weather in the given location.',
+    )
+
+    llm = OpenaiChatChoicesStreamService(
+        ApiKey(harness[HarnessSecrets].get_or_skip('openai_api_key').reveal()),
+    )
+
+    foo_req: ChatChoicesStreamRequest
+    foo_req = ChatChoicesStreamRequest(
+        [
+            SystemMessage("You are a helpful agent. Use any tools available to you to answer the user's questions."),
+            UserMessage('What is the weather in Seattle? Also, what is the weather in San Francisco?'),
         ],
         [
             Tool(tool_spec),
diff --git a/ommlds/minichain/drivers/ai/tools.py b/ommlds/minichain/drivers/ai/tools.py
@@ -47,28 +47,29 @@ async def generate_ai_chat(self, args: GenerateAiChatArgs) -> Chat:
                 dc.replace(args, chat=[*args.chat, *out]),
             )
 
-            cont = False
+            out.extend(new)
 
-            for msg in new:
-                out.append(msg)
+            tool_use_messages = [
+                msg
+                for msg in new
+                if isinstance(msg, ToolUseMessage)
+            ]
 
-                if isinstance(msg, ToolUseMessage):
-                    use = msg.tu
+            for msg in tool_use_messages:
+                use = msg.tu
 
-                    tce = self._catalog.by_name[check.non_empty_str(use.name)]
+                tce = self._catalog.by_name[check.non_empty_str(use.name)]
 
-                    trr = await self._executor.execute_tool_use(ToolUseExecution(
-                        msg.tu,
-                        tce,
-                    ))
+                trr = await self._executor.execute_tool_use(ToolUseExecution(
+                    msg.tu,
+                    tce,
+                ))
 
-                    trm = ToolUseResultMessage(trr)
+                trm = ToolUseResultMessage(trr)
 
-                    trm = check.isinstance(check.single(self._mt.transform(trm)), ToolUseResultMessage)
+                trm = check.isinstance(check.single(self._mt.transform(trm)), ToolUseResultMessage)
 
-                    out.append(trm)
+                out.append(trm)
 
-                    cont = True
-
-            if not cont:
+            if not tool_use_messages:
                 return out
diff --git a/ommlds/minichain/drivers/tests/test_dummy.py b/ommlds/minichain/drivers/tests/test_dummy.py
