diff --git a/libs/prebuilt/langgraph/prebuilt/chat_agent_executor.py b/libs/prebuilt/langgraph/prebuilt/chat_agent_executor.py index f2c472fcf..1914cc541 100644 --- a/libs/prebuilt/langgraph/prebuilt/chat_agent_executor.py +++ b/libs/prebuilt/langgraph/prebuilt/chat_agent_executor.py @@ -1,6 +1,7 @@ import inspect from typing import ( Any, + Awaitable, Callable, Literal, Optional, @@ -44,8 +45,10 @@ from langgraph.graph.state import CompiledStateGraph from langgraph.managed import IsLastStep, RemainingSteps from langgraph.prebuilt._internal import ToolCallWithContext 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.typing import ContextT from langgraph.warnings import LangGraphDeprecatedSinceV10 StructuredResponse = Union[dict, BaseModel] @@ -246,7 +249,12 @@ def _validate_chat_history( def create_react_agent( - model: Union[str, LanguageModelLike], + model: Union[ + str, + LanguageModelLike, + Callable[[StateSchema, Runtime[ContextT]], BaseChatModel], + Callable[[StateSchema, Runtime[ContextT]], Awaitable[BaseChatModel]], + ], tools: Union[Sequence[Union[BaseTool, Callable, dict[str, Any]]], ToolNode], *, prompt: Optional[Prompt] = None, @@ -271,7 +279,43 @@ def create_react_agent( For more details on using `create_react_agent`, visit [Agents](https://langchain-ai.github.io/langgraph/agents/overview/) documentation. Args: - model: The `LangChain` chat model that supports tool calling. + model: The language model for the agent. Supports static and dynamic + model selection. + + - **Static model**: A chat model instance (e.g., `ChatOpenAI()`) or + string identifier (e.g., `"openai:gpt-4"`) + - **Dynamic model**: A callable with signature + `(state, runtime) -> BaseChatModel` that returns different models + based on runtime context + + Dynamic functions receive graph state and runtime, enabling + context-dependent model selection. Must return a `BaseChatModel` + instance. For tool calling, bind tools using `.bind_tools()`. + Bound tools must be a subset of the `tools` parameter. + + Dynamic model example: + ```python + from dataclasses import dataclass + + @dataclass + class ModelContext: + model_name: str = "gpt-3.5-turbo" + + # Instantiate models globally + gpt4_model = ChatOpenAI(model="gpt-4") + gpt35_model = ChatOpenAI(model="gpt-3.5-turbo") + + def select_model(state: AgentState, runtime: Runtime[ModelContext]) -> ChatOpenAI: + model_name = runtime.context.model_name + model = gpt4_model if model_name == "gpt-4" else gpt35_model + return model.bind_tools(tools) + ``` + + !!! note "Dynamic Model Requirements" + Ensure returned models have appropriate tools bound via + `.bind_tools()` and support required functionality. Bound tools + must be a subset of those specified in the `tools` parameter. + tools: A list of tools or a ToolNode instance. If an empty list is provided, the agent will consist of a single LLM node without tool calling. prompt: An optional prompt for the LLM. Can take a few different forms: @@ -452,32 +496,63 @@ def create_react_agent( tool_node = ToolNode([t for t in tools if not isinstance(t, dict)]) tool_classes = list(tool_node.tools_by_name.values()) - if isinstance(model, str): - try: - from langchain.chat_models import ( # type: ignore[import-not-found] - init_chat_model, - ) - except ImportError: - raise ImportError( - "Please install langchain (`pip install langchain`) to use ':' string syntax for `model` parameter." - ) - - model = cast(BaseChatModel, init_chat_model(model)) + is_dynamic_model = not isinstance(model, (str, Runnable)) and callable(model) + is_async_dynamic_model = is_dynamic_model and inspect.iscoroutinefunction(model) tool_calling_enabled = len(tool_classes) > 0 - if ( - _should_bind_tools(model, tool_classes, num_builtin=len(llm_builtin_tools)) - and len(tool_classes + llm_builtin_tools) > 0 - ): - model = cast(BaseChatModel, model).bind_tools(tool_classes + llm_builtin_tools) # type: ignore[operator] + if not is_dynamic_model: + if isinstance(model, str): + try: + from langchain.chat_models import ( # type: ignore[import-not-found] + init_chat_model, + ) + except ImportError: + raise ImportError( + "Please install langchain (`pip install langchain`) to " + "use ':' string syntax for `model` parameter." + ) - model_runnable = _get_prompt_runnable(prompt) | model + model = cast(BaseChatModel, init_chat_model(model)) + + if ( + _should_bind_tools(model, tool_classes, num_builtin=len(llm_builtin_tools)) # type: ignore[arg-type] + and len(tool_classes + llm_builtin_tools) > 0 + ): + model = cast(BaseChatModel, model).bind_tools( + tool_classes + llm_builtin_tools # type: ignore[operator] + ) + + static_model: Optional[Runnable] = _get_prompt_runnable(prompt) | model # type: ignore[operator] + else: + # For dynamic models, we'll create the runnable at runtime + static_model = None # If any of the tools are configured to return_directly after running, # our graph needs to check if these were called should_return_direct = {t.name for t in tool_classes if t.return_direct} + def _resolve_model( + state: StateSchema, runtime: Runtime[ContextT] + ) -> LanguageModelLike: + """Resolve the model to use, handling both static and dynamic models.""" + if is_dynamic_model: + return _get_prompt_runnable(prompt) | model(state, runtime) # type: ignore[operator] + else: + return static_model + + async def _aresolve_model( + state: StateSchema, runtime: Runtime[ContextT] + ) -> LanguageModelLike: + """Async resolve the model to use, handling both static and dynamic models.""" + if is_async_dynamic_model: + resolved_model = await model(state, runtime) # type: ignore[misc,operator] + return _get_prompt_runnable(prompt) | resolved_model + elif is_dynamic_model: + return _get_prompt_runnable(prompt) | model(state, runtime) # type: ignore[operator] + else: + return static_model + def _are_more_steps_needed(state: StateSchema, response: BaseMessage) -> bool: has_tool_calls = isinstance(response, AIMessage) and response.tool_calls all_tools_return_direct = ( @@ -522,9 +597,26 @@ def create_react_agent( return state # Define the function that calls the model - def call_model(state: StateSchema, config: RunnableConfig) -> StateSchema: - state = _get_model_input_state(state) - response = cast(AIMessage, model_runnable.invoke(state, config)) + def call_model( + state: StateSchema, runtime: Runtime[ContextT], config: RunnableConfig + ) -> StateSchema: + if is_async_dynamic_model: + msg = ( + "Async model callable provided but agent invoked synchronously. " + "Use agent.ainvoke() or agent.astream(), or " + "provide a sync model callable." + ) + raise RuntimeError(msg) + + model_input = _get_model_input_state(state) + + if is_dynamic_model: + # Resolve dynamic model at runtime and apply prompt + dynamic_model = _resolve_model(state, runtime) + response = cast(AIMessage, dynamic_model.invoke(model_input, config)) # type: ignore[arg-type] + else: + response = cast(AIMessage, static_model.invoke(model_input, config)) # type: ignore[union-attr] + # add agent name to the AIMessage response.name = name @@ -540,9 +632,19 @@ def create_react_agent( # We return a list, because this will get added to the existing list return {"messages": [response]} - async def acall_model(state: StateSchema, config: RunnableConfig) -> StateSchema: - state = _get_model_input_state(state) - response = cast(AIMessage, await model_runnable.ainvoke(state, config)) + async def acall_model( + state: StateSchema, runtime: Runtime[ContextT], config: RunnableConfig + ) -> StateSchema: + model_input = _get_model_input_state(state) + + if is_dynamic_model: + # Resolve dynamic model at runtime and apply prompt + # (supports both sync and async) + dynamic_model = await _aresolve_model(state, runtime) + response = cast(AIMessage, await dynamic_model.ainvoke(model_input, config)) # type: ignore[arg-type] + else: + response = cast(AIMessage, await static_model.ainvoke(model_input, config)) # type: ignore[union-attr] + # add agent name to the AIMessage response.name = name if _are_more_steps_needed(state, response): @@ -579,22 +681,32 @@ def create_react_agent( input_schema = state_schema def generate_structured_response( - state: StateSchema, config: RunnableConfig + state: StateSchema, runtime: Runtime[ContextT], config: RunnableConfig ) -> StateSchema: + if is_async_dynamic_model: + msg = ( + "Async model callable provided but agent invoked synchronously. " + "Use agent.ainvoke() or agent.astream(), or provide a sync model callable." + ) + raise RuntimeError(msg) + messages = _get_state_value(state, "messages") 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( + resolved_model = _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, config: RunnableConfig + state: StateSchema, runtime: Runtime[ContextT], config: RunnableConfig ) -> StateSchema: messages = _get_state_value(state, "messages") structured_response_schema = response_format @@ -602,7 +714,10 @@ def create_react_agent( system_prompt, structured_response_schema = response_format messages = [SystemMessage(content=system_prompt)] + list(messages) - model_with_structured_output = _get_model(model).with_structured_output( + resolved_model = await _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) diff --git a/libs/prebuilt/tests/test_react_agent.py b/libs/prebuilt/tests/test_react_agent.py index b337cb325..f70e1a88b 100644 --- a/libs/prebuilt/tests/test_react_agent.py +++ b/libs/prebuilt/tests/test_react_agent.py @@ -13,10 +13,12 @@ from typing import ( ) import pytest +from langchain_core.language_models import BaseChatModel from langchain_core.messages import ( AIMessage, AnyMessage, HumanMessage, + MessageLikeRepresentation, RemoveMessage, SystemMessage, ToolCall, @@ -52,6 +54,7 @@ from langgraph.prebuilt.tool_node import ( _get_state_args, _infer_handled_types, ) +from langgraph.runtime import Runtime from langgraph.store.base import BaseStore from langgraph.store.memory import InMemoryStore from langgraph.types import Command, Interrupt, interrupt @@ -1092,7 +1095,7 @@ def test_inspect_react() -> None: @pytest.mark.parametrize("version", REACT_TOOL_CALL_VERSIONS) def test_react_with_subgraph_tools( - sync_checkpointer: BaseCheckpointSaver, version: str + sync_checkpointer: BaseCheckpointSaver, version: Literal["v1", "v2"] ) -> None: class State(TypedDict): a: int @@ -1367,6 +1370,376 @@ def test_get_model() -> None: _get_model(RunnableLambda(lambda message: message)) +@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): + # Return different models based on state + if "urgent" in state["messages"][-1].content: + return FakeToolCallingModel(tool_calls=[]) + else: + return FakeToolCallingModel(tool_calls=[]) + + agent = create_react_agent(dynamic_model, [], version=version) + + result = agent.invoke({"messages": [HumanMessage("hello")]}) + assert len(result["messages"]) == 2 + assert result["messages"][-1].content == "hello" + + result = agent.invoke({"messages": [HumanMessage("urgent help")]}) + assert len(result["messages"]) == 2 + assert result["messages"][-1].content == "urgent help" + + +@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 + def basic_tool(x: int) -> str: + """Basic tool.""" + return f"basic: {x}" + + @dec_tool + def advanced_tool(x: int) -> str: + """Advanced tool.""" + return f"advanced: {x}" + + def dynamic_model(state: dict, runtime: Runtime) -> BaseChatModel: + # Return model with different behaviors based on message content + if "advanced" in state["messages"][-1].content: + return FakeToolCallingModel( + tool_calls=[ + [{"args": {"x": 1}, "id": "1", "name": "advanced_tool"}], + [], + ] + ) + else: + return FakeToolCallingModel( + tool_calls=[[{"args": {"x": 1}, "id": "1", "name": "basic_tool"}], []] + ) + + agent = create_react_agent( + dynamic_model, [basic_tool, advanced_tool], version=version + ) + + # Test basic tool usage + result = agent.invoke({"messages": [HumanMessage("basic request")]}) + assert len(result["messages"]) == 3 + tool_message = result["messages"][-1] + assert tool_message.content == "basic: 1" + assert tool_message.name == "basic_tool" + + # Test advanced tool usage + result = agent.invoke({"messages": [HumanMessage("advanced request")]}) + assert len(result["messages"]) == 3 + tool_message = result["messages"][-1] + assert tool_message.content == "advanced: 1" + assert tool_message.name == "advanced_tool" + + +@dataclasses.dataclass +class Context: + user_id: str + + +@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]): + # Use context to determine model behavior + user_id = runtime.context.user_id + if user_id == "user_premium": + return FakeToolCallingModel(tool_calls=[]) + else: + return FakeToolCallingModel(tool_calls=[]) + + agent = create_react_agent( + dynamic_model, [], context_schema=Context, version=version + ) + + # Test with basic user + result = agent.invoke( + {"messages": [HumanMessage("hello")]}, + context=Context(user_id="user_basic"), + ) + assert len(result["messages"]) == 2 + + # Test with premium user + result = agent.invoke( + {"messages": [HumanMessage("hello")]}, + context=Context(user_id="user_premium"), + ) + assert len(result["messages"]) == 2 + + +@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): + model_preference: str = "default" + + def dynamic_model(state: CustomDynamicState, runtime: Runtime) -> BaseChatModel: + # Use custom state field to determine model + if state.get("model_preference") == "advanced": + return FakeToolCallingModel(tool_calls=[]) + else: + return FakeToolCallingModel(tool_calls=[]) + + agent = create_react_agent( + dynamic_model, [], state_schema=CustomDynamicState, version=version + ) + + result = agent.invoke( + {"messages": [HumanMessage("hello")], "model_preference": "advanced"} + ) + assert len(result["messages"]) == 2 + assert result["model_preference"] == "advanced" + + +@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", version=version) + result = agent.invoke({"messages": [HumanMessage("human_msg")]}) + assert result["messages"][-1].content == "system_msg-human_msg" + + # Test with callable prompt + def dynamic_prompt(state: AgentState) -> list[MessageLikeRepresentation]: + """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, version=version + ) + result = agent.invoke({"messages": [HumanMessage("human_msg")]}) + assert result["messages"][-1].content == "system_msg-human_msg" + + +async def test_dynamic_model_async() -> None: + """Test dynamic model with async operations.""" + + def dynamic_model(state: AgentState, runtime: Runtime) -> BaseChatModel: + return FakeToolCallingModel(tool_calls=[]) + + agent = create_react_agent(dynamic_model, []) + + result = await agent.ainvoke({"messages": [HumanMessage("hello async")]}) + assert len(result["messages"]) == 2 + assert result["messages"][-1].content == "hello async" + + +@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): + message: str + 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 + ) + + agent = create_react_agent( + dynamic_model, [], response_format=TestResponse, version=version + ) + + result = agent.invoke({"messages": [HumanMessage("hello")]}) + assert "structured_response" in result + assert result["structured_response"].message == "dynamic response" + assert result["structured_response"].confidence == 0.9 + + +def test_dynamic_model_with_checkpointer(sync_checkpointer): + """Test dynamic model with checkpointer.""" + call_count = 0 + + def dynamic_model(state: AgentState, runtime: Runtime) -> BaseChatModel: + nonlocal call_count + call_count += 1 + return FakeToolCallingModel( + tool_calls=[], + # Incrementing the call count as it is used to assign an id + # to the AIMessage. + # The default reducer semantics are to overwrite an existing message + # with the new one if the id matches. + index=call_count, + ) + + agent = create_react_agent(dynamic_model, [], checkpointer=sync_checkpointer) + config = {"configurable": {"thread_id": "test_dynamic"}} + + # First call + result1 = agent.invoke({"messages": [HumanMessage("hello")]}, config) + assert len(result1["messages"]) == 2 # Human + AI message + + # Second call - should load from checkpoint + result2 = agent.invoke({"messages": [HumanMessage("world")]}, config) + assert len(result2["messages"]) == 4 + + # Dynamic model should be called each time + assert call_count >= 2 + + +@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 + def tool_a(x: int) -> str: + """Tool A.""" + return f"A: {x}" + + @dec_tool + def tool_b(x: int) -> str: + """Tool B.""" + return f"B: {x}" + + def dynamic_model(state, runtime: Runtime): + # Switch tools based on message history + if any("use_b" in msg.content for msg in state["messages"]): + return FakeToolCallingModel( + tool_calls=[[{"args": {"x": 2}, "id": "1", "name": "tool_b"}], []] + ) + else: + return FakeToolCallingModel( + tool_calls=[[{"args": {"x": 1}, "id": "1", "name": "tool_a"}], []] + ) + + 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")]}) + last_message = result["messages"][-1] + assert isinstance(last_message, ToolMessage) + assert last_message.content == "B: 2" + + # Ask to use tool A + result = agent.invoke({"messages": [HumanMessage("hello")]}) + last_message = result["messages"][-1] + assert isinstance(last_message, ToolMessage) + assert last_message.content == "A: 1" + + +@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): + if "fail" in state["messages"][-1].content: + raise ValueError("Dynamic model failed") + return FakeToolCallingModel(tool_calls=[]) + + agent = create_react_agent(failing_dynamic_model, [], version=version) + + # Normal operation should work + result = agent.invoke({"messages": [HumanMessage("hello")]}) + assert len(result["messages"]) == 2 + + # Should propagate the error + with pytest.raises(ValueError, match="Dynamic model failed"): + agent.invoke({"messages": [HumanMessage("fail now")]}) + + +def test_dynamic_model_vs_static_model_behavior(): + """Test that dynamic and static models produce equivalent results when configured the same.""" + # Static model + static_model = FakeToolCallingModel(tool_calls=[]) + static_agent = create_react_agent(static_model, []) + + # Dynamic model returning the same model + def dynamic_model(state, runtime: Runtime): + return FakeToolCallingModel(tool_calls=[]) + + dynamic_agent = create_react_agent(dynamic_model, []) + + input_msg = {"messages": [HumanMessage("test message")]} + + static_result = static_agent.invoke(input_msg) + dynamic_result = dynamic_agent.invoke(input_msg) + + # Results should be equivalent (content-wise, IDs may differ) + assert len(static_result["messages"]) == len(dynamic_result["messages"]) + assert static_result["messages"][0].content == dynamic_result["messages"][0].content + assert static_result["messages"][1].content == dynamic_result["messages"][1].content + + +def test_dynamic_model_receives_correct_state(): + """Test that the dynamic model function receives the correct state, not the model input.""" + received_states = [] + + class CustomAgentState(AgentState): + custom_field: str + + def dynamic_model(state, runtime: Runtime) -> BaseChatModel: + # Capture the state that's passed to the dynamic model function + received_states.append(state) + return FakeToolCallingModel(tool_calls=[]) + + agent = create_react_agent(dynamic_model, [], state_schema=CustomAgentState) + + # Test with initial state + input_state = {"messages": [HumanMessage("hello")], "custom_field": "test_value"} + agent.invoke(input_state) + + # The dynamic model function should receive the original state, not the processed model input + assert len(received_states) == 1 + received_state = received_states[0] + + # Should have the custom field from original state + assert "custom_field" in received_state + assert received_state["custom_field"] == "test_value" + + # Should have the original messages + assert len(received_state["messages"]) == 1 + assert received_state["messages"][0].content == "hello" + + +async def test_dynamic_model_receives_correct_state_async(): + """Test that the async dynamic model function receives the correct state, not the model input.""" + received_states = [] + + class CustomAgentStateAsync(AgentState): + custom_field: str + + def dynamic_model(state, runtime: Runtime): + # Capture the state that's passed to the dynamic model function + received_states.append(state) + return FakeToolCallingModel(tool_calls=[]) + + agent = create_react_agent(dynamic_model, [], state_schema=CustomAgentStateAsync) + + # Test with initial state + input_state = { + "messages": [HumanMessage("hello async")], + "custom_field": "test_value_async", + } + await agent.ainvoke(input_state) + + # The dynamic model function should receive the original state, not the processed model input + assert len(received_states) == 1 + received_state = received_states[0] + + # Should have the custom field from original state + assert "custom_field" in received_state + assert received_state["custom_field"] == "test_value_async" + + # Should have the original messages + assert len(received_state["messages"]) == 1 + assert received_state["messages"][0].content == "hello async" + + def test_pre_model_hook() -> None: model = FakeToolCallingModel(tool_calls=[])