mirror of
https://github.com/langchain-ai/langgraph.git
synced 2026-09-09 19:27:54 +02:00
Compare commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
4dbc8da93d | ||
|
|
85ba90e474 | ||
|
|
3cae2f77ce | ||
|
|
907fefe519 | ||
|
|
9ec5b16107 |
@@ -98,6 +98,15 @@ _CHANNEL_BRANCH_TO = "branch:to:{}"
|
|||||||
_DEFAULT_ERROR_HANDLER_NODE = "__default_error_handler__"
|
_DEFAULT_ERROR_HANDLER_NODE = "__default_error_handler__"
|
||||||
|
|
||||||
|
|
||||||
|
class _CachePolicyUnset:
|
||||||
|
"""Sentinel: `cache_policy` was omitted on `add_node` (inherit defaults)."""
|
||||||
|
|
||||||
|
__slots__ = ()
|
||||||
|
|
||||||
|
|
||||||
|
_CACHE_POLICY_UNSET = _CachePolicyUnset()
|
||||||
|
|
||||||
|
|
||||||
@dataclass(slots=True)
|
@dataclass(slots=True)
|
||||||
class _NodeDefaults:
|
class _NodeDefaults:
|
||||||
"""Default node policies applied to every node at compile time."""
|
"""Default node policies applied to every node at compile time."""
|
||||||
@@ -263,6 +272,7 @@ class StateGraph(Generic[StateT, ContextT, InputT, OutputT]):
|
|||||||
self.context_schema = context_schema
|
self.context_schema = context_schema
|
||||||
|
|
||||||
self._node_defaults: _NodeDefaults = _NodeDefaults()
|
self._node_defaults: _NodeDefaults = _NodeDefaults()
|
||||||
|
self._cache_policy_opt_out: set[str] = set()
|
||||||
|
|
||||||
self._add_schema(self.state_schema)
|
self._add_schema(self.state_schema)
|
||||||
self._add_schema(self.input_schema, allow_managed=False)
|
self._add_schema(self.input_schema, allow_managed=False)
|
||||||
@@ -292,7 +302,8 @@ class StateGraph(Generic[StateT, ContextT, InputT, OutputT]):
|
|||||||
their own via `add_node(..., retry_policy=...)`. Also applies
|
their own via `add_node(..., retry_policy=...)`. Also applies
|
||||||
to error-handler nodes.
|
to error-handler nodes.
|
||||||
cache_policy: Default cache policy for nodes that don't specify
|
cache_policy: Default cache policy for nodes that don't specify
|
||||||
their own via `add_node(..., cache_policy=...)`. Does **not**
|
their own via `add_node(..., cache_policy=...)`. Per-node
|
||||||
|
`cache_policy=None` opts out of this default. Does **not**
|
||||||
apply to error-handler nodes.
|
apply to error-handler nodes.
|
||||||
error_handler: Default error handler invoked when any regular node
|
error_handler: Default error handler invoked when any regular node
|
||||||
raises and does not have its own `error_handler` set via
|
raises and does not have its own `error_handler` set via
|
||||||
@@ -668,7 +679,7 @@ class StateGraph(Generic[StateT, ContextT, InputT, OutputT]):
|
|||||||
metadata: dict[str, Any] | None = None,
|
metadata: dict[str, Any] | None = None,
|
||||||
input_schema: type[NodeInputT] | None = None,
|
input_schema: type[NodeInputT] | None = None,
|
||||||
retry_policy: RetryPolicy | Sequence[RetryPolicy] | None = None,
|
retry_policy: RetryPolicy | Sequence[RetryPolicy] | None = None,
|
||||||
cache_policy: CachePolicy | None = None,
|
cache_policy: CachePolicy | None | _CachePolicyUnset = _CACHE_POLICY_UNSET,
|
||||||
error_handler: StateNode[Any, ContextT] | None = None,
|
error_handler: StateNode[Any, ContextT] | None = None,
|
||||||
destinations: dict[str, str] | tuple[str, ...] | None = None,
|
destinations: dict[str, str] | tuple[str, ...] | None = None,
|
||||||
timeout: float | timedelta | TimeoutPolicy | None = None,
|
timeout: float | timedelta | TimeoutPolicy | None = None,
|
||||||
@@ -689,7 +700,10 @@ class StateGraph(Generic[StateT, ContextT, InputT, OutputT]):
|
|||||||
retry_policy: The retry policy for the node.
|
retry_policy: The retry policy for the node.
|
||||||
|
|
||||||
If a sequence is provided, the first matching policy will be applied.
|
If a sequence is provided, the first matching policy will be applied.
|
||||||
cache_policy: The cache policy for the node.
|
cache_policy: The cache policy for the node. Pass `None` explicitly to
|
||||||
|
disable caching for this node, including when graph defaults from
|
||||||
|
`set_node_defaults` would otherwise enable it. Omit the argument to
|
||||||
|
inherit graph defaults.
|
||||||
error_handler: Optional node-level error handler callable for this node.
|
error_handler: Optional node-level error handler callable for this node.
|
||||||
destinations: Destinations that indicate where a node can route to.
|
destinations: Destinations that indicate where a node can route to.
|
||||||
|
|
||||||
@@ -765,6 +779,15 @@ class StateGraph(Generic[StateT, ContextT, InputT, OutputT]):
|
|||||||
input_schema = cast(type[NodeInputT] | None, input_)
|
input_schema = cast(type[NodeInputT] | None, input_)
|
||||||
timeout = coerce_timeout_policy(timeout)
|
timeout = coerce_timeout_policy(timeout)
|
||||||
|
|
||||||
|
explicit_cache_opt_out = False
|
||||||
|
if isinstance(cache_policy, _CachePolicyUnset):
|
||||||
|
resolved_cache_policy: CachePolicy | None = None
|
||||||
|
elif isinstance(cache_policy, CachePolicy):
|
||||||
|
resolved_cache_policy = cache_policy
|
||||||
|
else:
|
||||||
|
explicit_cache_opt_out = True
|
||||||
|
resolved_cache_policy = None
|
||||||
|
|
||||||
if not isinstance(node, str):
|
if not isinstance(node, str):
|
||||||
action = node
|
action = node
|
||||||
if isinstance(action, Runnable):
|
if isinstance(action, Runnable):
|
||||||
@@ -789,6 +812,10 @@ class StateGraph(Generic[StateT, ContextT, InputT, OutputT]):
|
|||||||
)
|
)
|
||||||
if action is None:
|
if action is None:
|
||||||
raise RuntimeError
|
raise RuntimeError
|
||||||
|
if explicit_cache_opt_out:
|
||||||
|
self._cache_policy_opt_out.add(node)
|
||||||
|
else:
|
||||||
|
self._cache_policy_opt_out.discard(node)
|
||||||
if node in self.nodes:
|
if node in self.nodes:
|
||||||
raise ValueError(f"Node `{node}` already present.")
|
raise ValueError(f"Node `{node}` already present.")
|
||||||
if node == END or node == START:
|
if node == END or node == START:
|
||||||
@@ -875,7 +902,7 @@ class StateGraph(Generic[StateT, ContextT, InputT, OutputT]):
|
|||||||
metadata,
|
metadata,
|
||||||
input_schema=input_schema,
|
input_schema=input_schema,
|
||||||
retry_policy=retry_policy,
|
retry_policy=retry_policy,
|
||||||
cache_policy=cache_policy,
|
cache_policy=resolved_cache_policy,
|
||||||
error_handler_node=handler_node_name,
|
error_handler_node=handler_node_name,
|
||||||
ends=ends,
|
ends=ends,
|
||||||
defer=defer,
|
defer=defer,
|
||||||
@@ -887,7 +914,7 @@ class StateGraph(Generic[StateT, ContextT, InputT, OutputT]):
|
|||||||
metadata,
|
metadata,
|
||||||
input_schema=inferred_input_schema,
|
input_schema=inferred_input_schema,
|
||||||
retry_policy=retry_policy,
|
retry_policy=retry_policy,
|
||||||
cache_policy=cache_policy,
|
cache_policy=resolved_cache_policy,
|
||||||
error_handler_node=handler_node_name,
|
error_handler_node=handler_node_name,
|
||||||
ends=ends,
|
ends=ends,
|
||||||
defer=defer,
|
defer=defer,
|
||||||
@@ -899,7 +926,7 @@ class StateGraph(Generic[StateT, ContextT, InputT, OutputT]):
|
|||||||
metadata,
|
metadata,
|
||||||
input_schema=self.state_schema,
|
input_schema=self.state_schema,
|
||||||
retry_policy=retry_policy,
|
retry_policy=retry_policy,
|
||||||
cache_policy=cache_policy,
|
cache_policy=resolved_cache_policy,
|
||||||
error_handler_node=handler_node_name,
|
error_handler_node=handler_node_name,
|
||||||
ends=ends,
|
ends=ends,
|
||||||
defer=defer,
|
defer=defer,
|
||||||
@@ -1297,7 +1324,7 @@ class StateGraph(Generic[StateT, ContextT, InputT, OutputT]):
|
|||||||
)
|
)
|
||||||
|
|
||||||
# Apply builder defaults to node specs. Per-node values always win.
|
# Apply builder defaults to node specs. Per-node values always win.
|
||||||
for spec in self.nodes.values():
|
for node_name, spec in self.nodes.items():
|
||||||
# error_handler: regular nodes only — handlers must never
|
# error_handler: regular nodes only — handlers must never
|
||||||
# catch themselves or other handlers.
|
# catch themselves or other handlers.
|
||||||
if (
|
if (
|
||||||
@@ -1317,6 +1344,7 @@ class StateGraph(Generic[StateT, ContextT, InputT, OutputT]):
|
|||||||
not spec.is_error_handler
|
not spec.is_error_handler
|
||||||
and defaults.cache_policy is not None
|
and defaults.cache_policy is not None
|
||||||
and spec.cache_policy is None
|
and spec.cache_policy is None
|
||||||
|
and node_name not in self._cache_policy_opt_out
|
||||||
):
|
):
|
||||||
spec.cache_policy = defaults.cache_policy
|
spec.cache_policy = defaults.cache_policy
|
||||||
# timeout: all nodes — a stuck handler should be cancelled the
|
# timeout: all nodes — a stuck handler should be cancelled the
|
||||||
|
|||||||
@@ -58,6 +58,7 @@ from langgraph.pregel._retry import (
|
|||||||
from langgraph.pregel.protocol import StreamProtocol
|
from langgraph.pregel.protocol import StreamProtocol
|
||||||
from langgraph.runtime import DEFAULT_RUNTIME, ExecutionInfo, Runtime
|
from langgraph.runtime import DEFAULT_RUNTIME, ExecutionInfo, Runtime
|
||||||
from langgraph.types import (
|
from langgraph.types import (
|
||||||
|
CachePolicy,
|
||||||
Command,
|
Command,
|
||||||
PregelExecutableTask,
|
PregelExecutableTask,
|
||||||
RetryPolicy,
|
RetryPolicy,
|
||||||
@@ -2570,6 +2571,92 @@ def test_set_node_defaults_retry_policy_per_node_wins():
|
|||||||
assert attempts == 2
|
assert attempts == 2
|
||||||
|
|
||||||
|
|
||||||
|
def test_set_node_defaults_cache_policy_applies_to_nodes_without_own() -> None:
|
||||||
|
from langgraph.cache.memory import InMemoryCache
|
||||||
|
|
||||||
|
class State(TypedDict):
|
||||||
|
foo: str
|
||||||
|
|
||||||
|
calls = 0
|
||||||
|
|
||||||
|
def cached_node(state: State) -> State:
|
||||||
|
nonlocal calls
|
||||||
|
calls += 1
|
||||||
|
return {"foo": "ok"}
|
||||||
|
|
||||||
|
graph = (
|
||||||
|
StateGraph(State)
|
||||||
|
.set_node_defaults(cache_policy=CachePolicy(ttl=60))
|
||||||
|
.add_node("cached", cached_node)
|
||||||
|
.add_edge(START, "cached")
|
||||||
|
.compile(cache=InMemoryCache())
|
||||||
|
)
|
||||||
|
|
||||||
|
assert graph.nodes["cached"].cache_policy == CachePolicy(ttl=60)
|
||||||
|
|
||||||
|
graph.invoke({"foo": ""})
|
||||||
|
graph.invoke({"foo": ""})
|
||||||
|
assert calls == 1
|
||||||
|
|
||||||
|
|
||||||
|
def test_set_node_defaults_cache_policy_per_node_wins() -> None:
|
||||||
|
from langgraph.cache.memory import InMemoryCache
|
||||||
|
|
||||||
|
class State(TypedDict):
|
||||||
|
foo: str
|
||||||
|
|
||||||
|
graph = (
|
||||||
|
StateGraph(State)
|
||||||
|
.set_node_defaults(cache_policy=CachePolicy(ttl=60))
|
||||||
|
.add_node("a", lambda state: {"foo": "a"}, cache_policy=CachePolicy(ttl=5))
|
||||||
|
.add_node("b", lambda state: {"foo": "b"})
|
||||||
|
.add_edge(START, "a")
|
||||||
|
.add_edge("a", "b")
|
||||||
|
.compile(cache=InMemoryCache())
|
||||||
|
)
|
||||||
|
|
||||||
|
assert graph.nodes["a"].cache_policy == CachePolicy(ttl=5)
|
||||||
|
assert graph.nodes["b"].cache_policy == CachePolicy(ttl=60)
|
||||||
|
|
||||||
|
|
||||||
|
def test_set_node_defaults_cache_policy_explicit_none_opts_out() -> None:
|
||||||
|
from langgraph.cache.memory import InMemoryCache
|
||||||
|
|
||||||
|
class State(TypedDict):
|
||||||
|
foo: str
|
||||||
|
|
||||||
|
cached_calls = 0
|
||||||
|
uncached_calls = 0
|
||||||
|
|
||||||
|
def cached_node(state: State) -> State:
|
||||||
|
nonlocal cached_calls
|
||||||
|
cached_calls += 1
|
||||||
|
return {"foo": "cached"}
|
||||||
|
|
||||||
|
def uncached_node(state: State) -> State:
|
||||||
|
nonlocal uncached_calls
|
||||||
|
uncached_calls += 1
|
||||||
|
return {"foo": "uncached"}
|
||||||
|
|
||||||
|
graph = (
|
||||||
|
StateGraph(State)
|
||||||
|
.set_node_defaults(cache_policy=CachePolicy())
|
||||||
|
.add_node("cached", cached_node)
|
||||||
|
.add_node("uncached", uncached_node, cache_policy=None)
|
||||||
|
.add_edge(START, "cached")
|
||||||
|
.add_edge("cached", "uncached")
|
||||||
|
.compile(cache=InMemoryCache())
|
||||||
|
)
|
||||||
|
|
||||||
|
assert graph.nodes["cached"].cache_policy == CachePolicy()
|
||||||
|
assert graph.nodes["uncached"].cache_policy is None
|
||||||
|
|
||||||
|
graph.invoke({"foo": ""})
|
||||||
|
graph.invoke({"foo": ""})
|
||||||
|
assert cached_calls == 1
|
||||||
|
assert uncached_calls == 2
|
||||||
|
|
||||||
|
|
||||||
@pytest.mark.anyio
|
@pytest.mark.anyio
|
||||||
async def test_set_node_defaults_timeout():
|
async def test_set_node_defaults_timeout():
|
||||||
class State(TypedDict):
|
class State(TypedDict):
|
||||||
|
|||||||
Reference in New Issue
Block a user