Compare commits

..
Author SHA1 Message Date
Christian Bromann 741c6f8d50 cr 2026-03-18 15:46:35 -07:00
Christian Bromann 37a5504433 move to stream mode 2026-03-12 15:08:26 -07:00
Christian Bromann f6286dce38 feat(langgraph): add protocol improvements for better streaming 2026-03-12 15:08:25 -07:00
13 changed files with 497 additions and 182 deletions
@@ -328,9 +328,6 @@ async def _run(
store: weakref.ReferenceType[BaseStore], store: weakref.ReferenceType[BaseStore],
) -> None: ) -> None:
while item := await aqueue.get(): while item := await aqueue.get():
# don't run batch if the future is done (e.g. cancelled)
if item[0].done():
continue
# check if store is still alive # check if store is still alive
if s := store(): if s := store():
try: try:
@@ -338,9 +335,6 @@ async def _run(
items = [item] items = [item]
try: try:
while item := aqueue.get_nowait(): while item := aqueue.get_nowait():
# don't insert if the future is done (e.g. cancelled)
if item[0].done():
continue
items.append(item) items.append(item)
except asyncio.QueueEmpty: except asyncio.QueueEmpty:
pass pass
+1 -1
View File
@@ -1 +1 @@
__version__ = "0.4.18" __version__ = "0.4.15"
+20 -46
View File
@@ -721,13 +721,7 @@ def _deploy_base_options(
OPT_VERBOSE, OPT_VERBOSE,
OPT_HOST_URL, OPT_HOST_URL,
click.option("--image-name", hidden=True), click.option("--image-name", hidden=True),
click.option( click.option("--image-tag", default="latest", hidden=True),
"--tag",
"-t",
default="latest",
show_default=True,
help="Tag to use for the pushed deployment image.",
),
click.option( click.option(
"--config", "--config",
"-c", "-c",
@@ -764,7 +758,7 @@ def _deploy_base_options(
@cli.group( @cli.group(
cls=DeployGroup, cls=DeployGroup,
help=( help=(
"[Beta] Build and deploy a LangGraph image to LangSmith Deployment.\n\n" "[Beta] Build and deploy a LangGraph image to LangSmith Deployments.\n\n"
"This command is in beta and under active development. " "This command is in beta and under active development. "
"Expect frequent updates and improvements.\n\n" "Expect frequent updates and improvements.\n\n"
"Run from the root of your LangGraph project (where langgraph.json " "Run from the root of your LangGraph project (where langgraph.json "
@@ -801,7 +795,7 @@ def _deploy(
deployment_type: str, deployment_type: str,
name: str | None, name: str | None,
image_name: str | None, image_name: str | None,
tag: str, image_tag: str,
base_image: str | None, base_image: str | None,
install_command: str | None, install_command: str | None,
build_command: str | None, build_command: str | None,
@@ -982,7 +976,7 @@ def _deploy(
normalized_registry = normalized_registry.split("//", 1)[1] normalized_registry = normalized_registry.split("//", 1)[1]
repo_seed = image_name or name or config.parent.name repo_seed = image_name or name or config.parent.name
repo_name = _normalize_image_name(repo_seed) repo_name = _normalize_image_name(repo_seed)
tag_value = _normalize_image_tag(tag) tag_value = _normalize_image_tag(image_tag)
remote_image = f"{normalized_registry}/{repo_name}:{tag_value}" remote_image = f"{normalized_registry}/{repo_name}:{tag_value}"
registry_host = normalized_registry.split("/")[0] registry_host = normalized_registry.split("/")[0]
@@ -1122,7 +1116,7 @@ def _deploy(
) )
else: else:
click.secho( click.secho(
" Check status in the LangSmith Deployment dashboard.", " Check status in the LangSmith Deployments dashboard.",
fg="yellow", fg="yellow",
) )
@@ -1165,42 +1159,22 @@ def _call_host_backend_with_optional_tenant(
in-place so all subsequent calls through the same instance are in-place so all subsequent calls through the same instance are
tenant-aware. tenant-aware.
""" """
prompted_for_tenant = False try:
return operation(client)
while True: except HostBackendError as err:
try: if err.status_code == 403 and "requires workspace specification" in err.message:
click.secho(
"Your API key is org-scoped and requires a workspace ID.",
fg="yellow",
)
click.secho(
"Find your workspace ID in LangSmith under Settings > Workspaces.",
fg="yellow",
)
tenant_id = click.prompt("Workspace ID")
client._client.headers["X-Tenant-ID"] = tenant_id
return operation(client) return operation(client)
except HostBackendError as err: raise
if (
not prompted_for_tenant
and err.status_code == 403
and "requires workspace specification" in err.message
):
click.secho(
"Your API key is org-scoped and requires a workspace ID.",
fg="yellow",
)
click.secho(
"Find your workspace ID in LangSmith under Settings > Workspaces.",
fg="yellow",
)
client._client.headers["X-Tenant-ID"] = click.prompt("Workspace ID")
prompted_for_tenant = True
continue
if err.status_code == 403 and "not enabled" in err.message.lower():
from urllib.parse import urlparse
smith_host = "smith.langchain.com"
parsed = urlparse(client._base_url)
if (parsed.hostname or "").startswith("eu."):
smith_host = "eu.smith.langchain.com"
raise HostBackendError(
"LangSmith Deployment is not enabled for this organization. "
f"Enable it at https://{smith_host}/host/deployments"
" (ensure this matches the organization for your API key).",
status_code=403,
) from None
raise
@OPT_HOST_API_KEY @OPT_HOST_API_KEY
+55 -19
View File
@@ -8,20 +8,31 @@ from zipfile import ZipFile
import click import click
TEMPLATES: dict[str, dict[str, str]] = { TEMPLATES: dict[str, dict[str, str]] = {
"Deep Agent": {
"description": "An opinionated deployment template for a Deep Agent.",
"python": "https://github.com/langchain-ai/deep-agent-template/archive/refs/heads/main.zip",
"js": "https://github.com/langchain-ai/deep-agent-template-js/archive/refs/heads/main.zip",
},
"Agent": {
"description": "A simple agent that can be flexibly extended to many tools.",
"python": "https://github.com/langchain-ai/simple-agent-template/archive/refs/heads/main.zip",
},
"New LangGraph Project": { "New LangGraph Project": {
"description": "A simple, minimal chatbot with memory.", "description": "A simple, minimal chatbot with memory.",
"python": "https://github.com/langchain-ai/new-langgraph-project/archive/refs/heads/main.zip", "python": "https://github.com/langchain-ai/new-langgraph-project/archive/refs/heads/main.zip",
"js": "https://github.com/langchain-ai/new-langgraphjs-project/archive/refs/heads/main.zip", "js": "https://github.com/langchain-ai/new-langgraphjs-project/archive/refs/heads/main.zip",
}, },
"ReAct Agent": {
"description": "A simple agent that can be flexibly extended to many tools.",
"python": "https://github.com/langchain-ai/react-agent/archive/refs/heads/main.zip",
"js": "https://github.com/langchain-ai/react-agent-js/archive/refs/heads/main.zip",
},
"Memory Agent": {
"description": "A ReAct-style agent with an additional tool to store memories for use across conversational threads.",
"python": "https://github.com/langchain-ai/memory-agent/archive/refs/heads/main.zip",
"js": "https://github.com/langchain-ai/memory-agent-js/archive/refs/heads/main.zip",
},
"Retrieval Agent": {
"description": "An agent that includes a retrieval-based question-answering system.",
"python": "https://github.com/langchain-ai/retrieval-agent-template/archive/refs/heads/main.zip",
"js": "https://github.com/langchain-ai/retrieval-agent-template-js/archive/refs/heads/main.zip",
},
"Data-enrichment Agent": {
"description": "An agent that performs web searches and organizes its findings into a structured format.",
"python": "https://github.com/langchain-ai/data-enrichment/archive/refs/heads/main.zip",
"js": "https://github.com/langchain-ai/data-enrichment-js/archive/refs/heads/main.zip",
},
} }
# Generate TEMPLATE_IDS programmatically # Generate TEMPLATE_IDS programmatically
@@ -67,25 +78,19 @@ def _choose_template() -> str:
click.secho("❌ Invalid choice. Please try again.", fg="red") click.secho("❌ Invalid choice. Please try again.", fg="red")
return _choose_template() return _choose_template()
template_info = TEMPLATES[selected_template] # Prompt the user to choose between Python or JS/TS version
available_langs = [lang for lang in ("python", "js") if lang in template_info]
click.secho( click.secho(
f"\nYou selected: {selected_template} - {template_info['description']}", f"\nYou selected: {selected_template} - {TEMPLATES[selected_template]['description']}",
fg="green", fg="green",
) )
if len(available_langs) == 1:
return template_info[available_langs[0]]
version_choice: int = click.prompt( version_choice: int = click.prompt(
"Choose language (1 for Python 🐍, 2 for JS/TS 🌐)", type=int "Choose language (1 for Python 🐍, 2 for JS/TS 🌐)", type=int
) )
if version_choice == 1: if version_choice == 1:
return template_info["python"] return TEMPLATES[selected_template]["python"]
elif version_choice == 2: elif version_choice == 2:
return template_info["js"] return TEMPLATES[selected_template]["js"]
else: else:
click.secho("❌ Invalid choice. Please try again.", fg="red") click.secho("❌ Invalid choice. Please try again.", fg="red")
return _choose_template() return _choose_template()
@@ -125,6 +130,37 @@ def _download_repo_with_requests(repo_url: str, path: str) -> None:
sys.exit(1) sys.exit(1)
def _get_template_url(template_name: str) -> str | None:
"""
Retrieves the template URL based on the provided template name.
Args:
template_name: The name of the template.
Returns:
Optional[str]: The URL of the template if found, else None.
"""
if template_name in TEMPLATES:
click.secho(f"Template selected: {template_name}", fg="green")
version_choice: int = click.prompt(
"Choose version (1 for Python 🐍, 2 for JS/TS 🌐)", type=int
)
if version_choice == 1:
return TEMPLATES[template_name]["python"]
elif version_choice == 2:
return TEMPLATES[template_name]["js"]
else:
click.secho("❌ Invalid choice. Please try again.", fg="red")
return None
else:
click.secho(
f"Template '{template_name}' not found. Please select from the available options.",
fg="red",
)
return None
def create_new(path: str | None, template: str | None) -> None: def create_new(path: str | None, template: str | None) -> None:
"""Create a new LangGraph project at the specified PATH using the chosen TEMPLATE. """Create a new LangGraph project at the specified PATH using the chosen TEMPLATE.
@@ -3,17 +3,14 @@ import json
import os import os
import click import click
import httpx
import pytest import pytest
from langgraph_cli.cli import ( from langgraph_cli.cli import (
_call_host_backend_with_optional_tenant,
_docker_config_for_token, _docker_config_for_token,
_normalize_image_name, _normalize_image_name,
_normalize_image_tag, _normalize_image_tag,
_parse_env_from_config, _parse_env_from_config,
) )
from langgraph_cli.host_backend import HostBackendClient, HostBackendError
class TestDockerConfigForToken: class TestDockerConfigForToken:
@@ -135,92 +132,3 @@ class TestParseEnvFromConfig:
assert result["GOOD"] == "value" assert result["GOOD"] == "value"
# EMPTY= gives empty string, not None, so it should be present # EMPTY= gives empty string, not None, so it should be present
assert result["EMPTY"] == "" assert result["EMPTY"] == ""
class TestCallHostBackendWithOptionalTenant:
def _make_client(self, handler):
c = HostBackendClient("https://api.example.com", "test-key")
c._client = httpx.Client(
base_url="https://api.example.com",
transport=httpx.MockTransport(handler),
headers={"X-Api-Key": "test-key", "Accept": "application/json"},
timeout=30,
)
return c
def _make_eu_client(self, handler):
c = HostBackendClient("https://eu.api.host.langchain.com", "test-key")
c._client = httpx.Client(
base_url="https://eu.api.host.langchain.com",
transport=httpx.MockTransport(handler),
headers={"X-Api-Key": "test-key", "Accept": "application/json"},
timeout=30,
)
return c
def test_success_passes_through(self):
client = self._make_client(lambda req: httpx.Response(200, json={"ok": True}))
result = _call_host_backend_with_optional_tenant(
client, lambda c: c.list_deployments()
)
assert result == {"ok": True}
def test_403_not_enabled_gives_actionable_error(self):
detail = (
'{"detail":"LangSmith Deployment is not enabled for this organization"}'
)
client = self._make_client(lambda req: httpx.Response(403, text=detail))
with pytest.raises(HostBackendError, match="not enabled") as exc_info:
_call_host_backend_with_optional_tenant(
client, lambda c: c.list_deployments()
)
assert exc_info.value.status_code == 403
assert "smith.langchain.com" in exc_info.value.message
def test_403_not_enabled_eu_url(self):
detail = (
'{"detail":"LangSmith Deployment is not enabled for this organization"}'
)
client = self._make_eu_client(lambda req: httpx.Response(403, text=detail))
with pytest.raises(HostBackendError, match="not enabled") as exc_info:
_call_host_backend_with_optional_tenant(
client, lambda c: c.list_deployments()
)
assert "eu.smith.langchain.com" in exc_info.value.message
def test_workspace_retry_then_not_enabled_gives_actionable_error(self, monkeypatch):
requires_workspace = '{"detail":"requires workspace specification"}'
not_enabled = (
'{"detail":"LangSmith Deployment is not enabled for this organization"}'
)
seen_tenant_ids = []
def handler(req):
seen_tenant_ids.append(req.headers.get("X-Tenant-ID"))
if len(seen_tenant_ids) == 1:
return httpx.Response(403, text=requires_workspace)
if len(seen_tenant_ids) == 2:
return httpx.Response(403, text=not_enabled)
raise AssertionError("unexpected extra request")
monkeypatch.setattr(click, "prompt", lambda _text: "workspace-123")
client = self._make_client(handler)
with pytest.raises(HostBackendError, match="not enabled") as exc_info:
_call_host_backend_with_optional_tenant(
client, lambda c: c.list_deployments()
)
assert exc_info.value.status_code == 403
assert "smith.langchain.com" in exc_info.value.message
assert seen_tenant_ids == [None, "workspace-123"]
assert client._client.headers["X-Tenant-ID"] == "workspace-123"
def test_other_403_re_raises_original(self):
client = self._make_client(
lambda req: httpx.Response(403, text='{"detail":"some other error"}')
)
with pytest.raises(HostBackendError, match="some other error"):
_call_host_backend_with_optional_tenant(
client, lambda c: c.list_deployments()
)
+10 -1
View File
@@ -25,7 +25,7 @@ except ImportError:
_StreamingCallbackHandler = object # type: ignore _StreamingCallbackHandler = object # type: ignore
T = TypeVar("T") T = TypeVar("T")
Meta = tuple[tuple[str, ...], dict[str, Any]] Meta = tuple[tuple[str, ...], dict[str, Any] | None]
def _state_values(obj: Any) -> Sequence[Any]: def _state_values(obj: Any) -> Sequence[Any]:
@@ -56,6 +56,7 @@ class StreamMessagesHandler(BaseCallbackHandler, _StreamingCallbackHandler):
subgraphs: bool, subgraphs: bool,
*, *,
parent_ns: tuple[str, ...] | None = None, parent_ns: tuple[str, ...] | None = None,
dedupe_metadata: bool = False,
) -> None: ) -> None:
"""Configure the handler to stream messages from LLMs and nodes. """Configure the handler to stream messages from LLMs and nodes.
@@ -84,8 +85,10 @@ class StreamMessagesHandler(BaseCallbackHandler, _StreamingCallbackHandler):
self.stream = stream self.stream = stream
self.subgraphs = subgraphs self.subgraphs = subgraphs
self.metadata: dict[UUID, Meta] = {} self.metadata: dict[UUID, Meta] = {}
self.emitted_metadata: set[UUID] = set()
self.seen: set[int | str] = set() self.seen: set[int | str] = set()
self.parent_ns = parent_ns self.parent_ns = parent_ns
self.dedupe_metadata = dedupe_metadata
def _emit(self, meta: Meta, message: BaseMessage, *, dedupe: bool = False) -> None: def _emit(self, meta: Meta, message: BaseMessage, *, dedupe: bool = False) -> None:
if dedupe and message.id in self.seen: if dedupe and message.id in self.seen:
@@ -155,6 +158,10 @@ class StreamMessagesHandler(BaseCallbackHandler, _StreamingCallbackHandler):
if not isinstance(chunk, ChatGenerationChunk): if not isinstance(chunk, ChatGenerationChunk):
return return
if meta := self.metadata.get(run_id): if meta := self.metadata.get(run_id):
if self.dedupe_metadata and run_id in self.emitted_metadata:
meta = (meta[0], None)
else:
self.emitted_metadata.add(run_id)
self._emit(meta, chunk.message) self._emit(meta, chunk.message)
def on_llm_end( def on_llm_end(
@@ -170,6 +177,7 @@ class StreamMessagesHandler(BaseCallbackHandler, _StreamingCallbackHandler):
gen = response.generations[0][0] gen = response.generations[0][0]
if isinstance(gen, ChatGeneration): if isinstance(gen, ChatGeneration):
self._emit(meta, gen.message, dedupe=True) self._emit(meta, gen.message, dedupe=True)
self.emitted_metadata.discard(run_id)
self.metadata.pop(run_id, None) self.metadata.pop(run_id, None)
def on_llm_error( def on_llm_error(
@@ -180,6 +188,7 @@ class StreamMessagesHandler(BaseCallbackHandler, _StreamingCallbackHandler):
parent_run_id: UUID | None = None, parent_run_id: UUID | None = None,
**kwargs: Any, **kwargs: Any,
) -> Any: ) -> Any:
self.emitted_metadata.discard(run_id)
self.metadata.pop(run_id, None) self.metadata.pop(run_id, None)
def on_chain_start( def on_chain_start(
+2
View File
@@ -2614,6 +2614,7 @@ class Pregel(
stream.put, stream.put,
subgraphs, subgraphs,
parent_ns=tuple(ns_.split(NS_SEP)) if ns_ else None, parent_ns=tuple(ns_.split(NS_SEP)) if ns_ else None,
dedupe_metadata="compact" in stream_modes,
) )
) )
@@ -2965,6 +2966,7 @@ class Pregel(
stream_put, stream_put,
subgraphs, subgraphs,
parent_ns=tuple(ns_.split(NS_SEP)) if ns_ else None, parent_ns=tuple(ns_.split(NS_SEP)) if ns_ else None,
dedupe_metadata="compact" in stream_modes,
) )
) )
+59 -12
View File
@@ -109,6 +109,45 @@ class RemoteException(Exception):
pass pass
def _restore_message_metadata(
data: Any, metadata_by_message_id: dict[str, dict[str, Any]]
) -> Any:
"""Restore deduplicated message metadata using the message id as cache key."""
if not (isinstance(data, list) and len(data) == 2):
return data
message, metadata = data
if not isinstance(message, dict):
return data
message_id = message.get("id")
if isinstance(message_id, str):
if isinstance(metadata, dict):
metadata_by_message_id[message_id] = metadata
else:
metadata = metadata_by_message_id.get(message_id)
return (message, metadata)
def _merge_values_patch(
ns: tuple[str, ...],
mode: str,
data: Any,
values_by_ns: dict[tuple[str, ...], dict[str, Any]],
) -> tuple[str, Any]:
"""Merge `values-patch` events back into full values snapshots."""
if mode != "values-patch" or not isinstance(data, dict):
return mode, data
values = data.get("values")
if not isinstance(values, dict):
return "values", values if values is not None else {}
merged = dict(values_by_ns.get(ns, {}))
merged.update(values)
for key in data.get("deleted_keys", ()):
if isinstance(key, str):
merged.pop(key, None)
values_by_ns[ns] = merged
return "values", merged
class RemoteGraph(PregelProtocol): class RemoteGraph(PregelProtocol):
"""The `RemoteGraph` class is a client implementation for calling remote """The `RemoteGraph` class is a client implementation for calling remote
APIs that implement the LangGraph Server API specification. APIs that implement the LangGraph Server API specification.
@@ -766,6 +805,8 @@ class RemoteGraph(PregelProtocol):
else: else:
command = None command = None
thread_id = sanitized_config.get("configurable", {}).pop("thread_id", None) thread_id = sanitized_config.get("configurable", {}).pop("thread_id", None)
message_metadata_by_id: dict[str, dict[str, Any]] = {}
values_by_ns: dict[tuple[str, ...], dict[str, Any]] = {}
for chunk in sync_client.runs.stream( for chunk in sync_client.runs.stream(
thread_id=thread_id, thread_id=thread_id,
@@ -798,6 +839,11 @@ class RemoteGraph(PregelProtocol):
if caller_ns := (config or {}).get(CONF, {}).get(CONFIG_KEY_CHECKPOINT_NS): if caller_ns := (config or {}).get(CONF, {}).get(CONFIG_KEY_CHECKPOINT_NS):
caller_ns = tuple(caller_ns.split(NS_SEP)) caller_ns = tuple(caller_ns.split(NS_SEP))
ns = caller_ns + ns ns = caller_ns + ns
mode, data = _merge_values_patch(ns, mode, chunk.data, values_by_ns)
if mode != chunk.event:
chunk = chunk._replace(data=data)
elif data is not chunk.data:
chunk = chunk._replace(data=data)
# stream to parent stream # stream to parent stream
if stream is not None and mode in stream.modes: if stream is not None and mode in stream.modes:
stream((ns, mode, chunk.data)) stream((ns, mode, chunk.data))
@@ -815,7 +861,9 @@ class RemoteGraph(PregelProtocol):
continue continue
if chunk.event.startswith("messages"): if chunk.event.startswith("messages"):
chunk = chunk._replace(data=tuple(chunk.data)) chunk = chunk._replace(
data=_restore_message_metadata(chunk.data, message_metadata_by_id)
)
# emit chunk # emit chunk
if version == "v2": if version == "v2":
@@ -827,11 +875,6 @@ class RemoteGraph(PregelProtocol):
) )
yield {"type": mode, "ns": ns, "data": chunk.data, "interrupts": ints} yield {"type": mode, "ns": ns, "data": chunk.data, "interrupts": ints}
elif subgraphs: elif subgraphs:
if NS_SEP in chunk.event:
mode, ns_ = chunk.event.split(NS_SEP, 1)
ns = tuple(ns_.split(NS_SEP))
else:
mode, ns = chunk.event, ()
if req_single: if req_single:
yield ns, chunk.data yield ns, chunk.data
else: else:
@@ -921,6 +964,8 @@ class RemoteGraph(PregelProtocol):
else: else:
command = None command = None
thread_id = sanitized_config.get("configurable", {}).pop("thread_id", None) thread_id = sanitized_config.get("configurable", {}).pop("thread_id", None)
message_metadata_by_id: dict[str, dict[str, Any]] = {}
values_by_ns: dict[tuple[str, ...], dict[str, Any]] = {}
async for chunk in client.runs.stream( async for chunk in client.runs.stream(
thread_id=thread_id, thread_id=thread_id,
@@ -953,6 +998,11 @@ class RemoteGraph(PregelProtocol):
if caller_ns := (config or {}).get(CONF, {}).get(CONFIG_KEY_CHECKPOINT_NS): if caller_ns := (config or {}).get(CONF, {}).get(CONFIG_KEY_CHECKPOINT_NS):
caller_ns = tuple(caller_ns.split(NS_SEP)) caller_ns = tuple(caller_ns.split(NS_SEP))
ns = caller_ns + ns ns = caller_ns + ns
mode, data = _merge_values_patch(ns, mode, chunk.data, values_by_ns)
if mode != chunk.event:
chunk = chunk._replace(data=data)
elif data is not chunk.data:
chunk = chunk._replace(data=data)
# stream to parent stream # stream to parent stream
if stream is not None and mode in stream.modes: if stream is not None and mode in stream.modes:
stream((ns, mode, chunk.data)) stream((ns, mode, chunk.data))
@@ -970,7 +1020,9 @@ class RemoteGraph(PregelProtocol):
continue continue
if chunk.event.startswith("messages"): if chunk.event.startswith("messages"):
chunk = chunk._replace(data=tuple(chunk.data)) chunk = chunk._replace(
data=_restore_message_metadata(chunk.data, message_metadata_by_id)
)
# emit chunk # emit chunk
if version == "v2": if version == "v2":
@@ -982,11 +1034,6 @@ class RemoteGraph(PregelProtocol):
) )
yield {"type": mode, "ns": ns, "data": chunk.data, "interrupts": ints} yield {"type": mode, "ns": ns, "data": chunk.data, "interrupts": ints}
elif subgraphs: elif subgraphs:
if NS_SEP in chunk.event:
mode, ns_ = chunk.event.split(NS_SEP, 1)
ns = tuple(ns_.split(NS_SEP))
else:
mode, ns = chunk.event, ()
if req_single: if req_single:
yield ns, chunk.data yield ns, chunk.data
else: else:
+13 -5
View File
@@ -116,7 +116,14 @@ def ensure_valid_checkpointer(checkpointer: Checkpointer) -> Checkpointer:
StreamMode = Literal[ StreamMode = Literal[
"values", "updates", "checkpoints", "tasks", "debug", "messages", "custom" "values",
"updates",
"checkpoints",
"tasks",
"debug",
"messages",
"custom",
"compact",
] ]
"""How the stream method should emit outputs. """How the stream method should emit outputs.
@@ -275,13 +282,14 @@ class MessagesStreamPart(TypedDict):
"""Stream part emitted for `stream_mode="messages"`. """Stream part emitted for `stream_mode="messages"`.
`data` is a 2-tuple of `(message, metadata)` where `message` is a `data` is a 2-tuple of `(message, metadata)` where `message` is a
`BaseMessage` (e.g. `AIMessageChunk`) and `metadata` is a dict containing `BaseMessage` (e.g. `AIMessageChunk`) and `metadata` is either a dict containing
keys like `langgraph_step`, `langgraph_node`, `langgraph_triggers`, etc. keys like `langgraph_step`, `langgraph_node`, `langgraph_triggers`, etc. or
`None` for deduplicated follow-up chunks when `stream_mode` includes `"compact"`.
""" """
type: Literal["messages"] type: Literal["messages"]
ns: tuple[str, ...] ns: tuple[str, ...]
data: tuple[AnyMessage, dict[str, Any]] data: tuple[AnyMessage, dict[str, Any] | None]
class CustomStreamPart(TypedDict): class CustomStreamPart(TypedDict):
@@ -346,7 +354,7 @@ async for part in graph.astream(input, version="v2"):
if part["type"] == "values": if part["type"] == "values":
part["data"] # OutputT — full state (pydantic/dataclass/dict) part["data"] # OutputT — full state (pydantic/dataclass/dict)
elif part["type"] == "messages": elif part["type"] == "messages":
part["data"] # tuple[BaseMessage, dict] — (message, metadata) part["data"] # tuple[BaseMessage, dict | None] — (message, metadata)
elif part["type"] == "custom": elif part["type"] == "custom":
part["data"] # Any — user-defined part["data"] # Any — user-defined
``` ```
+151
View File
@@ -882,6 +882,80 @@ def test_stream_sanitizes_thread_id():
assert not passed_config["configurable"] assert not passed_config["configurable"]
def test_stream_restores_messages_and_merges_values_patch():
mock_sync_client = MagicMock()
mock_sync_client.runs.stream.return_value = [
StreamPart(
event="messages|tools:call_1",
data=[
{"id": "msg-1", "type": "AIMessageChunk", "content": "hel"},
{
"langgraph_checkpoint_ns": "tools:call_1",
"langgraph_node": "agent",
},
],
),
StreamPart(
event="messages|tools:call_1",
data=[
{"id": "msg-1", "type": "AIMessageChunk", "content": "lo"},
None,
],
),
StreamPart(
event="values|tools:call_1",
data={"messages": [{"type": "human", "content": "hi"}], "count": 1},
),
StreamPart(
event="values-patch|tools:call_1",
data={"values": {"count": 2}, "deleted_keys": ["messages"]},
),
]
remote_pregel = RemoteGraph("test_graph_id", sync_client=mock_sync_client)
parts = list(
remote_pregel.stream(
{"input": "data"},
config={"configurable": {"thread_id": "thread_1"}},
stream_mode=["messages", "values", "compact"],
subgraphs=True,
version="v2",
)
)
message_parts = [part for part in parts if part["type"] == "messages"]
assert message_parts[0]["data"][1] == {
"langgraph_checkpoint_ns": "tools:call_1",
"langgraph_node": "agent",
}
assert message_parts[1]["data"][1] == {
"langgraph_checkpoint_ns": "tools:call_1",
"langgraph_node": "agent",
}
value_parts = [part for part in parts if part["type"] == "values"]
assert value_parts == [
{
"type": "values",
"ns": ("tools:call_1",),
"data": {"messages": [{"type": "human", "content": "hi"}], "count": 1},
"interrupts": (),
},
{
"type": "values",
"ns": ("tools:call_1",),
"data": {"count": 2},
"interrupts": (),
},
]
_, kwargs = mock_sync_client.runs.stream.call_args
assert set(kwargs["stream_mode"]) == {
"messages-tuple",
"values",
"compact",
"updates",
}
@pytest.mark.anyio @pytest.mark.anyio
async def test_ainvoke(): async def test_ainvoke():
# set up test # set up test
@@ -1092,6 +1166,83 @@ def test_stream_context_base_model():
assert kwargs["context"] == ctx assert kwargs["context"] == ctx
@pytest.mark.anyio
async def test_astream_restores_messages_and_merges_values_patch():
mock_async_client = MagicMock()
async_iter = MagicMock()
async_iter.__aiter__.return_value = [
StreamPart(
event="messages|tools:call_1",
data=[
{"id": "msg-1", "type": "AIMessageChunk", "content": "hel"},
{
"langgraph_checkpoint_ns": "tools:call_1",
"langgraph_node": "agent",
},
],
),
StreamPart(
event="messages|tools:call_1",
data=[
{"id": "msg-1", "type": "AIMessageChunk", "content": "lo"},
None,
],
),
StreamPart(
event="values|tools:call_1",
data={"messages": [{"type": "human", "content": "hi"}], "count": 1},
),
StreamPart(
event="values-patch|tools:call_1",
data={"values": {"count": 2}, "deleted_keys": ["messages"]},
),
]
mock_async_client.runs.stream.return_value = async_iter
remote_pregel = RemoteGraph("test_graph_id", client=mock_async_client)
parts = []
async for part in remote_pregel.astream(
{"input": "data"},
config={"configurable": {"thread_id": "thread_1"}},
stream_mode=["messages", "values", "compact"],
subgraphs=True,
version="v2",
):
parts.append(part)
message_parts = [part for part in parts if part["type"] == "messages"]
assert message_parts[0]["data"][1] == {
"langgraph_checkpoint_ns": "tools:call_1",
"langgraph_node": "agent",
}
assert message_parts[1]["data"][1] == {
"langgraph_checkpoint_ns": "tools:call_1",
"langgraph_node": "agent",
}
value_parts = [part for part in parts if part["type"] == "values"]
assert value_parts == [
{
"type": "values",
"ns": ("tools:call_1",),
"data": {"messages": [{"type": "human", "content": "hi"}], "count": 1},
"interrupts": (),
},
{
"type": "values",
"ns": ("tools:call_1",),
"data": {"count": 2},
"interrupts": (),
},
]
_, kwargs = mock_async_client.runs.stream.call_args
assert set(kwargs["stream_mode"]) == {
"messages-tuple",
"values",
"compact",
"updates",
}
@pytest.mark.skip( @pytest.mark.skip(
"Unskip this test to manually test the LangSmith Deployment integration" "Unskip this test to manually test the LangSmith Deployment integration"
) )
+82
View File
@@ -90,6 +90,25 @@ def _make_messages_graph() -> StateGraph[
return builder return builder
def _make_streaming_messages_graph() -> StateGraph[
MessagesState, None, MessagesState, MessagesState
]:
model = FakeChatModel(messages=[AIMessage(content="hello world", id="ai-1")])
def call_model(state: MessagesState) -> dict[str, Any]:
streamed = model.stream(state["messages"])
message = next(streamed)
for chunk in streamed:
message += chunk
return {"messages": message}
builder = StateGraph(MessagesState, input_schema=MessagesState)
builder.add_node("call_model", call_model)
builder.add_edge(START, "call_model")
builder.add_edge("call_model", END)
return builder
def _make_custom_graph() -> Any: def _make_custom_graph() -> Any:
@entrypoint() @entrypoint()
def graph(inputs: Any, *, writer: StreamWriter) -> Any: def graph(inputs: Any, *, writer: StreamWriter) -> Any:
@@ -164,6 +183,26 @@ class TestV1BackwardsCompat:
ns, _data = chunk ns, _data = chunk
assert isinstance(ns, tuple) assert isinstance(ns, tuple)
def test_stream_v1_messages_keep_metadata_on_every_chunk(self) -> None:
graph = _make_streaming_messages_graph().compile()
chunks = list(graph.stream(_MSG_INPUT, stream_mode="messages"))
metadata = [meta for _message, meta in chunks]
assert len(metadata) >= 3
assert all(isinstance(meta, dict) for meta in metadata)
def test_stream_v1_messages_compact_dedupes_metadata(self) -> None:
graph = _make_streaming_messages_graph().compile()
chunks = list(graph.stream(_MSG_INPUT, stream_mode=["messages", "compact"]))
metadata = [
meta
for mode, payload in chunks
if mode == "messages"
for _message, meta in [payload]
]
assert len(metadata) >= 3
assert isinstance(metadata[0], dict)
assert all(meta is None for meta in metadata[1:])
# --- v2 sync stream --- # --- v2 sync stream ---
@@ -205,6 +244,26 @@ class TestV2Stream:
assert isinstance(metadata, dict) assert isinstance(metadata, dict)
assert "langgraph_node" in metadata assert "langgraph_node" in metadata
def test_messages_streaming_compact_dedupes_metadata(self) -> None:
graph = _make_streaming_messages_graph().compile()
chunks = list(
graph.stream(
_MSG_INPUT,
stream_mode=["messages", "compact"],
version="v2",
)
)
msg_chunks = [c for c in chunks if c["type"] == "messages"]
assert len(msg_chunks) >= 3
first_message, first_metadata = msg_chunks[0]["data"]
assert isinstance(first_message, BaseMessage)
assert isinstance(first_metadata, dict)
assert "langgraph_node" in first_metadata
for chunk in msg_chunks[1:]:
message, metadata = chunk["data"]
assert isinstance(message, BaseMessage)
assert metadata is None
def test_custom(self) -> None: def test_custom(self) -> None:
graph = _make_custom_graph() graph = _make_custom_graph()
chunks = list(graph.stream({"key": "val"}, stream_mode="custom", version="v2")) chunks = list(graph.stream({"key": "val"}, stream_mode="custom", version="v2"))
@@ -544,6 +603,29 @@ class TestV2StreamAsync:
assert isinstance(metadata, dict) assert isinstance(metadata, dict)
assert "langgraph_node" in metadata assert "langgraph_node" in metadata
@NEEDS_CONTEXTVARS
@pytest.mark.anyio
async def test_messages_streaming_compact_dedupes_metadata(self) -> None:
graph = _make_streaming_messages_graph().compile()
chunks = [
c
async for c in graph.astream(
_MSG_INPUT,
stream_mode=["messages", "compact"],
version="v2",
)
]
msg_chunks = [c for c in chunks if c["type"] == "messages"]
assert len(msg_chunks) >= 3
first_message, first_metadata = msg_chunks[0]["data"]
assert isinstance(first_message, BaseMessage)
assert isinstance(first_metadata, dict)
assert "langgraph_node" in first_metadata
for chunk in msg_chunks[1:]:
message, metadata = chunk["data"]
assert isinstance(message, BaseMessage)
assert metadata is None
@NEEDS_CONTEXTVARS @NEEDS_CONTEXTVARS
@pytest.mark.anyio @pytest.mark.anyio
async def test_custom(self) -> None: async def test_custom(self) -> None:
+23
View File
@@ -58,6 +58,7 @@ StreamMode = Literal[
"debug", "debug",
"custom", "custom",
"messages-tuple", "messages-tuple",
"compact",
] ]
""" """
Defines the mode of streaming: Defines the mode of streaming:
@@ -69,6 +70,7 @@ Defines the mode of streaming:
- "tasks": Stream task start and finish events. - "tasks": Stream task start and finish events.
- "debug": Stream detailed debug information. - "debug": Stream detailed debug information.
- "custom": Stream custom events. - "custom": Stream custom events.
- "compact": Enable compact streaming payloads for other selected modes.
""" """
DisconnectMode = Literal["cancel", "continue"] DisconnectMode = Literal["cancel", "continue"]
@@ -733,6 +735,26 @@ class ValuesStreamPart(TypedDict):
"""List of interrupts that occurred during this step.""" """List of interrupts that occurred during this step."""
class ValuesPatchPayload(TypedDict):
"""Incremental patch payload for subgraph `values` updates."""
values: dict[str, Any]
"""Only the changed fields since the previous `values` event for this namespace."""
deleted_keys: NotRequired[list[str]]
"""Optional list of keys that were removed from the previous values snapshot."""
class ValuesPatchStreamPart(TypedDict):
"""Stream part emitted for incremental subgraph value patches (`values-patch`)."""
type: Literal["values-patch"]
"""Stream part type discriminator."""
ns: list[str]
"""Namespace path of the emitting node (empty for root graph)."""
data: ValuesPatchPayload
"""Incremental state patch for the namespace."""
class UpdatesStreamPart(TypedDict): class UpdatesStreamPart(TypedDict):
"""Stream part emitted for `stream_mode="updates"`.""" """Stream part emitted for `stream_mode="updates"`."""
@@ -845,6 +867,7 @@ class MetadataStreamPart(TypedDict):
StreamPartV2 = ( StreamPartV2 = (
ValuesStreamPart ValuesStreamPart
| ValuesPatchStreamPart
| UpdatesStreamPart | UpdatesStreamPart
| MessagesPartialStreamPart | MessagesPartialStreamPart
| MessagesCompleteStreamPart | MessagesCompleteStreamPart
+81
View File
@@ -1,5 +1,6 @@
from __future__ import annotations from __future__ import annotations
import json
from collections.abc import Iterator, Sequence from collections.abc import Iterator, Sequence
from pathlib import Path from pathlib import Path
from typing import Any from typing import Any
@@ -8,7 +9,9 @@ import httpx
import pytest import pytest
from typing_extensions import assert_type from typing_extensions import assert_type
from langgraph_sdk._async.runs import RunsClient
from langgraph_sdk._shared.utilities import _sse_to_v2_dict from langgraph_sdk._shared.utilities import _sse_to_v2_dict
from langgraph_sdk._sync.runs import SyncRunsClient
from langgraph_sdk.client import HttpClient, SyncHttpClient from langgraph_sdk.client import HttpClient, SyncHttpClient
from langgraph_sdk.schema import ( from langgraph_sdk.schema import (
CheckpointPayload, CheckpointPayload,
@@ -24,6 +27,7 @@ from langgraph_sdk.schema import (
TaskResultPayload, TaskResultPayload,
TasksStreamPart, TasksStreamPart,
UpdatesStreamPart, UpdatesStreamPart,
ValuesPatchStreamPart,
ValuesStreamPart, ValuesStreamPart,
) )
from langgraph_sdk.sse import BytesLike, BytesLineDecoder, SSEDecoder from langgraph_sdk.sse import BytesLike, BytesLineDecoder, SSEDecoder
@@ -375,6 +379,81 @@ def test_sse_to_v2_dict_values_with_interrupts() -> None:
assert "__interrupt__" not in result["data"] assert "__interrupt__" not in result["data"]
def test_sse_to_v2_dict_values_patch() -> None:
payload = {"values": {"count": 2}, "deleted_keys": ["stale"]}
result = _sse_to_v2_dict("values-patch|tools:call_1", payload)
assert result is not None
_assert_v2_shape(result)
assert result == {
"type": "values-patch",
"ns": ["tools:call_1"],
"data": {"values": {"count": 2}, "deleted_keys": ["stale"]},
"interrupts": [],
}
@pytest.mark.asyncio
async def test_async_runs_stream_includes_compact_mode():
async def handler(request: httpx.Request) -> httpx.Response:
assert request.method == "POST"
assert request.url.path == "/runs/stream"
body = json.loads(request.content)
assert body["stream_mode"] == ["values", "compact"]
return httpx.Response(
200,
headers={"Content-Type": "text/event-stream"},
content=b"event: end\ndata: null\n\n",
)
transport = httpx.MockTransport(handler)
async with httpx.AsyncClient(
transport=transport, base_url="https://example.com"
) as client:
runs_client = RunsClient(HttpClient(client))
parts = [
part
async for part in runs_client.stream(
thread_id=None,
assistant_id="agent",
input={"messages": []},
stream_mode=["values", "compact"],
)
]
assert len(parts) == 1
assert parts[0].event == "end"
assert parts[0].data is None
def test_sync_runs_stream_includes_compact_mode():
def handler(request: httpx.Request) -> httpx.Response:
assert request.method == "POST"
assert request.url.path == "/runs/stream"
body = json.loads(request.content)
assert body["stream_mode"] == ["values", "compact"]
return httpx.Response(
200,
headers={"Content-Type": "text/event-stream"},
content=b"event: end\ndata: null\n\n",
)
transport = httpx.MockTransport(handler)
with httpx.Client(transport=transport, base_url="https://example.com") as client:
runs_client = SyncRunsClient(SyncHttpClient(client))
parts = list(
runs_client.stream(
thread_id=None,
assistant_id="agent",
input={"messages": []},
stream_mode=["values", "compact"],
)
)
assert len(parts) == 1
assert parts[0].event == "end"
assert parts[0].data is None
# --- client-side v2 stream wrapping --- # --- client-side v2 stream wrapping ---
@@ -448,6 +527,8 @@ def _check_v2_type_narrowing(part: StreamPartV2) -> None:
if part["type"] == "values": if part["type"] == "values":
assert_type(part, ValuesStreamPart) assert_type(part, ValuesStreamPart)
assert_type(part["data"], dict[str, Any]) assert_type(part["data"], dict[str, Any])
elif part["type"] == "values-patch":
assert_type(part, ValuesPatchStreamPart)
elif part["type"] == "updates": elif part["type"] == "updates":
assert_type(part, UpdatesStreamPart) assert_type(part, UpdatesStreamPart)
assert_type(part["data"], dict[str, Any]) assert_type(part["data"], dict[str, Any])