From b58a7fb2fe85abd6fd6877fa94033a0ee754bedc Mon Sep 17 00:00:00 2001 From: Eugene Yurtsev Date: Thu, 14 Aug 2025 23:34:46 -0400 Subject: [PATCH] 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) --- .../langgraph/prebuilt/chat_agent_executor.py | 293 ++++++++++++------ libs/prebuilt/langgraph/prebuilt/responses.py | 167 ++++++++++ .../__snapshots__/test_react_agent_graph.ambr | 220 +------------ libs/prebuilt/tests/test_react_agent.py | 140 +++++++-- libs/prebuilt/tests/test_react_agent_graph.py | 11 - libs/prebuilt/tests/test_responses.py | 225 ++++++++++++++ 6 files changed, 711 insertions(+), 345 deletions(-) create mode 100644 libs/prebuilt/langgraph/prebuilt/responses.py create mode 100644 libs/prebuilt/tests/test_responses.py diff --git a/libs/prebuilt/langgraph/prebuilt/chat_agent_executor.py b/libs/prebuilt/langgraph/prebuilt/chat_agent_executor.py index a8a7ffbfc..8bd8f33c8 100644 --- a/libs/prebuilt/langgraph/prebuilt/chat_agent_executor.py +++ b/libs/prebuilt/langgraph/prebuilt/chat_agent_executor.py @@ -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, diff --git a/libs/prebuilt/langgraph/prebuilt/responses.py b/libs/prebuilt/langgraph/prebuilt/responses.py new file mode 100644 index 000000000..3f6560413 --- /dev/null +++ b/libs/prebuilt/langgraph/prebuilt/responses.py @@ -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 diff --git a/libs/prebuilt/tests/__snapshots__/test_react_agent_graph.ambr b/libs/prebuilt/tests/__snapshots__/test_react_agent_graph.ambr index 2323534da..e10a7020a 100644 --- a/libs/prebuilt/tests/__snapshots__/test_react_agent_graph.ambr +++ b/libs/prebuilt/tests/__snapshots__/test_react_agent_graph.ambr @@ -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__; - - ''' -# --- diff --git a/libs/prebuilt/tests/test_react_agent.py b/libs/prebuilt/tests/test_react_agent.py index b5c4e25fe..5bcf582ca 100644 --- a/libs/prebuilt/tests/test_react_agent.py +++ b/libs/prebuilt/tests/test_react_agent.py @@ -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 diff --git a/libs/prebuilt/tests/test_react_agent_graph.py b/libs/prebuilt/tests/test_react_agent_graph.py index 6f76fe9e2..d6a95e075 100644 --- a/libs/prebuilt/tests/test_react_agent_graph.py +++ b/libs/prebuilt/tests/test_react_agent_graph.py @@ -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 diff --git a/libs/prebuilt/tests/test_responses.py b/libs/prebuilt/tests/test_responses.py new file mode 100644 index 000000000..d6f58cc94 --- /dev/null +++ b/libs/prebuilt/tests/test_responses.py @@ -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