feat(prebuilt): support ToolOutput response_format (#5915)

* Add support for ToolOutput response format.
* I don't love the name -- it's confusing unless you know that it's parameterizing a strategy.

We should determine if we want to support our old strategy for doing
things -- it has a higher latency (one extra LLM call), but it's a
reasonable built-in strategy as it doesn't do anything awkward with
conversation history. (Wouldn't surprising if it has overall better
performance than tool choice for longer conversations)
This commit is contained in:
Eugene Yurtsev
2025-08-14 23:34:46 -04:00
committed by GitHub
parent ebae60045f
commit b58a7fb2fe
6 changed files with 711 additions and 345 deletions
@@ -50,14 +50,21 @@ from langgraph.prebuilt._internal._typing import (
PreConfiguredChatModel,
SyncOrAsync,
)
from langgraph.prebuilt.responses import (
OutputToolBinding,
ResponseFormat,
SchemaSpec,
ToolOutput,
)
from langgraph.prebuilt.tool_node import ToolNode
from langgraph.runtime import Runtime
from langgraph.store.base import BaseStore
from langgraph.types import Checkpointer, Send
from langgraph.types import Checkpointer, Command, Send
from langgraph.warnings import LangGraphDeprecatedSinceV10
StructuredResponse = Union[dict, BaseModel]
StructuredResponseSchema = Union[dict, type[BaseModel]]
F = TypeVar("F", bound=Callable[..., Any])
@@ -96,6 +103,7 @@ class AgentStateWithStructuredResponsePydantic(AgentStatePydantic):
StateSchema = TypeVar("StateSchema", bound=Union[AgentState, AgentStatePydantic])
StateSchemaType = Type[StateSchema]
PROMPT_RUNNABLE_NAME = "Prompt"
Prompt = Union[
@@ -205,7 +213,7 @@ def _validate_chat_history(
class _AgentBuilder:
"""Internal builder class for constructing React agents with intuitive method-to-node mapping."""
"""Internal builder class for constructing and agent."""
def __init__(
self,
@@ -222,9 +230,7 @@ class _AgentBuilder:
tools: Union[Sequence[Union[BaseTool, Callable, dict[str, Any]]], ToolNode],
*,
prompt: Optional[Prompt] = None,
response_format: Optional[
Union[StructuredResponseSchema, tuple[str, StructuredResponseSchema]]
] = None,
response_format: Optional[ResponseFormat] = None,
pre_model_hook: Optional[RunnableLike] = None,
post_model_hook: Optional[RunnableLike] = None,
state_schema: Optional[StateSchemaType] = None,
@@ -279,6 +285,7 @@ class _AgentBuilder:
self._setup_tools()
self._setup_state_schema()
self._setup_structured_output_tools()
self._setup_model()
def _setup_tools(self) -> None:
@@ -299,6 +306,38 @@ class _AgentBuilder:
}
self._tool_calling_enabled = len(self._tool_classes) > 0
def _setup_structured_output_tools(self) -> None:
"""Set up structured output tools tracking for the tools strategy.
This method implements the "tools" strategy for structured output by:
1. Converting response format schemas to LangChain tools
2. Creating metadata for proper response reconstruction
3. Handling both Pydantic models and dict schemas
Future strategies (json_mode, guided) will have separate setup methods.
"""
self.structured_output_tools: dict[str, OutputToolBinding] = {}
if self.response_format is not None:
response_format = self.response_format
# Handle UsingToolStrategy wrapper
if isinstance(response_format, ToolOutput):
# Use tools strategy - process each ResponseSchema in the UsingToolStrategy
for response_schema in response_format.schemas:
# Use the factory method to create OutputToolBinding
structured_tool_info = OutputToolBinding.from_schema_spec(
response_schema
)
self.structured_output_tools[structured_tool_info.tool.name] = (
structured_tool_info
)
else:
# This shouldn't happen with the new ResponseFormat type, but keeping for safety
raise ValueError(
f"Unsupported response_format type: {type(response_format)}. "
f"Expected UsingToolStrategy."
)
def _setup_state_schema(self) -> None:
"""Setup state schema with validation."""
if self.state_schema is not None:
@@ -320,6 +359,57 @@ class _AgentBuilder:
else AgentState
)
def _handle_structured_response_tool_calls(
self, response: AIMessage
) -> Optional[Command]:
"""Handle tool calls that match structured output tools using the tools strategy.
Args:
response: The AI message containing potential tool calls
Returns:
Command with structured response update if found, None otherwise
Raises:
AssertionError: If multiple structured responses are returned
"""
if not response.tool_calls:
return None
structured_tool_calls = [
tool_call
for tool_call in response.tool_calls
if tool_call["name"] in self.structured_output_tools
]
if len(structured_tool_calls) > 1:
raise AssertionError(
"Model incorrectly returned multiple structured responses. "
"Behavior has not yet been defined in this case."
)
if len(structured_tool_calls) == 1:
tool_call = structured_tool_calls[0]
messages = [
response,
ToolMessage(
content="ok!",
tool_call_id=tool_call["id"],
name=tool_call["name"],
),
]
structured_tool_binding = self.structured_output_tools[tool_call["name"]]
structured_response = structured_tool_binding.parse(tool_call["args"])
return Command(
update={
"messages": messages,
"structured_response": structured_response,
}
)
return None
def _setup_model(self) -> None:
"""Setup model-related attributes."""
self._is_dynamic_model = not isinstance(
@@ -342,11 +432,30 @@ class _AgentBuilder:
)
model = init_chat_model(model)
if len(self._tool_classes + self._llm_builtin_tools) > 0:
model = cast(BaseChatModel, model).bind_tools(
self._tool_classes + self._llm_builtin_tools # type: ignore[operator]
)
# Collect all tools: regular tools + structured output tools
structured_output_tools = list(self.structured_output_tools.values())
all_tools = (
self._tool_classes
+ self._llm_builtin_tools
+ [info.tool for info in structured_output_tools]
)
if len(all_tools) > 0:
# Check if we need to force tool use for structured output
tool_choice = None
if (
self.response_format is not None
and isinstance(self.response_format, ToolOutput)
and self.response_format.tool_choice == "required"
):
tool_choice = "any"
if tool_choice:
model = cast(BaseChatModel, model).bind_tools(
all_tools, tool_choice=tool_choice
)
else:
model = cast(BaseChatModel, model).bind_tools(all_tools)
# Extract just the model part for direct invocation
self._static_model: Optional[Runnable] = model # type: ignore[assignment]
else:
@@ -404,7 +513,8 @@ class _AgentBuilder:
if isinstance(self._final_state_schema, type) and issubclass(
self._final_state_schema, BaseModel
):
# we're passing messages under `messages` key, as this is expected by the prompt
# we're passing messages under `messages` key, as this
# is expected by the prompt
state.messages = messages # type: ignore
else:
state["messages"] = messages # type: ignore
@@ -430,7 +540,8 @@ class _AgentBuilder:
def call_model(
state: StateSchema, runtime: Runtime[ContextT], config: RunnableConfig
) -> StateSchema:
) -> dict[str, Any] | Command:
"""Call the model with the current state and return the response."""
if self._is_async_dynamic_model:
raise RuntimeError(
"Async model callable provided but agent invoked synchronously. "
@@ -457,11 +568,18 @@ class _AgentBuilder:
)
]
}
# Check if any tool calls match structured output tools
structured_command = self._handle_structured_response_tool_calls(response)
if structured_command:
return structured_command
return {"messages": [response]}
async def acall_model(
state: StateSchema, runtime: Runtime[ContextT], config: RunnableConfig
) -> StateSchema:
) -> dict[str, Any] | Command:
"""Call the model with the current state and return the response."""
model_input = _get_model_input_state(state)
model = await self._aresolve_model(state, runtime)
@@ -485,6 +603,12 @@ class _AgentBuilder:
)
]
}
# Check if any tool calls match structured output tools
structured_command = self._handle_structured_response_tool_calls(response)
if structured_command:
return structured_command
return {"messages": [response]}
return RunnableCallable(call_model, acall_model)
@@ -511,57 +635,6 @@ class _AgentBuilder:
else:
return self._final_state_schema
def create_structured_response_node(self) -> Optional[RunnableCallable]:
"""Create the 'generate_structured_response' node if configured."""
if self.response_format is None:
return None
def generate_structured_response(
state: StateSchema, runtime: Runtime[ContextT], config: RunnableConfig
) -> StateSchema:
if self._is_async_dynamic_model:
raise RuntimeError(
"Async model callable provided but agent invoked synchronously. "
"Use agent.ainvoke() or agent.astream(), or provide a sync model callable."
)
messages = _get_state_value(state, "messages")
structured_response_schema = self.response_format
if isinstance(self.response_format, tuple):
system_prompt, structured_response_schema = self.response_format
messages = [SystemMessage(content=system_prompt)] + list(messages)
resolved_model = self._resolve_model(state, runtime)
model_with_structured_output = _get_model(
resolved_model
).with_structured_output(
cast(StructuredResponseSchema, structured_response_schema)
)
response = model_with_structured_output.invoke(messages, config)
return {"structured_response": response}
async def agenerate_structured_response(
state: StateSchema, runtime: Runtime[ContextT], config: RunnableConfig
) -> StateSchema:
messages = _get_state_value(state, "messages")
structured_response_schema = self.response_format
if isinstance(self.response_format, tuple):
system_prompt, structured_response_schema = self.response_format
messages = [SystemMessage(content=system_prompt)] + list(messages)
resolved_model = await self._aresolve_model(state, runtime)
model_with_structured_output = _get_model(
resolved_model
).with_structured_output(
cast(StructuredResponseSchema, structured_response_schema)
)
response = await model_with_structured_output.ainvoke(messages, config)
return {"structured_response": response}
return RunnableCallable(
generate_structured_response, agenerate_structured_response
)
def create_model_router(self) -> Callable[[StateSchema], Union[str, list[Send]]]:
"""Create routing function for model node conditional edges."""
@@ -569,11 +642,22 @@ class _AgentBuilder:
messages = _get_state_value(state, "messages")
last_message = messages[-1]
# Check if the last message is a ToolMessage from a structured tool.
# This condition exists to support structured output via tools.
# Once a tool has been called for structured output, we skip
# tool execution and go to END (if there is no post_model_hook).
if (
isinstance(last_message, ToolMessage)
and last_message.name in self.structured_output_tools
):
return END
if isinstance(last_message, ToolMessage):
return END
if not isinstance(last_message, AIMessage) or not last_message.tool_calls:
if self.post_model_hook is not None:
return "post_model_hook"
elif self.response_format is not None:
return "generate_structured_response"
else:
return END
else:
@@ -606,10 +690,22 @@ class _AgentBuilder:
def create_post_model_hook_router(
self,
) -> Callable[[StateSchema], Union[str, list[Send]]]:
"""Create routing function for post_model_hook node conditional edges."""
"""Create a routing function for post_model_hook node conditional edges."""
def post_model_hook_router(state: StateSchema) -> Union[str, list[Send]]:
messages = _get_state_value(state, "messages")
# Check if the last message is a ToolMessage from a structured tool.
# This condition exists to support structured output via tools.
# Once a tool has been called for structured output, we skip
# tool execution and go to END (if there is no post_model_hook).
last_message = messages[-1]
if (
isinstance(last_message, ToolMessage)
and last_message.name in self.structured_output_tools
):
return END
tool_messages = [
m.tool_call_id for m in messages if isinstance(m, ToolMessage)
]
@@ -639,15 +735,13 @@ class _AgentBuilder:
]
elif isinstance(messages[-1], ToolMessage):
return self._get_entry_point()
elif self.response_format is not None:
return "generate_structured_response"
else:
return END
return post_model_hook_router
def create_tools_router(self) -> Optional[Callable[[StateSchema], str]]:
"""Create routing function for tools node conditional edges."""
"""Create a routing function for tools node conditional edges."""
if not self._should_return_direct:
return None
@@ -696,11 +790,7 @@ class _AgentBuilder:
paths.extend([tool.name for tool in self._tool_classes])
else:
paths.append("tools")
if self.response_format:
paths.append("generate_structured_response")
else:
paths.append(END)
paths.append(END)
return paths
def _get_post_model_hook_paths(self) -> list[str]:
@@ -713,10 +803,7 @@ class _AgentBuilder:
]
else:
paths = [self._get_entry_point(), "tools"]
if self.response_format is not None:
paths.append("generate_structured_response")
else:
paths.append(END)
paths.append(END)
return paths
def build(
@@ -752,10 +839,6 @@ class _AgentBuilder:
if self.post_model_hook:
workflow.add_node("post_model_hook", self.post_model_hook) # type: ignore[arg-type]
structured_node = self.create_structured_response_node()
if structured_node:
workflow.add_node("generate_structured_response", structured_node)
# Add edges
if self.pre_model_hook:
workflow.add_edge("pre_model_hook", "agent")
@@ -829,9 +912,7 @@ def create_react_agent(
tools: Union[Sequence[Union[BaseTool, Callable, dict[str, Any]]], ToolNode],
*,
prompt: Optional[Prompt] = None,
response_format: Optional[
Union[StructuredResponseSchema, tuple[str, StructuredResponseSchema]]
] = None,
response_format: Optional[Union[ToolOutput, StructuredResponseSchema]] = None,
pre_model_hook: Optional[RunnableLike] = None,
post_model_hook: Optional[RunnableLike] = None,
state_schema: Optional[StateSchemaType] = None,
@@ -900,25 +981,27 @@ def create_react_agent(
- Callable: This function should take in full graph state and the output is then passed to the language model.
- Runnable: This runnable should take in full graph state and the output is then passed to the language model.
response_format: An optional schema for the final agent output.
response_format: An optional UsingToolStrategy configuration for structured responses.
If provided, output will be formatted to match the given schema and returned in the 'structured_response' state key.
If provided, the agent will handle structured output via tool calls during the normal conversation flow.
When the model calls a structured output tool, the response will be captured and returned in the 'structured_response' state key.
If not provided, `structured_response` will not be present in the output state.
Can be passed in as:
- an OpenAI function/tool schema,
- a JSON Schema,
- a TypedDict class,
- or a Pydantic class.
- a tuple (prompt, schema), where schema is one of the above.
The prompt will be used together with the model that is being used to generate the structured response.
The UsingToolStrategy should contain:
- schemas: A sequence of ResponseSchema objects that define the structured output format
- tool_choice: Either "required" or "auto" to control when structured output is used
Each ResponseSchema contains:
- schema: A Pydantic model that defines the structure
- name: Optional custom name for the tool (defaults to model name)
- description: Optional custom description (defaults to model docstring)
- strict: Whether to enforce strict validation
!!! Important
`response_format` requires the model to support `.with_structured_output`
`response_format` requires the model to support tool calling
!!! Note
The graph will make a separate call to the LLM to generate the structured response after the agent loop is finished.
This is not the only strategy to get structured responses, see more options in [this guide](https://langchain-ai.github.io/langgraph/how-tos/react-agent-structured-output/).
Structured responses are handled directly in the model call node via tool calls, eliminating the need for separate structured response nodes.
pre_model_hook: An optional node to add before the `agent` node (i.e., the node that calls the LLM).
Useful for managing long message histories (e.g., message trimming, summarization, etc.).
@@ -1055,6 +1138,22 @@ def create_react_agent(
f"Invalid version {version}. Supported versions are 'v1' and 'v2'."
)
if response_format and not isinstance(response_format, ToolOutput):
# Then it's a pydantic model or JSONSchema. We'll automatically convert
# it to the tool output strategy as it is widely supported.
response_format = ToolOutput(
schemas=[SchemaSpec(response_format)],
tool_choice="required",
)
elif isinstance(response_format, tuple):
if len(response_format) == 2:
raise ValueError(
"Passing a 2-tuple as response_format is no longer supported. "
)
else:
# Can only be a ToolOutput or None at this point.
response_format = cast(Optional[ToolOutput], response_format)
# Create and configure the agent builder
builder = _AgentBuilder(
model=model,
@@ -0,0 +1,167 @@
"""Types for setting agent response formats."""
from __future__ import annotations
from dataclasses import dataclass
from typing import Any, Generic, Literal, Optional, Sequence, Type, TypeVar
from langchain_core.tools import BaseTool
from langchain_core.tools import tool as create_tool
from pydantic import BaseModel
# For now, we support only Pydantic models as schemas.
SchemaT = TypeVar("SchemaT", bound=type[BaseModel])
Schema = TypeVar("Schema", bound=BaseModel)
ToolChoice = Literal["required", "auto"]
"""Required: model must call a tool; auto: model may respond with free text."""
@dataclass(init=False)
class SchemaSpec(Generic[SchemaT]):
"""Describes a structured output schema."""
schema: SchemaT
"""The schema for the response, can be a dict or a Pydantic model."""
name: Optional[str] = None
"""Name of the schema, used for tool calling.
If not provided, the name will be the model name.
"""
description: Optional[str] = None
"""Custom description of the schema.
If not provided, provided will use the model's docstring.
"""
strict: bool = False
"""Whether to enforce strict validation of the schema."""
def __init__(
self,
schema: SchemaT,
*,
name: Optional[str] = None,
description: Optional[str] = None,
strict: bool = False,
) -> None:
"""Initialize SchemaSpec with schema and optional parameters."""
self.schema = schema
self.name = name
self.description = description
self.strict = strict
@dataclass(init=False)
class ToolOutput:
"""Use a tool calling strategy for model responses."""
schemas: Sequence[SchemaSpec]
"""Schemas for the tool calls."""
tool_choice: ToolChoice
"""Whether to require tool calling or allow the model to choose.
- "required": The model must use tool calling.
- "auto": The model can choose whether to use tool calling.
Use `auto` if you want the agent to be able to respond with a non structured
response to ask clarifying questions.
"""
def __init__(
self, schemas: Sequence[SchemaSpec], *, tool_choice: ToolChoice = "required"
) -> None:
"""Initialize ToolOutput with schemas and tool choice."""
self.schemas = schemas
self.tool_choice = tool_choice
@dataclass
class OutputToolBinding(Generic[Schema]):
"""Information for tracking structured output tool metadata.
This contains all necessary information to handle structured responses
generated via tool calls, including the original schema, its type classification,
and the corresponding tool implementation used by the tools strategy.
"""
schema: Type[Schema]
"""The original schema provided for structured output (Pydantic model or dict schema)."""
schema_kind: Literal["pydantic"]
"""Classification of the schema type for proper response construction."""
tool: BaseTool
"""LangChain tool instance created from the schema for model binding."""
@classmethod
def from_schema_spec(
cls, schema_spec: SchemaSpec[Type[Schema]]
) -> "OutputToolBinding":
"""Create an OutputToolBinding instance from a SchemaSpec.
Args:
schema_spec: The SchemaSpec to convert
Returns:
An OutputToolBinding instance with the appropriate tool created
"""
# Extract the actual schema from SchemaSpec
schema = schema_spec.schema
kwargs = {}
schema_kind: Literal["pydantic"]
# Use custom description if provided
if schema_spec.description:
kwargs["description"] = schema_spec.description
# Determine schema kind
if isinstance(schema, type) and issubclass(schema, BaseModel):
schema_kind = "pydantic"
else:
raise ValueError(
f"Unsupported schema type: {type(schema)}. "
f"Only Pydantic models are supported."
)
if schema_spec.name is not None:
tool_creator = create_tool(schema_spec.name)
else:
tool_creator = create_tool # type: ignore[assignment]
tool = tool_creator(schema, **kwargs)
return cls(
schema=schema,
schema_kind=schema_kind,
tool=tool,
)
def parse(self, tool_args: dict[str, Any]) -> Schema:
"""Parse tool arguments according to the schema.
Args:
tool_args: The arguments from the tool call
Returns:
The parsed response according to the schema type
Raises:
ValueError: If parsing fails
"""
if self.schema_kind == "pydantic":
if not isinstance(self.schema, type) or not issubclass(
self.schema, BaseModel
):
raise ValueError(
f"Expected Pydantic model class for 'pydantic' kind, got {type(self.schema)}"
)
try:
return self.schema(**tool_args)
except Exception as e:
raise ValueError(
f"Failed to parse tool args to {self.schema.__name__}: {e}"
) from e
else:
raise ValueError(f"Unsupported schema kind: {self.schema_kind}")
# TODO: Add support for built-in structured responses (e.g., openai, grok)
ResponseFormat = ToolOutput
@@ -1,5 +1,5 @@
# serializer version: 1
# name: test_react_agent_graph_structure[None-None-None-tools0]
# name: test_react_agent_graph_structure[None-None-tools0]
'''
graph TD;
__start__ --> agent;
@@ -7,7 +7,7 @@
'''
# ---
# name: test_react_agent_graph_structure[None-None-None-tools1]
# name: test_react_agent_graph_structure[None-None-tools1]
'''
graph TD;
__start__ --> agent;
@@ -17,7 +17,7 @@
'''
# ---
# name: test_react_agent_graph_structure[None-None-pre_model_hook-tools0]
# name: test_react_agent_graph_structure[None-pre_model_hook-tools0]
'''
graph TD;
__start__ --> pre_model_hook;
@@ -26,7 +26,7 @@
'''
# ---
# name: test_react_agent_graph_structure[None-None-pre_model_hook-tools1]
# name: test_react_agent_graph_structure[None-pre_model_hook-tools1]
'''
graph TD;
__start__ --> pre_model_hook;
@@ -37,7 +37,7 @@
'''
# ---
# name: test_react_agent_graph_structure[None-post_model_hook-None-tools0]
# name: test_react_agent_graph_structure[post_model_hook-None-tools0]
'''
graph TD;
__start__ --> agent;
@@ -46,7 +46,7 @@
'''
# ---
# name: test_react_agent_graph_structure[None-post_model_hook-None-tools1]
# name: test_react_agent_graph_structure[post_model_hook-None-tools1]
'''
graph TD;
__start__ --> agent;
@@ -58,7 +58,7 @@
'''
# ---
# name: test_react_agent_graph_structure[None-post_model_hook-pre_model_hook-tools0]
# name: test_react_agent_graph_structure[post_model_hook-pre_model_hook-tools0]
'''
graph TD;
__start__ --> pre_model_hook;
@@ -68,7 +68,7 @@
'''
# ---
# name: test_react_agent_graph_structure[None-post_model_hook-pre_model_hook-tools1]
# name: test_react_agent_graph_structure[post_model_hook-pre_model_hook-tools1]
'''
graph TD;
__start__ --> pre_model_hook;
@@ -81,97 +81,7 @@
'''
# ---
# name: test_react_agent_graph_structure[ResponseFormat-None-None-tools0]
'''
graph TD;
__start__ --> agent;
agent --> generate_structured_response;
generate_structured_response --> __end__;
'''
# ---
# name: test_react_agent_graph_structure[ResponseFormat-None-None-tools1]
'''
graph TD;
__start__ --> agent;
agent -.-> generate_structured_response;
agent -.-> tools;
tools --> agent;
generate_structured_response --> __end__;
'''
# ---
# name: test_react_agent_graph_structure[ResponseFormat-None-pre_model_hook-tools0]
'''
graph TD;
__start__ --> pre_model_hook;
agent --> generate_structured_response;
pre_model_hook --> agent;
generate_structured_response --> __end__;
'''
# ---
# name: test_react_agent_graph_structure[ResponseFormat-None-pre_model_hook-tools1]
'''
graph TD;
__start__ --> pre_model_hook;
agent -.-> generate_structured_response;
agent -.-> tools;
pre_model_hook --> agent;
tools --> pre_model_hook;
generate_structured_response --> __end__;
'''
# ---
# name: test_react_agent_graph_structure[ResponseFormat-post_model_hook-None-tools0]
'''
graph TD;
__start__ --> agent;
agent --> post_model_hook;
post_model_hook --> generate_structured_response;
generate_structured_response --> __end__;
'''
# ---
# name: test_react_agent_graph_structure[ResponseFormat-post_model_hook-None-tools1]
'''
graph TD;
__start__ --> agent;
agent --> post_model_hook;
post_model_hook -.-> agent;
post_model_hook -.-> generate_structured_response;
post_model_hook -.-> tools;
tools --> agent;
generate_structured_response --> __end__;
'''
# ---
# name: test_react_agent_graph_structure[ResponseFormat-post_model_hook-pre_model_hook-tools0]
'''
graph TD;
__start__ --> pre_model_hook;
agent --> post_model_hook;
post_model_hook --> generate_structured_response;
pre_model_hook --> agent;
generate_structured_response --> __end__;
'''
# ---
# name: test_react_agent_graph_structure[ResponseFormat-post_model_hook-pre_model_hook-tools1]
'''
graph TD;
__start__ --> pre_model_hook;
agent --> post_model_hook;
post_model_hook -.-> generate_structured_response;
post_model_hook -.-> pre_model_hook;
post_model_hook -.-> tools;
pre_model_hook --> agent;
tools --> pre_model_hook;
generate_structured_response --> __end__;
'''
# ---
# name: test_react_agent_graph_structure_with_individual_nodes[no_response_format-no_post_hook-no_pre_hook-no_tools]
# name: test_react_agent_graph_structure_with_individual_nodes[no_post_hook-no_pre_hook-no_tools]
'''
graph TD;
__start__ --> agent;
@@ -179,7 +89,7 @@
'''
# ---
# name: test_react_agent_graph_structure_with_individual_nodes[no_response_format-no_post_hook-no_pre_hook-two_tools]
# name: test_react_agent_graph_structure_with_individual_nodes[no_post_hook-no_pre_hook-two_tools]
'''
graph TD;
__start__ --> agent;
@@ -191,7 +101,7 @@
'''
# ---
# name: test_react_agent_graph_structure_with_individual_nodes[no_response_format-no_post_hook-with_pre_hook-no_tools]
# name: test_react_agent_graph_structure_with_individual_nodes[no_post_hook-with_pre_hook-no_tools]
'''
graph TD;
__start__ --> pre_model_hook;
@@ -200,7 +110,7 @@
'''
# ---
# name: test_react_agent_graph_structure_with_individual_nodes[no_response_format-no_post_hook-with_pre_hook-two_tools]
# name: test_react_agent_graph_structure_with_individual_nodes[no_post_hook-with_pre_hook-two_tools]
'''
graph TD;
__start__ --> pre_model_hook;
@@ -213,7 +123,7 @@
'''
# ---
# name: test_react_agent_graph_structure_with_individual_nodes[no_response_format-with_post_hook-no_pre_hook-no_tools]
# name: test_react_agent_graph_structure_with_individual_nodes[with_post_hook-no_pre_hook-no_tools]
'''
graph TD;
__start__ --> agent;
@@ -222,7 +132,7 @@
'''
# ---
# name: test_react_agent_graph_structure_with_individual_nodes[no_response_format-with_post_hook-no_pre_hook-two_tools]
# name: test_react_agent_graph_structure_with_individual_nodes[with_post_hook-no_pre_hook-two_tools]
'''
graph TD;
__start__ --> agent;
@@ -236,7 +146,7 @@
'''
# ---
# name: test_react_agent_graph_structure_with_individual_nodes[no_response_format-with_post_hook-with_pre_hook-no_tools]
# name: test_react_agent_graph_structure_with_individual_nodes[with_post_hook-with_pre_hook-no_tools]
'''
graph TD;
__start__ --> pre_model_hook;
@@ -246,7 +156,7 @@
'''
# ---
# name: test_react_agent_graph_structure_with_individual_nodes[no_response_format-with_post_hook-with_pre_hook-two_tools]
# name: test_react_agent_graph_structure_with_individual_nodes[with_post_hook-with_pre_hook-two_tools]
'''
graph TD;
__start__ --> pre_model_hook;
@@ -261,101 +171,3 @@
'''
# ---
# name: test_react_agent_graph_structure_with_individual_nodes[with_response_format-no_post_hook-no_pre_hook-no_tools]
'''
graph TD;
__start__ --> agent;
agent --> generate_structured_response;
generate_structured_response --> __end__;
'''
# ---
# name: test_react_agent_graph_structure_with_individual_nodes[with_response_format-no_post_hook-no_pre_hook-two_tools]
'''
graph TD;
__start__ --> agent;
agent -.-> generate_structured_response;
agent -.-> tool;
agent -.-> tool2;
tool --> agent;
tool2 --> agent;
generate_structured_response --> __end__;
'''
# ---
# name: test_react_agent_graph_structure_with_individual_nodes[with_response_format-no_post_hook-with_pre_hook-no_tools]
'''
graph TD;
__start__ --> pre_model_hook;
agent --> generate_structured_response;
pre_model_hook --> agent;
generate_structured_response --> __end__;
'''
# ---
# name: test_react_agent_graph_structure_with_individual_nodes[with_response_format-no_post_hook-with_pre_hook-two_tools]
'''
graph TD;
__start__ --> pre_model_hook;
agent -.-> generate_structured_response;
agent -.-> tool;
agent -.-> tool2;
pre_model_hook --> agent;
tool --> pre_model_hook;
tool2 --> pre_model_hook;
generate_structured_response --> __end__;
'''
# ---
# name: test_react_agent_graph_structure_with_individual_nodes[with_response_format-with_post_hook-no_pre_hook-no_tools]
'''
graph TD;
__start__ --> agent;
agent --> post_model_hook;
post_model_hook --> generate_structured_response;
generate_structured_response --> __end__;
'''
# ---
# name: test_react_agent_graph_structure_with_individual_nodes[with_response_format-with_post_hook-no_pre_hook-two_tools]
'''
graph TD;
__start__ --> agent;
agent --> post_model_hook;
post_model_hook -.-> agent;
post_model_hook -.-> generate_structured_response;
post_model_hook -.-> tool;
post_model_hook -.-> tool2;
tool --> agent;
tool2 --> agent;
generate_structured_response --> __end__;
'''
# ---
# name: test_react_agent_graph_structure_with_individual_nodes[with_response_format-with_post_hook-with_pre_hook-no_tools]
'''
graph TD;
__start__ --> pre_model_hook;
agent --> post_model_hook;
post_model_hook --> generate_structured_response;
pre_model_hook --> agent;
generate_structured_response --> __end__;
'''
# ---
# name: test_react_agent_graph_structure_with_individual_nodes[with_response_format-with_post_hook-with_pre_hook-two_tools]
'''
graph TD;
__start__ --> pre_model_hook;
agent --> post_model_hook;
post_model_hook -.-> generate_structured_response;
post_model_hook -.-> pre_model_hook;
post_model_hook -.-> tool;
post_model_hook -.-> tool2;
pre_model_hook --> agent;
tool --> pre_model_hook;
tool2 --> pre_model_hook;
generate_structured_response --> __end__;
'''
# ---
+107 -33
View File
@@ -453,7 +453,10 @@ def test_react_agent_with_structured_response(version: str) -> None:
class WeatherResponse(BaseModel):
temperature: float = Field(description="The temperature in fahrenheit")
tool_calls = [[{"args": {}, "id": "1", "name": "get_weather"}], []]
tool_calls = [
[{"args": {}, "id": "1", "name": "get_weather"}],
[{"name": "WeatherResponse", "id": "2", "args": {"temperature": 75}}],
]
def get_weather():
"""Get the weather"""
@@ -463,17 +466,33 @@ def test_react_agent_with_structured_response(version: str) -> None:
model = FakeToolCallingModel(
tool_calls=tool_calls, structured_response=expected_structured_response
)
for response_format in (WeatherResponse, ("Meow", WeatherResponse)):
agent = create_react_agent(
model,
[get_weather],
response_format=response_format,
version=version,
)
response = agent.invoke({"messages": [HumanMessage("What's the weather?")]})
assert response["structured_response"] == expected_structured_response
assert len(response["messages"]) == 4
assert response["messages"][-2].content == "The weather is sunny and 75°F."
agent = create_react_agent(
model,
[get_weather],
response_format=WeatherResponse,
version=version,
)
response = agent.invoke({"messages": [HumanMessage("What's the weather?")]})
assert response["structured_response"] == expected_structured_response
assert len(response["messages"]) == 5
# Check message types in message history
msg_types = [m.type for m in response["messages"]]
assert msg_types == [
"human", # "What's the weather?"
"ai", # "What's the weather?"
"tool", # "The weather is sunny and 75°F."
"ai", # structured response
"tool", # ok!
]
assert [m.content for m in response["messages"]] == [
"What's the weather?",
"What's the weather?",
"The weather is sunny and 75°F.",
"What's the weather?-What's the weather?-The weather is sunny and 75°F.",
"ok!",
]
class CustomState(AgentState):
@@ -1410,9 +1429,17 @@ def test_dynamic_model_with_structured_response(version: str) -> None:
confidence: float
def dynamic_model(state, runtime: Runtime):
expected_response = TestResponse(message="dynamic response", confidence=0.9)
return FakeToolCallingModel(
tool_calls=[], structured_response=expected_response
tool_calls=[
[
ToolCall(
name="TestResponse",
args={"message": "dynamic response", "confidence": 0.9},
id="1",
type="tool_call",
)
]
],
)
agent = create_react_agent(
@@ -1677,16 +1704,16 @@ def test_post_model_hook_with_structured_output() -> None:
class WeatherResponse(BaseModel):
temperature: float = Field(description="The temperature in fahrenheit")
tool_calls = [[{"args": {}, "id": "1", "name": "get_weather"}]]
tool_calls: list[list[ToolCall]] = [
[{"args": {}, "id": "1", "name": "get_weather"}],
[{"args": {"temperature": 75}, "id": "2", "name": "WeatherResponse"}],
]
def get_weather():
"""Get the weather"""
return "The weather is sunny and 75°F."
expected_structured_response = WeatherResponse(temperature=75)
model = FakeToolCallingModel(
tool_calls=tool_calls, structured_response=expected_structured_response
)
class State(AgentState):
flag: bool
@@ -1695,6 +1722,7 @@ def test_post_model_hook_with_structured_output() -> None:
def post_model_hook(state: State) -> Union[dict[str, bool], Command]:
return {"flag": True}
model = FakeToolCallingModel(tool_calls=tool_calls)
agent = create_react_agent(
model,
[get_weather],
@@ -1704,7 +1732,6 @@ def test_post_model_hook_with_structured_output() -> None:
)
assert "post_model_hook" in agent.nodes
assert "generate_structured_response" in agent.nodes
response = agent.invoke(
{"messages": [HumanMessage("What's the weather?")], "flag": False}
@@ -1712,10 +1739,19 @@ def test_post_model_hook_with_structured_output() -> None:
assert response["flag"] is True
assert response["structured_response"] == expected_structured_response
# Reset the state of the model
model = FakeToolCallingModel(tool_calls=tool_calls)
agent = create_react_agent(
model,
[get_weather],
response_format=WeatherResponse,
post_model_hook=post_model_hook,
state_schema=State,
)
events = list(
agent.stream({"messages": [HumanMessage("What's the weather?")], "flag": False})
)
assert "generate_structured_response" in events[-1]
assert events == [
{
"agent": {
@@ -1724,7 +1760,7 @@ def test_post_model_hook_with_structured_output() -> None:
content="What's the weather?",
additional_kwargs={},
response_metadata={},
id="2",
id="0",
tool_calls=[
{
"name": "get_weather",
@@ -1745,7 +1781,7 @@ def test_post_model_hook_with_structured_output() -> None:
content="The weather is sunny and 75°F.",
name="get_weather",
tool_call_id="1",
),
)
]
}
},
@@ -1756,25 +1792,26 @@ def test_post_model_hook_with_structured_output() -> None:
content="What's the weather?-What's the weather?-The weather is sunny and 75°F.",
additional_kwargs={},
response_metadata={},
id="3",
id="1",
tool_calls=[
{
"name": "get_weather",
"args": {},
"id": "1",
"name": "WeatherResponse",
"args": {"temperature": 75},
"id": "2",
"type": "tool_call",
}
],
)
]
),
_AnyIdToolMessage(
content="ok!",
name="WeatherResponse",
tool_call_id="2",
),
],
"structured_response": WeatherResponse(temperature=75.0),
}
},
{"post_model_hook": {"flag": True}},
{
"generate_structured_response": {
"structured_response": WeatherResponse(temperature=75.0)
}
},
]
@@ -1837,3 +1874,40 @@ def test_create_react_agent_inject_vars_with_post_model_hook(
AIMessage("hi-hi-6", id="1"),
]
assert result["foo"] == 2
@pytest.mark.parametrize("version", REACT_TOOL_CALL_VERSIONS)
def test_response_format_using_tool_choice(version: Literal["v1", "v2"]) -> None:
"""Test response format using tool choice."""
class WeatherResponse(BaseModel):
temperature: float = Field(description="The temperature in fahrenheit")
tool_calls: list[list[ToolCall]] = [
[{"args": {}, "id": "1", "name": "get_weather"}],
[{"args": {"temperature": "75"}, "id": "2", "name": "WeatherResponse"}],
]
def get_weather() -> str:
"""Get the weather"""
return "The weather is sunny and 75°F."
expected_structured_response = WeatherResponse(temperature=75)
model = FakeToolCallingModel(tool_calls=tool_calls)
agent = create_react_agent(
model,
[get_weather],
response_format=WeatherResponse,
version=version,
)
response = agent.invoke(
{
"messages": [
{
"role": "user",
"content": "What's the weather?",
}
]
}
)
assert response.get("structured_response") == expected_structured_response
@@ -39,20 +39,17 @@ class ResponseFormat(BaseModel):
@pytest.mark.parametrize("tools", [[], [tool]])
@pytest.mark.parametrize("pre_model_hook", [None, pre_model_hook])
@pytest.mark.parametrize("post_model_hook", [None, post_model_hook])
@pytest.mark.parametrize("response_format", [None, ResponseFormat])
def test_react_agent_graph_structure(
snapshot: SnapshotAssertion,
tools: list[Callable],
pre_model_hook: Union[Callable, None],
post_model_hook: Union[Callable, None],
response_format: Union[type[BaseModel], None],
) -> None:
agent = create_react_agent(
model,
tools=tools,
pre_model_hook=pre_model_hook,
post_model_hook=post_model_hook,
response_format=response_format,
)
try:
assert agent.get_graph().draw_mermaid(with_styles=False) == snapshot
@@ -63,7 +60,6 @@ def test_react_agent_graph_structure(
f"tools: {tools}, "
f"pre_model_hook: {pre_model_hook}, "
f"post_model_hook: {post_model_hook}, "
f"response_format: {response_format}"
) from e
@@ -74,24 +70,17 @@ def test_react_agent_graph_structure(
@pytest.mark.parametrize(
"post_model_hook", [None, post_model_hook], ids=["no_post_hook", "with_post_hook"]
)
@pytest.mark.parametrize(
"response_format",
[None, ResponseFormat],
ids=["no_response_format", "with_response_format"],
)
def test_react_agent_graph_structure_with_individual_nodes(
snapshot: SnapshotAssertion,
tools: list[Callable],
pre_model_hook: Union[Callable, None],
post_model_hook: Union[Callable, None],
response_format: Union[type[BaseModel], None],
) -> None:
agent = create_react_agent(
model,
tools=tools,
pre_model_hook=pre_model_hook,
post_model_hook=post_model_hook,
response_format=response_format,
use_individual_tool_nodes=True,
)
assert agent.get_graph().draw_mermaid(with_styles=False) == snapshot
+225
View File
@@ -0,0 +1,225 @@
"""Unit tests for langgraph.prebuilt.responses module."""
import pytest
from pydantic import BaseModel
from langgraph.prebuilt.responses import (
OutputToolBinding,
ResponseFormat,
SchemaSpec,
ToolOutput,
)
class TestModel(BaseModel):
"""A test model for structured output."""
name: str
age: int
email: str = "default@example.com"
class CustomModel(BaseModel):
"""Custom model with a custom docstring."""
value: float
description: str
class EmptyDocModel(BaseModel):
# No custom docstring, should have no description in tool
data: str
class TestSchemaSpec:
"""Test SchemaSpec dataclass."""
def test_basic_creation(self):
"""Test basic SchemaSpec creation."""
schema = SchemaSpec(schema=TestModel)
assert schema.schema == TestModel
assert schema.name is None
assert schema.description is None
assert schema.strict is False
def test_creation_with_all_fields(self):
"""Test SchemaSpec creation with all fields."""
schema = SchemaSpec(
schema=TestModel,
name="custom_test_model",
description="A custom description",
strict=True,
)
assert schema.schema == TestModel
assert schema.name == "custom_test_model"
assert schema.description == "A custom description"
assert schema.strict is True
class TestUsingToolStrategy:
"""Test UsingToolStrategy dataclass."""
def test_basic_creation(self):
"""Test basic UsingToolStrategy creation."""
schema = SchemaSpec(schema=TestModel)
strategy = ToolOutput(schemas=[schema])
assert len(strategy.schemas) == 1
assert strategy.schemas[0] == schema
assert strategy.tool_choice == "required" # default
def test_creation_with_auto_tool_choice(self):
"""Test UsingToolStrategy creation with auto tool choice."""
schema = SchemaSpec(schema=TestModel)
strategy = ToolOutput(schemas=[schema], tool_choice="auto")
assert strategy.tool_choice == "auto"
def test_multiple_schemas(self):
"""Test UsingToolStrategy with multiple schemas."""
schema1 = SchemaSpec(schema=TestModel)
schema2 = SchemaSpec(schema=CustomModel)
strategy = ToolOutput(schemas=[schema1, schema2])
assert len(strategy.schemas) == 2
class TestOutputToolBinding:
"""Test OutputToolBinding dataclass and its methods."""
def test_from_schema_spec_basic(self):
"""Test basic OutputToolBinding creation from SchemaSpec."""
schema_spec = SchemaSpec(schema=TestModel)
tool_binding = OutputToolBinding.from_schema_spec(schema_spec)
assert tool_binding.schema == TestModel
assert tool_binding.schema_kind == "pydantic"
assert tool_binding.tool is not None
assert tool_binding.tool.name == "TestModel"
def test_from_schema_spec_with_custom_name(self):
"""Test OutputToolBinding creation with custom name."""
schema_spec = SchemaSpec(schema=TestModel, name="custom_tool_name")
tool_binding = OutputToolBinding.from_schema_spec(schema_spec)
assert tool_binding.tool.name == "custom_tool_name"
def test_from_schema_spec_with_custom_description(self):
"""Test OutputToolBinding creation with custom description."""
schema_spec = SchemaSpec(
schema=TestModel, description="Custom tool description"
)
tool_binding = OutputToolBinding.from_schema_spec(schema_spec)
assert tool_binding.tool.description == "Custom tool description"
def test_from_schema_spec_with_model_docstring(self):
"""Test OutputToolBinding creation using model docstring as description."""
schema_spec = SchemaSpec(schema=CustomModel)
tool_binding = OutputToolBinding.from_schema_spec(schema_spec)
assert tool_binding.tool.description == "Custom model with a custom docstring."
@pytest.mark.skip(
reason="Need to fix bug in langchain-core for inheritance of doc-strings."
)
def test_from_schema_spec_empty_docstring(self):
"""Test OutputToolBinding creation with model that has default docstring."""
# Create a model with the same docstring as BaseModel
class DefaultDocModel(BaseModel):
# This should have the same docstring as BaseModel
pass
schema_spec = SchemaSpec(schema=DefaultDocModel)
tool_binding = OutputToolBinding.from_schema_spec(schema_spec)
# Should use empty description when model has default BaseModel docstring
assert tool_binding.tool.description == ""
def test_parse_payload_pydantic_success(self):
"""Test successful parsing for Pydantic model."""
schema_spec = SchemaSpec(schema=TestModel)
tool_binding = OutputToolBinding.from_schema_spec(schema_spec)
tool_args = {"name": "John", "age": 30}
result = tool_binding.parse(tool_args)
assert isinstance(result, TestModel)
assert result.name == "John"
assert result.age == 30
assert result.email == "default@example.com" # default value
def test_parse_payload_pydantic_validation_error(self):
"""Test parsing failure for invalid Pydantic data."""
schema_spec = SchemaSpec(schema=TestModel)
tool_binding = OutputToolBinding.from_schema_spec(schema_spec)
# Missing required field 'name'
tool_args = {"age": 30}
with pytest.raises(ValueError, match="Failed to parse tool args to TestModel"):
tool_binding.parse(tool_args)
def test_parse_payload_invalid_kind(self):
"""Test parsing with invalid kind."""
from unittest.mock import Mock
mock_tool = Mock()
tool_binding = OutputToolBinding(
schema=TestModel,
schema_kind="invalid_kind", # type: ignore
tool=mock_tool,
)
with pytest.raises(ValueError, match="Unsupported schema kind: invalid_kind"):
tool_binding.parse({"name": "test", "age": 25})
def test_parse_payload_invalid_pydantic_schema(self):
"""Test parsing with invalid schema for pydantic kind."""
from unittest.mock import Mock
mock_tool = Mock()
# Create tool binding with dict schema but pydantic kind
tool_binding = OutputToolBinding(
schema={"type": "object"}, schema_kind="pydantic", tool=mock_tool
)
with pytest.raises(
ValueError, match="Expected Pydantic model class for 'pydantic' kind"
):
tool_binding.parse({"name": "test", "age": 25})
class TestResponseFormat:
"""Test ResponseFormat type alias."""
def test_response_format_is_using_tool_strategy(self):
"""Test that ResponseFormat is aliased to UsingToolStrategy."""
assert ResponseFormat is ToolOutput
def test_can_create_response_format(self):
"""Test that we can create ResponseFormat instances."""
schema = SchemaSpec(schema=TestModel)
response_format = ResponseFormat(schemas=[schema])
assert isinstance(response_format, ToolOutput)
assert len(response_format.schemas) == 1
class TestEdgeCases:
"""Test edge cases and error conditions."""
def test_empty_schemas_list(self) -> None:
"""Test UsingToolStrategy with empty schemas list."""
strategy = ToolOutput([SchemaSpec(EmptyDocModel)])
assert len(strategy.schemas) == 1
@pytest.mark.skip(
reason="Need to fix bug in langchain-core for inheritance of doc-strings."
)
def test_base_model_doc_constant(self) -> None:
"""Test that BASE_MODEL_DOC constant is set correctly."""
binding = OutputToolBinding.from_schema_spec(SchemaSpec(EmptyDocModel))
assert binding.tool.name == "EmptyDocModel"
assert (
binding.tool.description[:5] == ""
) # Should be empty for default docstring