Reduce perf impact of set_context (#4256)

- call it less often
- find the run from the run manager at callsite

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")

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],
+                            ]
                         )
                     ],
                 )

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,
             )
         )

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:

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:

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

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)

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

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)
 

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

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"),
         ],