diff --git a/libs/langgraph/langgraph/pregel/main.py b/libs/langgraph/langgraph/pregel/main.py index 37e8125f9..a6375884f 100644 --- a/libs/langgraph/langgraph/pregel/main.py +++ b/libs/langgraph/langgraph/pregel/main.py @@ -22,6 +22,7 @@ from inspect import isclass from typing import ( Any, Generic, + Literal, cast, get_type_hints, ) @@ -38,6 +39,7 @@ from langchain_core.runnables.config import ( get_callback_manager_for_config, ) from langchain_core.runnables.graph import Graph +from langchain_core.runnables.schema import CustomStreamEvent, StandardStreamEvent from langgraph.cache.base import BaseCache from langgraph.checkpoint.base import ( BaseCheckpointSaver, @@ -3021,6 +3023,62 @@ class Pregel( await asyncio.shield(run_manager.on_chain_error(e)) raise + async def astream_events( + self, + input: InputT | Command | None, + config: RunnableConfig | None = None, + *, + context: ContextT | None = None, + version: Literal["v1", "v2"], + include_names: Sequence[str] | None = None, + include_types: Sequence[str] | None = None, + include_tags: Sequence[str] | None = None, + exclude_names: Sequence[str] | None = None, + exclude_types: Sequence[str] | None = None, + exclude_tags: Sequence[str] | None = None, + **kwargs: Any, + ) -> AsyncIterator[StandardStreamEvent | CustomStreamEvent]: + """Stream events from the graph execution. + + This method extends the base Runnable.astream_events with support for + the LangGraph `context` parameter. + + Args: + input: The input to the graph. + config: The configuration to use for the run. + context: The static context to use for the run. + !!! version-added "Added in version 0.6.0" + version: The version of the event stream schema to use ("v1" or "v2"). + Use "v2" for the latest version. + include_names: Only include events from runnables with matching names. + include_types: Only include events from runnables with matching types. + include_tags: Only include events from runnables with matching tags. + exclude_names: Exclude events from runnables with matching names. + exclude_types: Exclude events from runnables with matching types. + exclude_tags: Exclude events from runnables with matching tags. + **kwargs: Additional arguments passed to the underlying stream. + + Yields: + Events from the graph execution. + """ + async with contextlib.aclosing( + super().astream_events( + input, + config, + version=version, + include_names=include_names, + include_types=include_types, + include_tags=include_tags, + exclude_names=exclude_names, + exclude_types=exclude_types, + exclude_tags=exclude_tags, + context=context, + **kwargs, + ) + ) as stream: # type: ignore[type-var] + async for event in stream: + yield event + def invoke( self, input: InputT | Command | None, diff --git a/libs/langgraph/langgraph/pregel/remote.py b/libs/langgraph/langgraph/pregel/remote.py index 2535d966a..ba8621a2a 100644 --- a/libs/langgraph/langgraph/pregel/remote.py +++ b/libs/langgraph/langgraph/pregel/remote.py @@ -907,6 +907,7 @@ class RemoteGraph(PregelProtocol): input: Any, config: RunnableConfig | None = None, *, + context: Any = None, version: Literal["v1", "v2"], include_names: Sequence[All] | None = None, include_types: Sequence[All] | None = None, diff --git a/libs/langgraph/tests/test_runtime.py b/libs/langgraph/tests/test_runtime.py index 0407b84d2..eb0622897 100644 --- a/libs/langgraph/tests/test_runtime.py +++ b/libs/langgraph/tests/test_runtime.py @@ -389,3 +389,45 @@ def test_context_coercion_pydantic_validation_errors() -> None: compiled.invoke( {"message": "test"}, context={"api_key": "sk_test", "timeout": "not_an_int"} ) + + +@pytest.mark.anyio +async def test_context_with_astream_events() -> None: + """Test that context is properly passed through astream_events.""" + + @dataclass + class Context: + api_key: str + + class State(TypedDict): + message: str + + def node_with_context(state: State, runtime: Runtime[Context]) -> dict[str, Any]: + return {"message": f"api_key: {runtime.context.api_key}"} + + graph = StateGraph(state_schema=State, context_schema=Context) + graph.add_node("node", node_with_context) + graph.add_edge(START, "node") + graph.add_edge("node", END) + compiled = graph.compile() + + events = [] + async for event in compiled.astream_events( + {"message": "test"}, + version="v2", + context=Context(api_key="sk_events_123"), + ): + events.append(event) + + # Verify we got events + assert len(events) > 0 + + # Find the final on_chain_end event (no parent_ids means it's the root) + end_events = [ + e for e in events if e["event"] == "on_chain_end" and not e.get("parent_ids") + ] + assert len(end_events) == 1 + + # Verify the output contains our context value + output = end_events[0]["data"]["output"] + assert output["message"] == "api_key: sk_events_123"