This commit is contained in:
Nuno Campos
2024-09-19 13:01:54 -07:00
parent 4f767cd2ca
commit 5de9b35416
20 changed files with 171 additions and 121 deletions
@@ -13,11 +13,11 @@ from typing import (
Sequence,
Tuple,
TypedDict,
TypeVar,
Union,
)
from langchain_core.runnables import ConfigurableFieldSpec, RunnableConfig
from typing_extensions import TypeVar
from langgraph.checkpoint.base.id import uuid6
from langgraph.checkpoint.serde.base import SerializerProtocol, maybe_add_typed_methods
@@ -29,7 +29,7 @@ from langgraph.checkpoint.serde.types import (
SendProtocol,
)
V = TypeVar("V", int, float, str, default=int)
V = TypeVar("V", int, float, str)
PendingWrite = Tuple[str, str, Any]
+1 -1
View File
@@ -75,7 +75,7 @@ lint lint_diff lint_package lint_tests:
[ "$(PYTHON_FILES)" = "" ] || poetry run ruff format $(PYTHON_FILES) --diff
[ "$(PYTHON_FILES)" = "" ] || poetry run ruff check --select I $(PYTHON_FILES)
[ "$(PYTHON_FILES)" = "" ] || mkdir -p $(MYPY_CACHE)
[ "$(PYTHON_FILES)" = "" ] || poetry run mypy $(PYTHON_FILES) --cache-dir $(MYPY_CACHE)
[ "$(PYTHON_FILES)" != "langgraph" ] || poetry run mypy $(PYTHON_FILES) --cache-dir $(MYPY_CACHE)
format format_diff:
poetry run ruff format $(PYTHON_FILES)
@@ -6,7 +6,7 @@ from langgraph.channels.base import BaseChannel, Value
from langgraph.errors import EmptyChannelError, InvalidUpdateError
class WaitForNames(NamedTuple):
class WaitForNames(NamedTuple, Generic[Value]):
names: set[Value]
@@ -1,4 +1,4 @@
from typing import Generic, Optional, Sequence, Type
from typing import Any, Generic, Optional, Sequence, Type
from typing_extensions import Self
@@ -11,7 +11,7 @@ class EphemeralValue(Generic[Value], BaseChannel[Value, Value, Value]):
__slots__ = ("value", "guard")
def __init__(self, typ: Type[Value], guard: bool = True) -> None:
def __init__(self, typ: Any, guard: bool = True) -> None:
super().__init__(typ)
self.guard = guard
+49 -30
View File
@@ -57,7 +57,7 @@ class Branch(NamedTuple):
def run(
self,
writer: Callable[
[list[Union[str, Send]], RunnableConfig], Optional[ChannelWrite]
[Sequence[Union[str, Send]], RunnableConfig], Optional[ChannelWrite]
],
reader: Optional[Callable[[RunnableConfig], Any]] = None,
) -> RunnableCallable:
@@ -79,7 +79,7 @@ class Branch(NamedTuple):
*,
reader: Optional[Callable[[RunnableConfig], Any]],
writer: Callable[
[list[Union[str, Send]], RunnableConfig], Optional[ChannelWrite]
[Sequence[Union[str, Send]], RunnableConfig], Optional[ChannelWrite]
],
) -> Runnable:
if reader:
@@ -100,7 +100,7 @@ class Branch(NamedTuple):
*,
reader: Optional[Callable[[RunnableConfig], Any]],
writer: Callable[
[list[Union[str, Send]], RunnableConfig], Optional[ChannelWrite]
[Sequence[Union[str, Send]], RunnableConfig], Optional[ChannelWrite]
],
) -> Runnable:
if reader:
@@ -117,18 +117,20 @@ class Branch(NamedTuple):
def _finish(
self,
writer: Callable[
[list[Union[str, Send]], RunnableConfig], Optional[ChannelWrite]
[Sequence[Union[str, Send]], RunnableConfig], Optional[ChannelWrite]
],
input: Any,
result: Any,
config: RunnableConfig,
):
) -> Union[Runnable, Any]:
if not isinstance(result, list):
result = [result]
if self.ends:
destinations = [r if isinstance(r, Send) else self.ends[r] for r in result]
destinations: Sequence[Union[Send, str]] = [
r if isinstance(r, Send) else self.ends[r] for r in result
]
else:
destinations = result
destinations = cast(Sequence[Union[Send, str]], result)
if any(dest is None or dest == START for dest in destinations):
raise ValueError("Branch did not return a valid destination")
if any(p.node == END for p in destinations if isinstance(p, Send)):
@@ -186,14 +188,20 @@ class Graph:
)
if not isinstance(node, str):
action = node
node = getattr(action, "name", action.__name__)
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
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[node] = NodeSpec(
coerce_to_runnable(action, name=node, trace=False), metadata
self.nodes[cast(str, node)] = NodeSpec(
coerce_to_runnable(action, name=cast(str, node), trace=False), metadata
)
def add_edge(self, start_key: str, end_key: str) -> None:
@@ -257,16 +265,20 @@ class Graph:
# coerce path_map to a dictionary
try:
if isinstance(path_map, dict):
path_map = path_map.copy()
path_map_ = path_map.copy()
elif isinstance(path_map, list):
path_map = {name: name for name in path_map}
elif rtn_type := get_type_hints(path.__call__).get(
"return"
) or get_type_hints(path).get("return"):
path_map_ = {name: name for name in path_map}
elif callable(path) and (
rtn_type := get_type_hints(path.__call__).get("return")
if hasattr(path, "__call__")
else get_type_hints(path).get("return")
):
if get_origin(rtn_type) is Literal:
path_map = {name: name for name in get_args(rtn_type)}
path_map_ = {name: name for name in get_args(rtn_type)}
else:
path_map_ = None
except Exception:
pass
path_map_ = None
# find a name for the condition
path = coerce_to_runnable(path, name=None, trace=True)
name = path.name or "condition"
@@ -276,7 +288,7 @@ class Graph:
f"Branch with name `{path.name}` already exists for node " f"`{source}`"
)
# save it
self.branches[source][name] = Branch(path, path_map, then)
self.branches[source][name] = Branch(path, path_map_, then)
def set_entry_point(self, key: str) -> None:
"""Specifies the first node to be called in the graph.
@@ -405,7 +417,6 @@ class Graph:
# create empty compiled graph
compiled = CompiledGraph(
builder=self,
nodes={},
channels={START: EphemeralValue(Any), END: EphemeralValue(Any)},
input_channels=START,
@@ -418,6 +429,7 @@ class Graph:
auto_validate=False,
debug=debug,
)
compiled.builder = self
# attach nodes, edges, and branches
for key, node in self.nodes.items():
@@ -437,10 +449,6 @@ class Graph:
class CompiledGraph(Pregel):
builder: Graph
def __init__(self, *, builder: Graph, **kwargs):
super().__init__(**kwargs)
self.builder = builder
def attach_node(self, key: str, node: NodeSpec) -> None:
self.channels[key] = EphemeralValue(Any)
self.nodes[key] = (
@@ -463,7 +471,7 @@ class CompiledGraph(Pregel):
def attach_branch(self, start: str, name: str, branch: Branch) -> None:
def branch_writer(
packets: list[Union[str, Send]], config: RunnableConfig
packets: Sequence[Union[str, Send]], config: RunnableConfig
) -> Optional[ChannelWrite]:
writes = [
(
@@ -473,7 +481,10 @@ class CompiledGraph(Pregel):
)
for p in packets
]
return ChannelWrite(writes, tags=[TAG_HIDDEN])
return ChannelWrite(
cast(Sequence[Union[ChannelWriteEntry, Send]], writes),
tags=[TAG_HIDDEN],
)
# add hidden start node
if start == START and start not in self.nodes:
@@ -489,7 +500,7 @@ class CompiledGraph(Pregel):
channel_name = f"branch:{start}:{name}:{end}"
self.channels[channel_name] = EphemeralValue(Any)
self.nodes[end].triggers.append(channel_name)
self.nodes[end].channels.append(channel_name)
cast(list[str], self.nodes[end].channels).append(channel_name)
def get_graph(
self,
@@ -504,17 +515,25 @@ class CompiledGraph(Pregel):
}
end_nodes: dict[str, DrawableNode] = {}
if xray:
subgraphs = dict(self.get_subgraphs())
subgraphs = {
k: v for k, v in self.get_subgraphs() if isinstance(v, CompiledGraph)
}
else:
subgraphs = {}
def add_edge(
start: str, end: str, label: Optional[str] = None, conditional: bool = False
start: str,
end: str,
label: Optional[Hashable] = None,
conditional: bool = False,
) -> None:
if end == END and END not in end_nodes:
end_nodes[END] = graph.add_node(self.get_output_schema(config), END)
return graph.add_edge(
start_nodes[start], end_nodes[end], label, conditional
start_nodes[start],
end_nodes[end],
str(label) if label is not None else None,
conditional,
)
for key, n in self.builder.nodes.items():
@@ -563,7 +582,7 @@ class CompiledGraph(Pregel):
elif branch.then is not None:
ends = {k: k for k in default_ends if k not in (END, branch.then)}
else:
ends = default_ends
ends = cast(dict[Hashable, str], default_ends)
for label, end in ends.items():
add_edge(
start,
+2 -2
View File
@@ -63,9 +63,9 @@ def add_messages(left: Messages, right: Messages) -> Messages:
"""
# coerce to list
if not isinstance(left, list):
left = [left]
left = [left] # type: ignore[assignment]
if not isinstance(right, list):
right = [right]
right = [right] # type: ignore[assignment]
# coerce to message
left = [
message_chunk_to_message(cast(BaseMessageChunk, m))
+28 -10
View File
@@ -7,11 +7,13 @@ from inspect import isclass, isfunction, signature
from typing import (
Any,
Callable,
Literal,
NamedTuple,
Optional,
Sequence,
Type,
Union,
cast,
get_origin,
get_type_hints,
overload,
@@ -122,7 +124,7 @@ class StateGraph(Graph):
>>> print(step1)
{'x': [0.5, 0.75]}"""
nodes: dict[str, StateNodeSpec]
nodes: dict[str, StateNodeSpec] # type: ignore[assignment]
channels: dict[str, BaseChannel]
managed: dict[str, ManagedValueSpec]
schemas: dict[Type[Any], dict[str, Union[BaseChannel, ManagedValueSpec]]]
@@ -302,7 +304,7 @@ class StateGraph(Graph):
if not isinstance(node, str):
action = node
if isinstance(action, Runnable):
node = action.name
node = action.get_name()
else:
node = getattr(action, "__name__", action.__class__.__name__)
if node is None:
@@ -323,13 +325,15 @@ class StateGraph(Graph):
raise ValueError(
"Node name must be provided if action is not a function"
)
if action is None:
raise RuntimeError
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.")
for character in (NS_SEP, NS_END):
if character in node:
if character in cast(str, node):
raise ValueError(
f"'{character}' is a reserved character and is not allowed in the node names."
)
@@ -349,8 +353,8 @@ class StateGraph(Graph):
pass
if input is not None:
self._add_schema(input)
self.nodes[node] = StateNodeSpec(
coerce_to_runnable(action, name=node, trace=False),
self.nodes[cast(str, node)] = StateNodeSpec(
coerce_to_runnable(action, name=cast(str, node), trace=False),
metadata,
input=input or self.schema,
retry_policy=retry,
@@ -449,7 +453,6 @@ class StateGraph(Graph):
)
compiled = CompiledStateGraph(
builder=self,
config_type=self.config_schema,
nodes={},
channels={
@@ -468,6 +471,7 @@ class StateGraph(Graph):
debug=debug,
store=store,
)
compiled.builder = self
compiled.attach_node(START, None)
for key, node in self.nodes.items():
@@ -618,7 +622,7 @@ class CompiledStateGraph(CompiledGraph):
def attach_branch(self, start: str, name: str, branch: Branch) -> None:
def branch_writer(
packets: list[Union[str, Send]], config: RunnableConfig
packets: Sequence[Union[str, Send]], config: RunnableConfig
) -> None:
if filtered := [p for p in packets if p != END]:
writes = [
@@ -638,7 +642,9 @@ class CompiledStateGraph(CompiledGraph):
),
)
)
ChannelWrite.do_write(config, writes)
ChannelWrite.do_write(
config, cast(Sequence[Union[Send, ChannelWriteEntry]], writes)
)
# attach branch publisher
schema = (
@@ -708,11 +714,23 @@ def _get_channels(
if name != "__slots__"
}
return (
{k: v for k, v in all_keys.items() if not is_managed_value(v)},
{k: v for k, v in all_keys.items() if isinstance(v, BaseChannel)},
{k: v for k, v in all_keys.items() if is_managed_value(v)},
)
@overload
def _get_channel(
name: str, annotation: Any, *, allow_managed: Literal[False]
) -> BaseChannel: ...
@overload
def _get_channel(
name: str, annotation: Any, *, allow_managed: Literal[True] = True
) -> Union[BaseChannel, ManagedValueSpec]: ...
def _get_channel(
name: str, annotation: Any, *, allow_managed: bool = True
) -> Union[BaseChannel, ManagedValueSpec]:
@@ -728,7 +746,7 @@ def _get_channel(
channel.key = name
return channel
fallback = LastValue(annotation)
fallback: LastValue = LastValue(annotation)
fallback.key = name
return fallback
@@ -419,7 +419,7 @@ def create_react_agent(
raise ValueError(f"Missing required key(s) {missing_keys} in state_schema")
if isinstance(tools, ToolExecutor):
tool_classes = tools.tools
tool_classes: Sequence[BaseTool] = tools.tools
tool_node = ToolNode(tool_classes)
elif isinstance(tools, ToolNode):
tool_classes = list(tools.tools_by_name.values())
@@ -382,7 +382,7 @@ def _get_state_args(tool: BaseTool) -> Dict[str, Optional[str]]:
full_schema = tool.get_input_schema()
tool_args_to_state_fields: Dict = {}
def _is_injection(type_arg: Any):
def _is_injection(type_arg: Any) -> bool:
if isinstance(type_arg, InjectedState) or (
isinstance(type_arg, type) and issubclass(type_arg, InjectedState)
):
+24 -20
View File
@@ -138,7 +138,7 @@ class Channel:
)
return PregelNode(
channels=cast(
Union[Mapping[None, str], Mapping[str, str]],
Union[list[str], Mapping[str, str]],
(
{key: channels}
if isinstance(channels, str) and key is not None
@@ -305,7 +305,9 @@ class Pregel(Runnable[Union[dict[str, Any], Any], Union[dict[str, Any], Any]]):
@property
def InputType(self) -> Any:
if isinstance(self.input_channels, str):
return self.channels[self.input_channels].UpdateType
channel = self.channels[self.input_channels]
if isinstance(channel, BaseChannel):
return channel.UpdateType
def get_input_schema(
self, config: Optional[RunnableConfig] = None
@@ -317,9 +319,9 @@ class Pregel(Runnable[Union[dict[str, Any], Any], Union[dict[str, Any], Any]]):
return create_model(
self.get_name("Input"),
field_definitions={
k: (self.channels[k].UpdateType, None)
k: (c.UpdateType, None)
for k in self.input_channels or self.channels.keys()
if isinstance(self.channels[k], BaseChannel)
if (c := self.channels[k]) and isinstance(c, BaseChannel)
},
)
@@ -335,7 +337,9 @@ class Pregel(Runnable[Union[dict[str, Any], Any], Union[dict[str, Any], Any]]):
@property
def OutputType(self) -> Any:
if isinstance(self.output_channels, str):
return self.channels[self.output_channels].ValueType
channel = self.channels[self.output_channels]
if isinstance(channel, BaseChannel):
return channel.ValueType
def get_output_schema(
self, config: Optional[RunnableConfig] = None
@@ -347,9 +351,9 @@ class Pregel(Runnable[Union[dict[str, Any], Any], Union[dict[str, Any], Any]]):
return create_model(
self.get_name("Output"),
field_definitions={
k: (self.channels[k].ValueType, None)
k: (c.ValueType, None)
for k in self.output_channels
if isinstance(self.channels[k], BaseChannel)
if (c := self.channels[k]) and isinstance(c, BaseChannel)
},
)
@@ -1050,8 +1054,8 @@ class Pregel(Runnable[Union[dict[str, Any], Any], Union[dict[str, Any], Any]]):
bool,
set[StreamMode],
Union[str, Sequence[str]],
Optional[Sequence[str]],
Optional[Sequence[str]],
Union[All, Sequence[str]],
Union[All, Sequence[str]],
Optional[BaseCheckpointSaver],
]:
debug = debug if debug is not None else self.debug
@@ -1199,8 +1203,8 @@ class Pregel(Runnable[Union[dict[str, Any], Any], Union[dict[str, Any], Any]]):
debug,
stream_modes,
output_keys,
interrupt_before,
interrupt_after,
interrupt_before_,
interrupt_after_,
checkpointer,
) = self._defaults(
config,
@@ -1253,7 +1257,7 @@ class Pregel(Runnable[Union[dict[str, Any], Any], Union[dict[str, Any], Any]]):
else:
return waiter
else:
get_waiter = None
get_waiter = None # type: ignore[assignment]
# Similarly to Bulk Synchronous Parallel / Pregel model
# computation proceeds in steps, while there are channel updates
# channel updates from step N are only visible in step N+1
@@ -1261,8 +1265,8 @@ class Pregel(Runnable[Union[dict[str, Any], Any], Union[dict[str, Any], Any]]):
# with channel updates applied only at the transition between steps
while loop.tick(
input_keys=self.input_channels,
interrupt_before=interrupt_before,
interrupt_after=interrupt_after,
interrupt_before=interrupt_before_,
interrupt_after=interrupt_after_,
manager=run_manager,
):
for _ in runner.tick(
@@ -1397,7 +1401,7 @@ class Pregel(Runnable[Union[dict[str, Any], Any], Union[dict[str, Any], Any]]):
# if running from astream_log() run each proc with streaming
do_stream = next(
(
h
cast(_StreamingCallbackHandler, h)
for h in run_manager.handlers
if isinstance(h, _StreamingCallbackHandler)
),
@@ -1415,8 +1419,8 @@ class Pregel(Runnable[Union[dict[str, Any], Any], Union[dict[str, Any], Any]]):
debug,
stream_modes,
output_keys,
interrupt_before,
interrupt_after,
interrupt_before_,
interrupt_after_,
checkpointer,
) = self._defaults(
config,
@@ -1457,7 +1461,7 @@ class Pregel(Runnable[Union[dict[str, Any], Any], Union[dict[str, Any], Any]]):
def get_waiter() -> asyncio.Task[None]:
return aioloop.create_task(stream.wait())
else:
get_waiter = None
get_waiter = None # type: ignore[assignment]
# Similarly to Bulk Synchronous Parallel / Pregel model
# computation proceeds in steps, while there are channel updates
# channel updates from step N are only visible in step N+1
@@ -1465,8 +1469,8 @@ class Pregel(Runnable[Union[dict[str, Any], Any], Union[dict[str, Any], Any]]):
# with channel updates applied only at the transition between steps
while loop.tick(
input_keys=self.input_channels,
interrupt_before=interrupt_before,
interrupt_after=interrupt_after,
interrupt_before=interrupt_before_,
interrupt_after=interrupt_after_,
manager=run_manager,
):
async for _ in runner.atick(
+6 -4
View File
@@ -52,6 +52,8 @@ from langgraph.pregel.read import PregelNode
from langgraph.pregel.types import All, PregelExecutableTask, PregelTask
from langgraph.utils.config import merge_configs, patch_config
GetNextVersion = Callable[[Optional[V], BaseChannel], V]
EMPTY_SEQ: tuple[str, ...] = tuple()
@@ -173,7 +175,7 @@ def apply_writes(
checkpoint: Checkpoint,
channels: Mapping[str, BaseChannel],
tasks: Iterable[WritesProtocol],
get_next_version: Optional[Callable[[Optional[V], BaseChannel], V]],
get_next_version: Optional[GetNextVersion],
) -> dict[str, list[Any]]:
# update seen versions
for task in tasks:
@@ -200,7 +202,7 @@ def apply_writes(
}:
if channels[chan].consume() and get_next_version is not None:
checkpoint["channel_versions"][chan] = get_next_version(
max_version, # type: ignore[arg-type]
max_version,
channels[chan],
)
@@ -234,7 +236,7 @@ def apply_writes(
if chan in channels:
if channels[chan].update(vals) and get_next_version is not None:
checkpoint["channel_versions"][chan] = get_next_version(
max_version, # type: ignore[arg-type]
max_version,
channels[chan],
)
updated_channels.add(chan)
@@ -244,7 +246,7 @@ def apply_writes(
if chan not in updated_channels:
if channels[chan].update([]) and get_next_version is not None:
checkpoint["channel_versions"][chan] = get_next_version(
max_version, # type: ignore[arg-type]
max_version,
channels[chan],
)
+2 -2
View File
@@ -82,7 +82,7 @@ TASK_NAMESPACE = UUID("6ba7b831-9dad-11d1-80b4-00c04fd430c8")
def map_debug_tasks(
step: int, tasks: list[PregelExecutableTask]
step: int, tasks: Iterable[PregelExecutableTask]
) -> Iterator[DebugOutputTask]:
ts = datetime.now(timezone.utc).isoformat()
for task in tasks:
@@ -132,7 +132,7 @@ def map_debug_checkpoint(
stream_channels: Union[str, Sequence[str]],
metadata: CheckpointMetadata,
checkpoint: Checkpoint,
tasks: list[PregelExecutableTask],
tasks: Iterable[PregelExecutableTask],
pending_writes: list[PendingWrite],
) -> Iterator[DebugOutputCheckpoint]:
yield {
+14 -8
View File
@@ -64,6 +64,7 @@ from langgraph.managed.base import (
WritableManagedValue,
)
from langgraph.pregel.algo import (
GetNextVersion,
PregelTaskWrites,
apply_writes,
increment,
@@ -92,7 +93,7 @@ from langgraph.pregel.io import (
)
from langgraph.pregel.manager import AsyncChannelsManager, ChannelsManager
from langgraph.pregel.read import PregelNode
from langgraph.pregel.types import PregelExecutableTask, StreamMode
from langgraph.pregel.types import All, PregelExecutableTask, StreamMode
from langgraph.pregel.utils import get_new_channel_versions
from langgraph.store.base import BaseStore
from langgraph.store.batch import AsyncBatchedStore
@@ -146,7 +147,7 @@ class PregelLoop:
skip_done_tasks: bool
is_nested: bool
checkpointer_get_next_version: Callable[[Optional[V]], V]
checkpointer_get_next_version: GetNextVersion
checkpointer_put_writes: Optional[
Callable[[RunnableConfig, Sequence[tuple[str, Any]], str], Any]
]
@@ -281,8 +282,8 @@ class PregelLoop:
self,
*,
input_keys: Union[str, Sequence[str]],
interrupt_after: Sequence[str] = EMPTY_SEQ,
interrupt_before: Sequence[str] = EMPTY_SEQ,
interrupt_after: Union[All, Sequence[str]] = EMPTY_SEQ,
interrupt_before: Union[All, Sequence[str]] = EMPTY_SEQ,
manager: Union[None, AsyncParentRunManager, ParentRunManager] = None,
) -> bool:
"""Execute a single iteration of the Pregel loop.
@@ -681,6 +682,10 @@ class SyncPregelLoop(PregelLoop, ContextManager):
if self.config.get("configurable", {}).get(
CONFIG_KEY_ENSURE_LATEST
) and self.checkpoint_config["configurable"].get("checkpoint_id"):
if self.checkpointer is None:
raise RuntimeError(
"Cannot ensure latest checkpoint without checkpointer"
)
saved = self.checkpointer.get_tuple(
patch_configurable(self.checkpoint_config, {"checkpoint_id": None})
)
@@ -771,7 +776,7 @@ class AsyncPregelLoop(PregelLoop, AsyncContextManager):
self.checkpointer_put_writes = checkpointer.aput_writes
else:
self.checkpointer_get_next_version = increment
self._checkpointer_put_after_previous = None # type: ignore[method-assign]
self._checkpointer_put_after_previous = None # type: ignore[assignment]
self.checkpointer_put_writes = None
async def _checkpointer_put_after_previous(
@@ -801,6 +806,10 @@ class AsyncPregelLoop(PregelLoop, AsyncContextManager):
if self.config.get("configurable", {}).get(
CONFIG_KEY_ENSURE_LATEST
) and self.checkpoint_config["configurable"].get("checkpoint_id"):
if self.checkpointer is None:
raise RuntimeError(
"Cannot ensure latest checkpoint without checkpointer"
)
saved = await self.checkpointer.aget_tuple(
patch_configurable(self.checkpoint_config, {"checkpoint_id": None})
)
@@ -858,6 +867,3 @@ class AsyncPregelLoop(PregelLoop, AsyncContextManager):
return await asyncio.shield(
self.stack.__aexit__(exc_type, exc_value, traceback)
)
EMPTY_SEQ = tuple()
+4 -9
View File
@@ -7,7 +7,6 @@ from typing import (
List,
Optional,
Sequence,
Tuple,
Union,
cast,
)
@@ -21,20 +20,16 @@ from langchain_core.tracers._streaming import T, _StreamingCallbackHandler
from langgraph.constants import NS_SEP
from langgraph.pregel.loop import StreamChunk
Meta = tuple[tuple[str, ...], dict[str, Any]]
class StreamMessagesHandler(BaseCallbackHandler, _StreamingCallbackHandler):
def __init__(self, stream: Callable[[StreamChunk], None]):
self.stream = stream
self.metadata: dict[UUID, tuple[tuple[str, ...], dict[str, Any]]] = {}
self.metadata: dict[UUID, Meta] = {}
self.seen: set[Union[int, str]] = set()
def _emit(
self,
meta: Tuple[str, dict[str, Any]],
message: BaseMessage,
*,
dedupe: bool = False,
) -> None:
def _emit(self, meta: Meta, message: BaseMessage, *, dedupe: bool = False) -> None:
ident = id(message)
if dedupe and message.id in self.seen:
return
+6 -6
View File
@@ -18,7 +18,7 @@ from langchain_core.runnables import (
RunnablePassthrough,
RunnableSerializable,
)
from langchain_core.runnables.base import Input, Other, Output, coerce_to_runnable
from langchain_core.runnables.base import Input, Other, coerce_to_runnable
from langchain_core.runnables.utils import ConfigurableFieldSpec
from langgraph.constants import CONFIG_KEY_READ
@@ -206,7 +206,7 @@ class PregelNode(Runnable):
Mapping[str, Runnable[Any, Other] | Callable[[Any], Other]],
],
) -> PregelNode:
if ChannelWrite.is_writer(other):
if isinstance(other, Runnable) and ChannelWrite.is_writer(other):
return self.copy(update=dict(writers=[*self.writers, other]))
elif self.bound is DEFAULT_BOUND:
return self.copy(update=dict(bound=coerce_to_runnable(other)))
@@ -237,7 +237,7 @@ class PregelNode(Runnable):
input: Input,
config: Optional[RunnableConfig] = None,
**kwargs: Optional[Any],
) -> Output:
) -> Any:
return self.bound.invoke(
input,
merge_configs({"metadata": self.metadata, "tags": self.tags}, config),
@@ -249,7 +249,7 @@ class PregelNode(Runnable):
input: Input,
config: Optional[RunnableConfig] = None,
**kwargs: Optional[Any],
) -> Output:
) -> Any:
return await self.bound.ainvoke(
input,
merge_configs({"metadata": self.metadata, "tags": self.tags}, config),
@@ -261,7 +261,7 @@ class PregelNode(Runnable):
input: Input,
config: Optional[RunnableConfig] = None,
**kwargs: Optional[Any],
) -> Iterator[Output]:
) -> Iterator[Any]:
yield from self.bound.stream(
input,
merge_configs({"metadata": self.metadata, "tags": self.tags}, config),
@@ -273,7 +273,7 @@ class PregelNode(Runnable):
input: Input,
config: Optional[RunnableConfig] = None,
**kwargs: Optional[Any],
) -> AsyncIterator[Output]:
) -> AsyncIterator[Any]:
async for item in self.bound.astream(
input,
merge_configs({"metadata": self.metadata, "tags": self.tags}, config),
+15 -10
View File
@@ -5,11 +5,13 @@ from typing import (
Any,
AsyncIterator,
Callable,
Iterable,
Iterator,
Optional,
Sequence,
Type,
Union,
cast,
)
from langgraph.constants import ERROR, INTERRUPT, NO_WRITES
@@ -33,7 +35,7 @@ class PregelRunner:
def tick(
self,
tasks: Sequence[PregelExecutableTask],
tasks: Iterable[PregelExecutableTask],
*,
reraise: bool = True,
timeout: Optional[float] = None,
@@ -106,7 +108,7 @@ class PregelRunner:
async def atick(
self,
tasks: Sequence[PregelExecutableTask],
tasks: Iterable[PregelExecutableTask],
*,
reraise: bool = True,
timeout: Optional[float] = None,
@@ -141,14 +143,17 @@ class PregelRunner:
for t in tasks:
if not t.writes:
futures[
self.submit(
arun_with_retry,
t,
retry_policy,
stream=self.use_astream,
__name__=t.name,
__cancel_on_exit__=True,
__reraise_on_exit__=reraise,
cast(
asyncio.Future,
self.submit(
arun_with_retry,
t,
retry_policy,
stream=self.use_astream,
__name__=t.name,
__cancel_on_exit__=True,
__reraise_on_exit__=reraise,
),
)
] = t
all_futures = futures.copy()
+8 -8
View File
@@ -1,4 +1,4 @@
from typing import Mapping, Optional, Sequence, Union
from typing import Any, Mapping, Optional, Sequence, Union
from langgraph.channels.base import BaseChannel
from langgraph.constants import RESERVED
@@ -65,18 +65,18 @@ def validate_graph(
raise ValueError(f"Output channel '{chan}' not in 'channels'")
if interrupt_after_nodes != "*":
for node in interrupt_after_nodes:
if node not in nodes:
raise ValueError(f"Node {node} not in nodes")
for n in interrupt_after_nodes:
if n not in nodes:
raise ValueError(f"Node {n} not in nodes")
if interrupt_before_nodes != "*":
for node in interrupt_before_nodes:
if node not in nodes:
raise ValueError(f"Node {node} not in nodes")
for n in interrupt_before_nodes:
if n not in nodes:
raise ValueError(f"Node {n} not in nodes")
def validate_keys(
keys: Optional[Union[str, Sequence[str]]],
channels: Mapping[str, BaseChannel],
channels: Mapping[str, Any],
) -> None:
if isinstance(keys, str):
if keys not in channels:
+3 -2
View File
@@ -9,6 +9,7 @@ from typing import (
Sequence,
TypeVar,
Union,
cast,
)
from langchain_core.runnables import Runnable, RunnableConfig
@@ -34,7 +35,7 @@ class ChannelWriteEntry(NamedTuple):
class ChannelWrite(RunnableCallable):
writes: Sequence[Union[ChannelWriteEntry, Send]]
writes: list[Union[ChannelWriteEntry, Send]]
"""
Sequence of write entries, each of which is a tuple of:
- channel name
@@ -54,7 +55,7 @@ class ChannelWrite(RunnableCallable):
require_at_least_one_of: Optional[Sequence[str]] = None,
):
super().__init__(func=self._write, afunc=self._awrite, name=None, tags=tags)
self.writes = writes
self.writes = cast(list[Union[ChannelWriteEntry, Send]], writes)
self.require_at_least_one_of = require_at_least_one_of
def get_name(
+1 -1
View File
@@ -269,7 +269,7 @@ class RunnableSeq(Runnable):
if isinstance(step, RunnableSequence):
steps_flat.extend(step.steps)
elif isinstance(step, RunnableSeq):
steps_flat.extend(step.steps) # type: ignore[has-type]
steps_flat.extend(step.steps)
else:
steps_flat.append(coerce_to_runnable(step, name=None, trace=True))
if len(steps_flat) < 2:
+1 -1
View File
@@ -57,7 +57,7 @@ warn_no_return = "False"
warn_unused_ignores = "True"
warn_redundant_casts = "True"
allow_redefinition = "True"
disable_error_code = "typeddict-item, return-value, override"
disable_error_code = "typeddict-item, return-value, override, has-type"
[tool.coverage.run]
omit = ["tests/*"]