Browse Source

feat: add provider-compatible tool transcripts

Problem: ChatAgent direct tools lacked a valid assistant tool-call and tool-result transcript, while the client implicitly forced every single tool. Risk: stricter local transcript validation may reject previously tolerated malformed internal histories; dual-agent Runtime still sends tools=[].
zhenyu.hu 2 weeks ago
parent
commit
65c63a5256

+ 5 - 0
src/agent_lab/application/event_agent.py

@@ -20,6 +20,7 @@ class EventAgentChatClient(Protocol):
         messages: list[ChatMessage],
         tools: list[dict[str, Any]],
         params: AgentParams,
+        tool_choice: dict[str, Any] | None = None,
     ) -> AsyncIterator[StreamItem]:
         ...
 
@@ -161,6 +162,10 @@ class EventAgent:
             messages=messages,
             tools=[tool],
             params=params,
+            tool_choice={
+                "type": "function",
+                "function": {"name": event.name},
+            },
         ):
             if item.kind == "raw_chunk" and item.raw_chunk is not None:
                 raw_chunks.append(item.raw_chunk)

+ 1 - 0
src/agent_lab/application/runtime.py

@@ -22,6 +22,7 @@ class ChatClient(Protocol):
         messages: list[ChatMessage],
         tools: list[dict[str, Any]],
         params: AgentParams,
+        tool_choice: dict[str, Any] | None = None,
     ) -> AsyncIterator[StreamItem]:
         ...
 

+ 10 - 1
src/agent_lab/domain/messages.py

@@ -1,7 +1,7 @@
 from dataclasses import dataclass
 from typing import Any, Literal
 
-from pydantic import BaseModel, ConfigDict
+from pydantic import BaseModel, ConfigDict, Field, model_validator
 
 from agent_lab.domain.events import ToolCallEvent
 
@@ -13,6 +13,15 @@ class ChatMessage(BaseModel):
     content: str
     name: str | None = None
     tool_call_id: str | None = None
+    tool_calls: list[ToolCallEvent] = Field(default_factory=list)
+
+    @model_validator(mode="after")
+    def validate_tool_fields(self) -> "ChatMessage":
+        if self.tool_calls and self.role != "assistant":
+            raise ValueError("tool_calls require assistant role")
+        if self.role == "tool" and not self.tool_call_id:
+            raise ValueError("tool messages require tool_call_id")
+        return self
 
 
 class TokenUsage(BaseModel):

+ 58 - 6
src/agent_lab/infrastructure/chat_client.py

@@ -33,7 +33,9 @@ class OpenAICompatibleChatClient:
         messages: list[ChatMessage],
         tools: list[dict[str, Any]],
         params: AgentParams,
+        tool_choice: dict[str, Any] | None = None,
     ) -> AsyncIterator[StreamItem]:
