mirror of
https://github.com/langchain-ai/langgraph.git
synced 2026-09-23 18:15:08 +02:00
feat(prebuilt): support ToolOutput response_format (#5915)
* Add support for ToolOutput response format. * I don't love the name -- it's confusing unless you know that it's parameterizing a strategy. We should determine if we want to support our old strategy for doing things -- it has a higher latency (one extra LLM call), but it's a reasonable built-in strategy as it doesn't do anything awkward with conversation history. (Wouldn't surprising if it has overall better performance than tool choice for longer conversations)
This commit is contained in:
Eugene Yurtsev
@@ -1,5 +1,5 @@
|
||||
# serializer version: 1
|
||||
# name: test_react_agent_graph_structure[None-None-None-tools0]
|
||||
# name: test_react_agent_graph_structure[None-None-tools0]
|
||||
'''
|
||||
graph TD;
|
||||
__start__ --> agent;
|
||||
@@ -7,7 +7,7 @@
|
||||
|
||||
'''
|
||||
# ---
|
||||
# name: test_react_agent_graph_structure[None-None-None-tools1]
|
||||
# name: test_react_agent_graph_structure[None-None-tools1]
|
||||
'''
|
||||
graph TD;
|
||||
__start__ --> agent;
|
||||
@@ -17,7 +17,7 @@
|
||||
|
||||
'''
|
||||
# ---
|
||||
# name: test_react_agent_graph_structure[None-None-pre_model_hook-tools0]
|
||||
# name: test_react_agent_graph_structure[None-pre_model_hook-tools0]
|
||||
'''
|
||||
graph TD;
|
||||
__start__ --> pre_model_hook;
|
||||
@@ -26,7 +26,7 @@
|
||||
|
||||
'''
|
||||
# ---
|
||||
# name: test_react_agent_graph_structure[None-None-pre_model_hook-tools1]
|
||||
# name: test_react_agent_graph_structure[None-pre_model_hook-tools1]
|
||||
'''
|
||||
graph TD;
|
||||
__start__ --> pre_model_hook;
|
||||
@@ -37,7 +37,7 @@
|
||||
|
||||
'''
|
||||
# ---
|
||||
# name: test_react_agent_graph_structure[None-post_model_hook-None-tools0]
|
||||
# name: test_react_agent_graph_structure[post_model_hook-None-tools0]
|
||||
'''
|
||||
graph TD;
|
||||
__start__ --> agent;
|
||||
@@ -46,7 +46,7 @@
|
||||
|
||||
'''
|
||||
# ---
|
||||
# name: test_react_agent_graph_structure[None-post_model_hook-None-tools1]
|
||||
# name: test_react_agent_graph_structure[post_model_hook-None-tools1]
|
||||
'''
|
||||
graph TD;
|
||||
__start__ --> agent;
|
||||
@@ -58,7 +58,7 @@
|
||||
|
||||
'''
|
||||
# ---
|
||||
# name: test_react_agent_graph_structure[None-post_model_hook-pre_model_hook-tools0]
|
||||
# name: test_react_agent_graph_structure[post_model_hook-pre_model_hook-tools0]
|
||||
'''
|
||||
graph TD;
|
||||
__start__ --> pre_model_hook;
|
||||
@@ -68,7 +68,7 @@
|
||||
|
||||
'''
|
||||
# ---
|
||||
# name: test_react_agent_graph_structure[None-post_model_hook-pre_model_hook-tools1]
|
||||
# name: test_react_agent_graph_structure[post_model_hook-pre_model_hook-tools1]
|
||||
'''
|
||||
graph TD;
|
||||
__start__ --> pre_model_hook;
|
||||
@@ -81,97 +81,7 @@
|
||||
|
||||
'''
|
||||
# ---
|
||||
# name: test_react_agent_graph_structure[ResponseFormat-None-None-tools0]
|
||||
'''
|
||||
graph TD;
|
||||
__start__ --> agent;
|
||||
agent --> generate_structured_response;
|
||||
generate_structured_response --> __end__;
|
||||
|
||||
'''
|
||||
# ---
|
||||
# name: test_react_agent_graph_structure[ResponseFormat-None-None-tools1]
|
||||
'''
|
||||
graph TD;
|
||||
__start__ --> agent;
|
||||
agent -.-> generate_structured_response;
|
||||
agent -.-> tools;
|
||||
tools --> agent;
|
||||
generate_structured_response --> __end__;
|
||||
|
||||
'''
|
||||
# ---
|
||||
# name: test_react_agent_graph_structure[ResponseFormat-None-pre_model_hook-tools0]
|
||||
'''
|
||||
graph TD;
|
||||
__start__ --> pre_model_hook;
|
||||
agent --> generate_structured_response;
|
||||
pre_model_hook --> agent;
|
||||
generate_structured_response --> __end__;
|
||||
|
||||
'''
|
||||
# ---
|
||||
# name: test_react_agent_graph_structure[ResponseFormat-None-pre_model_hook-tools1]
|
||||
'''
|
||||
graph TD;
|
||||
__start__ --> pre_model_hook;
|
||||
agent -.-> generate_structured_response;
|
||||
agent -.-> tools;
|
||||
pre_model_hook --> agent;
|
||||
tools --> pre_model_hook;
|
||||
generate_structured_response --> __end__;
|
||||
|
||||
'''
|
||||
# ---
|
||||
# name: test_react_agent_graph_structure[ResponseFormat-post_model_hook-None-tools0]
|
||||
'''
|
||||
graph TD;
|
||||
__start__ --> agent;
|
||||
agent --> post_model_hook;
|
||||
post_model_hook --> generate_structured_response;
|
||||
generate_structured_response --> __end__;
|
||||
|
||||
'''
|
||||
# ---
|
||||
# name: test_react_agent_graph_structure[ResponseFormat-post_model_hook-None-tools1]
|
||||
'''
|
||||
graph TD;
|
||||
__start__ --> agent;
|
||||
agent --> post_model_hook;
|
||||
post_model_hook -.-> agent;
|
||||
post_model_hook -.-> generate_structured_response;
|
||||
post_model_hook -.-> tools;
|
||||
tools --> agent;
|
||||
generate_structured_response --> __end__;
|
||||
|
||||
'''
|
||||
# ---
|
||||
# name: test_react_agent_graph_structure[ResponseFormat-post_model_hook-pre_model_hook-tools0]
|
||||
'''
|
||||
graph TD;
|
||||
__start__ --> pre_model_hook;
|
||||
agent --> post_model_hook;
|
||||
post_model_hook --> generate_structured_response;
|
||||
pre_model_hook --> agent;
|
||||
generate_structured_response --> __end__;
|
||||
|
||||
'''
|
||||
# ---
|
||||
# name: test_react_agent_graph_structure[ResponseFormat-post_model_hook-pre_model_hook-tools1]
|
||||
'''
|
||||
graph TD;
|
||||
__start__ --> pre_model_hook;
|
||||
agent --> post_model_hook;
|
||||
post_model_hook -.-> generate_structured_response;
|
||||
post_model_hook -.-> pre_model_hook;
|
||||
post_model_hook -.-> tools;
|
||||
pre_model_hook --> agent;
|
||||
tools --> pre_model_hook;
|
||||
generate_structured_response --> __end__;
|
||||
|
||||
'''
|
||||
# ---
|
||||
# name: test_react_agent_graph_structure_with_individual_nodes[no_response_format-no_post_hook-no_pre_hook-no_tools]
|
||||
# name: test_react_agent_graph_structure_with_individual_nodes[no_post_hook-no_pre_hook-no_tools]
|
||||
'''
|
||||
graph TD;
|
||||
__start__ --> agent;
|
||||
@@ -179,7 +89,7 @@
|
||||
|
||||
'''
|
||||
# ---
|
||||
# name: test_react_agent_graph_structure_with_individual_nodes[no_response_format-no_post_hook-no_pre_hook-two_tools]
|
||||
# name: test_react_agent_graph_structure_with_individual_nodes[no_post_hook-no_pre_hook-two_tools]
|
||||
'''
|
||||
graph TD;
|
||||
__start__ --> agent;
|
||||
@@ -191,7 +101,7 @@
|
||||
|
||||
'''
|
||||
# ---
|
||||
# name: test_react_agent_graph_structure_with_individual_nodes[no_response_format-no_post_hook-with_pre_hook-no_tools]
|
||||
# name: test_react_agent_graph_structure_with_individual_nodes[no_post_hook-with_pre_hook-no_tools]
|
||||
'''
|
||||
graph TD;
|
||||
__start__ --> pre_model_hook;
|
||||
@@ -200,7 +110,7 @@
|
||||
|
||||
'''
|
||||
# ---
|
||||
# name: test_react_agent_graph_structure_with_individual_nodes[no_response_format-no_post_hook-with_pre_hook-two_tools]
|
||||
# name: test_react_agent_graph_structure_with_individual_nodes[no_post_hook-with_pre_hook-two_tools]
|
||||
'''
|
||||
graph TD;
|
||||
__start__ --> pre_model_hook;
|
||||
@@ -213,7 +123,7 @@
|
||||
|
||||
'''
|
||||
# ---
|
||||
# name: test_react_agent_graph_structure_with_individual_nodes[no_response_format-with_post_hook-no_pre_hook-no_tools]
|
||||
# name: test_react_agent_graph_structure_with_individual_nodes[with_post_hook-no_pre_hook-no_tools]
|
||||
'''
|
||||
graph TD;
|
||||
__start__ --> agent;
|
||||
@@ -222,7 +132,7 @@
|
||||
|
||||
'''
|
||||
# ---
|
||||
# name: test_react_agent_graph_structure_with_individual_nodes[no_response_format-with_post_hook-no_pre_hook-two_tools]
|
||||
# name: test_react_agent_graph_structure_with_individual_nodes[with_post_hook-no_pre_hook-two_tools]
|
||||
'''
|
||||
graph TD;
|
||||
__start__ --> agent;
|
||||
@@ -236,7 +146,7 @@
|
||||
|
||||
'''
|
||||
# ---
|
||||
# name: test_react_agent_graph_structure_with_individual_nodes[no_response_format-with_post_hook-with_pre_hook-no_tools]
|
||||
# name: test_react_agent_graph_structure_with_individual_nodes[with_post_hook-with_pre_hook-no_tools]
|
||||
'''
|
||||
graph TD;
|
||||
__start__ --> pre_model_hook;
|
||||
@@ -246,7 +156,7 @@
|
||||
|
||||
'''
|
||||
# ---
|
||||
# name: test_react_agent_graph_structure_with_individual_nodes[no_response_format-with_post_hook-with_pre_hook-two_tools]
|
||||
# name: test_react_agent_graph_structure_with_individual_nodes[with_post_hook-with_pre_hook-two_tools]
|
||||
'''
|
||||
graph TD;
|
||||
__start__ --> pre_model_hook;
|
||||
@@ -261,101 +171,3 @@
|
||||
|
||||
'''
|
||||
# ---
|
||||
# name: test_react_agent_graph_structure_with_individual_nodes[with_response_format-no_post_hook-no_pre_hook-no_tools]
|
||||
'''
|
||||
graph TD;
|
||||
__start__ --> agent;
|
||||
agent --> generate_structured_response;
|
||||
generate_structured_response --> __end__;
|
||||
|
||||
'''
|
||||
# ---
|
||||
# name: test_react_agent_graph_structure_with_individual_nodes[with_response_format-no_post_hook-no_pre_hook-two_tools]
|
||||
'''
|
||||
graph TD;
|
||||
__start__ --> agent;
|
||||
agent -.-> generate_structured_response;
|
||||
agent -.-> tool;
|
||||
agent -.-> tool2;
|
||||
tool --> agent;
|
||||
tool2 --> agent;
|
||||
generate_structured_response --> __end__;
|
||||
|
||||
'''
|
||||
# ---
|
||||
# name: test_react_agent_graph_structure_with_individual_nodes[with_response_format-no_post_hook-with_pre_hook-no_tools]
|
||||
'''
|
||||
graph TD;
|
||||
__start__ --> pre_model_hook;
|
||||
agent --> generate_structured_response;
|
||||
pre_model_hook --> agent;
|
||||
generate_structured_response --> __end__;
|
||||
|
||||
'''
|
||||
# ---
|
||||
# name: test_react_agent_graph_structure_with_individual_nodes[with_response_format-no_post_hook-with_pre_hook-two_tools]
|
||||
'''
|
||||
graph TD;
|
||||
__start__ --> pre_model_hook;
|
||||
agent -.-> generate_structured_response;
|
||||
agent -.-> tool;
|
||||
agent -.-> tool2;
|
||||
pre_model_hook --> agent;
|
||||
tool --> pre_model_hook;
|
||||
tool2 --> pre_model_hook;
|
||||
generate_structured_response --> __end__;
|
||||
|
||||
'''
|
||||
# ---
|
||||
# name: test_react_agent_graph_structure_with_individual_nodes[with_response_format-with_post_hook-no_pre_hook-no_tools]
|
||||
'''
|
||||
graph TD;
|
||||
__start__ --> agent;
|
||||
agent --> post_model_hook;
|
||||
post_model_hook --> generate_structured_response;
|
||||
generate_structured_response --> __end__;
|
||||
|
||||
'''
|
||||
# ---
|
||||
# name: test_react_agent_graph_structure_with_individual_nodes[with_response_format-with_post_hook-no_pre_hook-two_tools]
|
||||
'''
|
||||
graph TD;
|
||||
__start__ --> agent;
|
||||
agent --> post_model_hook;
|
||||
post_model_hook -.-> agent;
|
||||
post_model_hook -.-> generate_structured_response;
|
||||
post_model_hook -.-> tool;
|
||||
post_model_hook -.-> tool2;
|
||||
tool --> agent;
|
||||
tool2 --> agent;
|
||||
generate_structured_response --> __end__;
|
||||
|
||||
'''
|
||||
# ---
|
||||
# name: test_react_agent_graph_structure_with_individual_nodes[with_response_format-with_post_hook-with_pre_hook-no_tools]
|
||||
'''
|
||||
graph TD;
|
||||
__start__ --> pre_model_hook;
|
||||
agent --> post_model_hook;
|
||||
post_model_hook --> generate_structured_response;
|
||||
pre_model_hook --> agent;
|
||||
generate_structured_response --> __end__;
|
||||
|
||||
'''
|
||||
# ---
|
||||
# name: test_react_agent_graph_structure_with_individual_nodes[with_response_format-with_post_hook-with_pre_hook-two_tools]
|
||||
'''
|
||||
graph TD;
|
||||
__start__ --> pre_model_hook;
|
||||
agent --> post_model_hook;
|
||||
post_model_hook -.-> generate_structured_response;
|
||||
post_model_hook -.-> pre_model_hook;
|
||||
post_model_hook -.-> tool;
|
||||
post_model_hook -.-> tool2;
|
||||
pre_model_hook --> agent;
|
||||
tool --> pre_model_hook;
|
||||
tool2 --> pre_model_hook;
|
||||
generate_structured_response --> __end__;
|
||||
|
||||
'''
|
||||
# ---
|
||||
|
||||
@@ -453,7 +453,10 @@ def test_react_agent_with_structured_response(version: str) -> None:
|
||||
class WeatherResponse(BaseModel):
|
||||
temperature: float = Field(description="The temperature in fahrenheit")
|
||||
|
||||
tool_calls = [[{"args": {}, "id": "1", "name": "get_weather"}], []]
|
||||
tool_calls = [
|
||||
[{"args": {}, "id": "1", "name": "get_weather"}],
|
||||
[{"name": "WeatherResponse", "id": "2", "args": {"temperature": 75}}],
|
||||
]
|
||||
|
||||
def get_weather():
|
||||
"""Get the weather"""
|
||||
@@ -463,17 +466,33 @@ def test_react_agent_with_structured_response(version: str) -> None:
|
||||
model = FakeToolCallingModel(
|
||||
tool_calls=tool_calls, structured_response=expected_structured_response
|
||||
)
|
||||
for response_format in (WeatherResponse, ("Meow", WeatherResponse)):
|
||||
agent = create_react_agent(
|
||||
model,
|
||||
[get_weather],
|
||||
response_format=response_format,
|
||||
version=version,
|
||||
)
|
||||
response = agent.invoke({"messages": [HumanMessage("What's the weather?")]})
|
||||
assert response["structured_response"] == expected_structured_response
|
||||
assert len(response["messages"]) == 4
|
||||
assert response["messages"][-2].content == "The weather is sunny and 75°F."
|
||||
agent = create_react_agent(
|
||||
model,
|
||||
[get_weather],
|
||||
response_format=WeatherResponse,
|
||||
version=version,
|
||||
)
|
||||
response = agent.invoke({"messages": [HumanMessage("What's the weather?")]})
|
||||
assert response["structured_response"] == expected_structured_response
|
||||
assert len(response["messages"]) == 5
|
||||
|
||||
# Check message types in message history
|
||||
msg_types = [m.type for m in response["messages"]]
|
||||
assert msg_types == [
|
||||
"human", # "What's the weather?"
|
||||
"ai", # "What's the weather?"
|
||||
"tool", # "The weather is sunny and 75°F."
|
||||
"ai", # structured response
|
||||
"tool", # ok!
|
||||
]
|
||||
|
||||
assert [m.content for m in response["messages"]] == [
|
||||
"What's the weather?",
|
||||
"What's the weather?",
|
||||
"The weather is sunny and 75°F.",
|
||||
"What's the weather?-What's the weather?-The weather is sunny and 75°F.",
|
||||
"ok!",
|
||||
]
|
||||
|
||||
|
||||
class CustomState(AgentState):
|
||||
@@ -1410,9 +1429,17 @@ def test_dynamic_model_with_structured_response(version: str) -> None:
|
||||
confidence: float
|
||||
|
||||
def dynamic_model(state, runtime: Runtime):
|
||||
expected_response = TestResponse(message="dynamic response", confidence=0.9)
|
||||
return FakeToolCallingModel(
|
||||
tool_calls=[], structured_response=expected_response
|
||||
tool_calls=[
|
||||
[
|
||||
ToolCall(
|
||||
name="TestResponse",
|
||||
args={"message": "dynamic response", "confidence": 0.9},
|
||||
id="1",
|
||||
type="tool_call",
|
||||
)
|
||||
]
|
||||
],
|
||||
)
|
||||
|
||||
agent = create_react_agent(
|
||||
@@ -1677,16 +1704,16 @@ def test_post_model_hook_with_structured_output() -> None:
|
||||
class WeatherResponse(BaseModel):
|
||||
temperature: float = Field(description="The temperature in fahrenheit")
|
||||
|
||||
tool_calls = [[{"args": {}, "id": "1", "name": "get_weather"}]]
|
||||
tool_calls: list[list[ToolCall]] = [
|
||||
[{"args": {}, "id": "1", "name": "get_weather"}],
|
||||
[{"args": {"temperature": 75}, "id": "2", "name": "WeatherResponse"}],
|
||||
]
|
||||
|
||||
def get_weather():
|
||||
"""Get the weather"""
|
||||
return "The weather is sunny and 75°F."
|
||||
|
||||
expected_structured_response = WeatherResponse(temperature=75)
|
||||
model = FakeToolCallingModel(
|
||||
tool_calls=tool_calls, structured_response=expected_structured_response
|
||||
)
|
||||
|
||||
class State(AgentState):
|
||||
flag: bool
|
||||
@@ -1695,6 +1722,7 @@ def test_post_model_hook_with_structured_output() -> None:
|
||||
def post_model_hook(state: State) -> Union[dict[str, bool], Command]:
|
||||
return {"flag": True}
|
||||
|
||||
model = FakeToolCallingModel(tool_calls=tool_calls)
|
||||
agent = create_react_agent(
|
||||
model,
|
||||
[get_weather],
|
||||
@@ -1704,7 +1732,6 @@ def test_post_model_hook_with_structured_output() -> None:
|
||||
)
|
||||
|
||||
assert "post_model_hook" in agent.nodes
|
||||
assert "generate_structured_response" in agent.nodes
|
||||
|
||||
response = agent.invoke(
|
||||
{"messages": [HumanMessage("What's the weather?")], "flag": False}
|
||||
@@ -1712,10 +1739,19 @@ def test_post_model_hook_with_structured_output() -> None:
|
||||
assert response["flag"] is True
|
||||
assert response["structured_response"] == expected_structured_response
|
||||
|
||||
# Reset the state of the model
|
||||
model = FakeToolCallingModel(tool_calls=tool_calls)
|
||||
agent = create_react_agent(
|
||||
model,
|
||||
[get_weather],
|
||||
response_format=WeatherResponse,
|
||||
post_model_hook=post_model_hook,
|
||||
state_schema=State,
|
||||
)
|
||||
|
||||
events = list(
|
||||
agent.stream({"messages": [HumanMessage("What's the weather?")], "flag": False})
|
||||
)
|
||||
assert "generate_structured_response" in events[-1]
|
||||
assert events == [
|
||||
{
|
||||
"agent": {
|
||||
@@ -1724,7 +1760,7 @@ def test_post_model_hook_with_structured_output() -> None:
|
||||
content="What's the weather?",
|
||||
additional_kwargs={},
|
||||
response_metadata={},
|
||||
id="2",
|
||||
id="0",
|
||||
tool_calls=[
|
||||
{
|
||||
"name": "get_weather",
|
||||
@@ -1745,7 +1781,7 @@ def test_post_model_hook_with_structured_output() -> None:
|
||||
content="The weather is sunny and 75°F.",
|
||||
name="get_weather",
|
||||
tool_call_id="1",
|
||||
),
|
||||
)
|
||||
]
|
||||
}
|
||||
},
|
||||
@@ -1756,25 +1792,26 @@ def test_post_model_hook_with_structured_output() -> None:
|
||||
content="What's the weather?-What's the weather?-The weather is sunny and 75°F.",
|
||||
additional_kwargs={},
|
||||
response_metadata={},
|
||||
id="3",
|
||||
id="1",
|
||||
tool_calls=[
|
||||
{
|
||||
"name": "get_weather",
|
||||
"args": {},
|
||||
"id": "1",
|
||||
"name": "WeatherResponse",
|
||||
"args": {"temperature": 75},
|
||||
"id": "2",
|
||||
"type": "tool_call",
|
||||
}
|
||||
],
|
||||
)
|
||||
]
|
||||
),
|
||||
_AnyIdToolMessage(
|
||||
content="ok!",
|
||||
name="WeatherResponse",
|
||||
tool_call_id="2",
|
||||
),
|
||||
],
|
||||
"structured_response": WeatherResponse(temperature=75.0),
|
||||
}
|
||||
},
|
||||
{"post_model_hook": {"flag": True}},
|
||||
{
|
||||
"generate_structured_response": {
|
||||
"structured_response": WeatherResponse(temperature=75.0)
|
||||
}
|
||||
},
|
||||
]
|
||||
|
||||
|
||||
@@ -1837,3 +1874,40 @@ def test_create_react_agent_inject_vars_with_post_model_hook(
|
||||
AIMessage("hi-hi-6", id="1"),
|
||||
]
|
||||
assert result["foo"] == 2
|
||||
|
||||
|
||||
@pytest.mark.parametrize("version", REACT_TOOL_CALL_VERSIONS)
|
||||
def test_response_format_using_tool_choice(version: Literal["v1", "v2"]) -> None:
|
||||
"""Test response format using tool choice."""
|
||||
|
||||
class WeatherResponse(BaseModel):
|
||||
temperature: float = Field(description="The temperature in fahrenheit")
|
||||
|
||||
tool_calls: list[list[ToolCall]] = [
|
||||
[{"args": {}, "id": "1", "name": "get_weather"}],
|
||||
[{"args": {"temperature": "75"}, "id": "2", "name": "WeatherResponse"}],
|
||||
]
|
||||
|
||||
def get_weather() -> str:
|
||||
"""Get the weather"""
|
||||
return "The weather is sunny and 75°F."
|
||||
|
||||
expected_structured_response = WeatherResponse(temperature=75)
|
||||
model = FakeToolCallingModel(tool_calls=tool_calls)
|
||||
agent = create_react_agent(
|
||||
model,
|
||||
[get_weather],
|
||||
response_format=WeatherResponse,
|
||||
version=version,
|
||||
)
|
||||
response = agent.invoke(
|
||||
{
|
||||
"messages": [
|
||||
{
|
||||
"role": "user",
|
||||
"content": "What's the weather?",
|
||||
}
|
||||
]
|
||||
}
|
||||
)
|
||||
assert response.get("structured_response") == expected_structured_response
|
||||
|
||||
@@ -39,20 +39,17 @@ class ResponseFormat(BaseModel):
|
||||
@pytest.mark.parametrize("tools", [[], [tool]])
|
||||
@pytest.mark.parametrize("pre_model_hook", [None, pre_model_hook])
|
||||
@pytest.mark.parametrize("post_model_hook", [None, post_model_hook])
|
||||
@pytest.mark.parametrize("response_format", [None, ResponseFormat])
|
||||
def test_react_agent_graph_structure(
|
||||
snapshot: SnapshotAssertion,
|
||||
tools: list[Callable],
|
||||
pre_model_hook: Union[Callable, None],
|
||||
post_model_hook: Union[Callable, None],
|
||||
response_format: Union[type[BaseModel], None],
|
||||
) -> None:
|
||||
agent = create_react_agent(
|
||||
model,
|
||||
tools=tools,
|
||||
pre_model_hook=pre_model_hook,
|
||||
post_model_hook=post_model_hook,
|
||||
response_format=response_format,
|
||||
)
|
||||
try:
|
||||
assert agent.get_graph().draw_mermaid(with_styles=False) == snapshot
|
||||
@@ -63,7 +60,6 @@ def test_react_agent_graph_structure(
|
||||
f"tools: {tools}, "
|
||||
f"pre_model_hook: {pre_model_hook}, "
|
||||
f"post_model_hook: {post_model_hook}, "
|
||||
f"response_format: {response_format}"
|
||||
) from e
|
||||
|
||||
|
||||
@@ -74,24 +70,17 @@ def test_react_agent_graph_structure(
|
||||
@pytest.mark.parametrize(
|
||||
"post_model_hook", [None, post_model_hook], ids=["no_post_hook", "with_post_hook"]
|
||||
)
|
||||
@pytest.mark.parametrize(
|
||||
"response_format",
|
||||
[None, ResponseFormat],
|
||||
ids=["no_response_format", "with_response_format"],
|
||||
)
|
||||
def test_react_agent_graph_structure_with_individual_nodes(
|
||||
snapshot: SnapshotAssertion,
|
||||
tools: list[Callable],
|
||||
pre_model_hook: Union[Callable, None],
|
||||
post_model_hook: Union[Callable, None],
|
||||
response_format: Union[type[BaseModel], None],
|
||||
) -> None:
|
||||
agent = create_react_agent(
|
||||
model,
|
||||
tools=tools,
|
||||
pre_model_hook=pre_model_hook,
|
||||
post_model_hook=post_model_hook,
|
||||
response_format=response_format,
|
||||
use_individual_tool_nodes=True,
|
||||
)
|
||||
assert agent.get_graph().draw_mermaid(with_styles=False) == snapshot
|
||||
|
||||
@@ -0,0 +1,225 @@
|
||||
"""Unit tests for langgraph.prebuilt.responses module."""
|
||||
|
||||
import pytest
|
||||
from pydantic import BaseModel
|
||||
|
||||
from langgraph.prebuilt.responses import (
|
||||
OutputToolBinding,
|
||||
ResponseFormat,
|
||||
SchemaSpec,
|
||||
ToolOutput,
|
||||
)
|
||||
|
||||
|
||||
class TestModel(BaseModel):
|
||||
"""A test model for structured output."""
|
||||
|
||||
name: str
|
||||
age: int
|
||||
email: str = "default@example.com"
|
||||
|
||||
|
||||
class CustomModel(BaseModel):
|
||||
"""Custom model with a custom docstring."""
|
||||
|
||||
value: float
|
||||
description: str
|
||||
|
||||
|
||||
class EmptyDocModel(BaseModel):
|
||||
# No custom docstring, should have no description in tool
|
||||
data: str
|
||||
|
||||
|
||||
class TestSchemaSpec:
|
||||
"""Test SchemaSpec dataclass."""
|
||||
|
||||
def test_basic_creation(self):
|
||||
"""Test basic SchemaSpec creation."""
|
||||
schema = SchemaSpec(schema=TestModel)
|
||||
assert schema.schema == TestModel
|
||||
assert schema.name is None
|
||||
assert schema.description is None
|
||||
assert schema.strict is False
|
||||
|
||||
def test_creation_with_all_fields(self):
|
||||
"""Test SchemaSpec creation with all fields."""
|
||||
schema = SchemaSpec(
|
||||
schema=TestModel,
|
||||
name="custom_test_model",
|
||||
description="A custom description",
|
||||
strict=True,
|
||||
)
|
||||
assert schema.schema == TestModel
|
||||
assert schema.name == "custom_test_model"
|
||||
assert schema.description == "A custom description"
|
||||
assert schema.strict is True
|
||||
|
||||
|
||||
class TestUsingToolStrategy:
|
||||
"""Test UsingToolStrategy dataclass."""
|
||||
|
||||
def test_basic_creation(self):
|
||||
"""Test basic UsingToolStrategy creation."""
|
||||
schema = SchemaSpec(schema=TestModel)
|
||||
strategy = ToolOutput(schemas=[schema])
|
||||
assert len(strategy.schemas) == 1
|
||||
assert strategy.schemas[0] == schema
|
||||
assert strategy.tool_choice == "required" # default
|
||||
|
||||
def test_creation_with_auto_tool_choice(self):
|
||||
"""Test UsingToolStrategy creation with auto tool choice."""
|
||||
schema = SchemaSpec(schema=TestModel)
|
||||
strategy = ToolOutput(schemas=[schema], tool_choice="auto")
|
||||
assert strategy.tool_choice == "auto"
|
||||
|
||||
def test_multiple_schemas(self):
|
||||
"""Test UsingToolStrategy with multiple schemas."""
|
||||
schema1 = SchemaSpec(schema=TestModel)
|
||||
schema2 = SchemaSpec(schema=CustomModel)
|
||||
strategy = ToolOutput(schemas=[schema1, schema2])
|
||||
assert len(strategy.schemas) == 2
|
||||
|
||||
|
||||
class TestOutputToolBinding:
|
||||
"""Test OutputToolBinding dataclass and its methods."""
|
||||
|
||||
def test_from_schema_spec_basic(self):
|
||||
"""Test basic OutputToolBinding creation from SchemaSpec."""
|
||||
schema_spec = SchemaSpec(schema=TestModel)
|
||||
tool_binding = OutputToolBinding.from_schema_spec(schema_spec)
|
||||
|
||||
assert tool_binding.schema == TestModel
|
||||
assert tool_binding.schema_kind == "pydantic"
|
||||
assert tool_binding.tool is not None
|
||||
assert tool_binding.tool.name == "TestModel"
|
||||
|
||||
def test_from_schema_spec_with_custom_name(self):
|
||||
"""Test OutputToolBinding creation with custom name."""
|
||||
schema_spec = SchemaSpec(schema=TestModel, name="custom_tool_name")
|
||||
tool_binding = OutputToolBinding.from_schema_spec(schema_spec)
|
||||
assert tool_binding.tool.name == "custom_tool_name"
|
||||
|
||||
def test_from_schema_spec_with_custom_description(self):
|
||||
"""Test OutputToolBinding creation with custom description."""
|
||||
schema_spec = SchemaSpec(
|
||||
schema=TestModel, description="Custom tool description"
|
||||
)
|
||||
tool_binding = OutputToolBinding.from_schema_spec(schema_spec)
|
||||
|
||||
assert tool_binding.tool.description == "Custom tool description"
|
||||
|
||||
def test_from_schema_spec_with_model_docstring(self):
|
||||
"""Test OutputToolBinding creation using model docstring as description."""
|
||||
schema_spec = SchemaSpec(schema=CustomModel)
|
||||
tool_binding = OutputToolBinding.from_schema_spec(schema_spec)
|
||||
|
||||
assert tool_binding.tool.description == "Custom model with a custom docstring."
|
||||
|
||||
@pytest.mark.skip(
|
||||
reason="Need to fix bug in langchain-core for inheritance of doc-strings."
|
||||
)
|
||||
def test_from_schema_spec_empty_docstring(self):
|
||||
"""Test OutputToolBinding creation with model that has default docstring."""
|
||||
|
||||
# Create a model with the same docstring as BaseModel
|
||||
class DefaultDocModel(BaseModel):
|
||||
# This should have the same docstring as BaseModel
|
||||
pass
|
||||
|
||||
schema_spec = SchemaSpec(schema=DefaultDocModel)
|
||||
tool_binding = OutputToolBinding.from_schema_spec(schema_spec)
|
||||
|
||||
# Should use empty description when model has default BaseModel docstring
|
||||
assert tool_binding.tool.description == ""
|
||||
|
||||
def test_parse_payload_pydantic_success(self):
|
||||
"""Test successful parsing for Pydantic model."""
|
||||
schema_spec = SchemaSpec(schema=TestModel)
|
||||
tool_binding = OutputToolBinding.from_schema_spec(schema_spec)
|
||||
|
||||
tool_args = {"name": "John", "age": 30}
|
||||
result = tool_binding.parse(tool_args)
|
||||
|
||||
assert isinstance(result, TestModel)
|
||||
assert result.name == "John"
|
||||
assert result.age == 30
|
||||
assert result.email == "default@example.com" # default value
|
||||
|
||||
def test_parse_payload_pydantic_validation_error(self):
|
||||
"""Test parsing failure for invalid Pydantic data."""
|
||||
schema_spec = SchemaSpec(schema=TestModel)
|
||||
tool_binding = OutputToolBinding.from_schema_spec(schema_spec)
|
||||
|
||||
# Missing required field 'name'
|
||||
tool_args = {"age": 30}
|
||||
|
||||
with pytest.raises(ValueError, match="Failed to parse tool args to TestModel"):
|
||||
tool_binding.parse(tool_args)
|
||||
|
||||
def test_parse_payload_invalid_kind(self):
|
||||
"""Test parsing with invalid kind."""
|
||||
from unittest.mock import Mock
|
||||
|
||||
mock_tool = Mock()
|
||||
|
||||
tool_binding = OutputToolBinding(
|
||||
schema=TestModel,
|
||||
schema_kind="invalid_kind", # type: ignore
|
||||
tool=mock_tool,
|
||||
)
|
||||
|
||||
with pytest.raises(ValueError, match="Unsupported schema kind: invalid_kind"):
|
||||
tool_binding.parse({"name": "test", "age": 25})
|
||||
|
||||
def test_parse_payload_invalid_pydantic_schema(self):
|
||||
"""Test parsing with invalid schema for pydantic kind."""
|
||||
from unittest.mock import Mock
|
||||
|
||||
mock_tool = Mock()
|
||||
|
||||
# Create tool binding with dict schema but pydantic kind
|
||||
tool_binding = OutputToolBinding(
|
||||
schema={"type": "object"}, schema_kind="pydantic", tool=mock_tool
|
||||
)
|
||||
|
||||
with pytest.raises(
|
||||
ValueError, match="Expected Pydantic model class for 'pydantic' kind"
|
||||
):
|
||||
tool_binding.parse({"name": "test", "age": 25})
|
||||
|
||||
|
||||
class TestResponseFormat:
|
||||
"""Test ResponseFormat type alias."""
|
||||
|
||||
def test_response_format_is_using_tool_strategy(self):
|
||||
"""Test that ResponseFormat is aliased to UsingToolStrategy."""
|
||||
assert ResponseFormat is ToolOutput
|
||||
|
||||
def test_can_create_response_format(self):
|
||||
"""Test that we can create ResponseFormat instances."""
|
||||
schema = SchemaSpec(schema=TestModel)
|
||||
response_format = ResponseFormat(schemas=[schema])
|
||||
|
||||
assert isinstance(response_format, ToolOutput)
|
||||
assert len(response_format.schemas) == 1
|
||||
|
||||
|
||||
class TestEdgeCases:
|
||||
"""Test edge cases and error conditions."""
|
||||
|
||||
def test_empty_schemas_list(self) -> None:
|
||||
"""Test UsingToolStrategy with empty schemas list."""
|
||||
strategy = ToolOutput([SchemaSpec(EmptyDocModel)])
|
||||
assert len(strategy.schemas) == 1
|
||||
|
||||
@pytest.mark.skip(
|
||||
reason="Need to fix bug in langchain-core for inheritance of doc-strings."
|
||||
)
|
||||
def test_base_model_doc_constant(self) -> None:
|
||||
"""Test that BASE_MODEL_DOC constant is set correctly."""
|
||||
binding = OutputToolBinding.from_schema_spec(SchemaSpec(EmptyDocModel))
|
||||
assert binding.tool.name == "EmptyDocModel"
|
||||
assert (
|
||||
binding.tool.description[:5] == ""
|
||||
) # Should be empty for default docstring
|
||||
Reference in New Issue
Block a user