diff --git a/libs/langgraph/langgraph/_internal/_runtime.py b/libs/langgraph/langgraph/_internal/_runtime.py deleted file mode 100644 index d3296e523..000000000 --- a/libs/langgraph/langgraph/_internal/_runtime.py +++ /dev/null @@ -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}), - ) diff --git a/libs/langgraph/langgraph/graph/state.py b/libs/langgraph/langgraph/graph/state.py index 40be8a8ee..745db69a9 100644 --- a/libs/langgraph/langgraph/graph/state.py +++ b/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: diff --git a/libs/langgraph/langgraph/pregel/_algo.py b/libs/langgraph/langgraph/pregel/_algo.py index adb8cdb64..8b9a433fb 100644 --- a/libs/langgraph/langgraph/pregel/_algo.py +++ b/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, ) diff --git a/libs/langgraph/langgraph/runtime.py b/libs/langgraph/langgraph/runtime.py index e793e007c..c819f1d4e 100644 --- a/libs/langgraph/langgraph/runtime.py +++ b/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(