mirror of
https://github.com/langchain-ai/langgraph.git
synced 2026-08-21 15:12:26 +02:00
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:
@@ -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__;
|
||||
|
||||
'''
|
||||
# ---
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
Reference in New Issue
Block a user