diff --git a/libs/langgraph/bench/fanout_to_subgraph.py b/libs/langgraph/bench/fanout_to_subgraph.py index 612d0fbf2..fc8c7133b 100644 --- a/libs/langgraph/bench/fanout_to_subgraph.py +++ b/libs/langgraph/bench/fanout_to_subgraph.py @@ -106,6 +106,7 @@ def fanout_to_subgraph_sync() -> StateGraph: if __name__ == "__main__": import asyncio import random + import time import uvloop @@ -123,4 +124,7 @@ if __name__ == "__main__": len([c async for c in graph.astream(input, config=config)]) uvloop.install() + start = time.time() asyncio.run(run()) + end = time.time() + print(f"Time taken: {end - start:.4f} seconds") diff --git a/libs/langgraph/langgraph/func/__init__.py b/libs/langgraph/langgraph/func/__init__.py index 8c734127c..92d8419d9 100644 --- a/libs/langgraph/langgraph/func/__init__.py +++ b/libs/langgraph/langgraph/func/__init__.py @@ -20,7 +20,7 @@ from typing import ( from langgraph.channels.ephemeral_value import EphemeralValue from langgraph.channels.last_value import LastValue from langgraph.checkpoint.base import BaseCheckpointSaver -from langgraph.constants import END, PREVIOUS, START, TAG_HIDDEN +from langgraph.constants import END, PREVIOUS, START from langgraph.pregel import Pregel from langgraph.pregel.call import ( P, @@ -434,8 +434,7 @@ class entrypoint: [ ChannelWriteEntry(END, mapper=_pluck_return_value), ChannelWriteEntry(PREVIOUS, mapper=_pluck_save_value), - ], - tags=[TAG_HIDDEN], + ] ) ], ) diff --git a/libs/langgraph/langgraph/graph/branch.py b/libs/langgraph/langgraph/graph/branch.py index fd2039f20..a1ae358d6 100644 --- a/libs/langgraph/langgraph/graph/branch.py +++ b/libs/langgraph/langgraph/graph/branch.py @@ -138,7 +138,6 @@ class Branch(NamedTuple): reader=reader, name=None, trace=False, - set_context=False, func_accepts_config=True, ) ) diff --git a/libs/langgraph/langgraph/graph/graph.py b/libs/langgraph/langgraph/graph/graph.py index d8bde68ec..fa28243fb 100644 --- a/libs/langgraph/langgraph/graph/graph.py +++ b/libs/langgraph/langgraph/graph/graph.py @@ -366,16 +366,14 @@ class CompiledGraph(Pregel): self.nodes[key] = ( PregelNode(channels=[], triggers=[], metadata=node.metadata) | node.runnable - | ChannelWrite([ChannelWriteEntry(key)], tags=[TAG_HIDDEN]) + | ChannelWrite([ChannelWriteEntry(key)]) ) cast(list[str], self.stream_channels).append(key) def attach_edge(self, start: str, end: str) -> None: if end == END: # publish to end channel - self.nodes[start].writers.append( - ChannelWrite([ChannelWriteEntry(END)], tags=[TAG_HIDDEN]) - ) + self.nodes[start].writers.append(ChannelWrite([ChannelWriteEntry(END)])) else: # subscribe to start channel self.nodes[end].triggers.append(start) @@ -393,10 +391,7 @@ class CompiledGraph(Pregel): ) for p in packets ] - return ChannelWrite( - cast(Sequence[Union[ChannelWriteEntry, Send]], writes), - tags=[TAG_HIDDEN], - ) + return ChannelWrite(cast(Sequence[Union[ChannelWriteEntry, Send]], writes)) # add hidden start node if start == START and start not in self.nodes: diff --git a/libs/langgraph/langgraph/graph/state.py b/libs/langgraph/langgraph/graph/state.py index 2316fc004..c4ac5ee54 100644 --- a/libs/langgraph/langgraph/graph/state.py +++ b/libs/langgraph/langgraph/graph/state.py @@ -808,7 +808,7 @@ class CompiledStateGraph(CompiledGraph): tags=[TAG_HIDDEN], triggers=[START], channels=[START], - writers=[ChannelWrite(write_entries, tags=[TAG_HIDDEN])], + writers=[ChannelWrite(write_entries)], ) elif node is not None: input_schema = node.input if node else self.builder.schema @@ -833,7 +833,7 @@ class CompiledStateGraph(CompiledGraph): # coerce state dict to schema class (eg. pydantic model) mapper=mapper, # publish to state keys - writers=[ChannelWrite(write_entries, tags=[TAG_HIDDEN])], + writers=[ChannelWrite(write_entries)], metadata=node.metadata, retry_policy=node.retry_policy, bound=node.runnable, @@ -859,9 +859,7 @@ class CompiledStateGraph(CompiledGraph): # publish to channel for start in starts: self.nodes[start].writers.append( - ChannelWrite( - (ChannelWriteEntry(channel_name, start),), tags=[TAG_HIDDEN] - ) + ChannelWrite((ChannelWriteEntry(channel_name, start),)) ) def attach_branch( @@ -933,9 +931,7 @@ class CompiledStateGraph(CompiledGraph): for end in ends: if end != END: self.nodes[end].writers.append( - ChannelWrite( - (ChannelWriteEntry(channel_name, end),), tags=[TAG_HIDDEN] - ) + ChannelWrite((ChannelWriteEntry(channel_name, end),)) ) def _migrate_checkpoint(self, checkpoint: Checkpoint) -> None: diff --git a/libs/langgraph/langgraph/pregel/__init__.py b/libs/langgraph/langgraph/pregel/__init__.py index c5d335d75..902317e20 100644 --- a/libs/langgraph/langgraph/pregel/__init__.py +++ b/libs/langgraph/langgraph/pregel/__init__.py @@ -2550,11 +2550,12 @@ class Pregel(PregelProtocol): do_stream = ( next( ( - cast(_StreamingCallbackHandler, h) + True for h in run_manager.handlers if isinstance(h, _StreamingCallbackHandler) + and not isinstance(h, StreamMessagesHandler) ), - None, + False, ) if _StreamingCallbackHandler is not None else False @@ -2624,7 +2625,7 @@ class Pregel(PregelProtocol): ), put_writes=weakref.WeakMethod(loop.put_writes), schedule_task=weakref.WeakMethod(loop.accept_push), - use_astream=do_stream is not None, + use_astream=do_stream, node_finished=config[CONF].get(CONFIG_KEY_NODE_FINISHED), ) # enable subgraph streaming diff --git a/libs/langgraph/langgraph/pregel/call.py b/libs/langgraph/langgraph/pregel/call.py index e1856aa80..a16fa36fd 100644 --- a/libs/langgraph/langgraph/pregel/call.py +++ b/libs/langgraph/langgraph/pregel/call.py @@ -10,7 +10,7 @@ from typing import Any, Callable, Generator, Generic, Optional, Sequence, TypeVa from langchain_core.runnables import Runnable from typing_extensions import ParamSpec -from langgraph.constants import CONF, CONFIG_KEY_CALL, RETURN, TAG_HIDDEN +from langgraph.constants import CONF, CONFIG_KEY_CALL, RETURN from langgraph.pregel.write import ChannelWrite, ChannelWriteEntry from langgraph.types import RetryPolicy from langgraph.utils.config import get_config @@ -197,7 +197,7 @@ def get_runnable_for_task(func: Callable[..., Any]) -> RunnableSeq: ) seq = RunnableSeq( run, - ChannelWrite([ChannelWriteEntry(RETURN)], tags=[TAG_HIDDEN]), + ChannelWrite([ChannelWriteEntry(RETURN)]), name=name, trace_inputs=functools.partial( _explode_args_trace_inputs, inspect.signature(func) diff --git a/libs/langgraph/langgraph/pregel/read.py b/libs/langgraph/langgraph/pregel/read.py index 5b345e811..e57620ec7 100644 --- a/libs/langgraph/langgraph/pregel/read.py +++ b/libs/langgraph/langgraph/pregel/read.py @@ -68,6 +68,7 @@ class ChannelRead(RunnableCallable): afunc=self._aread, tags=tags, name=None, + trace=False, func_accepts_config=True, ) self.fresh = fresh diff --git a/libs/langgraph/langgraph/pregel/write.py b/libs/langgraph/langgraph/pregel/write.py index 05fe5a388..234c1f5d7 100644 --- a/libs/langgraph/langgraph/pregel/write.py +++ b/libs/langgraph/langgraph/pregel/write.py @@ -54,14 +54,14 @@ class ChannelWrite(RunnableCallable): self, writes: Sequence[Union[ChannelWriteEntry, ChannelWriteTupleEntry, Send]], *, - tags: Optional[Sequence[str]] = None, + tags: Optional[Sequence[str]] = None, # ignored require_at_least_one_of: Optional[Sequence[str]] = None, # ignored ): super().__init__( func=self._write, afunc=self._awrite, name=None, - tags=tags, + trace=False, func_accepts_config=True, ) self.writes = cast( @@ -152,6 +152,8 @@ class ChannelWrite(RunnableCallable): tuples.append((w.channel, value)) else: raise ValueError(f"Invalid write entry: {w}") + # if we want to persist writes found before hitting a ParentCommand + # can move this to a finally block write: TYPE_SEND = config[CONF][CONFIG_KEY_SEND] write(tuples) diff --git a/libs/langgraph/langgraph/utils/runnable.py b/libs/langgraph/langgraph/utils/runnable.py index 2736cc7d5..c4ba527d6 100644 --- a/libs/langgraph/langgraph/utils/runnable.py +++ b/libs/langgraph/langgraph/utils/runnable.py @@ -36,6 +36,7 @@ from langchain_core.runnables.config import ( var_child_runnable_config, ) from langchain_core.runnables.utils import Input, Output +from langchain_core.tracers.langchain import LangChainTracer from typing_extensions import TypeGuard from langgraph.constants import ( @@ -60,58 +61,34 @@ except ImportError: def _set_config_context( - config: RunnableConfig, -) -> tuple[Token[Optional[RunnableConfig]], Optional[dict[str, Any]]]: + config: RunnableConfig, run: Any = None +) -> Token[Optional[RunnableConfig]]: """Set the child Runnable config + tracing context. Args: config (RunnableConfig): The config to set. """ - from langchain_core.tracers.langchain import LangChainTracer - config_token = var_child_runnable_config.set(config) - current_context = None - if ( - (callbacks := config.get("callbacks")) - and ( - parent_run_id := getattr(callbacks, "parent_run_id", None) - ) # Is callback manager - and ( - tracer := next( - ( - handler - for handler in getattr(callbacks, "handlers", []) - if isinstance(handler, LangChainTracer) - ), - None, - ) - ) - and (run := tracer.run_map.get(str(parent_run_id))) - ): - from langsmith.run_helpers import _set_tracing_context, get_tracing_context + if run is not None: + from langsmith.run_helpers import _set_tracing_context - current_context = get_tracing_context() _set_tracing_context({"parent": run}) - return config_token, current_context + return config_token -@contextmanager -def set_config_context(config: RunnableConfig) -> Generator[Context, None, None]: +def _unset_config_context( + token: Token[Optional[RunnableConfig]], run: Any = None +) -> None: """Set the child Runnable config + tracing context. Args: config (RunnableConfig): The config to set. """ - from langsmith.run_helpers import _set_tracing_context + var_child_runnable_config.reset(token) + if run is not None: + from langsmith.run_helpers import _set_tracing_context - ctx = copy_context() - config_token, _ = ctx.run(_set_config_context, config) - try: - yield ctx - finally: - ctx.run(var_child_runnable_config.reset, config_token) - ctx.run( - _set_tracing_context, + _set_tracing_context( { "parent": None, "project_name": None, @@ -119,10 +96,27 @@ def set_config_context(config: RunnableConfig) -> Generator[Context, None, None] "metadata": None, "enabled": None, "client": None, - }, + } ) +@contextmanager +def set_config_context( + config: RunnableConfig, run: Any = None +) -> Generator[Context, None, None]: + """Set the child Runnable config + tracing context. + + Args: + config (RunnableConfig): The config to set. + """ + ctx = copy_context() + config_token = ctx.run(_set_config_context, config, run) + try: + yield ctx + finally: + ctx.run(_unset_config_context, config_token, run) + + # Before Python 3.11 native StrEnum is not available class StrEnum(str, enum.Enum): """A string enum.""" @@ -254,7 +248,6 @@ class RunnableCallable(Runnable): tags: Optional[Sequence[str]] = None, trace: bool = True, recurse: bool = True, - set_context: bool = True, explode_args: bool = False, func_accepts_config: Optional[bool] = None, **kwargs: Any, @@ -278,7 +271,6 @@ class RunnableCallable(Runnable): self.kwargs = kwargs self.trace = trace self.recurse = recurse - self.set_context = set_context self.explode_args = explode_args # check signature if func is None and afunc is None: @@ -365,19 +357,21 @@ class RunnableCallable(Runnable): ) try: child_config = patch_config(config, callbacks=run_manager.get_child()) - if self.set_context: - with set_config_context(child_config) as context: - ret = context.run(self.func, *args, **kwargs) + # get the run + for h in run_manager.handlers: + if isinstance(h, LangChainTracer): + run = h.run_map.get(str(run_manager.run_id)) + break else: - ret = self.func(*args, **kwargs) + run = None + # run in context + with set_config_context(child_config, run) as context: + ret = context.run(self.func, *args, **kwargs) except BaseException as e: run_manager.on_chain_error(e) raise else: run_manager.on_chain_end(ret) - elif self.set_context: - with set_config_context(config) as context: - ret = context.run(self.func, *args, **kwargs) else: ret = self.func(*args, **kwargs) if self.recurse and isinstance(ret, Runnable): @@ -425,8 +419,14 @@ class RunnableCallable(Runnable): try: child_config = patch_config(config, callbacks=run_manager.get_child()) coro = cast(Coroutine[None, None, Any], self.afunc(*args, **kwargs)) - if ASYNCIO_ACCEPTS_CONTEXT and self.set_context: - with set_config_context(child_config) as context: + if ASYNCIO_ACCEPTS_CONTEXT: + for h in run_manager.handlers: + if isinstance(h, LangChainTracer): + run = h.run_map.get(str(run_manager.run_id)) + break + else: + run = None + with set_config_context(child_config, run) as context: ret = await asyncio.create_task(coro, context=context) else: ret = await coro @@ -435,10 +435,6 @@ class RunnableCallable(Runnable): raise else: await run_manager.on_chain_end(ret) - elif ASYNCIO_ACCEPTS_CONTEXT and self.set_context: - with set_config_context(config) as context: - coro = cast(Coroutine[None, None, Any], self.afunc(*args, **kwargs)) - ret = await asyncio.create_task(coro, context=context) else: ret = await self.afunc(*args, **kwargs) if self.recurse and isinstance(ret, Runnable): @@ -604,7 +600,6 @@ class RunnableSeq(Runnable): 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): @@ -612,8 +607,19 @@ class RunnableSeq(Runnable): config = patch_config( config, callbacks=run_manager.get_child(f"seq:step:{i + 1}") ) + # 1st step is the actual node, + # others are writers which don't need to be run in context if i == 0: - input = step.invoke(input, config, **kwargs) + # get the run object + for h in run_manager.handlers: + if isinstance(h, LangChainTracer): + run = h.run_map.get(str(run_manager.run_id)) + break + else: + run = None + # run in context + with set_config_context(config, run) as context: + input = context.run(step.invoke, input, config, **kwargs) else: input = step.invoke(input, config) # finish the root run @@ -649,8 +655,24 @@ class RunnableSeq(Runnable): config = patch_config( config, callbacks=run_manager.get_child(f"seq:step:{i + 1}") ) + # 1st step is the actual node, + # others are writers which don't need to be run in context if i == 0: - input = await step.ainvoke(input, config, **kwargs) + if ASYNCIO_ACCEPTS_CONTEXT: + # get the run object + for h in run_manager.handlers: + if isinstance(h, LangChainTracer): + run = h.run_map.get(str(run_manager.run_id)) + break + else: + run = None + # run in context + with set_config_context(config, run) as context: + input = await asyncio.create_task( + step.ainvoke(input, config, **kwargs), context=context + ) + else: + input = await step.ainvoke(input, config, **kwargs) else: input = await step.ainvoke(input, config) # finish the root run @@ -678,53 +700,48 @@ class RunnableSeq(Runnable): 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 _StreamingCallbackHandler is not None and ( - stream_handler := next( - ( - cast(_StreamingCallbackHandler, 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: Any = 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 + # get the run object + for h in run_manager.handlers: + if isinstance(h, LangChainTracer): + run = h.run_map.get(str(run_manager.run_id)) + break else: - run_manager.on_chain_end(output) + run = None + # create first step config + config = patch_config( + config, + callbacks=run_manager.get_child(f"seq:step:{1}"), + ) + # run all in context + with set_config_context(config, run) as context: + 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): + if idx == 0: + iterator = step.stream(input, config, **kwargs) + else: + config = patch_config( + config, + callbacks=run_manager.get_child(f"seq:step:{idx + 1}"), + ) + iterator = step.transform(iterator, config) + # populates streamed_output in astream_log() output if needed + if _StreamingCallbackHandler is not None: + for h in run_manager.handlers: + if isinstance(h, _StreamingCallbackHandler): + iterator = h.tap_output_iter(run_manager.run_id, iterator) + # consume into final output + output = context.run(_consume_iter, iterator) + # sequence doesn't emit output, yield to mark as generator + yield + except BaseException as e: + run_manager.on_chain_error(e) + raise + else: + run_manager.on_chain_end(output) async def astream( self, @@ -743,53 +760,121 @@ class RunnableSeq(Runnable): 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 _StreamingCallbackHandler is not None and ( - stream_handler := next( - ( - cast(_StreamingCallbackHandler, 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: Any = 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 + # 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 + if ASYNCIO_ACCEPTS_CONTEXT: + # get the run object + for h in run_manager.handlers: + if isinstance(h, LangChainTracer): + run = h.run_map.get(str(run_manager.run_id)) + break + else: + run = None + # create first step config + config = patch_config( + config, + callbacks=run_manager.get_child(f"seq:step:{1}"), + ) + # run all in context + with set_config_context(config, run) as context: + try: + async with AsyncExitStack() as stack: + for idx, step in enumerate(self.steps): + if idx == 0: + aiterator = step.astream(input, config, **kwargs) + else: + config = patch_config( + config, + callbacks=run_manager.get_child( + f"seq:step:{idx + 1}" + ), + ) + aiterator = step.atransform(aiterator, config) + if hasattr(aiterator, "aclose"): + stack.push_async_callback(aiterator.aclose) + # populates streamed_output in astream_log() output if needed + if _StreamingCallbackHandler is not None: + for h in run_manager.handlers: + if isinstance(h, _StreamingCallbackHandler): + aiterator = h.tap_output_aiter( + run_manager.run_id, aiterator + ) + # consume into final output + output = await asyncio.create_task( + _consume_aiter(aiterator), context=context + ) + # sequence doesn't emit output, yield to mark as generator + yield + except BaseException as e: + await run_manager.on_chain_error(e) + raise + else: + await run_manager.on_chain_end(output) else: - await run_manager.on_chain_end(output) + try: + async with AsyncExitStack() as stack: + 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) + # populates streamed_output in astream_log() output if needed + if _StreamingCallbackHandler is not None: + for h in run_manager.handlers: + if isinstance(h, _StreamingCallbackHandler): + aiterator = h.tap_output_aiter( + run_manager.run_id, aiterator + ) + # consume into final output + output = await _consume_aiter(aiterator) + # sequence doesn't emit output, yield to mark as generator + yield + except BaseException as e: + await run_manager.on_chain_error(e) + raise + else: + await run_manager.on_chain_end(output) + + +def _consume_iter(it: Iterator[Any]) -> Any: + """Consume an iterator.""" + output: Any = None + add_supported = False + for chunk in it: + # 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 + return output + + +async def _consume_aiter(it: AsyncIterator[Any]) -> Any: + """Consume an async iterator.""" + output: Any = None + add_supported = False + async for chunk in it: + # collect final output + if add_supported: + try: + output = output + chunk + except TypeError: + output = chunk + add_supported = False + else: + output = chunk + return output diff --git a/libs/langgraph/tests/test_large_cases.py b/libs/langgraph/tests/test_large_cases.py index bb8e124ff..7bba55ff5 100644 --- a/libs/langgraph/tests/test_large_cases.py +++ b/libs/langgraph/tests/test_large_cases.py @@ -4660,7 +4660,7 @@ def test_root_graph( content="result for query", name="search_api", tool_call_id="tool_call123", - id="00000000-0000-4000-8000-000000000033", + id="00000000-0000-4000-8000-000000000024", ) ] }, @@ -4683,7 +4683,7 @@ def test_root_graph( content="result for another", name="search_api", tool_call_id="tool_call456", - id="00000000-0000-4000-8000-000000000041", + id="00000000-0000-4000-8000-000000000030", ) ] }, @@ -5387,7 +5387,7 @@ def test_root_graph( "__root__": [ HumanMessage( content="what is weather in sf", - id="00000000-0000-4000-8000-000000000070", + id="00000000-0000-4000-8000-000000000051", ), AIMessage( content="", @@ -5407,7 +5407,7 @@ def test_root_graph( ), AIMessage(content="answer", id="ai2"), AIMessage( - content="an extra message", id="00000000-0000-4000-8000-000000000091" + content="an extra message", id="00000000-0000-4000-8000-000000000066" ), HumanMessage(content="what is weather in la"), ],