Remove non-state Graph

- Not used in any examples/docs
This commit is contained in:
Nuno Campos
2025-05-29 21:38:58 -04:00
committed by Sydney Runkle
parent 05f3904d09
commit afb83d2201
15 changed files with 595 additions and 2813 deletions
+1 -1
View File
@@ -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 graphs 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 graphs 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",
-35
View File
@@ -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 -2
View File
@@ -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",
-445
View File
@@ -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)
+153 -23
View File
@@ -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
+1 -1
View File
@@ -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)
-794
View File
@@ -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
+1 -173
View File
@@ -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:
+40 -106
View File
@@ -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
+6 -4
View File
@@ -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
+387 -358
View File
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.