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

Reduce from 9s to 4s on benchmark graph

Browse Source 
- Use a simpler version of RunnableSequence without tracing serialization
- Remove accepts_run_manager check in RunnableCallable
- Remove creation of ChannelWrite dynamically every time conditional edge runs
This commit is contained in:
Nuno Campos
2024-09-04 11:02:33 -07:00
parent f2e0dc1042
commit 535dbd5b06
12 changed files with 446 additions and 148 deletions
Download Patch File Download Diff File Expand all files Collapse all files
libs/langgraph/langgraph/graph/graph.py
+14 -8
View File
@@ -55,7 +55,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 +75,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 +86,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,7 +94,7 @@ 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)
@@ -105,10 +105,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 +124,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 +453,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
+9 -16
View File
@@ -45,10 +45,7 @@ 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.fields import get_field_default
from langgraph.utils.runnable import (
RunnableCallable,
coerce_to_runnable,
)
from langgraph.utils.runnable import coerce_to_runnable
logger = logging.getLogger(__name__)
@@ -534,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):
@@ -552,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
]
)
@@ -600,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):
@@ -630,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 = [
(
@@ -649,7 +642,7 @@ class CompiledStateGraph(CompiledGraph):
),
)
)
return ChannelWrite(writes, tags=[TAG_HIDDEN])
ChannelWrite.do_write(config, writes)
# attach branch publisher
schema = (
libs/langgraph/langgraph/pregel/__init__.py
+7 -15
View File
@@ -6,7 +6,6 @@ from functools import partial
from typing import (
Any,
AsyncIterator,
Awaitable,
Callable,
Dict,
Iterator,
@@ -28,7 +27,7 @@ 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,
@@ -101,12 +100,7 @@ from langgraph.utils.config import (
)
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:
@@ -171,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()
]
)
@@ -812,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(
@@ -976,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(
libs/langgraph/langgraph/pregel/algo.py
+2 -2
View File
@@ -297,7 +297,7 @@ def prepare_next_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}"
@@ -403,7 +403,7 @@ def prepare_next_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/read.py
+10 -8
View File
@@ -1,5 +1,6 @@
from __future__ import annotations
from functools import cached_property
from typing import (
Any,
AsyncIterator,
@@ -15,7 +16,6 @@ from langchain_core.runnables import (
Runnable,
RunnableConfig,
RunnablePassthrough,
RunnableSequence,
RunnableSerializable,
)
from langchain_core.runnables.base import Input, Other, Output, coerce_to_runnable
@@ -25,7 +25,7 @@ from langgraph.constants import CONFIG_KEY_READ
from langgraph.pregel.retry import RetryPolicy
from langgraph.pregel.write import ChannelWrite
from langgraph.utils.config import merge_configs
from langgraph.utils.runnable import RunnableCallable
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
+38 -56
View File
@@ -4,11 +4,9 @@ import asyncio
from typing import (
Any,
Callable,
List,
NamedTuple,
Optional,
Sequence,
Tuple,
TypeVar,
Union,
)
@@ -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/config.py
+3 -1
View File
@@ -13,6 +13,8 @@ def patch_configurable(
) -> RunnableConfig:
if config is None:
return {"configurable": patch}
elif "configurable" not in config:
return {**config, "configurable": patch}
else:
return {**config, "configurable": {**config["configurable"], **patch}}
@@ -130,7 +132,7 @@ def patch_config(
Returns:
RunnableConfig: The patched config.
"""
config = config or {}
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
libs/langgraph/langgraph/utils/runnable.py
+301 -20
View File
@@ -2,17 +2,18 @@ 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, Optional
from typing import Any, AsyncIterator, Awaitable, Callable, Iterator, Optional
from langchain_core.load.serializable import to_json_not_implemented
from langchain_core.runnables.base import (
Runnable,
RunnableConfig,
RunnableLambda,
RunnableLike,
RunnableParallel,
RunnableSequence,
)
from langchain_core.runnables.config import (
ensure_config,
@@ -21,7 +22,8 @@ from langchain_core.runnables.config import (
run_in_executor,
var_child_runnable_config,
)
from langchain_core.runnables.utils import accepts_config, accepts_run_manager
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
@@ -73,16 +75,13 @@ class RunnableCallable(Runnable):
self.func = func
if func is not None:
self.func_accepts_config = accepts_config(func)
self.func_accepts_run_manager = accepts_run_manager(func)
self.afunc = afunc
if afunc is not None:
self.afunc_accepts_config = accepts_config(afunc)
self.afunc_accepts_run_manager = accepts_run_manager(afunc)
self.config: Optional[RunnableConfig] = {"tags": tags} if tags else None
self.kwargs = kwargs
self.trace = trace
self.recurse = recurse
self.serialized = to_json_not_implemented(self)
def __repr__(self) -> str:
repr_args = {
@@ -102,13 +101,15 @@ class RunnableCallable(Runnable):
" 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(
self.serialized,
None,
input,
name=config.get("run_name") or self.get_name(),
run_id=config.pop("run_id", None),
@@ -117,10 +118,6 @@ class RunnableCallable(Runnable):
child_config = patch_config(config, callbacks=run_manager.get_child())
context = copy_context()
context.run(_set_config_context, child_config)
if self.func_accepts_config:
kwargs["config"] = config
if self.func_accepts_run_manager:
kwargs["run_manager"] = run_manager
ret = context.run(self.func, input, **kwargs)
except BaseException as e:
run_manager.on_chain_error(e)
@@ -129,8 +126,6 @@ class RunnableCallable(Runnable):
run_manager.on_chain_end(ret)
else:
context.run(_set_config_context, config)
if self.func_accepts_config:
kwargs["config"] = config
ret = context.run(self.func, input, **kwargs)
if isinstance(ret, Runnable) and self.recurse:
return ret.invoke(input, config)
@@ -142,12 +137,14 @@ class RunnableCallable(Runnable):
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(
self.serialized,
None,
input,
name=config.get("run_name") or self.name,
run_id=config.pop("run_id", None),
@@ -155,10 +152,6 @@ class RunnableCallable(Runnable):
try:
child_config = patch_config(config, callbacks=run_manager.get_child())
context.run(_set_config_context, child_config)
if self.afunc_accepts_config:
kwargs["config"] = config
if self.afunc_accepts_run_manager:
kwargs["run_manager"] = run_manager
coro = self.afunc(input, **kwargs)
if ASYNCIO_ACCEPTS_CONTEXT:
ret = await asyncio.create_task(coro, context=context)
@@ -171,8 +164,6 @@ class RunnableCallable(Runnable):
await run_manager.on_chain_end(ret)
else:
context.run(_set_config_context, config)
if self.afunc_accepts_config:
kwargs["config"] = config
if ASYNCIO_ACCEPTS_CONTEXT:
ret = await asyncio.create_task(
self.afunc(input, **kwargs), context=context
@@ -236,3 +227,293 @@ def coerce_to_runnable(thing: RunnableLike, *, name: str, trace: bool) -> Runnab
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
+1 -3
View File
@@ -2087,9 +2087,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")