mirror of
https://github.com/langchain-ai/langgraph.git
synced 2026-08-17 21:25:46 +02:00
feat(prebuilt): Add dynamic model to create_react_agent (#5651)
This PR allows a developer to change the model configuration at run time based on context. This includes that list of tools available to the model to call.
```python
def create_react_agent(
model: Union[
str,
LanguageModelLike,
Callable[[SateLike, Runtime...], BaseChatModel], # <--- New
],
tools: Union[
Sequence[Union[BaseTool, Callable, dict[str, Any]]], ToolNode]
],
*,
....
llm = init_chat_model(...)
def prepare_model(state, runtime):
selected_tool_names = func(state, context)
return llm.bind(tools=selected_tool_names)
create_react_agent(
prepare_model,
tools=all_known_tools
)
```
## Semantics
1. `tools` = are the known tools, used to configure ToolNode and will
configure:
1. model provided as string
2. model provided as BaseChatModel (if it has no tools bound to it)
2. If a user provides a dynamic model (callable), the user is
responsible for binding tools
Alternative considered:
1. Passing `Callable[[SateLike, Config...], list[BaseTool]]` to tools
2. Passing `Callable[[SateLike, Config...], list[str]]` to a tool
selector
Both have the issue that there's non obvious interplay between tool
selection and dynamic models. (i.e., if we want to introduce dynamic
models at in the future, the API will become tricky to explain)
---------
Co-authored-by: Sydney Runkle <54324534+sydney-runkle@users.noreply.github.com>
This commit is contained in:
Eugene Yurtsev
co-authored by
Sydney Runkle
Sydney Runkle
parent
8495f6f95d
commit
f6aa19709e
@@ -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 '<provider>:<model>' 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 '<provider>:<model>' 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)
|
||||
|
||||
@@ -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=[])
|
||||
|
||||
|
||||
Reference in New Issue
Block a user