WEBDIENSTE/langgraph
1
0
Fork 0
mirror of https://github.com/langchain-ai/langgraph.git synced 2026-08-25 17:12:26 +02:00
Code Issues Packages Projects Releases Wiki Activity

context in remote

Browse Source
This commit is contained in:
Sydney Runkle
2026-03-12 10:52:10 -04:00
parent b7b052e66c
commit 04076c4bb5
2 changed files with 201 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
libs/langgraph/langgraph/pregel/remote.py
+17
View File
@@ -31,6 +31,7 @@ from langgraph_sdk.client import (
)
from langgraph_sdk.schema import (
Checkpoint,
Context,
QueryParamTypes,
ThreadState,
)
@@ -691,6 +692,7 @@ class RemoteGraph(PregelProtocol):
input: dict[str, Any] | Any,
config: RunnableConfig | None = None,
*,
context: Context | None = None,
stream_mode: StreamMode | list[StreamMode] | None = None,
interrupt_before: All | Sequence[str] | None = None,
interrupt_after: All | Sequence[str] | None = None,
@@ -707,6 +709,7 @@ class RemoteGraph(PregelProtocol):
input: dict[str, Any] | Any,
config: RunnableConfig | None = None,
*,
context: Context | None = None,
stream_mode: StreamMode | list[StreamMode] | None = None,
interrupt_before: All | Sequence[str] | None = None,
interrupt_after: All | Sequence[str] | None = None,
@@ -722,6 +725,7 @@ class RemoteGraph(PregelProtocol):
input: dict[str, Any] | Any,
config: RunnableConfig | None = None,
*,
context: Context | None = None,
stream_mode: StreamMode | list[StreamMode] | None = None,
interrupt_before: All | Sequence[str] | None = None,
interrupt_after: All | Sequence[str] | None = None,
@@ -769,6 +773,7 @@ class RemoteGraph(PregelProtocol):
input=input,
command=command,
config=sanitized_config,
context=context,
stream_mode=stream_modes,
interrupt_before=interrupt_before,
interrupt_after=interrupt_after,
@@ -842,6 +847,7 @@ class RemoteGraph(PregelProtocol):
input: dict[str, Any] | Any,
config: RunnableConfig | None = None,
*,
context: Context | None = None,
stream_mode: StreamMode | list[StreamMode] | None = None,
interrupt_before: All | Sequence[str] | None = None,
interrupt_after: All | Sequence[str] | None = None,
@@ -858,6 +864,7 @@ class RemoteGraph(PregelProtocol):
input: dict[str, Any] | Any,
config: RunnableConfig | None = None,
*,
context: Context | None = None,
stream_mode: StreamMode | list[StreamMode] | None = None,
interrupt_before: All | Sequence[str] | None = None,
interrupt_after: All | Sequence[str] | None = None,
@@ -873,6 +880,7 @@ class RemoteGraph(PregelProtocol):
input: dict[str, Any] | Any,
config: RunnableConfig | None = None,
*,
context: Context | None = None,
stream_mode: StreamMode | list[StreamMode] | None = None,
interrupt_before: All | Sequence[str] | None = None,
interrupt_after: All | Sequence[str] | None = None,
@@ -920,6 +928,7 @@ class RemoteGraph(PregelProtocol):
input=input,
command=command,
config=sanitized_config,
context=context,
stream_mode=stream_modes,
interrupt_before=interrupt_before,
interrupt_after=interrupt_after,
@@ -1009,6 +1018,7 @@ class RemoteGraph(PregelProtocol):
input: dict[str, Any] | Any,
config: RunnableConfig | None = None,
*,
context: Context | None = None,
interrupt_before: All | Sequence[str] | None = None,
interrupt_after: All | Sequence[str] | None = None,
headers: dict[str, str] | None = None,
@@ -1023,6 +1033,7 @@ class RemoteGraph(PregelProtocol):
input: dict[str, Any] | Any,
config: RunnableConfig | None = None,
*,
context: Context | None = None,
interrupt_before: All | Sequence[str] | None = None,
interrupt_after: All | Sequence[str] | None = None,
headers: dict[str, str] | None = None,
@@ -1036,6 +1047,7 @@ class RemoteGraph(PregelProtocol):
input: dict[str, Any] | Any,
config: RunnableConfig | None = None,
*,
context: Context | None = None,
interrupt_before: All | Sequence[str] | None = None,
interrupt_after: All | Sequence[str] | None = None,
headers: dict[str, str] | None = None,
@@ -1061,6 +1073,7 @@ class RemoteGraph(PregelProtocol):
for chunk in self.stream( # type: ignore[misc, call-overload]
input,
config=config,
context=context,
interrupt_before=interrupt_before,
interrupt_after=interrupt_after,
headers=headers,
@@ -1087,6 +1100,7 @@ class RemoteGraph(PregelProtocol):
input: dict[str, Any] | Any,
config: RunnableConfig | None = None,
*,
context: Context | None = None,
interrupt_before: All | Sequence[str] | None = None,
interrupt_after: All | Sequence[str] | None = None,
headers: dict[str, str] | None = None,
@@ -1101,6 +1115,7 @@ class RemoteGraph(PregelProtocol):
input: dict[str, Any] | Any,
config: RunnableConfig | None = None,
*,
context: Context | None = None,
interrupt_before: All | Sequence[str] | None = None,
interrupt_after: All | Sequence[str] | None = None,
headers: dict[str, str] | None = None,
@@ -1114,6 +1129,7 @@ class RemoteGraph(PregelProtocol):
input: dict[str, Any] | Any,
config: RunnableConfig | None = None,
*,
context: Context | None = None,
interrupt_before: All | Sequence[str] | None = None,
interrupt_after: All | Sequence[str] | None = None,
headers: dict[str, str] | None = None,
@@ -1139,6 +1155,7 @@ class RemoteGraph(PregelProtocol):
async for chunk in self.astream( # type: ignore[misc, call-overload]
input,
config=config,
context=context,
interrupt_before=interrupt_before,
interrupt_after=interrupt_after,
headers=headers,
libs/langgraph/tests/test_remote_graph.py
+184
View File
@@ -1,5 +1,6 @@
import re
import sys
from dataclasses import dataclass
from typing import Annotated
from unittest.mock import AsyncMock, MagicMock
@@ -10,6 +11,7 @@ from langchain_core.runnables import RunnableConfig
from langchain_core.runnables.graph import Edge as DrawableEdge
from langchain_core.runnables.graph import Node as DrawableNode
from langgraph_sdk.schema import StreamPart
from pydantic import BaseModel
from typing_extensions import TypedDict
from langgraph.errors import GraphInterrupt
@@ -908,6 +910,188 @@ async def test_ainvoke():
assert result == {"messages": [{"type": "human", "content": "world"}]}
def test_stream_context():
"""Test that context is passed through to the SDK client in stream."""
mock_sync_client = MagicMock()
mock_sync_client.runs.stream.return_value = [
StreamPart(event="values", data={"chunk": "data1"}),
]
remote_pregel = RemoteGraph(
"test_graph_id",
sync_client=mock_sync_client,
)
config = {"configurable": {"thread_id": "thread_1"}}
context = {"model_name": "anthropic", "user_id": "123"}
stream_parts = list(
remote_pregel.stream(
{"input": "data"},
config,
context=context,
stream_mode="values",
)
)
assert stream_parts == [{"chunk": "data1"}]
_, kwargs = mock_sync_client.runs.stream.call_args
assert kwargs["context"] == {"model_name": "anthropic", "user_id": "123"}
def test_stream_context_none():
"""Test that context defaults to None when not provided."""
mock_sync_client = MagicMock()
mock_sync_client.runs.stream.return_value = [
StreamPart(event="values", data={"chunk": "data1"}),
]
remote_pregel = RemoteGraph(
"test_graph_id",
sync_client=mock_sync_client,
)
config = {"configurable": {"thread_id": "thread_1"}}
list(remote_pregel.stream({"input": "data"}, config, stream_mode="values"))
_, kwargs = mock_sync_client.runs.stream.call_args
assert kwargs["context"] is None
@pytest.mark.anyio
async def test_astream_context():
"""Test that context is passed through to the SDK client in astream."""
mock_async_client = MagicMock()
async_iter = MagicMock()
async_iter.__aiter__.return_value = [
StreamPart(event="values", data={"chunk": "data1"}),
]
mock_async_client.runs.stream.return_value = async_iter
remote_pregel = RemoteGraph(
"test_graph_id",
client=mock_async_client,
)
config = {"configurable": {"thread_id": "thread_1"}}
context = {"model_name": "anthropic"}
chunks = []
async for chunk in remote_pregel.astream(
{"input": "data"},
config,
context=context,
stream_mode="values",
):
chunks.append(chunk)
assert chunks == [{"chunk": "data1"}]
_, kwargs = mock_async_client.runs.stream.call_args
assert kwargs["context"] == {"model_name": "anthropic"}
def test_invoke_context():
"""Test that context is passed through to the SDK client in invoke."""
mock_sync_client = MagicMock()
mock_sync_client.runs.stream.return_value = [
StreamPart(event="values", data={"result": "done"}),
]
remote_pregel = RemoteGraph(
"test_graph_id",
sync_client=mock_sync_client,
)
config = {"configurable": {"thread_id": "thread_1"}}
context = {"model_name": "openai"}
result = remote_pregel.invoke({"input": "data"}, config, context=context)
assert result == {"result": "done"}
_, kwargs = mock_sync_client.runs.stream.call_args
assert kwargs["context"] == {"model_name": "openai"}
@pytest.mark.anyio
async def test_ainvoke_context():
"""Test that context is passed through to the SDK client in ainvoke."""
mock_async_client = MagicMock()
async_iter = MagicMock()
async_iter.__aiter__.return_value = [
StreamPart(event="values", data={"result": "done"}),
]
mock_async_client.runs.stream.return_value = async_iter
remote_pregel = RemoteGraph(
"test_graph_id",
client=mock_async_client,
)
config = {"configurable": {"thread_id": "thread_1"}}
context = {"user_id": "456"}
result = await remote_pregel.ainvoke({"input": "data"}, config, context=context)
assert result == {"result": "done"}
_, kwargs = mock_async_client.runs.stream.call_args
assert kwargs["context"] == {"user_id": "456"}
def test_stream_context_dataclass():
"""Test that a dataclass context is passed through to the SDK client."""
@dataclass
class MyContext:
model_name: str
user_id: str
mock_sync_client = MagicMock()
mock_sync_client.runs.stream.return_value = [
StreamPart(event="values", data={"chunk": "data1"}),
]
remote_pregel = RemoteGraph(
"test_graph_id",
sync_client=mock_sync_client,
)
config = {"configurable": {"thread_id": "thread_1"}}
ctx = MyContext(model_name="anthropic", user_id="123")
list(
remote_pregel.stream(
{"input": "data"}, config, context=ctx, stream_mode="values"
)
)
_, kwargs = mock_sync_client.runs.stream.call_args
assert kwargs["context"] == ctx
def test_stream_context_base_model():
"""Test that a BaseModel context is passed through to the SDK client."""
class MyContext(BaseModel):
model_name: str
user_id: str
mock_sync_client = MagicMock()
mock_sync_client.runs.stream.return_value = [
StreamPart(event="values", data={"chunk": "data1"}),
]
remote_pregel = RemoteGraph(
"test_graph_id",
sync_client=mock_sync_client,
)
config = {"configurable": {"thread_id": "thread_1"}}
ctx = MyContext(model_name="anthropic", user_id="123")
list(
remote_pregel.stream(
{"input": "data"}, config, context=ctx, stream_mode="values"
)
)
_, kwargs = mock_sync_client.runs.stream.call_args
assert kwargs["context"] == ctx
@pytest.mark.skip(
"Unskip this test to manually test the LangSmith Deployment integration"
)