fix(langgraph): Remote Baggage (#5964)

Fix baggage propagation for opt-in distributed tracing when no
additional headers are provided.
This commit is contained in:
William FH
2025-08-20 01:47:12 +00:00
committed by GitHub
parent a566f1f892
commit a692e24a58
5 changed files with 31 additions and 26 deletions
+9 -9
View File
@@ -733,9 +733,9 @@ class RemoteGraph(PregelProtocol):
interrupt_after=interrupt_after,
stream_subgraphs=subgraphs or stream is not None,
if_not_exists="create",
headers=_merge_tracing_headers(headers)
if self.distributed_tracing
else headers,
headers=(
_merge_tracing_headers(headers) if self.distributed_tracing else headers
),
params=params,
**kwargs,
):
@@ -842,9 +842,9 @@ class RemoteGraph(PregelProtocol):
interrupt_after=interrupt_after,
stream_subgraphs=subgraphs or stream is not None,
if_not_exists="create",
headers=_merge_tracing_headers(headers)
if self.distributed_tracing
else headers,
headers=(
_merge_tracing_headers(headers) if self.distributed_tracing else headers
),
params=params,
**kwargs,
):
@@ -996,11 +996,11 @@ class RemoteGraph(PregelProtocol):
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
tracing_headers["baggage"] = (
f"{headers['baggage']},{tracing_headers['baggage']}"
)
headers.update(tracing_headers)
else:
headers = tracing_headers
+1 -1
View File
@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
[project]
name = "langgraph"
version = "0.6.5"
version = "0.6.6"
description = "Building stateful, multi-actor applications with LLMs"
authors = []
requires-python = ">=3.9"
+13 -8
View File
@@ -1,6 +1,6 @@
import re
import sys
from typing import Annotated, Union
from typing import Annotated, Optional, Union
from unittest.mock import AsyncMock, MagicMock
import langsmith as ls
@@ -1183,7 +1183,10 @@ async def test_remote_graph_stream_messages_tuple(
@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):
@pytest.mark.parametrize("headers", [None, {"foo": "bar"}])
async def test_include_headers(
distributed_tracing: bool, stream: bool, headers: Optional[dict[str, str]]
):
mock_async_client = MagicMock()
async_iter = MagicMock()
return_value = [
@@ -1213,7 +1216,7 @@ async def test_include_headers(distributed_tracing: bool, stream: bool):
async for _ in remote_pregel.astream(
{"input": {"messages": [{"type": "human", "content": "hello"}]}},
config,
headers={"foo": "bar"},
headers=headers,
):
pass
@@ -1221,12 +1224,14 @@ async def test_include_headers(distributed_tracing: bool, stream: bool):
await remote_pregel.ainvoke(
{"input": {"messages": [{"type": "human", "content": "hello"}]}},
config,
headers={"foo": "bar"},
headers=headers,
)
expected = {"foo": "bar"}
expected = headers.copy() if headers else None
if distributed_tracing:
if expected is None:
expected = {}
expected["langsmith-trace"] = AnyStr()
expected["baggage"] = AnyStr()
expected["baggage"] = AnyStr("langsmith-metadata=")
assert astream_mock.call_args.kwargs["headers"] == expected
stream_mock.assert_not_called()
@@ -1237,7 +1242,7 @@ async def test_include_headers(distributed_tracing: bool, stream: bool):
for _ in remote_pregel.stream(
{"input": {"messages": [{"type": "human", "content": "hello"}]}},
config,
headers={"foo": "bar"},
headers=headers,
):
pass
@@ -1245,6 +1250,6 @@ async def test_include_headers(distributed_tracing: bool, stream: bool):
remote_pregel.invoke(
{"input": {"messages": [{"type": "human", "content": "hello"}]}},
config,
headers={"foo": "bar"},
headers=headers,
)
assert stream_mock.call_args.kwargs["headers"] == expected
+7 -7
View File
@@ -1269,7 +1269,7 @@ wheels = [
[[package]]
name = "langgraph"
version = "0.6.5"
version = "0.6.6"
source = { editable = "." }
dependencies = [
{ name = "langchain-core" },
@@ -1350,7 +1350,7 @@ dev = [
[[package]]
name = "langgraph-api"
version = "0.2.134"
version = "0.2.135"
source = { registry = "https://pypi.org/simple" }
dependencies = [
{ name = "cloudpickle", marker = "python_full_version >= '3.11'" },
@@ -1373,9 +1373,9 @@ dependencies = [
{ name = "uvicorn", marker = "python_full_version >= '3.11'" },
{ name = "watchfiles", marker = "python_full_version >= '3.11'" },
]
sdist = { url = "https://files.pythonhosted.org/packages/3d/34/a7f7560b00973c66e11a962a5cb3f878c412adb5dba52faa8e272d6afe4b/langgraph_api-0.2.134.tar.gz", hash = "sha256:75a098c9241dc29c0968b69e249c984f4bd3d77d91716837e490f80a282d93dd", size = 253684, upload-time = "2025-08-18T09:01:53.279Z" }
sdist = { url = "https://files.pythonhosted.org/packages/8f/1e/5ca12f00852e2f54075c65dc2d4eca4d2d8facd404d5c0d17803fe620eae/langgraph_api-0.2.135.tar.gz", hash = "sha256:8cbcd5ff7be89a72d01dd3b32b0873bc5114ab9ff206fc94c2781ba56d09182a", size = 255563, upload-time = "2025-08-20T01:02:01.661Z" }
wheels = [
{ url = "https://files.pythonhosted.org/packages/59/95/81610cccfa280d04d084dfcc8a9bdf9924ffe9a43ed4a4c788f7721748e4/langgraph_api-0.2.134-py3-none-any.whl", hash = "sha256:f30fb474a346ddf40c9e8b2aad485bcce9be33d294fea68769c5a6315f0457e8", size = 205842, upload-time = "2025-08-18T09:01:51.796Z" },
{ url = "https://files.pythonhosted.org/packages/a6/98/18c475de127927622079e087b7cfe1ab97ed44fe1a2425f8c8de07ea0e64/langgraph_api-0.2.135-py3-none-any.whl", hash = "sha256:ac78035075230ccb816f48b242a76c5f8e607c074fe7802928425ba76a11b3f4", size = 205849, upload-time = "2025-08-20T01:02:00.156Z" },
]
[[package]]
@@ -2774,7 +2774,7 @@ wheels = [
[[package]]
name = "requests"
version = "2.32.4"
version = "2.32.5"
source = { registry = "https://pypi.org/simple" }
dependencies = [
{ name = "certifi" },
@@ -2782,9 +2782,9 @@ dependencies = [
{ name = "idna" },
{ name = "urllib3" },
]
sdist = { url = "https://files.pythonhosted.org/packages/e1/0a/929373653770d8a0d7ea76c37de6e41f11eb07559b103b1c02cafb3f7cf8/requests-2.32.4.tar.gz", hash = "sha256:27d0316682c8a29834d3264820024b62a36942083d52caf2f14c0591336d3422", size = 135258, upload-time = "2025-06-09T16:43:07.34Z" }
sdist = { url = "https://files.pythonhosted.org/packages/c9/74/b3ff8e6c8446842c3f5c837e9c3dfcfe2018ea6ecef224c710c85ef728f4/requests-2.32.5.tar.gz", hash = "sha256:dbba0bac56e100853db0ea71b82b4dfd5fe2bf6d3754a8893c3af500cec7d7cf", size = 134517, upload-time = "2025-08-18T20:46:02.573Z" }
wheels = [
{ url = "https://files.pythonhosted.org/packages/7c/e4/56027c4a6b4ae70ca9de302488c5ca95ad4a39e190093d6c1a8ace08341b/requests-2.32.4-py3-none-any.whl", hash = "sha256:27babd3cda2a6d50b30443204ee89830707d396671944c998b5975b031ac2b2c", size = 64847, upload-time = "2025-06-09T16:43:05.728Z" },
{ url = "https://files.pythonhosted.org/packages/1e/db/4254e3eabe8020b458f1a747140d32277ec7a271daf1d235b70dc0b4e6e3/requests-2.32.5-py3-none-any.whl", hash = "sha256:2462f94637a34fd532264295e186976db0f5d453d1cdd31473c85a6a161affb6", size = 64738, upload-time = "2025-08-18T20:46:00.542Z" },
]
[[package]]
+1 -1
View File
@@ -316,7 +316,7 @@ wheels = [
[[package]]
name = "langgraph"
version = "0.6.5"
version = "0.6.6"
source = { editable = "../langgraph" }
dependencies = [
{ name = "langchain-core" },