mirror of
https://github.com/langchain-ai/langgraph.git
synced 2026-08-25 09:02:25 +02:00
open-swe[bot]GitHubopen-swe[bot] <open-swe@users.noreply.github.com>Sydney RunkleSydney RunkleEugene YurtsevFixes #5735 Implement context coercion functionality for LangGraph runtime to improve API usability. Key changes: - Added `_coerce_context` function in `pregel/main.py` - Supports coercion for: - Pydantic BaseModel - Dataclasses - TypedDict - Comprehensive test coverage added in `tests/test_runtime.py` - Handles edge cases like None context and missing fields The implementation allows users to pass dictionaries as context, which will be automatically converted to the expected schema type, making the API more flexible and user-friendly. --------- Co-authored-by: open-swe[bot] <open-swe@users.noreply.github.com> Co-authored-by: Sydney Runkle <sydneymarierunkle@gmail.com> Co-authored-by: Sydney Runkle <54324534+sydney-runkle@users.noreply.github.com> Co-authored-by: Eugene Yurtsev <eyurtsev@gmail.com>
392 lines
12 KiB
Python
392 lines
12 KiB
Python
from dataclasses import dataclass
|
|
from typing import Any
|
|
|
|
import pytest
|
|
from pydantic import BaseModel, ValidationError
|
|
from typing_extensions import TypedDict
|
|
|
|
from langgraph.graph import END, START, StateGraph
|
|
from langgraph.runtime import Runtime, get_runtime
|
|
|
|
|
|
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}"}
|
|
|
|
graph = StateGraph(state_schema=State, context_schema=Context)
|
|
graph.add_node("injected_runtime", injected_runtime)
|
|
graph.add_edge(START, "injected_runtime")
|
|
graph.add_edge("injected_runtime", END)
|
|
compiled = graph.compile()
|
|
result = compiled.invoke(
|
|
{"message": "hello world"}, context=Context(api_key="sk_123456")
|
|
)
|
|
assert result == {"message": "api key: sk_123456"}
|
|
|
|
|
|
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}"}
|
|
|
|
graph = StateGraph(state_schema=State, context_schema=Context)
|
|
graph.add_node("context_runtime", context_runtime)
|
|
graph.add_edge(START, "context_runtime")
|
|
graph.add_edge("context_runtime", END)
|
|
compiled = graph.compile()
|
|
result = compiled.invoke(
|
|
{"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
|
|
|
|
prev = Runtime(context=Context(api_key="abc"))
|
|
new = prev.override(context=Context(api_key="def"))
|
|
assert new.override(context=Context(api_key="def")).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"))
|
|
runtime3 = Runtime(context=None)
|
|
|
|
assert runtime1.merge(runtime2).context.api_key == "def"
|
|
# override only applies to non-falsy values
|
|
assert runtime1.merge(runtime3).context.api_key == "abc" # type: ignore
|
|
|
|
|
|
def test_runtime_propogated_to_subgraph() -> None:
|
|
@dataclass
|
|
class Context:
|
|
username: str
|
|
|
|
class State(TypedDict, total=False):
|
|
subgraph: str
|
|
main: str
|
|
|
|
def subgraph_node_1(state: State, runtime: Runtime[Context]):
|
|
return {"subgraph": f"{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 {"main": f"{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 == {"subgraph": "Alice!", "main": "Alice!"}
|
|
|
|
|
|
def test_context_coercion_dataclass() -> None:
|
|
"""Test that dict context is coerced to dataclass."""
|
|
|
|
@dataclass
|
|
class Context:
|
|
api_key: str
|
|
timeout: int = 30
|
|
|
|
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}, timeout: {runtime.context.timeout}"
|
|
}
|
|
|
|
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()
|
|
|
|
# Test dict coercion with all fields
|
|
result = compiled.invoke(
|
|
{"message": "test"}, context={"api_key": "sk_test", "timeout": 60}
|
|
)
|
|
assert result == {"message": "api_key: sk_test, timeout: 60"}
|
|
|
|
# Test dict coercion with default field
|
|
result = compiled.invoke({"message": "test"}, context={"api_key": "sk_test2"})
|
|
assert result == {"message": "api_key: sk_test2, timeout: 30"}
|
|
|
|
# Test with actual dataclass instance (should still work)
|
|
result = compiled.invoke(
|
|
{"message": "test"}, context=Context(api_key="sk_test3", timeout=90)
|
|
)
|
|
assert result == {"message": "api_key: sk_test3, timeout: 90"}
|
|
|
|
|
|
def test_context_coercion_pydantic() -> None:
|
|
"""Test that dict context is coerced to Pydantic model."""
|
|
|
|
class Context(BaseModel):
|
|
api_key: str
|
|
timeout: int = 30
|
|
tags: list[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}, timeout: {runtime.context.timeout}, tags: {runtime.context.tags}"
|
|
}
|
|
|
|
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()
|
|
|
|
# Test dict coercion with all fields
|
|
result = compiled.invoke(
|
|
{"message": "test"},
|
|
context={"api_key": "sk_test", "timeout": 60, "tags": ["prod", "v2"]},
|
|
)
|
|
assert result == {"message": "api_key: sk_test, timeout: 60, tags: ['prod', 'v2']"}
|
|
|
|
# Test dict coercion with defaults
|
|
result = compiled.invoke({"message": "test"}, context={"api_key": "sk_test2"})
|
|
assert result == {"message": "api_key: sk_test2, timeout: 30, tags: []"}
|
|
|
|
# Test with actual Pydantic instance (should still work)
|
|
result = compiled.invoke(
|
|
{"message": "test"},
|
|
context=Context(api_key="sk_test3", timeout=90, tags=["test"]),
|
|
)
|
|
assert result == {"message": "api_key: sk_test3, timeout: 90, tags: ['test']"}
|
|
|
|
|
|
def test_context_coercion_typeddict() -> None:
|
|
"""Test that dict context with TypedDict schema passes through as-is."""
|
|
|
|
class Context(TypedDict):
|
|
api_key: str
|
|
timeout: int
|
|
|
|
class State(TypedDict):
|
|
message: str
|
|
|
|
def node_with_context(state: State, runtime: Runtime[Context]) -> dict[str, Any]:
|
|
# TypedDict context is just a dict at runtime
|
|
return {
|
|
"message": f"api_key: {runtime.context['api_key']}, timeout: {runtime.context['timeout']}"
|
|
}
|
|
|
|
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()
|
|
|
|
# Test dict passes through for TypedDict
|
|
result = compiled.invoke(
|
|
{"message": "test"}, context={"api_key": "sk_test", "timeout": 60}
|
|
)
|
|
assert result == {"message": "api_key: sk_test, timeout: 60"}
|
|
|
|
|
|
def test_context_coercion_none() -> None:
|
|
"""Test that None context is handled properly."""
|
|
|
|
@dataclass
|
|
class Context:
|
|
api_key: str
|
|
|
|
class State(TypedDict):
|
|
message: str
|
|
|
|
def node_without_context(state: State, runtime: Runtime[Context]) -> dict[str, Any]:
|
|
# Should be None when no context provided
|
|
return {"message": f"context is None: {runtime.context is None}"}
|
|
|
|
graph = StateGraph(state_schema=State, context_schema=Context)
|
|
graph.add_node("node", node_without_context)
|
|
graph.add_edge(START, "node")
|
|
graph.add_edge("node", END)
|
|
compiled = graph.compile()
|
|
|
|
# Test with None context
|
|
result = compiled.invoke({"message": "test"}, context=None)
|
|
assert result == {"message": "context is None: True"}
|
|
|
|
# Test without context parameter (defaults to None)
|
|
result = compiled.invoke({"message": "test"})
|
|
assert result == {"message": "context is None: True"}
|
|
|
|
|
|
def test_context_coercion_errors() -> None:
|
|
"""Test error handling for invalid context."""
|
|
|
|
@dataclass
|
|
class Context:
|
|
api_key: str # Required field
|
|
|
|
class State(TypedDict):
|
|
message: str
|
|
|
|
def node_with_context(state: State, runtime: Runtime[Context]) -> dict[str, Any]:
|
|
return {"message": "should not reach here"}
|
|
|
|
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()
|
|
|
|
# Test missing required field
|
|
with pytest.raises(TypeError):
|
|
compiled.invoke({"message": "test"}, context={"timeout": 60})
|
|
|
|
# Test invalid dict keys
|
|
with pytest.raises(TypeError):
|
|
compiled.invoke(
|
|
{"message": "test"}, context={"api_key": "test", "invalid_field": "value"}
|
|
)
|
|
|
|
|
|
@pytest.mark.anyio
|
|
async def test_context_coercion_async() -> None:
|
|
"""Test context coercion with async methods."""
|
|
|
|
@dataclass
|
|
class Context:
|
|
api_key: str
|
|
async_mode: bool = True
|
|
|
|
class State(TypedDict):
|
|
message: str
|
|
|
|
async def async_node(state: State, runtime: Runtime[Context]) -> dict[str, Any]:
|
|
return {
|
|
"message": f"async api_key: {runtime.context.api_key}, async_mode: {runtime.context.async_mode}"
|
|
}
|
|
|
|
graph = StateGraph(state_schema=State, context_schema=Context)
|
|
graph.add_node("node", async_node)
|
|
graph.add_edge(START, "node")
|
|
graph.add_edge("node", END)
|
|
compiled = graph.compile()
|
|
|
|
# Test dict coercion with ainvoke
|
|
result = await compiled.ainvoke(
|
|
{"message": "test"}, context={"api_key": "sk_async", "async_mode": False}
|
|
)
|
|
assert result == {"message": "async api_key: sk_async, async_mode: False"}
|
|
|
|
# Test dict coercion with astream
|
|
chunks = []
|
|
async for chunk in compiled.astream(
|
|
{"message": "test"}, context={"api_key": "sk_stream"}
|
|
):
|
|
chunks.append(chunk)
|
|
|
|
# Find the chunk with our node output
|
|
node_output = None
|
|
for chunk in chunks:
|
|
if "node" in chunk:
|
|
node_output = chunk["node"]
|
|
break
|
|
|
|
assert node_output == {"message": "async api_key: sk_stream, async_mode: True"}
|
|
|
|
|
|
def test_context_coercion_stream() -> None:
|
|
"""Test context coercion with sync stream method."""
|
|
|
|
@dataclass
|
|
class Context:
|
|
api_key: str
|
|
stream_mode: str = "default"
|
|
|
|
class State(TypedDict):
|
|
message: str
|
|
|
|
def node_with_context(state: State, runtime: Runtime[Context]) -> dict[str, Any]:
|
|
return {
|
|
"message": f"stream api_key: {runtime.context.api_key}, mode: {runtime.context.stream_mode}"
|
|
}
|
|
|
|
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()
|
|
|
|
# Test dict coercion with stream
|
|
chunks = []
|
|
for chunk in compiled.stream(
|
|
{"message": "test"}, context={"api_key": "sk_stream", "stream_mode": "fast"}
|
|
):
|
|
chunks.append(chunk)
|
|
|
|
# Find the chunk with our node output
|
|
node_output = None
|
|
for chunk in chunks:
|
|
if "node" in chunk:
|
|
node_output = chunk["node"]
|
|
break
|
|
|
|
assert node_output == {"message": "stream api_key: sk_stream, mode: fast"}
|
|
|
|
|
|
def test_context_coercion_pydantic_validation_errors() -> None:
|
|
"""Test that Pydantic validation errors are raised."""
|
|
|
|
class Context(BaseModel):
|
|
api_key: str
|
|
timeout: int
|
|
|
|
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}, timeout: {runtime.context.timeout}"
|
|
}
|
|
|
|
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()
|
|
|
|
with pytest.raises(ValidationError):
|
|
compiled.invoke(
|
|
{"message": "test"}, context={"api_key": "sk_test", "timeout": "not_an_int"}
|
|
)
|