diff --git a/libs/langgraph/tests/__snapshots__/test_large_cases.ambr b/libs/langgraph/tests/__snapshots__/test_large_cases.ambr index 3df0b4159..34e4ed64c 100644 --- a/libs/langgraph/tests/__snapshots__/test_large_cases.ambr +++ b/libs/langgraph/tests/__snapshots__/test_large_cases.ambr @@ -175,10 +175,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"}, {"additionalProperties": true, "type": "object"}]}, "type": "array"}], "title": "Content"}, "additional_kwargs": {"additionalProperties": true, "title": "Additional Kwargs", "type": "object"}, "response_metadata": {"additionalProperties": true, "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"}}, "description": "The state of the agent.", "properties": {"messages": {"items": {"$ref": "#/$defs/BaseMessage"}, "title": "Messages", "type": "array"}, "remaining_steps": {"title": "Remaining Steps", "type": "integer"}, "structured_response": {"title": "Structured Response", "type": "null"}}, "required": ["messages"], "title": "AgentState", "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"}, {"additionalProperties": true, "type": "object"}]}, "type": "array"}], "title": "Content"}, "additional_kwargs": {"additionalProperties": true, "title": "Additional Kwargs", "type": "object"}, "response_metadata": {"additionalProperties": true, "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"}}, "description": "The state of the agent.", "properties": {"messages": {"items": {"$ref": "#/$defs/BaseMessage"}, "title": "Messages", "type": "array"}, "remaining_steps": {"title": "Remaining Steps", "type": "integer"}}, "required": ["messages"], "title": "AgentState", "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"}, {"additionalProperties": true, "type": "object"}]}, "type": "array"}], "title": "Content"}, "additional_kwargs": {"additionalProperties": true, "title": "Additional Kwargs", "type": "object"}, "response_metadata": {"additionalProperties": true, "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"}}, "description": "The state of the agent.", "properties": {"messages": {"items": {"$ref": "#/$defs/BaseMessage"}, "title": "Messages", "type": "array"}, "remaining_steps": {"title": "Remaining Steps", "type": "integer"}, "structured_response": {"title": "Structured Response", "type": "null"}}, "required": ["messages"], "title": "AgentState", "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"}, {"additionalProperties": true, "type": "object"}]}, "type": "array"}], "title": "Content"}, "additional_kwargs": {"additionalProperties": true, "title": "Additional Kwargs", "type": "object"}, "response_metadata": {"additionalProperties": true, "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"}}, "description": "The state of the agent.", "properties": {"messages": {"items": {"$ref": "#/$defs/BaseMessage"}, "title": "Messages", "type": "array"}, "remaining_steps": {"title": "Remaining Steps", "type": "integer"}}, "required": ["messages"], "title": "AgentState", "type": "object"}' # --- # name: test_prebuilt_tool_chat.2 ''' @@ -198,7 +198,7 @@ } }, { - "id": "model", + "id": "agent", "type": "runnable", "data": { "id": [ @@ -207,7 +207,7 @@ "_runnable", "RunnableCallable" ], - "name": "model" + "name": "agent" } }, { @@ -230,21 +230,21 @@ "edges": [ { "source": "__start__", - "target": "model" + "target": "agent" }, { - "source": "model", + "source": "agent", "target": "__end__", "conditional": true }, { - "source": "model", + "source": "agent", "target": "tools", "conditional": true }, { "source": "tools", - "target": "model" + "target": "agent" } ] } @@ -253,10 +253,10 @@ # name: test_prebuilt_tool_chat.3 ''' graph TD; - __start__ --> model; - model -.-> __end__; - model -.-> tools; - tools --> model; + __start__ --> agent; + agent -.-> __end__; + agent -.-> tools; + tools --> agent; ''' # --- diff --git a/libs/langgraph/tests/conftest.py b/libs/langgraph/tests/conftest.py index 465cdc9d5..4ef30b918 100644 --- a/libs/langgraph/tests/conftest.py +++ b/libs/langgraph/tests/conftest.py @@ -69,7 +69,7 @@ def cache(request: pytest.FixtureRequest) -> Iterator[BaseCache]: elif request.param == "redis": # Get worker ID for parallel test isolation worker_id = getattr(request.config, "workerinput", {}).get("workerid", "master") - + redis_client = redis.Redis( host="localhost", port=6379, db=0, decode_responses=False ) diff --git a/libs/langgraph/tests/test_large_cases.py b/libs/langgraph/tests/test_large_cases.py index b367043e3..d626d05ed 100644 --- a/libs/langgraph/tests/test_large_cases.py +++ b/libs/langgraph/tests/test_large_cases.py @@ -1390,11 +1390,11 @@ def test_prebuilt_tool_chat(snapshot: SnapshotAssertion) -> None: ), { "langgraph_step": 1, - "langgraph_node": "model", - "langgraph_triggers": ("branch:to:model",), - "langgraph_path": (PULL, "model"), - "langgraph_checkpoint_ns": AnyStr("model:"), - "checkpoint_ns": AnyStr("model:"), + "langgraph_node": "agent", + "langgraph_triggers": ("branch:to:agent",), + "langgraph_path": (PULL, "agent"), + "langgraph_checkpoint_ns": AnyStr("agent:"), + "checkpoint_ns": AnyStr("agent:"), "ls_provider": "fakechatmodel", "ls_model_type": "chat", }, @@ -1449,11 +1449,11 @@ def test_prebuilt_tool_chat(snapshot: SnapshotAssertion) -> None: ), { "langgraph_step": 3, - "langgraph_node": "model", - "langgraph_triggers": ("branch:to:model",), - "langgraph_path": (PULL, "model"), - "langgraph_checkpoint_ns": AnyStr("model:"), - "checkpoint_ns": AnyStr("model:"), + "langgraph_node": "agent", + "langgraph_triggers": ("branch:to:agent",), + "langgraph_path": (PULL, "agent"), + "langgraph_checkpoint_ns": AnyStr("agent:"), + "checkpoint_ns": AnyStr("agent:"), "ls_provider": "fakechatmodel", "ls_model_type": "chat", }, @@ -1497,11 +1497,11 @@ def test_prebuilt_tool_chat(snapshot: SnapshotAssertion) -> None: ), { "langgraph_step": 5, - "langgraph_node": "model", - "langgraph_triggers": ("branch:to:model",), - "langgraph_path": (PULL, "model"), - "langgraph_checkpoint_ns": AnyStr("model:"), - "checkpoint_ns": AnyStr("model:"), + "langgraph_node": "agent", + "langgraph_triggers": ("branch:to:agent",), + "langgraph_path": (PULL, "agent"), + "langgraph_checkpoint_ns": AnyStr("agent:"), + "checkpoint_ns": AnyStr("agent:"), "ls_provider": "fakechatmodel", "ls_model_type": "chat", }, @@ -1533,7 +1533,7 @@ def test_prebuilt_tool_chat(snapshot: SnapshotAssertion) -> None: for output in (invoke_updates_events, stream_updates_events): assert output[:3] == [ { - "model": { + "agent": { "messages": [ _AnyIdAIMessage( content="", @@ -1560,7 +1560,7 @@ def test_prebuilt_tool_chat(snapshot: SnapshotAssertion) -> None: } }, { - "model": { + "agent": { "messages": [ _AnyIdAIMessage( content="", @@ -1606,7 +1606,7 @@ def test_prebuilt_tool_chat(snapshot: SnapshotAssertion) -> None: }, ) assert output[5:] == [ - {"model": {"messages": [_AnyIdAIMessage(content="answer")]}} + {"agent": {"messages": [_AnyIdAIMessage(content="answer")]}} ] diff --git a/libs/langgraph/tests/test_large_cases_async.py b/libs/langgraph/tests/test_large_cases_async.py index 60da9acca..022034e66 100644 --- a/libs/langgraph/tests/test_large_cases_async.py +++ b/libs/langgraph/tests/test_large_cases_async.py @@ -1143,11 +1143,11 @@ async def test_prebuilt_tool_chat() -> None: ), { "langgraph_step": 1, - "langgraph_node": "model", - "langgraph_triggers": ("branch:to:model",), - "langgraph_path": (PULL, "model"), - "langgraph_checkpoint_ns": AnyStr("model:"), - "checkpoint_ns": AnyStr("model:"), + "langgraph_node": "agent", + "langgraph_triggers": ("branch:to:agent",), + "langgraph_path": (PULL, "agent"), + "langgraph_checkpoint_ns": AnyStr("agent:"), + "checkpoint_ns": AnyStr("agent:"), "ls_provider": "fakechatmodel", "ls_model_type": "chat", }, @@ -1202,11 +1202,11 @@ async def test_prebuilt_tool_chat() -> None: ), { "langgraph_step": 3, - "langgraph_node": "model", - "langgraph_triggers": ("branch:to:model",), - "langgraph_path": (PULL, "model"), - "langgraph_checkpoint_ns": AnyStr("model:"), - "checkpoint_ns": AnyStr("model:"), + "langgraph_node": "agent", + "langgraph_triggers": ("branch:to:agent",), + "langgraph_path": (PULL, "agent"), + "langgraph_checkpoint_ns": AnyStr("agent:"), + "checkpoint_ns": AnyStr("agent:"), "ls_provider": "fakechatmodel", "ls_model_type": "chat", }, @@ -1250,11 +1250,11 @@ async def test_prebuilt_tool_chat() -> None: ), { "langgraph_step": 5, - "langgraph_node": "model", - "langgraph_triggers": ("branch:to:model",), - "langgraph_path": (PULL, "model"), - "langgraph_checkpoint_ns": AnyStr("model:"), - "checkpoint_ns": AnyStr("model:"), + "langgraph_node": "agent", + "langgraph_triggers": ("branch:to:agent",), + "langgraph_path": (PULL, "agent"), + "langgraph_checkpoint_ns": AnyStr("agent:"), + "checkpoint_ns": AnyStr("agent:"), "ls_provider": "fakechatmodel", "ls_model_type": "chat", }, @@ -1269,7 +1269,7 @@ async def test_prebuilt_tool_chat() -> None: ] assert stream_updates_events[:3] == [ { - "model": { + "agent": { "messages": [ _AnyIdAIMessage( content="", @@ -1296,7 +1296,7 @@ async def test_prebuilt_tool_chat() -> None: } }, { - "model": { + "agent": { "messages": [ _AnyIdAIMessage( content="", @@ -1342,7 +1342,7 @@ async def test_prebuilt_tool_chat() -> None: }, ) assert stream_updates_events[5:] == [ - {"model": {"messages": [_AnyIdAIMessage(content="answer")]}} + {"agent": {"messages": [_AnyIdAIMessage(content="answer")]}} ] diff --git a/libs/prebuilt/langgraph/prebuilt/chat_agent_executor.py b/libs/prebuilt/langgraph/prebuilt/chat_agent_executor.py index d28f6249b..6ce3725d5 100644 --- a/libs/prebuilt/langgraph/prebuilt/chat_agent_executor.py +++ b/libs/prebuilt/langgraph/prebuilt/chat_agent_executor.py @@ -7,12 +7,14 @@ from typing import ( Awaitable, Callable, Generic, + Literal, Optional, Sequence, Union, cast, get_type_hints, ) +from warnings import warn from langchain_core.language_models import ( BaseChatModel, @@ -21,6 +23,7 @@ from langchain_core.language_models import ( ) from langchain_core.messages import ( AIMessage, + AnyMessage, BaseMessage, SystemMessage, ToolCall, @@ -35,6 +38,7 @@ from pydantic import BaseModel from typing_extensions import Annotated, NotRequired, TypedDict, TypeVar from langgraph._internal._runnable import RunnableCallable, RunnableLike +from langgraph._internal._typing import MISSING from langgraph.errors import ErrorCode, create_error_message from langgraph.graph import END, StateGraph from langgraph.graph.message import add_messages @@ -58,22 +62,39 @@ from langgraph.store.base import BaseStore from langgraph.types import Checkpointer, Command, Send from langgraph.typing import ContextT, StateT -StructuredResponseT = TypeVar( - "StructuredResponseT", - default=None, -) +StructuredResponseT = TypeVar("StructuredResponseT", default=None) STRUCTURED_OUTPUT_ERROR_TEMPLATE = "Error: {error}\n Please fix your mistakes." -class AgentState(TypedDict, Generic[StructuredResponseT]): +class AgentState(TypedDict): """The state of the agent.""" messages: Annotated[Sequence[BaseMessage], add_messages] remaining_steps: NotRequired[RemainingSteps] - structured_response: NotRequired[StructuredResponseT] + +class AgentStatePydantic(BaseModel): + """The state of the agent.""" + + messages: Annotated[Sequence[BaseMessage], add_messages] + + remaining_steps: RemainingSteps = 25 + + +class AgentStateWithStructuredResponse(AgentState, Generic[StructuredResponseT]): + """The state of the agent with a structured response.""" + + structured_response: StructuredResponseT + + +class AgentStateWithStructuredResponsePydantic( + AgentStatePydantic, Generic[StructuredResponseT] +): + """The state of the agent with a structured response.""" + + structured_response: StructuredResponseT PROMPT_RUNNABLE_NAME = "Prompt" @@ -161,17 +182,15 @@ def _validate_chat_history( raise ValueError(error_message) -class _AgentBuilder(Generic[StateT, ContextT, StructuredResponseT]): +class _AgentBuilder(Generic[StructuredResponseT]): """Internal builder class for constructing and agent.""" - _final_state_schema: type[StateT] - def __init__( self, model: Union[ str, BaseChatModel, - SyncOrAsync[[StateT, Runtime[ContextT]], BaseChatModel], + SyncOrAsync[[StateT, Runtime[ContextT]], BaseModel], ], tools: Union[Sequence[Union[BaseTool, Callable, dict[str, Any]]], ToolNode], *, @@ -181,6 +200,7 @@ class _AgentBuilder(Generic[StateT, ContextT, StructuredResponseT]): post_model_hook: Optional[RunnableLike] = None, state_schema: Optional[type[StateT]] = None, context_schema: Optional[type[ContextT]] = None, + version: Literal["v1", "v2"] = "v2", name: Optional[str] = None, store: Optional[BaseStore] = None, ): @@ -192,6 +212,7 @@ class _AgentBuilder(Generic[StateT, ContextT, StructuredResponseT]): self.post_model_hook = post_model_hook self.state_schema = state_schema self.context_schema = context_schema + self.version = version self.name = name self.store = store @@ -283,7 +304,11 @@ class _AgentBuilder(Generic[StateT, ContextT, StructuredResponseT]): self._final_state_schema = self.state_schema else: - self._final_state_schema = cast(type[StateT], AgentState) + self._final_state_schema = ( + AgentStateWithStructuredResponse # type: ignore[assignment] + if self.response_format is not None + else AgentState + ) def _handle_structured_response_tool_calls( self, response: AIMessage @@ -344,7 +369,7 @@ class _AgentBuilder(Generic[StateT, ContextT, StructuredResponseT]): return Command( update={"messages": [response, *tool_messages]}, - goto="model", + goto="agent", ) def _handle_single_structured_output( @@ -406,7 +431,7 @@ class _AgentBuilder(Generic[StateT, ContextT, StructuredResponseT]): ), ], }, - goto="model", + goto="agent", ) def _handle_structured_output_error( @@ -558,10 +583,20 @@ class _AgentBuilder(Generic[StateT, ContextT, StructuredResponseT]): """Create the 'agent' node that calls the LLM.""" def _get_model_input_state(state: StateT) -> StateT: - messages = _get_state_value(state, "messages") - error_msg = ( - f"Expected input to call_model to have 'messages' key, but got {state}" - ) + if self.pre_model_hook is not None: + messages = _get_state_value( + state, "llm_input_messages" + ) or _get_state_value(state, "messages") + error_msg = ( + f"Expected input to call_model to have 'llm_input_messages' " + f"or 'messages' key, but got {state}" + ) + else: + messages = _get_state_value(state, "messages") + error_msg = ( + f"Expected input to call_model to " + f"have 'messages' key, but got {state}" + ) if messages is None: raise ValueError(error_msg) @@ -681,6 +716,28 @@ class _AgentBuilder(Generic[StateT, ContextT, StructuredResponseT]): return RunnableCallable(call_model, acall_model) + def _get_input_schema(self) -> type[StateT]: + """Get input schema for model node.""" + if self.pre_model_hook is not None: + if isinstance(self._final_state_schema, type) and issubclass( + self._final_state_schema, BaseModel + ): + from pydantic import create_model + + return create_model( + "CallModelInputSchema", + llm_input_messages=(list[AnyMessage], ...), + __base__=self._final_state_schema, + ) + else: + + class CallModelInputSchema(self._final_state_schema): # type: ignore + llm_input_messages: list[AnyMessage] + + return CallModelInputSchema + else: + return self._final_state_schema + def create_model_router(self) -> Callable[[StateT], Union[str, list[Send]]]: """Create routing function for model node conditional edges.""" @@ -707,13 +764,16 @@ class _AgentBuilder(Generic[StateT, ContextT, StructuredResponseT]): else: return END else: - if self.post_model_hook is not None: - return "post_model_hook" - tool_calls = [ - self._tool_node.inject_tool_args(call, state, self.store) # type: ignore[arg-type] - for call in last_message.tool_calls - ] - return [Send("tools", [tool_call]) for tool_call in tool_calls] + if self.version == "v1": + return "tools" + elif self.version == "v2": + if self.post_model_hook is not None: + return "post_model_hook" + tool_calls = [ + self._tool_node.inject_tool_args(call, state, self.store) # type: ignore[arg-type] + for call in last_message.tool_calls + ] + return [Send("tools", [tool_call]) for tool_call in tool_calls] return should_continue @@ -784,7 +844,7 @@ class _AgentBuilder(Generic[StateT, ContextT, StructuredResponseT]): def _get_entry_point(self) -> str: """Get the workflow entry point.""" - return "pre_model_hook" if self.pre_model_hook else "model" + return "pre_model_hook" if self.pre_model_hook else "agent" def _get_model_paths(self) -> list[str]: """Get possible edge destinations from model node.""" @@ -817,7 +877,9 @@ class _AgentBuilder(Generic[StateT, ContextT, StructuredResponseT]): workflow.set_entry_point(self._get_entry_point()) # Add nodes - workflow.add_node("model", self.create_model_node()) + workflow.add_node( + "agent", self.create_model_node(), input_schema=self._get_input_schema() + ) if self._tool_calling_enabled: workflow.add_node("tools", self._tool_node) @@ -830,10 +892,10 @@ class _AgentBuilder(Generic[StateT, ContextT, StructuredResponseT]): # Add edges if self.pre_model_hook: - workflow.add_edge("pre_model_hook", "model") + workflow.add_edge("pre_model_hook", "agent") if self.post_model_hook: - workflow.add_edge("model", "post_model_hook") + workflow.add_edge("agent", "post_model_hook") post_hook_paths = self._get_post_model_hook_paths() if len(post_hook_paths) == 1: # No need for a conditional edge if there's only one path @@ -848,10 +910,10 @@ class _AgentBuilder(Generic[StateT, ContextT, StructuredResponseT]): model_paths = self._get_model_paths() if len(model_paths) == 1: # No need for a conditional edge if there's only one path - workflow.add_edge("model", model_paths[0]) + workflow.add_edge("agent", model_paths[0]) else: workflow.add_conditional_edges( - "model", + "agent", self.create_model_router(), path_map=model_paths, ) @@ -903,7 +965,7 @@ def create_react_agent( model: Union[ str, BaseChatModel, - SyncOrAsync[[StateT, Runtime[ContextT]], BaseChatModel], + SyncOrAsync[[StateT, Runtime[ContextT]], BaseModel], ], tools: Union[Sequence[Union[BaseTool, Callable, dict[str, Any]]], ToolNode], *, @@ -924,7 +986,9 @@ def create_react_agent( interrupt_before: Optional[list[str]] = None, interrupt_after: Optional[list[str]] = None, debug: bool = False, + version: Literal["v1", "v2"] = "v2", name: Optional[str] = None, + **deprecated_kwargs: Any, ) -> CompiledStateGraph: """Creates an agent graph that calls tools in a loop until a stopping condition is met. @@ -1006,17 +1070,24 @@ def create_react_agent( Useful for managing long message histories (e.g., message trimming, summarization, etc.). Pre-model hook must be a callable or a runnable that takes in current graph state and returns a state update in the form of ```python - # Where `messages` MUST be provided + # At least one of `messages` or `llm_input_messages` MUST be provided { - # will UPDATE the `messages` in the state + # If provided, will UPDATE the `messages` in the state "messages": [RemoveMessage(id=REMOVE_ALL_MESSAGES), ...], + # If provided, will be used as the input to the LLM, + # and will NOT UPDATE `messages` in the state + "llm_input_messages": [...], # Any other state keys that need to be propagated ... } ``` + !!! Important + At least one of `messages` or `llm_input_messages` MUST be provided and will be used as an input to the `agent` node. + The rest of the keys will be added to the graph state. + !!! Warning - you should OVERWRITE the `messages` key by doing the following: + If you are returning `messages` in the pre-model hook, you should OVERWRITE the `messages` key by doing the following: ```python { @@ -1027,6 +1098,9 @@ def create_react_agent( post_model_hook: An optional node to add after the `agent` node (i.e., the node that calls the LLM). Useful for implementing human-in-the-loop, guardrails, validation, or other post-processing. Post-model hook must be a callable or a runnable that takes in current graph state and returns a state update. + + !!! Note + Only available with `version="v2"`. state_schema: An optional state schema that defines graph state. Must have `messages` and `remaining_steps` keys. Defaults to `AgentState` that defines those two keys. @@ -1036,20 +1110,34 @@ def create_react_agent( store: An optional store object. This is used for persisting data across multiple threads (e.g., multiple conversations / users). interrupt_before: An optional list of node names to interrupt before. - Should be one of the following: "model", "tools". + Should be one of the following: "agent", "tools". This is useful if you want to add a user confirmation or other interrupt before taking an action. interrupt_after: An optional list of node names to interrupt after. - Should be one of the following: "model", "tools". + Should be one of the following: "agent", "tools". This is useful if you want to return directly or run additional processing on an output. debug: A flag indicating whether to enable debug mode. + version: Determines the version of the graph to create. + Can be one of: + + - `"v1"`: The tool node processes a single message. All tool + calls in the message are executed in parallel within the tool node. + - `"v2"`: The tool node processes a tool call. + Tool calls are distributed across multiple instances of the tool + node using the [Send](https://langchain-ai.github.io/langgraph/concepts/low_level/#send) + API. name: An optional name for the CompiledStateGraph. This name will be automatically used when adding ReAct agent graph to another graph as a subgraph node - particularly useful for building multi-agent systems. - Returns: - A CompiledStateGraph that can be used for chat interactions. + !!! warning "`config_schema` Deprecated" + The `config_schema` parameter is deprecated in v0.6.0 and support will be removed in v2.0.0. + Please use `context_schema` instead to specify the schema for run-scoped context. - The "model" node calls the language model with the messages list (after applying the prompt). + + Returns: + A compiled LangChain runnable that can be used for chat interactions. + + The "agent" node calls the language model with the messages list (after applying the prompt). If the resulting AIMessage contains `tool_calls`, the graph will then call the ["tools"][langgraph.prebuilt.tool_node.ToolNode]. The "tools" node executes the tools (1 tool per `tool_call`) and adds the responses to the messages list as `ToolMessage` objects. The agent node then calls the language model again. @@ -1088,8 +1176,24 @@ def create_react_agent( print(chunk) ``` """ + # Handle deprecated config_schema parameter + if ( + config_schema := deprecated_kwargs.pop("config_schema", MISSING) + ) is not MISSING: + warn( + "`config_schema` is deprecated and will be removed. Please use `context_schema` instead.", + category=DeprecationWarning, + ) + if context_schema is None: + context_schema = config_schema + + if len(deprecated_kwargs) > 0: + raise TypeError( + f"create_react_agent() got unexpected keyword arguments: {deprecated_kwargs}" + ) + if response_format and not isinstance(response_format, (ToolOutput, NativeOutput)): - if _supports_native_structured_output(model): + if _supports_native_structured_output(model): # type: ignore[arg-type] response_format = NativeOutput( schema=response_format, ) @@ -1104,7 +1208,7 @@ def create_react_agent( ) # Create and configure the agent builder - builder = _AgentBuilder[StateT, ContextT, StructuredResponseT]( + builder = _AgentBuilder( model=model, tools=tools, prompt=prompt, @@ -1115,6 +1219,7 @@ def create_react_agent( post_model_hook=post_model_hook, state_schema=state_schema, context_schema=context_schema, + version=version, name=name, store=store, ) @@ -1134,4 +1239,7 @@ def create_react_agent( __all__ = [ "create_react_agent", "AgentState", + "AgentStatePydantic", + "AgentStateWithStructuredResponse", + "AgentStateWithStructuredResponsePydantic", ] diff --git a/libs/prebuilt/tests/__snapshots__/test_react_agent_graph.ambr b/libs/prebuilt/tests/__snapshots__/test_react_agent_graph.ambr index 9f8ed33c8..49e40649e 100644 --- a/libs/prebuilt/tests/__snapshots__/test_react_agent_graph.ambr +++ b/libs/prebuilt/tests/__snapshots__/test_react_agent_graph.ambr @@ -2,18 +2,18 @@ # name: test_react_agent_graph_structure[None-None-tools0] ''' graph TD; - __start__ --> model; - model --> __end__; + __start__ --> agent; + agent --> __end__; ''' # --- # name: test_react_agent_graph_structure[None-None-tools1] ''' graph TD; - __start__ --> model; - model -.-> __end__; - model -.-> tools; - tools --> model; + __start__ --> agent; + agent -.-> __end__; + agent -.-> tools; + tools --> agent; ''' # --- @@ -21,8 +21,8 @@ ''' graph TD; __start__ --> pre_model_hook; - pre_model_hook --> model; - model --> __end__; + pre_model_hook --> agent; + agent --> __end__; ''' # --- @@ -30,9 +30,9 @@ ''' graph TD; __start__ --> pre_model_hook; - model -.-> __end__; - model -.-> tools; - pre_model_hook --> model; + agent -.-> __end__; + agent -.-> tools; + pre_model_hook --> agent; tools --> pre_model_hook; ''' @@ -40,8 +40,8 @@ # name: test_react_agent_graph_structure[post_model_hook-None-tools0] ''' graph TD; - __start__ --> model; - model --> post_model_hook; + __start__ --> agent; + agent --> post_model_hook; post_model_hook --> __end__; ''' @@ -49,12 +49,12 @@ # name: test_react_agent_graph_structure[post_model_hook-None-tools1] ''' graph TD; - __start__ --> model; - model --> post_model_hook; + __start__ --> agent; + agent --> post_model_hook; post_model_hook -.-> __end__; - post_model_hook -.-> model; + post_model_hook -.-> agent; post_model_hook -.-> tools; - tools --> model; + tools --> agent; ''' # --- @@ -62,8 +62,8 @@ ''' graph TD; __start__ --> pre_model_hook; - model --> post_model_hook; - pre_model_hook --> model; + agent --> post_model_hook; + pre_model_hook --> agent; post_model_hook --> __end__; ''' @@ -72,11 +72,11 @@ ''' graph TD; __start__ --> pre_model_hook; - model --> post_model_hook; + agent --> post_model_hook; post_model_hook -.-> __end__; post_model_hook -.-> pre_model_hook; post_model_hook -.-> tools; - pre_model_hook --> model; + pre_model_hook --> agent; tools --> pre_model_hook; ''' diff --git a/libs/prebuilt/tests/test_deprecation.py b/libs/prebuilt/tests/test_deprecation.py new file mode 100644 index 000000000..686445dc2 --- /dev/null +++ b/libs/prebuilt/tests/test_deprecation.py @@ -0,0 +1,41 @@ +import pytest +from typing_extensions import TypedDict + +from langgraph.prebuilt import create_react_agent +from langgraph.warnings import LangGraphDeprecatedSinceV10 +from tests.model import FakeToolCallingModel + + +class Config(TypedDict): + model: str + + +@pytest.mark.filterwarnings("ignore:`config_schema` is deprecated") +@pytest.mark.filterwarnings("ignore:`get_config_jsonschema` is deprecated") +def test_config_schema_deprecation() -> None: + with pytest.warns( + DeprecationWarning, + match="`config_schema` is deprecated and will be removed. Please use `context_schema` instead.", + ): + agent = create_react_agent(FakeToolCallingModel(), [], config_schema=Config) + assert agent.context_schema == Config + + with pytest.warns( + LangGraphDeprecatedSinceV10, + match="`config_schema` is deprecated. Use `get_context_jsonschema` for the relevant schema instead.", + ): + assert agent.config_schema() is not None + + with pytest.warns( + LangGraphDeprecatedSinceV10, + match="`get_config_jsonschema` is deprecated. Use `get_context_jsonschema` instead.", + ): + assert agent.get_config_jsonschema() is not None + + +def test_extra_kwargs_deprecation() -> None: + with pytest.raises( + TypeError, + match="create_react_agent\(\) got unexpected keyword arguments: \{'extra': 'extra'\}", + ): + create_react_agent(FakeToolCallingModel(), [], extra="extra") diff --git a/libs/prebuilt/tests/test_react_agent.py b/libs/prebuilt/tests/test_react_agent.py index 54b6289a7..16f12053f 100644 --- a/libs/prebuilt/tests/test_react_agent.py +++ b/libs/prebuilt/tests/test_react_agent.py @@ -2,6 +2,8 @@ import dataclasses import inspect from typing import ( Annotated, + Literal, + Optional, Union, ) @@ -31,6 +33,8 @@ from langgraph.prebuilt import ( ) from langgraph.prebuilt.chat_agent_executor import ( AgentState, + AgentStatePydantic, + StateT, _validate_chat_history, ) from langgraph.prebuilt.tool_node import ( @@ -49,14 +53,18 @@ from tests.model import FakeToolCallingModel pytestmark = pytest.mark.anyio +REACT_TOOL_CALL_VERSIONS = ["v1", "v2"] -def test_no_prompt(sync_checkpointer: BaseCheckpointSaver) -> None: + +@pytest.mark.parametrize("version", REACT_TOOL_CALL_VERSIONS) +def test_no_prompt(sync_checkpointer: BaseCheckpointSaver, version: str) -> None: model = FakeToolCallingModel() agent = create_react_agent( model, [], checkpointer=sync_checkpointer, + version=version, ) inputs = [HumanMessage("hi?")] thread = {"configurable": {"thread_id": "123"}} @@ -164,7 +172,8 @@ def test_runnable_prompt(): assert response == expected_response -def test_prompt_with_store(): +@pytest.mark.parametrize("version", REACT_TOOL_CALL_VERSIONS) +def test_prompt_with_store(version: Literal["v1", "v2"]): def add(a: int, b: int): """Adds a and b""" return a + b @@ -189,6 +198,7 @@ def test_prompt_with_store(): [add], prompt=prompt, store=in_memory_store, + version=version, ) response = agent.invoke( {"messages": [("user", "hi")]}, {"configurable": {"user_id": "1"}} @@ -201,6 +211,7 @@ def test_prompt_with_store(): [add], prompt=prompt_no_store, store=in_memory_store, + version=version, ) response = agent.invoke( {"messages": [("user", "hi")]}, {"configurable": {"user_id": "2"}} @@ -241,7 +252,9 @@ async def test_prompt_with_store_async(): assert response["messages"][-1].content == "User name is Alice-hi" # test state modifier that doesn't use store works - agent = create_react_agent(model, [add], prompt=prompt_no_store, store=in_memory_store) + agent = create_react_agent( + model, [add], prompt=prompt_no_store, store=in_memory_store + ) response = await agent.ainvoke( {"messages": [("user", "hi")]}, {"configurable": {"user_id": "2"}} ) @@ -249,8 +262,9 @@ async def test_prompt_with_store_async(): @pytest.mark.parametrize("tool_style", ["openai", "anthropic"]) +@pytest.mark.parametrize("version", REACT_TOOL_CALL_VERSIONS) @pytest.mark.parametrize("include_builtin", [True, False]) -def test_model_with_tools(tool_style: str, include_builtin: bool) -> None: +def test_model_with_tools(tool_style: str, version: str, include_builtin: bool) -> None: model = FakeToolCallingModel(tool_style=tool_style) @dec_tool @@ -285,6 +299,7 @@ def test_model_with_tools(tool_style: str, include_builtin: bool) -> None: create_react_agent( model.bind_tools(tools), tools, + version=version, ) @@ -414,7 +429,8 @@ def test__infer_handled_types() -> None: _infer_handled_types(handler) -def test_react_agent_with_structured_response() -> None: +@pytest.mark.parametrize("version", REACT_TOOL_CALL_VERSIONS) +def test_react_agent_with_structured_response(version: str) -> None: class WeatherResponse(BaseModel): temperature: float = Field(description="The temperature in fahrenheit") @@ -435,6 +451,7 @@ def test_react_agent_with_structured_response() -> None: model, [get_weather], response_format=WeatherResponse, + version=version, ) response = agent.invoke({"messages": [HumanMessage("What's the weather?")]}) assert response["structured_response"] == expected_structured_response @@ -463,8 +480,16 @@ class CustomState(AgentState): user_name: str +class CustomStatePydantic(AgentStatePydantic): + user_name: Optional[str] = None + + +@pytest.mark.parametrize("version", REACT_TOOL_CALL_VERSIONS) +@pytest.mark.parametrize("state_schema", [CustomState, CustomStatePydantic]) def test_react_agent_update_state( sync_checkpointer: BaseCheckpointSaver, + version: Literal["v1", "v2"], + state_schema: StateT, ) -> None: @dec_tool def get_user_name(tool_call_id: Annotated[str, InjectedToolCallId]): @@ -481,22 +506,34 @@ def test_react_agent_update_state( } ) - def prompt(state: CustomState): - user_name = state.get("user_name") - if user_name is None: - return state["messages"] + if issubclass(state_schema, AgentStatePydantic): - system_msg = f"User name is {user_name}" - return [{"role": "system", "content": system_msg}] + state["messages"] + def prompt(state: CustomStatePydantic): + user_name = state.user_name + if user_name is None: + return state.messages + + system_msg = f"User name is {user_name}" + return [{"role": "system", "content": system_msg}] + state.messages + else: + + def prompt(state: CustomState): + user_name = state.get("user_name") + if user_name is None: + return state["messages"] + + system_msg = f"User name is {user_name}" + return [{"role": "system", "content": system_msg}] + state["messages"] tool_calls = [[{"args": {}, "id": "1", "name": "get_user_name"}]] model = FakeToolCallingModel(tool_calls=tool_calls) agent = create_react_agent( model, [get_user_name], - state_schema=CustomState, + state_schema=state_schema, prompt=prompt, checkpointer=sync_checkpointer, + version=version, ) config = {"configurable": {"thread_id": "1"}} # Run until interrupted @@ -512,8 +549,9 @@ def test_react_agent_update_state( assert tool_message.name == "get_user_name" +@pytest.mark.parametrize("version", REACT_TOOL_CALL_VERSIONS) def test_react_agent_parallel_tool_calls( - sync_checkpointer: BaseCheckpointSaver, + sync_checkpointer: BaseCheckpointSaver, version: str ) -> None: human_assistance_execution_count = 0 @@ -546,6 +584,7 @@ def test_react_agent_parallel_tool_calls( model, [human_assistance, get_weather], checkpointer=sync_checkpointer, + version=version, ) config = {"configurable": {"thread_id": "1"}} query = "Get user assistance and also check the weather" @@ -556,11 +595,17 @@ def test_react_agent_parallel_tool_calls( if messages := event.get("messages"): message_types.append([m.type for m in messages]) - assert message_types == [ - ["human"], - ["human", "ai"], - ["human", "ai", "tool"], - ] + if version == "v1": + assert message_types == [ + ["human"], + ["human", "ai"], + ] + elif version == "v2": + assert message_types == [ + ["human"], + ["human", "ai"], + ["human", "ai", "tool"], + ] # Resume message_types = [] @@ -576,28 +621,54 @@ def test_react_agent_parallel_tool_calls( ["human", "ai", "tool", "tool", "ai"], ] - assert human_assistance_execution_count == 1 - assert get_weather_execution_count == 1 + if version == "v1": + assert human_assistance_execution_count == 1 + assert get_weather_execution_count == 2 + elif version == "v2": + assert human_assistance_execution_count == 1 + assert get_weather_execution_count == 1 class AgentStateExtraKey(AgentState): foo: int -def test_create_react_agent_inject_vars() -> None: +class AgentStateExtraKeyPydantic(AgentStatePydantic): + foo: int + + +@pytest.mark.parametrize("version", REACT_TOOL_CALL_VERSIONS) +@pytest.mark.parametrize( + "state_schema", [AgentStateExtraKey, AgentStateExtraKeyPydantic] +) +def test_create_react_agent_inject_vars( + version: Literal["v1", "v2"], state_schema: StateT +) -> None: """Test that the agent can inject state and store into tool functions.""" store = InMemoryStore() namespace = ("test",) store.put(namespace, "test_key", {"bar": 3}) - def tool1( - some_val: int, - state: Annotated[dict, InjectedState], - store: Annotated[BaseStore, InjectedStore()], - ) -> str: - """Tool 1 docstring.""" - store_val = store.get(namespace, "test_key").value["bar"] - return some_val + state["foo"] + store_val + if issubclass(state_schema, AgentStatePydantic): + + def tool1( + some_val: int, + state: Annotated[AgentStateExtraKeyPydantic, InjectedState], + store: Annotated[BaseStore, InjectedStore()], + ) -> str: + """Tool 1 docstring.""" + store_val = store.get(namespace, "test_key").value["bar"] + return some_val + state.foo + store_val + else: + + def tool1( + some_val: int, + state: Annotated[dict, InjectedState], + store: Annotated[BaseStore, InjectedStore()], + ) -> str: + """Tool 1 docstring.""" + store_val = store.get(namespace, "test_key").value["bar"] + return some_val + state["foo"] + store_val tool_call = { "name": "tool1", @@ -609,8 +680,9 @@ def test_create_react_agent_inject_vars() -> None: agent = create_react_agent( model, ToolNode([tool1], handle_tool_errors=False), - state_schema=AgentStateExtraKey, + state_schema=state_schema, store=store, + version=version, ) result = agent.invoke({"messages": [{"role": "user", "content": "hi"}], "foo": 2}) assert result["messages"] == [ @@ -622,7 +694,8 @@ def test_create_react_agent_inject_vars() -> None: assert result["foo"] == 2 -async def test_return_direct() -> None: +@pytest.mark.parametrize("version", REACT_TOOL_CALL_VERSIONS) +async def test_return_direct(version: str) -> None: @dec_tool(return_direct=True) def tool_return_direct(input: str) -> str: """A tool that returns directly.""" @@ -649,6 +722,7 @@ async def test_return_direct() -> None: agent = create_react_agent( model, [tool_return_direct, tool_normal], + version=version, ) # Test direct return for tool_return_direct @@ -673,7 +747,9 @@ async def test_return_direct() -> None: ), ] model = FakeToolCallingModel(tool_calls=[second_tool_call, []]) - agent = create_react_agent(model, [tool_return_direct, tool_normal]) + agent = create_react_agent( + model, [tool_return_direct, tool_normal], version=version + ) result = agent.invoke( {"messages": [HumanMessage(content="Test normal", id="hum1")]} ) @@ -702,7 +778,9 @@ async def test_return_direct() -> None: ), ] model = FakeToolCallingModel(tool_calls=[both_tool_calls, []]) - agent = create_react_agent(model, [tool_return_direct, tool_normal]) + agent = create_react_agent( + model, [tool_return_direct, tool_normal], version=version + ) result = agent.invoke({"messages": [HumanMessage(content="Test both", id="hum2")]}) assert result["messages"] == [ HumanMessage(content="Test both", id="hum2"), @@ -740,11 +818,12 @@ def test__get_state_args() -> None: def test_inspect_react() -> None: model = FakeToolCallingModel(tool_calls=[]) agent = create_react_agent(model, []) - inspect.getclosurevars(agent.nodes["model"].bound.func) + inspect.getclosurevars(agent.nodes["agent"].bound.func) +@pytest.mark.parametrize("version", REACT_TOOL_CALL_VERSIONS) def test_react_with_subgraph_tools( - sync_checkpointer: BaseCheckpointSaver, + sync_checkpointer: BaseCheckpointSaver, version: Literal["v1", "v2"] ) -> None: class State(TypedDict): a: int @@ -800,6 +879,7 @@ def test_react_with_subgraph_tools( model, tool_node, checkpointer=sync_checkpointer, + version=version, ) result = agent.invoke( {"messages": [HumanMessage(content="What's 2 + 3 and 2 * 3?")]}, @@ -830,7 +910,8 @@ def test_react_with_subgraph_tools( ] -def test_react_agent_subgraph_streaming_sync() -> None: +@pytest.mark.parametrize("version", REACT_TOOL_CALL_VERSIONS) +def test_react_agent_subgraph_streaming_sync(version: Literal["v1", "v2"]) -> None: """Test React agent streaming when used as a subgraph node sync version""" @dec_tool @@ -850,6 +931,7 @@ def test_react_agent_subgraph_streaming_sync() -> None: model, tools=[get_weather], prompt="You are a helpful travel assistant.", + version=version, ) # Create a subgraph that uses the React agent as a node @@ -919,7 +1001,8 @@ def test_react_agent_subgraph_streaming_sync() -> None: assert msg.content.startswith("The weather of Tokyo is sunny.") -async def test_react_agent_subgraph_streaming() -> None: +@pytest.mark.parametrize("version", REACT_TOOL_CALL_VERSIONS) +async def test_react_agent_subgraph_streaming(version: Literal["v1", "v2"]) -> None: """Test React agent streaming when used as a subgraph node.""" @dec_tool @@ -939,6 +1022,7 @@ async def test_react_agent_subgraph_streaming() -> None: model, tools=[get_weather], prompt="You are a helpful travel assistant.", + version=version, ) # Create a subgraph that uses the React agent as a node @@ -1010,8 +1094,9 @@ async def test_react_agent_subgraph_streaming() -> None: assert msg.content.startswith("The weather of Tokyo is sunny.") +@pytest.mark.parametrize("version", REACT_TOOL_CALL_VERSIONS) def test_tool_node_node_interrupt( - sync_checkpointer: BaseCheckpointSaver, + sync_checkpointer: BaseCheckpointSaver, version: str ) -> None: def tool_normal(some_val: int) -> str: """Tool docstring.""" @@ -1037,6 +1122,7 @@ def test_tool_node_node_interrupt( model, [tool_interrupt, tool_normal], checkpointer=sync_checkpointer, + version=version, ) result = agent.invoke({"messages": [HumanMessage("hi?")]}, config) expected_messages = [ @@ -1061,7 +1147,11 @@ def test_tool_node_node_interrupt( ), _AnyIdToolMessage(content="normal", name="tool_normal", tool_call_id="2"), ] - assert result["messages"] == expected_messages + if version == "v1": + # Interrupt blocks second tool result + assert result["messages"] == expected_messages[:-1] + elif version == "v2": + assert result["messages"] == expected_messages state = agent.get_state(config) assert state.next == ("tools",) @@ -1075,7 +1165,8 @@ def test_tool_node_node_interrupt( ) -def test_dynamic_model_basic() -> None: +@pytest.mark.parametrize("version", REACT_TOOL_CALL_VERSIONS) +def test_dynamic_model_basic(version: str) -> None: """Test basic dynamic model functionality.""" def dynamic_model(state, runtime: Runtime): @@ -1085,7 +1176,7 @@ def test_dynamic_model_basic() -> None: else: return FakeToolCallingModel(tool_calls=[]) - agent = create_react_agent(dynamic_model, []) + agent = create_react_agent(dynamic_model, [], version=version) result = agent.invoke({"messages": [HumanMessage("hello")]}) assert len(result["messages"]) == 2 @@ -1096,7 +1187,8 @@ def test_dynamic_model_basic() -> None: assert result["messages"][-1].content == "urgent help" -def test_dynamic_model_with_tools() -> None: +@pytest.mark.parametrize("version", REACT_TOOL_CALL_VERSIONS) +def test_dynamic_model_with_tools(version: Literal["v1", "v2"]) -> None: """Test dynamic model with tool calling.""" @dec_tool @@ -1123,7 +1215,9 @@ def test_dynamic_model_with_tools() -> None: tool_calls=[[{"args": {"x": 1}, "id": "1", "name": "basic_tool"}], []] ) - agent = create_react_agent(dynamic_model, [basic_tool, advanced_tool]) + agent = create_react_agent( + dynamic_model, [basic_tool, advanced_tool], version=version + ) # Test basic tool usage result = agent.invoke({"messages": [HumanMessage("basic request")]}) @@ -1145,7 +1239,8 @@ class Context: user_id: str -def test_dynamic_model_with_context() -> None: +@pytest.mark.parametrize("version", REACT_TOOL_CALL_VERSIONS) +def test_dynamic_model_with_context(version: str) -> None: """Test dynamic model using config parameters.""" def dynamic_model(state, runtime: Runtime[Context]): @@ -1156,7 +1251,9 @@ def test_dynamic_model_with_context() -> None: else: return FakeToolCallingModel(tool_calls=[]) - agent = create_react_agent(dynamic_model, [], context_schema=Context) + agent = create_react_agent( + dynamic_model, [], context_schema=Context, version=version + ) # Test with basic user result = agent.invoke( @@ -1173,7 +1270,8 @@ def test_dynamic_model_with_context() -> None: assert len(result["messages"]) == 2 -def test_dynamic_model_with_state_schema() -> None: +@pytest.mark.parametrize("version", REACT_TOOL_CALL_VERSIONS) +def test_dynamic_model_with_state_schema(version: Literal["v1", "v2"]) -> None: """Test dynamic model with custom state schema.""" class CustomDynamicState(AgentState): @@ -1186,7 +1284,9 @@ def test_dynamic_model_with_state_schema() -> None: else: return FakeToolCallingModel(tool_calls=[]) - agent = create_react_agent(dynamic_model, [], state_schema=CustomDynamicState) + agent = create_react_agent( + dynamic_model, [], state_schema=CustomDynamicState, version=version + ) result = agent.invoke( {"messages": [HumanMessage("hello")], "model_preference": "advanced"} @@ -1195,14 +1295,15 @@ def test_dynamic_model_with_state_schema() -> None: assert result["model_preference"] == "advanced" -def test_dynamic_model_with_prompt() -> None: +@pytest.mark.parametrize("version", REACT_TOOL_CALL_VERSIONS) +def test_dynamic_model_with_prompt(version: Literal["v1", "v2"]) -> None: """Test dynamic model with different prompt types.""" def dynamic_model(state: AgentState, runtime: Runtime) -> BaseChatModel: return FakeToolCallingModel(tool_calls=[]) # Test with string prompt - agent = create_react_agent(dynamic_model, [], prompt="system_msg") + agent = create_react_agent(dynamic_model, [], prompt="system_msg", version=version) result = agent.invoke({"messages": [HumanMessage("human_msg")]}) assert result["messages"][-1].content == "system_msg-human_msg" @@ -1211,7 +1312,9 @@ def test_dynamic_model_with_prompt() -> None: """Generate a dynamic system message based on state.""" return [{"role": "system", "content": "system_msg"}] + list(state["messages"]) - agent = create_react_agent(dynamic_model, [], prompt=dynamic_prompt) + agent = create_react_agent( + dynamic_model, [], prompt=dynamic_prompt, version=version + ) result = agent.invoke({"messages": [HumanMessage("human_msg")]}) assert result["messages"][-1].content == "system_msg-human_msg" @@ -1229,7 +1332,8 @@ async def test_dynamic_model_async() -> None: assert result["messages"][-1].content == "hello async" -def test_dynamic_model_with_structured_response() -> None: +@pytest.mark.parametrize("version", REACT_TOOL_CALL_VERSIONS) +def test_dynamic_model_with_structured_response(version: str) -> None: """Test dynamic model with structured response format.""" class TestResponse(BaseModel): @@ -1250,7 +1354,9 @@ def test_dynamic_model_with_structured_response() -> None: ], ) - agent = create_react_agent(dynamic_model, [], response_format=TestResponse) + agent = create_react_agent( + dynamic_model, [], response_format=TestResponse, version=version + ) result = agent.invoke({"messages": [HumanMessage("hello")]}) assert "structured_response" in result @@ -1289,7 +1395,8 @@ def test_dynamic_model_with_checkpointer(sync_checkpointer): assert call_count >= 2 -def test_dynamic_model_state_dependent_tools() -> None: +@pytest.mark.parametrize("version", REACT_TOOL_CALL_VERSIONS) +def test_dynamic_model_state_dependent_tools(version: Literal["v1", "v2"]) -> None: """Test dynamic model that changes available tools based on state.""" @dec_tool @@ -1313,7 +1420,7 @@ def test_dynamic_model_state_dependent_tools() -> None: tool_calls=[[{"args": {"x": 1}, "id": "1", "name": "tool_a"}], []] ) - agent = create_react_agent(dynamic_model, [tool_a, tool_b]) + agent = create_react_agent(dynamic_model, [tool_a, tool_b], version=version) # Ask to use tool B result = agent.invoke({"messages": [HumanMessage("use_b please")]}) @@ -1328,7 +1435,8 @@ def test_dynamic_model_state_dependent_tools() -> None: assert last_message.content == "A: 1" -def test_dynamic_model_error_handling() -> None: +@pytest.mark.parametrize("version", REACT_TOOL_CALL_VERSIONS) +def test_dynamic_model_error_handling(version: Literal["v1", "v2"]) -> None: """Test error handling in dynamic model.""" def failing_dynamic_model(state, runtime: Runtime): @@ -1336,7 +1444,7 @@ def test_dynamic_model_error_handling() -> None: raise ValueError("Dynamic model failed") return FakeToolCallingModel(tool_calls=[]) - agent = create_react_agent(failing_dynamic_model, []) + agent = create_react_agent(failing_dynamic_model, [], version=version) # Normal operation should work result = agent.invoke({"messages": [HumanMessage("hello")]}) @@ -1435,7 +1543,6 @@ async def test_dynamic_model_receives_correct_state_async(): assert received_state["messages"][0].content == "hello async" -@pytest.mark.skip(reason="TODO: support with prepare call") def test_pre_model_hook() -> None: model = FakeToolCallingModel(tool_calls=[]) @@ -1490,7 +1597,7 @@ def test_post_model_hook() -> None: events = list(pmh_agent.stream({"messages": [HumanMessage("hi?")], "flag": False})) assert events == [ { - "model": { + "agent": { "messages": [ AIMessage( content="hi?", @@ -1559,7 +1666,7 @@ def test_post_model_hook_with_structured_output() -> None: ) assert events == [ { - "model": { + "agent": { "messages": [ AIMessage( content="What's the weather?", @@ -1591,7 +1698,7 @@ def test_post_model_hook_with_structured_output() -> None: } }, { - "model": { + "agent": { "messages": [ AIMessage( content="What's the weather?-What's the weather?-The weather is sunny and 75°F.", @@ -1620,19 +1727,36 @@ def test_post_model_hook_with_structured_output() -> None: ] -def test_create_react_agent_inject_vars_with_post_model_hook() -> None: +@pytest.mark.parametrize( + "state_schema", [AgentStateExtraKey, AgentStateExtraKeyPydantic] +) +def test_create_react_agent_inject_vars_with_post_model_hook( + state_schema: StateT, +) -> None: store = InMemoryStore() namespace = ("test",) store.put(namespace, "test_key", {"bar": 3}) - def tool1( - some_val: int, - state: Annotated[dict, InjectedState], - store: Annotated[BaseStore, InjectedStore()], - ) -> str: - """Tool 1 docstring.""" - store_val = store.get(namespace, "test_key").value["bar"] - return some_val + state["foo"] + store_val + if issubclass(state_schema, AgentStatePydantic): + + def tool1( + some_val: int, + state: Annotated[AgentStateExtraKeyPydantic, InjectedState], + store: Annotated[BaseStore, InjectedStore()], + ) -> str: + """Tool 1 docstring.""" + store_val = store.get(namespace, "test_key").value["bar"] + return some_val + state.foo + store_val + else: + + def tool1( + some_val: int, + state: Annotated[dict, InjectedState], + store: Annotated[BaseStore, InjectedStore()], + ) -> str: + """Tool 1 docstring.""" + store_val = store.get(namespace, "test_key").value["bar"] + return some_val + state["foo"] + store_val tool_call = { "name": "tool1", @@ -1649,7 +1773,7 @@ def test_create_react_agent_inject_vars_with_post_model_hook() -> None: agent = create_react_agent( model, ToolNode([tool1], handle_tool_errors=False), - state_schema=AgentStateExtraKey, + state_schema=state_schema, store=store, post_model_hook=post_model_hook, ) @@ -1664,7 +1788,8 @@ def test_create_react_agent_inject_vars_with_post_model_hook() -> None: assert result["foo"] == 2 -def test_response_format_using_tool_choice() -> None: +@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): @@ -1685,6 +1810,7 @@ def test_response_format_using_tool_choice() -> None: model, [get_weather], response_format=WeatherResponse, + version=version, ) response = agent.invoke( { diff --git a/libs/prebuilt/tests/test_response_format.py b/libs/prebuilt/tests/test_response_format.py index 0d7ce0816..fe23a9c25 100644 --- a/libs/prebuilt/tests/test_response_format.py +++ b/libs/prebuilt/tests/test_response_format.py @@ -120,7 +120,9 @@ class TestResponseFormatAsModel: model = FakeToolCallingModel(tool_calls=tool_calls) - agent = create_react_agent(model, [get_weather], response_format=WeatherBaseModel) + agent = create_react_agent( + model, [get_weather], response_format=WeatherBaseModel + ) response = agent.invoke({"messages": [HumanMessage("What's the weather?")]}) assert response["structured_response"] == EXPECTED_WEATHER_PYDANTIC @@ -141,7 +143,9 @@ class TestResponseFormatAsModel: model = FakeToolCallingModel(tool_calls=tool_calls) - agent = create_react_agent(model, [get_weather], response_format=WeatherDataclass) + agent = create_react_agent( + model, [get_weather], response_format=WeatherDataclass + ) response = agent.invoke({"messages": [HumanMessage("What's the weather?")]}) assert response["structured_response"] == EXPECTED_WEATHER_DATACLASS @@ -162,7 +166,9 @@ class TestResponseFormatAsModel: model = FakeToolCallingModel(tool_calls=tool_calls) - agent = create_react_agent(model, [get_weather], response_format=WeatherTypedDict) + agent = create_react_agent( + model, [get_weather], response_format=WeatherTypedDict + ) response = agent.invoke({"messages": [HumanMessage("What's the weather?")]}) assert response["structured_response"] == EXPECTED_WEATHER_DICT @@ -183,7 +189,9 @@ class TestResponseFormatAsModel: model = FakeToolCallingModel(tool_calls=tool_calls) - agent = create_react_agent(model, [get_weather], response_format=weather_json_schema) + agent = create_react_agent( + model, [get_weather], response_format=weather_json_schema + ) response = agent.invoke({"messages": [HumanMessage("What's the weather?")]}) assert response["structured_response"] == EXPECTED_WEATHER_DICT