diff --git a/libs/langgraph/langgraph/pregel/remote.py b/libs/langgraph/langgraph/pregel/remote.py index 7b9602d81..16ffd86a0 100644 --- a/libs/langgraph/langgraph/pregel/remote.py +++ b/libs/langgraph/langgraph/pregel/remote.py @@ -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, diff --git a/libs/langgraph/tests/test_remote_graph.py b/libs/langgraph/tests/test_remote_graph.py index e0eacf006..4edb9325e 100644 --- a/libs/langgraph/tests/test_remote_graph.py +++ b/libs/langgraph/tests/test_remote_graph.py @@ -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" )