diff --git a/artemis/services/llm.py b/artemis/services/llm.py index ec91e6fa0..79fc2d167 100644 --- a/artemis/services/llm.py +++ b/artemis/services/llm.py @@ -624,6 +624,12 @@ def _endpoint_key(self) -> str: return type(self.base_model).__name__ def _traced_call(self, args: tuple, kwargs: dict) -> tuple[tuple, dict, Any]: + # Every request leaves through here: keep parallel tool results + # contiguous so OpenAI-compatible providers accept image carriers. + from artemis.tools.tool_wrapper import order_tool_batch_messages + + if args: + args = (order_tool_batch_messages(args[0]), *args[1:]) trace_id = None if self.ctx and self.ctx.data_engine: trace_id = CURRENT_TRACE_ID.get() diff --git a/artemis/tools/tool_wrapper.py b/artemis/tools/tool_wrapper.py index 89f0908a6..c2e3932cb 100644 --- a/artemis/tools/tool_wrapper.py +++ b/artemis/tools/tool_wrapper.py @@ -20,7 +20,7 @@ from typing import Annotated, Any, Literal, get_args, get_origin import uuid -from langchain_core.messages import BaseMessage, HumanMessage, ToolMessage +from langchain_core.messages import AIMessage, BaseMessage, HumanMessage, ToolMessage from langchain_core.tools import BaseTool from langchain_core.tools.base import InjectedToolCallId from langgraph.prebuilt import InjectedState @@ -268,3 +268,33 @@ def tool_result_messages( ToolMessage(tool_call_id=tool_call_id, content=text, status=status, **extra), HumanMessage(content=[{"type": "text", "text": caption}, *images]), ] + + +def order_tool_batch_messages(messages: list) -> list: + """Move "human" image carriers after every tool result of their turn. + + Agent loops extend history with ``tool_result_messages`` once per tool + call, so a turn with parallel screenshot calls becomes Tool, Human, Tool, + Human. OpenAI-compatible APIs require all tool results to directly follow + the assistant turn and reject that with a 400, so each run of tool results + and carriers after a tool-calling AIMessage is reordered to all + ToolMessages first, then the HumanMessages in their original order. + """ + if not isinstance(messages, list): + return messages + ordered: list = [] + i = 0 + while i < len(messages): + msg = messages[i] + ordered.append(msg) + i += 1 + if not (isinstance(msg, AIMessage) and msg.tool_calls): + continue + tools: list = [] + humans: list = [] + while i < len(messages) and isinstance(messages[i], (ToolMessage, HumanMessage)): + (tools if isinstance(messages[i], ToolMessage) else humans).append(messages[i]) + i += 1 + ordered.extend(tools) + ordered.extend(humans) + return ordered diff --git a/tests/unit/tools/history/test_result_messages.py b/tests/unit/tools/history/test_result_messages.py index 59829339d..6424afa7e 100644 --- a/tests/unit/tools/history/test_result_messages.py +++ b/tests/unit/tools/history/test_result_messages.py @@ -118,3 +118,80 @@ def test_explicit_carrier_wins_over_the_llm(): assert len(msgs) == 2 msgs = tool_result_messages("tc", SHOT, image_carrier="tool", llm=_llm("openai")) assert len(msgs) == 1 + + +# --- Parallel tool calls: image carriers must follow the whole tool batch ------- + + +def _parallel_turn(n): + from langchain_core.messages import AIMessage + + calls = [{"name": "get_step_screenshot", "args": {}, "id": f"call_{i}"} for i in range(n)] + return AIMessage(content="", tool_calls=calls) + + +def test_order_tool_batch_moves_image_carriers_after_all_tool_results(): + """OpenAI-compatible APIs reject a user turn between the tool results of + one assistant turn (400: tool_call_ids did not have response messages).""" + from artemis.tools.tool_wrapper import order_tool_batch_messages + + ai = _parallel_turn(2) + per_call = [ + tool_result_messages(f"call_{i}", SHOT, name="get_step_screenshot", llm=_llm("openai")) + for i in range(2) + ] + interleaved = [ai, *per_call[0], *per_call[1]] + + ordered = order_tool_batch_messages(interleaved) + + assert [type(m) for m in ordered] == [ + type(ai), + ToolMessage, + ToolMessage, + HumanMessage, + HumanMessage, + ] + assert [m.tool_call_id for m in ordered[1:3]] == ["call_0", "call_1"] + + +def test_order_tool_batch_keeps_valid_history_unchanged(): + from artemis.tools.tool_wrapper import order_tool_batch_messages + + ai = _parallel_turn(2) + history = [ + HumanMessage(content="check"), + ai, + ToolMessage(tool_call_id="call_0", content="a"), + ToolMessage(tool_call_id="call_1", content="b"), + HumanMessage(content="This is your final iteration."), + ] + + assert order_tool_batch_messages(history) == history + + +@pytest.mark.asyncio +async def test_wrapper_sends_tool_batch_before_image_carriers(): + from artemis.services.llm import RobustChatModelWrapper + + seen = {} + + class _Model: + async def ainvoke(self, messages, **kwargs): + seen["messages"] = messages + return "ok" + + ai = _parallel_turn(2) + interleaved = [ai] + for i in range(2): + interleaved += tool_result_messages( + f"call_{i}", SHOT, name="get_step_screenshot", llm=_llm("openai") + ) + + await RobustChatModelWrapper(_Model()).ainvoke(interleaved) + + assert [type(m).__name__ for m in seen["messages"][1:]] == [ + "ToolMessage", + "ToolMessage", + "HumanMessage", + "HumanMessage", + ]