mirror of
https://github.com/langchain-ai/langgraph.git
synced 2026-08-17 21:25:46 +02:00
Remove non-state Graph
- Not used in any examples/docs
This commit is contained in:
committed by
Sydney Runkle
parent
05f3904d09
commit
afb83d2201
@@ -95,7 +95,7 @@ my-app/
|
||||
|
||||
## Define Graphs
|
||||
|
||||
Implement your graphs! Graphs can be defined in a single file or multiple files. Make note of the variable names of each [CompiledGraph][langgraph.graph.graph.CompiledGraph] to be included in the LangGraph application. The variable names will be used later when creating the [LangGraph configuration file](../reference/cli.md#configuration-file).
|
||||
Implement your graphs! Graphs can be defined in a single file or multiple files. Make note of the variable names of each [CompiledStateGraph][langgraph.graph.state.CompiledStateGraph] to be included in the LangGraph application. The variable names will be used later when creating the [LangGraph configuration file](../reference/cli.md#configuration-file).
|
||||
|
||||
Example `agent.py` file, which shows how to import from other modules you define (code for the modules is not shown here, please see [this repository](https://github.com/langchain-ai/langgraph-example) to see their implementation):
|
||||
|
||||
|
||||
@@ -108,7 +108,7 @@ my-app/
|
||||
|
||||
## Define Graphs
|
||||
|
||||
Implement your graphs! Graphs can be defined in a single file or multiple files. Make note of the variable names of each [CompiledGraph][langgraph.graph.graph.CompiledGraph] to be included in the LangGraph application. The variable names will be used later when creating the [LangGraph configuration file](../reference/cli.md#configuration-file).
|
||||
Implement your graphs! Graphs can be defined in a single file or multiple files. Make note of the variable names of each [CompiledStateGraph][langgraph.graph.state.CompiledStateGraph] to be included in the LangGraph application. The variable names will be used later when creating the [LangGraph configuration file](../reference/cli.md#configuration-file).
|
||||
|
||||
Example `agent.py` file, which shows how to import from other modules you define (code for the modules is not shown here, please see [this repository](https://github.com/langchain-ai/langgraph-example-pyproject) to see their implementation):
|
||||
|
||||
|
||||
@@ -12,9 +12,9 @@
|
||||
"\n",
|
||||
"\n",
|
||||
"1. **Run the graph** with initial inputs using `invoke` or `stream` APIs.\n",
|
||||
"2. **Identify a checkpoint in an existing thread**: Use the [`get_state_history()`][langgraph.graph.graph.CompiledGraph.get_state_history] method to retrieve the execution history for a specific `thread_id` and locate the desired `checkpoint_id`. \n",
|
||||
"2. **Identify a checkpoint in an existing thread**: Use the [`get_state_history()`][langgraph.graph.state.CompiledStateGraph.get_state_history] method to retrieve the execution history for a specific `thread_id` and locate the desired `checkpoint_id`. \n",
|
||||
" Alternatively, set a [breakpoint](../../../concepts/breakpoints/) before the node(s) where you want execution to pause. You can then find the most recent checkpoint recorded up to that breakpoint.\n",
|
||||
"3. **(Optional) modify the graph state**: Use the [`update_state`][langgraph.graph.graph.CompiledGraph.update_state] method to modify the graph’s state at the checkpoint and resume execution from alternative state.\n",
|
||||
"3. **(Optional) modify the graph state**: Use the [`update_state`][langgraph.graph.state.CompiledStateGraph.update_state] method to modify the graph’s state at the checkpoint and resume execution from alternative state.\n",
|
||||
"4. **Resume execution from the checkpoint**: Use the `invoke` or `stream` APIs with an input of `None` and a configuration containing the appropriate `thread_id` and `checkpoint_id`.\n",
|
||||
"\n",
|
||||
"## Example\n",
|
||||
|
||||
@@ -36,41 +36,6 @@
|
||||
- aget_subgraphs
|
||||
- with_config
|
||||
|
||||
::: langgraph.graph.graph.Graph
|
||||
options:
|
||||
show_if_no_docstring: true
|
||||
show_root_heading: true
|
||||
show_root_full_path: false
|
||||
members:
|
||||
- add_node
|
||||
- add_edge
|
||||
- add_conditional_edges
|
||||
- compile
|
||||
|
||||
::: langgraph.graph.graph.CompiledGraph
|
||||
options:
|
||||
show_if_no_docstring: true
|
||||
show_root_heading: true
|
||||
show_root_full_path: false
|
||||
members:
|
||||
- stream
|
||||
- astream
|
||||
- invoke
|
||||
- ainvoke
|
||||
- get_state
|
||||
- aget_state
|
||||
- get_state_history
|
||||
- aget_state_history
|
||||
- update_state
|
||||
- aupdate_state
|
||||
- bulk_update_state
|
||||
- abulk_update_state
|
||||
- get_graph
|
||||
- aget_graph
|
||||
- get_subgraphs
|
||||
- aget_subgraphs
|
||||
- with_config
|
||||
|
||||
::: langgraph.graph.message
|
||||
options:
|
||||
members:
|
||||
|
||||
@@ -1,11 +1,10 @@
|
||||
from langgraph.graph.graph import END, START, Graph
|
||||
from langgraph.constants import END, START
|
||||
from langgraph.graph.message import MessageGraph, MessagesState, add_messages
|
||||
from langgraph.graph.state import StateGraph
|
||||
|
||||
__all__ = [
|
||||
"END",
|
||||
"START",
|
||||
"Graph",
|
||||
"StateGraph",
|
||||
"MessageGraph",
|
||||
"add_messages",
|
||||
|
||||
@@ -1,445 +0,0 @@
|
||||
import logging
|
||||
from collections import defaultdict
|
||||
from collections.abc import Awaitable, Hashable, Sequence
|
||||
from typing import (
|
||||
Any,
|
||||
Callable,
|
||||
NamedTuple,
|
||||
Optional,
|
||||
Union,
|
||||
cast,
|
||||
overload,
|
||||
)
|
||||
|
||||
from langchain_core.runnables import Runnable
|
||||
from typing_extensions import Self
|
||||
|
||||
from langgraph.cache.base import BaseCache
|
||||
from langgraph.channels.ephemeral_value import EphemeralValue
|
||||
from langgraph.constants import (
|
||||
EMPTY_SEQ,
|
||||
END,
|
||||
NS_END,
|
||||
NS_SEP,
|
||||
START,
|
||||
TAG_HIDDEN,
|
||||
Send,
|
||||
)
|
||||
from langgraph.graph.branch import Branch
|
||||
from langgraph.pregel import NodeBuilder, Pregel
|
||||
from langgraph.pregel.read import PregelNode
|
||||
from langgraph.pregel.write import ChannelWrite, ChannelWriteEntry
|
||||
from langgraph.store.base import BaseStore
|
||||
from langgraph.types import All, Checkpointer
|
||||
from langgraph.utils.runnable import RunnableLike, coerce_to_runnable
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class NodeSpec(NamedTuple):
|
||||
runnable: Runnable
|
||||
metadata: Optional[dict[str, Any]] = None
|
||||
ends: Optional[Union[tuple[str, ...], dict[str, str]]] = EMPTY_SEQ
|
||||
|
||||
|
||||
class Graph:
|
||||
def __init__(self) -> None:
|
||||
self.nodes: dict[str, NodeSpec] = {}
|
||||
self.edges = set[tuple[str, str]]()
|
||||
self.branches: defaultdict[str, dict[str, Branch]] = defaultdict(dict)
|
||||
self.support_multiple_edges = False
|
||||
self.compiled = False
|
||||
|
||||
@property
|
||||
def _all_edges(self) -> set[tuple[str, str]]:
|
||||
return self.edges
|
||||
|
||||
@overload
|
||||
def add_node(
|
||||
self,
|
||||
node: RunnableLike,
|
||||
*,
|
||||
metadata: Optional[dict[str, Any]] = None,
|
||||
) -> Self: ...
|
||||
|
||||
@overload
|
||||
def add_node(
|
||||
self,
|
||||
node: str,
|
||||
action: RunnableLike,
|
||||
*,
|
||||
metadata: Optional[dict[str, Any]] = None,
|
||||
) -> Self: ...
|
||||
|
||||
def add_node(
|
||||
self,
|
||||
node: Union[str, RunnableLike],
|
||||
action: Optional[RunnableLike] = None,
|
||||
*,
|
||||
metadata: Optional[dict[str, Any]] = None,
|
||||
) -> Self:
|
||||
"""Add a new node to the graph.
|
||||
|
||||
Args:
|
||||
node: The function or runnable this node will run.
|
||||
If a string is provided, it will be used as the node name, and action will be used as the function or runnable.
|
||||
action: The action associated with the node. (default: None)
|
||||
Will be used as the node function or runnable if `node` is a string (node name).
|
||||
metadata: The metadata associated with the node. (default: None)
|
||||
"""
|
||||
if isinstance(node, str):
|
||||
for character in (NS_SEP, NS_END):
|
||||
if character in node:
|
||||
raise ValueError(
|
||||
f"'{character}' is a reserved character and is not allowed in the node names."
|
||||
)
|
||||
|
||||
if self.compiled:
|
||||
logger.warning(
|
||||
"Adding a node to a graph that has already been compiled. This will "
|
||||
"not be reflected in the compiled graph."
|
||||
)
|
||||
if not isinstance(node, str):
|
||||
action = node
|
||||
node = getattr(action, "name", getattr(action, "__name__"))
|
||||
if node is None:
|
||||
raise ValueError(
|
||||
"Node name must be provided if action is not a function"
|
||||
)
|
||||
if action is None:
|
||||
raise RuntimeError(
|
||||
"Expected a function or Runnable action in add_node. Received None."
|
||||
)
|
||||
if node in self.nodes:
|
||||
raise ValueError(f"Node `{node}` already present.")
|
||||
if node == END or node == START:
|
||||
raise ValueError(f"Node `{node}` is reserved.")
|
||||
|
||||
self.nodes[cast(str, node)] = NodeSpec(
|
||||
coerce_to_runnable(action, name=cast(str, node), trace=False), metadata
|
||||
)
|
||||
return self
|
||||
|
||||
def add_edge(self, start_key: str, end_key: str) -> Self:
|
||||
"""Add a directed edge from the start node to the end node.
|
||||
|
||||
Args:
|
||||
start_key: The key of the start node of the edge.
|
||||
end_key: The key of the end node of the edge.
|
||||
"""
|
||||
if self.compiled:
|
||||
logger.warning(
|
||||
"Adding an edge to a graph that has already been compiled. This will "
|
||||
"not be reflected in the compiled graph."
|
||||
)
|
||||
if start_key == END:
|
||||
raise ValueError("END cannot be a start node")
|
||||
if end_key == START:
|
||||
raise ValueError("START cannot be an end node")
|
||||
|
||||
# run this validation only for non-StateGraph graphs
|
||||
if not hasattr(self, "channels") and start_key in set(
|
||||
start for start, _ in self.edges
|
||||
):
|
||||
raise ValueError(
|
||||
f"Already found path for node '{start_key}'.\n"
|
||||
"For multiple edges, use StateGraph with an Annotated state key."
|
||||
)
|
||||
|
||||
self.edges.add((start_key, end_key))
|
||||
return self
|
||||
|
||||
def add_conditional_edges(
|
||||
self,
|
||||
source: str,
|
||||
path: Union[
|
||||
Callable[..., Union[Hashable, list[Hashable]]],
|
||||
Callable[..., Awaitable[Union[Hashable, list[Hashable]]]],
|
||||
Runnable[Any, Union[Hashable, list[Hashable]]],
|
||||
],
|
||||
path_map: Optional[Union[dict[Hashable, str], list[str]]] = None,
|
||||
then: Optional[str] = None,
|
||||
) -> Self:
|
||||
"""Add a conditional edge from the starting node to any number of destination nodes.
|
||||
|
||||
Args:
|
||||
source: The starting node. This conditional edge will run when
|
||||
exiting this node.
|
||||
path: The callable that determines the next
|
||||
node or nodes. If not specifying `path_map` it should return one or
|
||||
more nodes. If it returns END, the graph will stop execution.
|
||||
path_map: Optional mapping of paths to node
|
||||
names. If omitted the paths returned by `path` should be node names.
|
||||
then: The name of a node to execute after the nodes
|
||||
selected by `path`.
|
||||
|
||||
Returns:
|
||||
Self: The instance of the graph, allowing for method chaining.
|
||||
|
||||
Note: Without typehints on the `path` function's return value (e.g., `-> Literal["foo", "__end__"]:`)
|
||||
or a path_map, the graph visualization assumes the edge could transition to any node in the graph.
|
||||
|
||||
""" # noqa: E501
|
||||
if self.compiled:
|
||||
logger.warning(
|
||||
"Adding an edge to a graph that has already been compiled. This will "
|
||||
"not be reflected in the compiled graph."
|
||||
)
|
||||
|
||||
# find a name for the condition
|
||||
path = coerce_to_runnable(path, name=None, trace=True)
|
||||
name = path.name or "condition"
|
||||
# validate the condition
|
||||
if name in self.branches[source]:
|
||||
raise ValueError(
|
||||
f"Branch with name `{path.name}` already exists for node `{source}`"
|
||||
)
|
||||
# save it
|
||||
self.branches[source][name] = Branch.from_path(path, path_map, then, False)
|
||||
return self
|
||||
|
||||
def set_entry_point(self, key: str) -> Self:
|
||||
"""Specifies the first node to be called in the graph.
|
||||
|
||||
Equivalent to calling `add_edge(START, key)`.
|
||||
|
||||
Parameters:
|
||||
key (str): The key of the node to set as the entry point.
|
||||
|
||||
Returns:
|
||||
Self: The instance of the graph, allowing for method chaining.
|
||||
"""
|
||||
return self.add_edge(START, key)
|
||||
|
||||
def set_conditional_entry_point(
|
||||
self,
|
||||
path: Union[
|
||||
Callable[..., Union[Hashable, list[Hashable]]],
|
||||
Callable[..., Awaitable[Union[Hashable, list[Hashable]]]],
|
||||
Runnable[Any, Union[Hashable, list[Hashable]]],
|
||||
],
|
||||
path_map: Optional[Union[dict[Hashable, str], list[str]]] = None,
|
||||
then: Optional[str] = None,
|
||||
) -> Self:
|
||||
"""Sets a conditional entry point in the graph.
|
||||
|
||||
Args:
|
||||
path: The callable that determines the next
|
||||
node or nodes. If not specifying `path_map` it should return one or
|
||||
more nodes. If it returns END, the graph will stop execution.
|
||||
path_map: Optional mapping of paths to node
|
||||
names. If omitted the paths returned by `path` should be node names.
|
||||
then: The name of a node to execute after the nodes
|
||||
selected by `path`.
|
||||
|
||||
Returns:
|
||||
Self: The instance of the graph, allowing for method chaining.
|
||||
"""
|
||||
return self.add_conditional_edges(START, path, path_map, then)
|
||||
|
||||
def set_finish_point(self, key: str) -> Self:
|
||||
"""Marks a node as a finish point of the graph.
|
||||
|
||||
If the graph reaches this node, it will cease execution.
|
||||
|
||||
Parameters:
|
||||
key (str): The key of the node to set as the finish point.
|
||||
|
||||
Returns:
|
||||
Self: The instance of the graph, allowing for method chaining.
|
||||
"""
|
||||
return self.add_edge(key, END)
|
||||
|
||||
def validate(self, interrupt: Optional[Sequence[str]] = None) -> Self:
|
||||
# assemble sources
|
||||
all_sources = {src for src, _ in self._all_edges}
|
||||
for start, branches in self.branches.items():
|
||||
all_sources.add(start)
|
||||
for cond, branch in branches.items():
|
||||
if branch.then is not None:
|
||||
if branch.ends is not None:
|
||||
for end in branch.ends.values():
|
||||
if end != END:
|
||||
all_sources.add(end)
|
||||
else:
|
||||
for node in self.nodes:
|
||||
if node != start and node != branch.then:
|
||||
all_sources.add(node)
|
||||
for name, spec in self.nodes.items():
|
||||
if spec.ends:
|
||||
all_sources.add(name)
|
||||
# validate sources
|
||||
for source in all_sources:
|
||||
if source not in self.nodes and source != START:
|
||||
raise ValueError(f"Found edge starting at unknown node '{source}'")
|
||||
|
||||
if START not in all_sources:
|
||||
raise ValueError(
|
||||
"Graph must have an entrypoint: add at least one edge from START to another node"
|
||||
)
|
||||
|
||||
# assemble targets
|
||||
all_targets = {end for _, end in self._all_edges}
|
||||
for start, branches in self.branches.items():
|
||||
for cond, branch in branches.items():
|
||||
if branch.then is not None:
|
||||
all_targets.add(branch.then)
|
||||
if branch.ends is not None:
|
||||
for end in branch.ends.values():
|
||||
if end not in self.nodes and end != END:
|
||||
raise ValueError(
|
||||
f"At '{start}' node, '{cond}' branch found unknown target '{end}'"
|
||||
)
|
||||
all_targets.add(end)
|
||||
else:
|
||||
all_targets.add(END)
|
||||
for node in self.nodes:
|
||||
if node != start and node != branch.then:
|
||||
all_targets.add(node)
|
||||
for name, spec in self.nodes.items():
|
||||
if spec.ends:
|
||||
all_targets.update(spec.ends)
|
||||
for target in all_targets:
|
||||
if target not in self.nodes and target != END:
|
||||
raise ValueError(f"Found edge ending at unknown node `{target}`")
|
||||
# validate interrupts
|
||||
if interrupt:
|
||||
for node in interrupt:
|
||||
if node not in self.nodes:
|
||||
raise ValueError(f"Interrupt node `{node}` not found")
|
||||
|
||||
self.compiled = True
|
||||
return self
|
||||
|
||||
def compile(
|
||||
self,
|
||||
checkpointer: Checkpointer = None,
|
||||
interrupt_before: Optional[Union[All, list[str]]] = None,
|
||||
interrupt_after: Optional[Union[All, list[str]]] = None,
|
||||
debug: bool = False,
|
||||
name: Optional[str] = None,
|
||||
*,
|
||||
cache: Optional[BaseCache] = None,
|
||||
store: Optional[BaseStore] = None,
|
||||
) -> "CompiledGraph":
|
||||
"""Compiles the graph into a `CompiledGraph` object.
|
||||
|
||||
The compiled graph implements the `Runnable` interface and can be invoked,
|
||||
streamed, batched, and run asynchronously.
|
||||
|
||||
Args:
|
||||
checkpointer: A checkpoint saver object or flag.
|
||||
If provided, this Checkpointer serves as a fully versioned "short-term memory" for the graph,
|
||||
allowing it to be paused, resumed, and replayed from any point.
|
||||
If None, it may inherit the parent graph's checkpointer when used as a subgraph.
|
||||
If False, it will not use or inherit any checkpointer.
|
||||
interrupt_before: An optional list of node names to interrupt before.
|
||||
interrupt_after: An optional list of node names to interrupt after.
|
||||
debug: A flag indicating whether to enable debug mode.
|
||||
name: The name to use for the compiled graph.
|
||||
|
||||
Returns:
|
||||
CompiledGraph: The compiled graph.
|
||||
"""
|
||||
# assign default values
|
||||
interrupt_before = interrupt_before or []
|
||||
interrupt_after = interrupt_after or []
|
||||
|
||||
# validate the graph
|
||||
self.validate(
|
||||
interrupt=(
|
||||
(interrupt_before if interrupt_before != "*" else []) + interrupt_after
|
||||
if interrupt_after != "*"
|
||||
else []
|
||||
)
|
||||
)
|
||||
|
||||
# create empty compiled graph
|
||||
compiled = CompiledGraph(
|
||||
builder=self,
|
||||
nodes={},
|
||||
channels={START: EphemeralValue(Any), END: EphemeralValue(Any)},
|
||||
input_channels=START,
|
||||
output_channels=END,
|
||||
stream_mode="values",
|
||||
stream_channels=[],
|
||||
checkpointer=checkpointer,
|
||||
interrupt_before_nodes=interrupt_before,
|
||||
interrupt_after_nodes=interrupt_after,
|
||||
auto_validate=False,
|
||||
debug=debug,
|
||||
name=name or "LangGraph",
|
||||
cache=cache,
|
||||
store=store,
|
||||
)
|
||||
|
||||
# attach nodes, edges, and branches
|
||||
for key, node in self.nodes.items():
|
||||
compiled.attach_node(key, node)
|
||||
|
||||
for start, end in self.edges:
|
||||
compiled.attach_edge(start, end)
|
||||
|
||||
for start, branches in self.branches.items():
|
||||
for name, branch in branches.items():
|
||||
compiled.attach_branch(start, name, branch)
|
||||
|
||||
# validate the compiled graph
|
||||
return compiled.validate()
|
||||
|
||||
|
||||
class CompiledGraph(Pregel):
|
||||
builder: Graph
|
||||
|
||||
def __init__(self, *, builder: Graph, **kwargs: Any) -> None:
|
||||
super().__init__(**kwargs)
|
||||
self.builder = builder
|
||||
|
||||
def attach_node(self, key: str, node: NodeSpec) -> None:
|
||||
self.channels[key] = EphemeralValue(Any)
|
||||
self.nodes[key] = (
|
||||
PregelNode(channels=[], triggers=[], metadata=node.metadata)
|
||||
| node.runnable
|
||||
| ChannelWrite([ChannelWriteEntry(key)])
|
||||
)
|
||||
cast(list[str], self.stream_channels).append(key)
|
||||
|
||||
def attach_edge(self, start: str, end: str) -> None:
|
||||
if end == END:
|
||||
# publish to end channel
|
||||
self.nodes[start].writers.append(ChannelWrite([ChannelWriteEntry(END)]))
|
||||
else:
|
||||
# subscribe to start channel
|
||||
self.nodes[end].triggers.append(start)
|
||||
cast(list[str], self.nodes[end].channels).append(start)
|
||||
|
||||
def attach_branch(self, start: str, name: str, branch: Branch) -> None:
|
||||
def get_writes(
|
||||
packets: Sequence[Union[str, Send]], static: bool = False
|
||||
) -> Sequence[Union[ChannelWriteEntry, Send]]:
|
||||
return [
|
||||
(
|
||||
ChannelWriteEntry(f"branch:{start}:{name}:{p}" if p != END else END)
|
||||
if not isinstance(p, Send)
|
||||
else p
|
||||
)
|
||||
for p in packets
|
||||
]
|
||||
|
||||
# add hidden start node
|
||||
if start == START and start not in self.nodes:
|
||||
self.nodes[start] = (
|
||||
NodeBuilder().subscribe_only(START).meta(TAG_HIDDEN).build()
|
||||
)
|
||||
|
||||
# attach branch writer
|
||||
self.nodes[start] |= branch.run(get_writes)
|
||||
|
||||
# attach branch readers
|
||||
ends = branch.ends.values() if branch.ends else [node for node in self.nodes]
|
||||
for end in ends:
|
||||
if end != END:
|
||||
channel_name = f"branch:{start}:{name}:{end}"
|
||||
self.channels[channel_name] = EphemeralValue(Any)
|
||||
self.nodes[end].triggers.append(channel_name)
|
||||
cast(list[str], self.nodes[end].channels).append(channel_name)
|
||||
@@ -43,10 +43,12 @@ from langgraph.channels.named_barrier_value import (
|
||||
from langgraph.checkpoint.base import Checkpoint
|
||||
from langgraph.constants import (
|
||||
EMPTY_SEQ,
|
||||
END,
|
||||
INTERRUPT,
|
||||
MISSING,
|
||||
NS_END,
|
||||
NS_SEP,
|
||||
START,
|
||||
TAG_HIDDEN,
|
||||
TASKS,
|
||||
)
|
||||
@@ -57,17 +59,11 @@ from langgraph.errors import (
|
||||
create_error_message,
|
||||
)
|
||||
from langgraph.graph.branch import Branch
|
||||
from langgraph.graph.graph import (
|
||||
END,
|
||||
START,
|
||||
CompiledGraph,
|
||||
Graph,
|
||||
Send,
|
||||
)
|
||||
from langgraph.managed.base import (
|
||||
ManagedValueSpec,
|
||||
is_managed_value,
|
||||
)
|
||||
from langgraph.pregel import Pregel
|
||||
from langgraph.pregel.read import ChannelRead, PregelNode
|
||||
from langgraph.pregel.write import (
|
||||
ChannelWrite,
|
||||
@@ -75,7 +71,7 @@ from langgraph.pregel.write import (
|
||||
ChannelWriteTupleEntry,
|
||||
)
|
||||
from langgraph.store.base import BaseStore
|
||||
from langgraph.types import All, CachePolicy, Checkpointer, Command, RetryPolicy
|
||||
from langgraph.types import All, CachePolicy, Checkpointer, Command, RetryPolicy, Send
|
||||
from langgraph.utils.fields import get_field_default, get_update_as_tuples
|
||||
from langgraph.utils.pydantic import create_model
|
||||
from langgraph.utils.runnable import RunnableLike, coerce_to_runnable
|
||||
@@ -114,7 +110,7 @@ class StateNodeSpec(NamedTuple):
|
||||
defer: bool = False
|
||||
|
||||
|
||||
class StateGraph(Graph):
|
||||
class StateGraph:
|
||||
"""A graph whose nodes communicate by reading and writing to a shared state.
|
||||
The signature of each node is State -> Partial<State>.
|
||||
|
||||
@@ -166,7 +162,9 @@ class StateGraph(Graph):
|
||||
```
|
||||
"""
|
||||
|
||||
nodes: dict[str, StateNodeSpec] # type: ignore[assignment]
|
||||
edges: set[tuple[str, str]]
|
||||
nodes: dict[str, StateNodeSpec]
|
||||
branches: defaultdict[str, dict[str, Branch]]
|
||||
channels: dict[str, BaseChannel]
|
||||
managed: dict[str, ManagedValueSpec]
|
||||
schemas: dict[type[Any], dict[str, Union[BaseChannel, ManagedValueSpec]]]
|
||||
@@ -179,7 +177,6 @@ class StateGraph(Graph):
|
||||
input: Optional[type[Any]] = None,
|
||||
output: Optional[type[Any]] = None,
|
||||
) -> None:
|
||||
super().__init__()
|
||||
if state_schema is None:
|
||||
if input is None or output is None:
|
||||
raise ValueError("Must provide state_schema or input and output")
|
||||
@@ -195,10 +192,14 @@ class StateGraph(Graph):
|
||||
input = state_schema
|
||||
if output is None:
|
||||
output = state_schema
|
||||
self.nodes = {}
|
||||
self.edges = set[tuple[str, str]]()
|
||||
self.branches = defaultdict(dict)
|
||||
self.support_multiple_edges = False
|
||||
self.compiled = False
|
||||
self.schemas = {}
|
||||
self.channels = {}
|
||||
self.managed = {}
|
||||
self.type_hints: dict[type[Any], dict[str, Any]] = {}
|
||||
self.schema = state_schema
|
||||
self.input = input
|
||||
self.output = output
|
||||
@@ -226,7 +227,6 @@ class StateGraph(Graph):
|
||||
" Managed channels are not permitted in Input/Output schema."
|
||||
)
|
||||
self.schemas[schema] = {**channels, **managed}
|
||||
self.type_hints[schema] = type_hints
|
||||
for key, channel in channels.items():
|
||||
if key in self.channels:
|
||||
if self.channels[key] != channel:
|
||||
@@ -373,7 +373,7 @@ class StateGraph(Graph):
|
||||
raise ValueError(f"Node `{node}` is reserved.")
|
||||
|
||||
for character in (NS_SEP, NS_END):
|
||||
if character in cast(str, node):
|
||||
if character in node:
|
||||
raise ValueError(
|
||||
f"'{character}' is a reserved character and is not allowed in the node names."
|
||||
)
|
||||
@@ -428,8 +428,8 @@ class StateGraph(Graph):
|
||||
|
||||
if input is not None:
|
||||
self._add_schema(input)
|
||||
self.nodes[cast(str, node)] = StateNodeSpec(
|
||||
coerce_to_runnable(action, name=cast(str, node), trace=False),
|
||||
self.nodes[node] = StateNodeSpec(
|
||||
coerce_to_runnable(action, name=node, trace=False),
|
||||
metadata,
|
||||
input=input or self.schema,
|
||||
retry_policy=retry,
|
||||
@@ -456,14 +456,30 @@ class StateGraph(Graph):
|
||||
Returns:
|
||||
Self: The instance of the state graph, allowing for method chaining.
|
||||
"""
|
||||
if isinstance(start_key, str):
|
||||
return super().add_edge(start_key, end_key)
|
||||
|
||||
if self.compiled:
|
||||
logger.warning(
|
||||
"Adding an edge to a graph that has already been compiled. This will "
|
||||
"not be reflected in the compiled graph."
|
||||
)
|
||||
|
||||
if isinstance(start_key, str):
|
||||
if start_key == END:
|
||||
raise ValueError("END cannot be a start node")
|
||||
if end_key == START:
|
||||
raise ValueError("START cannot be an end node")
|
||||
|
||||
# run this validation only for non-StateGraph graphs
|
||||
if not hasattr(self, "channels") and start_key in set(
|
||||
start for start, _ in self.edges
|
||||
):
|
||||
raise ValueError(
|
||||
f"Already found path for node '{start_key}'.\n"
|
||||
"For multiple edges, use StateGraph with an Annotated state key."
|
||||
)
|
||||
|
||||
self.edges.add((start_key, end_key))
|
||||
return self
|
||||
|
||||
for start in start_key:
|
||||
if start == END:
|
||||
raise ValueError("END cannot be a start node")
|
||||
@@ -570,6 +586,119 @@ class StateGraph(Graph):
|
||||
|
||||
return self
|
||||
|
||||
def set_entry_point(self, key: str) -> Self:
|
||||
"""Specifies the first node to be called in the graph.
|
||||
|
||||
Equivalent to calling `add_edge(START, key)`.
|
||||
|
||||
Parameters:
|
||||
key (str): The key of the node to set as the entry point.
|
||||
|
||||
Returns:
|
||||
Self: The instance of the graph, allowing for method chaining.
|
||||
"""
|
||||
return self.add_edge(START, key)
|
||||
|
||||
def set_conditional_entry_point(
|
||||
self,
|
||||
path: Union[
|
||||
Callable[..., Union[Hashable, list[Hashable]]],
|
||||
Callable[..., Awaitable[Union[Hashable, list[Hashable]]]],
|
||||
Runnable[Any, Union[Hashable, list[Hashable]]],
|
||||
],
|
||||
path_map: Optional[Union[dict[Hashable, str], list[str]]] = None,
|
||||
then: Optional[str] = None,
|
||||
) -> Self:
|
||||
"""Sets a conditional entry point in the graph.
|
||||
|
||||
Args:
|
||||
path: The callable that determines the next
|
||||
node or nodes. If not specifying `path_map` it should return one or
|
||||
more nodes. If it returns END, the graph will stop execution.
|
||||
path_map: Optional mapping of paths to node
|
||||
names. If omitted the paths returned by `path` should be node names.
|
||||
then: The name of a node to execute after the nodes
|
||||
selected by `path`.
|
||||
|
||||
Returns:
|
||||
Self: The instance of the graph, allowing for method chaining.
|
||||
"""
|
||||
return self.add_conditional_edges(START, path, path_map, then)
|
||||
|
||||
def set_finish_point(self, key: str) -> Self:
|
||||
"""Marks a node as a finish point of the graph.
|
||||
|
||||
If the graph reaches this node, it will cease execution.
|
||||
|
||||
Parameters:
|
||||
key (str): The key of the node to set as the finish point.
|
||||
|
||||
Returns:
|
||||
Self: The instance of the graph, allowing for method chaining.
|
||||
"""
|
||||
return self.add_edge(key, END)
|
||||
|
||||
def validate(self, interrupt: Optional[Sequence[str]] = None) -> Self:
|
||||
# assemble sources
|
||||
all_sources = {src for src, _ in self._all_edges}
|
||||
for start, branches in self.branches.items():
|
||||
all_sources.add(start)
|
||||
for cond, branch in branches.items():
|
||||
if branch.then is not None:
|
||||
if branch.ends is not None:
|
||||
for end in branch.ends.values():
|
||||
if end != END:
|
||||
all_sources.add(end)
|
||||
else:
|
||||
for node in self.nodes:
|
||||
if node != start and node != branch.then:
|
||||
all_sources.add(node)
|
||||
for name, spec in self.nodes.items():
|
||||
if spec.ends:
|
||||
all_sources.add(name)
|
||||
# validate sources
|
||||
for source in all_sources:
|
||||
if source not in self.nodes and source != START:
|
||||
raise ValueError(f"Found edge starting at unknown node '{source}'")
|
||||
|
||||
if START not in all_sources:
|
||||
raise ValueError(
|
||||
"Graph must have an entrypoint: add at least one edge from START to another node"
|
||||
)
|
||||
|
||||
# assemble targets
|
||||
all_targets = {end for _, end in self._all_edges}
|
||||
for start, branches in self.branches.items():
|
||||
for cond, branch in branches.items():
|
||||
if branch.then is not None:
|
||||
all_targets.add(branch.then)
|
||||
if branch.ends is not None:
|
||||
for end in branch.ends.values():
|
||||
if end not in self.nodes and end != END:
|
||||
raise ValueError(
|
||||
f"At '{start}' node, '{cond}' branch found unknown target '{end}'"
|
||||
)
|
||||
all_targets.add(end)
|
||||
else:
|
||||
all_targets.add(END)
|
||||
for node in self.nodes:
|
||||
if node != start and node != branch.then:
|
||||
all_targets.add(node)
|
||||
for name, spec in self.nodes.items():
|
||||
if spec.ends:
|
||||
all_targets.update(spec.ends)
|
||||
for target in all_targets:
|
||||
if target not in self.nodes and target != END:
|
||||
raise ValueError(f"Found edge ending at unknown node `{target}`")
|
||||
# validate interrupts
|
||||
if interrupt:
|
||||
for node in interrupt:
|
||||
if node not in self.nodes:
|
||||
raise ValueError(f"Interrupt node `{node}` not found")
|
||||
|
||||
self.compiled = True
|
||||
return self
|
||||
|
||||
def compile(
|
||||
self,
|
||||
checkpointer: Checkpointer = None,
|
||||
@@ -680,17 +809,19 @@ class StateGraph(Graph):
|
||||
return compiled.validate()
|
||||
|
||||
|
||||
class CompiledStateGraph(CompiledGraph):
|
||||
class CompiledStateGraph(Pregel):
|
||||
builder: StateGraph
|
||||
schema_to_mapper: dict[type[Any], Optional[Callable[[Any], Any]]]
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
builder: StateGraph,
|
||||
schema_to_mapper: dict[type[Any], Optional[Callable[[Any], Any]]],
|
||||
**kwargs: Any,
|
||||
) -> None:
|
||||
super().__init__(**kwargs)
|
||||
self.builder = builder
|
||||
self.schema_to_mapper = schema_to_mapper
|
||||
|
||||
def get_input_schema(
|
||||
@@ -794,7 +925,6 @@ class CompiledStateGraph(CompiledGraph):
|
||||
mapper = _pick_mapper(
|
||||
list(input_values),
|
||||
input_schema,
|
||||
self.builder.type_hints[input_schema],
|
||||
)
|
||||
self.schema_to_mapper[input_schema] = mapper
|
||||
|
||||
@@ -890,7 +1020,7 @@ class CompiledStateGraph(CompiledGraph):
|
||||
if schema in self.schema_to_mapper:
|
||||
mapper = self.schema_to_mapper[schema]
|
||||
else:
|
||||
mapper = _pick_mapper(channels, schema, self.builder.type_hints[schema])
|
||||
mapper = _pick_mapper(channels, schema)
|
||||
self.schema_to_mapper[schema] = mapper
|
||||
# create reader
|
||||
reader: Optional[Callable[[RunnableConfig], Any]] = partial(
|
||||
@@ -1031,7 +1161,7 @@ class CompiledStateGraph(CompiledGraph):
|
||||
|
||||
|
||||
def _pick_mapper(
|
||||
state_keys: Sequence[str], schema: type[Any], type_hints: Optional[dict[str, Any]]
|
||||
state_keys: Sequence[str], schema: type[Any]
|
||||
) -> Optional[Callable[[Any], Any]]:
|
||||
if state_keys == ["__root__"]:
|
||||
return None
|
||||
|
||||
@@ -46,7 +46,7 @@ def get_fields(
|
||||
return model.model_fields
|
||||
|
||||
if hasattr(model, "__fields__"):
|
||||
return model.__fields__ # type: ignore[return-value]
|
||||
return model.__fields__
|
||||
msg = f"Expected a Pydantic model. Got {type(model)}"
|
||||
raise TypeError(msg)
|
||||
|
||||
|
||||
@@ -18,7 +18,6 @@ from langgraph.checkpoint.memory import InMemorySaver
|
||||
from langgraph.constants import END, PULL, PUSH, START
|
||||
from langgraph.errors import NodeInterrupt
|
||||
from langgraph.graph import StateGraph
|
||||
from langgraph.graph.graph import Graph
|
||||
from langgraph.graph.message import MessageGraph, MessagesState, add_messages
|
||||
from langgraph.prebuilt.chat_agent_executor import create_react_agent
|
||||
from langgraph.prebuilt.tool_node import ToolNode
|
||||
@@ -496,799 +495,6 @@ def test_fork_always_re_runs_nodes(
|
||||
]
|
||||
|
||||
|
||||
def test_conditional_graph(
|
||||
snapshot: SnapshotAssertion, sync_checkpointer: BaseCheckpointSaver
|
||||
) -> None:
|
||||
from langchain_core.language_models.fake import FakeStreamingListLLM
|
||||
from langchain_core.prompts import PromptTemplate
|
||||
from langchain_core.runnables import RunnablePassthrough
|
||||
from langchain_core.tools import tool
|
||||
|
||||
# Assemble the tools
|
||||
@tool()
|
||||
def search_api(query: str) -> str:
|
||||
"""Searches the API for the query."""
|
||||
return f"result for {query}"
|
||||
|
||||
tools = [search_api]
|
||||
|
||||
# Construct the agent
|
||||
prompt = PromptTemplate.from_template("Hello!")
|
||||
|
||||
llm = FakeStreamingListLLM(
|
||||
responses=[
|
||||
"tool:search_api:query",
|
||||
"tool:search_api:another",
|
||||
"finish:answer",
|
||||
]
|
||||
)
|
||||
|
||||
def agent_parser(input: str) -> Union[AgentAction, AgentFinish]:
|
||||
if input.startswith("finish"):
|
||||
_, answer = input.split(":")
|
||||
return AgentFinish(return_values={"answer": answer}, log=input)
|
||||
else:
|
||||
_, tool_name, tool_input = input.split(":")
|
||||
return AgentAction(tool=tool_name, tool_input=tool_input, log=input)
|
||||
|
||||
agent = RunnablePassthrough.assign(agent_outcome=prompt | llm | agent_parser)
|
||||
|
||||
# Define tool execution logic
|
||||
def execute_tools(data: dict) -> dict:
|
||||
data = data.copy()
|
||||
agent_action: AgentAction = data.pop("agent_outcome")
|
||||
observation = {t.name: t for t in tools}[agent_action.tool].invoke(
|
||||
agent_action.tool_input
|
||||
)
|
||||
if data.get("intermediate_steps") is None:
|
||||
data["intermediate_steps"] = []
|
||||
else:
|
||||
data["intermediate_steps"] = data["intermediate_steps"].copy()
|
||||
data["intermediate_steps"].append([agent_action, observation])
|
||||
return data
|
||||
|
||||
# Define decision-making logic
|
||||
def should_continue(data: dict) -> str:
|
||||
# Logic to decide whether to continue in the loop or exit
|
||||
if isinstance(data["agent_outcome"], AgentFinish):
|
||||
return "exit"
|
||||
else:
|
||||
return "continue"
|
||||
|
||||
# Define a new graph
|
||||
workflow = Graph()
|
||||
|
||||
workflow.add_node("agent", agent)
|
||||
workflow.add_node(
|
||||
"tools",
|
||||
execute_tools,
|
||||
metadata={"parents": {}, "version": 2, "variant": "b"},
|
||||
)
|
||||
|
||||
workflow.set_entry_point("agent")
|
||||
|
||||
workflow.add_conditional_edges(
|
||||
"agent", should_continue, {"continue": "tools", "exit": END}
|
||||
)
|
||||
|
||||
workflow.add_edge("tools", "agent")
|
||||
|
||||
app = workflow.compile()
|
||||
|
||||
if isinstance(sync_checkpointer, InMemorySaver):
|
||||
assert json.dumps(app.get_graph().to_json(), indent=2) == snapshot
|
||||
assert app.get_graph().draw_mermaid(with_styles=False) == snapshot
|
||||
assert app.get_graph().draw_mermaid() == snapshot
|
||||
|
||||
assert app.invoke({"input": "what is weather in sf"}) == {
|
||||
"input": "what is weather in sf",
|
||||
"intermediate_steps": [
|
||||
[
|
||||
AgentAction(
|
||||
tool="search_api",
|
||||
tool_input="query",
|
||||
log="tool:search_api:query",
|
||||
),
|
||||
"result for query",
|
||||
],
|
||||
[
|
||||
AgentAction(
|
||||
tool="search_api",
|
||||
tool_input="another",
|
||||
log="tool:search_api:another",
|
||||
),
|
||||
"result for another",
|
||||
],
|
||||
],
|
||||
"agent_outcome": AgentFinish(
|
||||
return_values={"answer": "answer"}, log="finish:answer"
|
||||
),
|
||||
}
|
||||
|
||||
assert [c for c in app.stream({"input": "what is weather in sf"})] == [
|
||||
{
|
||||
"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",
|
||||
"intermediate_steps": [
|
||||
[
|
||||
AgentAction(
|
||||
tool="search_api",
|
||||
tool_input="query",
|
||||
log="tool:search_api:query",
|
||||
),
|
||||
"result for query",
|
||||
]
|
||||
],
|
||||
}
|
||||
},
|
||||
{
|
||||
"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",
|
||||
"intermediate_steps": [
|
||||
[
|
||||
AgentAction(
|
||||
tool="search_api",
|
||||
tool_input="query",
|
||||
log="tool:search_api:query",
|
||||
),
|
||||
"result for query",
|
||||
],
|
||||
[
|
||||
AgentAction(
|
||||
tool="search_api",
|
||||
tool_input="another",
|
||||
log="tool:search_api:another",
|
||||
),
|
||||
"result for another",
|
||||
],
|
||||
],
|
||||
}
|
||||
},
|
||||
{
|
||||
"agent": {
|
||||
"input": "what is weather in sf",
|
||||
"intermediate_steps": [
|
||||
[
|
||||
AgentAction(
|
||||
tool="search_api",
|
||||
tool_input="query",
|
||||
log="tool:search_api:query",
|
||||
),
|
||||
"result for query",
|
||||
],
|
||||
[
|
||||
AgentAction(
|
||||
tool="search_api",
|
||||
tool_input="another",
|
||||
log="tool:search_api:another",
|
||||
),
|
||||
"result for another",
|
||||
],
|
||||
],
|
||||
"agent_outcome": AgentFinish(
|
||||
return_values={"answer": "answer"}, log="finish:answer"
|
||||
),
|
||||
}
|
||||
},
|
||||
]
|
||||
|
||||
# test state get/update methods with interrupt_after
|
||||
|
||||
app_w_interrupt = workflow.compile(
|
||||
checkpointer=sync_checkpointer,
|
||||
interrupt_after=["agent"],
|
||||
)
|
||||
config = {"configurable": {"thread_id": "1"}}
|
||||
|
||||
assert [
|
||||
c for c in app_w_interrupt.stream({"input": "what is weather in sf"}, config)
|
||||
] == [
|
||||
{
|
||||
"agent": {
|
||||
"input": "what is weather in sf",
|
||||
"agent_outcome": AgentAction(
|
||||
tool="search_api", tool_input="query", log="tool:search_api:query"
|
||||
),
|
||||
}
|
||||
}
|
||||
]
|
||||
|
||||
assert app_w_interrupt.get_state(config) == StateSnapshot(
|
||||
values={
|
||||
"agent": {
|
||||
"input": "what is weather in sf",
|
||||
"agent_outcome": AgentAction(
|
||||
tool="search_api", tool_input="query", log="tool:search_api:query"
|
||||
),
|
||||
},
|
||||
},
|
||||
tasks=(PregelTask(AnyStr(), "tools", (PULL, "tools")),),
|
||||
next=("tools",),
|
||||
created_at=AnyStr(),
|
||||
config={
|
||||
"configurable": {
|
||||
"thread_id": "1",
|
||||
"checkpoint_ns": "",
|
||||
"checkpoint_id": AnyStr(),
|
||||
}
|
||||
},
|
||||
metadata={
|
||||
"parents": {},
|
||||
"source": "loop",
|
||||
"step": 0,
|
||||
"thread_id": "1",
|
||||
},
|
||||
parent_config=None,
|
||||
interrupts=(),
|
||||
)
|
||||
assert (
|
||||
app_w_interrupt.checkpointer.get_tuple(config).config["configurable"][
|
||||
"checkpoint_id"
|
||||
]
|
||||
is not None
|
||||
)
|
||||
|
||||
app_w_interrupt.update_state(
|
||||
config,
|
||||
{
|
||||
"agent_outcome": AgentAction(
|
||||
tool="search_api",
|
||||
tool_input="query",
|
||||
log="tool:search_api:a different query",
|
||||
),
|
||||
"input": "what is weather in sf",
|
||||
},
|
||||
)
|
||||
|
||||
assert app_w_interrupt.get_state(config) == StateSnapshot(
|
||||
values={
|
||||
"agent": {
|
||||
"agent_outcome": AgentAction(
|
||||
tool="search_api",
|
||||
tool_input="query",
|
||||
log="tool:search_api:a different query",
|
||||
),
|
||||
"input": "what is weather in sf",
|
||||
},
|
||||
},
|
||||
tasks=(PregelTask(AnyStr(), "tools", (PULL, "tools")),),
|
||||
next=("tools",),
|
||||
config={
|
||||
"configurable": {
|
||||
"thread_id": "1",
|
||||
"checkpoint_ns": "",
|
||||
"checkpoint_id": AnyStr(),
|
||||
}
|
||||
},
|
||||
created_at=AnyStr(),
|
||||
metadata={
|
||||
"parents": {},
|
||||
"source": "update",
|
||||
"step": 1,
|
||||
"thread_id": "1",
|
||||
},
|
||||
parent_config=(
|
||||
list(app_w_interrupt.checkpointer.list(config, limit=2))[-1].config
|
||||
),
|
||||
interrupts=(),
|
||||
)
|
||||
|
||||
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",
|
||||
"intermediate_steps": [
|
||||
[
|
||||
AgentAction(
|
||||
tool="search_api",
|
||||
tool_input="query",
|
||||
log="tool:search_api:a different query",
|
||||
),
|
||||
"result for query",
|
||||
]
|
||||
],
|
||||
}
|
||||
},
|
||||
{
|
||||
"agent": {
|
||||
"input": "what is weather in sf",
|
||||
"intermediate_steps": [
|
||||
[
|
||||
AgentAction(
|
||||
tool="search_api",
|
||||
tool_input="query",
|
||||
log="tool:search_api:a different query",
|
||||
),
|
||||
"result for query",
|
||||
]
|
||||
],
|
||||
"agent_outcome": AgentAction(
|
||||
tool="search_api",
|
||||
tool_input="another",
|
||||
log="tool:search_api:another",
|
||||
),
|
||||
}
|
||||
},
|
||||
]
|
||||
|
||||
app_w_interrupt.update_state(
|
||||
config,
|
||||
{
|
||||
"input": "what is weather in sf",
|
||||
"intermediate_steps": [
|
||||
[
|
||||
AgentAction(
|
||||
tool="search_api",
|
||||
tool_input="query",
|
||||
log="tool:search_api:a different query",
|
||||
),
|
||||
"result for query",
|
||||
]
|
||||
],
|
||||
"agent_outcome": AgentFinish(
|
||||
return_values={"answer": "a really nice answer"},
|
||||
log="finish:a really nice answer",
|
||||
),
|
||||
},
|
||||
)
|
||||
|
||||
assert app_w_interrupt.get_state(config) == StateSnapshot(
|
||||
values={
|
||||
"agent": {
|
||||
"input": "what is weather in sf",
|
||||
"intermediate_steps": [
|
||||
[
|
||||
AgentAction(
|
||||
tool="search_api",
|
||||
tool_input="query",
|
||||
log="tool:search_api:a different query",
|
||||
),
|
||||
"result for query",
|
||||
]
|
||||
],
|
||||
"agent_outcome": AgentFinish(
|
||||
return_values={"answer": "a really nice answer"},
|
||||
log="finish:a really nice answer",
|
||||
),
|
||||
},
|
||||
},
|
||||
tasks=(),
|
||||
next=(),
|
||||
config={
|
||||
"configurable": {
|
||||
"thread_id": "1",
|
||||
"checkpoint_ns": "",
|
||||
"checkpoint_id": AnyStr(),
|
||||
}
|
||||
},
|
||||
created_at=AnyStr(),
|
||||
metadata={
|
||||
"parents": {},
|
||||
"source": "update",
|
||||
"step": 4,
|
||||
"thread_id": "1",
|
||||
},
|
||||
parent_config=(
|
||||
list(app_w_interrupt.checkpointer.list(config, limit=2))[-1].config
|
||||
),
|
||||
interrupts=(),
|
||||
)
|
||||
|
||||
# test state get/update methods with interrupt_before
|
||||
|
||||
app_w_interrupt = workflow.compile(
|
||||
checkpointer=sync_checkpointer,
|
||||
interrupt_before=["tools"],
|
||||
)
|
||||
config = {"configurable": {"thread_id": "2"}}
|
||||
llm.i = 0 # reset the llm
|
||||
|
||||
assert [
|
||||
c for c in app_w_interrupt.stream({"input": "what is weather in sf"}, config)
|
||||
] == [
|
||||
{
|
||||
"agent": {
|
||||
"input": "what is weather in sf",
|
||||
"agent_outcome": AgentAction(
|
||||
tool="search_api", tool_input="query", log="tool:search_api:query"
|
||||
),
|
||||
}
|
||||
}
|
||||
]
|
||||
|
||||
assert app_w_interrupt.get_state(config) == StateSnapshot(
|
||||
values={
|
||||
"agent": {
|
||||
"input": "what is weather in sf",
|
||||
"agent_outcome": AgentAction(
|
||||
tool="search_api", tool_input="query", log="tool:search_api:query"
|
||||
),
|
||||
},
|
||||
},
|
||||
tasks=(PregelTask(AnyStr(), "tools", (PULL, "tools")),),
|
||||
next=("tools",),
|
||||
config={
|
||||
"configurable": {
|
||||
"thread_id": "2",
|
||||
"checkpoint_ns": "",
|
||||
"checkpoint_id": AnyStr(),
|
||||
}
|
||||
},
|
||||
created_at=AnyStr(),
|
||||
metadata={
|
||||
"parents": {},
|
||||
"source": "loop",
|
||||
"step": 0,
|
||||
"thread_id": "2",
|
||||
},
|
||||
parent_config=None,
|
||||
interrupts=(),
|
||||
)
|
||||
|
||||
app_w_interrupt.update_state(
|
||||
config,
|
||||
{
|
||||
"agent_outcome": AgentAction(
|
||||
tool="search_api",
|
||||
tool_input="query",
|
||||
log="tool:search_api:a different query",
|
||||
),
|
||||
"input": "what is weather in sf",
|
||||
},
|
||||
)
|
||||
|
||||
assert app_w_interrupt.get_state(config) == StateSnapshot(
|
||||
values={
|
||||
"agent": {
|
||||
"agent_outcome": AgentAction(
|
||||
tool="search_api",
|
||||
tool_input="query",
|
||||
log="tool:search_api:a different query",
|
||||
),
|
||||
"input": "what is weather in sf",
|
||||
},
|
||||
},
|
||||
tasks=(PregelTask(AnyStr(), "tools", (PULL, "tools")),),
|
||||
next=("tools",),
|
||||
config={
|
||||
"configurable": {
|
||||
"thread_id": "2",
|
||||
"checkpoint_ns": "",
|
||||
"checkpoint_id": AnyStr(),
|
||||
}
|
||||
},
|
||||
created_at=AnyStr(),
|
||||
metadata={
|
||||
"parents": {},
|
||||
"source": "update",
|
||||
"step": 1,
|
||||
"thread_id": "2",
|
||||
},
|
||||
parent_config=(
|
||||
list(app_w_interrupt.checkpointer.list(config, limit=2))[-1].config
|
||||
),
|
||||
interrupts=(),
|
||||
)
|
||||
|
||||
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",
|
||||
"intermediate_steps": [
|
||||
[
|
||||
AgentAction(
|
||||
tool="search_api",
|
||||
tool_input="query",
|
||||
log="tool:search_api:a different query",
|
||||
),
|
||||
"result for query",
|
||||
]
|
||||
],
|
||||
}
|
||||
},
|
||||
{
|
||||
"agent": {
|
||||
"input": "what is weather in sf",
|
||||
"intermediate_steps": [
|
||||
[
|
||||
AgentAction(
|
||||
tool="search_api",
|
||||
tool_input="query",
|
||||
log="tool:search_api:a different query",
|
||||
),
|
||||
"result for query",
|
||||
]
|
||||
],
|
||||
"agent_outcome": AgentAction(
|
||||
tool="search_api",
|
||||
tool_input="another",
|
||||
log="tool:search_api:another",
|
||||
),
|
||||
}
|
||||
},
|
||||
]
|
||||
|
||||
app_w_interrupt.update_state(
|
||||
config,
|
||||
{
|
||||
"input": "what is weather in sf",
|
||||
"intermediate_steps": [
|
||||
[
|
||||
AgentAction(
|
||||
tool="search_api",
|
||||
tool_input="query",
|
||||
log="tool:search_api:a different query",
|
||||
),
|
||||
"result for query",
|
||||
]
|
||||
],
|
||||
"agent_outcome": AgentFinish(
|
||||
return_values={"answer": "a really nice answer"},
|
||||
log="finish:a really nice answer",
|
||||
),
|
||||
},
|
||||
)
|
||||
|
||||
assert app_w_interrupt.get_state(config) == StateSnapshot(
|
||||
values={
|
||||
"agent": {
|
||||
"input": "what is weather in sf",
|
||||
"intermediate_steps": [
|
||||
[
|
||||
AgentAction(
|
||||
tool="search_api",
|
||||
tool_input="query",
|
||||
log="tool:search_api:a different query",
|
||||
),
|
||||
"result for query",
|
||||
]
|
||||
],
|
||||
"agent_outcome": AgentFinish(
|
||||
return_values={"answer": "a really nice answer"},
|
||||
log="finish:a really nice answer",
|
||||
),
|
||||
},
|
||||
},
|
||||
tasks=(),
|
||||
next=(),
|
||||
config={
|
||||
"configurable": {
|
||||
"thread_id": "2",
|
||||
"checkpoint_ns": "",
|
||||
"checkpoint_id": AnyStr(),
|
||||
}
|
||||
},
|
||||
created_at=AnyStr(),
|
||||
metadata={
|
||||
"parents": {},
|
||||
"source": "update",
|
||||
"step": 4,
|
||||
"thread_id": "2",
|
||||
},
|
||||
parent_config=(
|
||||
list(app_w_interrupt.checkpointer.list(config, limit=2))[-1].config
|
||||
),
|
||||
interrupts=(),
|
||||
)
|
||||
|
||||
# test re-invoke to continue with interrupt_before
|
||||
|
||||
app_w_interrupt = workflow.compile(
|
||||
checkpointer=sync_checkpointer,
|
||||
interrupt_before=["tools"],
|
||||
)
|
||||
config = {"configurable": {"thread_id": "3"}}
|
||||
llm.i = 0 # reset the llm
|
||||
|
||||
assert [
|
||||
c for c in app_w_interrupt.stream({"input": "what is weather in sf"}, config)
|
||||
] == [
|
||||
{
|
||||
"agent": {
|
||||
"input": "what is weather in sf",
|
||||
"agent_outcome": AgentAction(
|
||||
tool="search_api", tool_input="query", log="tool:search_api:query"
|
||||
),
|
||||
}
|
||||
}
|
||||
]
|
||||
|
||||
assert app_w_interrupt.get_state(config) == StateSnapshot(
|
||||
values={
|
||||
"agent": {
|
||||
"input": "what is weather in sf",
|
||||
"agent_outcome": AgentAction(
|
||||
tool="search_api", tool_input="query", log="tool:search_api:query"
|
||||
),
|
||||
},
|
||||
},
|
||||
tasks=(PregelTask(AnyStr(), "tools", (PULL, "tools")),),
|
||||
next=("tools",),
|
||||
config={
|
||||
"configurable": {
|
||||
"thread_id": "3",
|
||||
"checkpoint_ns": "",
|
||||
"checkpoint_id": AnyStr(),
|
||||
}
|
||||
},
|
||||
created_at=AnyStr(),
|
||||
metadata={
|
||||
"parents": {},
|
||||
"source": "loop",
|
||||
"step": 0,
|
||||
"thread_id": "3",
|
||||
},
|
||||
parent_config=None,
|
||||
interrupts=(),
|
||||
)
|
||||
|
||||
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",
|
||||
"intermediate_steps": [
|
||||
[
|
||||
AgentAction(
|
||||
tool="search_api",
|
||||
tool_input="query",
|
||||
log="tool:search_api:query",
|
||||
),
|
||||
"result for query",
|
||||
]
|
||||
],
|
||||
}
|
||||
},
|
||||
{
|
||||
"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",
|
||||
),
|
||||
}
|
||||
},
|
||||
]
|
||||
|
||||
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",
|
||||
"intermediate_steps": [
|
||||
[
|
||||
AgentAction(
|
||||
tool="search_api",
|
||||
tool_input="query",
|
||||
log="tool:search_api:query",
|
||||
),
|
||||
"result for query",
|
||||
],
|
||||
[
|
||||
AgentAction(
|
||||
tool="search_api",
|
||||
tool_input="another",
|
||||
log="tool:search_api:another",
|
||||
),
|
||||
"result for another",
|
||||
],
|
||||
],
|
||||
}
|
||||
},
|
||||
{
|
||||
"agent": {
|
||||
"input": "what is weather in sf",
|
||||
"intermediate_steps": [
|
||||
[
|
||||
AgentAction(
|
||||
tool="search_api",
|
||||
tool_input="query",
|
||||
log="tool:search_api:query",
|
||||
),
|
||||
"result for query",
|
||||
],
|
||||
[
|
||||
AgentAction(
|
||||
tool="search_api",
|
||||
tool_input="another",
|
||||
log="tool:search_api:another",
|
||||
),
|
||||
"result for another",
|
||||
],
|
||||
],
|
||||
"agent_outcome": AgentFinish(
|
||||
return_values={"answer": "answer"}, log="finish:answer"
|
||||
),
|
||||
}
|
||||
},
|
||||
]
|
||||
|
||||
|
||||
def test_conditional_state_graph(
|
||||
snapshot: SnapshotAssertion,
|
||||
sync_checkpointer: BaseCheckpointSaver,
|
||||
|
||||
@@ -20,7 +20,6 @@ from langgraph.channels.ephemeral_value import EphemeralValue
|
||||
from langgraph.channels.last_value import LastValue
|
||||
from langgraph.checkpoint.base import BaseCheckpointSaver
|
||||
from langgraph.constants import END, PULL, PUSH, START
|
||||
from langgraph.graph.graph import Graph
|
||||
from langgraph.graph.message import MessageGraph, add_messages
|
||||
from langgraph.graph.state import StateGraph
|
||||
from langgraph.prebuilt.chat_agent_executor import create_react_agent
|
||||
@@ -495,871 +494,6 @@ async def test_fork_always_re_runs_nodes(
|
||||
]
|
||||
|
||||
|
||||
async def test_conditional_graph(async_checkpointer: BaseCheckpointSaver) -> None:
|
||||
from langchain_core.agents import AgentAction, AgentFinish
|
||||
from langchain_core.language_models.fake import FakeStreamingListLLM
|
||||
from langchain_core.prompts import PromptTemplate
|
||||
from langchain_core.runnables import RunnablePassthrough
|
||||
from langchain_core.tools import tool
|
||||
|
||||
# Assemble the tools
|
||||
@tool()
|
||||
def search_api(query: str) -> str:
|
||||
"""Searches the API for the query."""
|
||||
return f"result for {query}"
|
||||
|
||||
tools = [search_api]
|
||||
|
||||
# Construct the agent
|
||||
prompt = PromptTemplate.from_template("Hello!")
|
||||
|
||||
llm = FakeStreamingListLLM(
|
||||
responses=[
|
||||
"tool:search_api:query",
|
||||
"tool:search_api:another",
|
||||
"finish:answer",
|
||||
]
|
||||
)
|
||||
|
||||
async def agent_parser(input: str) -> Union[AgentAction, AgentFinish]:
|
||||
if input.startswith("finish"):
|
||||
_, answer = input.split(":")
|
||||
return AgentFinish(return_values={"answer": answer}, log=input)
|
||||
else:
|
||||
_, tool_name, tool_input = input.split(":")
|
||||
return AgentAction(tool=tool_name, tool_input=tool_input, log=input)
|
||||
|
||||
agent = RunnablePassthrough.assign(agent_outcome=prompt | llm | agent_parser)
|
||||
|
||||
# Define tool execution logic
|
||||
async def execute_tools(data: dict) -> dict:
|
||||
data = data.copy()
|
||||
agent_action: AgentAction = data.pop("agent_outcome")
|
||||
observation = await {t.name: t for t in tools}[agent_action.tool].ainvoke(
|
||||
agent_action.tool_input
|
||||
)
|
||||
if data.get("intermediate_steps") is None:
|
||||
data["intermediate_steps"] = []
|
||||
else:
|
||||
data["intermediate_steps"] = data["intermediate_steps"].copy()
|
||||
data["intermediate_steps"].append([agent_action, observation])
|
||||
return data
|
||||
|
||||
# Define decision-making logic
|
||||
async def should_continue(data: dict, config: RunnableConfig) -> str:
|
||||
# Logic to decide whether to continue in the loop or exit
|
||||
if isinstance(data["agent_outcome"], AgentFinish):
|
||||
return "exit"
|
||||
else:
|
||||
return "continue"
|
||||
|
||||
# Define a new graph
|
||||
workflow = Graph()
|
||||
|
||||
workflow.add_node("agent", agent)
|
||||
workflow.add_node("tools", execute_tools)
|
||||
|
||||
workflow.set_entry_point("agent")
|
||||
|
||||
workflow.add_conditional_edges(
|
||||
"agent", should_continue, {"continue": "tools", "exit": END}
|
||||
)
|
||||
|
||||
workflow.add_edge("tools", "agent")
|
||||
|
||||
app = workflow.compile()
|
||||
|
||||
assert await app.ainvoke({"input": "what is weather in sf"}) == {
|
||||
"input": "what is weather in sf",
|
||||
"intermediate_steps": [
|
||||
[
|
||||
AgentAction(
|
||||
tool="search_api",
|
||||
tool_input="query",
|
||||
log="tool:search_api:query",
|
||||
),
|
||||
"result for query",
|
||||
],
|
||||
[
|
||||
AgentAction(
|
||||
tool="search_api",
|
||||
tool_input="another",
|
||||
log="tool:search_api:another",
|
||||
),
|
||||
"result for another",
|
||||
],
|
||||
],
|
||||
"agent_outcome": AgentFinish(
|
||||
return_values={"answer": "answer"}, log="finish:answer"
|
||||
),
|
||||
}
|
||||
|
||||
assert [c async for c in app.astream({"input": "what is weather in sf"})] == [
|
||||
{
|
||||
"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",
|
||||
"intermediate_steps": [
|
||||
[
|
||||
AgentAction(
|
||||
tool="search_api",
|
||||
tool_input="query",
|
||||
log="tool:search_api:query",
|
||||
),
|
||||
"result for query",
|
||||
]
|
||||
],
|
||||
}
|
||||
},
|
||||
{
|
||||
"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",
|
||||
"intermediate_steps": [
|
||||
[
|
||||
AgentAction(
|
||||
tool="search_api",
|
||||
tool_input="query",
|
||||
log="tool:search_api:query",
|
||||
),
|
||||
"result for query",
|
||||
],
|
||||
[
|
||||
AgentAction(
|
||||
tool="search_api",
|
||||
tool_input="another",
|
||||
log="tool:search_api:another",
|
||||
),
|
||||
"result for another",
|
||||
],
|
||||
],
|
||||
}
|
||||
},
|
||||
{
|
||||
"agent": {
|
||||
"input": "what is weather in sf",
|
||||
"intermediate_steps": [
|
||||
[
|
||||
AgentAction(
|
||||
tool="search_api",
|
||||
tool_input="query",
|
||||
log="tool:search_api:query",
|
||||
),
|
||||
"result for query",
|
||||
],
|
||||
[
|
||||
AgentAction(
|
||||
tool="search_api",
|
||||
tool_input="another",
|
||||
log="tool:search_api:another",
|
||||
),
|
||||
"result for another",
|
||||
],
|
||||
],
|
||||
"agent_outcome": AgentFinish(
|
||||
return_values={"answer": "answer"}, log="finish:answer"
|
||||
),
|
||||
}
|
||||
},
|
||||
]
|
||||
|
||||
patches = [c async for c in app.astream_log({"input": "what is weather in sf"})]
|
||||
patch_paths = {op["path"] for log in patches for op in log.ops}
|
||||
|
||||
# Check that agent (one of the nodes) has its output streamed to the logs
|
||||
assert "/logs/agent/streamed_output/-" in patch_paths
|
||||
assert "/logs/agent:2/streamed_output/-" in patch_paths
|
||||
assert "/logs/agent:3/streamed_output/-" in patch_paths
|
||||
# Check that agent (one of the nodes) has its final output set in the logs
|
||||
assert "/logs/agent/final_output" in patch_paths
|
||||
assert "/logs/agent:2/final_output" in patch_paths
|
||||
assert "/logs/agent:3/final_output" in patch_paths
|
||||
assert [
|
||||
p["value"]
|
||||
for log in patches
|
||||
for p in log.ops
|
||||
if p["path"] == "/logs/agent/final_output"
|
||||
or p["path"] == "/logs/agent:2/final_output"
|
||||
or p["path"] == "/logs/agent:3/final_output"
|
||||
] == [
|
||||
{
|
||||
"input": "what is weather in sf",
|
||||
"agent_outcome": AgentAction(
|
||||
tool="search_api", tool_input="query", log="tool:search_api:query"
|
||||
),
|
||||
},
|
||||
{
|
||||
"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",
|
||||
),
|
||||
},
|
||||
{
|
||||
"input": "what is weather in sf",
|
||||
"intermediate_steps": [
|
||||
[
|
||||
AgentAction(
|
||||
tool="search_api",
|
||||
tool_input="query",
|
||||
log="tool:search_api:query",
|
||||
),
|
||||
"result for query",
|
||||
],
|
||||
[
|
||||
AgentAction(
|
||||
tool="search_api",
|
||||
tool_input="another",
|
||||
log="tool:search_api:another",
|
||||
),
|
||||
"result for another",
|
||||
],
|
||||
],
|
||||
"agent_outcome": AgentFinish(
|
||||
return_values={"answer": "answer"}, log="finish:answer"
|
||||
),
|
||||
},
|
||||
]
|
||||
|
||||
# test state get/update methods with interrupt_after
|
||||
|
||||
app_w_interrupt = workflow.compile(
|
||||
checkpointer=async_checkpointer,
|
||||
interrupt_after=["agent"],
|
||||
)
|
||||
config = {"configurable": {"thread_id": "1"}}
|
||||
|
||||
assert [
|
||||
c
|
||||
async for c in app_w_interrupt.astream(
|
||||
{"input": "what is weather in sf"}, config
|
||||
)
|
||||
] == [
|
||||
{
|
||||
"agent": {
|
||||
"input": "what is weather in sf",
|
||||
"agent_outcome": AgentAction(
|
||||
tool="search_api",
|
||||
tool_input="query",
|
||||
log="tool:search_api:query",
|
||||
),
|
||||
}
|
||||
}
|
||||
]
|
||||
|
||||
assert await app_w_interrupt.aget_state(config) == StateSnapshot(
|
||||
values={
|
||||
"agent": {
|
||||
"input": "what is weather in sf",
|
||||
"agent_outcome": AgentAction(
|
||||
tool="search_api",
|
||||
tool_input="query",
|
||||
log="tool:search_api:query",
|
||||
),
|
||||
},
|
||||
},
|
||||
tasks=(PregelTask(AnyStr(), "tools", (PULL, "tools")),),
|
||||
next=("tools",),
|
||||
config=(await app_w_interrupt.checkpointer.aget_tuple(config)).config,
|
||||
created_at=(await app_w_interrupt.checkpointer.aget_tuple(config)).checkpoint[
|
||||
"ts"
|
||||
],
|
||||
metadata={
|
||||
"parents": {},
|
||||
"source": "loop",
|
||||
"step": 0,
|
||||
"thread_id": "1",
|
||||
},
|
||||
parent_config=None,
|
||||
interrupts=(),
|
||||
)
|
||||
|
||||
await app_w_interrupt.aupdate_state(
|
||||
config,
|
||||
{
|
||||
"agent_outcome": AgentAction(
|
||||
tool="search_api",
|
||||
tool_input="query",
|
||||
log="tool:search_api:a different query",
|
||||
),
|
||||
"input": "what is weather in sf",
|
||||
},
|
||||
)
|
||||
|
||||
assert await app_w_interrupt.aget_state(config) == StateSnapshot(
|
||||
values={
|
||||
"agent": {
|
||||
"agent_outcome": AgentAction(
|
||||
tool="search_api",
|
||||
tool_input="query",
|
||||
log="tool:search_api:a different query",
|
||||
),
|
||||
"input": "what is weather in sf",
|
||||
},
|
||||
},
|
||||
tasks=(PregelTask(AnyStr(), "tools", (PULL, "tools")),),
|
||||
next=("tools",),
|
||||
config={
|
||||
"configurable": {
|
||||
"thread_id": "1",
|
||||
"checkpoint_ns": "",
|
||||
"checkpoint_id": AnyStr(),
|
||||
}
|
||||
},
|
||||
created_at=AnyStr(),
|
||||
metadata={
|
||||
"parents": {},
|
||||
"source": "update",
|
||||
"step": 1,
|
||||
"thread_id": "1",
|
||||
},
|
||||
parent_config=[
|
||||
c async for c in app_w_interrupt.checkpointer.alist(config, limit=2)
|
||||
][-1].config,
|
||||
interrupts=(),
|
||||
)
|
||||
|
||||
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",
|
||||
"intermediate_steps": [
|
||||
[
|
||||
AgentAction(
|
||||
tool="search_api",
|
||||
tool_input="query",
|
||||
log="tool:search_api:a different query",
|
||||
),
|
||||
"result for query",
|
||||
]
|
||||
],
|
||||
}
|
||||
},
|
||||
{
|
||||
"agent": {
|
||||
"input": "what is weather in sf",
|
||||
"intermediate_steps": [
|
||||
[
|
||||
AgentAction(
|
||||
tool="search_api",
|
||||
tool_input="query",
|
||||
log="tool:search_api:a different query",
|
||||
),
|
||||
"result for query",
|
||||
]
|
||||
],
|
||||
"agent_outcome": AgentAction(
|
||||
tool="search_api",
|
||||
tool_input="another",
|
||||
log="tool:search_api:another",
|
||||
),
|
||||
}
|
||||
},
|
||||
]
|
||||
|
||||
await app_w_interrupt.aupdate_state(
|
||||
config,
|
||||
{
|
||||
"input": "what is weather in sf",
|
||||
"intermediate_steps": [
|
||||
[
|
||||
AgentAction(
|
||||
tool="search_api",
|
||||
tool_input="query",
|
||||
log="tool:search_api:a different query",
|
||||
),
|
||||
"result for query",
|
||||
]
|
||||
],
|
||||
"agent_outcome": AgentFinish(
|
||||
return_values={"answer": "a really nice answer"},
|
||||
log="finish:a really nice answer",
|
||||
),
|
||||
},
|
||||
)
|
||||
|
||||
assert await app_w_interrupt.aget_state(config) == StateSnapshot(
|
||||
values={
|
||||
"agent": {
|
||||
"input": "what is weather in sf",
|
||||
"intermediate_steps": [
|
||||
[
|
||||
AgentAction(
|
||||
tool="search_api",
|
||||
tool_input="query",
|
||||
log="tool:search_api:a different query",
|
||||
),
|
||||
"result for query",
|
||||
]
|
||||
],
|
||||
"agent_outcome": AgentFinish(
|
||||
return_values={"answer": "a really nice answer"},
|
||||
log="finish:a really nice answer",
|
||||
),
|
||||
},
|
||||
},
|
||||
tasks=(),
|
||||
next=(),
|
||||
config={
|
||||
"configurable": {
|
||||
"thread_id": "1",
|
||||
"checkpoint_ns": "",
|
||||
"checkpoint_id": AnyStr(),
|
||||
}
|
||||
},
|
||||
created_at=AnyStr(),
|
||||
metadata={
|
||||
"parents": {},
|
||||
"source": "update",
|
||||
"step": 4,
|
||||
"thread_id": "1",
|
||||
},
|
||||
parent_config=[
|
||||
c async for c in app_w_interrupt.checkpointer.alist(config, limit=2)
|
||||
][-1].config,
|
||||
interrupts=(),
|
||||
)
|
||||
|
||||
# test state get/update methods with interrupt_before
|
||||
|
||||
app_w_interrupt = workflow.compile(
|
||||
checkpointer=async_checkpointer,
|
||||
interrupt_before=["tools"],
|
||||
)
|
||||
config = {"configurable": {"thread_id": "2"}}
|
||||
llm.i = 0
|
||||
|
||||
assert [
|
||||
c
|
||||
async for c in app_w_interrupt.astream(
|
||||
{"input": "what is weather in sf"}, config
|
||||
)
|
||||
] == [
|
||||
{
|
||||
"agent": {
|
||||
"input": "what is weather in sf",
|
||||
"agent_outcome": AgentAction(
|
||||
tool="search_api",
|
||||
tool_input="query",
|
||||
log="tool:search_api:query",
|
||||
),
|
||||
}
|
||||
}
|
||||
]
|
||||
|
||||
assert await app_w_interrupt.aget_state(config) == StateSnapshot(
|
||||
values={
|
||||
"agent": {
|
||||
"input": "what is weather in sf",
|
||||
"agent_outcome": AgentAction(
|
||||
tool="search_api",
|
||||
tool_input="query",
|
||||
log="tool:search_api:query",
|
||||
),
|
||||
},
|
||||
},
|
||||
tasks=(PregelTask(AnyStr(), "tools", (PULL, "tools")),),
|
||||
next=("tools",),
|
||||
config={
|
||||
"configurable": {
|
||||
"thread_id": "2",
|
||||
"checkpoint_ns": "",
|
||||
"checkpoint_id": AnyStr(),
|
||||
}
|
||||
},
|
||||
created_at=AnyStr(),
|
||||
metadata={
|
||||
"parents": {},
|
||||
"source": "loop",
|
||||
"step": 0,
|
||||
"thread_id": "2",
|
||||
},
|
||||
parent_config=None,
|
||||
interrupts=(),
|
||||
)
|
||||
|
||||
await app_w_interrupt.aupdate_state(
|
||||
config,
|
||||
{
|
||||
"agent_outcome": AgentAction(
|
||||
tool="search_api",
|
||||
tool_input="query",
|
||||
log="tool:search_api:a different query",
|
||||
),
|
||||
"input": "what is weather in sf",
|
||||
},
|
||||
)
|
||||
|
||||
assert await app_w_interrupt.aget_state(config) == StateSnapshot(
|
||||
values={
|
||||
"agent": {
|
||||
"agent_outcome": AgentAction(
|
||||
tool="search_api",
|
||||
tool_input="query",
|
||||
log="tool:search_api:a different query",
|
||||
),
|
||||
"input": "what is weather in sf",
|
||||
},
|
||||
},
|
||||
tasks=(PregelTask(AnyStr(), "tools", (PULL, "tools")),),
|
||||
next=("tools",),
|
||||
config={
|
||||
"configurable": {
|
||||
"thread_id": "2",
|
||||
"checkpoint_ns": "",
|
||||
"checkpoint_id": AnyStr(),
|
||||
}
|
||||
},
|
||||
created_at=AnyStr(),
|
||||
metadata={
|
||||
"parents": {},
|
||||
"source": "update",
|
||||
"step": 1,
|
||||
"thread_id": "2",
|
||||
},
|
||||
parent_config=[
|
||||
c async for c in app_w_interrupt.checkpointer.alist(config, limit=2)
|
||||
][-1].config,
|
||||
interrupts=(),
|
||||
)
|
||||
|
||||
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",
|
||||
"intermediate_steps": [
|
||||
[
|
||||
AgentAction(
|
||||
tool="search_api",
|
||||
tool_input="query",
|
||||
log="tool:search_api:a different query",
|
||||
),
|
||||
"result for query",
|
||||
]
|
||||
],
|
||||
}
|
||||
},
|
||||
{
|
||||
"agent": {
|
||||
"input": "what is weather in sf",
|
||||
"intermediate_steps": [
|
||||
[
|
||||
AgentAction(
|
||||
tool="search_api",
|
||||
tool_input="query",
|
||||
log="tool:search_api:a different query",
|
||||
),
|
||||
"result for query",
|
||||
]
|
||||
],
|
||||
"agent_outcome": AgentAction(
|
||||
tool="search_api",
|
||||
tool_input="another",
|
||||
log="tool:search_api:another",
|
||||
),
|
||||
}
|
||||
},
|
||||
]
|
||||
|
||||
await app_w_interrupt.aupdate_state(
|
||||
config,
|
||||
{
|
||||
"input": "what is weather in sf",
|
||||
"intermediate_steps": [
|
||||
[
|
||||
AgentAction(
|
||||
tool="search_api",
|
||||
tool_input="query",
|
||||
log="tool:search_api:a different query",
|
||||
),
|
||||
"result for query",
|
||||
]
|
||||
],
|
||||
"agent_outcome": AgentFinish(
|
||||
return_values={"answer": "a really nice answer"},
|
||||
log="finish:a really nice answer",
|
||||
),
|
||||
},
|
||||
)
|
||||
|
||||
assert await app_w_interrupt.aget_state(config) == StateSnapshot(
|
||||
values={
|
||||
"agent": {
|
||||
"input": "what is weather in sf",
|
||||
"intermediate_steps": [
|
||||
[
|
||||
AgentAction(
|
||||
tool="search_api",
|
||||
tool_input="query",
|
||||
log="tool:search_api:a different query",
|
||||
),
|
||||
"result for query",
|
||||
]
|
||||
],
|
||||
"agent_outcome": AgentFinish(
|
||||
return_values={"answer": "a really nice answer"},
|
||||
log="finish:a really nice answer",
|
||||
),
|
||||
},
|
||||
},
|
||||
tasks=(),
|
||||
next=(),
|
||||
config={
|
||||
"configurable": {
|
||||
"thread_id": "2",
|
||||
"checkpoint_ns": "",
|
||||
"checkpoint_id": AnyStr(),
|
||||
}
|
||||
},
|
||||
created_at=AnyStr(),
|
||||
metadata={
|
||||
"parents": {},
|
||||
"source": "update",
|
||||
"step": 4,
|
||||
"thread_id": "2",
|
||||
},
|
||||
parent_config=[
|
||||
c async for c in app_w_interrupt.checkpointer.alist(config, limit=2)
|
||||
][-1].config,
|
||||
interrupts=(),
|
||||
)
|
||||
|
||||
# test re-invoke to continue with interrupt_before
|
||||
|
||||
app_w_interrupt = workflow.compile(
|
||||
checkpointer=async_checkpointer,
|
||||
interrupt_before=["tools"],
|
||||
)
|
||||
config = {"configurable": {"thread_id": "3"}}
|
||||
llm.i = 0 # reset the llm
|
||||
|
||||
assert [
|
||||
c
|
||||
async for c in app_w_interrupt.astream(
|
||||
{"input": "what is weather in sf"}, config
|
||||
)
|
||||
] == [
|
||||
{
|
||||
"agent": {
|
||||
"input": "what is weather in sf",
|
||||
"agent_outcome": AgentAction(
|
||||
tool="search_api",
|
||||
tool_input="query",
|
||||
log="tool:search_api:query",
|
||||
),
|
||||
}
|
||||
}
|
||||
]
|
||||
|
||||
assert await app_w_interrupt.aget_state(config) == StateSnapshot(
|
||||
values={
|
||||
"agent": {
|
||||
"input": "what is weather in sf",
|
||||
"agent_outcome": AgentAction(
|
||||
tool="search_api",
|
||||
tool_input="query",
|
||||
log="tool:search_api:query",
|
||||
),
|
||||
},
|
||||
},
|
||||
tasks=(PregelTask(AnyStr(), "tools", (PULL, "tools")),),
|
||||
next=("tools",),
|
||||
config={
|
||||
"configurable": {
|
||||
"thread_id": "3",
|
||||
"checkpoint_ns": "",
|
||||
"checkpoint_id": AnyStr(),
|
||||
}
|
||||
},
|
||||
created_at=AnyStr(),
|
||||
metadata={
|
||||
"parents": {},
|
||||
"source": "loop",
|
||||
"step": 0,
|
||||
"thread_id": "3",
|
||||
},
|
||||
parent_config=None,
|
||||
interrupts=(),
|
||||
)
|
||||
|
||||
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",
|
||||
"intermediate_steps": [
|
||||
[
|
||||
AgentAction(
|
||||
tool="search_api",
|
||||
tool_input="query",
|
||||
log="tool:search_api:query",
|
||||
),
|
||||
"result for query",
|
||||
]
|
||||
],
|
||||
}
|
||||
},
|
||||
{
|
||||
"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",
|
||||
),
|
||||
}
|
||||
},
|
||||
]
|
||||
|
||||
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",
|
||||
"intermediate_steps": [
|
||||
[
|
||||
AgentAction(
|
||||
tool="search_api",
|
||||
tool_input="query",
|
||||
log="tool:search_api:query",
|
||||
),
|
||||
"result for query",
|
||||
],
|
||||
[
|
||||
AgentAction(
|
||||
tool="search_api",
|
||||
tool_input="another",
|
||||
log="tool:search_api:another",
|
||||
),
|
||||
"result for another",
|
||||
],
|
||||
],
|
||||
}
|
||||
},
|
||||
{
|
||||
"agent": {
|
||||
"input": "what is weather in sf",
|
||||
"intermediate_steps": [
|
||||
[
|
||||
AgentAction(
|
||||
tool="search_api",
|
||||
tool_input="query",
|
||||
log="tool:search_api:query",
|
||||
),
|
||||
"result for query",
|
||||
],
|
||||
[
|
||||
AgentAction(
|
||||
tool="search_api",
|
||||
tool_input="another",
|
||||
log="tool:search_api:another",
|
||||
),
|
||||
"result for another",
|
||||
],
|
||||
],
|
||||
"agent_outcome": AgentFinish(
|
||||
return_values={"answer": "answer"}, log="finish:answer"
|
||||
),
|
||||
}
|
||||
},
|
||||
]
|
||||
|
||||
|
||||
async def test_conditional_graph_state(async_checkpointer: BaseCheckpointSaver) -> None:
|
||||
from langchain_core.agents import AgentAction, AgentFinish
|
||||
from langchain_core.language_models.fake import FakeStreamingListLLM
|
||||
|
||||
@@ -45,7 +45,7 @@ from langgraph.config import get_stream_writer
|
||||
from langgraph.constants import CONFIG_KEY_NODE_FINISHED, ERROR, PULL, START
|
||||
from langgraph.errors import InvalidUpdateError
|
||||
from langgraph.func import entrypoint, task
|
||||
from langgraph.graph import END, Graph, StateGraph
|
||||
from langgraph.graph import END, StateGraph
|
||||
from langgraph.graph.message import MessageGraph, MessagesState, add_messages
|
||||
from langgraph.prebuilt.tool_node import ToolNode
|
||||
from langgraph.pregel import (
|
||||
@@ -83,81 +83,6 @@ logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def test_graph_validation() -> None:
|
||||
def logic(inp: str) -> str:
|
||||
return ""
|
||||
|
||||
workflow = Graph()
|
||||
workflow.add_node("agent", logic)
|
||||
workflow.set_entry_point("agent")
|
||||
workflow.set_finish_point("agent")
|
||||
assert workflow.compile(), "valid graph"
|
||||
|
||||
# Accept a dead-end
|
||||
workflow = Graph()
|
||||
workflow.add_node("agent", logic)
|
||||
workflow.set_entry_point("agent")
|
||||
workflow.compile()
|
||||
|
||||
workflow = Graph()
|
||||
workflow.add_node("agent", logic)
|
||||
workflow.set_finish_point("agent")
|
||||
with pytest.raises(ValueError, match="must have an entrypoint"):
|
||||
workflow.compile()
|
||||
|
||||
workflow = Graph()
|
||||
workflow.add_node("agent", logic)
|
||||
workflow.add_node("tools", logic)
|
||||
workflow.set_entry_point("agent")
|
||||
workflow.add_conditional_edges("agent", logic, {"continue": "tools", "exit": END})
|
||||
workflow.add_edge("tools", "agent")
|
||||
assert workflow.compile(), "valid graph"
|
||||
|
||||
workflow = Graph()
|
||||
workflow.add_node("agent", logic)
|
||||
workflow.add_node("tools", logic)
|
||||
workflow.set_entry_point("tools")
|
||||
workflow.add_conditional_edges("agent", logic, {"continue": "tools", "exit": END})
|
||||
workflow.add_edge("tools", "agent")
|
||||
assert workflow.compile(), "valid graph"
|
||||
|
||||
workflow = Graph()
|
||||
workflow.set_entry_point("tools")
|
||||
workflow.add_conditional_edges("agent", logic, {"continue": "tools", "exit": END})
|
||||
workflow.add_edge("tools", "agent")
|
||||
workflow.add_node("agent", logic)
|
||||
workflow.add_node("tools", logic)
|
||||
assert workflow.compile(), "valid graph"
|
||||
|
||||
workflow = Graph()
|
||||
workflow.set_entry_point("tools")
|
||||
workflow.add_conditional_edges(
|
||||
"agent", logic, {"continue": "tools", "exit": END, "hmm": "extra"}
|
||||
)
|
||||
workflow.add_edge("tools", "agent")
|
||||
workflow.add_node("agent", logic)
|
||||
workflow.add_node("tools", logic)
|
||||
with pytest.raises(ValueError, match="unknown"): # extra is not defined
|
||||
workflow.compile()
|
||||
|
||||
workflow = Graph()
|
||||
workflow.set_entry_point("agent")
|
||||
workflow.add_conditional_edges("agent", logic, {"continue": "tools", "exit": END})
|
||||
workflow.add_edge("tools", "extra")
|
||||
workflow.add_node("agent", logic)
|
||||
workflow.add_node("tools", logic)
|
||||
with pytest.raises(ValueError, match="unknown"): # extra is not defined
|
||||
workflow.compile()
|
||||
|
||||
workflow = Graph()
|
||||
workflow.add_node("agent", logic)
|
||||
workflow.add_node("tools", logic)
|
||||
workflow.add_node("extra", logic)
|
||||
workflow.set_entry_point("agent")
|
||||
workflow.add_conditional_edges("agent", logic)
|
||||
workflow.add_edge("tools", "agent")
|
||||
# Accept, even though extra is dead-end
|
||||
workflow.compile()
|
||||
|
||||
class State(TypedDict):
|
||||
hello: str
|
||||
|
||||
@@ -490,11 +415,6 @@ def test_invoke_single_process_in_out(mocker: MockerFixture) -> None:
|
||||
input_channels="input",
|
||||
output_channels="output",
|
||||
)
|
||||
graph = Graph()
|
||||
graph.add_node("add_one", add_one)
|
||||
graph.set_entry_point("add_one")
|
||||
graph.set_finish_point("add_one")
|
||||
gapp = graph.compile()
|
||||
|
||||
assert app.input_schema.model_json_schema() == {
|
||||
"title": "LangGraphInput",
|
||||
@@ -516,21 +436,6 @@ def test_invoke_single_process_in_out(mocker: MockerFixture) -> None:
|
||||
assert app.invoke(2, output_keys=["output"]) == {"output": 3}
|
||||
assert repr(app), "does not raise recursion error"
|
||||
|
||||
assert gapp.invoke(2, debug=True) == 3
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"falsy_value",
|
||||
[None, False, 0, "", [], {}, set(), frozenset(), 0.0, 0j],
|
||||
)
|
||||
def test_invoke_single_process_in_out_falsy_values(falsy_value: Any) -> None:
|
||||
graph = Graph()
|
||||
graph.add_node("return_falsy_const", lambda *args, **kwargs: falsy_value)
|
||||
graph.set_entry_point("return_falsy_const")
|
||||
graph.set_finish_point("return_falsy_const")
|
||||
gapp = graph.compile()
|
||||
assert gapp.invoke(1) == falsy_value
|
||||
|
||||
|
||||
def test_invoke_single_process_in_write_kwargs(mocker: MockerFixture) -> None:
|
||||
add_one = mocker.Mock(side_effect=lambda x: x + 1)
|
||||
@@ -644,29 +549,6 @@ def test_invoke_two_processes_in_out(mocker: MockerFixture) -> None:
|
||||
with pytest.raises(GraphRecursionError):
|
||||
app.invoke(2, {"recursion_limit": 1}, debug=1)
|
||||
|
||||
graph = Graph()
|
||||
graph.add_node("add_one", add_one)
|
||||
graph.add_node("add_one_more", add_one)
|
||||
graph.set_entry_point("add_one")
|
||||
graph.set_finish_point("add_one_more")
|
||||
graph.add_edge("add_one", "add_one_more")
|
||||
gapp = graph.compile()
|
||||
|
||||
assert gapp.invoke(2) == 4
|
||||
|
||||
for step, values in enumerate(gapp.stream(2, debug=1), start=1):
|
||||
if step == 1:
|
||||
assert values == {
|
||||
"add_one": 3,
|
||||
}
|
||||
elif step == 2:
|
||||
assert values == {
|
||||
"add_one_more": 4,
|
||||
}
|
||||
else:
|
||||
assert 0, f"{step}:{values}"
|
||||
assert step == 2
|
||||
|
||||
|
||||
def test_run_from_checkpoint_id_retains_previous_writes(
|
||||
sync_checkpointer: BaseCheckpointSaver,
|
||||
@@ -773,16 +655,6 @@ def test_batch_two_processes_in_out() -> None:
|
||||
{"output": 7},
|
||||
]
|
||||
|
||||
graph = Graph()
|
||||
graph.add_node("add_one", add_one_with_delay)
|
||||
graph.add_node("add_one_more", add_one_with_delay)
|
||||
graph.set_entry_point("add_one")
|
||||
graph.set_finish_point("add_one_more")
|
||||
graph.add_edge("add_one", "add_one_more")
|
||||
gapp = graph.compile()
|
||||
|
||||
assert gapp.batch([3, 2, 1, 3, 5]) == [5, 4, 3, 5, 7]
|
||||
|
||||
|
||||
def test_invoke_many_processes_in_out(mocker: MockerFixture) -> None:
|
||||
test_size = 100
|
||||
@@ -1827,50 +1699,6 @@ def test_invoke_two_processes_no_in(mocker: MockerFixture) -> None:
|
||||
Pregel(nodes={"one": one, "two": two})
|
||||
|
||||
|
||||
def test_conditional_entrypoint_graph(snapshot: SnapshotAssertion) -> None:
|
||||
def left(data: str) -> str:
|
||||
return data + "->left"
|
||||
|
||||
def right(data: str) -> str:
|
||||
return data + "->right"
|
||||
|
||||
def should_start(data: str) -> str:
|
||||
# Logic to decide where to start
|
||||
if len(data) > 10:
|
||||
return "go-right"
|
||||
else:
|
||||
return "go-left"
|
||||
|
||||
# Define a new graph
|
||||
workflow = Graph()
|
||||
|
||||
workflow.add_node("left", left)
|
||||
workflow.add_node("right", right)
|
||||
|
||||
workflow.set_conditional_entry_point(
|
||||
should_start, {"go-left": "left", "go-right": "right"}
|
||||
)
|
||||
|
||||
workflow.add_conditional_edges("left", lambda data: END, {END: END})
|
||||
workflow.add_edge("right", END)
|
||||
|
||||
app = workflow.compile()
|
||||
|
||||
assert json.dumps(app.get_input_schema().model_json_schema()) == snapshot
|
||||
assert json.dumps(app.get_output_schema().model_json_schema()) == snapshot
|
||||
assert json.dumps(app.get_graph().to_json(), indent=2) == snapshot
|
||||
assert app.get_graph().draw_mermaid(with_styles=False) == snapshot
|
||||
|
||||
assert (
|
||||
app.invoke("what is weather in sf", debug=True)
|
||||
== "what is weather in sf->right"
|
||||
)
|
||||
|
||||
assert [*app.stream("what is weather in sf")] == [
|
||||
{"right": "what is weather in sf->right"},
|
||||
]
|
||||
|
||||
|
||||
def test_conditional_entrypoint_to_multiple_state_graph(
|
||||
snapshot: SnapshotAssertion,
|
||||
) -> None:
|
||||
|
||||
@@ -43,7 +43,7 @@ from langgraph.checkpoint.memory import InMemorySaver
|
||||
from langgraph.constants import CONFIG_KEY_NODE_FINISHED, ERROR, PULL, PUSH, START
|
||||
from langgraph.errors import InvalidUpdateError, NodeInterrupt
|
||||
from langgraph.func import entrypoint, task
|
||||
from langgraph.graph import END, Graph, StateGraph
|
||||
from langgraph.graph import END, StateGraph
|
||||
from langgraph.graph.message import MessagesState, add_messages
|
||||
from langgraph.prebuilt.tool_node import ToolNode
|
||||
from langgraph.pregel import GraphRecursionError, NodeBuilder, Pregel, StateSnapshot
|
||||
@@ -262,7 +262,10 @@ async def test_checkpoint_put_after_cancellation() -> None:
|
||||
finally:
|
||||
logs.append("awhile.end")
|
||||
|
||||
builder = Graph()
|
||||
class State(TypedDict):
|
||||
hello: str
|
||||
|
||||
builder = StateGraph(State)
|
||||
builder.add_node("agent", awhile)
|
||||
builder.set_entry_point("agent")
|
||||
builder.set_finish_point("agent")
|
||||
@@ -271,7 +274,7 @@ async def test_checkpoint_put_after_cancellation() -> None:
|
||||
thread1 = {"configurable": {"thread_id": "1"}}
|
||||
|
||||
# start the task
|
||||
t = asyncio.create_task(graph.ainvoke(1, thread1))
|
||||
t = asyncio.create_task(graph.ainvoke({"hello": "world"}, thread1))
|
||||
# cancel after 0.2 seconds
|
||||
await asyncio.sleep(0.2)
|
||||
t.cancel()
|
||||
@@ -325,7 +328,10 @@ async def test_checkpoint_put_after_cancellation_stream_anext() -> None:
|
||||
finally:
|
||||
logs.append("awhile.end")
|
||||
|
||||
builder = Graph()
|
||||
class State(TypedDict):
|
||||
hello: str
|
||||
|
||||
builder = StateGraph(State)
|
||||
builder.add_node("agent", awhile)
|
||||
builder.set_entry_point("agent")
|
||||
builder.set_finish_point("agent")
|
||||
@@ -334,7 +340,7 @@ async def test_checkpoint_put_after_cancellation_stream_anext() -> None:
|
||||
thread1 = {"configurable": {"thread_id": "1"}}
|
||||
|
||||
# start the task
|
||||
s = graph.astream(1, thread1)
|
||||
s = graph.astream({"hello": "world"}, thread1)
|
||||
t = asyncio.create_task(s.__anext__())
|
||||
# cancel after 0.2 seconds
|
||||
await asyncio.sleep(0.2)
|
||||
@@ -389,7 +395,10 @@ async def test_checkpoint_put_after_cancellation_stream_events_anext() -> None:
|
||||
finally:
|
||||
logs.append("awhile.end")
|
||||
|
||||
builder = Graph()
|
||||
class State(TypedDict):
|
||||
hello: str
|
||||
|
||||
builder = StateGraph(State)
|
||||
builder.add_node("agent", awhile)
|
||||
builder.set_entry_point("agent")
|
||||
builder.set_finish_point("agent")
|
||||
@@ -398,7 +407,9 @@ async def test_checkpoint_put_after_cancellation_stream_events_anext() -> None:
|
||||
thread1 = {"configurable": {"thread_id": "1"}}
|
||||
|
||||
# start the task
|
||||
s = graph.astream_events(1, thread1, version="v2", include_names=["LangGraph"])
|
||||
s = graph.astream_events(
|
||||
{"hello": "world"}, thread1, version="v2", include_names=["LangGraph"]
|
||||
)
|
||||
# skip first event (happens right away)
|
||||
await s.__anext__()
|
||||
# start the task for 2nd event
|
||||
@@ -436,7 +447,10 @@ async def test_node_cancellation_on_external_cancel() -> None:
|
||||
inner_task_cancelled = True
|
||||
raise
|
||||
|
||||
builder = Graph()
|
||||
class State(TypedDict):
|
||||
hello: str
|
||||
|
||||
builder = StateGraph(State)
|
||||
builder.add_node("agent", awhile)
|
||||
builder.set_entry_point("agent")
|
||||
builder.set_finish_point("agent")
|
||||
@@ -444,7 +458,7 @@ async def test_node_cancellation_on_external_cancel() -> None:
|
||||
graph = builder.compile()
|
||||
|
||||
with pytest.raises(asyncio.TimeoutError):
|
||||
await asyncio.wait_for(graph.ainvoke(1), 0.5)
|
||||
await asyncio.wait_for(graph.ainvoke({"hello": "world"}), 0.5)
|
||||
|
||||
assert inner_task_cancelled
|
||||
|
||||
@@ -463,7 +477,10 @@ async def test_node_cancellation_on_other_node_exception() -> None:
|
||||
async def iambad(input: Any) -> None:
|
||||
raise ValueError("I am bad")
|
||||
|
||||
builder = Graph()
|
||||
class State(TypedDict):
|
||||
hello: str
|
||||
|
||||
builder = StateGraph(State)
|
||||
builder.add_node("agent", awhile)
|
||||
builder.add_node("bad", iambad)
|
||||
builder.set_conditional_entry_point(lambda _: ["agent", "bad"], then=END)
|
||||
@@ -472,7 +489,7 @@ async def test_node_cancellation_on_other_node_exception() -> None:
|
||||
|
||||
with pytest.raises(ValueError, match="I am bad"):
|
||||
# This will raise ValueError, not TimeoutError
|
||||
await asyncio.wait_for(graph.ainvoke(1), 0.5)
|
||||
await asyncio.wait_for(graph.ainvoke({"hello": "world"}), 0.5)
|
||||
|
||||
assert inner_task_cancelled
|
||||
|
||||
@@ -484,7 +501,10 @@ async def test_node_cancellation_on_other_node_exception_two() -> None:
|
||||
async def iambad(input: Any) -> None:
|
||||
raise ValueError("I am bad")
|
||||
|
||||
builder = Graph()
|
||||
class State(TypedDict):
|
||||
hello: str
|
||||
|
||||
builder = StateGraph(State)
|
||||
builder.add_node("agent", awhile)
|
||||
builder.add_node("bad", iambad)
|
||||
builder.set_conditional_entry_point(lambda _: ["agent", "bad"], then=END)
|
||||
@@ -493,7 +513,7 @@ async def test_node_cancellation_on_other_node_exception_two() -> None:
|
||||
|
||||
with pytest.raises(ValueError, match="I am bad"):
|
||||
# This will raise ValueError, not CancelledError
|
||||
await graph.ainvoke(1)
|
||||
await graph.ainvoke({"hello": "world"})
|
||||
|
||||
|
||||
@NEEDS_CONTEXTVARS
|
||||
@@ -1113,9 +1133,12 @@ async def test_step_timeout_on_stream_hang(stream_hang_s: float) -> None:
|
||||
|
||||
async def alittlewhile(input: Any) -> None:
|
||||
await asyncio.sleep(0.6)
|
||||
return "1"
|
||||
return {"hello": "1"}
|
||||
|
||||
builder = Graph()
|
||||
class State(TypedDict):
|
||||
hello: str
|
||||
|
||||
builder = StateGraph(State)
|
||||
builder.add_node(awhile)
|
||||
builder.add_node(alittlewhile)
|
||||
builder.set_conditional_entry_point(lambda _: ["awhile", "alittlewhile"], then=END)
|
||||
@@ -1123,8 +1146,8 @@ async def test_step_timeout_on_stream_hang(stream_hang_s: float) -> None:
|
||||
graph.step_timeout = 1
|
||||
|
||||
with pytest.raises(asyncio.TimeoutError):
|
||||
async for chunk in graph.astream(1, stream_mode="updates"):
|
||||
assert chunk == {"alittlewhile": {"alittlewhile": "1"}}
|
||||
async for chunk in graph.astream({"hello": "world"}, stream_mode="updates"):
|
||||
assert chunk == {"alittlewhile": {"hello": "1"}}
|
||||
await asyncio.sleep(stream_hang_s)
|
||||
|
||||
assert inner_task_cancelled
|
||||
@@ -1382,11 +1405,6 @@ async def test_invoke_single_process_in_out(mocker: MockerFixture) -> None:
|
||||
input_channels="input",
|
||||
output_channels="output",
|
||||
)
|
||||
graph = Graph()
|
||||
graph.add_node("add_one", add_one)
|
||||
graph.set_entry_point("add_one")
|
||||
graph.set_finish_point("add_one")
|
||||
gapp = graph.compile()
|
||||
|
||||
assert app.input_schema.model_json_schema() == {
|
||||
"title": "LangGraphInput",
|
||||
@@ -1399,21 +1417,6 @@ async def test_invoke_single_process_in_out(mocker: MockerFixture) -> None:
|
||||
assert await app.ainvoke(2) == 3
|
||||
assert await app.ainvoke(2, output_keys=["output"]) == {"output": 3}
|
||||
|
||||
assert await gapp.ainvoke(2) == 3
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"falsy_value",
|
||||
[None, False, 0, "", [], {}, set(), frozenset(), 0.0, 0j],
|
||||
)
|
||||
async def test_invoke_single_process_in_out_falsy_values(falsy_value: Any) -> None:
|
||||
graph = Graph()
|
||||
graph.add_node("return_falsy_const", lambda *args, **kwargs: falsy_value)
|
||||
graph.set_entry_point("return_falsy_const")
|
||||
graph.set_finish_point("return_falsy_const")
|
||||
gapp = graph.compile()
|
||||
assert falsy_value == await gapp.ainvoke(1)
|
||||
|
||||
|
||||
async def test_invoke_single_process_in_write_kwargs(mocker: MockerFixture) -> None:
|
||||
add_one = mocker.Mock(side_effect=lambda x: x + 1)
|
||||
@@ -1543,29 +1546,6 @@ async def test_invoke_two_processes_in_out(mocker: MockerFixture) -> None:
|
||||
}
|
||||
assert step == 2
|
||||
|
||||
graph = Graph()
|
||||
graph.add_node("add_one", add_one)
|
||||
graph.add_node("add_one_more", add_one)
|
||||
graph.set_entry_point("add_one")
|
||||
graph.set_finish_point("add_one_more")
|
||||
graph.add_edge("add_one", "add_one_more")
|
||||
gapp = graph.compile()
|
||||
|
||||
assert await gapp.ainvoke(2) == 4
|
||||
|
||||
step = 0
|
||||
async for values in gapp.astream(2):
|
||||
step += 1
|
||||
if step == 1:
|
||||
assert values == {
|
||||
"add_one": 3,
|
||||
}
|
||||
elif step == 2:
|
||||
assert values == {
|
||||
"add_one_more": 4,
|
||||
}
|
||||
assert step == 2
|
||||
|
||||
|
||||
async def test_batch_two_processes_in_out() -> None:
|
||||
async def add_one_with_delay(inp: int) -> int:
|
||||
@@ -1595,16 +1575,6 @@ async def test_batch_two_processes_in_out() -> None:
|
||||
{"output": 7},
|
||||
]
|
||||
|
||||
graph = Graph()
|
||||
graph.add_node("add_one", add_one_with_delay)
|
||||
graph.add_node("add_one_more", add_one_with_delay)
|
||||
graph.set_entry_point("add_one")
|
||||
graph.set_finish_point("add_one_more")
|
||||
graph.add_edge("add_one", "add_one_more")
|
||||
gapp = graph.compile()
|
||||
|
||||
assert await gapp.abatch([3, 2, 1, 3, 5]) == [5, 4, 3, 5, 7]
|
||||
|
||||
|
||||
async def test_invoke_many_processes_in_out(mocker: MockerFixture) -> None:
|
||||
test_size = 100
|
||||
@@ -3845,42 +3815,6 @@ async def test_invoke_two_processes_no_out(mocker: MockerFixture) -> None:
|
||||
assert await app.ainvoke(2) is None
|
||||
|
||||
|
||||
async def test_conditional_entrypoint_graph() -> None:
|
||||
async def left(data: str) -> str:
|
||||
return data + "->left"
|
||||
|
||||
async def right(data: str) -> str:
|
||||
return data + "->right"
|
||||
|
||||
def should_start(data: str) -> str:
|
||||
# Logic to decide where to start
|
||||
if len(data) > 10:
|
||||
return "go-right"
|
||||
else:
|
||||
return "go-left"
|
||||
|
||||
# Define a new graph
|
||||
workflow = Graph()
|
||||
|
||||
workflow.add_node("left", left)
|
||||
workflow.add_node("right", right)
|
||||
|
||||
workflow.set_conditional_entry_point(
|
||||
should_start, {"go-left": "left", "go-right": "right"}
|
||||
)
|
||||
|
||||
workflow.add_conditional_edges("left", lambda data: END)
|
||||
workflow.add_edge("right", END)
|
||||
|
||||
app = workflow.compile()
|
||||
|
||||
assert await app.ainvoke("what is weather in sf") == "what is weather in sf->right"
|
||||
|
||||
assert [c async for c in app.astream("what is weather in sf")] == [
|
||||
{"right": "what is weather in sf->right"},
|
||||
]
|
||||
|
||||
|
||||
async def test_conditional_entrypoint_graph_state() -> None:
|
||||
class AgentState(TypedDict, total=False):
|
||||
input: str
|
||||
|
||||
@@ -18,7 +18,7 @@ import pytest
|
||||
from typing_extensions import NotRequired, Required, TypedDict
|
||||
|
||||
from langgraph.graph import END, StateGraph
|
||||
from langgraph.graph.graph import CompiledGraph
|
||||
from langgraph.graph.state import CompiledStateGraph
|
||||
from langgraph.utils.config import _is_not_empty
|
||||
from langgraph.utils.fields import (
|
||||
_is_optional_type,
|
||||
@@ -103,7 +103,7 @@ def test_is_generator() -> None:
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def rt_graph() -> CompiledGraph:
|
||||
def rt_graph() -> CompiledStateGraph:
|
||||
class State(TypedDict):
|
||||
foo: int
|
||||
node_run_id: int
|
||||
@@ -120,7 +120,7 @@ def rt_graph() -> CompiledGraph:
|
||||
return graph.compile()
|
||||
|
||||
|
||||
def test_runnable_callable_tracing_nested(rt_graph: CompiledGraph) -> None:
|
||||
def test_runnable_callable_tracing_nested(rt_graph: CompiledStateGraph) -> None:
|
||||
with patch("langsmith.client.Client", spec=langsmith.Client) as mock_client:
|
||||
with patch("langchain_core.tracers.langchain.get_client") as mock_get_client:
|
||||
mock_get_client.return_value = mock_client
|
||||
@@ -133,7 +133,9 @@ def test_runnable_callable_tracing_nested(rt_graph: CompiledGraph) -> None:
|
||||
sys.version_info < (3, 11),
|
||||
reason="Python 3.11+ is required for async contextvars support",
|
||||
)
|
||||
async def test_runnable_callable_tracing_nested_async(rt_graph: CompiledGraph) -> None:
|
||||
async def test_runnable_callable_tracing_nested_async(
|
||||
rt_graph: CompiledStateGraph,
|
||||
) -> None:
|
||||
with patch("langsmith.client.Client", spec=langsmith.Client) as mock_client:
|
||||
with patch("langchain_core.tracers.langchain.get_client") as mock_get_client:
|
||||
mock_get_client.return_value = mock_client
|
||||
|
||||
Generated
+387
-358
File diff suppressed because it is too large
Load Diff
@@ -36,8 +36,8 @@ from typing_extensions import Annotated, TypedDict
|
||||
|
||||
from langgraph.errors import ErrorCode, create_error_message
|
||||
from langgraph.graph import END, StateGraph
|
||||
from langgraph.graph.graph import CompiledGraph
|
||||
from langgraph.graph.message import add_messages
|
||||
from langgraph.graph.state import CompiledStateGraph
|
||||
from langgraph.managed import IsLastStep, RemainingSteps
|
||||
from langgraph.prebuilt.tool_node import ToolNode
|
||||
from langgraph.store.base import BaseStore
|
||||
@@ -257,7 +257,7 @@ def create_react_agent(
|
||||
debug: bool = False,
|
||||
version: Literal["v1", "v2"] = "v2",
|
||||
name: Optional[str] = None,
|
||||
) -> CompiledGraph:
|
||||
) -> CompiledStateGraph:
|
||||
"""Creates an agent graph that calls tools in a loop until a stopping condition is met.
|
||||
|
||||
For more details on using `create_react_agent`, visit [Agents](https://langchain-ai.github.io/langgraph/agents/overview/) documentation.
|
||||
|
||||
Reference in New Issue
Block a user