diff --git a/libs/langgraph/langgraph/pregel/main.py b/libs/langgraph/langgraph/pregel/main.py index 45d6399ff..295d12e29 100644 --- a/libs/langgraph/langgraph/pregel/main.py +++ b/libs/langgraph/langgraph/pregel/main.py @@ -12,6 +12,7 @@ from dataclasses import is_dataclass from functools import partial from inspect import isclass from typing import Any, Callable, Generic, Optional, Union, cast, get_type_hints +from unittest.mock import DEFAULT from uuid import UUID, uuid5 from langchain_core.globals import get_debug @@ -116,7 +117,7 @@ from langgraph.pregel._validate import validate_graph, validate_keys from langgraph.pregel._write import ChannelWrite, ChannelWriteEntry from langgraph.pregel.debug import get_bolded_text, get_colored_text, tasks_w_writes from langgraph.pregel.protocol import PregelProtocol, StreamChunk, StreamProtocol -from langgraph.runtime import Runtime +from langgraph.runtime import DEFAULT_RUNTIME, Runtime from langgraph.store.base import BaseStore from langgraph.types import ( All, @@ -2570,12 +2571,16 @@ class Pregel( if durability is not None or deprecated_checkpoint_during is not None: config[CONF][CONFIG_KEY_DURABILITY] = durability_ - config[CONF][CONFIG_KEY_RUNTIME] = Runtime( + runtime = Runtime( context=context, store=store, stream_writer=stream_writer, previous=None, ) + parent_runtime = config[CONF].get(CONFIG_KEY_RUNTIME, DEFAULT_RUNTIME) + runtime = parent_runtime.merge(runtime) + config[CONF][CONFIG_KEY_RUNTIME] = runtime + with SyncPregelLoop( input, stream=StreamProtocol(stream.put, stream_modes), @@ -2861,12 +2866,16 @@ class Pregel( if durability is not None or deprecated_checkpoint_during is not None: config[CONF][CONFIG_KEY_DURABILITY] = durability_ - config[CONF][CONFIG_KEY_RUNTIME] = Runtime( + runtime = Runtime( context=context, store=store, stream_writer=stream_writer, previous=None, ) + parent_runtime = config[CONF].get(CONFIG_KEY_RUNTIME, DEFAULT_RUNTIME) + runtime = parent_runtime.merge(runtime) + config[CONF][CONFIG_KEY_RUNTIME] = runtime + async with AsyncPregelLoop( input, stream=StreamProtocol(stream.put_nowait, stream_modes), diff --git a/libs/langgraph/tests/test_runtime.py b/libs/langgraph/tests/test_runtime.py index cfbf361b7..f7f30b0fd 100644 --- a/libs/langgraph/tests/test_runtime.py +++ b/libs/langgraph/tests/test_runtime.py @@ -7,16 +7,14 @@ from langgraph.graph import END, START, StateGraph from langgraph.runtime import Runtime, get_runtime -@dataclass -class Context: - api_key: str - - -class State(TypedDict): - message: str - - def test_injected_runtime() -> None: + @dataclass + class Context: + api_key: str + + class State(TypedDict): + message: str + def injected_runtime(state: State, runtime: Runtime[Context]) -> dict[str, Any]: return {"message": f"api key: {runtime.context.api_key}"} @@ -32,6 +30,13 @@ def test_injected_runtime() -> None: def test_context_runtime() -> None: + @dataclass + class Context: + api_key: str + + class State(TypedDict): + message: str + def context_runtime(state: State) -> dict[str, Any]: runtime = get_runtime(Context) return {"message": f"api key: {runtime.context.api_key}"} @@ -45,3 +50,56 @@ def test_context_runtime() -> None: {"message": "hello world"}, context=Context(api_key="sk_123456") ) assert result == {"message": "api key: sk_123456"} + + +def test_override_runtime() -> None: + @dataclass + class Context: + api_key: str + + runtime1 = Runtime(context=Context(api_key="abc")) + runtime = runtime1.override(context=Context(api_key="def")) + assert runtime.context.api_key == "def" + + +def test_merge_runtime() -> None: + @dataclass + class Context: + api_key: str + + runtime1 = Runtime(context=Context(api_key="abc")) + runtime2 = Runtime(context=Context(api_key="def")) + runtime = runtime1.merge(runtime2) + assert runtime.context.api_key == "def" + + +def test_runtime_propogated_to_subgraph() -> None: + @dataclass + class Context: + username: str + + class State(TypedDict, total=False): + foo: str + bar: str + + def subgraph_node_1(state: State, runtime: Runtime[Context]): + return {"bar": f"hi {runtime.context.username}!"} + + subgraph_builder = StateGraph(State, context_schema=Context) + subgraph_builder.add_node(subgraph_node_1) + subgraph_builder.set_entry_point("subgraph_node_1") + subgraph = subgraph_builder.compile() + + def main_node(state: State, runtime: Runtime[Context]): + return {"foo": f"hello {runtime.context.username}!"} + + builder = StateGraph(State, context_schema=Context) + builder.add_node(main_node) + builder.add_node("node_1", subgraph) + builder.set_entry_point("main_node") + builder.add_edge("main_node", "node_1") + graph = builder.compile() + + context = Context(username="Alice") + result = graph.invoke({}, context=context) + assert result == {"foo": "hello Alice!", "bar": "hi Alice!"}