Compare commits

..
Author SHA1 Message Date
William Fu-Hinthorn ec6ce0a0b0 update docstring 2025-12-30 20:56:34 +09:00
William Fu-Hinthorn bd52e723a6 feat: Support context in astream_events 2025-12-30 20:47:47 +09:00
4 changed files with 102 additions and 26 deletions
+57
View File
@@ -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,61 @@ 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["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 1.0.6"
version: The version of the event stream schema to use ("v2").
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,
@@ -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,
+42
View File
@@ -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"
+2 -26
View File
@@ -14,7 +14,6 @@ import functools
import logging
import os
import re
import ssl
import sys
import warnings
from collections.abc import AsyncIterator, Callable, Iterator, Mapping, Sequence
@@ -181,7 +180,6 @@ def get_client(
api_key: str | None = NOT_PROVIDED,
headers: Mapping[str, str] | None = None,
timeout: TimeoutTypes | None = None,
verify: VerifyTypes | None = NOT_PROVIDED,
) -> LangGraphClient:
"""Create and configure a LangGraphClient.
@@ -210,13 +208,6 @@ def get_client(
- float (total seconds)
- tuple `(connect, read, write, pool)` in seconds
Defaults: connect=5, read=300, write=300, pool=5.
verify:
SSL certificate verification. May be:
- `True`: verify SSL certificates using system CA bundle
- `False`: disable SSL verification (not recommended for production)
- `str`: path to a CA bundle file or directory
- `ssl.SSLContext`: custom SSL context for full control
- Not provided (default): uses httpx default
Returns:
LangGraphClient:
@@ -283,10 +274,7 @@ def get_client(
_registered_transports.append(transport)
if transport is None:
transport_kwargs: dict[str, Any] = {"retries": 5}
if verify is not NOT_PROVIDED:
transport_kwargs["verify"] = verify
transport = httpx.AsyncHTTPTransport(**transport_kwargs)
transport = httpx.AsyncHTTPTransport(retries=5)
client = httpx.AsyncClient(
base_url=url,
transport=transport,
@@ -3583,7 +3571,6 @@ def get_sync_client(
api_key: str | None = NOT_PROVIDED,
headers: Mapping[str, str] | None = None,
timeout: TimeoutTypes | None = None,
verify: VerifyTypes | None = NOT_PROVIDED,
) -> SyncLangGraphClient:
"""Get a synchronous LangGraphClient instance.
@@ -3601,12 +3588,6 @@ def get_sync_client(
Accepts an httpx.Timeout instance, a float (seconds), or a tuple of timeouts.
Tuple format is (connect, read, write, pool)
If not provided, defaults to connect=5s, read=300s, write=300s, and pool=5s.
verify: SSL certificate verification. Can be:
- `True`: verify SSL certificates using system CA bundle
- `False`: disable SSL verification (not recommended for production)
- `str`: path to a CA bundle file or directory
- `ssl.SSLContext`: custom SSL context for full control
- Not provided (default): uses httpx default
Returns:
SyncLangGraphClient: The top-level synchronous client for accessing AssistantsClient,
ThreadsClient, RunsClient, and CronClient.
@@ -3639,10 +3620,7 @@ def get_sync_client(
if url is None:
url = "http://localhost:8123"
transport_kwargs: dict[str, Any] = {"retries": 5}
if verify is not NOT_PROVIDED:
transport_kwargs["verify"] = verify
transport = httpx.HTTPTransport(**transport_kwargs)
transport = httpx.HTTPTransport(retries=5)
client = httpx.Client(
base_url=url,
transport=transport,
@@ -6907,5 +6885,3 @@ TimeoutTypes = (
| tuple[float | None, float | None, float | None, float | None]
| httpx.Timeout
)
VerifyTypes = bool | str | ssl.SSLContext