+        self._validate_tool_transcript(messages)
         payload: dict[str, Any] = {
             "model": self._resolve_model(params.model),
             "messages": [self._serialize_message(message) for message in messages],
@@ -45,11 +47,8 @@ class OpenAICompatibleChatClient:
             payload["stream_options"] = {"include_usage": True}
         if tools:
             payload["tools"] = tools
-            if len(tools) == 1:
-                payload["tool_choice"] = {
-                    "type": "function",
-                    "function": {"name": tools[0]["function"]["name"]},
-                }
+        if tool_choice is not None:
+            payload["tool_choice"] = tool_choice
         if params.extra_body:
             payload.update(params.extra_body)
 
@@ -95,12 +94,65 @@ class OpenAICompatibleChatClient:
         requested_model = model.strip() if model else ""
         return requested_model or self.default_model
 
+    def _validate_tool_transcript(self, messages: list[ChatMessage]) -> None:
+        pending: dict[str, None] = {}
+        replied_ids: set[str] = set()
+
+        for message in messages:
+            if pending:
+                if message.role != "tool":
+                    unresolved = ", ".join(pending)
+                    raise ValueError(
+                        f"unresolved tool calls before {message.role} message: "
+                        f"{unresolved}"
+                    )
+
+                tool_call_id = message.tool_call_id
+                assert tool_call_id is not None
+                if tool_call_id in replied_ids:
+                    raise ValueError(f"duplicate tool reply: {tool_call_id}")
+                if tool_call_id not in pending:
+                    raise ValueError(f"mismatched tool_call_id: {tool_call_id}")
+                del pending[tool_call_id]
+                replied_ids.add(tool_call_id)
+                continue
+
+            if message.role == "tool":
+                tool_call_id = message.tool_call_id
+                assert tool_call_id is not None
+                if tool_call_id in replied_ids:
+                    raise ValueError(f"duplicate tool reply: {tool_call_id}")
+                raise ValueError(f"orphan tool reply: {tool_call_id}")
+
+            if message.tool_calls:
+                pending = {tool_call.id: None for tool_call in message.tool_calls}
+
+        if pending:
+            unresolved = ", ".join(pending)
+            raise ValueError(f"unresolved tool calls at end: {unresolved}")
+
     def _serialize_message(self, message: ChatMessage) -> dict[str, Any]:
         if message.role == "tool":
-            raise ValueError("tool messages are internal and cannot be sent to ChatAgent")
+            return {
+                "role": "tool",
+                "content": message.content,
+                "tool_call_id": message.tool_call_id,
+            }
 
         payload = {
             "role": message.role,
             "content": message.content,
         }
+        if message.tool_calls:
+            payload["tool_calls"] = [
+                {
+                    "id": tool_call.id,
+                    "type": "function",
+                    "function": {
+                        "name": tool_call.name,
+                        "arguments": tool_call.raw_arguments,
+                    },
+                }
+                for tool_call in message.tool_calls
+            ]
         return payload

+ 12 - 0
tests/test_debug_runtime.py

@@ -120,6 +120,7 @@ class FakeChatClient:
         messages: list[ChatMessage],
         tools: list[dict],
         params: AgentParams,
+        tool_choice: dict[str, Any] | None = None,
     ) -> AsyncIterator[StreamItem]:
         if tools:
             yield StreamItem.raw_response_chunk(
@@ -178,6 +179,7 @@ class WrongSourceEventChatClient:
         messages: list[ChatMessage],
         tools: list[dict],
         params: AgentParams,
+        tool_choice: dict[str, Any] | None = None,
     ) -> AsyncIterator[StreamItem]:
         if tools:
             yield _event_tool_call_from_tools(tools, messages)
@@ -209,6 +211,7 @@ class StrictHistoryChatClient:
         messages: list[ChatMessage],
         tools: list[dict],
         params: AgentParams,
+        tool_choice: dict[str, Any] | None = None,
     ) -> AsyncIterator[StreamItem]:
         if tools:
             yield StreamItem.raw_response_chunk(
@@ -266,6 +269,7 @@ class MultiEventChatClient:
         messages: list[ChatMessage],
         tools: list[dict],
         params: AgentParams,
+        tool_choice: dict[str, Any] | None = None,
     ) -> AsyncIterator[StreamItem]:
         if tools:
             yield _event_tool_call_from_tools(tools, messages)
@@ -305,6 +309,7 @@ class ToolCapturingChatClient:
         messages: list[ChatMessage],
         tools: list[dict],
         params: AgentParams,
+        tool_choice: dict[str, Any] | None = None,
     ) -> AsyncIterator[StreamItem]:
         self.messages = list(messages)
         self.tools = list(tools)
@@ -322,6 +327,7 @@ class EventLoopLimitChatClient:
         messages: list[ChatMessage],
         tools: list[dict],
         params: AgentParams,
+        tool_choice: dict[str, Any] | None = None,
     ) -> AsyncIterator[StreamItem]:
         if tools:
             yield _event_tool_call_from_tools(tools, messages)
