WEBDIENSTE/langgraph
1
0
Fork 0
mirror of https://github.com/langchain-ai/langgraph.git synced 2026-08-25 09:02:25 +02:00
Code Issues Packages Projects Releases Wiki Activity

Merge pull request #1597 from langchain-ai/nc/3sep/perf-optimizations

Browse Source 
Reduce cpu time spent on langchain-core utilities
This commit is contained in:
Nuno Campos
2024-09-04 12:45:20 -07:00
committed by GitHub
parent 6bf367300b 3a1b90c882
commit af780f97a4
24 changed files with 963 additions and 478 deletions
Download Patch File Download Diff File Expand all files Collapse all files
libs/checkpoint/langgraph/checkpoint/base/id.py
+3 -3
View File
@@ -3,7 +3,7 @@ https://github.com/oittaa/uuid6-python/blob/main/src/uuid6/__init__.py#L95
Bundled in to avoid install issues with uuid6 package
"""
import secrets
import random
import time
import uuid
from typing import Optional, Tuple
@@ -96,9 +96,9 @@ def uuid6(node: Optional[int] = None, clock_seq: Optional[int] = None) -> UUID:
timestamp = _last_v6_timestamp + 1
_last_v6_timestamp = timestamp
if clock_seq is None:
clock_seq = secrets.randbits(14) # instead of stable storage
clock_seq = random.getrandbits(14) # instead of stable storage
if node is None:
node = secrets.randbits(48)
node = random.getrandbits(48)
time_high_and_time_mid = (timestamp >> 12) & 0xFFFFFFFFFFFF
time_low_and_version = timestamp & 0x0FFF
uuid_int = time_high_and_time_mid << 80
libs/langgraph/langgraph/graph/graph.py
+17 -10
View File
@@ -1,3 +1,4 @@
import asyncio
import logging
from collections import defaultdict
from typing import (
@@ -38,7 +39,7 @@ from langgraph.pregel import Channel, Pregel
from langgraph.pregel.read import PregelNode
from langgraph.pregel.types import All
from langgraph.pregel.write import ChannelWrite, ChannelWriteEntry
from langgraph.utils import RunnableCallable, coerce_to_runnable
from langgraph.utils.runnable import RunnableCallable, coerce_to_runnable
logger = logging.getLogger(__name__)
@@ -55,7 +56,7 @@ class Branch(NamedTuple):
def run(
self,
writer: Callable[[list[str]], Optional[Runnable]],
writer: Callable[[list[str], RunnableConfig], None],
reader: Optional[Callable[[RunnableConfig], Any]] = None,
) -> None:
return ChannelWrite.register_writer(
@@ -75,7 +76,7 @@ class Branch(NamedTuple):
config: RunnableConfig,
*,
reader: Optional[Callable[[], Any]],
writer: Callable[[list[str]], Optional[Runnable]],
writer: Callable[[list[str], RunnableConfig], None],
) -> Runnable:
if reader:
value = reader(config)
@@ -86,7 +87,7 @@ class Branch(NamedTuple):
else:
value = input
result = self.path.invoke(value, config)
return self._finish(writer, input, result)
return self._finish(writer, input, result, config)
async def _aroute(
self,
@@ -94,10 +95,10 @@ class Branch(NamedTuple):
config: RunnableConfig,
*,
reader: Optional[Callable[[], Any]],
writer: Callable[[list[str]], Optional[Runnable]],
writer: Callable[[list[str], RunnableConfig], Optional[Runnable]],
) -> Runnable:
if reader:
value = reader(config)
value = await asyncio.to_thread(reader, config)
# passthrough additional keys from node to branch
# only doable when using dict states
if isinstance(value, dict) and isinstance(input, dict):
@@ -105,10 +106,14 @@ class Branch(NamedTuple):
else:
value = input
result = await self.path.ainvoke(value, config)
return self._finish(writer, input, result)
return self._finish(writer, input, result, config)
def _finish(
self, writer: Callable[[list[str]], Optional[Runnable]], input: Any, result: Any
self,
writer: Callable[[list[str], RunnableConfig], None],
input: Any,
result: Any,
config: RunnableConfig,
):
if not isinstance(result, list):
result = [result]
@@ -120,7 +125,7 @@ class Branch(NamedTuple):
raise ValueError("Branch did not return a valid destination")
if any(p.node == END for p in destinations if isinstance(p, Send)):
raise InvalidUpdateError("Cannot send a packet to the END node")
return writer(destinations) or input
return writer(destinations, config) or input
class Graph:
@@ -449,7 +454,9 @@ class CompiledGraph(Pregel):
self.nodes[end].channels.append(start)
def attach_branch(self, start: str, name: str, branch: Branch) -> None:
def branch_writer(packets: list[Union[str, Send]]) -> Optional[ChannelWrite]:
def branch_writer(
packets: list[Union[str, Send]], config: RunnableConfig
) -> Optional[ChannelWrite]:
writes = [
(
ChannelWriteEntry(f"branch:{start}:{name}:{p}" if p != END else END)
libs/langgraph/langgraph/graph/state.py
+10 -13
View File
@@ -44,7 +44,8 @@ from langgraph.pregel.read import ChannelRead, PregelNode
from langgraph.pregel.types import All, RetryPolicy
from langgraph.pregel.write import SKIP_WRITE, ChannelWrite, ChannelWriteEntry
from langgraph.store.base import BaseStore
from langgraph.utils import RunnableCallable, coerce_to_runnable, get_field_default
from langgraph.utils.fields import get_field_default
from langgraph.utils.runnable import coerce_to_runnable
logger = logging.getLogger(__name__)
@@ -530,9 +531,7 @@ class CompiledStateGraph(CompiledGraph):
if is_writable_managed_value(v)
]
def _get_state_key(
input: Union[None, dict, Any], config: RunnableConfig, *, key: str
) -> Any:
def _get_state_key(input: Union[None, dict, Any], *, key: str) -> Any:
if input is None:
return SKIP_WRITE
elif isinstance(input, dict):
@@ -548,12 +547,7 @@ class CompiledStateGraph(CompiledGraph):
[ChannelWriteEntry("__root__", skip_none=True)]
if output_keys == ["__root__"]
else [
ChannelWriteEntry(
key,
mapper=RunnableCallable(
_get_state_key, key=key, trace=False, recurse=False
),
)
ChannelWriteEntry(key, mapper=partial(_get_state_key, key=key))
for key in output_keys
]
)
@@ -596,7 +590,8 @@ class CompiledStateGraph(CompiledGraph):
],
metadata=node.metadata,
retry_policy=node.retry_policy,
).pipe(node.runnable)
bound=node.runnable,
)
def attach_edge(self, starts: Union[str, Sequence[str]], end: str) -> None:
if isinstance(starts, str):
@@ -626,7 +621,9 @@ class CompiledStateGraph(CompiledGraph):
)
def attach_branch(self, start: str, name: str, branch: Branch) -> None:
def branch_writer(packets: list[Union[str, Send]]) -> Optional[ChannelWrite]:
def branch_writer(
packets: list[Union[str, Send]], config: RunnableConfig
) -> Optional[ChannelWrite]:
if filtered := [p for p in packets if p != END]:
writes = [
(
@@ -645,7 +642,7 @@ class CompiledStateGraph(CompiledGraph):
),
)
)
return ChannelWrite(writes, tags=[TAG_HIDDEN])
ChannelWrite.do_write(config, writes)
# attach branch publisher
schema = (
libs/langgraph/langgraph/prebuilt/tool_executor.py
+1 -1
View File
@@ -6,7 +6,7 @@ from langchain_core.tools import BaseTool
from langchain_core.tools import tool as create_tool
from langgraph._api.deprecation import deprecated
from langgraph.utils import RunnableCallable
from langgraph.utils.runnable import RunnableCallable
INVALID_TOOL_MSG_TEMPLATE = (
"{requested_tool_name} is not a valid tool, "
libs/langgraph/langgraph/prebuilt/tool_node.py
+1 -1
View File
@@ -21,7 +21,7 @@ from langchain_core.tools import BaseTool, InjectedToolArg
from langchain_core.tools import tool as create_tool
from typing_extensions import get_args
from langgraph.utils import RunnableCallable
from langgraph.utils.runnable import RunnableCallable
INVALID_TOOL_NAME_ERROR_TEMPLATE = (
"Error: {requested_tool} is not a valid tool, try one of [{available_tools}]."
libs/langgraph/langgraph/prebuilt/tool_validator.py
+1 -1
View File
@@ -33,7 +33,7 @@ from langchain_core.tools import BaseTool, create_schema_from_function
from pydantic import BaseModel as BaseModelV2
from pydantic import ValidationError as ValidationErrorV2
from langgraph.utils import RunnableCallable
from langgraph.utils.runnable import RunnableCallable
def _default_format_error(
libs/langgraph/langgraph/pregel/__init__.py
+16 -20
View File
@@ -6,7 +6,6 @@ from functools import partial
from typing import (
Any,
AsyncIterator,
Awaitable,
Callable,
Dict,
Iterator,
@@ -28,14 +27,12 @@ from langchain_core.runnables import (
RunnableLambda,
RunnableSequence,
)
from langchain_core.runnables.base import Input, Output, coerce_to_runnable
from langchain_core.runnables.base import Input, Output
from langchain_core.runnables.config import (
RunnableConfig,
ensure_config,
get_async_callback_manager_for_config,
get_callback_manager_for_config,
merge_configs,
patch_config,
)
from langchain_core.runnables.utils import (
ConfigurableFieldSpec,
@@ -76,7 +73,6 @@ from langgraph.pregel.algo import (
local_write,
prepare_next_tasks,
)
from langgraph.pregel.config import patch_checkpoint_map, patch_configurable
from langgraph.pregel.debug import tasks_w_writes
from langgraph.pregel.io import read_channels
from langgraph.pregel.loop import AsyncPregelLoop, SyncPregelLoop
@@ -96,14 +92,15 @@ from langgraph.pregel.utils import (
from langgraph.pregel.validate import validate_graph, validate_keys
from langgraph.pregel.write import ChannelWrite, ChannelWriteEntry
from langgraph.store.base import BaseStore
from langgraph.utils import RunnableCallable
from langgraph.utils.config import (
merge_configs,
patch_checkpoint_map,
patch_config,
patch_configurable,
)
from langgraph.utils.runnable import RunnableCallable
WriteValue = Union[
Runnable[Input, Output],
Callable[[Input], Output],
Callable[[Input], Awaitable[Output]],
Any,
]
WriteValue = Union[Callable[[Input], Output], Any]
class Channel:
@@ -168,11 +165,9 @@ class Channel:
return ChannelWrite(
[ChannelWriteEntry(c) for c in channels]
+ [
(
ChannelWriteEntry(k, skip_none=True, mapper=coerce_to_runnable(v))
if isinstance(v, Runnable) or callable(v)
else ChannelWriteEntry(k, value=v)
)
ChannelWriteEntry(k, mapper=v)
if callable(v)
else ChannelWriteEntry(k, value=v)
for k, v in kwargs.items()
]
)
@@ -809,7 +804,7 @@ class Pregel(Runnable[Union[dict[str, Any], Any], Union[dict[str, Any], Any]]):
managed,
):
# create task to run all writers of the chosen node
writers = self.nodes[as_node].get_writers()
writers = self.nodes[as_node].flat_writers
if not writers:
raise InvalidUpdateError(f"Node {as_node} has no writers")
task = PregelExecutableTask(
@@ -973,7 +968,7 @@ class Pregel(Runnable[Union[dict[str, Any], Any], Union[dict[str, Any], Any]]):
managed,
):
# create task to run all writers of the chosen node
writers = self.nodes[as_node].get_writers()
writers = self.nodes[as_node].flat_writers
if not writers:
raise InvalidUpdateError(f"Node {as_node} has no writers")
task = PregelExecutableTask(
@@ -1424,7 +1419,8 @@ class Pregel(Runnable[Union[dict[str, Any], Any], Union[dict[str, Any], Any]]):
# channel updates from step N are only visible in step N+1
# channels are guaranteed to be immutable for the duration of the step,
# with channel updates applied only at the transition between steps
while loop.tick(
while await asyncio.to_thread(
loop.tick,
input_keys=self.input_channels,
interrupt_before=interrupt_before,
interrupt_after=interrupt_after,
libs/langgraph/langgraph/pregel/algo.py
+19 -13
View File
@@ -1,4 +1,3 @@
import json
from collections import defaultdict, deque
from functools import partial
from typing import (
@@ -17,11 +16,7 @@ from typing import (
from uuid import UUID, uuid5
from langchain_core.callbacks.manager import AsyncParentRunManager, ParentRunManager
from langchain_core.runnables.config import (
RunnableConfig,
merge_configs,
patch_config,
)
from langchain_core.runnables.config import RunnableConfig
from langgraph.channels.base import BaseChannel
from langgraph.checkpoint.base import (
@@ -51,6 +46,7 @@ from langgraph.pregel.log import logger
from langgraph.pregel.manager import ChannelsManager
from langgraph.pregel.read import PregelNode
from langgraph.pregel.types import All, PregelExecutableTask, PregelTask
from langgraph.utils.config import merge_configs, patch_config
class WritesProtocol(Protocol):
@@ -273,16 +269,19 @@ def prepare_next_tasks(
checkpointer: Optional[BaseCheckpointSaver] = None,
manager: Union[None, ParentRunManager, AsyncParentRunManager] = None,
) -> Union[list[PregelTask], list[PregelExecutableTask]]:
checkpoint_id = UUID(checkpoint["id"])
configurable = config.get("configurable", {})
parent_ns = configurable.get("checkpoint_ns", "")
tasks: Union[list[PregelTask], list[PregelExecutableTask]] = []
# Consume pending packets
for packet in checkpoint["pending_sends"]:
if not isinstance(packet, Send):
logger.warn(f"Ignoring invalid packet type {type(packet)} in pending sends")
logger.warning(
f"Ignoring invalid packet type {type(packet)} in pending sends"
)
continue
if packet.node not in processes:
logger.warn(f"Ignoring unknown node name {packet.node} in pending sends")
logger.warning(f"Ignoring unknown node name {packet.node} in pending sends")
continue
# create task id
triggers = [TASKS]
@@ -296,11 +295,16 @@ def prepare_next_tasks(
f"{parent_ns}{NS_SEP}{packet.node}" if parent_ns else packet.node
)
task_id = str(
uuid5(UUID(checkpoint["id"]), json.dumps((checkpoint_ns, metadata)))
uuid5(
checkpoint_id,
"".join(
(checkpoint_ns, str(step), packet.node, *triggers, str(len(tasks)))
),
)
)
if for_execution:
proc = processes[packet.node]
if node := proc.get_node():
if node := proc.node:
managed.replace_runtime_placeholders(step, packet.arg)
writes = deque()
task_checkpoint_ns = f"{checkpoint_ns}:{task_id}"
@@ -400,13 +404,15 @@ def prepare_next_tasks(
checkpoint_ns = f"{parent_ns}{NS_SEP}{name}" if parent_ns else name
task_id = str(
uuid5(
UUID(checkpoint["id"]),
json.dumps((checkpoint_ns, metadata)),
checkpoint_id,
"".join(
(checkpoint_ns, str(step), name, *triggers, str(len(tasks)))
),
)
)
if for_execution:
if node := proc.get_node():
if node := proc.node:
writes = deque()
task_checkpoint_ns = f"{checkpoint_ns}:{task_id}"
tasks.append(
libs/langgraph/langgraph/pregel/config.py
-34
View File
@@ -1,34 +0,0 @@
from typing import Any, Optional
from langchain_core.runnables import RunnableConfig
from langgraph.checkpoint.base import CheckpointMetadata
from langgraph.constants import CONFIG_KEY_CHECKPOINT_MAP
def patch_configurable(
config: Optional[RunnableConfig], patch: dict[str, Any]
) -> RunnableConfig:
if config is None:
return {"configurable": patch}
else:
return {**config, "configurable": {**config["configurable"], **patch}}
def patch_checkpoint_map(
config: RunnableConfig, metadata: Optional[CheckpointMetadata]
) -> RunnableConfig:
if parents := (metadata.get("parents") if metadata else None):
return patch_configurable(
config,
{
CONFIG_KEY_CHECKPOINT_MAP: {
**parents,
config["configurable"]["checkpoint_ns"]: config["configurable"][
"checkpoint_id"
],
},
},
)
else:
return config
libs/langgraph/langgraph/pregel/executor.py
+3 -2
View File
@@ -99,6 +99,7 @@ class AsyncBackgroundExecutor(AsyncContextManager):
self.context_not_supported = sys.version_info < (3, 11)
self.tasks: dict[asyncio.Task, bool] = {}
self.sentinel = object()
self.loop = asyncio.get_running_loop()
def submit(
self,
@@ -110,9 +111,9 @@ class AsyncBackgroundExecutor(AsyncContextManager):
) -> asyncio.Task[T]:
coro = fn(*args, **kwargs)
if self.context_not_supported:
task = asyncio.create_task(coro, name=__name__)
task = self.loop.create_task(coro, name=__name__)
else:
task = asyncio.create_task(coro, name=__name__, context=copy_context())
task = self.loop.create_task(coro, name=__name__, context=copy_context())
self.tasks[task] = __cancel_on_exit__
task.add_done_callback(self.done)
return task
libs/langgraph/langgraph/pregel/loop.py
+1 -1
View File
@@ -59,7 +59,6 @@ from langgraph.pregel.algo import (
prepare_next_tasks,
should_interrupt,
)
from langgraph.pregel.config import patch_configurable
from langgraph.pregel.debug import (
map_debug_checkpoint,
map_debug_task_results,
@@ -86,6 +85,7 @@ from langgraph.pregel.types import PregelExecutableTask
from langgraph.pregel.utils import get_new_channel_versions
from langgraph.store.base import BaseStore
from langgraph.store.batch import AsyncBatchedStore
from langgraph.utils.config import patch_configurable
V = TypeVar("V")
INPUT_DONE = object()
libs/langgraph/langgraph/pregel/manager.py
+4 -3
View File
@@ -2,7 +2,7 @@ import asyncio
from contextlib import AsyncExitStack, ExitStack, asynccontextmanager, contextmanager
from typing import AsyncIterator, Iterator, Mapping, Optional, Union
from langchain_core.runnables import RunnableConfig, patch_config
from langchain_core.runnables import RunnableConfig
from langgraph.channels.base import BaseChannel
from langgraph.checkpoint.base import Checkpoint
@@ -14,6 +14,7 @@ from langgraph.managed.base import (
)
from langgraph.managed.context import Context
from langgraph.store.base import BaseStore
from langgraph.utils.config import patch_configurable
@contextmanager
@@ -26,7 +27,7 @@ def ChannelsManager(
skip_context: bool = False,
) -> Iterator[tuple[Mapping[str, BaseChannel], ManagedValueMapping]]:
"""Manage channels for the lifetime of a Pregel invocation (multiple steps)."""
config_for_managed = patch_config(config, configurable={CONFIG_KEY_STORE: store})
config_for_managed = patch_configurable(config, {CONFIG_KEY_STORE: store})
channel_specs: Mapping[str, BaseChannel] = {}
managed_specs: Mapping[str, ManagedValueSpec] = {}
for k, v in specs.items():
@@ -69,7 +70,7 @@ async def AsyncChannelsManager(
skip_context: bool = False,
) -> AsyncIterator[Mapping[str, BaseChannel]]:
"""Manage channels for the lifetime of a Pregel invocation (multiple steps)."""
config_for_managed = patch_config(config, configurable={CONFIG_KEY_STORE: store})
config_for_managed = patch_configurable(config, {CONFIG_KEY_STORE: store})
channel_specs: Mapping[str, BaseChannel] = {}
managed_specs: Mapping[str, ManagedValueSpec] = {}
for k, v in specs.items():
libs/langgraph/langgraph/pregel/read.py
+11 -9
View File
@@ -1,5 +1,6 @@
from __future__ import annotations
from functools import cached_property
from typing import (
Any,
AsyncIterator,
@@ -15,17 +16,16 @@ from langchain_core.runnables import (
Runnable,
RunnableConfig,
RunnablePassthrough,
RunnableSequence,
RunnableSerializable,
)
from langchain_core.runnables.base import Input, Other, Output, coerce_to_runnable
from langchain_core.runnables.config import merge_configs
from langchain_core.runnables.utils import ConfigurableFieldSpec
from langgraph.constants import CONFIG_KEY_READ
from langgraph.pregel.retry import RetryPolicy
from langgraph.pregel.write import ChannelWrite
from langgraph.utils import RunnableCallable
from langgraph.utils.config import merge_configs
from langgraph.utils.runnable import RunnableCallable, RunnableSeq
READ_TYPE = Callable[[str, bool], Union[Any, dict[str, Any]]]
@@ -149,7 +149,8 @@ class PregelNode(Runnable):
attrs = {**self.__dict__, **update}
return PregelNode(**attrs)
def get_writers(self) -> list[Runnable]:
@cached_property
def flat_writers(self) -> list[Runnable]:
"""Get writers with optimizations applied."""
writers = self.writers.copy()
while (
@@ -167,16 +168,17 @@ class PregelNode(Runnable):
writers.pop()
return writers
def get_node(self) -> Optional[Runnable[Any, Any]]:
writers = self.get_writers()
@cached_property
def node(self) -> Optional[Runnable[Any, Any]]:
writers = self.flat_writers
if self.bound is DEFAULT_BOUND and not writers:
return None
elif self.bound is DEFAULT_BOUND and len(writers) == 1:
return writers[0]
elif self.bound is DEFAULT_BOUND:
return RunnableSequence(*writers)
return RunnableSeq(*writers)
elif writers:
return RunnableSequence(self.bound, *writers)
return RunnableSeq(self.bound, *writers)
else:
return self.bound
@@ -209,7 +211,7 @@ class PregelNode(Runnable):
elif self.bound is DEFAULT_BOUND:
return self.copy(update=dict(bound=coerce_to_runnable(other)))
else:
return self.copy(update=dict(bound=self.bound | other))
return self.copy(update=dict(bound=RunnableSeq(self.bound, other)))
def pipe(
self,
libs/langgraph/langgraph/pregel/write.py
+39 -57
View File
@@ -4,11 +4,9 @@ import asyncio
from typing import (
Any,
Callable,
List,
NamedTuple,
Optional,
Sequence,
Tuple,
TypeVar,
Union,
)
@@ -18,7 +16,7 @@ from langchain_core.runnables.utils import ConfigurableFieldSpec
from langgraph.constants import CONFIG_KEY_SEND, TASKS, Send
from langgraph.errors import InvalidUpdateError
from langgraph.utils import RunnableCallable
from langgraph.utils.runnable import RunnableCallable
TYPE_SEND = Callable[[Sequence[tuple[str, Any]]], None]
R = TypeVar("R", bound=Runnable)
@@ -32,7 +30,7 @@ class ChannelWriteEntry(NamedTuple):
channel: str
value: Any = PASSTHROUGH
skip_none: bool = False
mapper: Optional[Runnable] = None
mapper: Optional[Callable] = None
class ChannelWrite(RunnableCallable):
@@ -59,9 +57,6 @@ class ChannelWrite(RunnableCallable):
self.writes = writes
self.require_at_least_one_of = require_at_least_one_of
def __repr_args__(self) -> Any:
return [("writes", self.writes)]
def get_name(
self, suffix: Optional[str] = None, *, name: Optional[str] = None
) -> str:
@@ -82,65 +77,29 @@ class ChannelWrite(RunnableCallable):
]
def _write(self, input: Any, config: RunnableConfig) -> None:
# split packets and entries
writes = [(TASKS, packet) for packet in self.writes if isinstance(packet, Send)]
entries = [
write for write in self.writes if isinstance(write, ChannelWriteEntry)
writes = [
ChannelWriteEntry(write.channel, input, write.skip_none, write.mapper)
if isinstance(write, ChannelWriteEntry) and write.value is PASSTHROUGH
else write
for write in self.writes
]
for entry in entries:
if entry.channel == TASKS:
raise InvalidUpdateError("Cannot write to the reserved channel TASKS")
# process entries into values
values = [
input if write.value is PASSTHROUGH else write.value for write in entries
]
values = [
val if write.mapper is None else write.mapper.invoke(val, config)
for val, write in zip(values, entries)
]
values = [
(write.channel, val)
for val, write in zip(values, entries)
if not write.skip_none or val is not None
]
# write packets and values
self.do_write(
config,
writes + values,
writes,
self.require_at_least_one_of if input is not None else None,
)
return input
async def _awrite(self, input: Any, config: RunnableConfig) -> None:
# split packets and entries
writes = [(TASKS, packet) for packet in self.writes if isinstance(packet, Send)]
entries = [
write for write in self.writes if isinstance(write, ChannelWriteEntry)
writes = [
ChannelWriteEntry(write.channel, input, write.skip_none, write.mapper)
if isinstance(write, ChannelWriteEntry) and write.value is PASSTHROUGH
else write
for write in self.writes
]
for entry in entries:
if entry.channel == TASKS:
raise InvalidUpdateError("Cannot write to the reserved channel TASKS")
# process entries into values
values = [
input if write.value is PASSTHROUGH else write.value for write in entries
]
values = await asyncio.gather(
*(
_mk_future(val)
if write.mapper is None
else write.mapper.ainvoke(val, config)
for val, write in zip(values, entries)
)
)
values = [
(write.channel, val)
for val, write in zip(values, entries)
if not write.skip_none or val is not None
]
# write packets and values
self.do_write(
config,
writes + values,
writes,
self.require_at_least_one_of if input is not None else None,
)
return input
@@ -148,9 +107,32 @@ class ChannelWrite(RunnableCallable):
@staticmethod
def do_write(
config: RunnableConfig,
values: List[Tuple[str, Any]],
writes: Sequence[Union[ChannelWriteEntry, Send]],
require_at_least_one_of: Optional[Sequence[str]] = None,
) -> None:
# validate
for w in writes:
if isinstance(w, ChannelWriteEntry):
if w.channel == TASKS:
raise InvalidUpdateError(
"Cannot write to the reserved channel TASKS"
)
if w.value is PASSTHROUGH:
raise InvalidUpdateError("PASSTHROUGH value must be replaced")
# split packets and entries
sends = [(TASKS, packet) for packet in writes if isinstance(packet, Send)]
entries = [write for write in writes if isinstance(write, ChannelWriteEntry)]
# process entries into values
values = [
write.mapper(write.value) if write.mapper is not None else write.value
for write in entries
]
values = [
(write.channel, val)
for val, write in zip(values, entries)
if not write.skip_none or val is not None
]
# filter out SKIP_WRITE values
filtered = [(chan, val) for chan, val in values if val is not SKIP_WRITE]
if require_at_least_one_of is not None:
if not {chan for chan, _ in filtered} & set(require_at_least_one_of):
@@ -158,7 +140,7 @@ class ChannelWrite(RunnableCallable):
f"Must write to at least one of {require_at_least_one_of}"
)
write: TYPE_SEND = config["configurable"][CONFIG_KEY_SEND]
write(filtered)
write(sends + filtered)
@staticmethod
def is_writer(runnable: Runnable) -> bool:
libs/langgraph/langgraph/utils.py
-282
View File
@@ -1,282 +0,0 @@
import asyncio
import enum
import inspect
import sys
from contextvars import copy_context
from functools import partial, wraps
from typing import Any, AsyncIterator, Awaitable, Callable, Optional, Type, Union
from langchain_core.runnables.base import (
Runnable,
RunnableConfig,
RunnableLambda,
RunnableLike,
RunnableParallel,
)
from langchain_core.runnables.config import (
merge_configs,
run_in_executor,
var_child_runnable_config,
)
from langchain_core.runnables.utils import accepts_config
from typing_extensions import (
Annotated,
NotRequired,
ReadOnly,
Required,
TypeGuard,
get_origin,
)
try:
from langchain_core.runnables.config import _set_config_context
except ImportError:
# For forwards compatibility
def _set_config_context(context: RunnableConfig) -> None: # type: ignore
"""Set the context for the current thread."""
var_child_runnable_config.set(context)
# Before Python 3.11 native StrEnum is not available
class StrEnum(str, enum.Enum):
"""A string enum."""
class RunnableCallable(Runnable):
"""A much simpler version of RunnableLambda that requires sync and async functions."""
def __init__(
self,
func: Callable[..., Optional[Runnable]],
afunc: Optional[Callable[..., Awaitable[Optional[Runnable]]]] = None,
*,
name: Optional[str] = None,
tags: Optional[list[str]] = None,
trace: bool = True,
recurse: bool = True,
**kwargs: Any,
) -> None:
if name is not None:
self.name = name
elif func:
try:
if func.__name__ != "<lambda>":
self.name = func.__name__
except AttributeError:
pass
elif afunc:
try:
self.name = afunc.__name__
except AttributeError:
pass
self.func = func
self.afunc = afunc
self.config: Optional[RunnableConfig] = {"tags": tags} if tags else None
self.kwargs = kwargs
self.trace = trace
self.recurse = recurse
def __repr__(self) -> str:
repr_args = {
k: v
for k, v in self.__dict__.items()
if k not in {"name", "func", "afunc", "config", "kwargs", "trace"}
}
return f"{self.get_name()}({', '.join(f'{k}={v!r}' for k, v in repr_args.items())})"
def invoke(
self, input: Any, config: Optional[RunnableConfig] = None, **kwargs: Any
) -> Any:
if self.func is None:
raise TypeError(
f'No synchronous function provided to "{self.name}".'
"\nEither initialize with a synchronous function or invoke"
" via the async API (ainvoke, astream, etc.)"
)
kwargs = {**self.kwargs, **kwargs}
if self.trace:
ret = self._call_with_config(
self.func, input, merge_configs(self.config, config), **kwargs
)
else:
config = merge_configs(self.config, config)
context = copy_context()
context.run(_set_config_context, config)
if accepts_config(self.func):
kwargs["config"] = config
ret = context.run(self.func, input, **kwargs)
if isinstance(ret, Runnable) and self.recurse:
return ret.invoke(input, config)
return ret
async def ainvoke(
self, input: Any, config: Optional[RunnableConfig] = None, **kwargs: Any
) -> Any:
if not self.afunc:
return self.invoke(input, config)
kwargs = {**self.kwargs, **kwargs}
if self.trace:
ret = await self._acall_with_config(
self.afunc, input, merge_configs(self.config, config), **kwargs
)
else:
config = merge_configs(self.config, config)
context = copy_context()
context.run(_set_config_context, config)
if accepts_config(self.afunc):
kwargs["config"] = config
if sys.version_info >= (3, 11):
ret = await asyncio.create_task(
self.afunc(input, **kwargs), context=context
)
else:
ret = await self.afunc(input, **kwargs)
if isinstance(ret, Runnable) and self.recurse:
return await ret.ainvoke(input, config)
return ret
def is_async_callable(
func: Any,
) -> TypeGuard[Callable[..., Awaitable]]:
"""Check if a function is async."""
return (
asyncio.iscoroutinefunction(func)
or hasattr(func, "__call__")
and asyncio.iscoroutinefunction(func.__call__)
)
def is_async_generator(
func: Any,
) -> TypeGuard[Callable[..., AsyncIterator]]:
"""Check if a function is an async generator."""
return (
inspect.isasyncgenfunction(func)
or hasattr(func, "__call__")
and inspect.isasyncgenfunction(func.__call__)
)
def coerce_to_runnable(thing: RunnableLike, *, name: str, trace: bool) -> Runnable:
"""Coerce a runnable-like object into a Runnable.
Args:
thing: A runnable-like object.
Returns:
A Runnable.
"""
if isinstance(thing, Runnable):
return thing
elif is_async_generator(thing) or inspect.isgeneratorfunction(thing):
return RunnableLambda(thing, name=name)
elif callable(thing):
if is_async_callable(thing):
return RunnableCallable(None, thing, name=name, trace=trace)
else:
return RunnableCallable(
thing,
wraps(thing)(partial(run_in_executor, None, thing)),
name=name,
trace=trace,
)
elif isinstance(thing, dict):
return RunnableParallel(thing)
else:
raise TypeError(
f"Expected a Runnable, callable or dict."
f"Instead got an unsupported type: {type(thing)}"
)
def _is_optional_type(type_: Any) -> bool:
"""Check if a type is Optional."""
if hasattr(type_, "__origin__") and hasattr(type_, "__args__"):
origin = get_origin(type_)
if origin is Optional:
return True
if origin is Union:
return any(
arg is type(None) or _is_optional_type(arg) for arg in type_.__args__
)
if origin is Annotated:
return _is_optional_type(type_.__args__[0])
return origin is None
if hasattr(type_, "__bound__") and type_.__bound__ is not None:
return _is_optional_type(type_.__bound__)
return type_ is None
def _is_required_type(type_: Any) -> Optional[bool]:
"""Check if an annotation is marked as Required/NotRequired.
Returns:
- True if required
- False if not required
- None if not annotated with either
"""
origin = get_origin(type_)
if origin is Required:
return True
if origin is NotRequired:
return False
if origin is Annotated or getattr(origin, "__args__", None):
# See https://typing.readthedocs.io/en/latest/spec/typeddict.html#interaction-with-annotated
return _is_required_type(type_.__args__[0])
return None
def _is_readonly_type(type_: Any) -> bool:
"""Check if an annotation is marked as ReadOnly.
Returns:
- True if is read only
- False if not read only
"""
# See: https://typing.readthedocs.io/en/latest/spec/typeddict.html#typing-readonly-type-qualifier
origin = get_origin(type_)
if origin is Annotated:
return _is_readonly_type(type_.__args__[0])
if origin is ReadOnly:
return True
return False
_DEFAULT_KEYS = frozenset()
def get_field_default(name: str, type_: Any, schema: Type[Any]) -> Any:
"""Determine the default value for a field in a state schema.
This is based on:
If TypedDict:
- Required/NotRequired
- total=False -> everything optional
- Type annotation (Optional/Union[None])
"""
optional_keys = getattr(schema, "__optional_keys__", _DEFAULT_KEYS)
irq = _is_required_type(type_)
if name in optional_keys:
# Either total=False or explicit NotRequired.
# No type annotation trumps this.
if irq:
# Unless it's earlier versions of python & explicit Required
return ...
return None
if irq is not None:
if irq:
# Handle Required[<type>]
# (we already handled NotRequired and total=False)
return ...
# Handle NotRequired[<type>] for earlier versions of python
return None
# Note, we ignore ReadOnly attributes,
# as they don't make much sense. (we don't care if you mutate the state in your node)
# and mutating state in your node has no effect on our graph state.
# Base case is the annotation
if _is_optional_type(type_):
return None
return ...
libs/langgraph/langgraph/utils/__init__.py
View File
libs/langgraph/langgraph/utils/config.py
+152
View File
@@ -0,0 +1,152 @@
from typing import Any, Optional
from langchain_core.callbacks import Callbacks
from langchain_core.runnables import RunnableConfig
from langchain_core.runnables.config import COPIABLE_KEYS, DEFAULT_RECURSION_LIMIT
from langgraph.checkpoint.base import CheckpointMetadata
from langgraph.constants import CONFIG_KEY_CHECKPOINT_MAP
def patch_configurable(
config: Optional[RunnableConfig], patch: dict[str, Any]
) -> RunnableConfig:
if config is None:
return {"configurable": patch}
elif "configurable" not in config:
return {**config, "configurable": patch}
else:
return {**config, "configurable": {**config["configurable"], **patch}}
def patch_checkpoint_map(
config: RunnableConfig, metadata: Optional[CheckpointMetadata]
) -> RunnableConfig:
if parents := (metadata.get("parents") if metadata else None):
return patch_configurable(
config,
{
CONFIG_KEY_CHECKPOINT_MAP: {
**parents,
config["configurable"]["checkpoint_ns"]: config["configurable"][
"checkpoint_id"
],
},
},
)
else:
return config
def merge_configs(*configs: Optional[RunnableConfig]) -> RunnableConfig:
"""Merge multiple configs into one.
Args:
*configs (Optional[RunnableConfig]): The configs to merge.
Returns:
RunnableConfig: The merged config.
"""
base: RunnableConfig = {}
# Even though the keys aren't literals, this is correct
# because both dicts are the same type
for config in configs:
if config is None:
continue
for key in config:
if key == "metadata":
base[key] = { # type: ignore
**base.get(key, {}), # type: ignore
**(config.get(key) or {}), # type: ignore
}
elif key == "tags":
base[key] = sorted( # type: ignore
set(base.get(key, []) + (config.get(key) or [])), # type: ignore
)
elif key == "configurable":
base[key] = { # type: ignore
**base.get(key, {}), # type: ignore
**(config.get(key) or {}), # type: ignore
}
elif key == "callbacks":
base_callbacks = base.get("callbacks")
these_callbacks = config["callbacks"]
# callbacks can be either None, list[handler] or manager
# so merging two callbacks values has 6 cases
if isinstance(these_callbacks, list):
if base_callbacks is None:
base["callbacks"] = these_callbacks.copy()
elif isinstance(base_callbacks, list):
base["callbacks"] = base_callbacks + these_callbacks
else:
# base_callbacks is a manager
mngr = base_callbacks.copy()
for callback in these_callbacks:
mngr.add_handler(callback, inherit=True)
base["callbacks"] = mngr
elif these_callbacks is not None:
# these_callbacks is a manager
if base_callbacks is None:
base["callbacks"] = these_callbacks.copy()
elif isinstance(base_callbacks, list):
mngr = these_callbacks.copy()
for callback in base_callbacks:
mngr.add_handler(callback, inherit=True)
base["callbacks"] = mngr
else:
# base_callbacks is also a manager
base["callbacks"] = base_callbacks.merge(these_callbacks)
elif key == "recursion_limit":
if config["recursion_limit"] != DEFAULT_RECURSION_LIMIT:
base["recursion_limit"] = config["recursion_limit"]
elif key in COPIABLE_KEYS and config[key] is not None: # type: ignore[literal-required]
base[key] = config[key].copy() # type: ignore[literal-required]
else:
base[key] = config[key] or base.get(key) # type: ignore
return base
def patch_config(
config: Optional[RunnableConfig],
*,
callbacks: Optional[Callbacks] = None,
recursion_limit: Optional[int] = None,
max_concurrency: Optional[int] = None,
run_name: Optional[str] = None,
configurable: Optional[dict[str, Any]] = None,
) -> RunnableConfig:
"""Patch a config with new values.
Args:
config (Optional[RunnableConfig]): The config to patch.
callbacks (Optional[BaseCallbackManager], optional): The callbacks to set.
Defaults to None.
recursion_limit (Optional[int], optional): The recursion limit to set.
Defaults to None.
max_concurrency (Optional[int], optional): The max concurrency to set.
Defaults to None.
run_name (Optional[str], optional): The run name to set. Defaults to None.
configurable (Optional[Dict[str, Any]], optional): The configurable to set.
Defaults to None.
Returns:
RunnableConfig: The patched config.
"""
config = config.copy() or {}
if callbacks is not None:
# If we're replacing callbacks, we need to unset run_name
# As that should apply only to the same run as the original callbacks
config["callbacks"] = callbacks
if "run_name" in config:
del config["run_name"]
if "run_id" in config:
del config["run_id"]
if recursion_limit is not None:
config["recursion_limit"] = recursion_limit
if max_concurrency is not None:
config["max_concurrency"] = max_concurrency
if run_name is not None:
config["run_name"] = run_name
if configurable is not None:
config["configurable"] = {**config.get("configurable", {}), **configurable}
return config
libs/langgraph/langgraph/utils/fields.py
+101
View File
@@ -0,0 +1,101 @@
from typing import Any, Optional, Type, Union
from typing_extensions import (
Annotated,
NotRequired,
ReadOnly,
Required,
get_origin,
)
def _is_optional_type(type_: Any) -> bool:
"""Check if a type is Optional."""
if hasattr(type_, "__origin__") and hasattr(type_, "__args__"):
origin = get_origin(type_)
if origin is Optional:
return True
if origin is Union:
return any(
arg is type(None) or _is_optional_type(arg) for arg in type_.__args__
)
if origin is Annotated:
return _is_optional_type(type_.__args__[0])
return origin is None
if hasattr(type_, "__bound__") and type_.__bound__ is not None:
return _is_optional_type(type_.__bound__)
return type_ is None
def _is_required_type(type_: Any) -> Optional[bool]:
"""Check if an annotation is marked as Required/NotRequired.
Returns:
- True if required
- False if not required
- None if not annotated with either
"""
origin = get_origin(type_)
if origin is Required:
return True
if origin is NotRequired:
return False
if origin is Annotated or getattr(origin, "__args__", None):
# See https://typing.readthedocs.io/en/latest/spec/typeddict.html#interaction-with-annotated
return _is_required_type(type_.__args__[0])
return None
def _is_readonly_type(type_: Any) -> bool:
"""Check if an annotation is marked as ReadOnly.
Returns:
- True if is read only
- False if not read only
"""
# See: https://typing.readthedocs.io/en/latest/spec/typeddict.html#typing-readonly-type-qualifier
origin = get_origin(type_)
if origin is Annotated:
return _is_readonly_type(type_.__args__[0])
if origin is ReadOnly:
return True
return False
_DEFAULT_KEYS = frozenset()
def get_field_default(name: str, type_: Any, schema: Type[Any]) -> Any:
"""Determine the default value for a field in a state schema.
This is based on:
If TypedDict:
- Required/NotRequired
- total=False -> everything optional
- Type annotation (Optional/Union[None])
"""
optional_keys = getattr(schema, "__optional_keys__", _DEFAULT_KEYS)
irq = _is_required_type(type_)
if name in optional_keys:
# Either total=False or explicit NotRequired.
# No type annotation trumps this.
if irq:
# Unless it's earlier versions of python & explicit Required
return ...
return None
if irq is not None:
if irq:
# Handle Required[<type>]
# (we already handled NotRequired and total=False)
return ...
# Handle NotRequired[<type>] for earlier versions of python
return None
# Note, we ignore ReadOnly attributes,
# as they don't make much sense. (we don't care if you mutate the state in your node)
# and mutating state in your node has no effect on our graph state.
# Base case is the annotation
if _is_optional_type(type_):
return None
return ...
libs/langgraph/langgraph/utils/runnable.py
+519
View File
@@ -0,0 +1,519 @@
import asyncio
import enum
import inspect
import sys
from contextlib import AsyncExitStack
from contextvars import copy_context
from functools import partial, wraps
from typing import Any, AsyncIterator, Awaitable, Callable, Iterator, Optional
from langchain_core.runnables.base import (
Runnable,
RunnableConfig,
RunnableLambda,
RunnableLike,
RunnableParallel,
RunnableSequence,
)
from langchain_core.runnables.config import (
ensure_config,
get_async_callback_manager_for_config,
get_callback_manager_for_config,
run_in_executor,
var_child_runnable_config,
)
from langchain_core.runnables.utils import Input, Output, accepts_config
from langchain_core.tracers._streaming import _StreamingCallbackHandler
from typing_extensions import TypeGuard
from langgraph.utils.config import merge_configs, patch_config
try:
from langchain_core.runnables.config import _set_config_context
except ImportError:
# For forwards compatibility
def _set_config_context(context: RunnableConfig) -> None: # type: ignore
"""Set the context for the current thread."""
var_child_runnable_config.set(context)
# Before Python 3.11 native StrEnum is not available
class StrEnum(str, enum.Enum):
"""A string enum."""
ASYNCIO_ACCEPTS_CONTEXT = sys.version_info >= (3, 11)
class RunnableCallable(Runnable):
"""A much simpler version of RunnableLambda that requires sync and async functions."""
def __init__(
self,
func: Callable[..., Optional[Runnable]],
afunc: Optional[Callable[..., Awaitable[Optional[Runnable]]]] = None,
*,
name: Optional[str] = None,
tags: Optional[list[str]] = None,
trace: bool = True,
recurse: bool = True,
**kwargs: Any,
) -> None:
if name is not None:
self.name = name
elif func:
try:
if func.__name__ != "<lambda>":
self.name = func.__name__
except AttributeError:
pass
elif afunc:
try:
self.name = afunc.__name__
except AttributeError:
pass
self.func = func
if func is not None:
self.func_accepts_config = accepts_config(func)
self.afunc = afunc
if afunc is not None:
self.afunc_accepts_config = accepts_config(afunc)
self.config: Optional[RunnableConfig] = {"tags": tags} if tags else None
self.kwargs = kwargs
self.trace = trace
self.recurse = recurse
def __repr__(self) -> str:
repr_args = {
k: v
for k, v in self.__dict__.items()
if k not in {"name", "func", "afunc", "config", "kwargs", "trace"}
}
return f"{self.get_name()}({', '.join(f'{k}={v!r}' for k, v in repr_args.items())})"
def invoke(
self, input: Any, config: Optional[RunnableConfig] = None, **kwargs: Any
) -> Any:
if self.func is None:
raise TypeError(
f'No synchronous function provided to "{self.name}".'
"\nEither initialize with a synchronous function or invoke"
" via the async API (ainvoke, astream, etc.)"
)
kwargs = {**self.kwargs, **kwargs}
if self.func_accepts_config:
kwargs["config"] = config
config = ensure_config(merge_configs(self.config, config))
context = copy_context()
if self.trace:
config = ensure_config(config)
callback_manager = get_callback_manager_for_config(config)
run_manager = callback_manager.on_chain_start(
None,
input,
name=config.get("run_name") or self.get_name(),
run_id=config.pop("run_id", None),
)
try:
child_config = patch_config(config, callbacks=run_manager.get_child())
context = copy_context()
context.run(_set_config_context, child_config)
ret = context.run(self.func, input, **kwargs)
except BaseException as e:
run_manager.on_chain_error(e)
raise
else:
run_manager.on_chain_end(ret)
else:
context.run(_set_config_context, config)
ret = context.run(self.func, input, **kwargs)
if isinstance(ret, Runnable) and self.recurse:
return ret.invoke(input, config)
return ret
async def ainvoke(
self, input: Any, config: Optional[RunnableConfig] = None, **kwargs: Any
) -> Any:
if not self.afunc:
return self.invoke(input, config)
kwargs = {**self.kwargs, **kwargs}
if self.afunc_accepts_config:
kwargs["config"] = config
config = ensure_config(merge_configs(self.config, config))
context = copy_context()
if self.trace:
callback_manager = get_async_callback_manager_for_config(config)
run_manager = await callback_manager.on_chain_start(
None,
input,
name=config.get("run_name") or self.name,
run_id=config.pop("run_id", None),
)
try:
child_config = patch_config(config, callbacks=run_manager.get_child())
context.run(_set_config_context, child_config)
coro = self.afunc(input, **kwargs)
if ASYNCIO_ACCEPTS_CONTEXT:
ret = await asyncio.create_task(coro, context=context)
else:
ret = await coro
except BaseException as e:
await run_manager.on_chain_error(e)
raise
else:
await run_manager.on_chain_end(ret)
else:
context.run(_set_config_context, config)
if ASYNCIO_ACCEPTS_CONTEXT:
ret = await asyncio.create_task(
self.afunc(input, **kwargs), context=context
)
else:
ret = await self.afunc(input, **kwargs)
if isinstance(ret, Runnable) and self.recurse:
return await ret.ainvoke(input, config)
return ret
def is_async_callable(
func: Any,
) -> TypeGuard[Callable[..., Awaitable]]:
"""Check if a function is async."""
return (
asyncio.iscoroutinefunction(func)
or hasattr(func, "__call__")
and asyncio.iscoroutinefunction(func.__call__)
)
def is_async_generator(
func: Any,
) -> TypeGuard[Callable[..., AsyncIterator]]:
"""Check if a function is an async generator."""
return (
inspect.isasyncgenfunction(func)
or hasattr(func, "__call__")
and inspect.isasyncgenfunction(func.__call__)
)
def coerce_to_runnable(thing: RunnableLike, *, name: str, trace: bool) -> Runnable:
"""Coerce a runnable-like object into a Runnable.
Args:
thing: A runnable-like object.
Returns:
A Runnable.
"""
if isinstance(thing, Runnable):
return thing
elif is_async_generator(thing) or inspect.isgeneratorfunction(thing):
return RunnableLambda(thing, name=name)
elif callable(thing):
if is_async_callable(thing):
return RunnableCallable(None, thing, name=name, trace=trace)
else:
return RunnableCallable(
thing,
wraps(thing)(partial(run_in_executor, None, thing)),
name=name,
trace=trace,
)
elif isinstance(thing, dict):
return RunnableParallel(thing)
else:
raise TypeError(
f"Expected a Runnable, callable or dict."
f"Instead got an unsupported type: {type(thing)}"
)
class RunnableSeq(Runnable):
"""A simpler version of RunnableSequence."""
def __init__(
self,
*steps: RunnableLike,
name: Optional[str] = None,
) -> None:
"""Create a new RunnableSequence.
Args:
steps: The steps to include in the sequence.
name: The name of the Runnable. Defaults to None.
first: The first Runnable in the sequence. Defaults to None.
middle: The middle Runnables in the sequence. Defaults to None.
last: The last Runnable in the sequence. Defaults to None.
Raises:
ValueError: If the sequence has less than 2 steps.
"""
steps_flat: list[Runnable] = []
for step in steps:
if isinstance(step, RunnableSequence):
steps_flat.extend(step.steps)
elif isinstance(step, RunnableSeq):
steps_flat.extend(step.steps)
else:
steps_flat.append(coerce_to_runnable(step, name=None, trace=True))
if len(steps_flat) < 2:
raise ValueError(
f"RunnableSeq must have at least 2 steps, got {len(steps_flat)}"
)
self.steps = steps_flat
self.name = name
def __or__(
self,
other: Any,
) -> Runnable:
if isinstance(other, RunnableSequence):
return RunnableSeq(
*self.steps,
other.first,
*other.middle,
other.last,
name=self.name or other.name,
)
elif isinstance(other, RunnableSeq):
return RunnableSeq(
*self.steps,
*other.steps,
name=self.name or other.name,
)
else:
return RunnableSeq(
*self.steps,
coerce_to_runnable(other),
name=self.name,
)
def __ror__(
self,
other: Any,
) -> Runnable:
if isinstance(other, RunnableSequence):
return RunnableSequence(
other.first,
*other.middle,
other.last,
*self.steps,
name=other.name or self.name,
)
elif isinstance(other, RunnableSeq):
return RunnableSeq(
*other.steps,
*self.steps,
name=other.name or self.name,
)
else:
return RunnableSequence(
coerce_to_runnable(other),
*self.steps,
name=self.name,
)
def invoke(
self, input: Input, config: Optional[RunnableConfig] = None, **kwargs: Any
) -> Output:
# setup callbacks and context
config = ensure_config(config)
callback_manager = get_callback_manager_for_config(config)
# start the root run
run_manager = callback_manager.on_chain_start(
None,
input,
name=config.get("run_name") or self.get_name(),
run_id=config.pop("run_id", None),
)
# invoke all steps in sequence
try:
for i, step in enumerate(self.steps):
# mark each step as a child run
config = patch_config(
config, callbacks=run_manager.get_child(f"seq:step:{i+1}")
)
context = copy_context()
context.run(_set_config_context, config)
if i == 0:
input = context.run(step.invoke, input, config, **kwargs)
else:
input = context.run(step.invoke, input, config)
# finish the root run
except BaseException as e:
run_manager.on_chain_error(e)
raise
else:
run_manager.on_chain_end(input)
return input
async def ainvoke(
self,
input: Input,
config: Optional[RunnableConfig] = None,
**kwargs: Optional[Any],
) -> Output:
# setup callbacks
config = ensure_config(config)
callback_manager = get_async_callback_manager_for_config(config)
# start the root run
run_manager = await callback_manager.on_chain_start(
None,
input,
name=config.get("run_name") or self.get_name(),
run_id=config.pop("run_id", None),
)
# invoke all steps in sequence
try:
for i, step in enumerate(self.steps):
# mark each step as a child run
config = patch_config(
config, callbacks=run_manager.get_child(f"seq:step:{i+1}")
)
context = copy_context()
context.run(_set_config_context, config)
if i == 0:
coro = step.ainvoke(input, config, **kwargs)
else:
coro = step.ainvoke(input, config)
if ASYNCIO_ACCEPTS_CONTEXT:
input = await asyncio.create_task(coro, context=context)
else:
input = await asyncio.create_task(coro)
# finish the root run
except BaseException as e:
await run_manager.on_chain_error(e)
raise
else:
await run_manager.on_chain_end(input)
return input
def stream(
self,
input: Input,
config: Optional[RunnableConfig] = None,
**kwargs: Optional[Any],
) -> Iterator[Output]:
# setup callbacks
config = ensure_config(config)
callback_manager = get_callback_manager_for_config(config)
# start the root run
run_manager = callback_manager.on_chain_start(
None,
input,
name=config.get("run_name") or self.get_name(),
run_id=config.pop("run_id", None),
)
try:
# stream the last steps
# transform the input stream of each step with the next
# steps that don't natively support transforming an input stream will
# buffer input in memory until all available, and then start emitting output
for idx, step in enumerate(self.steps):
config = patch_config(
config,
callbacks=run_manager.get_child(f"seq:step:{idx+1}"),
)
if idx == 0:
iterator = step.stream(input, config, **kwargs)
else:
iterator = step.transform(iterator, config)
if stream_handler := next(
(
h
for h in run_manager.handlers
if isinstance(h, _StreamingCallbackHandler)
),
None,
):
# populates streamed_output in astream_log() output if needed
iterator = stream_handler.tap_output_iter(run_manager.run_id, iterator)
output: Output = None
add_supported = False
for chunk in iterator:
yield chunk
# collect final output
if output is None:
output = chunk
elif add_supported:
try:
output = output + chunk
except TypeError:
output = chunk
add_supported = False
else:
output = chunk
except BaseException as e:
run_manager.on_chain_error(e)
raise
else:
run_manager.on_chain_end(output)
async def astream(
self,
input: Input,
config: Optional[RunnableConfig] = None,
**kwargs: Optional[Any],
) -> AsyncIterator[Output]:
# setup callbacks
config = ensure_config(config)
callback_manager = get_async_callback_manager_for_config(config)
# start the root run
run_manager = await callback_manager.on_chain_start(
None,
input,
name=config.get("run_name") or self.get_name(),
run_id=config.pop("run_id", None),
)
try:
async with AsyncExitStack() as stack:
# stream the last steps
# transform the input stream of each step with the next
# steps that don't natively support transforming an input stream will
# buffer input in memory until all available, and then start emitting output
for idx, step in enumerate(self.steps):
config = patch_config(
config,
callbacks=run_manager.get_child(f"seq:step:{idx+1}"),
)
if idx == 0:
aiterator = step.astream(input, config, **kwargs)
else:
aiterator = step.atransform(aiterator, config)
if hasattr(aiterator, "aclose"):
stack.push_async_callback(aiterator.aclose)
if stream_handler := next(
(
h
for h in run_manager.handlers
if isinstance(h, _StreamingCallbackHandler)
),
None,
):
# populates streamed_output in astream_log() output if needed
aiterator = stream_handler.tap_output_aiter(
run_manager.run_id, aiterator
)
output: Output = None
add_supported = False
async for chunk in aiterator:
yield chunk
# collect final output
if add_supported:
try:
output = output + chunk
except TypeError:
output = chunk
add_supported = False
else:
output = chunk
except BaseException as e:
await run_manager.on_chain_error(e)
raise
else:
await run_manager.on_chain_end(output)
libs/langgraph/poetry.lock
Generated
+10 -5
View File
@@ -1,4 +1,4 @@
# This file is automatically @generated by Poetry 1.8.3 and should not be changed by hand.
# This file is automatically @generated by Poetry 1.8.2 and should not be changed by hand.
[[package]]
name = "aiohappyeyeballs"
@@ -1884,18 +1884,22 @@ url = "../checkpoint-sqlite"
[[package]]
name = "langsmith"
version = "0.1.79"
version = "0.1.111"
description = "Client library to connect to the LangSmith LLM Tracing and Evaluation Platform."
optional = false
python-versions = "<4.0,>=3.8.1"
files = [
{file = "langsmith-0.1.79-py3-none-any.whl", hash = "sha256:c7f2c23981917713b5515b773f37c84ff68a7adf803476e2ebb5adcb36a04202"},
{file = "langsmith-0.1.79.tar.gz", hash = "sha256:d215718cfdcdf4a011126b7a3d4a37eee96d887e59ac1e628a57e24b2bfa3163"},
{file = "langsmith-0.1.111-py3-none-any.whl", hash = "sha256:e5c702764911193c9812fe55136ae01cd0b9ddf5dff0b068ce6fd60eeddbcb40"},
{file = "langsmith-0.1.111.tar.gz", hash = "sha256:bab24fd6125685f588d682693c4a3253e163804242829b1ff902e1a3e984a94c"},
]
[package.dependencies]
httpx = ">=0.23.0,<1"
orjson = ">=3.9.14,<4.0.0"
pydantic = ">=1,<3"
pydantic = [
{version = ">=1,<3", markers = "python_full_version < \"3.12.4\""},
{version = ">=2.7.4,<3.0.0", markers = "python_full_version >= \"3.12.4\""},
]
requests = ">=2,<3"
[[package]]
@@ -3051,6 +3055,7 @@ files = [
{file = "PyYAML-6.0.1-cp311-cp311-win_amd64.whl", hash = "sha256:bf07ee2fef7014951eeb99f56f39c9bb4af143d8aa3c21b1677805985307da34"},
{file = "PyYAML-6.0.1-cp312-cp312-macosx_10_9_x86_64.whl", hash = "sha256:855fb52b0dc35af121542a76b9a84f8d1cd886ea97c84703eaa6d88e37a2ad28"},
{file = "PyYAML-6.0.1-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:40df9b996c2b73138957fe23a16a4f0ba614f4c0efce1e9406a184b6d07fa3a9"},
{file = "PyYAML-6.0.1-cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:a08c6f0fe150303c1c6b71ebcd7213c2858041a7e01975da3a99aed1e7a378ef"},
{file = "PyYAML-6.0.1-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:6c22bec3fbe2524cde73d7ada88f6566758a8f7227bfbf93a408a9d86bcc12a0"},
{file = "PyYAML-6.0.1-cp312-cp312-musllinux_1_1_x86_64.whl", hash = "sha256:8d4e9c88387b0f5c7d5f281e55304de64cf7f9c0021a3525bd3b1c542da3b0e4"},
{file = "PyYAML-6.0.1-cp312-cp312-win32.whl", hash = "sha256:d483d2cdf104e7c9fa60c544d92981f12ad66a457afae824d146093b8c294c54"},
libs/langgraph/tests/__snapshots__/test_pregel.ambr
+39
View File
@@ -225,6 +225,7 @@
"id": [
"langgraph",
"utils",
"runnable",
"RunnableCallable"
],
"name": "left"
@@ -237,6 +238,7 @@
"id": [
"langgraph",
"utils",
"runnable",
"RunnableCallable"
],
"name": "right"
@@ -306,6 +308,7 @@
"id": [
"langgraph",
"utils",
"runnable",
"RunnableCallable"
],
"name": "left"
@@ -318,6 +321,7 @@
"id": [
"langgraph",
"utils",
"runnable",
"RunnableCallable"
],
"name": "right"
@@ -387,6 +391,7 @@
"id": [
"langgraph",
"utils",
"runnable",
"RunnableCallable"
],
"name": "get_weather"
@@ -800,6 +805,7 @@
"id": [
"langgraph",
"utils",
"runnable",
"RunnableCallable"
],
"name": "tools"
@@ -949,6 +955,7 @@
"id": [
"langgraph",
"utils",
"runnable",
"RunnableCallable"
],
"name": "tools"
@@ -1081,6 +1088,7 @@
'id': list([
'langgraph',
'utils',
'runnable',
'RunnableCallable',
]),
'name': 'tools',
@@ -1151,6 +1159,7 @@
"id": [
"langgraph",
"utils",
"runnable",
"RunnableCallable"
],
"name": "tools"
@@ -1300,6 +1309,7 @@
"id": [
"langgraph",
"utils",
"runnable",
"RunnableCallable"
],
"name": "tools"
@@ -1432,6 +1442,7 @@
'id': list([
'langgraph',
'utils',
'runnable',
'RunnableCallable',
]),
'name': 'tools',
@@ -1502,6 +1513,7 @@
"id": [
"langgraph",
"utils",
"runnable",
"RunnableCallable"
],
"name": "tools"
@@ -1651,6 +1663,7 @@
"id": [
"langgraph",
"utils",
"runnable",
"RunnableCallable"
],
"name": "tools"
@@ -1783,6 +1796,7 @@
'id': list([
'langgraph',
'utils',
'runnable',
'RunnableCallable',
]),
'name': 'tools',
@@ -1853,6 +1867,7 @@
"id": [
"langgraph",
"utils",
"runnable",
"RunnableCallable"
],
"name": "tools"
@@ -2002,6 +2017,7 @@
"id": [
"langgraph",
"utils",
"runnable",
"RunnableCallable"
],
"name": "tools"
@@ -2134,6 +2150,7 @@
'id': list([
'langgraph',
'utils',
'runnable',
'RunnableCallable',
]),
'name': 'tools',
@@ -2204,6 +2221,7 @@
"id": [
"langgraph",
"utils",
"runnable",
"RunnableCallable"
],
"name": "tools"
@@ -2353,6 +2371,7 @@
"id": [
"langgraph",
"utils",
"runnable",
"RunnableCallable"
],
"name": "tools"
@@ -2485,6 +2504,7 @@
'id': list([
'langgraph',
'utils',
'runnable',
'RunnableCallable',
]),
'name': 'tools',
@@ -2642,6 +2662,7 @@
"id": [
"langgraph",
"utils",
"runnable",
"RunnableCallable"
],
"name": "tools"
@@ -2723,6 +2744,7 @@
"id": [
"langgraph",
"utils",
"runnable",
"RunnableCallable"
],
"name": "tools"
@@ -2804,6 +2826,7 @@
"id": [
"langgraph",
"utils",
"runnable",
"RunnableCallable"
],
"name": "tools"
@@ -2885,6 +2908,7 @@
"id": [
"langgraph",
"utils",
"runnable",
"RunnableCallable"
],
"name": "tools"
@@ -2966,6 +2990,7 @@
"id": [
"langgraph",
"utils",
"runnable",
"RunnableCallable"
],
"name": "tools"
@@ -3028,6 +3053,7 @@
"id": [
"langgraph",
"utils",
"runnable",
"RunnableCallable"
],
"name": "A"
@@ -3040,6 +3066,7 @@
"id": [
"langgraph",
"utils",
"runnable",
"RunnableCallable"
],
"name": "B"
@@ -4661,6 +4688,7 @@
'id': list([
'langgraph',
'utils',
'runnable',
'RunnableCallable',
]),
'name': 'tool_one',
@@ -4678,6 +4706,7 @@
'id': list([
'langgraph',
'utils',
'runnable',
'RunnableCallable',
]),
'name': 'tool_two:tool_two_slow',
@@ -4690,6 +4719,7 @@
'id': list([
'langgraph',
'utils',
'runnable',
'RunnableCallable',
]),
'name': 'tool_two:tool_two_fast',
@@ -4707,6 +4737,7 @@
'id': list([
'langgraph',
'utils',
'runnable',
'RunnableCallable',
]),
'name': 'tool_three',
@@ -5264,6 +5295,7 @@
'id': list([
'langgraph',
'utils',
'runnable',
'RunnableCallable',
]),
'name': 'ask_question',
@@ -5276,6 +5308,7 @@
'id': list([
'langgraph',
'utils',
'runnable',
'RunnableCallable',
]),
'name': 'answer_question',
@@ -5323,6 +5356,7 @@
'id': list([
'langgraph',
'utils',
'runnable',
'RunnableCallable',
]),
'name': 'generate_analysts',
@@ -5348,6 +5382,7 @@
'id': list([
'langgraph',
'utils',
'runnable',
'RunnableCallable',
]),
'name': 'generate_sections',
@@ -5413,6 +5448,7 @@
'id': list([
'langgraph',
'utils',
'runnable',
'RunnableCallable',
]),
'name': 'generate_analysts',
@@ -5430,6 +5466,7 @@
'id': list([
'langgraph',
'utils',
'runnable',
'RunnableCallable',
]),
'name': 'conduct_interview:ask_question',
@@ -5442,6 +5479,7 @@
'id': list([
'langgraph',
'utils',
'runnable',
'RunnableCallable',
]),
'name': 'conduct_interview:answer_question',
@@ -5459,6 +5497,7 @@
'id': list([
'langgraph',
'utils',
'runnable',
'RunnableCallable',
]),
'name': 'generate_sections',
libs/langgraph/tests/test_pregel.py
+12 -14
View File
@@ -1848,9 +1848,7 @@ def test_invoke_two_processes_one_in_two_out(mocker: MockerFixture) -> None:
add_one = mocker.Mock(side_effect=lambda x: x + 1)
one = (
Channel.subscribe_to("input")
| add_one
| Channel.write_to(output=RunnablePassthrough(), between=RunnablePassthrough())
Channel.subscribe_to("input") | add_one | Channel.write_to("output", "between")
)
two = Channel.subscribe_to("between") | add_one | Channel.write_to("output")
@@ -4824,7 +4822,7 @@ def test_message_graph(
content="result for query",
name="search_api",
tool_call_id="tool_call123",
id="00000000-0000-4000-8000-000000000011",
id="00000000-0000-4000-8000-000000000010",
),
AIMessage(
content="",
@@ -4841,7 +4839,7 @@ def test_message_graph(
content="result for another",
name="search_api",
tool_call_id="tool_call456",
id="00000000-0000-4000-8000-000000000020",
id="00000000-0000-4000-8000-000000000018",
),
AIMessage(content="answer", id="ai3"),
]
@@ -4866,7 +4864,7 @@ def test_message_graph(
content="result for query",
name="search_api",
tool_call_id="tool_call123",
id="00000000-0000-4000-8000-000000000036",
id="00000000-0000-4000-8000-000000000033",
)
]
},
@@ -4889,7 +4887,7 @@ def test_message_graph(
content="result for another",
name="search_api",
tool_call_id="tool_call456",
id="00000000-0000-4000-8000-000000000045",
id="00000000-0000-4000-8000-000000000041",
)
]
},
@@ -5558,7 +5556,7 @@ def test_root_graph(
content="result for query",
name="search_api",
tool_call_id="tool_call123",
id="00000000-0000-4000-8000-000000000011",
id="00000000-0000-4000-8000-000000000010",
),
AIMessage(
content="",
@@ -5575,7 +5573,7 @@ def test_root_graph(
content="result for another",
name="search_api",
tool_call_id="tool_call456",
id="00000000-0000-4000-8000-000000000020",
id="00000000-0000-4000-8000-000000000018",
),
AIMessage(content="answer", id="ai3"),
]
@@ -5600,7 +5598,7 @@ def test_root_graph(
content="result for query",
name="search_api",
tool_call_id="tool_call123",
id="00000000-0000-4000-8000-000000000036",
id="00000000-0000-4000-8000-000000000033",
)
]
},
@@ -5623,7 +5621,7 @@ def test_root_graph(
content="result for another",
name="search_api",
tool_call_id="tool_call456",
id="00000000-0000-4000-8000-000000000045",
id="00000000-0000-4000-8000-000000000041",
)
]
},
@@ -6223,7 +6221,7 @@ def test_root_graph(
"__root__": [
HumanMessage(
content="what is weather in sf",
id="00000000-0000-4000-8000-000000000077",
id="00000000-0000-4000-8000-000000000070",
),
AIMessage(
content="",
@@ -6239,12 +6237,12 @@ def test_root_graph(
ToolMessage(
content="result for a different query",
name="search_api",
id="00000000-0000-4000-8000-000000000091",
id="00000000-0000-4000-8000-000000000082",
tool_call_id="tool_call123",
),
AIMessage(content="answer", id="ai2"),
AIMessage(
content="an extra message", id="00000000-0000-4000-8000-000000000101"
content="an extra message", id="00000000-0000-4000-8000-000000000091"
),
HumanMessage(content="what is weather in la"),
],
libs/langgraph/tests/test_pregel_async.py
+2 -3
View File
@@ -514,6 +514,7 @@ async def test_cancel_graph_astream_events_v2(checkpointer_name: Optional[str])
if chunk["event"] == "on_chain_stream" and not chunk["parent_ids"]:
got_event = True
assert chunk["data"]["chunk"] == {"alittlewhile": {"value": 2}}
await asyncio.sleep(0.1)
break
# did break
@@ -2087,9 +2088,7 @@ async def test_invoke_two_processes_one_in_two_out(mocker: MockerFixture) -> Non
add_one = mocker.Mock(side_effect=lambda x: x + 1)
one = (
Channel.subscribe_to("input")
| add_one
| Channel.write_to(output=RunnablePassthrough(), between=RunnablePassthrough())
Channel.subscribe_to("input") | add_one | Channel.write_to("output", "between")
)
two = Channel.subscribe_to("between") | add_one | Channel.write_to("output")
libs/langgraph/tests/test_utils.py
+2 -6
View File
@@ -21,12 +21,8 @@ from typing_extensions import Annotated, NotRequired, Required
from langgraph.graph import END, StateGraph
from langgraph.graph.graph import CompiledGraph
from langgraph.utils import (
_is_optional_type,
get_field_default,
is_async_callable,
is_async_generator,
)
from langgraph.utils.fields import _is_optional_type, get_field_default
from langgraph.utils.runnable import is_async_callable, is_async_generator
pytestmark = pytest.mark.anyio