refactor(langgraph): improve Runtime interface re patch/overrides (#5546)

libs/langgraph/langgraph/_internal/_runtime.py

@@ -1,39 +0,0 @@
-"""Internal utilities for the Runtime class."""
-
-from __future__ import annotations
-
-from dataclasses import replace
-from typing import Any, cast
-
-from typing_extensions import TypedDict, Unpack
-
-from langgraph.runtime import Runtime
-from langgraph.store.base import BaseStore
-from langgraph.types import StreamWriter
-
-
-class RuntimePatch(TypedDict, total=False):
-    """Patch structure for the Runtime class."""
-
-    context: Any
-    store: BaseStore | None
-    stream_writer: StreamWriter
-    previous: Any
-
-
-def patch_runtime(runtime: Runtime, **overrides: Unpack[RuntimePatch]) -> Runtime:
-    """Patch the runtime with the given overrides, returning a new instance."""
-    return replace(runtime, **overrides)
-
-
-def patch_runtime_non_null(
-    runtime: Runtime, **overrides: Unpack[RuntimePatch]
-) -> Runtime:
-    """Patch the runtime with the given overrides, returning a new instance.
-
-    Only patch fields with overrides that are not None.
-    """
-    return replace(
-        runtime,
-        **cast(dict[str, Any], {k: v for k, v in overrides.items() if v is not None}),
-    )

libs/langgraph/langgraph/graph/state.py

@@ -135,7 +135,7 @@ class StateGraph(Generic[StateT, ContextT, InputT, OutputT]):
         from typing_extensions import Annotated, TypedDict
         from langgraph.checkpoint.memory import MemorySaver
         from langgraph.graph import StateGraph
-        from langgraph.rumtime import Runtime
+        from langgraph.runtime import Runtime
 
         def reducer(a: list, b: int | None) -> list:
             if b is not None:

libs/langgraph/langgraph/pregel/_algo.py

@@ -53,7 +53,6 @@ from langgraph._internal._constants import (
     RETURN,
     TASKS,
 )
-from langgraph._internal._runtime import patch_runtime_non_null
 from langgraph._internal._typing import EMPTY_SEQ, MISSING
 from langgraph.channels.base import BaseChannel
 from langgraph.channels.topic import Topic
@@ -71,7 +70,7 @@ from langgraph.pregel._io import read_channels
 from langgraph.pregel._log import logger
 from langgraph.pregel._read import INPUT_CACHE_KEY_TYPE, PregelNode
 from langgraph.pregel._scratchpad import PregelScratchpad
-from langgraph.runtime import DEFAULT_RUNTIME
+from langgraph.runtime import DEFAULT_RUNTIME, Runtime
 from langgraph.store.base import BaseStore
 from langgraph.types import (
     All,
@@ -583,10 +582,10 @@ def prepare_single_task(
                 step,
                 stop,
             )
-            runtime = patch_runtime_non_null(
-                configurable.get(CONFIG_KEY_RUNTIME, DEFAULT_RUNTIME),
-                store=store,
+            runtime = cast(
+                Runtime, configurable.get(CONFIG_KEY_RUNTIME, DEFAULT_RUNTIME)
             )
+            runtime = runtime.override(store=store)
             return PregelExecutableTask(
                 name,
                 call.input,
@@ -713,10 +712,11 @@ def prepare_single_task(
                 step,
                 stop,
             )
-            runtime = patch_runtime_non_null(
-                configurable.get(CONFIG_KEY_RUNTIME, DEFAULT_RUNTIME),
-                store=store,
-                previous=checkpoint["channel_values"].get(PREVIOUS, None),
+            runtime = cast(
+                Runtime, configurable.get(CONFIG_KEY_RUNTIME, DEFAULT_RUNTIME)
+            )
+            runtime = runtime.override(
+                store=store, previous=checkpoint["channel_values"].get(PREVIOUS, None)
             )
             return PregelExecutableTask(
                 packet.node,
@@ -852,8 +852,10 @@ def prepare_single_task(
                         )
                     else:
                         cache_key = None
-                    runtime = patch_runtime_non_null(
-                        configurable.get(CONFIG_KEY_RUNTIME, DEFAULT_RUNTIME),
+                    runtime = cast(
+                        Runtime, configurable.get(CONFIG_KEY_RUNTIME, DEFAULT_RUNTIME)
+                    )
+                    runtime = runtime.override(
                         previous=checkpoint["channel_values"].get(PREVIOUS, None),
                         store=store,
                     )

libs/langgraph/langgraph/runtime.py

@@ -1,8 +1,10 @@
 from __future__ import annotations
 
-from dataclasses import dataclass
+from dataclasses import dataclass, field, replace
 from typing import Any, Generic, cast
 
+from typing_extensions import TypedDict, Unpack
+
 from langgraph._internal._constants import CONF, CONFIG_KEY_RUNTIME
 from langgraph.config import get_config
 from langgraph.store.base import BaseStore
@@ -13,6 +15,13 @@ from langgraph.typing import ContextT
 def _no_op_stream_writer(_: Any) -> None: ...
 
 
+class _RuntimeOverrides(TypedDict, Generic[ContextT], total=False):
+    context: ContextT
+    store: BaseStore | None
+    stream_writer: StreamWriter
+    previous: Any
+
+
 @dataclass(**_DC_KWARGS)
 class Runtime(Generic[ContextT]):
     """Convenience class that bundles run-scoped context and graph configuration.
@@ -20,21 +29,42 @@ class Runtime(Generic[ContextT]):
     !!! version-added "Added in version 1.0.0."
     """
 
-    context: ContextT
+    context: ContextT = field(default=None)  # type: ignore[assignment]
     """Static context for the graph run, like user_id, db_conn, etc.
     
     Can also be thought of as 'run dependencies'."""
 
-    store: BaseStore | None
+    store: BaseStore | None = field(default=None)
     """Store for the graph run, enabling persistence and memory."""
 
-    stream_writer: StreamWriter
+    stream_writer: StreamWriter = field(default=_no_op_stream_writer)
     """Function that writes to the custom stream."""
 
-    previous: Any | None
+    previous: Any = field(default=None)
     """The previous return value for the given thread.
     
-    Only available with the functional API when a checkpointer is provided."""
+    Only available with the functional API when a checkpointer is provided.
+    """
+
+    def merge(self, other: Runtime[ContextT]) -> Runtime[ContextT]:
+        """Merge two runtimes together.
+
+        If a value is not provided in the other runtime, the value from the current runtime is used.
+        """
+        return Runtime(
+            context=other.context or self.context,
+            store=other.store or self.store,
+            stream_writer=other.stream_writer
+            if other.stream_writer is not _no_op_stream_writer
+            else self.stream_writer,
+            previous=other.previous or self.previous,
+        )
+
+    def override(
+        self, **overrides: Unpack[_RuntimeOverrides[ContextT]]
+    ) -> Runtime[ContextT]:
+        """Replace the runtime with a new runtime with the given overrides."""
+        return replace(self, **overrides)
 
 
 DEFAULT_RUNTIME = Runtime(