WEBDIENSTE/langgraph
1
0
Fork 0
mirror of https://github.com/langchain-ai/langgraph.git synced 2026-09-06 17:57:49 +02:00
Code Issues Packages Projects Releases Wiki Activity

chore: typing for headers in remote graph (#5653)

Browse Source
This commit is contained in:
William FH
2025-07-24 18:00:44 -07:00
committed by GitHub
parent 2e5445c565
commit cdaa7ba003
2 changed files with 104 additions and 19 deletions
Download Patch File Download Diff File Expand all files Collapse all files
libs/langgraph/langgraph/pregel/remote.py
+25 -7
View File
@@ -633,6 +633,7 @@ class RemoteGraph(PregelProtocol):
interrupt_before: All | Sequence[str] | None = None,
interrupt_after: All | Sequence[str] | None = None,
subgraphs: bool = False,
headers: dict[str, str] | None = None,
**kwargs: Any,
) -> Iterator[dict[str, Any] | Any]:
"""Create a run and stream the results.
@@ -648,6 +649,7 @@ class RemoteGraph(PregelProtocol):
interrupt_before: Interrupt the graph before these nodes.
interrupt_after: Interrupt the graph after these nodes.
subgraphs: Stream from subgraphs.
headers: Additional headers to pass to the request.
**kwargs: Additional params to pass to client.runs.stream.
Yields:
@@ -676,7 +678,9 @@ class RemoteGraph(PregelProtocol):
interrupt_after=interrupt_after,
stream_subgraphs=subgraphs or stream is not None,
if_not_exists="create",
headers=self._merge_tracing_headers(kwargs.pop("headers", None) or {}),
headers=_merge_tracing_headers(headers)
if self.distributed_tracing
else headers,
**kwargs,
):
# split mode and ns
@@ -736,6 +740,7 @@ class RemoteGraph(PregelProtocol):
interrupt_before: All | Sequence[str] | None = None,
interrupt_after: All | Sequence[str] | None = None,
subgraphs: bool = False,
headers: dict[str, str] | None = None,
**kwargs: Any,
) -> AsyncIterator[dict[str, Any] | Any]:
"""Create a run and stream the results.
@@ -751,6 +756,7 @@ class RemoteGraph(PregelProtocol):
interrupt_before: Interrupt the graph before these nodes.
interrupt_after: Interrupt the graph after these nodes.
subgraphs: Stream from subgraphs.
headers: Additional headers to pass to the request.
**kwargs: Additional params to pass to client.runs.stream.
Yields:
@@ -779,7 +785,9 @@ class RemoteGraph(PregelProtocol):
interrupt_after=interrupt_after,
stream_subgraphs=subgraphs or stream is not None,
if_not_exists="create",
headers=self._merge_tracing_headers(kwargs.pop("headers", None) or {}),
headers=_merge_tracing_headers(headers)
if self.distributed_tracing
else headers,
**kwargs,
):
# split mode and ns
@@ -853,6 +861,7 @@ class RemoteGraph(PregelProtocol):
*,
interrupt_before: All | Sequence[str] | None = None,
interrupt_after: All | Sequence[str] | None = None,
headers: dict[str, str] | None = None,
**kwargs: Any,
) -> dict[str, Any] | Any:
"""Create a run, wait until it finishes and return the final state.
@@ -862,6 +871,7 @@ class RemoteGraph(PregelProtocol):
config: A `RunnableConfig` for graph invocation.
interrupt_before: Interrupt the graph before these nodes.
interrupt_after: Interrupt the graph after these nodes.
headers: Additional headers to pass to the request.
**kwargs: Additional params to pass to RemoteGraph.stream.
Returns:
@@ -872,6 +882,7 @@ class RemoteGraph(PregelProtocol):
config=config,
interrupt_before=interrupt_before,
interrupt_after=interrupt_after,
headers=headers,
stream_mode="values",
**kwargs,
):
@@ -888,6 +899,7 @@ class RemoteGraph(PregelProtocol):
*,
interrupt_before: All | Sequence[str] | None = None,
interrupt_after: All | Sequence[str] | None = None,
headers: dict[str, str] | None = None,
**kwargs: Any,
) -> dict[str, Any] | Any:
"""Create a run, wait until it finishes and return the final state.
@@ -897,6 +909,7 @@ class RemoteGraph(PregelProtocol):
config: A `RunnableConfig` for graph invocation.
interrupt_before: Interrupt the graph before these nodes.
interrupt_after: Interrupt the graph after these nodes.
headers: Additional headers to pass to the request.
**kwargs: Additional params to pass to RemoteGraph.astream.
Returns:
@@ -907,6 +920,7 @@ class RemoteGraph(PregelProtocol):
config=config,
interrupt_before=interrupt_before,
interrupt_after=interrupt_after,
headers=headers,
stream_mode="values",
**kwargs,
):
@@ -916,12 +930,16 @@ class RemoteGraph(PregelProtocol):
except UnboundLocalError:
return None
def _merge_tracing_headers(self, headers: dict[str, str]) -> dict[str, str]:
if rt := ls.get_current_run_tree():
tracing_headers = rt.to_headers()
baggage = tracing_headers.pop("baggage")
def _merge_tracing_headers(headers: dict[str, str] | None) -> dict[str, str] | None:
if rt := ls.get_current_run_tree():
tracing_headers = rt.to_headers()
baggage = tracing_headers.pop("baggage")
if headers:
if "baggage" in headers:
baggage = headers["baggage"] + "," + baggage
tracing_headers["baggage"] = baggage
headers.update(tracing_headers)
return headers
else:
headers = tracing_headers
return headers
libs/langgraph/tests/test_remote_graph.py
+79 -12
View File
@@ -3,6 +3,7 @@ import sys
from typing import Annotated, Union
from unittest.mock import AsyncMock, MagicMock
import langsmith as ls
import pytest
from langchain_core.messages import AnyMessage, BaseMessage
from langchain_core.runnables import RunnableConfig
@@ -899,21 +900,20 @@ async def test_langgraph_cloud_integration():
}
# test invoke
response = app.invoke(
app.invoke(
input,
config={"configurable": {"thread_id": "39a6104a-34e7-4f83-929c-d9eb163003c9"}},
interrupt_before=["agent"],
)
print("response:", response["messages"][-1].content)
# test stream
async for chunk in app.astream(
async for _ in app.astream(
input,
config={"configurable": {"thread_id": "2dc3e3e7-39ac-4597-aa57-4404b944e82a"}},
subgraphs=True,
stream_mode=["debug", "messages"],
):
print("chunk:", chunk)
pass
# test stream events
async for chunk in remote_pregel.astream_events(
@@ -923,17 +923,16 @@ async def test_langgraph_cloud_integration():
subgraphs=True,
stream_mode=[],
):
print("chunk:", chunk)
pass
# test get state
state_snapshot = await remote_pregel.aget_state(
await remote_pregel.aget_state(
config={"configurable": {"thread_id": "2dc3e3e7-39ac-4597-aa57-4404b944e82a"}},
subgraphs=True,
)
print("state snapshot:", state_snapshot)
# test update state
response = await remote_pregel.aupdate_state(
await remote_pregel.aupdate_state(
config={"configurable": {"thread_id": "6645e002-ed50-4022-92a3-d0d186fdf812"}},
values={
"messages": [
@@ -944,18 +943,16 @@ async def test_langgraph_cloud_integration():
]
},
)
print("response:", response)
# test get history
async for state in remote_pregel.aget_state_history(
config={"configurable": {"thread_id": "2dc3e3e7-39ac-4597-aa57-4404b944e82a"}},
):
print("state snapshot:", state)
pass
# test get graph
remote_pregel.graph_id = "fe096781-5601-53d2-b2f6-0d3403f7e9ca" # must be UUID
graph = await remote_pregel.aget_graph(xray=True)
print("graph:", graph)
await remote_pregel.aget_graph(xray=True)
def test_sanitize_config():
@@ -1181,3 +1178,73 @@ async def test_remote_graph_stream_messages_tuple(
assert coerced_events == coerced_inmem_events
# TODO: Fix the namespace matching in the next api release.
# assert namespaces == inmem_namespaces
@pytest.mark.anyio
@pytest.mark.parametrize("distributed_tracing", [False, True])
@pytest.mark.parametrize("stream", [False, True])
async def test_include_headers(distributed_tracing: bool, stream: bool):
mock_async_client = MagicMock()
async_iter = MagicMock()
return_value = [
StreamPart(event="values", data={"chunk": "data1"}),
]
async_iter.__aiter__.return_value = return_value
astream_mock = mock_async_client.runs.stream
astream_mock.return_value = async_iter
mock_sync_client = MagicMock()
sync_iter = MagicMock()
sync_iter.__iter__.return_value = return_value
stream_mock = mock_sync_client.runs.stream
stream_mock.return_value = async_iter
remote_pregel = RemoteGraph(
"test_graph_id",
client=mock_async_client,
sync_client=mock_sync_client,
distributed_tracing=distributed_tracing,
)
config = {"configurable": {"thread_id": "thread_1"}}
with ls.tracing_context(enabled=True, client=MagicMock()):
with ls.trace("foo"):
if stream:
async for _ in remote_pregel.astream(
{"input": {"messages": [{"type": "human", "content": "hello"}]}},
config,
headers={"foo": "bar"},
):
pass
else:
await remote_pregel.ainvoke(
{"input": {"messages": [{"type": "human", "content": "hello"}]}},
config,
headers={"foo": "bar"},
)
expected = {"foo": "bar"}
if distributed_tracing:
expected["langsmith-trace"] = AnyStr()
expected["baggage"] = AnyStr()
assert astream_mock.call_args.kwargs["headers"] == expected
stream_mock.assert_not_called()
with ls.tracing_context(enabled=True, client=MagicMock()):
with ls.trace("foo"):
if stream:
for _ in remote_pregel.stream(
{"input": {"messages": [{"type": "human", "content": "hello"}]}},
config,
headers={"foo": "bar"},
):
pass
else:
remote_pregel.invoke(
{"input": {"messages": [{"type": "human", "content": "hello"}]}},
config,
headers={"foo": "bar"},
)
assert stream_mock.call_args.kwargs["headers"] == expected