structured output w/ multiple pending tool calls

Author: sydney-runkle

libs/langchain_v1/langchain/agents/middleware_agent.py

@@ -5,7 +5,7 @@
 from typing import Any, cast
 
 from langchain_core.language_models.chat_models import BaseChatModel
-from langchain_core.messages import AIMessage, SystemMessage, ToolMessage
+from langchain_core.messages import AIMessage, AnyMessage, SystemMessage, ToolMessage
 from langchain_core.runnables import Runnable
 from langchain_core.tools import BaseTool
 from langgraph.constants import END, START
@@ -219,8 +219,6 @@ def _handle_model_output(state: dict[str, Any], output: AIMessage) -> dict[str,
             if not output.tool_calls and native_output_binding:
                 structured_response = native_output_binding.parse(output)
                 return {"messages": [output], "response": structured_response}
-            if state.get("response") is not None:
-                return {"messages": [output], "response": None}
             return {"messages": [output]}
 
         # Handle structured output with tools strategy
@@ -418,7 +416,7 @@ async def amodel_request(state: dict[str, Any]) -> dict[str, Any]:
     if tool_node is not None:
         graph.add_conditional_edges(
             "tools",
-            _make_tools_to_model_edge(tool_node, first_node),
+            _make_tools_to_model_edge(tool_node, first_node, structured_output_tools),
             [first_node, END],
         )
         graph.add_conditional_edges(
@@ -482,38 +480,30 @@ def _resolve_jump(jump_to: JumpTo | None, first_node: str) -> str | None:
     return None
 
 
+def _fetch_last_ai_and_tool_messages(
+    messages: list[AnyMessage],
+) -> tuple[AIMessage, list[ToolMessage]]:
+    last_ai_index: int
+    last_ai_message: AIMessage
+
+    for i in range(len(messages) - 1, -1, -1):
+        if isinstance(messages[i], AIMessage):
+            last_ai_index = i
+            last_ai_message = cast("AIMessage", messages[i])
+            break
+
+    tool_messages = [m for m in messages[last_ai_index + 1 :] if isinstance(m, ToolMessage)]
+    return last_ai_message, tool_messages
+
+
 def _make_model_to_tools_edge(
     first_node: str, structured_output_tools: dict[str, OutputToolBinding], tool_node: ToolNode
 ) -> Callable[[AgentState], str | list[Send] | None]:
     def model_to_tools(state: AgentState) -> str | list[Send] | None:
         if jump_to := state.get("jump_to"):
             return _resolve_jump(jump_to, first_node)
 
-        last_message = state["messages"][-1]
-
-        # Check if this is a ToolMessage from structured output - if so, end
-        # interesting, should we be auto ending here? should we execute other tools?
-        if isinstance(last_message, ToolMessage) and last_message.name in structured_output_tools:
-            return END
-
-        # Find the last AI message and all tool messages since said AI message
-        last_ai_index = None
-        last_ai_message: AIMessage
-        for i in range(len(state["messages"]) - 1, -1, -1):
-            if isinstance(state["messages"][i], AIMessage):
-                last_ai_index = i
-                last_ai_message = cast("AIMessage", state["messages"][i])
-                break
-
-        tool_messages = (
-            [
-                m.tool_call_id
-                for m in state["messages"][last_ai_index + 1 :]
-                if isinstance(m, ToolMessage)
-            ]
-            if last_ai_index is not None
-            else []
-        )
+        last_ai_message, tool_messages = _fetch_last_ai_and_tool_messages(state["messages"])
 
         pending_tool_calls = [
             c
@@ -538,17 +528,21 @@ def model_to_tools(state: AgentState) -> str | list[Send] | None:
 
 
 def _make_tools_to_model_edge(
-    tool_node: ToolNode, next_node: str
+    tool_node: ToolNode, next_node: str, structured_output_tools: dict[str, OutputToolBinding]
 ) -> Callable[[AgentState], str | None]:
     def tools_to_model(state: AgentState) -> str | None:
-        ai_message = [m for m in state["messages"] if isinstance(m, AIMessage)][-1]
+        last_ai_message, tool_messages = _fetch_last_ai_and_tool_messages(state["messages"])
+
         if all(
             tool_node.tools_by_name[c["name"]].return_direct
-            for c in ai_message.tool_calls
+            for c in last_ai_message.tool_calls
             if c["name"] in tool_node.tools_by_name
         ):
             return END
 
+        if any(t.name in structured_output_tools for t in tool_messages):
+            return END
+
         return next_node
 
     return tools_to_model

libs/langchain_v1/tests/unit_tests/agents/test_middleware_agent.py

@@ -4,6 +4,7 @@
 
 from syrupy.assertion import SnapshotAssertion
 
+from pydantic import BaseModel, Field
 from langchain_core.language_models import BaseChatModel
 from langchain_core.language_models.chat_models import BaseChatModel
 from langchain_core.messages import (
@@ -28,7 +29,7 @@
 from langgraph.checkpoint.memory import InMemorySaver
 from langgraph.constants import END
 from langgraph.graph.message import REMOVE_ALL_MESSAGES
-from langgraph.prebuilt.interrupt import ActionRequest
+from langchain.agents.structured_output import ToolStrategy
 
 from .messages import _AnyIdHumanMessage, _AnyIdToolMessage
 from .model import FakeToolCallingModel
@@ -1275,3 +1276,118 @@ def modify_model_request(self, request: ModelRequest, state: AgentState) -> Mode
     assert (
         result["messages"][2].content == "You are a helpful assistant.-Hello-remember to be nice!"
     )
+
+
+def test_tools_to_model_edge_with_structured_and_regular_tool_calls():
+    """Test that when there are both structured and regular tool calls, we execute regular and jump to END."""
+
+    class WeatherResponse(BaseModel):
+        """Weather response."""
+
+        temperature: float = Field(description="Temperature in fahrenheit")
+        condition: str = Field(description="Weather condition")
+
+    @tool
+    def regular_tool(query: str) -> str:
+        """A regular tool that returns a string."""
+        return f"Regular tool result for: {query}"
+
+    # Create a fake model that returns both structured and regular tool calls
+    class FakeModelWithBothToolCalls(FakeToolCallingModel):
+        def __init__(self):
+            super().__init__()
+            self.tool_calls = [
+                [
+                    ToolCall(
+                        name="WeatherResponse",
+                        args={"temperature": 72.0, "condition": "sunny"},
+                        id="structured_call_1",
+                    ),
+                    ToolCall(
+                        name="regular_tool", args={"query": "test query"}, id="regular_call_1"
+                    ),
+                ]
+            ]
+
+    # Create agent with both structured output and regular tools
+    agent = create_agent(
+        model=FakeModelWithBothToolCalls(),
+        tools=[regular_tool],
+        response_format=ToolStrategy(schema=WeatherResponse),
+    )
+
+    # Compile and invoke the agent
+    compiled_agent = agent.compile()
+    result = compiled_agent.invoke(
+        {"messages": [HumanMessage("What's the weather and help me with a query?")]}
+    )
+
+    # Verify that we have the expected messages:
+    # 1. Human message
+    # 2. AI message with both tool calls
+    # 3. Tool message from structured tool call
+    # 4. Tool message from regular tool call
+
+    messages = result["messages"]
+    assert len(messages) >= 4
+
+    # Check that we have the AI message with both tool calls
+    ai_message = messages[1]
+    assert isinstance(ai_message, AIMessage)
+    assert len(ai_message.tool_calls) == 2
+
+    # Check that we have a tool message from the regular tool
+    tool_messages = [m for m in messages if isinstance(m, ToolMessage)]
+    assert len(tool_messages) >= 1
+
+    # The regular tool should have been executed
+    regular_tool_message = next((m for m in tool_messages if m.name == "regular_tool"), None)
+    assert regular_tool_message is not None
+    assert "Regular tool result for: test query" in regular_tool_message.content
+
+    # Verify that the structured response is available in the result
+    assert "response" in result
+    assert result["response"] is not None
+    assert hasattr(result["response"], "temperature")
+    assert result["response"].temperature == 72.0
+    assert result["response"].condition == "sunny"
+
+
+def test_human_in_the_loop_middleware_with_structured_response() -> None:
+    """Test that we can get structured response with human in the loop middleware."""
+
+    class WeatherBaseModel(BaseModel):
+        temperature: float = Field(description="Temperature in fahrenheit")
+        condition: str = Field(description="Weather condition")
+
+    tool_calls = [
+        [
+            {"args": {"a": 1, "b": 2}, "id": "1", "name": "add_numbers"},
+            {
+                "name": "WeatherBaseModel",
+                "id": "2",
+                "args": {"temperature": 72.0, "condition": "sunny"},
+            },
+        ],
+    ]
+
+    @tool
+    def add_numbers(a: int, b: int) -> int:
+        """Add two numbers."""
+        return a + b
+
+    model = FakeToolCallingModel(tool_calls=tool_calls)
+
+    agent = create_agent(model=model, tools=[add_numbers], response_format=WeatherBaseModel)
+    agent = agent.compile()
+    response = agent.invoke(
+        {"messages": [HumanMessage("Add 1 and 2, then return the weather forecast.")]}
+    )
+
+    assert response["response"] == WeatherBaseModel(temperature=72.0, condition="sunny")
+    messages = response["messages"]
+    assert len(messages) == 4
+    assert messages[0].content == "Add 1 and 2, then return the weather forecast."
+    assert len(messages[1].tool_calls) == 2
+    assert messages[2].name == "WeatherBaseModel"
+    assert messages[3].name == "add_numbers"