@@ -349,6 +355,7 @@ class RoundStatsChatClient:
         messages: list[ChatMessage],
         tools: list[dict],
         params: AgentParams,
+        tool_choice: dict[str, Any] | None = None,
     ) -> AsyncIterator[StreamItem]:
         yield StreamItem.message_delta("hello")
         yield StreamItem.usage_item(
@@ -370,6 +377,7 @@ class EventRoundStatsChatClient:
         messages: list[ChatMessage],
         tools: list[dict],
         params: AgentParams,
+        tool_choice: dict[str, Any] | None = None,
     ) -> AsyncIterator[StreamItem]:
         if tools:
             yield StreamItem.raw_response_chunk(
@@ -446,6 +454,7 @@ class ContextBoundaryChatClient:
         messages: list[ChatMessage],
         tools: list[dict],
         params: AgentParams,
+        tool_choice: dict[str, Any] | None = None,
     ) -> AsyncIterator[StreamItem]:
         if tools:
             self.event_agent_messages.append(list(messages))
@@ -476,6 +485,7 @@ class TwoTurnSessionChatClient:
         messages: list[ChatMessage],
         tools: list[dict],
         params: AgentParams,
+        tool_choice: dict[str, Any] | None = None,
     ) -> AsyncIterator[StreamItem]:
         if tools:
             yield _event_tool_call_from_tools(tools, messages)
@@ -507,6 +517,7 @@ class SlowAfterEventChatClient:
         messages: list[ChatMessage],
         tools: list[dict],
         params: AgentParams,
+        tool_choice: dict[str, Any] | None = None,
     ) -> AsyncIterator[StreamItem]:
         if tools:
             yield _event_tool_call_from_tools(tools, messages)
@@ -766,6 +777,7 @@ async def test_runtime_audit_includes_model_params_prompts_results_and_usage():
         "content": "You are a debugger.",
         "name": None,
         "tool_call_id": None,
+        "tool_calls": [],
     }
 
     chat_response = next(

+ 9 - 0
tests/test_event_agent.py

@@ -20,12 +20,14 @@ class ToolCallingChatClient:
         messages: list[ChatMessage],
         tools: list[dict],
         params: AgentParams,
+        tool_choice: dict | None = None,
     ) -> AsyncIterator[StreamItem]:
         self.calls.append(
             {
                 "messages": list(messages),
                 "tools": list(tools),
                 "params": params,
+                "tool_choice": tool_choice,
             }
         )
         tool_name = tools[0]["function"]["name"]
@@ -65,12 +67,14 @@ class NoToolCallChatClient:
         messages: list[ChatMessage],
         tools: list[dict],
         params: AgentParams,
+        tool_choice: dict | None = None,
     ) -> AsyncIterator[StreamItem]:
         self.calls.append(
             {
                 "messages": list(messages),
                 "tools": list(tools),
                 "params": params,
+                "tool_choice": tool_choice,
             }
         )
         yield StreamItem.message_delta("I should have called the tool.")
@@ -85,6 +89,7 @@ class WrongSourceToolCallChatClient:
         messages: list[ChatMessage],
         tools: list[dict],
         params: AgentParams,
+        tool_choice: dict | None = None,
     ) -> AsyncIterator[StreamItem]:
         yield self.item
 
@@ -133,6 +138,10 @@ async def test_event_agent_resolves_tool_arguments_with_llm_tool_call():
             },
         }
     ]
