From 10d46acc60d19db426b6508b8b81de96aa1bab6d Mon Sep 17 00:00:00 2001 From: Vadym Barda Date: Fri, 10 Jan 2025 11:06:59 -0500 Subject: [PATCH] langgraph: add structured output to create_react_agent (#2848) ```python class WeatherResponse(BaseModel): """Respond to the user with this""" temperature: float = Field(description="The temperature in fahrenheit") wind_direction: str = Field( description="The direction of the wind in abbreviated form" ) wind_speed: float = Field(description="The speed of the wind in mph") @tool def get_weather(city: Literal["nyc", "sf"]): """Use this to get weather information.""" if city == "nyc": return "It is cloudy in NYC, with 5 mph winds in the North-East direction and a temperature of 70 degrees" elif city == "sf": return "It is 75 degrees and sunny in SF, with 3 mph winds in the South-East direction" else: raise AssertionError("Unknown city") model = ChatOpenAI() tools = [get_weather] agent_with_structured_output = create_react_agent(model, tools, response_format=WeatherResponse) agent_with_structured_output.invoke({"messages": [("user", "what's the weather in nyc?")]}) ``` ```pycon { 'messages': [...], 'structured_response': WeatherResponse(temperature=70.0, wind_directon='NE', wind_speed=5.0) } ``` --- .../langgraph/prebuilt/chat_agent_executor.py | 127 ++++++++++++++++-- .../tests/__snapshots__/test_large_cases.ambr | 4 +- libs/langgraph/tests/test_prebuilt.py | 45 ++++++- 3 files changed, 163 insertions(+), 13 deletions(-) diff --git a/libs/langgraph/langgraph/prebuilt/chat_agent_executor.py b/libs/langgraph/langgraph/prebuilt/chat_agent_executor.py index 1b4797b80..7172c8e10 100644 --- a/libs/langgraph/langgraph/prebuilt/chat_agent_executor.py +++ b/libs/langgraph/langgraph/prebuilt/chat_agent_executor.py @@ -1,4 +1,13 @@ -from typing import Callable, Literal, Optional, Sequence, Type, TypeVar, Union, cast +from typing import ( + Callable, + Literal, + Optional, + Sequence, + Type, + TypeVar, + Union, + cast, +) from langchain_core.language_models import BaseChatModel, LanguageModelLike from langchain_core.messages import AIMessage, BaseMessage, SystemMessage, ToolMessage @@ -8,11 +17,12 @@ from langchain_core.runnables import ( RunnableConfig, ) from langchain_core.tools import BaseTool +from pydantic import BaseModel from typing_extensions import Annotated, TypedDict from langgraph._api.deprecation import deprecated_parameter from langgraph.errors import ErrorCode, create_error_message -from langgraph.graph import StateGraph +from langgraph.graph import END, StateGraph from langgraph.graph.graph import CompiledGraph from langgraph.graph.message import add_messages from langgraph.managed import IsLastStep, RemainingSteps @@ -22,11 +32,14 @@ from langgraph.store.base import BaseStore from langgraph.types import Checkpointer from langgraph.utils.runnable import RunnableCallable +StructuredResponse = Union[dict, BaseModel] +StructuredResponseSchema = Union[dict, type[BaseModel]] + # We create the AgentState that we will pass around # This simply involves a list of messages # We want steps to return messages to append to the list -# So we annotate the messages attribute with operator.add +# So we annotate the messages attribute with `add_messages` reducer class AgentState(TypedDict): """The state of the agent.""" @@ -36,6 +49,8 @@ class AgentState(TypedDict): remaining_steps: RemainingSteps + structured_response: StructuredResponse + StateSchema = TypeVar("StateSchema", bound=AgentState) StateSchemaType = Type[StateSchema] @@ -162,6 +177,19 @@ def _should_bind_tools(model: LanguageModelLike, tools: Sequence[BaseTool]) -> b return False +def _get_model(model: LanguageModelLike) -> BaseChatModel: + """Get the underlying model from a RunnableBinding or return the model itself.""" + if isinstance(model, RunnableBinding): + model = model.bound + + if not isinstance(model, BaseChatModel): + raise TypeError( + f"Expected `model` to be a ChatModel or RunnableBinding (e.g. model.bind_tools(...)), got {type(model)}" + ) + + return model + + def _validate_chat_history( messages: Sequence[BaseMessage], ) -> None: @@ -201,6 +229,9 @@ def create_react_agent( state_schema: Optional[StateSchemaType] = None, messages_modifier: Optional[MessagesModifier] = None, state_modifier: Optional[StateModifier] = None, + response_format: Optional[ + Union[StructuredResponseSchema, tuple[str, StructuredResponseSchema]] + ] = None, checkpointer: Optional[Checkpointer] = None, store: Optional[BaseStore] = None, interrupt_before: Optional[list[str]] = None, @@ -236,6 +267,25 @@ def create_react_agent( - str: This is converted to a SystemMessage and added to the beginning of the list of messages in state["messages"]. - 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. + + If provided, output will be formatted to match the given schema 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. + + !!! Important + `response_format` requires the model to support `.with_structured_output` + + !!! 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/). checkpointer: An optional checkpoint saver object. This is used for persisting the state of the graph (e.g., as chat memory) for a single thread (e.g., a single conversation). store: An optional store object. This is used for persisting data @@ -527,9 +577,11 @@ def create_react_agent( """ if state_schema is not None: - if missing_keys := {"messages", "is_last_step"} - set( - state_schema.__annotations__ - ): + required_keys = {"messages", "remaining_steps"} + if response_format is not None: + required_keys.add("structured_response") + + if missing_keys := required_keys - set(state_schema.__annotations__): raise ValueError(f"Missing required key(s) {missing_keys} in state_schema") if isinstance(tools, ToolExecutor): @@ -633,11 +685,54 @@ def create_react_agent( # We return a list, because this will get added to the existing list return {"messages": [response]} + def generate_structured_response( + state: AgentState, config: RunnableConfig + ) -> AgentState: + # NOTE: we exclude the last message because there is enough information + # for the LLM to generate the structured response + messages = state["messages"][:-1] + structured_response_schema = response_format + if isinstance(response_format, tuple): + system_prompt, structured_response_schema = response_format + messages = [SystemMessage(content=system_prompt)] + list(messages) + + model_with_structured_output = _get_model(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: AgentState, config: RunnableConfig + ) -> AgentState: + # NOTE: we exclude the last message because there is enough information + # for the LLM to generate the structured response + messages = state["messages"][:-1] + structured_response_schema = response_format + if isinstance(response_format, tuple): + system_prompt, structured_response_schema = response_format + messages = [SystemMessage(content=system_prompt)] + list(messages) + + model_with_structured_output = _get_model(model).with_structured_output( + cast(StructuredResponseSchema, structured_response_schema) + ) + response = await model_with_structured_output.ainvoke(messages, config) + return {"structured_response": response} + if not tool_calling_enabled: # Define a new graph workflow = StateGraph(state_schema or AgentState) workflow.add_node("agent", RunnableCallable(call_model, acall_model)) workflow.set_entry_point("agent") + if response_format is not None: + workflow.add_node( + "generate_structured_response", + RunnableCallable( + generate_structured_response, agenerate_structured_response + ), + ) + workflow.add_edge("agent", "generate_structured_response") + return workflow.compile( checkpointer=checkpointer, store=store, @@ -647,12 +742,12 @@ def create_react_agent( ) # Define the function that determines whether to continue or not - def should_continue(state: AgentState) -> Literal["tools", "__end__"]: + def should_continue(state: AgentState) -> str: messages = state["messages"] last_message = messages[-1] # If there is no function call, then we finish if not isinstance(last_message, AIMessage) or not last_message.tool_calls: - return "__end__" + return END if response_format is None else "generate_structured_response" # Otherwise if there is, we continue else: return "tools" @@ -668,6 +763,19 @@ def create_react_agent( # This means that this node is the first one called workflow.set_entry_point("agent") + # Add a structured output node if response_format is provided + if response_format is not None: + workflow.add_node( + "generate_structured_response", + RunnableCallable( + generate_structured_response, agenerate_structured_response + ), + ) + workflow.add_edge("generate_structured_response", END) + should_continue_destinations = ["tools", "generate_structured_response"] + else: + should_continue_destinations = ["tools", END] + # We now add a conditional edge workflow.add_conditional_edges( # First, we define the start node. We use `agent`. @@ -675,6 +783,7 @@ def create_react_agent( "agent", # Next, we pass in the function that will determine which node is called next. should_continue, + path_map=should_continue_destinations, ) def route_tool_responses(state: AgentState) -> Literal["agent", "__end__"]: @@ -682,7 +791,7 @@ def create_react_agent( if not isinstance(m, ToolMessage): break if m.name in should_return_direct: - return "__end__" + return END return "agent" if should_return_direct: diff --git a/libs/langgraph/tests/__snapshots__/test_large_cases.ambr b/libs/langgraph/tests/__snapshots__/test_large_cases.ambr index 017b1aa9b..c272ae1d6 100644 --- a/libs/langgraph/tests/__snapshots__/test_large_cases.ambr +++ b/libs/langgraph/tests/__snapshots__/test_large_cases.ambr @@ -2832,10 +2832,10 @@ ''' # --- # name: test_prebuilt_tool_chat - '{"$defs": {"BaseMessage": {"additionalProperties": true, "description": "Base abstract message class.\\n\\nMessages are the inputs and outputs of ChatModels.", "properties": {"content": {"anyOf": [{"type": "string"}, {"items": {"anyOf": [{"type": "string"}, {"type": "object"}]}, "type": "array"}], "title": "Content"}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"title": "Type", "type": "string"}, "name": {"anyOf": [{"type": "string"}, {"type": "null"}], "default": null, "title": "Name"}, "id": {"anyOf": [{"type": "string"}, {"type": "null"}], "default": null, "title": "Id"}}, "required": ["content", "type"], "title": "BaseMessage", "type": "object"}}, "properties": {"messages": {"items": {"$ref": "#/$defs/BaseMessage"}, "title": "Messages", "type": "array"}}, "required": ["messages"], "title": "LangGraphInput", "type": "object"}' + '{"$defs": {"BaseMessage": {"additionalProperties": true, "description": "Base abstract message class.\\n\\nMessages are the inputs and outputs of ChatModels.", "properties": {"content": {"anyOf": [{"type": "string"}, {"items": {"anyOf": [{"type": "string"}, {"type": "object"}]}, "type": "array"}], "title": "Content"}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"title": "Type", "type": "string"}, "name": {"anyOf": [{"type": "string"}, {"type": "null"}], "default": null, "title": "Name"}, "id": {"anyOf": [{"type": "string"}, {"type": "null"}], "default": null, "title": "Id"}}, "required": ["content", "type"], "title": "BaseMessage", "type": "object"}, "BaseModel": {"properties": {}, "title": "BaseModel", "type": "object"}}, "properties": {"messages": {"items": {"$ref": "#/$defs/BaseMessage"}, "title": "Messages", "type": "array"}, "structured_response": {"anyOf": [{"type": "object"}, {"$ref": "#/$defs/BaseModel"}], "title": "Structured Response"}}, "required": ["messages", "structured_response"], "title": "LangGraphInput", "type": "object"}' # --- # name: test_prebuilt_tool_chat.1 - '{"$defs": {"BaseMessage": {"additionalProperties": true, "description": "Base abstract message class.\\n\\nMessages are the inputs and outputs of ChatModels.", "properties": {"content": {"anyOf": [{"type": "string"}, {"items": {"anyOf": [{"type": "string"}, {"type": "object"}]}, "type": "array"}], "title": "Content"}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"title": "Type", "type": "string"}, "name": {"anyOf": [{"type": "string"}, {"type": "null"}], "default": null, "title": "Name"}, "id": {"anyOf": [{"type": "string"}, {"type": "null"}], "default": null, "title": "Id"}}, "required": ["content", "type"], "title": "BaseMessage", "type": "object"}}, "properties": {"messages": {"items": {"$ref": "#/$defs/BaseMessage"}, "title": "Messages", "type": "array"}}, "required": ["messages"], "title": "LangGraphOutput", "type": "object"}' + '{"$defs": {"BaseMessage": {"additionalProperties": true, "description": "Base abstract message class.\\n\\nMessages are the inputs and outputs of ChatModels.", "properties": {"content": {"anyOf": [{"type": "string"}, {"items": {"anyOf": [{"type": "string"}, {"type": "object"}]}, "type": "array"}], "title": "Content"}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"title": "Type", "type": "string"}, "name": {"anyOf": [{"type": "string"}, {"type": "null"}], "default": null, "title": "Name"}, "id": {"anyOf": [{"type": "string"}, {"type": "null"}], "default": null, "title": "Id"}}, "required": ["content", "type"], "title": "BaseMessage", "type": "object"}, "BaseModel": {"properties": {}, "title": "BaseModel", "type": "object"}}, "properties": {"messages": {"items": {"$ref": "#/$defs/BaseMessage"}, "title": "Messages", "type": "array"}, "structured_response": {"anyOf": [{"type": "object"}, {"$ref": "#/$defs/BaseModel"}], "title": "Structured Response"}}, "required": ["messages", "structured_response"], "title": "LangGraphOutput", "type": "object"}' # --- # name: test_prebuilt_tool_chat.2 ''' diff --git a/libs/langgraph/tests/test_prebuilt.py b/libs/langgraph/tests/test_prebuilt.py index 3186cac9e..2a7f57001 100644 --- a/libs/langgraph/tests/test_prebuilt.py +++ b/libs/langgraph/tests/test_prebuilt.py @@ -32,7 +32,7 @@ from langchain_core.outputs import ChatGeneration, ChatResult from langchain_core.runnables import Runnable, RunnableLambda from langchain_core.tools import BaseTool, ToolException from langchain_core.tools import tool as dec_tool -from pydantic import BaseModel, ValidationError +from pydantic import BaseModel, Field, ValidationError from pydantic.v1 import BaseModel as BaseModelV1 from pydantic.v1 import ValidationError as ValidationErrorV1 from typing_extensions import TypedDict @@ -47,7 +47,11 @@ from langgraph.prebuilt import ( create_react_agent, tools_condition, ) -from langgraph.prebuilt.chat_agent_executor import AgentState, _validate_chat_history +from langgraph.prebuilt.chat_agent_executor import ( + AgentState, + StructuredResponse, + _validate_chat_history, +) from langgraph.prebuilt.tool_node import ( TOOL_CALL_ERROR_TEMPLATE, InjectedState, @@ -71,6 +75,7 @@ pytestmark = pytest.mark.anyio class FakeToolCallingModel(BaseChatModel): tool_calls: Optional[list[list[ToolCall]]] = None + structured_response: Optional[StructuredResponse] = None index: int = 0 tool_style: Literal["openai", "anthropic"] = "openai" @@ -98,6 +103,14 @@ class FakeToolCallingModel(BaseChatModel): def _llm_type(self) -> str: return "fake-tool-call-model" + def with_structured_output( + self, schema: Type[BaseModel] + ) -> Runnable[LanguageModelInput, StructuredResponse]: + if self.structured_response is None: + raise ValueError("Structured response is not set") + + return RunnableLambda(lambda x: self.structured_response) + def bind_tools( self, tools: Sequence[Union[Dict[str, Any], Type[BaseModel], Callable, BaseTool]], @@ -511,6 +524,34 @@ def test__infer_handled_types() -> None: _infer_handled_types(handler) +@pytest.mark.skipif( + not IS_LANGCHAIN_CORE_030_OR_GREATER, + reason="Pydantic v1 is required for this test to pass in langchain-core < 0.3", +) +def test_react_agent_with_structured_response() -> None: + class WeatherResponse(BaseModel): + temperature: float = Field(description="The temperature in fahrenheit") + + tool_calls = [[{"args": {}, "id": "1", "name": "get_weather"}], []] + + 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 + ) + for response_format in (WeatherResponse, ("Meow", WeatherResponse)): + agent = create_react_agent( + model, [get_weather], response_format=response_format + ) + 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." + + # tools for testing Too def tool1(some_val: int, some_other_val: str) -> str: """Tool 1 docstring."""