mirror of
https://github.com/langchain-ai/langgraph.git
synced 2026-08-22 07:32:25 +02:00
Add test from harrison, always produce output when resuming graph
This commit is contained in:
@@ -18,20 +18,18 @@ start-postgres:
|
||||
docker compose -f tests/compose-postgres.yml up -V --force-recreate --wait
|
||||
|
||||
stop-postgres:
|
||||
docker compose -f tests/compose-postgres.yml down
|
||||
docker compose -f tests/compose-postgres.yml down -v
|
||||
|
||||
TEST_PATH ?= .
|
||||
|
||||
test:
|
||||
make start-postgres; \
|
||||
poetry run pytest $(TEST_PATH); \
|
||||
make start-postgres && poetry run pytest $(TEST_PATH); \
|
||||
EXIT_CODE=$$?; \
|
||||
make stop-postgres; \
|
||||
exit $$EXIT_CODE
|
||||
|
||||
test_watch:
|
||||
make start-postgres; \
|
||||
poetry run ptw . -- --ff -v -x -n auto --dist worksteal --snapshot-update --tb short $(TEST_PATH); \
|
||||
make start-postgres && poetry run ptw . -- --ff -v -x -n auto --dist worksteal --snapshot-update --tb short $(TEST_PATH); \
|
||||
EXIT_CODE=$$?; \
|
||||
make stop-postgres; \
|
||||
exit $$EXIT_CODE
|
||||
|
||||
@@ -73,15 +73,19 @@ class AddableValuesDict(AddableDict):
|
||||
|
||||
def map_output_values(
|
||||
output_channels: Union[str, Sequence[str]],
|
||||
pending_writes: Sequence[tuple[str, Any]],
|
||||
pending_writes: Union[True, Sequence[tuple[str, Any]]],
|
||||
channels: Mapping[str, BaseChannel],
|
||||
) -> Iterator[Union[dict[str, Any], Any]]:
|
||||
"""Map pending writes (a sequence of tuples (channel, value)) to output chunk."""
|
||||
if isinstance(output_channels, str):
|
||||
if any(chan == output_channels for chan, _ in pending_writes):
|
||||
if pending_writes is True or any(
|
||||
chan == output_channels for chan, _ in pending_writes
|
||||
):
|
||||
yield read_channel(channels, output_channels)
|
||||
else:
|
||||
if {c for c, _ in pending_writes if c in output_channels}:
|
||||
if pending_writes is True or {
|
||||
c for c, _ in pending_writes if c in output_channels
|
||||
}:
|
||||
yield AddableValuesDict(read_channels(channels, output_channels))
|
||||
|
||||
|
||||
|
||||
@@ -401,6 +401,11 @@ class PregelLoop:
|
||||
if k in self.checkpoint["channel_versions"]:
|
||||
version = self.checkpoint["channel_versions"][k]
|
||||
self.checkpoint["versions_seen"][INTERRUPT][k] = version
|
||||
# produce values output
|
||||
self.stream.extend(
|
||||
(self.config["configurable"].get("checkpoint_ns", ""), "values", v)
|
||||
for v in map_output_values(self.output_keys, True, self.channels)
|
||||
)
|
||||
# map inputs to channel updates
|
||||
elif input_writes := deque(map_input(input_keys, self.input)):
|
||||
# discard any unfinished tasks from previous checkpoint
|
||||
|
||||
@@ -5099,6 +5099,131 @@
|
||||
# name: test_state_graph_w_config_inherited_state_keys.2
|
||||
'{"title": "LangGraphOutput", "type": "object", "properties": {"input": {"title": "Input", "type": "string"}, "agent_outcome": {"title": "Agent Outcome", "anyOf": [{"$ref": "#/definitions/AgentAction"}, {"$ref": "#/definitions/AgentFinish"}]}, "intermediate_steps": {"title": "Intermediate Steps", "type": "array", "items": {"type": "array", "minItems": 2, "maxItems": 2, "items": [{"$ref": "#/definitions/AgentAction"}, {"type": "string"}]}}}, "definitions": {"AgentAction": {"title": "AgentAction", "description": "Represents a request to execute an action by an agent.\\n\\nThe action consists of the name of the tool to execute and the input to pass\\nto the tool. The log is used to pass along extra information about the action.", "type": "object", "properties": {"tool": {"title": "Tool", "type": "string"}, "tool_input": {"title": "Tool Input", "anyOf": [{"type": "string"}, {"type": "object"}]}, "log": {"title": "Log", "type": "string"}, "type": {"title": "Type", "default": "AgentAction", "enum": ["AgentAction"], "type": "string"}}, "required": ["tool", "tool_input", "log"]}, "AgentFinish": {"title": "AgentFinish", "description": "Final return value of an ActionAgent.\\n\\nAgents return an AgentFinish when they have reached a stopping condition.", "type": "object", "properties": {"return_values": {"title": "Return Values", "type": "object"}, "log": {"title": "Log", "type": "string"}, "type": {"title": "Type", "default": "AgentFinish", "enum": ["AgentFinish"], "type": "string"}}, "required": ["return_values", "log"]}}}'
|
||||
# ---
|
||||
# name: test_weather_subgraph[memory]
|
||||
'''
|
||||
%%{init: {'flowchart': {'curve': 'linear'}}}%%
|
||||
graph TD;
|
||||
__start__([__start__]):::first
|
||||
router_node(router_node)
|
||||
normal_llm_node(normal_llm_node)
|
||||
weather_graph_model_node(model_node)
|
||||
weather_graph_weather_node(weather_node<hr/><small><em>__interrupt = before</em></small>)
|
||||
__end__([__end__]):::last
|
||||
subgraph weather_graph
|
||||
weather_graph_model_node --> weather_graph_weather_node;
|
||||
end
|
||||
__start__ --> router_node;
|
||||
normal_llm_node --> __end__;
|
||||
weather_graph_weather_node --> __end__;
|
||||
router_node -.-> normal_llm_node;
|
||||
router_node -.-> weather_graph_model_node;
|
||||
router_node -.-> __end__;
|
||||
classDef default fill:#f2f0ff,line-height:1.2
|
||||
classDef first fill-opacity:0
|
||||
classDef last fill:#bfb6fc
|
||||
|
||||
'''
|
||||
# ---
|
||||
# name: test_weather_subgraph[postgres]
|
||||
'''
|
||||
%%{init: {'flowchart': {'curve': 'linear'}}}%%
|
||||
graph TD;
|
||||
__start__([__start__]):::first
|
||||
router_node(router_node)
|
||||
normal_llm_node(normal_llm_node)
|
||||
weather_graph_model_node(model_node)
|
||||
weather_graph_weather_node(weather_node<hr/><small><em>__interrupt = before</em></small>)
|
||||
__end__([__end__]):::last
|
||||
subgraph weather_graph
|
||||
weather_graph_model_node --> weather_graph_weather_node;
|
||||
end
|
||||
__start__ --> router_node;
|
||||
normal_llm_node --> __end__;
|
||||
weather_graph_weather_node --> __end__;
|
||||
router_node -.-> normal_llm_node;
|
||||
router_node -.-> weather_graph_model_node;
|
||||
router_node -.-> __end__;
|
||||
classDef default fill:#f2f0ff,line-height:1.2
|
||||
classDef first fill-opacity:0
|
||||
classDef last fill:#bfb6fc
|
||||
|
||||
'''
|
||||
# ---
|
||||
# name: test_weather_subgraph[postgres_pipe]
|
||||
'''
|
||||
%%{init: {'flowchart': {'curve': 'linear'}}}%%
|
||||
graph TD;
|
||||
__start__([__start__]):::first
|
||||
router_node(router_node)
|
||||
normal_llm_node(normal_llm_node)
|
||||
weather_graph_model_node(model_node)
|
||||
weather_graph_weather_node(weather_node<hr/><small><em>__interrupt = before</em></small>)
|
||||
__end__([__end__]):::last
|
||||
subgraph weather_graph
|
||||
weather_graph_model_node --> weather_graph_weather_node;
|
||||
end
|
||||
__start__ --> router_node;
|
||||
normal_llm_node --> __end__;
|
||||
weather_graph_weather_node --> __end__;
|
||||
router_node -.-> normal_llm_node;
|
||||
router_node -.-> weather_graph_model_node;
|
||||
router_node -.-> __end__;
|
||||
classDef default fill:#f2f0ff,line-height:1.2
|
||||
classDef first fill-opacity:0
|
||||
classDef last fill:#bfb6fc
|
||||
|
||||
'''
|
||||
# ---
|
||||
# name: test_weather_subgraph[postgres_pool]
|
||||
'''
|
||||
%%{init: {'flowchart': {'curve': 'linear'}}}%%
|
||||
graph TD;
|
||||
__start__([__start__]):::first
|
||||
router_node(router_node)
|
||||
normal_llm_node(normal_llm_node)
|
||||
weather_graph_model_node(model_node)
|
||||
weather_graph_weather_node(weather_node<hr/><small><em>__interrupt = before</em></small>)
|
||||
__end__([__end__]):::last
|
||||
subgraph weather_graph
|
||||
weather_graph_model_node --> weather_graph_weather_node;
|
||||
end
|
||||
__start__ --> router_node;
|
||||
normal_llm_node --> __end__;
|
||||
weather_graph_weather_node --> __end__;
|
||||
router_node -.-> normal_llm_node;
|
||||
router_node -.-> weather_graph_model_node;
|
||||
router_node -.-> __end__;
|
||||
classDef default fill:#f2f0ff,line-height:1.2
|
||||
classDef first fill-opacity:0
|
||||
classDef last fill:#bfb6fc
|
||||
|
||||
'''
|
||||
# ---
|
||||
# name: test_weather_subgraph[sqlite]
|
||||
'''
|
||||
%%{init: {'flowchart': {'curve': 'linear'}}}%%
|
||||
graph TD;
|
||||
__start__([__start__]):::first
|
||||
router_node(router_node)
|
||||
normal_llm_node(normal_llm_node)
|
||||
weather_graph_model_node(model_node)
|
||||
weather_graph_weather_node(weather_node<hr/><small><em>__interrupt = before</em></small>)
|
||||
__end__([__end__]):::last
|
||||
subgraph weather_graph
|
||||
weather_graph_model_node --> weather_graph_weather_node;
|
||||
end
|
||||
__start__ --> router_node;
|
||||
normal_llm_node --> __end__;
|
||||
weather_graph_weather_node --> __end__;
|
||||
router_node -.-> normal_llm_node;
|
||||
router_node -.-> weather_graph_model_node;
|
||||
router_node -.-> __end__;
|
||||
classDef default fill:#f2f0ff,line-height:1.2
|
||||
classDef first fill-opacity:0
|
||||
classDef last fill:#bfb6fc
|
||||
|
||||
'''
|
||||
# ---
|
||||
# name: test_xray_issue
|
||||
'''
|
||||
%%{init: {'flowchart': {'curve': 'linear'}}}%%
|
||||
|
||||
@@ -20,6 +20,7 @@ from typing import (
|
||||
Tuple,
|
||||
TypedDict,
|
||||
Union,
|
||||
cast,
|
||||
get_type_hints,
|
||||
)
|
||||
|
||||
@@ -2293,6 +2294,16 @@ def test_conditional_graph(
|
||||
)
|
||||
|
||||
assert [c for c in app_w_interrupt.stream(None, config)] == [
|
||||
{
|
||||
"agent": {
|
||||
"agent_outcome": AgentAction(
|
||||
tool="search_api",
|
||||
tool_input="query",
|
||||
log="tool:search_api:a different query",
|
||||
),
|
||||
"input": "what is weather in sf",
|
||||
}
|
||||
},
|
||||
{
|
||||
"tools": {
|
||||
"input": "what is weather in sf",
|
||||
@@ -2503,6 +2514,16 @@ def test_conditional_graph(
|
||||
)
|
||||
|
||||
assert [c for c in app_w_interrupt.stream(None, config)] == [
|
||||
{
|
||||
"agent": {
|
||||
"agent_outcome": AgentAction(
|
||||
tool="search_api",
|
||||
tool_input="query",
|
||||
log="tool:search_api:a different query",
|
||||
),
|
||||
"input": "what is weather in sf",
|
||||
},
|
||||
},
|
||||
{
|
||||
"tools": {
|
||||
"input": "what is weather in sf",
|
||||
@@ -2668,6 +2689,14 @@ def test_conditional_graph(
|
||||
)
|
||||
|
||||
assert [c for c in app_w_interrupt.stream(None, config)] == [
|
||||
{
|
||||
"agent": {
|
||||
"input": "what is weather in sf",
|
||||
"agent_outcome": AgentAction(
|
||||
tool="search_api", tool_input="query", log="tool:search_api:query"
|
||||
),
|
||||
},
|
||||
},
|
||||
{
|
||||
"tools": {
|
||||
"input": "what is weather in sf",
|
||||
@@ -2706,6 +2735,26 @@ def test_conditional_graph(
|
||||
]
|
||||
|
||||
assert [c for c in app_w_interrupt.stream(None, config)] == [
|
||||
{
|
||||
"agent": {
|
||||
"input": "what is weather in sf",
|
||||
"intermediate_steps": [
|
||||
[
|
||||
AgentAction(
|
||||
tool="search_api",
|
||||
tool_input="query",
|
||||
log="tool:search_api:query",
|
||||
),
|
||||
"result for query",
|
||||
]
|
||||
],
|
||||
"agent_outcome": AgentAction(
|
||||
tool="search_api",
|
||||
tool_input="another",
|
||||
log="tool:search_api:another",
|
||||
),
|
||||
}
|
||||
},
|
||||
{
|
||||
"tools": {
|
||||
"input": "what is weather in sf",
|
||||
@@ -8396,17 +8445,12 @@ def test_nested_graph_interrupts_parallel(
|
||||
# test stream values w/ nested interrupt
|
||||
config = {"configurable": {"thread_id": "3"}}
|
||||
assert [*app.stream({"my_key": ""}, config, stream_mode="values")] == [
|
||||
{
|
||||
"my_key": "",
|
||||
},
|
||||
{"my_key": ""},
|
||||
]
|
||||
assert [*app.stream(None, config, stream_mode="values")] == [
|
||||
{
|
||||
"my_key": "got here and there and parallel",
|
||||
},
|
||||
{
|
||||
"my_key": "got here and there and parallel and back again",
|
||||
},
|
||||
{"my_key": ""},
|
||||
{"my_key": "got here and there and parallel"},
|
||||
{"my_key": "got here and there and parallel and back again"},
|
||||
]
|
||||
|
||||
# test interrupts BEFORE the parallel node
|
||||
@@ -8416,14 +8460,13 @@ def test_nested_graph_interrupts_parallel(
|
||||
{"my_key": ""}
|
||||
]
|
||||
# while we're waiting for the node w/ interrupt inside to finish
|
||||
assert [*app.stream(None, config, stream_mode="values")] == []
|
||||
assert [*app.stream(None, config, stream_mode="values")] == [
|
||||
{
|
||||
"my_key": "got here and there and parallel",
|
||||
},
|
||||
{
|
||||
"my_key": "got here and there and parallel and back again",
|
||||
},
|
||||
{"my_key": ""},
|
||||
]
|
||||
assert [*app.stream(None, config, stream_mode="values")] == [
|
||||
{"my_key": ""},
|
||||
{"my_key": "got here and there and parallel"},
|
||||
{"my_key": "got here and there and parallel and back again"},
|
||||
]
|
||||
|
||||
# test interrupts AFTER the parallel node
|
||||
@@ -8433,12 +8476,12 @@ def test_nested_graph_interrupts_parallel(
|
||||
{"my_key": ""}
|
||||
]
|
||||
assert [*app.stream(None, config, stream_mode="values")] == [
|
||||
{"my_key": ""},
|
||||
{"my_key": "got here and there and parallel"},
|
||||
]
|
||||
assert [*app.stream(None, config, stream_mode="values")] == [
|
||||
{
|
||||
"my_key": "got here and there and parallel and back again",
|
||||
},
|
||||
{"my_key": "got here and there and parallel"},
|
||||
{"my_key": "got here and there and parallel and back again"},
|
||||
]
|
||||
|
||||
|
||||
@@ -8520,20 +8563,13 @@ def test_doubly_nested_graph_interrupts(
|
||||
# test stream values w/ nested interrupt
|
||||
config = {"configurable": {"thread_id": "3"}}
|
||||
assert [*app.stream({"my_key": "my value"}, config, stream_mode="values")] == [
|
||||
{
|
||||
"my_key": "my value",
|
||||
},
|
||||
{
|
||||
"my_key": "hi my value",
|
||||
},
|
||||
{"my_key": "my value"},
|
||||
{"my_key": "hi my value"},
|
||||
]
|
||||
assert [*app.stream(None, config, stream_mode="values")] == [
|
||||
{
|
||||
"my_key": "hi my value here and there",
|
||||
},
|
||||
{
|
||||
"my_key": "hi my value here and there and back again",
|
||||
},
|
||||
{"my_key": "hi my value"},
|
||||
{"my_key": "hi my value here and there"},
|
||||
{"my_key": "hi my value here and there and back again"},
|
||||
]
|
||||
|
||||
|
||||
@@ -10102,6 +10138,389 @@ def test_send_to_nested_graphs(
|
||||
assert actual_history == expected_history
|
||||
|
||||
|
||||
@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC)
|
||||
def test_weather_subgraph(
|
||||
request: pytest.FixtureRequest, checkpointer_name: str, snapshot: SnapshotAssertion
|
||||
) -> None:
|
||||
from langchain_core.language_models.fake_chat_models import (
|
||||
FakeMessagesListChatModel,
|
||||
)
|
||||
from langchain_core.messages import AIMessage, HumanMessage, ToolCall
|
||||
from langchain_core.tools import tool
|
||||
|
||||
from langgraph.graph import MessagesState
|
||||
|
||||
checkpointer = request.getfixturevalue("checkpointer_" + checkpointer_name)
|
||||
|
||||
# setup subgraph
|
||||
|
||||
@tool
|
||||
def get_weather(city: str):
|
||||
"""Get the weather for a specific city"""
|
||||
return f"I'ts sunny in {city}!"
|
||||
|
||||
weather_model = FakeMessagesListChatModel(
|
||||
responses=[
|
||||
AIMessage(
|
||||
content="",
|
||||
tool_calls=[
|
||||
ToolCall(
|
||||
id="tool_call123",
|
||||
name="get_weather",
|
||||
args={"city": "San Francisco"},
|
||||
)
|
||||
],
|
||||
)
|
||||
]
|
||||
)
|
||||
|
||||
class SubGraphState(MessagesState):
|
||||
city: str
|
||||
|
||||
def model_node(state: SubGraphState):
|
||||
result = weather_model.invoke(state["messages"])
|
||||
return {"city": cast(AIMessage, result).tool_calls[0]["args"]["city"]}
|
||||
|
||||
def weather_node(state: SubGraphState):
|
||||
result = get_weather.invoke({"city": state["city"]})
|
||||
return {"messages": [{"role": "assistant", "content": result}]}
|
||||
|
||||
subgraph = StateGraph(SubGraphState)
|
||||
subgraph.add_node(model_node)
|
||||
subgraph.add_node(weather_node)
|
||||
subgraph.add_edge(START, "model_node")
|
||||
subgraph.add_edge("model_node", "weather_node")
|
||||
subgraph.add_edge("weather_node", END)
|
||||
subgraph = subgraph.compile(interrupt_before=["weather_node"])
|
||||
|
||||
# setup main graph
|
||||
|
||||
class RouterState(MessagesState):
|
||||
route: Literal["weather", "other"]
|
||||
|
||||
class Router(TypedDict):
|
||||
route: Literal["weather", "other"]
|
||||
|
||||
router_model = FakeMessagesListChatModel(
|
||||
responses=[
|
||||
AIMessage(
|
||||
content="",
|
||||
tool_calls=[
|
||||
ToolCall(
|
||||
id="tool_call123",
|
||||
name="router",
|
||||
args={"dest": "weather"},
|
||||
)
|
||||
],
|
||||
)
|
||||
]
|
||||
)
|
||||
|
||||
def router_node(state: RouterState):
|
||||
system_message = "Classify the incoming query as either about weather or not."
|
||||
messages = [{"role": "system", "content": system_message}] + state["messages"]
|
||||
route = router_model.invoke(messages)
|
||||
return {"route": cast(AIMessage, route).tool_calls[0]["args"]["dest"]}
|
||||
|
||||
def normal_llm_node(state: RouterState):
|
||||
return {"messages": [AIMessage("Hello!")]}
|
||||
|
||||
def route_after_prediction(state: RouterState):
|
||||
if state["route"] == "weather":
|
||||
return "weather_graph"
|
||||
else:
|
||||
return "normal_llm_node"
|
||||
|
||||
graph = StateGraph(RouterState)
|
||||
graph.add_node(router_node)
|
||||
graph.add_node(normal_llm_node)
|
||||
graph.add_node("weather_graph", subgraph)
|
||||
graph.add_edge(START, "router_node")
|
||||
graph.add_conditional_edges("router_node", route_after_prediction)
|
||||
graph.add_edge("normal_llm_node", END)
|
||||
graph.add_edge("weather_graph", END)
|
||||
graph = graph.compile(checkpointer=checkpointer)
|
||||
|
||||
assert graph.get_graph(xray=1).draw_mermaid() == snapshot
|
||||
|
||||
config = {"configurable": {"thread_id": "1"}}
|
||||
inputs = {"messages": [{"role": "user", "content": "what's the weather in sf"}]}
|
||||
|
||||
# run until interrupt
|
||||
assert [
|
||||
c
|
||||
for c in graph.stream(
|
||||
inputs, config=config, stream_mode="updates", subgraphs=True
|
||||
)
|
||||
] == [
|
||||
((), {"router_node": {"route": "weather"}}),
|
||||
((AnyStr("weather_graph:"),), {"model_node": {"city": "San Francisco"}}),
|
||||
]
|
||||
|
||||
# check current state
|
||||
state = graph.get_state(config)
|
||||
assert state == StateSnapshot(
|
||||
values={
|
||||
"messages": [HumanMessage(content="what's the weather in sf", id=AnyStr())],
|
||||
"route": "weather",
|
||||
},
|
||||
next=("weather_graph",),
|
||||
config={
|
||||
"configurable": {
|
||||
"thread_id": "1",
|
||||
"checkpoint_ns": "",
|
||||
"checkpoint_id": AnyStr(),
|
||||
}
|
||||
},
|
||||
metadata={
|
||||
"source": "loop",
|
||||
"writes": {"router_node": {"route": "weather"}},
|
||||
"step": 1,
|
||||
"parents": {},
|
||||
},
|
||||
created_at=AnyStr(),
|
||||
parent_config={
|
||||
"configurable": {
|
||||
"thread_id": "1",
|
||||
"checkpoint_ns": "",
|
||||
"checkpoint_id": AnyStr(),
|
||||
}
|
||||
},
|
||||
tasks=(
|
||||
PregelTask(
|
||||
id=AnyStr(),
|
||||
name="weather_graph",
|
||||
state={
|
||||
"configurable": {
|
||||
"thread_id": "1",
|
||||
"checkpoint_ns": AnyStr("weather_graph:"),
|
||||
}
|
||||
},
|
||||
),
|
||||
),
|
||||
)
|
||||
|
||||
# update
|
||||
graph.update_state(state.tasks[0].state, {"city": "la"})
|
||||
|
||||
# run after update
|
||||
assert [
|
||||
c
|
||||
for c in graph.stream(
|
||||
None, config=config, stream_mode="updates", subgraphs=True
|
||||
)
|
||||
] == [
|
||||
(
|
||||
(AnyStr("weather_graph:"),),
|
||||
{
|
||||
"weather_node": {
|
||||
"messages": [{"role": "assistant", "content": "I'ts sunny in la!"}]
|
||||
}
|
||||
},
|
||||
),
|
||||
(
|
||||
(),
|
||||
{
|
||||
"weather_graph": {
|
||||
"messages": [
|
||||
HumanMessage(content="what's the weather in sf", id=AnyStr()),
|
||||
AIMessage(content="I'ts sunny in la!", id=AnyStr()),
|
||||
]
|
||||
}
|
||||
},
|
||||
),
|
||||
]
|
||||
|
||||
# try updating acting as weather node
|
||||
config = {"configurable": {"thread_id": "14"}}
|
||||
inputs = {"messages": [{"role": "user", "content": "what's the weather in sf"}]}
|
||||
assert [
|
||||
c
|
||||
for c in graph.stream(
|
||||
inputs, config=config, stream_mode="updates", subgraphs=True
|
||||
)
|
||||
] == [
|
||||
((), {"router_node": {"route": "weather"}}),
|
||||
((AnyStr("weather_graph:"),), {"model_node": {"city": "San Francisco"}}),
|
||||
]
|
||||
state = graph.get_state(config, subgraphs=True)
|
||||
assert state == StateSnapshot(
|
||||
values={
|
||||
"messages": [HumanMessage(content="what's the weather in sf", id=AnyStr())],
|
||||
"route": "weather",
|
||||
},
|
||||
next=("weather_graph",),
|
||||
config={
|
||||
"configurable": {
|
||||
"thread_id": "14",
|
||||
"checkpoint_ns": "",
|
||||
"checkpoint_id": AnyStr(),
|
||||
}
|
||||
},
|
||||
metadata={
|
||||
"source": "loop",
|
||||
"writes": {"router_node": {"route": "weather"}},
|
||||
"step": 1,
|
||||
"parents": {},
|
||||
},
|
||||
created_at=AnyStr(),
|
||||
parent_config={
|
||||
"configurable": {
|
||||
"thread_id": "14",
|
||||
"checkpoint_ns": "",
|
||||
"checkpoint_id": AnyStr(),
|
||||
}
|
||||
},
|
||||
tasks=(
|
||||
PregelTask(
|
||||
id=AnyStr(),
|
||||
name="weather_graph",
|
||||
state=StateSnapshot(
|
||||
values={
|
||||
"messages": [
|
||||
HumanMessage(
|
||||
content="what's the weather in sf", id=AnyStr()
|
||||
)
|
||||
],
|
||||
"city": "San Francisco",
|
||||
},
|
||||
next=("weather_node",),
|
||||
config={
|
||||
"configurable": {
|
||||
"thread_id": "14",
|
||||
"checkpoint_ns": AnyStr("weather_graph:"),
|
||||
"checkpoint_id": AnyStr(),
|
||||
"checkpoint_map": AnyDict(
|
||||
{
|
||||
"": AnyStr(),
|
||||
AnyStr("weather_graph:"): AnyStr(),
|
||||
}
|
||||
),
|
||||
}
|
||||
},
|
||||
metadata={
|
||||
"source": "loop",
|
||||
"writes": {"model_node": {"city": "San Francisco"}},
|
||||
"step": 1,
|
||||
"parents": {"": AnyStr()},
|
||||
},
|
||||
created_at=AnyStr(),
|
||||
parent_config={
|
||||
"configurable": {
|
||||
"thread_id": "14",
|
||||
"checkpoint_ns": AnyStr("weather_graph:"),
|
||||
"checkpoint_id": AnyStr(),
|
||||
}
|
||||
},
|
||||
tasks=(PregelTask(id=AnyStr(), name="weather_node"),),
|
||||
),
|
||||
),
|
||||
),
|
||||
)
|
||||
graph.update_state(
|
||||
state.tasks[0].state.config,
|
||||
{"messages": [{"role": "assistant", "content": "rainy"}]},
|
||||
as_node="weather_node",
|
||||
)
|
||||
state = graph.get_state(config, subgraphs=True)
|
||||
assert state == StateSnapshot(
|
||||
values={
|
||||
"messages": [HumanMessage(content="what's the weather in sf", id=AnyStr())],
|
||||
"route": "weather",
|
||||
},
|
||||
next=("weather_graph",),
|
||||
config={
|
||||
"configurable": {
|
||||
"thread_id": "14",
|
||||
"checkpoint_ns": "",
|
||||
"checkpoint_id": AnyStr(),
|
||||
}
|
||||
},
|
||||
metadata={
|
||||
"source": "loop",
|
||||
"writes": {"router_node": {"route": "weather"}},
|
||||
"step": 1,
|
||||
"parents": {},
|
||||
},
|
||||
created_at=AnyStr(),
|
||||
parent_config={
|
||||
"configurable": {
|
||||
"thread_id": "14",
|
||||
"checkpoint_ns": "",
|
||||
"checkpoint_id": AnyStr(),
|
||||
}
|
||||
},
|
||||
tasks=(
|
||||
PregelTask(
|
||||
id=AnyStr(),
|
||||
name="weather_graph",
|
||||
state=StateSnapshot(
|
||||
values={
|
||||
"messages": [
|
||||
HumanMessage(
|
||||
content="what's the weather in sf", id=AnyStr()
|
||||
),
|
||||
AIMessage(content="rainy", id=AnyStr()),
|
||||
],
|
||||
"city": "San Francisco",
|
||||
},
|
||||
next=(),
|
||||
config={
|
||||
"configurable": {
|
||||
"thread_id": "14",
|
||||
"checkpoint_ns": AnyStr("weather_graph:"),
|
||||
"checkpoint_id": AnyStr(),
|
||||
"checkpoint_map": AnyDict(
|
||||
{
|
||||
"": AnyStr(),
|
||||
AnyStr("weather_graph:"): AnyStr(),
|
||||
}
|
||||
),
|
||||
}
|
||||
},
|
||||
metadata={
|
||||
"source": "update",
|
||||
"step": 2,
|
||||
"writes": {
|
||||
"weather_node": {
|
||||
"messages": [{"role": "assistant", "content": "rainy"}]
|
||||
}
|
||||
},
|
||||
"parents": {"": AnyStr()},
|
||||
},
|
||||
created_at=AnyStr(),
|
||||
parent_config={
|
||||
"configurable": {
|
||||
"thread_id": "14",
|
||||
"checkpoint_ns": AnyStr("weather_graph:"),
|
||||
"checkpoint_id": AnyStr(),
|
||||
}
|
||||
},
|
||||
tasks=(),
|
||||
),
|
||||
),
|
||||
),
|
||||
)
|
||||
assert [
|
||||
c
|
||||
for c in graph.stream(
|
||||
None, config=config, stream_mode="updates", subgraphs=True
|
||||
)
|
||||
] == [
|
||||
(
|
||||
(),
|
||||
{
|
||||
"weather_graph": {
|
||||
"messages": [
|
||||
HumanMessage(content="what's the weather in sf", id=AnyStr()),
|
||||
AIMessage(content="rainy", id=AnyStr()),
|
||||
]
|
||||
}
|
||||
},
|
||||
),
|
||||
]
|
||||
|
||||
|
||||
def test_repeat_condition(snapshot: SnapshotAssertion) -> None:
|
||||
class AgentState(TypedDict):
|
||||
hello: str
|
||||
|
||||
@@ -347,7 +347,7 @@ async def test_node_not_cancelled_on_other_node_interrupted(
|
||||
assert not inner_task_cancelled
|
||||
assert awhiles == 1
|
||||
|
||||
assert await graph.ainvoke(None, thread, debug=True) is None
|
||||
assert await graph.ainvoke(None, thread, debug=True) == {"hello": "world"}
|
||||
|
||||
assert not inner_task_cancelled
|
||||
assert awhiles == 1
|
||||
@@ -2572,6 +2572,16 @@ async def test_conditional_graph(
|
||||
)
|
||||
|
||||
assert [c async for c in app_w_interrupt.astream(None, config)] == [
|
||||
{
|
||||
"agent": {
|
||||
"agent_outcome": AgentAction(
|
||||
tool="search_api",
|
||||
tool_input="query",
|
||||
log="tool:search_api:a different query",
|
||||
),
|
||||
"input": "what is weather in sf",
|
||||
},
|
||||
},
|
||||
{
|
||||
"tools": {
|
||||
"input": "what is weather in sf",
|
||||
@@ -2797,6 +2807,16 @@ async def test_conditional_graph(
|
||||
)
|
||||
|
||||
assert [c async for c in app_w_interrupt.astream(None, config)] == [
|
||||
{
|
||||
"agent": {
|
||||
"agent_outcome": AgentAction(
|
||||
tool="search_api",
|
||||
tool_input="query",
|
||||
log="tool:search_api:a different query",
|
||||
),
|
||||
"input": "what is weather in sf",
|
||||
},
|
||||
},
|
||||
{
|
||||
"tools": {
|
||||
"input": "what is weather in sf",
|
||||
@@ -2973,6 +2993,14 @@ async def test_conditional_graph(
|
||||
)
|
||||
|
||||
assert [c async for c in app_w_interrupt.astream(None, config)] == [
|
||||
{
|
||||
"agent": {
|
||||
"input": "what is weather in sf",
|
||||
"agent_outcome": AgentAction(
|
||||
tool="search_api", tool_input="query", log="tool:search_api:query"
|
||||
),
|
||||
},
|
||||
},
|
||||
{
|
||||
"tools": {
|
||||
"input": "what is weather in sf",
|
||||
@@ -3011,6 +3039,26 @@ async def test_conditional_graph(
|
||||
]
|
||||
|
||||
assert [c async for c in app_w_interrupt.astream(None, config)] == [
|
||||
{
|
||||
"agent": {
|
||||
"input": "what is weather in sf",
|
||||
"intermediate_steps": [
|
||||
[
|
||||
AgentAction(
|
||||
tool="search_api",
|
||||
tool_input="query",
|
||||
log="tool:search_api:query",
|
||||
),
|
||||
"result for query",
|
||||
]
|
||||
],
|
||||
"agent_outcome": AgentAction(
|
||||
tool="search_api",
|
||||
tool_input="another",
|
||||
log="tool:search_api:another",
|
||||
),
|
||||
}
|
||||
},
|
||||
{
|
||||
"tools": {
|
||||
"input": "what is weather in sf",
|
||||
@@ -6891,6 +6939,7 @@ async def test_nested_graph_interrupts_parallel(
|
||||
{"my_key": ""},
|
||||
]
|
||||
assert [c async for c in app.astream(None, config, stream_mode="values")] == [
|
||||
{"my_key": ""},
|
||||
{"my_key": "got here and there and parallel"},
|
||||
{"my_key": "got here and there and parallel and back again"},
|
||||
]
|
||||
@@ -6904,8 +6953,11 @@ async def test_nested_graph_interrupts_parallel(
|
||||
{"my_key": ""},
|
||||
]
|
||||
# while we're waiting for the node w/ interrupt inside to finish
|
||||
assert [c async for c in app.astream(None, config, stream_mode="values")] == []
|
||||
assert [c async for c in app.astream(None, config, stream_mode="values")] == [
|
||||
{"my_key": ""},
|
||||
]
|
||||
assert [c async for c in app.astream(None, config, stream_mode="values")] == [
|
||||
{"my_key": ""},
|
||||
{"my_key": "got here and there and parallel"},
|
||||
{"my_key": "got here and there and parallel and back again"},
|
||||
]
|
||||
@@ -6919,9 +6971,11 @@ async def test_nested_graph_interrupts_parallel(
|
||||
{"my_key": ""},
|
||||
]
|
||||
assert [c async for c in app.astream(None, config, stream_mode="values")] == [
|
||||
{"my_key": ""},
|
||||
{"my_key": "got here and there and parallel"},
|
||||
]
|
||||
assert [c async for c in app.astream(None, config, stream_mode="values")] == [
|
||||
{"my_key": "got here and there and parallel"},
|
||||
{"my_key": "got here and there and parallel and back again"},
|
||||
]
|
||||
|
||||
@@ -7007,20 +7061,13 @@ async def test_doubly_nested_graph_interrupts(
|
||||
c
|
||||
async for c in app.astream({"my_key": "my value"}, config, stream_mode="values")
|
||||
] == [
|
||||
{
|
||||
"my_key": "my value",
|
||||
},
|
||||
{
|
||||
"my_key": "hi my value",
|
||||
},
|
||||
{"my_key": "my value"},
|
||||
{"my_key": "hi my value"},
|
||||
]
|
||||
assert [c async for c in app.astream(None, config, stream_mode="values")] == [
|
||||
{
|
||||
"my_key": "hi my value here and there",
|
||||
},
|
||||
{
|
||||
"my_key": "hi my value here and there and back again",
|
||||
},
|
||||
{"my_key": "hi my value"},
|
||||
{"my_key": "hi my value here and there"},
|
||||
{"my_key": "hi my value here and there and back again"},
|
||||
]
|
||||
|
||||
|
||||
|
||||
Reference in New Issue
Block a user