+    assert chat_client.calls[0]["tool_choice"] == {
+        "type": "function",
+        "function": {"name": "handoff_note"},
+    }
     assert chat_client.calls[0]["params"].model == "event-model"
     assert agent.raw_model_chunks([event]) == [
         {

+ 320 - 17
tests/test_websocket_api.py

@@ -425,6 +425,63 @@ def test_chat_message_allows_internal_tool_replies():
     assert message.tool_call_id == "call_1"
 
 
+def test_chat_message_preserves_assistant_tool_calls():
+    message = ChatMessage(
+        role="assistant",
+        content="I will inspect both sources.",
+        tool_calls=[
+            ToolCallEvent(
+                id="call_1",
+                name="mock_search",
+                arguments={"query": "latency docs"},
+                raw_arguments='{"query":"latency docs"}',
+            ),
+            ToolCallEvent(
+                id="call_2",
+                name="handoff_note",
+                arguments={"message": "inspect provider behavior"},
+                raw_arguments='{"message":"inspect provider behavior"}',
+            ),
+        ],
+    )
+
+    assert message.model_dump()["tool_calls"] == [
+        {
+            "id": "call_1",
+            "name": "mock_search",
+            "arguments": {"query": "latency docs"},
+            "raw_arguments": '{"query":"latency docs"}',
+        },
+        {
+            "id": "call_2",
+            "name": "handoff_note",
+            "arguments": {"message": "inspect provider behavior"},
+            "raw_arguments": '{"message":"inspect provider behavior"}',
+        },
+    ]
+
+
+def test_chat_message_rejects_tool_calls_on_non_assistant_role():
+    with pytest.raises(ValidationError, match="tool_calls require assistant role"):
+        ChatMessage(
+            role="user",
+            content="invalid",
+            tool_calls=[
+                ToolCallEvent(
+                    id="call_1",
+                    name="mock_search",
+                    arguments={},
+                    raw_arguments="{}",
+                )
+            ],
+        )
+
+
+def test_chat_message_rejects_tool_role_without_tool_call_id():
+    with pytest.raises(ValidationError, match="tool messages require tool_call_id"):
+        ChatMessage(role="tool", content="orphan result")
+
+
 class HistoryCapturingChatClient:
     def __init__(self) -> None:
         self.calls = 0
@@ -497,10 +554,7 @@ async def test_openai_chat_client_streams_sse_chunks_through_parser():
         assert payload["model"] == "model-x"
         assert payload["messages"] == [{"role": "user", "content": "hi"}]
         assert payload["tools"][0]["function"]["name"] == "handoff_note"
-        assert payload["tool_choice"] == {
-            "type": "function",
-            "function": {"name": "handoff_note"},
-        }
+        assert "tool_choice" not in payload
         assert payload["temperature"] == 0.3
         assert payload["max_tokens"] == 50
         assert request.headers["authorization"] == "Bearer test-key"
@@ -544,6 +598,47 @@ async def test_openai_chat_client_streams_sse_chunks_through_parser():
     ]
 
 
+@pytest.mark.asyncio
+async def test_openai_chat_client_serializes_explicit_tool_choice():
+    captured_payloads: list[dict] = []
+    forced_choice = {
+        "type": "function",
+        "function": {"name": "handoff_note"},
+    }
+
+    def handler(request: httpx.Request) -> httpx.Response:
+        captured_payloads.append(json.loads(request.content))
+        return httpx.Response(200, content=b"data: [DONE]\n\n")
+
+    async with httpx.AsyncClient(
+        transport=httpx.MockTransport(handler),
+        base_url="https://llm.test/v1",
+    ) as http_client:
+        client = OpenAICompatibleChatClient(
+            api_key="",
+            base_url="https://llm.test/v1",
+            default_model="provider-default",
+            request_timeout_seconds=5,
+            http_client=http_client,
+        )
+        [
+            item
+            async for item in client.stream_chat(
+                messages=[ChatMessage(role="user", content="prepare handoff")],
+                tools=[
+                    {
+                        "type": "function",
+                        "function": {"name": "handoff_note", "parameters": {}},
+                    }
+                ],
+                params=AgentParams(model="provider-default"),
+                tool_choice=forced_choice,
+            )
+        ]
+
+    assert captured_payloads[0]["tool_choice"] == forced_choice
+
+
 @pytest.mark.asyncio
 async def test_openai_chat_client_flushes_one_provider_tool_call_at_done():
     def handler(request: httpx.Request) -> httpx.Response:
@@ -633,17 +728,122 @@ async def test_openai_chat_client_uses_default_model_when_request_model_is_blank
 
 
 @pytest.mark.asyncio
-async def test_openai_chat_client_rejects_internal_tool_messages_before_network():
-    network_called = False
+async def test_openai_chat_client_serializes_provider_tool_transcript():
+    captured_payloads: list[dict] = []
 
     def handler(request: httpx.Request) -> httpx.Response:
-        nonlocal network_called
-        network_called = True
+        captured_payloads.append(json.loads(request.content))
         return httpx.Response(
             200,
             content=b'data: {"choices":[{"delta":{"content":"done"},"finish_reason":null}]}\n\n',
         )
 
+    async with httpx.AsyncClient(
+        transport=httpx.MockTransport(handler),
+        base_url="https://llm.test/v1",
+    ) as http_client:
+        client = OpenAICompatibleChatClient(
+            api_key="",
+            base_url="https://llm.test/v1",
+            default_model="provider-default",
+            request_timeout_seconds=5,
+            http_client=http_client,
+        )
+        [
+            item
+            async for item in client.stream_chat(
+                messages=[
+                    ChatMessage(role="user", content="inspect both"),
+                    ChatMessage(
+                        role="assistant",
+                        content="I will inspect both sources.",
+                        tool_calls=[
+                            ToolCallEvent(
+                                id="call_1",
+                                name="mock_search",
+                                arguments={"query": "latency docs"},
+                                raw_arguments='{"query":"latency docs"}',
+                            ),
+                            ToolCallEvent(
+                                id="call_2",
+                                name="handoff_note",
+                                arguments={"message": "inspect provider behavior"},
+                                raw_arguments='{"message":"inspect provider behavior"}',
+                            ),
+                        ],
+                    ),
+                    ChatMessage(
+                        role="tool",
+                        content='{"results":[]}',
+                        name="mock_search",
+                        tool_call_id="call_1",
+                    ),
+                    ChatMessage(
+                        role="tool",
+                        content='{"message":"handled"}',
+                        name="handoff_note",
+                        tool_call_id="call_2",
+                    ),
+                    ChatMessage(role="user", content="continue"),
+                ],
+                tools=[],
+                params=AgentParams(
+                    model="provider-default",
+                    temperature=0.3,
+                    max_tokens=50,
+                ),
+            )
+        ]
+
+    assert captured_payloads[0]["messages"] == [
+        {"role": "user", "content": "inspect both"},
+        {
+            "role": "assistant",
+            "content": "I will inspect both sources.",
+            "tool_calls": [
+                {
+                    "id": "call_1",
+                    "type": "function",
+                    "function": {
+                        "name": "mock_search",
+                        "arguments": '{"query":"latency docs"}',
+                    },
+                },
+                {
+                    "id": "call_2",
+                    "type": "function",
+                    "function": {
+                        "name": "handoff_note",
+                        "arguments": '{"message":"inspect provider behavior"}',
+                    },
+                },
+            ],
+        },
+        {
+            "role": "tool",
+            "content": '{"results":[]}',
+            "tool_call_id": "call_1",
+        },
+        {
+            "role": "tool",
+            "content": '{"message":"handled"}',
+            "tool_call_id": "call_2",
+        },
+        {"role": "user", "content": "continue"},
+    ]
+
+
+async def _assert_tool_transcript_rejected(
+    messages: list[ChatMessage],
+    error_match: str,
+) -> None:
+    network_called = False
+
+    def handler(request: httpx.Request) -> httpx.Response:
+        nonlocal network_called
+        network_called = True
+        return httpx.Response(200)
+
     async with httpx.AsyncClient(
         transport=httpx.MockTransport(handler),
         base_url="https://llm.test/v1",
@@ -656,25 +856,128 @@ async def test_openai_chat_client_rejects_internal_tool_messages_before_network(
             http_client=http_client,
         )
 
-        with pytest.raises(ValueError, match="tool messages are internal"):
+        with pytest.raises(ValueError, match=error_match):
             [
                 item
                 async for item in client.stream_chat(
-                    messages=[
-                        ChatMessage(
-                            role="tool",
-                            content="internal tool result",
-                            tool_call_id="call_1",
-                        )
-                    ],
+                    messages=messages,
                     tools=[],
-                    params=AgentParams(model="provider-default", temperature=0.3, max_tokens=50),
+                    params=AgentParams(model="provider-default"),
                 )
             ]
 
     assert network_called is False
 
 
+@pytest.mark.asyncio
+async def test_openai_chat_client_rejects_orphan_tool_reply_before_network():
+    await _assert_tool_transcript_rejected(
+        [
+            ChatMessage(
+                role="tool",
+                content="orphan",
+                tool_call_id="call_1",
+            )
+        ],
+        "orphan tool reply: call_1",
+    )
+
+
+@pytest.mark.asyncio
+async def test_openai_chat_client_rejects_mismatched_tool_call_id_before_network():
+    await _assert_tool_transcript_rejected(
+        [
+            ChatMessage(
+                role="assistant",
+                content="",
+                tool_calls=[
+                    ToolCallEvent(
+                        id="call_1",
+                        name="mock_search",
+                        arguments={},
+                        raw_arguments="{}",
+                    )
+                ],
+            ),
+            ChatMessage(role="tool", content="wrong", tool_call_id="call_2"),
+        ],
+        "mismatched tool_call_id: call_2",
+    )
+
+
+@pytest.mark.asyncio
+async def test_openai_chat_client_rejects_duplicate_tool_reply_before_network():
+    await _assert_tool_transcript_rejected(
+        [
+            ChatMessage(
+                role="assistant",
+                content="",
+                tool_calls=[
+                    ToolCallEvent(
+                        id="call_1",
+                        name="mock_search",
+                        arguments={},
+                        raw_arguments="{}",
+                    )
+                ],
+            ),
+            ChatMessage(role="tool", content="first", tool_call_id="call_1"),
+            ChatMessage(role="tool", content="duplicate", tool_call_id="call_1"),
+        ],
+        "duplicate tool reply: call_1",
+    )
+
+
+@pytest.mark.asyncio
+async def test_openai_chat_client_rejects_non_tool_message_with_unresolved_calls():
+    await _assert_tool_transcript_rejected(
+        [
+            ChatMessage(
+                role="assistant",
+                content="checking",
+                tool_calls=[
+                    ToolCallEvent(
+                        id="call_1",
+                        name="mock_search",
+                        arguments={},
+                        raw_arguments="{}",
+                    ),
+                    ToolCallEvent(
+                        id="call_2",
+                        name="handoff_note",
+                        arguments={},
+                        raw_arguments="{}",
+                    ),
+                ],
+            ),
+            ChatMessage(role="tool", content="done", tool_call_id="call_1"),
+            ChatMessage(role="user", content="continue too early"),
+        ],
+        "unresolved tool calls before user message: call_2",
+    )
+
+
+@pytest.mark.asyncio
+async def test_openai_chat_client_rejects_unfinished_tool_calls_at_end():
+    await _assert_tool_transcript_rejected(
+        [
+            ChatMessage(
+                role="assistant",
+                content="checking",
+                tool_calls=[
+                    ToolCallEvent(
+                        id="call_1",
+                        name="mock_search",
+                        arguments={},
+                        raw_arguments="{}",
+                    )
+                ],
+            )
+        ],
+        "unresolved tool calls at end: call_1",
+    )
+
+
 def test_static_pre_message_role_selector_does_not_offer_tool():
     html = Path("src/agent_lab/presentation/static/index.html").read_text()