mirror of
https://github.com/langchain-ai/langgraph.git
synced 2026-11-04 05:11:20 +01:00
Compare commits
11
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
bc1b23ffdc | ||
|
|
95cdb86d28 | ||
|
|
0a05f15e21 | ||
|
|
2a30615bd7 | ||
|
|
5c96bce8d3 | ||
|
|
2f46bb3ca4 | ||
|
|
4b7f2f5230 | ||
|
|
5a3e5dd739 | ||
|
|
eb8ec111b3 | ||
|
|
8c177f7388 | ||
|
|
55013146ae |
No files matched your search
@@ -86,6 +86,47 @@ langgraph dockerfile SAVE_PATH [OPTIONS]
|
||||
-c, --config FILE Config file path
|
||||
```
|
||||
|
||||
### `langgraph deploy metrics`
|
||||
|
||||
Fetch deployment monitoring time series as JSON through the LangSmith control plane.
|
||||
Use `langgraph deploy list` to find the deployment ID. Authentication uses the same
|
||||
API key and workspace configuration as the other deployment commands.
|
||||
|
||||
```bash
|
||||
langgraph deploy metrics DEPLOYMENT_ID \
|
||||
--start-time 2026-10-28T00:00:00Z \
|
||||
--end-time 2026-10-28T01:00:00Z \
|
||||
--metrics cpu_usage,memory_usage,worker_count
|
||||
```
|
||||
|
||||
Omit `--metrics` to fetch all available metrics. Time bounds default to the start of
|
||||
the current hour through now, and the backend picks the resolution from the time
|
||||
range. Output contains
|
||||
charts and their time-series data, not traces. This command never prompts; for an
|
||||
organization-scoped key, set `LANGSMITH_TENANT_ID` to the workspace ID.
|
||||
|
||||
### `langgraph deploy scale`
|
||||
|
||||
Change a deployment's compute tier, dedicated Postgres tier, or both. Omitted
|
||||
settings remain unchanged. The backend enforces deployment permissions, supported
|
||||
tiers, and platform restrictions.
|
||||
|
||||
```bash
|
||||
langgraph deploy scale DEPLOYMENT_ID \
|
||||
--compute-tier DEDICATED_M --database-tier DEDICATED_L
|
||||
```
|
||||
|
||||
Compute tiers are `SERVERLESS_S`, `SERVERLESS_M`, `SERVERLESS_L`, `DEDICATED_S`,
|
||||
`DEDICATED_M`, and `DEDICATED_L`. Database tiers are available only for dedicated
|
||||
deployments: `DEDICATED_S`, `DEDICATED_M`, and `DEDICATED_L`. Availability depends
|
||||
on the deployment's platform and enabled features.
|
||||
|
||||
Scaling prompts for confirmation because it creates a revision, may interrupt
|
||||
service, and may change billing. Pass `--yes` to confirm without prompting, with
|
||||
credentials and workspace configured for non-interactive use. The command reports
|
||||
that scaling was requested; it does not wait for the revision to become healthy.
|
||||
Arbitrary worker replica overrides are not supported by this command.
|
||||
|
||||
## Configuration
|
||||
|
||||
The CLI uses a `langgraph.json` configuration file with these key settings:
|
||||
|
||||
@@ -2148,6 +2148,6 @@ yocto-queue@^0.1.0:
|
||||
integrity sha512-rVksvsnNCdJ/ohGc6xgPwyN8eheCxsiLM8mxuE/t/mOVqJewPuO1miLpTHQiRgTKCLexL4MeAFVagts7HmNZ2Q==
|
||||
|
||||
"zod@^3.25.76 || ^4":
|
||||
version "4.3.6"
|
||||
resolved "https://registry.yarnpkg.com/zod/-/zod-4.3.6.tgz#89c56e0aa7d2b05107d894412227087885ab112a"
|
||||
integrity sha512-rftlrkhHZOcjDwkGlnUtZZkvaPHCsDATp4pGpuOOMDaTdDDXF91wuVDJoWoPsKX/3YPQ5fHuF3STjcYyKr+Qhg==
|
||||
version "4.6.7"
|
||||
resolved "https://registry.yarnpkg.com/zod/-/zod-4.6.7.tgz#e806cf64a7f4e37635ce9f2f52c04c3766f75e5b"
|
||||
integrity sha512-g8W2D2kU6a8untc053fX1uVBKzM7yqAxW/3zwk83+6RyF2+SKOuV+c0wBDhUwzV1PaGapmZH+rmcUzPf+3CGwg==
|
||||
@@ -13,6 +13,7 @@ from contextlib import contextmanager
|
||||
from dataclasses import asdict, dataclass, field
|
||||
from datetime import datetime, timezone
|
||||
from typing import Generic, Protocol, TypeVar
|
||||
from uuid import UUID
|
||||
|
||||
import click
|
||||
import click.exceptions
|
||||
@@ -2635,6 +2636,109 @@ def deploy_list(
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@OPT_HOST_API_KEY
|
||||
@OPT_HOST_URL
|
||||
@click.argument("deployment_id", type=click.UUID)
|
||||
@click.option(
|
||||
"--start-time",
|
||||
help="ISO8601 start time. Defaults to the start of the current hour.",
|
||||
)
|
||||
@click.option("--end-time", help="ISO8601 end time. Defaults to now.")
|
||||
@click.option(
|
||||
"--metrics", help="Comma-separated metric names. Defaults to all metrics."
|
||||
)
|
||||
@deploy.command(
|
||||
"metrics", help="[Beta] Fetch deployment monitoring time series as JSON."
|
||||
)
|
||||
def deploy_metrics(
|
||||
api_key: str | None,
|
||||
host_url: str | None,
|
||||
deployment_id: UUID,
|
||||
start_time: str | None,
|
||||
end_time: str | None,
|
||||
metrics: str | None,
|
||||
) -> None:
|
||||
params = {
|
||||
key: value
|
||||
for key, value in {
|
||||
"start_time": start_time,
|
||||
"end_time": end_time,
|
||||
"metrics": metrics,
|
||||
}.items()
|
||||
if value is not None
|
||||
}
|
||||
client = _create_host_backend_client(host_url, api_key, no_input=True)
|
||||
result = _call_host_backend_with_optional_tenant(
|
||||
client,
|
||||
lambda c: c.get_deployment_metrics(str(deployment_id), params),
|
||||
no_input=True,
|
||||
)
|
||||
click.echo(json_mod.dumps(result))
|
||||
|
||||
|
||||
@OPT_HOST_API_KEY
|
||||
@OPT_HOST_URL
|
||||
@click.argument("deployment_id", type=click.UUID)
|
||||
@click.option(
|
||||
"--compute-tier",
|
||||
type=click.Choice(
|
||||
[
|
||||
"SERVERLESS_S",
|
||||
"SERVERLESS_M",
|
||||
"SERVERLESS_L",
|
||||
"DEDICATED_S",
|
||||
"DEDICATED_M",
|
||||
"DEDICATED_L",
|
||||
]
|
||||
),
|
||||
help="Compute resource tier. Omit to leave compute unchanged.",
|
||||
)
|
||||
@click.option(
|
||||
"--database-tier",
|
||||
type=click.Choice(["DEDICATED_S", "DEDICATED_M", "DEDICATED_L"]),
|
||||
help="Postgres resource tier for dedicated deployments. Omit to leave unchanged.",
|
||||
)
|
||||
@click.option(
|
||||
"--yes", is_flag=True, help="Confirm the scaling change without prompting."
|
||||
)
|
||||
@deploy.command(
|
||||
"scale",
|
||||
help="[Beta] Change compute or Postgres tiers. Creates a revision and may change billing.",
|
||||
)
|
||||
def deploy_scale(
|
||||
api_key: str | None,
|
||||
host_url: str | None,
|
||||
deployment_id: UUID,
|
||||
compute_tier: str | None,
|
||||
database_tier: str | None,
|
||||
yes: bool,
|
||||
) -> None:
|
||||
payload = {
|
||||
key: value
|
||||
for key, value in {
|
||||
"compute_tier": compute_tier,
|
||||
"database_tier": database_tier,
|
||||
}.items()
|
||||
if value is not None
|
||||
}
|
||||
if not payload:
|
||||
raise click.UsageError("Specify --compute-tier or --database-tier.")
|
||||
if not yes:
|
||||
changes = ", ".join(f"{key}={value}" for key, value in payload.items())
|
||||
click.confirm(
|
||||
f"Scale deployment {deployment_id} ({changes})? "
|
||||
"This creates a revision and may change billing or interrupt service.",
|
||||
abort=True,
|
||||
)
|
||||
client = _create_host_backend_client(host_url, api_key, no_input=yes)
|
||||
_call_host_backend_with_optional_tenant(
|
||||
client,
|
||||
lambda c: c.scale_deployment(str(deployment_id), payload),
|
||||
no_input=yes,
|
||||
)
|
||||
click.echo(f"Scaling requested for deployment {deployment_id}.")
|
||||
|
||||
|
||||
@deploy.group(
|
||||
"listeners",
|
||||
cls=NestedHelpGroup,
|
||||
|
||||
@@ -362,6 +362,18 @@ class HostBackendClient:
|
||||
f"/v2/deployments/{deployment_id}/revisions/{revision_id}",
|
||||
)
|
||||
|
||||
def get_deployment_metrics(
|
||||
self, deployment_id: str, params: dict[str, Any]
|
||||
) -> dict[str, Any]:
|
||||
return self._request(
|
||||
"GET", f"/v2/deployments/{deployment_id}/metrics", params=params
|
||||
)
|
||||
|
||||
def scale_deployment(self, deployment_id: str, payload: dict[str, str]) -> None:
|
||||
self._request(
|
||||
"PATCH", f"/v2/deployments/{deployment_id}/resource-tiers", payload
|
||||
)
|
||||
|
||||
def get_build_logs(
|
||||
self, project_id: str, revision_id: str, payload: dict[str, Any]
|
||||
) -> Any:
|
||||
|
||||
@@ -0,0 +1,130 @@
|
||||
import json
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
from click.testing import CliRunner
|
||||
|
||||
import langgraph_cli.deploy as deploy_module
|
||||
from langgraph_cli.cli import cli
|
||||
from langgraph_cli.host_backend import HostBackendClient
|
||||
|
||||
DEPLOYMENT_ID = "019a0000-0000-7000-8000-000000000001"
|
||||
|
||||
|
||||
def invoke(monkeypatch, args, handler, input=None):
|
||||
client = HostBackendClient(
|
||||
"https://api.example.com",
|
||||
"key",
|
||||
tenant_id="workspace-id",
|
||||
transport=httpx.MockTransport(handler),
|
||||
)
|
||||
monkeypatch.setattr(deploy_module, "HostBackendClient", lambda *a, **kw: client)
|
||||
return CliRunner().invoke(
|
||||
cli,
|
||||
["deploy", *args, "--api-key", "key"],
|
||||
input=input,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("filtered", [False, True])
|
||||
def test_metrics(monkeypatch, filtered):
|
||||
charts = {"id": "metrics", "charts": [], "sub_sections": []}
|
||||
params = (
|
||||
{
|
||||
"start_time": "2026-10-28T00:00:00Z",
|
||||
"end_time": "2026-10-28T01:00:00Z",
|
||||
"metrics": "cpu_usage,memory_usage",
|
||||
}
|
||||
if filtered
|
||||
else {}
|
||||
)
|
||||
|
||||
def handler(request):
|
||||
assert request.method == "GET"
|
||||
assert request.url.path == f"/v2/deployments/{DEPLOYMENT_ID}/metrics"
|
||||
assert dict(request.url.params) == params
|
||||
assert request.headers["x-tenant-id"] == "workspace-id"
|
||||
assert request.headers["x-api-key"] == "key"
|
||||
return httpx.Response(200, json=charts)
|
||||
|
||||
args = ["metrics", DEPLOYMENT_ID]
|
||||
for key, value in params.items():
|
||||
args.extend([f"--{key.replace('_', '-')}", value])
|
||||
result = invoke(monkeypatch, args, handler)
|
||||
assert result.exit_code == 0, result.output
|
||||
assert json.loads(result.output) == charts
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"tiers",
|
||||
[
|
||||
{"compute_tier": "SERVERLESS_M"},
|
||||
{"database_tier": "DEDICATED_L"},
|
||||
{"compute_tier": "DEDICATED_M", "database_tier": "DEDICATED_L"},
|
||||
],
|
||||
)
|
||||
def test_scale(monkeypatch, tiers):
|
||||
def handler(request):
|
||||
assert request.method == "PATCH"
|
||||
assert request.url.path == f"/v2/deployments/{DEPLOYMENT_ID}/resource-tiers"
|
||||
assert request.headers["x-tenant-id"] == "workspace-id"
|
||||
assert json.loads(request.content) == tiers
|
||||
return httpx.Response(200, json={"secrets": [{"value": "secret-value"}]})
|
||||
|
||||
args = ["scale", DEPLOYMENT_ID, "--yes"]
|
||||
for key, value in tiers.items():
|
||||
args.extend([f"--{key.replace('_', '-')}", value])
|
||||
result = invoke(monkeypatch, args, handler)
|
||||
assert result.exit_code == 0, result.output
|
||||
assert "Scaling requested" in result.output
|
||||
assert "secret-value" not in result.output
|
||||
|
||||
|
||||
@pytest.mark.parametrize("answer,expected_code", [("y\n", 0), ("n\n", 1)])
|
||||
def test_scale_confirmation(monkeypatch, answer, expected_code):
|
||||
requests = []
|
||||
|
||||
def handler(request):
|
||||
requests.append(request)
|
||||
return httpx.Response(200, json={})
|
||||
|
||||
result = invoke(
|
||||
monkeypatch,
|
||||
["scale", DEPLOYMENT_ID, "--compute-tier", "DEDICATED_S"],
|
||||
handler,
|
||||
input=answer,
|
||||
)
|
||||
assert result.exit_code == expected_code, result.output
|
||||
assert len(requests) == (1 if expected_code == 0 else 0)
|
||||
assert "may change billing or interrupt service" in result.output
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"args",
|
||||
[
|
||||
["scale", DEPLOYMENT_ID, "--yes"],
|
||||
["scale", DEPLOYMENT_ID, "--database-tier", "SERVERLESS_M", "--yes"],
|
||||
["metrics", DEPLOYMENT_ID, "--step-seconds", "60"],
|
||||
["metrics", "invalid-id"],
|
||||
],
|
||||
)
|
||||
def test_invalid_options_do_not_send_requests(monkeypatch, args):
|
||||
def handler(request):
|
||||
pytest.fail("Invalid options must not send a request")
|
||||
|
||||
result = invoke(monkeypatch, args, handler)
|
||||
assert result.exit_code == 2, result.output
|
||||
|
||||
|
||||
@pytest.mark.parametrize("command", ["metrics", "scale"])
|
||||
@pytest.mark.parametrize("status", [403, 422, 429])
|
||||
def test_backend_errors(monkeypatch, command, status):
|
||||
def handler(request):
|
||||
return httpx.Response(status, json={"detail": "Operation rejected"})
|
||||
|
||||
args = [command, DEPLOYMENT_ID]
|
||||
if command == "scale":
|
||||
args.extend(["--compute-tier", "DEDICATED_M", "--yes"])
|
||||
result = invoke(monkeypatch, args, handler)
|
||||
assert result.exit_code == 1
|
||||
assert "Operation rejected" in result.output
|
||||
@@ -255,11 +255,6 @@ R = TypeVar("R")
|
||||
S = TypeVar("S")
|
||||
|
||||
|
||||
class _EntrypointPregel(Pregel):
|
||||
"""The graph that `@entrypoint` builds. `Pregel` itself cannot be constructed
|
||||
directly."""
|
||||
|
||||
|
||||
# The decorator was wrapped in a class to support the `final` attribute.
|
||||
# In this form, the `final` attribute should play nicely with IDE autocompletion,
|
||||
# and type checking tools.
|
||||
@@ -578,7 +573,7 @@ class entrypoint(Generic[ContextT]):
|
||||
else:
|
||||
output_type = save_type = sig.return_annotation
|
||||
|
||||
graph: Pregel[Any, ContextT, Any, Any] = _EntrypointPregel(
|
||||
graph: Pregel[Any, ContextT, Any, Any] = Pregel(
|
||||
nodes={
|
||||
func.__name__: PregelNode(
|
||||
bound=bound,
|
||||
|
||||
@@ -1,3 +1,3 @@
|
||||
from langgraph.pregel.main import Pregel
|
||||
from langgraph.pregel.main import NodeBuilder, Pregel
|
||||
|
||||
__all__ = ("Pregel",)
|
||||
__all__ = ("Pregel", "NodeBuilder")
|
||||
@@ -216,7 +216,7 @@ class StreamMessagesHandler(BaseCallbackHandler, _StreamingCallbackHandler):
|
||||
ns = tuple(cast(str, metadata["langgraph_checkpoint_ns"]).split(NS_SEP))[
|
||||
:-1
|
||||
]
|
||||
if not self.subgraphs and len(ns) > 0:
|
||||
if not self.subgraphs and len(ns) > 0 and ns != self.parent_ns:
|
||||
return
|
||||
self.metadata[run_id] = (ns, metadata)
|
||||
for value in _state_values(inputs):
|
||||
|
||||
@@ -47,7 +47,9 @@ def validate_graph(
|
||||
f"not in known channels: '{repr(sorted(channels))[:100]}'"
|
||||
)
|
||||
else:
|
||||
raise TypeError(f"Invalid node type {type(node)}, expected PregelNode")
|
||||
raise TypeError(
|
||||
f"Invalid node type {type(node)}, expected PregelNode or NodeBuilder"
|
||||
)
|
||||
|
||||
for chan in subscribed_channels:
|
||||
if chan not in channels:
|
||||
|
||||
@@ -17,6 +17,7 @@ from collections.abc import (
|
||||
Sequence,
|
||||
)
|
||||
from dataclasses import is_dataclass, replace
|
||||
from datetime import timedelta
|
||||
from functools import partial
|
||||
from inspect import isclass
|
||||
from typing import (
|
||||
@@ -34,6 +35,7 @@ from langchain_core.globals import get_debug
|
||||
from langchain_core.runnables import (
|
||||
RunnableSequence,
|
||||
)
|
||||
from langchain_core.runnables.base import Input, Output
|
||||
from langchain_core.runnables.config import (
|
||||
RunnableConfig,
|
||||
get_async_callback_manager_for_config,
|
||||
@@ -91,6 +93,13 @@ from langgraph._internal._queue import ( # type: ignore[attr-defined]
|
||||
AsyncQueue,
|
||||
SyncQueue,
|
||||
)
|
||||
from langgraph._internal._runnable import (
|
||||
Runnable,
|
||||
RunnableLike,
|
||||
RunnableSeq,
|
||||
coerce_to_runnable,
|
||||
)
|
||||
from langgraph._internal._timeout import coerce_timeout_policy
|
||||
from langgraph._internal._typing import MISSING, DeprecatedKwargs
|
||||
from langgraph.callbacks import (
|
||||
GraphInterruptEvent,
|
||||
@@ -146,7 +155,7 @@ from langgraph.pregel._messages import (
|
||||
StreamMessagesHandlerV2,
|
||||
ensure_message_ids,
|
||||
)
|
||||
from langgraph.pregel._read import PregelNode
|
||||
from langgraph.pregel._read import DEFAULT_BOUND, PregelNode
|
||||
from langgraph.pregel._retry import RetryPolicy
|
||||
from langgraph.pregel._runner import PregelRunner
|
||||
from langgraph.pregel._task_status import read_task_statuses
|
||||
@@ -156,6 +165,7 @@ from langgraph.pregel._utils import (
|
||||
validate_timeout_supported,
|
||||
)
|
||||
from langgraph.pregel._validate import validate_graph, validate_keys
|
||||
from langgraph.pregel._write import ChannelWrite, ChannelWriteEntry
|
||||
from langgraph.pregel.debug import get_bolded_text, get_colored_text, tasks_w_writes
|
||||
from langgraph.pregel.protocol import PregelProtocol, StreamChunk, StreamProtocol
|
||||
from langgraph.runtime import (
|
||||
@@ -187,6 +197,7 @@ from langgraph.types import (
|
||||
StateUpdate,
|
||||
StreamMode,
|
||||
StreamPart,
|
||||
TimeoutPolicy,
|
||||
ensure_valid_checkpointer,
|
||||
)
|
||||
from langgraph.typing import ContextT, InputT, OutputT, StateT
|
||||
@@ -197,7 +208,182 @@ try:
|
||||
except ImportError:
|
||||
_StreamingCallbackHandler = None # type: ignore
|
||||
|
||||
__all__ = ("Pregel",)
|
||||
__all__ = ("NodeBuilder", "Pregel")
|
||||
|
||||
_WriteValue = Callable[[Input], Output] | Any
|
||||
|
||||
|
||||
class NodeBuilder:
|
||||
__slots__ = (
|
||||
"_channels",
|
||||
"_triggers",
|
||||
"_tags",
|
||||
"_metadata",
|
||||
"_writes",
|
||||
"_bound",
|
||||
"_retry_policy",
|
||||
"_cache_policy",
|
||||
"_timeout",
|
||||
)
|
||||
|
||||
_channels: str | list[str]
|
||||
_triggers: list[str]
|
||||
_tags: list[str]
|
||||
_metadata: dict[str, Any]
|
||||
_writes: list[ChannelWriteEntry]
|
||||
_bound: Runnable
|
||||
_retry_policy: list[RetryPolicy]
|
||||
_cache_policy: CachePolicy | None
|
||||
_timeout: TimeoutPolicy | None
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
) -> None:
|
||||
self._channels = []
|
||||
self._triggers = []
|
||||
self._tags = []
|
||||
self._metadata = {}
|
||||
self._writes = []
|
||||
self._bound = DEFAULT_BOUND
|
||||
self._retry_policy = []
|
||||
self._cache_policy = None
|
||||
self._timeout = None
|
||||
|
||||
def subscribe_only(
|
||||
self,
|
||||
channel: str,
|
||||
) -> Self:
|
||||
"""Subscribe to a single channel."""
|
||||
if not self._channels:
|
||||
self._channels = channel
|
||||
else:
|
||||
raise ValueError(
|
||||
"Cannot subscribe to single channels when other channels are already subscribed to"
|
||||
)
|
||||
|
||||
self._triggers.append(channel)
|
||||
|
||||
return self
|
||||
|
||||
def subscribe_to(
|
||||
self,
|
||||
*channels: str,
|
||||
read: bool = True,
|
||||
) -> Self:
|
||||
"""Add channels to subscribe to.
|
||||
|
||||
Node will be invoked when any of these channels are updated, with a dict of the
|
||||
channel values as input.
|
||||
|
||||
Args:
|
||||
channels: Channel name(s) to subscribe to
|
||||
read: If `True`, the channels will be included in the input to the node.
|
||||
Otherwise, they will trigger the node without being sent in input.
|
||||
|
||||
Returns:
|
||||
Self for chaining
|
||||
"""
|
||||
if isinstance(self._channels, str):
|
||||
raise ValueError(
|
||||
"Cannot subscribe to channels when subscribed to a single channel"
|
||||
)
|
||||
if read:
|
||||
if not self._channels:
|
||||
self._channels = list(channels)
|
||||
else:
|
||||
self._channels.extend(channels)
|
||||
|
||||
if isinstance(channels, str):
|
||||
self._triggers.append(channels)
|
||||
else:
|
||||
self._triggers.extend(channels)
|
||||
|
||||
return self
|
||||
|
||||
def read_from(
|
||||
self,
|
||||
*channels: str,
|
||||
) -> Self:
|
||||
"""Adds the specified channels to read from, without subscribing to them."""
|
||||
assert isinstance(self._channels, list), (
|
||||
"Cannot read additional channels when subscribed to single channels"
|
||||
)
|
||||
self._channels.extend(channels)
|
||||
return self
|
||||
|
||||
def do(
|
||||
self,
|
||||
node: RunnableLike,
|
||||
) -> Self:
|
||||
"""Adds the specified node."""
|
||||
if self._bound is not DEFAULT_BOUND:
|
||||
self._bound = RunnableSeq(
|
||||
self._bound, coerce_to_runnable(node, name=None, trace=True)
|
||||
)
|
||||
else:
|
||||
self._bound = coerce_to_runnable(node, name=None, trace=True)
|
||||
return self
|
||||
|
||||
def write_to(
|
||||
self,
|
||||
*channels: str | ChannelWriteEntry,
|
||||
**kwargs: _WriteValue,
|
||||
) -> Self:
|
||||
"""Add channel writes.
|
||||
|
||||
Args:
|
||||
*channels: Channel names to write to.
|
||||
**kwargs: Channel name and value mappings.
|
||||
|
||||
Returns:
|
||||
Self for chaining
|
||||
"""
|
||||
self._writes.extend(
|
||||
ChannelWriteEntry(c) if isinstance(c, str) else c for c in channels
|
||||
)
|
||||
self._writes.extend(
|
||||
ChannelWriteEntry(k, mapper=v)
|
||||
if callable(v)
|
||||
else ChannelWriteEntry(k, value=v)
|
||||
for k, v in kwargs.items()
|
||||
)
|
||||
|
||||
return self
|
||||
|
||||
def meta(self, *tags: str, **metadata: Any) -> Self:
|
||||
"""Add tags or metadata to the node."""
|
||||
self._tags.extend(tags)
|
||||
self._metadata.update(metadata)
|
||||
return self
|
||||
|
||||
def add_retry_policies(self, *policies: RetryPolicy) -> Self:
|
||||
"""Adds retry policies to the node."""
|
||||
self._retry_policy.extend(policies)
|
||||
return self
|
||||
|
||||
def add_cache_policy(self, policy: CachePolicy) -> Self:
|
||||
"""Adds cache policies to the node."""
|
||||
self._cache_policy = policy
|
||||
return self
|
||||
|
||||
def set_timeout(self, timeout: float | timedelta | TimeoutPolicy | None) -> Self:
|
||||
"""Set the per-attempt timeout policy for this node."""
|
||||
self._timeout = coerce_timeout_policy(timeout)
|
||||
return self
|
||||
|
||||
def build(self) -> PregelNode:
|
||||
"""Builds the node."""
|
||||
return PregelNode(
|
||||
channels=self._channels,
|
||||
triggers=self._triggers,
|
||||
tags=self._tags,
|
||||
metadata=self._metadata,
|
||||
writers=[ChannelWrite(self._writes)],
|
||||
bound=self._bound,
|
||||
retry_policy=self._retry_policy,
|
||||
cache_policy=self._cache_policy,
|
||||
timeout=self._timeout,
|
||||
)
|
||||
|
||||
|
||||
# Kwargs that ``stream_events(version="v3")`` / ``astream_events(version="v3")``
|
||||
@@ -275,22 +461,255 @@ class Pregel(
|
||||
PregelProtocol[StateT, ContextT, InputT, OutputT],
|
||||
Generic[StateT, ContextT, InputT, OutputT],
|
||||
):
|
||||
"""The compiled graph that runs a LangGraph application.
|
||||
"""Pregel manages the runtime behavior for LangGraph applications.
|
||||
|
||||
`StateGraph.compile()` and the `@entrypoint` decorator return an instance of
|
||||
this class. It provides the methods that run and inspect the application,
|
||||
such as `invoke`, `stream`, `get_state` and `update_state`.
|
||||
## Overview
|
||||
|
||||
The runtime is named after Google's
|
||||
[Pregel algorithm](https://research.google/pubs/pub37252/). It runs the
|
||||
application in steps. In each step it plans which nodes to run, runs them in
|
||||
parallel, and then applies their writes to the graph's state, so writes
|
||||
become visible only in the next step. The run ends when no node is
|
||||
scheduled, or when the recursion limit is reached.
|
||||
Pregel combines [**actors**](https://en.wikipedia.org/wiki/Actor_model)
|
||||
and **channels** into a single application.
|
||||
**Actors** read data from channels and write data to channels.
|
||||
Pregel organizes the execution of the application into multiple steps,
|
||||
following the **Pregel Algorithm**/**Bulk Synchronous Parallel** model.
|
||||
|
||||
`Pregel` is not constructed directly. Build graphs with
|
||||
[`StateGraph`][langgraph.graph.StateGraph] or
|
||||
[`@entrypoint`][langgraph.func.entrypoint].
|
||||
Each step consists of three phases:
|
||||
|
||||
- **Plan**: Determine which **actors** to execute in this step. For example,
|
||||
in the first step, select the **actors** that subscribe to the special
|
||||
**input** channels; in subsequent steps,
|
||||
select the **actors** that subscribe to channels updated in the previous step.
|
||||
- **Execution**: Execute all selected **actors** in parallel,
|
||||
until all complete, or one fails, or a timeout is reached. During this
|
||||
phase, channel updates are invisible to actors until the next step.
|
||||
- **Update**: Update the channels with the values written by the **actors**
|
||||
in this step.
|
||||
|
||||
Repeat until no **actors** are selected for execution, or a maximum number of
|
||||
steps is reached.
|
||||
|
||||
## Actors
|
||||
|
||||
An **actor** is a `PregelNode`.
|
||||
It subscribes to channels, reads data from them, and writes data to them.
|
||||
It can be thought of as an **actor** in the Pregel algorithm.
|
||||
`PregelNodes` implement LangChain's
|
||||
Runnable interface.
|
||||
|
||||
## Channels
|
||||
|
||||
Channels are used to communicate between actors (`PregelNodes`).
|
||||
Each channel has a value type, an update type, and an update function – which
|
||||
takes a sequence of updates and
|
||||
modifies the stored value. Channels can be used to send data from one chain to
|
||||
another, or to send data from a chain to itself in a future step. LangGraph
|
||||
provides a number of built-in channels:
|
||||
|
||||
### Basic channels: LastValue and Topic
|
||||
|
||||
- `LastValue`: The default channel, stores the last value sent to the channel,
|
||||
useful for input and output values, or for sending data from one step to the next
|
||||
- `Topic`: A configurable PubSub Topic, useful for sending multiple values
|
||||
between *actors*, or for accumulating output. Can be configured to deduplicate
|
||||
values, and/or to accumulate values over the course of multiple steps.
|
||||
|
||||
### Advanced channels: Context and BinaryOperatorAggregate
|
||||
|
||||
- `Context`: exposes the value of a context manager, managing its lifecycle.
|
||||
Useful for accessing external resources that require setup and/or teardown. e.g.
|
||||
`client = Context(httpx.Client)`
|
||||
- `BinaryOperatorAggregate`: stores a persistent value, updated by applying
|
||||
a binary operator to the current value and each update
|
||||
sent to the channel, useful for computing aggregates over multiple steps. e.g.
|
||||
`total = BinaryOperatorAggregate(int, operator.add)`
|
||||
|
||||
## Examples
|
||||
|
||||
Most users will interact with Pregel via a
|
||||
[StateGraph (Graph API)][langgraph.graph.StateGraph] or via an
|
||||
[entrypoint (Functional API)][langgraph.func.entrypoint].
|
||||
|
||||
However, for **advanced** use cases, Pregel can be used directly. If you're
|
||||
not sure whether you need to use Pregel directly, then the answer is probably no
|
||||
- you should use the Graph API or Functional API instead. These are higher-level
|
||||
interfaces that will compile down to Pregel under the hood.
|
||||
|
||||
Here are some examples to give you a sense of how it works:
|
||||
|
||||
Example: Single node application
|
||||
```python
|
||||
from langgraph.channels import EphemeralValue
|
||||
from langgraph.pregel import Pregel, NodeBuilder
|
||||
|
||||
node1 = (
|
||||
NodeBuilder().subscribe_only("a")
|
||||
.do(lambda x: x + x)
|
||||
.write_to("b")
|
||||
)
|
||||
|
||||
app = Pregel(
|
||||
nodes={"node1": node1},
|
||||
channels={
|
||||
"a": EphemeralValue(str),
|
||||
"b": EphemeralValue(str),
|
||||
},
|
||||
input_channels=["a"],
|
||||
output_channels=["b"],
|
||||
)
|
||||
|
||||
app.invoke({"a": "foo"})
|
||||
```
|
||||
|
||||
```con
|
||||
{'b': 'foofoo'}
|
||||
```
|
||||
|
||||
Example: Using multiple nodes and multiple output channels
|
||||
```python
|
||||
from langgraph.channels import LastValue, EphemeralValue
|
||||
from langgraph.pregel import Pregel, NodeBuilder
|
||||
|
||||
node1 = (
|
||||
NodeBuilder().subscribe_only("a")
|
||||
.do(lambda x: x + x)
|
||||
.write_to("b")
|
||||
)
|
||||
|
||||
node2 = (
|
||||
NodeBuilder().subscribe_to("b")
|
||||
.do(lambda x: x["b"] + x["b"])
|
||||
.write_to("c")
|
||||
)
|
||||
|
||||
|
||||
app = Pregel(
|
||||
nodes={"node1": node1, "node2": node2},
|
||||
channels={
|
||||
"a": EphemeralValue(str),
|
||||
"b": LastValue(str),
|
||||
"c": EphemeralValue(str),
|
||||
},
|
||||
input_channels=["a"],
|
||||
output_channels=["b", "c"],
|
||||
)
|
||||
|
||||
app.invoke({"a": "foo"})
|
||||
```
|
||||
|
||||
```con
|
||||
{'b': 'foofoo', 'c': 'foofoofoofoo'}
|
||||
```
|
||||
|
||||
Example: Using a Topic channel
|
||||
```python
|
||||
from langgraph.channels import LastValue, EphemeralValue, Topic
|
||||
from langgraph.pregel import Pregel, NodeBuilder
|
||||
|
||||
node1 = (
|
||||
NodeBuilder().subscribe_only("a")
|
||||
.do(lambda x: x + x)
|
||||
.write_to("b", "c")
|
||||
)
|
||||
|
||||
node2 = (
|
||||
NodeBuilder().subscribe_only("b")
|
||||
.do(lambda x: x + x)
|
||||
.write_to("c")
|
||||
)
|
||||
|
||||
|
||||
app = Pregel(
|
||||
nodes={"node1": node1, "node2": node2},
|
||||
channels={
|
||||
"a": EphemeralValue(str),
|
||||
"b": EphemeralValue(str),
|
||||
"c": Topic(str, accumulate=True),
|
||||
},
|
||||
input_channels=["a"],
|
||||
output_channels=["c"],
|
||||
)
|
||||
|
||||
app.invoke({"a": "foo"})
|
||||
```
|
||||
|
||||
```pycon
|
||||
{"c": ["foofoo", "foofoofoofoo"]}
|
||||
```
|
||||
|
||||
Example: Using a `BinaryOperatorAggregate` channel
|
||||
```python
|
||||
from langgraph.channels import EphemeralValue, BinaryOperatorAggregate
|
||||
from langgraph.pregel import Pregel, NodeBuilder
|
||||
|
||||
|
||||
node1 = (
|
||||
NodeBuilder().subscribe_only("a")
|
||||
.do(lambda x: x + x)
|
||||
.write_to("b", "c")
|
||||
)
|
||||
|
||||
node2 = (
|
||||
NodeBuilder().subscribe_only("b")
|
||||
.do(lambda x: x + x)
|
||||
.write_to("c")
|
||||
)
|
||||
|
||||
|
||||
def reducer(current, update):
|
||||
if current:
|
||||
return current + " | " + update
|
||||
else:
|
||||
return update
|
||||
|
||||
|
||||
app = Pregel(
|
||||
nodes={"node1": node1, "node2": node2},
|
||||
channels={
|
||||
"a": EphemeralValue(str),
|
||||
"b": EphemeralValue(str),
|
||||
"c": BinaryOperatorAggregate(str, operator=reducer),
|
||||
},
|
||||
input_channels=["a"],
|
||||
output_channels=["c"],
|
||||
)
|
||||
|
||||
app.invoke({"a": "foo"})
|
||||
```
|
||||
|
||||
```con
|
||||
{'c': 'foofoo | foofoofoofoo'}
|
||||
```
|
||||
|
||||
Example: Introducing a cycle
|
||||
This example demonstrates how to introduce a cycle in the graph, by having
|
||||
a chain write to a channel it subscribes to.
|
||||
|
||||
Execution will continue until a `None` value is written to the channel.
|
||||
|
||||
```python
|
||||
from langgraph.channels import EphemeralValue
|
||||
from langgraph.pregel import Pregel, NodeBuilder, ChannelWriteEntry
|
||||
|
||||
example_node = (
|
||||
NodeBuilder()
|
||||
.subscribe_only("value")
|
||||
.do(lambda x: x + x if len(x) < 10 else None)
|
||||
.write_to(ChannelWriteEntry(channel="value", skip_none=True))
|
||||
)
|
||||
|
||||
app = Pregel(
|
||||
nodes={"example_node": example_node},
|
||||
channels={
|
||||
"value": EphemeralValue(str),
|
||||
},
|
||||
input_channels=["value"],
|
||||
output_channels=["value"],
|
||||
)
|
||||
|
||||
app.invoke({"value": "a"})
|
||||
```
|
||||
|
||||
```con
|
||||
{'value': 'aaaaaaaaaaaaaaaa'}
|
||||
```
|
||||
"""
|
||||
|
||||
nodes: dict[str, PregelNode]
|
||||
@@ -349,7 +768,7 @@ class Pregel(
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
nodes: dict[str, PregelNode],
|
||||
nodes: dict[str, PregelNode | NodeBuilder],
|
||||
channels: dict[str, BaseChannel | ManagedValueSpec] | None,
|
||||
auto_validate: bool = True,
|
||||
stream_mode: StreamMode = "values",
|
||||
@@ -374,13 +793,6 @@ class Pregel(
|
||||
stream_transformers: Sequence[Callable[[tuple[str, ...]], Any]] | None = None,
|
||||
**deprecated_kwargs: Unpack[DeprecatedKwargs],
|
||||
) -> None:
|
||||
# LangGraph builds graphs through subclasses: `CompiledStateGraph` and
|
||||
# the graph `@entrypoint` returns.
|
||||
if type(self) is Pregel:
|
||||
raise TypeError(
|
||||
"`Pregel` cannot be constructed directly. "
|
||||
"Build graphs with `StateGraph` or `@entrypoint`."
|
||||
)
|
||||
if (
|
||||
config_type := deprecated_kwargs.get("config_type", MISSING)
|
||||
) is not MISSING:
|
||||
@@ -395,7 +807,9 @@ class Pregel(
|
||||
|
||||
checkpointer = ensure_valid_checkpointer(checkpointer)
|
||||
|
||||
self.nodes = dict(nodes)
|
||||
self.nodes = {
|
||||
k: v.build() if isinstance(v, NodeBuilder) else v for k, v in nodes.items()
|
||||
}
|
||||
self.channels = channels or {}
|
||||
if TASKS in self.channels and not isinstance(self.channels[TASKS], Topic):
|
||||
raise ValueError(
|
||||
|
||||
@@ -847,6 +847,11 @@ class RemoteGraph(PregelProtocol):
|
||||
# filter for what was actually requested
|
||||
if mode not in requested:
|
||||
continue
|
||||
# subgraph events may have been requested only for the parent stream.
|
||||
# messages are exempt: the server dedupes them across namespaces, so
|
||||
# dropping them here would also drop the root-level message
|
||||
if not subgraphs and NS_SEP in chunk.event and mode != "messages":
|
||||
continue
|
||||
|
||||
if chunk.event.startswith("messages"):
|
||||
chunk = chunk._replace(data=tuple(chunk.data))
|
||||
@@ -997,6 +1002,11 @@ class RemoteGraph(PregelProtocol):
|
||||
# filter for what was actually requested
|
||||
if mode not in requested:
|
||||
continue
|
||||
# subgraph events may have been requested only for the parent stream.
|
||||
# messages are exempt: the server dedupes them across namespaces, so
|
||||
# dropping them here would also drop the root-level message
|
||||
if not subgraphs and NS_SEP in chunk.event and mode != "messages":
|
||||
continue
|
||||
|
||||
if chunk.event.startswith("messages"):
|
||||
chunk = chunk._replace(data=tuple(chunk.data))
|
||||
|
||||
@@ -1,222 +0,0 @@
|
||||
"""Build graphs directly on the Pregel engine, for tests of engine behavior.
|
||||
|
||||
LangGraph does not expose this: applications build graphs with `StateGraph` or
|
||||
`@entrypoint`. Tests use it to set up channels, triggers and writes that those
|
||||
APIs do not produce on their own.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Callable
|
||||
from datetime import timedelta
|
||||
from typing import Any
|
||||
|
||||
from langchain_core.runnables.base import Input, Output
|
||||
from typing_extensions import Self
|
||||
|
||||
from langgraph._internal._runnable import (
|
||||
Runnable,
|
||||
RunnableLike,
|
||||
RunnableSeq,
|
||||
coerce_to_runnable,
|
||||
)
|
||||
from langgraph._internal._timeout import coerce_timeout_policy
|
||||
from langgraph.pregel import Pregel
|
||||
from langgraph.pregel._read import DEFAULT_BOUND, PregelNode
|
||||
from langgraph.pregel._write import ChannelWrite, ChannelWriteEntry
|
||||
from langgraph.types import CachePolicy, RetryPolicy, TimeoutPolicy
|
||||
|
||||
__all__ = ("NodeBuilder", "RawPregel")
|
||||
|
||||
_WriteValue = Callable[[Input], Output] | Any
|
||||
|
||||
|
||||
class NodeBuilder:
|
||||
"""Builds a `PregelNode` that reads and writes named channels."""
|
||||
|
||||
__slots__ = (
|
||||
"_channels",
|
||||
"_triggers",
|
||||
"_tags",
|
||||
"_metadata",
|
||||
"_writes",
|
||||
"_bound",
|
||||
"_retry_policy",
|
||||
"_cache_policy",
|
||||
"_timeout",
|
||||
)
|
||||
|
||||
_channels: str | list[str]
|
||||
_triggers: list[str]
|
||||
_tags: list[str]
|
||||
_metadata: dict[str, Any]
|
||||
_writes: list[ChannelWriteEntry]
|
||||
_bound: Runnable
|
||||
_retry_policy: list[RetryPolicy]
|
||||
_cache_policy: CachePolicy | None
|
||||
_timeout: TimeoutPolicy | None
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
) -> None:
|
||||
self._channels = []
|
||||
self._triggers = []
|
||||
self._tags = []
|
||||
self._metadata = {}
|
||||
self._writes = []
|
||||
self._bound = DEFAULT_BOUND
|
||||
self._retry_policy = []
|
||||
self._cache_policy = None
|
||||
self._timeout = None
|
||||
|
||||
def subscribe_only(
|
||||
self,
|
||||
channel: str,
|
||||
) -> Self:
|
||||
"""Subscribe to a single channel."""
|
||||
if not self._channels:
|
||||
self._channels = channel
|
||||
else:
|
||||
raise ValueError(
|
||||
"Cannot subscribe to single channels when other channels are already subscribed to"
|
||||
)
|
||||
|
||||
self._triggers.append(channel)
|
||||
|
||||
return self
|
||||
|
||||
def subscribe_to(
|
||||
self,
|
||||
*channels: str,
|
||||
read: bool = True,
|
||||
) -> Self:
|
||||
"""Add channels to subscribe to.
|
||||
|
||||
Node will be invoked when any of these channels are updated, with a dict of the
|
||||
channel values as input.
|
||||
|
||||
Args:
|
||||
channels: Channel name(s) to subscribe to
|
||||
read: If `True`, the channels will be included in the input to the node.
|
||||
Otherwise, they will trigger the node without being sent in input.
|
||||
|
||||
Returns:
|
||||
Self for chaining
|
||||
"""
|
||||
if isinstance(self._channels, str):
|
||||
raise ValueError(
|
||||
"Cannot subscribe to channels when subscribed to a single channel"
|
||||
)
|
||||
if read:
|
||||
if not self._channels:
|
||||
self._channels = list(channels)
|
||||
else:
|
||||
self._channels.extend(channels)
|
||||
|
||||
if isinstance(channels, str):
|
||||
self._triggers.append(channels)
|
||||
else:
|
||||
self._triggers.extend(channels)
|
||||
|
||||
return self
|
||||
|
||||
def read_from(
|
||||
self,
|
||||
*channels: str,
|
||||
) -> Self:
|
||||
"""Adds the specified channels to read from, without subscribing to them."""
|
||||
assert isinstance(self._channels, list), (
|
||||
"Cannot read additional channels when subscribed to single channels"
|
||||
)
|
||||
self._channels.extend(channels)
|
||||
return self
|
||||
|
||||
def do(
|
||||
self,
|
||||
node: RunnableLike,
|
||||
) -> Self:
|
||||
"""Adds the specified node."""
|
||||
if self._bound is not DEFAULT_BOUND:
|
||||
self._bound = RunnableSeq(
|
||||
self._bound, coerce_to_runnable(node, name=None, trace=True)
|
||||
)
|
||||
else:
|
||||
self._bound = coerce_to_runnable(node, name=None, trace=True)
|
||||
return self
|
||||
|
||||
def write_to(
|
||||
self,
|
||||
*channels: str | ChannelWriteEntry,
|
||||
**kwargs: _WriteValue,
|
||||
) -> Self:
|
||||
"""Add channel writes.
|
||||
|
||||
Args:
|
||||
*channels: Channel names to write to.
|
||||
**kwargs: Channel name and value mappings.
|
||||
|
||||
Returns:
|
||||
Self for chaining
|
||||
"""
|
||||
self._writes.extend(
|
||||
ChannelWriteEntry(c) if isinstance(c, str) else c for c in channels
|
||||
)
|
||||
self._writes.extend(
|
||||
ChannelWriteEntry(k, mapper=v)
|
||||
if callable(v)
|
||||
else ChannelWriteEntry(k, value=v)
|
||||
for k, v in kwargs.items()
|
||||
)
|
||||
|
||||
return self
|
||||
|
||||
def meta(self, *tags: str, **metadata: Any) -> Self:
|
||||
"""Add tags or metadata to the node."""
|
||||
self._tags.extend(tags)
|
||||
self._metadata.update(metadata)
|
||||
return self
|
||||
|
||||
def add_retry_policies(self, *policies: RetryPolicy) -> Self:
|
||||
"""Adds retry policies to the node."""
|
||||
self._retry_policy.extend(policies)
|
||||
return self
|
||||
|
||||
def add_cache_policy(self, policy: CachePolicy) -> Self:
|
||||
"""Adds cache policies to the node."""
|
||||
self._cache_policy = policy
|
||||
return self
|
||||
|
||||
def set_timeout(self, timeout: float | timedelta | TimeoutPolicy | None) -> Self:
|
||||
"""Set the per-attempt timeout policy for this node."""
|
||||
self._timeout = coerce_timeout_policy(timeout)
|
||||
return self
|
||||
|
||||
def build(self) -> PregelNode:
|
||||
"""Builds the node."""
|
||||
return PregelNode(
|
||||
channels=self._channels,
|
||||
triggers=self._triggers,
|
||||
tags=self._tags,
|
||||
metadata=self._metadata,
|
||||
writers=[ChannelWrite(self._writes)],
|
||||
bound=self._bound,
|
||||
retry_policy=self._retry_policy,
|
||||
cache_policy=self._cache_policy,
|
||||
timeout=self._timeout,
|
||||
)
|
||||
|
||||
|
||||
class RawPregel(Pregel):
|
||||
"""A `Pregel` built from hand-written nodes and channels.
|
||||
|
||||
Accepts `NodeBuilder` values in `nodes`, and builds them.
|
||||
"""
|
||||
|
||||
def __init__(self, *, nodes: dict[str, PregelNode | NodeBuilder], **kwargs: Any):
|
||||
super().__init__(
|
||||
nodes={
|
||||
k: v.build() if isinstance(v, NodeBuilder) else v
|
||||
for k, v in nodes.items()
|
||||
},
|
||||
**kwargs,
|
||||
)
|
||||
@@ -19,8 +19,8 @@ from langgraph.channels.delta import DeltaChannel
|
||||
from langgraph.channels.last_value import LastValue
|
||||
from langgraph.graph import END, START, StateGraph
|
||||
from langgraph.graph.message import _messages_delta_reducer
|
||||
from langgraph.pregel import NodeBuilder, Pregel
|
||||
from langgraph.types import Command, Durability, Overwrite
|
||||
from tests.raw_pregel import NodeBuilder, RawPregel
|
||||
|
||||
pytestmark = pytest.mark.anyio
|
||||
|
||||
@@ -279,7 +279,7 @@ def _pregel_graph(seen: list[list[str | None]]) -> Any:
|
||||
.do(look)
|
||||
.write_to("messages")
|
||||
)
|
||||
return RawPregel(
|
||||
return Pregel(
|
||||
nodes={"look": node},
|
||||
channels={
|
||||
"messages": DeltaChannel(_messages_delta_reducer),
|
||||
|
||||
@@ -7,9 +7,8 @@ from langgraph.checkpoint.memory import InMemorySaver
|
||||
from langgraph.channels.binop import BinaryOperatorAggregate
|
||||
from langgraph.channels.delta import DeltaChannel
|
||||
from langgraph.channels.last_value import LastValue
|
||||
from langgraph.pregel import Pregel
|
||||
from langgraph.pregel import NodeBuilder, Pregel
|
||||
from langgraph.types import Durability
|
||||
from tests.raw_pregel import NodeBuilder, RawPregel
|
||||
|
||||
pytestmark = pytest.mark.anyio
|
||||
|
||||
@@ -20,7 +19,7 @@ def _sorted_extend(current: list, writes: list) -> list:
|
||||
|
||||
def _delta_input_graph() -> Pregel:
|
||||
node = NodeBuilder().subscribe_only("go").do(lambda _: [2]).write_to("log", "plain")
|
||||
return RawPregel(
|
||||
return Pregel(
|
||||
nodes={"n": node},
|
||||
channels={
|
||||
"log": DeltaChannel(_sorted_extend),
|
||||
|
||||
@@ -3,8 +3,7 @@ from langgraph.checkpoint.memory import InMemorySaver
|
||||
|
||||
from langgraph.channels.delta import DeltaChannel
|
||||
from langgraph.channels.last_value import LastValue
|
||||
from langgraph.pregel import Pregel
|
||||
from tests.raw_pregel import NodeBuilder, RawPregel
|
||||
from langgraph.pregel import NodeBuilder, Pregel
|
||||
|
||||
pytestmark = pytest.mark.anyio
|
||||
|
||||
@@ -18,7 +17,7 @@ def _extend(current: list, writes: list) -> list:
|
||||
def _graph(reads: list) -> Pregel:
|
||||
writer = NodeBuilder().subscribe_only("a").do(lambda _: [1]).write_to("d")
|
||||
reader = NodeBuilder().subscribe_only("d").do(lambda d: reads.append(list(d)))
|
||||
return RawPregel(
|
||||
return Pregel(
|
||||
nodes={"writer": writer, "reader": reader},
|
||||
channels={"a": LastValue(str), "d": DeltaChannel(_extend)},
|
||||
input_channels="a",
|
||||
|
||||
@@ -13,9 +13,8 @@ from langgraph.channels.delta import DeltaChannel
|
||||
from langgraph.channels.last_value import LastValue
|
||||
from langgraph.graph import START, StateGraph
|
||||
from langgraph.graph.message import _messages_delta_reducer
|
||||
from langgraph.pregel import Pregel
|
||||
from langgraph.pregel import NodeBuilder, Pregel
|
||||
from langgraph.types import Overwrite
|
||||
from tests.raw_pregel import NodeBuilder, RawPregel
|
||||
|
||||
pytestmark = pytest.mark.anyio
|
||||
|
||||
@@ -37,7 +36,7 @@ def _extend(current: list, writes: list) -> list:
|
||||
|
||||
def _delta_input_graph(saver: InMemorySaver) -> Pregel:
|
||||
node = NodeBuilder().subscribe_only("go").do(lambda _: [2]).write_to("log")
|
||||
return RawPregel(
|
||||
return Pregel(
|
||||
nodes={"n": node},
|
||||
channels={"log": DeltaChannel(_extend), "go": LastValue(int)},
|
||||
input_channels=["log", "go"],
|
||||
|
||||
@@ -30,8 +30,8 @@ from langgraph.channels.delta import DeltaChannel
|
||||
from langgraph.channels.last_value import LastValue
|
||||
from langgraph.graph import START, StateGraph
|
||||
from langgraph.graph.message import _messages_delta_reducer
|
||||
from langgraph.pregel import NodeBuilder, Pregel
|
||||
from langgraph.types import Overwrite, StateSnapshot, StateUpdate
|
||||
from tests.raw_pregel import NodeBuilder, RawPregel
|
||||
|
||||
pytestmark = pytest.mark.anyio
|
||||
|
||||
@@ -549,7 +549,7 @@ def _sorted_extend(current: list, writes: list) -> list:
|
||||
|
||||
def _delta_input_graph(snapshot_frequency: int = 1000) -> Any:
|
||||
node = NodeBuilder().subscribe_only("go").do(lambda _: [2]).write_to("log", "plain")
|
||||
return RawPregel(
|
||||
return Pregel(
|
||||
nodes={"n": node},
|
||||
channels={
|
||||
"log": DeltaChannel(_sorted_extend, snapshot_frequency=snapshot_frequency),
|
||||
@@ -649,7 +649,7 @@ def _messages_input_graph() -> Any:
|
||||
.do(lambda _: [HumanMessage("n", id="n")])
|
||||
.write_to("messages")
|
||||
)
|
||||
return RawPregel(
|
||||
return Pregel(
|
||||
nodes={"n": node},
|
||||
channels={
|
||||
"messages": DeltaChannel(_messages_delta_reducer),
|
||||
|
||||
@@ -13,13 +13,13 @@ from langgraph.errors import NodeInterrupt
|
||||
from langgraph.func import entrypoint, task
|
||||
from langgraph.graph import StateGraph
|
||||
from langgraph.graph.message import MessageGraph
|
||||
from langgraph.pregel import NodeBuilder, Pregel
|
||||
from langgraph.types import GraphOutput, Interrupt, RetryPolicy
|
||||
from langgraph.warnings import (
|
||||
LangGraphDeprecatedSinceV05,
|
||||
LangGraphDeprecatedSinceV10,
|
||||
LangGraphDeprecatedSinceV11,
|
||||
)
|
||||
from tests.raw_pregel import NodeBuilder, RawPregel
|
||||
|
||||
|
||||
class PlainState(TypedDict): ...
|
||||
@@ -157,7 +157,7 @@ def test_config_type_deprecation_pregel(mocker: MockerFixture) -> None:
|
||||
LangGraphDeprecatedSinceV10,
|
||||
match="`config_type` is deprecated and will be removed. Please use `context_schema` instead.",
|
||||
):
|
||||
instance = RawPregel(
|
||||
instance = Pregel(
|
||||
nodes={
|
||||
"one": chain,
|
||||
},
|
||||
|
||||
@@ -39,6 +39,7 @@ from langgraph.channels.untracked_value import UntrackedValue
|
||||
from langgraph.constants import END, START
|
||||
from langgraph.graph import StateGraph
|
||||
from langgraph.graph.message import MessagesState, add_messages
|
||||
from langgraph.pregel import NodeBuilder, Pregel
|
||||
from langgraph.types import (
|
||||
Command,
|
||||
Durability,
|
||||
@@ -61,7 +62,6 @@ from tests.messages import (
|
||||
_AnyIdHumanMessage,
|
||||
_AnyIdToolMessage,
|
||||
)
|
||||
from tests.raw_pregel import NodeBuilder, RawPregel
|
||||
|
||||
|
||||
def test_invoke_two_processes_in_out_interrupt(
|
||||
@@ -72,7 +72,7 @@ def test_invoke_two_processes_in_out_interrupt(
|
||||
one = NodeBuilder().subscribe_only("input").do(add_one).write_to("inbox")
|
||||
two = NodeBuilder().subscribe_only("inbox").do(add_one).write_to("output")
|
||||
|
||||
app = RawPregel(
|
||||
app = Pregel(
|
||||
nodes={"one": one, "two": two},
|
||||
channels={
|
||||
"inbox": LastValue(int),
|
||||
|
||||
@@ -39,6 +39,7 @@ from langgraph.constants import END, START
|
||||
from langgraph.graph import MessagesState
|
||||
from langgraph.graph.message import add_messages
|
||||
from langgraph.graph.state import StateGraph
|
||||
from langgraph.pregel import NodeBuilder, Pregel
|
||||
from langgraph.types import PregelTask, Send, StateSnapshot, StreamWriter
|
||||
from tests.any_int import AnyInt
|
||||
from tests.any_str import AnyDict, AnyStr, UnsortedSequence
|
||||
@@ -50,7 +51,6 @@ from tests.messages import (
|
||||
_AnyIdHumanMessage,
|
||||
_AnyIdToolMessage,
|
||||
)
|
||||
from tests.raw_pregel import NodeBuilder, RawPregel
|
||||
|
||||
pytestmark = pytest.mark.anyio
|
||||
|
||||
@@ -61,7 +61,7 @@ async def test_invoke_two_processes_in_out_interrupt(
|
||||
add_one = mocker.Mock(side_effect=lambda x: x + 1)
|
||||
one = NodeBuilder().subscribe_only("input").do(add_one).write_to("inbox")
|
||||
two = NodeBuilder().subscribe_only("inbox").do(add_one).write_to("output")
|
||||
app = RawPregel(
|
||||
app = Pregel(
|
||||
nodes={"one": one, "two": two},
|
||||
channels={
|
||||
"inbox": LastValue(int),
|
||||
|
||||
@@ -66,7 +66,10 @@ from langgraph.errors import GraphRecursionError, InvalidUpdateError, ParentComm
|
||||
from langgraph.func import entrypoint, task
|
||||
from langgraph.graph import END, START, StateGraph
|
||||
from langgraph.graph.message import MessagesState, _messages_delta_reducer, add_messages
|
||||
from langgraph.pregel import Pregel
|
||||
from langgraph.pregel import (
|
||||
NodeBuilder,
|
||||
Pregel,
|
||||
)
|
||||
from langgraph.pregel._loop import PregelLoop, SyncPregelLoop
|
||||
from langgraph.pregel._runner import PregelRunner
|
||||
from langgraph.runtime import RunControl
|
||||
@@ -92,7 +95,6 @@ from tests.messages import (
|
||||
_AnyIdHumanMessage,
|
||||
_AnyIdToolMessage,
|
||||
)
|
||||
from tests.raw_pregel import NodeBuilder, RawPregel
|
||||
|
||||
pytestmark = pytest.mark.anyio
|
||||
|
||||
@@ -255,7 +257,7 @@ def test_context_json_schema() -> None:
|
||||
x: int
|
||||
y: str = field(default="foo")
|
||||
|
||||
app = RawPregel(
|
||||
app = Pregel(
|
||||
nodes={
|
||||
"one": chain,
|
||||
},
|
||||
@@ -447,7 +449,7 @@ def test_invoke_single_process_in_out(mocker: MockerFixture) -> None:
|
||||
add_one = mocker.Mock(side_effect=lambda x: x + 1)
|
||||
chain = NodeBuilder().subscribe_only("input").do(add_one).write_to("output")
|
||||
|
||||
app = RawPregel(
|
||||
app = Pregel(
|
||||
nodes={
|
||||
"one": chain,
|
||||
},
|
||||
@@ -483,7 +485,7 @@ def test_invoke_single_process_in_write_kwargs(mocker: MockerFixture) -> None:
|
||||
.write_to("output", fixed=5, output_plus_one=lambda x: x + 1)
|
||||
)
|
||||
|
||||
app = RawPregel(
|
||||
app = Pregel(
|
||||
nodes={"one": chain},
|
||||
channels={
|
||||
"input": LastValue(int),
|
||||
@@ -519,7 +521,7 @@ def test_invoke_single_process_in_out_dict(mocker: MockerFixture) -> None:
|
||||
add_one = mocker.Mock(side_effect=lambda x: x + 1)
|
||||
chain = NodeBuilder().subscribe_only("input").do(add_one).write_to("output")
|
||||
|
||||
app = RawPregel(
|
||||
app = Pregel(
|
||||
nodes={"one": chain},
|
||||
channels={"input": LastValue(int), "output": LastValue(int)},
|
||||
input_channels="input",
|
||||
@@ -544,7 +546,7 @@ def test_invoke_single_process_in_dict_out_dict(mocker: MockerFixture) -> None:
|
||||
add_one = mocker.Mock(side_effect=lambda x: x + 1)
|
||||
chain = NodeBuilder().subscribe_only("input").do(add_one).write_to("output")
|
||||
|
||||
app = RawPregel(
|
||||
app = Pregel(
|
||||
nodes={"one": chain},
|
||||
channels={"input": LastValue(int), "output": LastValue(int)},
|
||||
input_channels=["input"],
|
||||
@@ -570,7 +572,7 @@ def test_invoke_two_processes_in_out(mocker: MockerFixture) -> None:
|
||||
one = NodeBuilder().subscribe_only("input").do(add_one).write_to("inbox")
|
||||
two = NodeBuilder().subscribe_only("inbox").do(add_one).write_to("output")
|
||||
|
||||
app = RawPregel(
|
||||
app = Pregel(
|
||||
nodes={"one": one, "two": two},
|
||||
channels={
|
||||
"inbox": LastValue(int),
|
||||
@@ -675,7 +677,7 @@ def test_batch_two_processes_in_out() -> None:
|
||||
one = NodeBuilder().subscribe_only("input").do(add_one_with_delay).write_to("one")
|
||||
two = NodeBuilder().subscribe_only("one").do(add_one_with_delay).write_to("output")
|
||||
|
||||
app = RawPregel(
|
||||
app = Pregel(
|
||||
nodes={"one": one, "two": two},
|
||||
channels={
|
||||
"one": LastValue(int),
|
||||
@@ -707,7 +709,7 @@ def test_invoke_many_processes_in_out(mocker: MockerFixture) -> None:
|
||||
)
|
||||
nodes["last"] = NodeBuilder().subscribe_only(str(i)).do(add_one).write_to("output")
|
||||
|
||||
app = RawPregel(
|
||||
app = Pregel(
|
||||
nodes=nodes,
|
||||
channels={str(i): LastValue(int) for i in range(-1, test_size - 2)}
|
||||
| {"input": LastValue(int), "output": LastValue(int)},
|
||||
@@ -735,7 +737,7 @@ def test_batch_many_processes_in_out(mocker: MockerFixture) -> None:
|
||||
)
|
||||
nodes["last"] = NodeBuilder().subscribe_only(str(i)).do(add_one).write_to("output")
|
||||
|
||||
app = RawPregel(
|
||||
app = Pregel(
|
||||
nodes=nodes,
|
||||
channels={str(i): LastValue(int) for i in range(-1, test_size - 2)}
|
||||
| {"input": LastValue(int), "output": LastValue(int)},
|
||||
@@ -768,7 +770,7 @@ def test_invoke_two_processes_two_in_two_out_invalid(mocker: MockerFixture) -> N
|
||||
one = NodeBuilder().subscribe_only("input").do(add_one).write_to("output")
|
||||
two = NodeBuilder().subscribe_only("input").do(add_one).write_to("output")
|
||||
|
||||
app = RawPregel(
|
||||
app = Pregel(
|
||||
nodes={"one": one, "two": two},
|
||||
channels={"output": LastValue(int), "input": LastValue(int)},
|
||||
input_channels="input",
|
||||
@@ -801,7 +803,7 @@ def test_invoke_two_processes_two_in_two_out_valid(mocker: MockerFixture) -> Non
|
||||
one = NodeBuilder().subscribe_only("input").do(add_one).write_to("output")
|
||||
two = NodeBuilder().subscribe_only("input").do(add_one).write_to("output")
|
||||
|
||||
app = RawPregel(
|
||||
app = Pregel(
|
||||
nodes={"one": one, "two": two},
|
||||
channels={
|
||||
"input": LastValue(int),
|
||||
@@ -842,7 +844,7 @@ def test_invoke_checkpoint_two(
|
||||
.do(raise_if_above_10)
|
||||
)
|
||||
|
||||
app = RawPregel(
|
||||
app = Pregel(
|
||||
nodes={"one": one},
|
||||
channels={
|
||||
"total": BinaryOperatorAggregate(int, operator.add),
|
||||
@@ -1449,7 +1451,7 @@ def test_invoke_checkpoint_three(
|
||||
.do(raise_if_above_10)
|
||||
)
|
||||
|
||||
app = RawPregel(
|
||||
app = Pregel(
|
||||
nodes={"one": one},
|
||||
channels={
|
||||
"total": BinaryOperatorAggregate(int, operator.add),
|
||||
@@ -1574,7 +1576,7 @@ def test_invoke_two_processes_two_in_join_two_out(mocker: MockerFixture) -> None
|
||||
NodeBuilder().subscribe_only("inbox").do(add_10_each).write_to("output")
|
||||
)
|
||||
|
||||
app = RawPregel(
|
||||
app = Pregel(
|
||||
nodes={
|
||||
"one": one,
|
||||
"chain_three": chain_three,
|
||||
@@ -1605,7 +1607,7 @@ def test_invoke_join_then_call_other_pregel(
|
||||
add_one = mocker.Mock(side_effect=lambda x: x + 1)
|
||||
add_10_each = mocker.Mock(side_effect=lambda x: [y + 10 for y in x])
|
||||
|
||||
inner_app = RawPregel(
|
||||
inner_app = Pregel(
|
||||
nodes={
|
||||
"one": NodeBuilder().subscribe_only("input").do(add_one).write_to("output")
|
||||
},
|
||||
@@ -1626,7 +1628,7 @@ def test_invoke_join_then_call_other_pregel(
|
||||
)
|
||||
chain_three = NodeBuilder().subscribe_only("outbox_one").do(sum).write_to("output")
|
||||
|
||||
app = RawPregel(
|
||||
app = Pregel(
|
||||
nodes={
|
||||
"one": one,
|
||||
"two": two,
|
||||
@@ -1667,7 +1669,7 @@ def test_invoke_two_processes_one_in_two_out(mocker: MockerFixture) -> None:
|
||||
)
|
||||
two = NodeBuilder().subscribe_only("between").do(add_one).write_to("output")
|
||||
|
||||
app = RawPregel(
|
||||
app = Pregel(
|
||||
nodes={"one": one, "two": two},
|
||||
channels={
|
||||
"input": LastValue(int),
|
||||
@@ -1694,7 +1696,7 @@ def test_invoke_two_processes_no_out(mocker: MockerFixture) -> None:
|
||||
one = NodeBuilder().subscribe_only("input").do(add_one).write_to("between")
|
||||
two = NodeBuilder().subscribe_only("between").do(add_one)
|
||||
|
||||
app = RawPregel(
|
||||
app = Pregel(
|
||||
nodes={"one": one, "two": two},
|
||||
channels={
|
||||
"input": LastValue(int),
|
||||
@@ -1717,7 +1719,7 @@ def test_invoke_two_processes_no_in(mocker: MockerFixture) -> None:
|
||||
two = NodeBuilder().subscribe_only("between").do(add_one)
|
||||
|
||||
with pytest.raises(TypeError):
|
||||
RawPregel(nodes={"one": one, "two": two})
|
||||
Pregel(nodes={"one": one, "two": two})
|
||||
|
||||
|
||||
def test_conditional_entrypoint_to_multiple_state_graph(
|
||||
@@ -4126,7 +4128,7 @@ def test_xray_lance(snapshot: SnapshotAssertion):
|
||||
def test_channel_values(sync_checkpointer: BaseCheckpointSaver) -> None:
|
||||
config = {"configurable": {"thread_id": "1"}}
|
||||
chain = NodeBuilder().subscribe_only("input").write_to("output")
|
||||
app = RawPregel(
|
||||
app = Pregel(
|
||||
nodes={
|
||||
"one": chain,
|
||||
},
|
||||
@@ -8948,6 +8950,39 @@ def test_subgraph_streaming_sync() -> None:
|
||||
assert result["num_chunks"] == 9
|
||||
|
||||
|
||||
def test_subgraph_streaming_node_output_messages_sync() -> None:
|
||||
"""Messages returned by a node (not streamed by an LLM) are emitted when a
|
||||
graph is explicitly streamed with `stream_mode="messages"` inside a node."""
|
||||
|
||||
def respond(state: MessagesState) -> dict:
|
||||
return {"messages": [AIMessage("hi there", id="ai1")]}
|
||||
|
||||
subgraph = (
|
||||
StateGraph(MessagesState)
|
||||
.add_node("respond", respond)
|
||||
.add_edge(START, "respond")
|
||||
.compile()
|
||||
)
|
||||
msgs = {"messages": [("user", "hello")]}
|
||||
|
||||
top_level = [msg for msg, _ in subgraph.stream(msgs, stream_mode="messages")]
|
||||
|
||||
def parent_node(state: MessagesState, config: RunnableConfig) -> dict:
|
||||
nonlocal in_node
|
||||
in_node = [
|
||||
msg for msg, _ in subgraph.stream(msgs, config, stream_mode="messages")
|
||||
]
|
||||
return {}
|
||||
|
||||
in_node: list = []
|
||||
StateGraph(MessagesState).add_node("parent", parent_node).add_edge(
|
||||
START, "parent"
|
||||
).compile().invoke(msgs)
|
||||
|
||||
assert [m.content for m in top_level] == ["hi there"]
|
||||
assert [m.content for m in in_node] == ["hi there"]
|
||||
|
||||
|
||||
def test_get_graph_nonterminal_last_step_source(snapshot: SnapshotAssertion) -> None:
|
||||
class State(TypedDict):
|
||||
messages: list[str]
|
||||
@@ -9505,7 +9540,7 @@ class _IntVersionSaver(InMemorySaver):
|
||||
|
||||
|
||||
def _build_chain_after_a_delta_channel() -> Pregel:
|
||||
return RawPregel(
|
||||
return Pregel(
|
||||
nodes={
|
||||
"a": NodeBuilder().subscribe_only("inp").do(lambda _: ["a"]).write_to("d"),
|
||||
"b": NodeBuilder().subscribe_only("d").do(lambda _: "b").write_to("x"),
|
||||
@@ -9550,7 +9585,7 @@ async def test_aupdate_state_after_an_exit_snapshot_infers_the_last_writer() ->
|
||||
|
||||
|
||||
def test_update_state_after_a_fork_seal_infers_the_subscriber_that_ran() -> None:
|
||||
graph = RawPregel(
|
||||
graph = Pregel(
|
||||
nodes={
|
||||
"a": NodeBuilder()
|
||||
.subscribe_only("inp")
|
||||
@@ -9585,7 +9620,7 @@ def test_update_state_after_a_fork_seal_infers_the_subscriber_that_ran() -> None
|
||||
def test_update_state_after_a_fork_seal_infers_the_node_whose_read_it_advanced() -> (
|
||||
None
|
||||
):
|
||||
graph = RawPregel(
|
||||
graph = Pregel(
|
||||
nodes={
|
||||
"a": NodeBuilder().subscribe_only("inp").do(lambda _: ["a"]).write_to("d"),
|
||||
"b": NodeBuilder().subscribe_only("d").do(lambda _: "b").write_to("out"),
|
||||
|
||||
@@ -69,6 +69,7 @@ from langgraph.errors import (
|
||||
from langgraph.func import entrypoint, task
|
||||
from langgraph.graph import END, START, StateGraph
|
||||
from langgraph.graph.message import MessagesState, _messages_delta_reducer, add_messages
|
||||
from langgraph.pregel import NodeBuilder, Pregel
|
||||
from langgraph.pregel._loop import AsyncPregelLoop, PregelLoop
|
||||
from langgraph.pregel._runner import PregelRunner
|
||||
from langgraph.runtime import RunControl
|
||||
@@ -94,7 +95,6 @@ from tests.messages import (
|
||||
_AnyIdHumanMessage,
|
||||
_AnyIdToolMessage,
|
||||
)
|
||||
from tests.raw_pregel import NodeBuilder, RawPregel
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
@@ -1446,7 +1446,7 @@ async def test_invoke_single_process_in_out(mocker: MockerFixture) -> None:
|
||||
add_one = mocker.Mock(side_effect=lambda x: x + 1)
|
||||
chain = NodeBuilder().subscribe_only("input").do(add_one).write_to("output")
|
||||
|
||||
app = RawPregel(
|
||||
app = Pregel(
|
||||
nodes={
|
||||
"one": chain,
|
||||
},
|
||||
@@ -1479,7 +1479,7 @@ async def test_invoke_single_process_in_write_kwargs(mocker: MockerFixture) -> N
|
||||
.write_to("output", fixed=5, output_plus_one=lambda x: x + 1)
|
||||
)
|
||||
|
||||
app = RawPregel(
|
||||
app = Pregel(
|
||||
nodes={"one": chain},
|
||||
channels={
|
||||
"input": LastValue(int),
|
||||
@@ -1515,7 +1515,7 @@ async def test_invoke_single_process_in_out_dict(mocker: MockerFixture) -> None:
|
||||
add_one = mocker.Mock(side_effect=lambda x: x + 1)
|
||||
chain = NodeBuilder().subscribe_only("input").do(add_one).write_to("output")
|
||||
|
||||
app = RawPregel(
|
||||
app = Pregel(
|
||||
nodes={"one": chain},
|
||||
channels={"input": LastValue(int), "output": LastValue(int)},
|
||||
input_channels="input",
|
||||
@@ -1540,7 +1540,7 @@ async def test_invoke_single_process_in_dict_out_dict(mocker: MockerFixture) ->
|
||||
add_one = mocker.Mock(side_effect=lambda x: x + 1)
|
||||
chain = NodeBuilder().subscribe_only("input").do(add_one).write_to("output")
|
||||
|
||||
app = RawPregel(
|
||||
app = Pregel(
|
||||
nodes={"one": chain},
|
||||
channels={"input": LastValue(int), "output": LastValue(int)},
|
||||
input_channels=["input"],
|
||||
@@ -1567,7 +1567,7 @@ async def test_invoke_two_processes_in_out(mocker: MockerFixture) -> None:
|
||||
one = NodeBuilder().subscribe_only("input").do(add_one).write_to("inbox")
|
||||
two = NodeBuilder().subscribe_only("inbox").do(add_one).write_to("output")
|
||||
|
||||
app = RawPregel(
|
||||
app = Pregel(
|
||||
nodes={"one": one, "two": two},
|
||||
channels={
|
||||
"inbox": LastValue(int),
|
||||
@@ -1607,7 +1607,7 @@ async def test_batch_two_processes_in_out() -> None:
|
||||
one = NodeBuilder().subscribe_only("input").do(add_one_with_delay).write_to("one")
|
||||
two = NodeBuilder().subscribe_only("one").do(add_one_with_delay).write_to("output")
|
||||
|
||||
app = RawPregel(
|
||||
app = Pregel(
|
||||
nodes={"one": one, "two": two},
|
||||
channels={
|
||||
"one": LastValue(int),
|
||||
@@ -1639,7 +1639,7 @@ async def test_invoke_many_processes_in_out(mocker: MockerFixture) -> None:
|
||||
)
|
||||
nodes["last"] = NodeBuilder().subscribe_only(str(i)).do(add_one).write_to("output")
|
||||
|
||||
app = RawPregel(
|
||||
app = Pregel(
|
||||
nodes=nodes,
|
||||
channels={str(i): LastValue(int) for i in range(-1, test_size - 2)}
|
||||
| {"input": LastValue(int), "output": LastValue(int)},
|
||||
@@ -1668,7 +1668,7 @@ async def test_batch_many_processes_in_out(mocker: MockerFixture) -> None:
|
||||
)
|
||||
nodes["last"] = NodeBuilder().subscribe_only(str(i)).do(add_one).write_to("output")
|
||||
|
||||
app = RawPregel(
|
||||
app = Pregel(
|
||||
nodes=nodes,
|
||||
channels={str(i): LastValue(int) for i in range(-1, test_size - 2)}
|
||||
| {"input": LastValue(int), "output": LastValue(int)},
|
||||
@@ -1704,7 +1704,7 @@ async def test_invoke_two_processes_two_in_two_out_invalid(
|
||||
one = NodeBuilder().subscribe_only("input").do(add_one).write_to("output")
|
||||
two = NodeBuilder().subscribe_only("input").do(add_one).write_to("output")
|
||||
|
||||
app = RawPregel(
|
||||
app = Pregel(
|
||||
nodes={"one": one, "two": two},
|
||||
channels={"output": LastValue(int), "input": LastValue(int)},
|
||||
input_channels="input",
|
||||
@@ -1722,7 +1722,7 @@ async def test_invoke_two_processes_two_in_two_out_valid(mocker: MockerFixture)
|
||||
one = NodeBuilder().subscribe_only("input").do(add_one).write_to("output")
|
||||
two = NodeBuilder().subscribe_only("input").do(add_one).write_to("output")
|
||||
|
||||
app = RawPregel(
|
||||
app = Pregel(
|
||||
nodes={"one": one, "two": two},
|
||||
channels={
|
||||
"input": LastValue(int),
|
||||
@@ -1763,7 +1763,7 @@ async def test_invoke_checkpoint(
|
||||
.do(raise_if_above_10)
|
||||
)
|
||||
|
||||
app = RawPregel(
|
||||
app = Pregel(
|
||||
nodes={"one": one},
|
||||
channels={
|
||||
"total": BinaryOperatorAggregate(int, operator.add),
|
||||
@@ -3641,7 +3641,7 @@ async def test_invoke_checkpoint_three(
|
||||
.do(raise_if_above_10)
|
||||
)
|
||||
|
||||
app = RawPregel(
|
||||
app = Pregel(
|
||||
nodes={"one": one},
|
||||
channels={
|
||||
"total": BinaryOperatorAggregate(int, operator.add),
|
||||
@@ -3765,7 +3765,7 @@ async def test_invoke_two_processes_two_in_join_two_out(mocker: MockerFixture) -
|
||||
NodeBuilder().subscribe_only("inbox").do(add_10_each).write_to("output")
|
||||
)
|
||||
|
||||
app = RawPregel(
|
||||
app = Pregel(
|
||||
nodes={
|
||||
"one": one,
|
||||
"chain_three": chain_three,
|
||||
@@ -3799,7 +3799,7 @@ async def test_invoke_two_processes_one_in_two_out(mocker: MockerFixture) -> Non
|
||||
)
|
||||
two = NodeBuilder().subscribe_only("between").do(add_one).write_to("output")
|
||||
|
||||
app = RawPregel(
|
||||
app = Pregel(
|
||||
nodes={"one": one, "two": two},
|
||||
channels={
|
||||
"input": LastValue(int),
|
||||
@@ -3823,7 +3823,7 @@ async def test_invoke_two_processes_no_out(mocker: MockerFixture) -> None:
|
||||
one = NodeBuilder().subscribe_only("input").do(add_one).write_to("between")
|
||||
two = NodeBuilder().subscribe_only("between").do(add_one)
|
||||
|
||||
app = RawPregel(
|
||||
app = Pregel(
|
||||
nodes={"one": one, "two": two},
|
||||
channels={
|
||||
"input": LastValue(int),
|
||||
@@ -9503,6 +9503,44 @@ async def test_subgraph_streaming_async() -> None:
|
||||
assert result["num_chunks"] == 9
|
||||
|
||||
|
||||
async def test_subgraph_streaming_node_output_messages_async() -> None:
|
||||
"""Messages returned by a node (not streamed by an LLM) are emitted when a
|
||||
graph is explicitly streamed with `stream_mode="messages"` inside a node."""
|
||||
|
||||
async def respond(state: MessagesState) -> dict:
|
||||
return {"messages": [AIMessage("hi there", id="ai1")]}
|
||||
|
||||
subgraph = (
|
||||
StateGraph(MessagesState)
|
||||
.add_node("respond", respond)
|
||||
.add_edge(START, "respond")
|
||||
.compile()
|
||||
)
|
||||
msgs = {"messages": [("user", "hello")]}
|
||||
|
||||
top_level = [msg async for msg, _ in subgraph.astream(msgs, stream_mode="messages")]
|
||||
|
||||
async def parent_node(state: MessagesState, config: RunnableConfig) -> dict:
|
||||
nonlocal in_node
|
||||
in_node = [
|
||||
msg
|
||||
async for msg, _ in subgraph.astream(msgs, config, stream_mode="messages")
|
||||
]
|
||||
return {}
|
||||
|
||||
in_node: list = []
|
||||
await (
|
||||
StateGraph(MessagesState)
|
||||
.add_node("parent", parent_node)
|
||||
.add_edge(START, "parent")
|
||||
.compile()
|
||||
.ainvoke(msgs)
|
||||
)
|
||||
|
||||
assert [m.content for m in top_level] == ["hi there"]
|
||||
assert [m.content for m in in_node] == ["hi there"]
|
||||
|
||||
|
||||
@NEEDS_CONTEXTVARS
|
||||
async def test_null_resume_disallowed_with_multiple_interrupts(
|
||||
async_checkpointer: BaseCheckpointSaver,
|
||||
|
||||
@@ -1,56 +0,0 @@
|
||||
"""`Pregel` is the type of compiled graphs, built only by `StateGraph` and
|
||||
`@entrypoint`."""
|
||||
|
||||
import pytest
|
||||
from langchain_core.runnables import RunnableConfig
|
||||
from langgraph.checkpoint.memory import InMemorySaver
|
||||
from typing_extensions import TypedDict
|
||||
|
||||
import langgraph.pregel
|
||||
from langgraph.channels.ephemeral_value import EphemeralValue
|
||||
from langgraph.func import entrypoint
|
||||
from langgraph.graph import START, StateGraph
|
||||
from langgraph.pregel import Pregel
|
||||
from langgraph.pregel._read import PregelNode
|
||||
|
||||
|
||||
def test_pregel_cannot_be_constructed_directly() -> None:
|
||||
with pytest.raises(TypeError, match="`Pregel` cannot be constructed directly"):
|
||||
Pregel(
|
||||
nodes={"node": PregelNode(channels="a", triggers=["a"])},
|
||||
channels={"a": EphemeralValue(str)},
|
||||
input_channels="a",
|
||||
output_channels="a",
|
||||
)
|
||||
|
||||
|
||||
def test_node_builder_is_not_exported() -> None:
|
||||
assert langgraph.pregel.__all__ == ("Pregel",)
|
||||
assert not hasattr(langgraph.pregel, "NodeBuilder")
|
||||
|
||||
|
||||
def test_compiled_graphs_are_pregel_instances() -> None:
|
||||
class State(TypedDict):
|
||||
value: int
|
||||
|
||||
def increment(state: State) -> State:
|
||||
return {"value": state["value"] + 1}
|
||||
|
||||
builder = StateGraph(State)
|
||||
builder.add_node("increment", increment)
|
||||
builder.add_edge(START, "increment")
|
||||
graph = builder.compile(checkpointer=InMemorySaver())
|
||||
|
||||
@entrypoint(checkpointer=InMemorySaver())
|
||||
def workflow(value: int) -> int:
|
||||
return value + 1
|
||||
|
||||
assert isinstance(graph, Pregel)
|
||||
assert isinstance(workflow, Pregel)
|
||||
|
||||
# `with_config` copies the graph by calling its constructor again.
|
||||
config: RunnableConfig = {"configurable": {"thread_id": "1"}}
|
||||
assert graph.with_config(recursion_limit=5).invoke({"value": 1}, config) == {
|
||||
"value": 2
|
||||
}
|
||||
assert workflow.with_config(recursion_limit=5).invoke(1, config) == 2
|
||||
@@ -16,9 +16,11 @@ from langgraph_sdk.schema import StreamPart
|
||||
from pydantic import BaseModel
|
||||
from typing_extensions import TypedDict
|
||||
|
||||
from langgraph._internal._constants import CONFIG_KEY_STREAM
|
||||
from langgraph.errors import GraphInterrupt
|
||||
from langgraph.graph import END, START, MessagesState, StateGraph, add_messages
|
||||
from langgraph.pregel import Pregel
|
||||
from langgraph.pregel.protocol import StreamProtocol
|
||||
from langgraph.pregel.remote import RemoteGraph
|
||||
from langgraph.types import Interrupt, StateSnapshot
|
||||
from tests.any_str import AnyStr
|
||||
@@ -863,6 +865,87 @@ async def test_astream_subgraphs_keeps_caller_namespace():
|
||||
assert stream_parts == _CALLER_PREFIXED
|
||||
|
||||
|
||||
# Subgraph events requested only on behalf of a streaming parent graph
|
||||
_PARENT_STREAMED_PARTS = [
|
||||
StreamPart(event="updates|sub:abc", data={"inner": {"chunk": "inner"}}),
|
||||
StreamPart(event="messages|sub:abc", data=[{"content": "token"}, {}]),
|
||||
StreamPart(event="updates", data={"sub": {"chunk": "root"}}),
|
||||
]
|
||||
|
||||
|
||||
def _parent_streamed_config(parent_parts: list) -> RunnableConfig:
|
||||
return {
|
||||
"configurable": {
|
||||
"thread_id": "thread_1",
|
||||
"checkpoint_ns": "outer:123",
|
||||
CONFIG_KEY_STREAM: StreamProtocol(
|
||||
parent_parts.append, {"updates", "messages"}
|
||||
),
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
def test_stream_without_subgraphs_skips_parent_requested_subgraph_events():
|
||||
mock_sync_client = MagicMock()
|
||||
mock_sync_client.runs.stream.return_value = _PARENT_STREAMED_PARTS
|
||||
remote_pregel = RemoteGraph("test_graph_id", sync_client=mock_sync_client)
|
||||
|
||||
for stream_mode, expected in [
|
||||
("updates", [{"sub": {"chunk": "root"}}]),
|
||||
# messages are not filtered: the server dedupes them across namespaces
|
||||
("messages", [({"content": "token"}, {})]),
|
||||
]:
|
||||
parent_parts: list = []
|
||||
stream_parts = list(
|
||||
remote_pregel.stream(
|
||||
{"input": "data"},
|
||||
_parent_streamed_config(parent_parts),
|
||||
stream_mode=stream_mode,
|
||||
)
|
||||
)
|
||||
|
||||
assert stream_parts == expected
|
||||
# subgraph events are still requested and forwarded to the parent
|
||||
assert mock_sync_client.runs.stream.call_args.kwargs["stream_subgraphs"]
|
||||
assert [ns for ns, _, _ in parent_parts] == [
|
||||
("outer:123", "sub:abc"),
|
||||
("outer:123", "sub:abc"),
|
||||
("outer:123",),
|
||||
]
|
||||
|
||||
|
||||
async def test_astream_without_subgraphs_skips_parent_requested_subgraph_events():
|
||||
mock_async_client = MagicMock()
|
||||
async_iter = MagicMock()
|
||||
async_iter.__aiter__.return_value = _PARENT_STREAMED_PARTS
|
||||
mock_async_client.runs.stream.return_value = async_iter
|
||||
remote_pregel = RemoteGraph("test_graph_id", client=mock_async_client)
|
||||
|
||||
for stream_mode, expected in [
|
||||
("updates", [{"sub": {"chunk": "root"}}]),
|
||||
# messages are not filtered: the server dedupes them across namespaces
|
||||
("messages", [({"content": "token"}, {})]),
|
||||
]:
|
||||
parent_parts: list = []
|
||||
stream_parts = [
|
||||
part
|
||||
async for part in remote_pregel.astream(
|
||||
{"input": "data"},
|
||||
_parent_streamed_config(parent_parts),
|
||||
stream_mode=stream_mode,
|
||||
)
|
||||
]
|
||||
|
||||
assert stream_parts == expected
|
||||
# subgraph events are still requested and forwarded to the parent
|
||||
assert mock_async_client.runs.stream.call_args.kwargs["stream_subgraphs"]
|
||||
assert [ns for ns, _, _ in parent_parts] == [
|
||||
("outer:123", "sub:abc"),
|
||||
("outer:123", "sub:abc"),
|
||||
("outer:123",),
|
||||
]
|
||||
|
||||
|
||||
def test_invoke():
|
||||
# set up test
|
||||
mock_sync_client = MagicMock()
|
||||
|
||||
@@ -47,6 +47,7 @@ from langgraph.errors import (
|
||||
)
|
||||
from langgraph.func import entrypoint, task
|
||||
from langgraph.graph import END, START, StateGraph, add_messages
|
||||
from langgraph.pregel import NodeBuilder, Pregel
|
||||
from langgraph.pregel._read import PregelNode
|
||||
from langgraph.pregel._retry import (
|
||||
_checkpoint_ns_for_parent_command,
|
||||
@@ -66,7 +67,6 @@ from langgraph.types import (
|
||||
TimeoutPolicy,
|
||||
interrupt,
|
||||
)
|
||||
from tests.raw_pregel import NodeBuilder, RawPregel
|
||||
|
||||
NEEDS_CONTEXTVARS = pytest.mark.skipif(
|
||||
sys.version_info < (3, 11),
|
||||
@@ -1260,7 +1260,7 @@ def test_pregel_validate_rejects_sync_writer_timeout():
|
||||
return value
|
||||
|
||||
with pytest.raises(ValueError, match="only supported for async nodes"):
|
||||
RawPregel(
|
||||
Pregel(
|
||||
nodes={
|
||||
"slow": PregelNode(
|
||||
channels="input",
|
||||
@@ -1284,7 +1284,7 @@ def test_pregel_validate_rejects_wrapped_sync_runnable_lambda_timeout():
|
||||
return value + 1
|
||||
|
||||
with pytest.raises(ValueError, match="only supported for async nodes"):
|
||||
RawPregel(
|
||||
Pregel(
|
||||
nodes={
|
||||
"slow": (
|
||||
NodeBuilder()
|
||||
@@ -1307,7 +1307,7 @@ def test_pregel_validate_accepts_wrapped_async_runnable_lambda_timeout():
|
||||
async def slow(value: int) -> int:
|
||||
return value + 1
|
||||
|
||||
RawPregel(
|
||||
Pregel(
|
||||
nodes={
|
||||
"slow": (
|
||||
NodeBuilder()
|
||||
@@ -1334,7 +1334,7 @@ def test_pregel_validate_rejects_parallel_sync_branch_timeout():
|
||||
return value + 1
|
||||
|
||||
with pytest.raises(ValueError, match="only supported for async nodes"):
|
||||
RawPregel(
|
||||
Pregel(
|
||||
nodes={
|
||||
"parallel": (
|
||||
NodeBuilder()
|
||||
@@ -1363,7 +1363,7 @@ def test_pregel_validate_rejects_sync_node_timeout():
|
||||
return value + 1
|
||||
|
||||
with pytest.raises(ValueError, match="only supported for async nodes"):
|
||||
RawPregel(
|
||||
Pregel(
|
||||
nodes={
|
||||
"slow": (
|
||||
NodeBuilder()
|
||||
@@ -1388,7 +1388,7 @@ async def test_pregel_validate_accepts_async_runnable_lambda_timeout():
|
||||
await asyncio.sleep(0.2)
|
||||
return value + 1
|
||||
|
||||
graph = RawPregel(
|
||||
graph = Pregel(
|
||||
nodes={
|
||||
"slow": (
|
||||
NodeBuilder()
|
||||
@@ -1418,7 +1418,7 @@ async def test_pregel_validate_accepts_runnable_callable_with_sync_and_async_tim
|
||||
await asyncio.sleep(0.2)
|
||||
return value + 1
|
||||
|
||||
graph = RawPregel(
|
||||
graph = Pregel(
|
||||
nodes={
|
||||
"slow": (
|
||||
NodeBuilder()
|
||||
@@ -1632,7 +1632,7 @@ async def test_node_builder_timeout_e2e():
|
||||
await asyncio.sleep(0.2)
|
||||
return value + 1
|
||||
|
||||
graph = RawPregel(
|
||||
graph = Pregel(
|
||||
nodes={
|
||||
"slow": (
|
||||
NodeBuilder()
|
||||
|
||||
Generated
+3
-3
@@ -3419,11 +3419,11 @@ wheels = [
|
||||
|
||||
[[package]]
|
||||
name = "setuptools"
|
||||
version = "83.0.0"
|
||||
version = "84.1.0"
|
||||
source = { registry = "https://pypi.org/simple" }
|
||||
sdist = { url = "https://files.pythonhosted.org/packages/34/26/f5d29e25ffdb535afef2d35cdb55b325298f96debd670da4c325e08d70f4/setuptools-83.0.0.tar.gz", hash = "sha256:025bccbbf0fa05b6192bc64ae1e7b16e001fd6d6d4d5de03c97b1c1ade523bef", size = 1154254, upload-time = "2026-07-04T15:31:22.699Z" }
|
||||
sdist = { url = "https://files.pythonhosted.org/packages/73/eb/7e5f628ee2eb168fecbf8026de7856b2543d6d609cc85f086b9a91ac93d6/setuptools-84.1.0.tar.gz", hash = "sha256:23657205b63d4d31395ab917815247f694aafad28cd21316bf5fd45a51048ea0", size = 1171783, upload-time = "2026-10-19T20:11:43.489Z" }
|
||||
wheels = [
|
||||
{ url = "https://files.pythonhosted.org/packages/5d/40/e1e72872c6354b306daef1703549e8e83b4d43cfea356311bf722a043752/setuptools-83.0.0-py3-none-any.whl", hash = "sha256:29b23c360f22f414dc7336bb39178cc7bcbf6021ed2733cde173f09dba19abb3", size = 1008090, upload-time = "2026-07-04T15:31:20.885Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/18/02/b4bbb90f8f29e334898ff45c6d9b1fffaa6f378017052b3023a06262b5f7/setuptools-84.1.0-py3-none-any.whl", hash = "sha256:5231689c6ac8156cd155cfb98c3fd6fb252c3e7e9bd3d78b64d3c55f132624c9", size = 819380, upload-time = "2026-10-19T20:11:41.684Z" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
|
||||
@@ -14,9 +14,14 @@ import orjson
|
||||
|
||||
from langgraph_sdk._shared.utilities import (
|
||||
_orjson_default,
|
||||
_reconnect_delay,
|
||||
_validate_reconnect_location,
|
||||
)
|
||||
from langgraph_sdk.errors import _araise_for_status_typed
|
||||
from langgraph_sdk.errors import (
|
||||
InternalServerError,
|
||||
RateLimitError,
|
||||
_araise_for_status_typed,
|
||||
)
|
||||
from langgraph_sdk.schema import QueryParamTypes, StreamPart
|
||||
from langgraph_sdk.sse import SSEDecoder, aiter_lines_raw
|
||||
|
||||
@@ -230,65 +235,80 @@ class HttpClient:
|
||||
current_content = content if reconnect_path is None else None
|
||||
current_params = params if reconnect_path is None else None
|
||||
|
||||
retry = False
|
||||
async with self.client.stream(
|
||||
current_method,
|
||||
reconnect_path or path,
|
||||
headers=current_headers,
|
||||
content=current_content,
|
||||
params=current_params,
|
||||
) as res:
|
||||
if reconnect_path is None and on_response:
|
||||
on_response(res)
|
||||
# check status
|
||||
await _araise_for_status_typed(res)
|
||||
# check content type
|
||||
content_type = res.headers.get("content-type", "").partition(";")[0]
|
||||
if "text/event-stream" not in content_type:
|
||||
raise httpx.TransportError(
|
||||
"Expected response header Content-Type to contain 'text/event-stream', "
|
||||
f"got {content_type!r}"
|
||||
)
|
||||
reconnect_cause: httpx.HTTPError | None = None
|
||||
try:
|
||||
async with self.client.stream(
|
||||
current_method,
|
||||
reconnect_path or path,
|
||||
headers=current_headers,
|
||||
content=current_content,
|
||||
params=current_params,
|
||||
) as res:
|
||||
if reconnect_path is None and on_response:
|
||||
on_response(res)
|
||||
# check status
|
||||
await _araise_for_status_typed(res)
|
||||
# check content type
|
||||
content_type = res.headers.get("content-type", "").partition(";")[0]
|
||||
if "text/event-stream" not in content_type:
|
||||
raise httpx.TransportError(
|
||||
"Expected response header Content-Type to contain 'text/event-stream', "
|
||||
f"got {content_type!r}"
|
||||
)
|
||||
|
||||
reconnect_location = res.headers.get("location")
|
||||
if reconnect_location:
|
||||
_validate_reconnect_location(
|
||||
self.client.base_url, reconnect_location
|
||||
)
|
||||
reconnect_path = reconnect_location
|
||||
reconnect_location = res.headers.get("location")
|
||||
if reconnect_location:
|
||||
_validate_reconnect_location(
|
||||
self.client.base_url, reconnect_location
|
||||
)
|
||||
reconnect_path = reconnect_location
|
||||
|
||||
# parse SSE
|
||||
decoder = SSEDecoder()
|
||||
try:
|
||||
async for line in aiter_lines_raw(res):
|
||||
sse = decoder.decode(line=cast("bytes", line).rstrip(b"\n"))
|
||||
if sse is not None:
|
||||
# parse SSE
|
||||
decoder = SSEDecoder(last_event_id=last_event_id)
|
||||
try:
|
||||
async for line in aiter_lines_raw(res):
|
||||
sse = decoder.decode(line=cast("bytes", line).rstrip(b"\n"))
|
||||
if sse is not None:
|
||||
if decoder.last_event_id is not None:
|
||||
last_event_id = decoder.last_event_id
|
||||
if sse.event or decoder.has_data:
|
||||
# Progress since the last drop earns a fresh
|
||||
# set of reconnect attempts.
|
||||
reconnect_attempts = 0
|
||||
yield sse
|
||||
except httpx.HTTPError as err:
|
||||
# httpx.TransportError inherits from HTTPError, so transient
|
||||
# disconnects during streaming land here.
|
||||
if reconnect_path is None:
|
||||
raise
|
||||
reconnect_cause = err
|
||||
else:
|
||||
if sse := decoder.decode(b""):
|
||||
if decoder.last_event_id is not None:
|
||||
last_event_id = decoder.last_event_id
|
||||
if sse.event or sse.data is not None:
|
||||
if sse.event or decoder.has_data:
|
||||
# decoder.decode(b"") flushes the in-flight event and may
|
||||
# return an empty placeholder when there is no pending
|
||||
# message. Skip these no-op events so the stream doesn't
|
||||
# emit a trailing blank item after reconnects.
|
||||
yield sse
|
||||
except httpx.HTTPError:
|
||||
# httpx.TransportError inherits from HTTPError, so transient
|
||||
# disconnects during streaming land here.
|
||||
if reconnect_path is None:
|
||||
raise
|
||||
retry = True
|
||||
else:
|
||||
if sse := decoder.decode(b""):
|
||||
if decoder.last_event_id is not None:
|
||||
last_event_id = decoder.last_event_id
|
||||
if sse.event or sse.data is not None:
|
||||
# decoder.decode(b"") flushes the in-flight event and may
|
||||
# return an empty placeholder when there is no pending
|
||||
# message. Skip these no-op events so the stream doesn't
|
||||
# emit a trailing blank item after reconnects.
|
||||
yield sse
|
||||
if retry:
|
||||
except (httpx.TransportError, RateLimitError, InternalServerError) as err:
|
||||
# A reconnect that can't open (connection refused, 429, 5xx or a
|
||||
# non-SSE body) usually means the server is restarting or
|
||||
# mid-deploy, so retry it like a drop. Before the server sends a
|
||||
# Location there is nothing to reconnect to, so raise.
|
||||
if reconnect_path is None:
|
||||
raise
|
||||
reconnect_cause = err
|
||||
if reconnect_cause is not None:
|
||||
reconnect_attempts += 1
|
||||
if reconnect_attempts > max_reconnect_attempts:
|
||||
raise httpx.TransportError(
|
||||
"Exceeded maximum SSE reconnection attempts"
|
||||
)
|
||||
) from reconnect_cause
|
||||
await asyncio.sleep(
|
||||
_reconnect_delay(reconnect_attempts - 1, base=1.0, cap=5.0)
|
||||
)
|
||||
continue
|
||||
break
|
||||
|
||||
|
||||
@@ -16,7 +16,6 @@ from __future__ import annotations
|
||||
import asyncio
|
||||
import contextlib
|
||||
import logging
|
||||
import random
|
||||
from collections.abc import AsyncGenerator, AsyncIterator, Generator, Mapping
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Any, Literal, TypedDict, cast
|
||||
@@ -25,7 +24,7 @@ from langchain_core.language_models.chat_model_stream import AsyncChatModelStrea
|
||||
from langchain_protocol import Event, SubscribeParams
|
||||
|
||||
from langgraph_sdk._async.http import HttpClient
|
||||
from langgraph_sdk._shared.utilities import _quote_path_param
|
||||
from langgraph_sdk._shared.utilities import _quote_path_param, _reconnect_delay
|
||||
from langgraph_sdk.schema import LangSmithTracing, QueryParamTypes
|
||||
from langgraph_sdk.stream.controller import _SeenEventIds
|
||||
from langgraph_sdk.stream.decoders import (
|
||||
@@ -1691,11 +1690,13 @@ class AsyncThreadStream:
|
||||
|
||||
async def _reconnect_sleep(self, attempt: int) -> None:
|
||||
"""Sleep with exponential backoff and jitter for reconnect attempt `attempt`."""
|
||||
base = self._shared_reconnect_backoff_base
|
||||
cap = self._shared_reconnect_backoff_cap
|
||||
delay = min(cap, base * (2**attempt))
|
||||
jitter = random.uniform(0, delay * 0.25)
|
||||
await asyncio.sleep(delay + jitter)
|
||||
await asyncio.sleep(
|
||||
_reconnect_delay(
|
||||
attempt,
|
||||
base=self._shared_reconnect_backoff_base,
|
||||
cap=self._shared_reconnect_backoff_cap,
|
||||
)
|
||||
)
|
||||
|
||||
async def _reconnect_shared_stream(self) -> bool:
|
||||
"""Attempt to reopen the shared stream after a post-ready transport drop.
|
||||
|
||||
@@ -4,6 +4,7 @@ from __future__ import annotations
|
||||
|
||||
import functools
|
||||
import os
|
||||
import random
|
||||
import re
|
||||
from collections.abc import Mapping
|
||||
from datetime import tzinfo
|
||||
@@ -194,6 +195,15 @@ def _validate_reconnect_location(base_url: httpx.URL, location: str) -> str:
|
||||
return location
|
||||
|
||||
|
||||
def _reconnect_delay(attempt: int, *, base: float, cap: float) -> float:
|
||||
"""Seconds to wait before reconnect attempt `attempt` (0-based).
|
||||
|
||||
Exponential backoff capped at `cap`, plus up to 25% jitter.
|
||||
"""
|
||||
delay = min(cap, base * (2**attempt))
|
||||
return delay + random.uniform(0, delay * 0.25)
|
||||
|
||||
|
||||
def _provided_vals(d: Mapping[str, Any]) -> dict[str, Any]:
|
||||
return {k: v for k, v in d.items() if v is not None}
|
||||
|
||||
|
||||
@@ -4,6 +4,7 @@ from __future__ import annotations
|
||||
|
||||
import logging
|
||||
import sys
|
||||
import time
|
||||
import warnings
|
||||
from collections.abc import Callable, Iterator, Mapping
|
||||
from typing import Any, cast
|
||||
@@ -13,9 +14,14 @@ import orjson
|
||||
|
||||
from langgraph_sdk._shared.utilities import (
|
||||
_orjson_default,
|
||||
_reconnect_delay,
|
||||
_validate_reconnect_location,
|
||||
)
|
||||
from langgraph_sdk.errors import _raise_for_status_typed
|
||||
from langgraph_sdk.errors import (
|
||||
InternalServerError,
|
||||
RateLimitError,
|
||||
_raise_for_status_typed,
|
||||
)
|
||||
from langgraph_sdk.schema import QueryParamTypes, StreamPart
|
||||
from langgraph_sdk.sse import SSEDecoder, iter_lines_raw
|
||||
|
||||
@@ -232,62 +238,75 @@ class SyncHttpClient:
|
||||
current_content = content if reconnect_path is None else None
|
||||
current_params = params if reconnect_path is None else None
|
||||
|
||||
retry = False
|
||||
with self.client.stream(
|
||||
current_method,
|
||||
reconnect_path or path,
|
||||
headers=current_headers,
|
||||
content=current_content,
|
||||
params=current_params,
|
||||
) as res:
|
||||
if reconnect_path is None and on_response:
|
||||
on_response(res)
|
||||
# check status
|
||||
_raise_for_status_typed(res)
|
||||
# check content type
|
||||
content_type = res.headers.get("content-type", "").partition(";")[0]
|
||||
if "text/event-stream" not in content_type:
|
||||
raise httpx.TransportError(
|
||||
"Expected response header Content-Type to contain 'text/event-stream', "
|
||||
f"got {content_type!r}"
|
||||
)
|
||||
reconnect_cause: httpx.HTTPError | None = None
|
||||
try:
|
||||
with self.client.stream(
|
||||
current_method,
|
||||
reconnect_path or path,
|
||||
headers=current_headers,
|
||||
content=current_content,
|
||||
params=current_params,
|
||||
) as res:
|
||||
if reconnect_path is None and on_response:
|
||||
on_response(res)
|
||||
# check status
|
||||
_raise_for_status_typed(res)
|
||||
# check content type
|
||||
content_type = res.headers.get("content-type", "").partition(";")[0]
|
||||
if "text/event-stream" not in content_type:
|
||||
raise httpx.TransportError(
|
||||
"Expected response header Content-Type to contain 'text/event-stream', "
|
||||
f"got {content_type!r}"
|
||||
)
|
||||
|
||||
reconnect_location = res.headers.get("location")
|
||||
if reconnect_location:
|
||||
_validate_reconnect_location(
|
||||
self.client.base_url, reconnect_location
|
||||
)
|
||||
reconnect_path = reconnect_location
|
||||
reconnect_location = res.headers.get("location")
|
||||
if reconnect_location:
|
||||
_validate_reconnect_location(
|
||||
self.client.base_url, reconnect_location
|
||||
)
|
||||
reconnect_path = reconnect_location
|
||||
|
||||
decoder = SSEDecoder()
|
||||
try:
|
||||
for line in iter_lines_raw(res):
|
||||
sse = decoder.decode(cast(bytes, line).rstrip(b"\n"))
|
||||
if sse is not None:
|
||||
decoder = SSEDecoder(last_event_id=last_event_id)
|
||||
try:
|
||||
for line in iter_lines_raw(res):
|
||||
sse = decoder.decode(cast(bytes, line).rstrip(b"\n"))
|
||||
if sse is not None:
|
||||
if decoder.last_event_id is not None:
|
||||
last_event_id = decoder.last_event_id
|
||||
if sse.event or decoder.has_data:
|
||||
# Progress since the last drop earns a fresh
|
||||
# set of reconnect attempts.
|
||||
reconnect_attempts = 0
|
||||
yield sse
|
||||
except httpx.HTTPError as err:
|
||||
# httpx.TransportError inherits from HTTPError, so transient
|
||||
# disconnects during streaming land here.
|
||||
if reconnect_path is None:
|
||||
raise
|
||||
reconnect_cause = err
|
||||
else:
|
||||
if sse := decoder.decode(b""):
|
||||
if decoder.last_event_id is not None:
|
||||
last_event_id = decoder.last_event_id
|
||||
if sse.event or sse.data is not None:
|
||||
if sse.event or decoder.has_data:
|
||||
# See async stream implementation for rationale on
|
||||
# skipping empty flush events.
|
||||
yield sse
|
||||
except httpx.HTTPError:
|
||||
# httpx.TransportError inherits from HTTPError, so transient
|
||||
# disconnects during streaming land here.
|
||||
if reconnect_path is None:
|
||||
raise
|
||||
retry = True
|
||||
else:
|
||||
if sse := decoder.decode(b""):
|
||||
if decoder.last_event_id is not None:
|
||||
last_event_id = decoder.last_event_id
|
||||
if sse.event or sse.data is not None:
|
||||
# See async stream implementation for rationale on
|
||||
# skipping empty flush events.
|
||||
yield sse
|
||||
if retry:
|
||||
except (httpx.TransportError, RateLimitError, InternalServerError) as err:
|
||||
# A reconnect that can't open (connection refused, 429, 5xx or a
|
||||
# non-SSE body) usually means the server is restarting or
|
||||
# mid-deploy, so retry it like a drop. Before the server sends a
|
||||
# Location there is nothing to reconnect to, so raise.
|
||||
if reconnect_path is None:
|
||||
raise
|
||||
reconnect_cause = err
|
||||
if reconnect_cause is not None:
|
||||
reconnect_attempts += 1
|
||||
if reconnect_attempts > max_reconnect_attempts:
|
||||
raise httpx.TransportError(
|
||||
"Exceeded maximum SSE reconnection attempts"
|
||||
)
|
||||
) from reconnect_cause
|
||||
time.sleep(_reconnect_delay(reconnect_attempts - 1, base=1.0, cap=5.0))
|
||||
continue
|
||||
break
|
||||
|
||||
|
||||
@@ -811,7 +811,7 @@ class _On:
|
||||
def __call__(
|
||||
self,
|
||||
*,
|
||||
resources: str | Sequence[str],
|
||||
resources: str | Sequence[str] | None = None,
|
||||
actions: str | Sequence[str] | None = None,
|
||||
) -> Callable[[AHO], AHO]: ...
|
||||
|
||||
|
||||
@@ -30,18 +30,14 @@ class BytesLineDecoder:
|
||||
# See https://docs.python.org/3/glossary.html#term-universal-newlines
|
||||
NEWLINE_CHARS = b"\n\r"
|
||||
|
||||
# We always push a trailing `\r` into the next decode iteration.
|
||||
if self.trailing_cr:
|
||||
text = b"\r" + text
|
||||
self.trailing_cr = False
|
||||
if text.endswith(b"\r"):
|
||||
self.trailing_cr = True
|
||||
text = text[:-1]
|
||||
if not text:
|
||||
return []
|
||||
if self.trailing_cr and text.startswith(b"\n"):
|
||||
text = text[1:]
|
||||
self.trailing_cr = text.endswith(b"\r")
|
||||
|
||||
if not text:
|
||||
# NOTE: the edge case input of empty text doesn't occur in practice,
|
||||
# because other httpx internals filter out this value
|
||||
return [] # pragma: no cover
|
||||
return []
|
||||
|
||||
trailing_newline = text[-1] in NEWLINE_CHARS
|
||||
lines = cast(list[BytesLike], text.splitlines())
|
||||
@@ -66,7 +62,7 @@ class BytesLineDecoder:
|
||||
return lines
|
||||
|
||||
def flush(self) -> list[BytesLike]:
|
||||
if not self.buffer and not self.trailing_cr:
|
||||
if not self.buffer:
|
||||
return []
|
||||
|
||||
lines: list[BytesLike] = [self.buffer]
|
||||
@@ -76,18 +72,25 @@ class BytesLineDecoder:
|
||||
|
||||
|
||||
class SSEDecoder:
|
||||
def __init__(self) -> None:
|
||||
def __init__(self, *, last_event_id: str | None = None) -> None:
|
||||
self._event = ""
|
||||
self._data = bytearray()
|
||||
self._last_event_id: str | None = None
|
||||
self._last_event_id = last_event_id
|
||||
self._retry: int | None = None
|
||||
self._at_stream_start = True
|
||||
self._has_data = False
|
||||
self._id_changed = False
|
||||
|
||||
@property
|
||||
def has_data(self) -> bool:
|
||||
"""Whether the last decoded block contained a data field."""
|
||||
return self._has_data
|
||||
|
||||
@property
|
||||
def last_event_id(self) -> str | None:
|
||||
"""Return the last event identifier that was seen.
|
||||
|
||||
`None` until the stream sends an `id` field, and `""` after an empty
|
||||
`None` until an ID is supplied or received, and `""` after an empty
|
||||
one, which resets the reconnect cursor.
|
||||
"""
|
||||
|
||||
@@ -101,17 +104,26 @@ class SSEDecoder:
|
||||
line = line.removeprefix(b"\xef\xbb\xbf")
|
||||
|
||||
if not line:
|
||||
self._has_data = bool(self._data)
|
||||
if (
|
||||
not self._event
|
||||
and not self._data
|
||||
and self._last_event_id is None
|
||||
and not self._id_changed
|
||||
and self._retry is None
|
||||
):
|
||||
return None
|
||||
|
||||
body = self._data[:-1]
|
||||
data = None
|
||||
if body and not body.isspace():
|
||||
try:
|
||||
data = orjson.loads(body)
|
||||
except orjson.JSONDecodeError:
|
||||
data = orjson.loads(body.decode(errors="replace"))
|
||||
|
||||
sse = StreamPart(
|
||||
event=self._event,
|
||||
data=orjson.loads(self._data) if self._data else None, # ty: ignore[invalid-argument-type]
|
||||
data=data, # ty: ignore[invalid-argument-type]
|
||||
id=self._last_event_id or None,
|
||||
)
|
||||
|
||||
@@ -119,6 +131,7 @@ class SSEDecoder:
|
||||
self._event = ""
|
||||
self._data = bytearray()
|
||||
self._retry = None
|
||||
self._id_changed = False
|
||||
|
||||
return sse
|
||||
|
||||
@@ -131,16 +144,16 @@ class SSEDecoder:
|
||||
value = value[1:]
|
||||
|
||||
if fieldname == b"event":
|
||||
self._event = value.decode()
|
||||
self._event = value.decode(errors="replace")
|
||||
elif fieldname == b"data":
|
||||
if self._data:
|
||||
self._data.extend(b"\n")
|
||||
self._data.extend(value)
|
||||
self._data.extend(b"\n")
|
||||
elif fieldname == b"id":
|
||||
if b"\0" in value:
|
||||
pass
|
||||
else:
|
||||
self._last_event_id = value.decode()
|
||||
self._last_event_id = value.decode(errors="replace")
|
||||
self._id_changed = True
|
||||
elif fieldname == b"retry":
|
||||
with contextlib.suppress(TypeError, ValueError):
|
||||
self._retry = int(value)
|
||||
|
||||
@@ -15,7 +15,6 @@ from __future__ import annotations
|
||||
import asyncio
|
||||
import contextlib
|
||||
import logging
|
||||
import random
|
||||
from collections import OrderedDict
|
||||
from collections.abc import AsyncGenerator, AsyncIterator, Awaitable, Callable
|
||||
from dataclasses import dataclass, field
|
||||
@@ -23,6 +22,7 @@ from typing import Any
|
||||
|
||||
from langchain_protocol import Event, SubscribeParams
|
||||
|
||||
from langgraph_sdk._shared.utilities import _reconnect_delay
|
||||
from langgraph_sdk.stream.transport import AsyncProtocolTransport, EventStreamHandle
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
@@ -275,11 +275,13 @@ class StreamController:
|
||||
|
||||
async def _reconnect_sleep(self, attempt: int) -> None:
|
||||
"""Sleep with exponential backoff and jitter for reconnect attempt *attempt*."""
|
||||
base = self._reconnect_backoff_base
|
||||
cap = self._reconnect_backoff_cap
|
||||
delay = min(cap, base * (2**attempt))
|
||||
jitter = random.uniform(0, delay * 0.25)
|
||||
await asyncio.sleep(delay + jitter)
|
||||
await asyncio.sleep(
|
||||
_reconnect_delay(
|
||||
attempt,
|
||||
base=self._reconnect_backoff_base,
|
||||
cap=self._reconnect_backoff_cap,
|
||||
)
|
||||
)
|
||||
|
||||
async def _reconnect_shared_stream(self) -> bool:
|
||||
"""Attempt to reopen the shared stream after a transport drop.
|
||||
|
||||
@@ -4,7 +4,6 @@ from __future__ import annotations
|
||||
|
||||
import contextlib
|
||||
import logging
|
||||
import random
|
||||
import threading
|
||||
import time
|
||||
from dataclasses import dataclass, field
|
||||
@@ -13,6 +12,7 @@ from typing import Any
|
||||
|
||||
from langchain_protocol import Event, SubscribeParams
|
||||
|
||||
from langgraph_sdk._shared.utilities import _reconnect_delay
|
||||
from langgraph_sdk.stream.errors import (
|
||||
is_non_retryable_stream_error,
|
||||
stream_error_frame,
|
||||
@@ -282,11 +282,13 @@ class SyncStreamController:
|
||||
|
||||
def _reconnect_sleep(self, attempt: int) -> None:
|
||||
"""Sleep with exponential backoff + jitter before a reconnect attempt."""
|
||||
base = self._reconnect_backoff_base
|
||||
cap = self._reconnect_backoff_cap
|
||||
delay = min(cap, base * (2**attempt))
|
||||
jitter = random.uniform(0, delay * 0.25)
|
||||
time.sleep(delay + jitter)
|
||||
time.sleep(
|
||||
_reconnect_delay(
|
||||
attempt,
|
||||
base=self._reconnect_backoff_base,
|
||||
cap=self._reconnect_backoff_cap,
|
||||
)
|
||||
)
|
||||
|
||||
def _reconnect_shared_stream(self) -> bool:
|
||||
"""Attempt to reopen the shared stream after a transport drop.
|
||||
|
||||
@@ -1,8 +1,49 @@
|
||||
import ast
|
||||
import inspect
|
||||
|
||||
import pytest
|
||||
|
||||
from langgraph_sdk import Auth
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"kwargs",
|
||||
[{}, {"actions": "read"}, {"resources": None}],
|
||||
)
|
||||
def test_global_handler_factory_accepts_optional_resources(kwargs) -> None:
|
||||
auth = Auth()
|
||||
tree = ast.parse(inspect.getsource(type(auth.on)))
|
||||
cls = tree.body[0]
|
||||
assert isinstance(cls, ast.ClassDef)
|
||||
factory_overload = next(
|
||||
node
|
||||
for node in cls.body
|
||||
if isinstance(node, ast.FunctionDef)
|
||||
and node.name == "__call__"
|
||||
and node.decorator_list
|
||||
and any(arg.arg == "resources" for arg in node.args.kwonlyargs)
|
||||
)
|
||||
resources = next(
|
||||
arg for arg in factory_overload.args.kwonlyargs if arg.arg == "resources"
|
||||
)
|
||||
assert resources.annotation is not None
|
||||
assert ast.unparse(resources.annotation) == "str | Sequence[str] | None"
|
||||
default = factory_overload.args.kw_defaults[0]
|
||||
assert isinstance(default, ast.Constant) and default.value is None
|
||||
|
||||
@auth.on(**kwargs)
|
||||
async def handler(ctx, value):
|
||||
del ctx, value
|
||||
return None
|
||||
|
||||
if kwargs.get("actions"):
|
||||
assert auth._handlers == {("*", "read"): [handler]}
|
||||
assert auth._global_handlers == []
|
||||
else:
|
||||
assert auth._global_handlers == [handler]
|
||||
assert auth._handlers == {}
|
||||
|
||||
|
||||
def test_handler_multiple_resources_and_actions() -> None:
|
||||
auth = Auth()
|
||||
|
||||
|
||||
@@ -1,6 +1,8 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Iterator, Sequence
|
||||
import asyncio
|
||||
import time
|
||||
from collections.abc import Callable, Iterator, Sequence
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
@@ -13,6 +15,7 @@ from langgraph_sdk._async.runs import _wrap_stream_v2
|
||||
from langgraph_sdk._shared.utilities import _sse_to_v2_dict
|
||||
from langgraph_sdk._sync.runs import _wrap_stream_v2_sync
|
||||
from langgraph_sdk.client import HttpClient, SyncHttpClient
|
||||
from langgraph_sdk.errors import InternalServerError, NotFoundError
|
||||
from langgraph_sdk.schema import (
|
||||
CheckpointPayload,
|
||||
CheckpointsStreamPart,
|
||||
@@ -67,6 +70,19 @@ class ListByteStream(httpx.ByteStream):
|
||||
return None
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def no_sleep(monkeypatch):
|
||||
"""Record backoff sleeps without actually waiting."""
|
||||
sleeps: list[float] = []
|
||||
|
||||
async def fake_async_sleep(d):
|
||||
sleeps.append(d)
|
||||
|
||||
monkeypatch.setattr(asyncio, "sleep", fake_async_sleep)
|
||||
monkeypatch.setattr(time, "sleep", lambda d: sleeps.append(d))
|
||||
return sleeps
|
||||
|
||||
|
||||
def iter_lines_raw(payload: list[bytes]) -> Iterator[BytesLike]:
|
||||
decoder = BytesLineDecoder()
|
||||
for part in payload:
|
||||
@@ -130,6 +146,110 @@ def test_sse_decoder_strips_a_leading_bom():
|
||||
assert decoder.decode(b"") == StreamPart(event="values", data=1, id=None) # ty: ignore[invalid-argument-type]
|
||||
|
||||
|
||||
def test_control_blocks_do_not_repeat_after_a_cursor_update():
|
||||
decoder = SSEDecoder()
|
||||
decoder.decode(b"id: cursor")
|
||||
assert decoder.decode(b"") == StreamPart(event="", data=None, id="cursor") # ty: ignore[invalid-argument-type]
|
||||
decoder.decode(b": heartbeat")
|
||||
assert decoder.decode(b"") is None
|
||||
assert decoder.decode(b"") is None
|
||||
decoder.decode(b"id:")
|
||||
assert decoder.decode(b"") == StreamPart(event="", data=None, id=None) # ty: ignore[invalid-argument-type]
|
||||
assert decoder.last_event_id == ""
|
||||
assert decoder.decode(b"") is None
|
||||
|
||||
|
||||
def test_cr_lines_are_emitted_without_waiting_for_another_chunk():
|
||||
decoder = BytesLineDecoder()
|
||||
assert decoder.decode(b"data: 1\r\r") == [b"data: 1", b""]
|
||||
assert decoder.decode(b"\n") == []
|
||||
assert decoder.decode(b"data: 2\r") == [b"data: 2"]
|
||||
assert decoder.decode(b"") == []
|
||||
assert decoder.decode(b"\n\r") == [b""]
|
||||
assert decoder.flush() == []
|
||||
|
||||
|
||||
def test_empty_data_lines_preserve_their_newlines():
|
||||
decoder = SSEDecoder()
|
||||
for line in (b"data", b"data: 1", b"data:"):
|
||||
decoder.decode(line)
|
||||
assert decoder._data == b"\n1\n\n"
|
||||
assert decoder.decode(b"") == StreamPart(event="", data=1, id=None) # ty: ignore[invalid-argument-type]
|
||||
decoder.decode(b"data:")
|
||||
decoder.decode(b"data:")
|
||||
assert decoder.decode(b"") == StreamPart(event="", data=None, id=None) # ty: ignore[invalid-argument-type]
|
||||
assert decoder.has_data
|
||||
assert decoder.decode(b"") is None
|
||||
assert not decoder.has_data
|
||||
|
||||
|
||||
def test_malformed_utf8_is_replaced_in_fields_and_json_strings():
|
||||
decoder = SSEDecoder()
|
||||
for line in (b"event: \xff", b"id: \xff", b'data: "\xff"'):
|
||||
decoder.decode(line)
|
||||
assert decoder.decode(b"") == StreamPart(event="\ufffd", data="\ufffd", id="\ufffd") # ty: ignore[invalid-argument-type]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("asynchronous", [False, True])
|
||||
async def test_null_and_empty_data_are_not_dropped(asynchronous: bool):
|
||||
def handler(_request: httpx.Request) -> httpx.Response:
|
||||
return httpx.Response(
|
||||
200,
|
||||
headers={"Content-Type": "text/event-stream"},
|
||||
content=b"data: null\n\ndata:\n\n",
|
||||
)
|
||||
|
||||
transport = httpx.MockTransport(handler)
|
||||
if asynchronous:
|
||||
async with httpx.AsyncClient(
|
||||
transport=transport, base_url="https://example.com"
|
||||
) as client:
|
||||
parts = [part async for part in HttpClient(client).stream("/stream", "GET")]
|
||||
else:
|
||||
with httpx.Client(
|
||||
transport=transport, base_url="https://example.com"
|
||||
) as client:
|
||||
parts = list(SyncHttpClient(client).stream("/stream", "GET"))
|
||||
assert parts == [StreamPart(event="", data=None), StreamPart(event="", data=None)] # ty: ignore[invalid-argument-type]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("asynchronous", [False, True])
|
||||
async def test_cursor_is_inherited_after_reconnect(asynchronous: bool):
|
||||
requests: list[httpx.Request] = []
|
||||
|
||||
def handler(request: httpx.Request) -> httpx.Response:
|
||||
requests.append(request)
|
||||
if len(requests) == 1:
|
||||
stream_cls = AsyncListByteStream if asynchronous else ListByteStream
|
||||
return httpx.Response(
|
||||
200,
|
||||
headers={"Content-Type": "text/event-stream", "Location": "/reconnect"},
|
||||
stream=stream_cls(
|
||||
[b"id: cursor\nevent: values\ndata: 1\n\n"],
|
||||
httpx.RemoteProtocolError("incomplete chunked read"),
|
||||
),
|
||||
)
|
||||
assert (b"Last-Event-ID", b"cursor") in request.headers.raw
|
||||
return httpx.Response(
|
||||
200,
|
||||
headers={"Content-Type": "text/event-stream"},
|
||||
content=b"event: values\ndata: 2\n\n",
|
||||
)
|
||||
|
||||
transport = httpx.MockTransport(handler)
|
||||
if asynchronous:
|
||||
async with httpx.AsyncClient(
|
||||
transport=transport, base_url="https://example.com"
|
||||
) as client:
|
||||
parts = [part async for part in HttpClient(client).stream("/stream", "GET")]
|
||||
else:
|
||||
with httpx.Client(
|
||||
transport=transport, base_url="https://example.com"
|
||||
) as client:
|
||||
parts = list(SyncHttpClient(client).stream("/stream", "GET"))
|
||||
assert [part.id for part in parts] == ["cursor", "cursor"]
|
||||
|
||||
|
||||
# --- HTTP client streaming ---
|
||||
|
||||
|
||||
@@ -176,6 +296,7 @@ def test_sync_http_client_stream_flushes_trailing_event():
|
||||
assert parts == [StreamPart(event="foo", data={"bar": 1})]
|
||||
|
||||
|
||||
@pytest.mark.usefixtures("no_sleep")
|
||||
def test_sync_http_client_stream_recovers_after_disconnect():
|
||||
reconnect_path = "/reconnect"
|
||||
first_chunks = [
|
||||
@@ -241,6 +362,7 @@ def test_sync_http_client_stream_recovers_after_disconnect():
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.usefixtures("no_sleep")
|
||||
async def test_http_client_stream_recovers_after_disconnect():
|
||||
reconnect_path = "/reconnect"
|
||||
first_chunks = [
|
||||
@@ -375,6 +497,7 @@ _CALLER_CURSORS = pytest.mark.parametrize(
|
||||
|
||||
@_EMPTY_IDS
|
||||
@_CALLER_CURSORS
|
||||
@pytest.mark.usefixtures("no_sleep")
|
||||
def test_sync_stream_reconnects_without_a_cursor_after_an_empty_id(
|
||||
empty_id: bytes, headers: dict[str, str] | None
|
||||
):
|
||||
@@ -387,6 +510,7 @@ def test_sync_stream_reconnects_without_a_cursor_after_an_empty_id(
|
||||
|
||||
@_EMPTY_IDS
|
||||
@_CALLER_CURSORS
|
||||
@pytest.mark.usefixtures("no_sleep")
|
||||
async def test_async_stream_reconnects_without_a_cursor_after_an_empty_id(
|
||||
empty_id: bytes, headers: dict[str, str] | None
|
||||
):
|
||||
@@ -397,6 +521,7 @@ async def test_async_stream_reconnects_without_a_cursor_after_an_empty_id(
|
||||
assert "last-event-id" not in reconnect.headers
|
||||
|
||||
|
||||
@pytest.mark.usefixtures("no_sleep")
|
||||
def test_sync_stream_id_replaces_a_lowercase_caller_cursor():
|
||||
reconnect = _sync_reconnect(
|
||||
[b"id: 1\nevent: values\ndata: 1\n\n"], {"last-event-id": "5"}
|
||||
@@ -405,6 +530,7 @@ def test_sync_stream_id_replaces_a_lowercase_caller_cursor():
|
||||
assert reconnect.headers.get_list("last-event-id") == ["1"]
|
||||
|
||||
|
||||
@pytest.mark.usefixtures("no_sleep")
|
||||
async def test_async_stream_id_replaces_a_lowercase_caller_cursor():
|
||||
reconnect = await _async_reconnect(
|
||||
[b"id: 1\nevent: values\ndata: 1\n\n"], {"last-event-id": "5"}
|
||||
@@ -413,12 +539,14 @@ async def test_async_stream_id_replaces_a_lowercase_caller_cursor():
|
||||
assert reconnect.headers.get_list("last-event-id") == ["1"]
|
||||
|
||||
|
||||
@pytest.mark.usefixtures("no_sleep")
|
||||
def test_sync_stream_keeps_the_caller_cursor_until_the_stream_sends_an_id():
|
||||
reconnect = _sync_reconnect([b"event: values\ndata: 1\n\n"], {"Last-Event-ID": "5"})
|
||||
|
||||
assert reconnect.headers.get_list("last-event-id") == ["5"]
|
||||
|
||||
|
||||
@pytest.mark.usefixtures("no_sleep")
|
||||
async def test_async_stream_keeps_the_caller_cursor_until_the_stream_sends_an_id():
|
||||
reconnect = await _async_reconnect(
|
||||
[b"event: values\ndata: 1\n\n"], {"Last-Event-ID": "5"}
|
||||
@@ -427,6 +555,236 @@ async def test_async_stream_keeps_the_caller_cursor_until_the_stream_sends_an_id
|
||||
assert reconnect.headers.get_list("last-event-id") == ["5"]
|
||||
|
||||
|
||||
# --- reconnects survive a server restart or a deploy ---
|
||||
|
||||
SSE = {"Content-Type": "text/event-stream"}
|
||||
DROP = httpx.RemoteProtocolError("incomplete chunked read")
|
||||
|
||||
|
||||
def _event(i: int) -> bytes:
|
||||
return f"id: {i}\nevent: values\ndata: {i}\n\n".encode()
|
||||
|
||||
|
||||
def _scripted(steps: list[Callable[[httpx.Request], httpx.Response]]):
|
||||
"""Handler that plays one step per request and records each request."""
|
||||
seen: list[httpx.Request] = []
|
||||
|
||||
def handler(request: httpx.Request) -> httpx.Response:
|
||||
seen.append(request)
|
||||
if len(seen) > len(steps):
|
||||
raise AssertionError("unexpected request")
|
||||
return steps[len(seen) - 1](request)
|
||||
|
||||
return handler, seen
|
||||
|
||||
|
||||
def _ok(stream_cls, chunks, *, drop=False, location=None):
|
||||
headers = dict(SSE)
|
||||
if location:
|
||||
headers["Location"] = location
|
||||
|
||||
def step(_request):
|
||||
return httpx.Response(
|
||||
200, headers=headers, stream=stream_cls(chunks, DROP if drop else None)
|
||||
)
|
||||
|
||||
return step
|
||||
|
||||
|
||||
def _status(code):
|
||||
return lambda request: httpx.Response(
|
||||
code, json={"detail": "upstream"}, request=request
|
||||
)
|
||||
|
||||
|
||||
def _connect_error(request):
|
||||
raise httpx.ConnectError("[Errno 61] Connection refused", request=request)
|
||||
|
||||
|
||||
def _not_sse(request):
|
||||
return httpx.Response(
|
||||
200, headers={"Content-Type": "text/html"}, text="<html>", request=request
|
||||
)
|
||||
|
||||
|
||||
def _transient_steps(stream_cls, failure):
|
||||
"""POST drops after event 1; first reconnect fails; second succeeds."""
|
||||
return [
|
||||
_ok(stream_cls, [_event(1)], drop=True, location="/reconnect"),
|
||||
failure,
|
||||
_ok(stream_cls, [_event(2), b"event: end\ndata: null\n\n"]),
|
||||
]
|
||||
|
||||
|
||||
def _long_stream_steps(stream_cls, drops):
|
||||
"""POST + `drops` reconnects, each delivering one event then dropping."""
|
||||
steps = [_ok(stream_cls, [_event(0)], drop=True, location="/reconnect")]
|
||||
for i in range(1, drops):
|
||||
steps.append(_ok(stream_cls, [_event(i)], drop=True))
|
||||
steps.append(_ok(stream_cls, [_event(drops), b"event: end\ndata: null\n\n"]))
|
||||
return steps
|
||||
|
||||
|
||||
def _run_sync(steps):
|
||||
handler, seen = _scripted(steps)
|
||||
with httpx.Client(
|
||||
transport=httpx.MockTransport(handler), base_url="https://example.com"
|
||||
) as client:
|
||||
parts = list(SyncHttpClient(client).stream("/stream", "POST", json={}))
|
||||
return parts, seen
|
||||
|
||||
|
||||
async def _run_async(steps):
|
||||
sync_handler, seen = _scripted(steps)
|
||||
|
||||
async def handler(request):
|
||||
return sync_handler(request)
|
||||
|
||||
async with httpx.AsyncClient(
|
||||
transport=httpx.MockTransport(handler), base_url="https://example.com"
|
||||
) as client:
|
||||
parts = [p async for p in HttpClient(client).stream("/stream", "POST", json={})]
|
||||
return parts, seen
|
||||
|
||||
|
||||
FAILURES = {
|
||||
"429": _status(429),
|
||||
"502": _status(502),
|
||||
"503": _status(503),
|
||||
"connect_error": _connect_error,
|
||||
"not_sse": _not_sse,
|
||||
}
|
||||
|
||||
|
||||
# A reconnect that hits a 429 / 502 / 503 / connect error / non-SSE body, then
|
||||
# succeeds, should recover.
|
||||
|
||||
|
||||
@pytest.mark.usefixtures("no_sleep")
|
||||
@pytest.mark.parametrize("failure", FAILURES.values(), ids=FAILURES.keys())
|
||||
def test_sync_reconnect_recovers_from_transient_failure(failure):
|
||||
parts, seen = _run_sync(_transient_steps(ListByteStream, failure))
|
||||
assert [p.event for p in parts] == ["values", "values", "end"]
|
||||
assert len(seen) == 3
|
||||
|
||||
|
||||
@pytest.mark.usefixtures("no_sleep")
|
||||
@pytest.mark.parametrize("failure", FAILURES.values(), ids=FAILURES.keys())
|
||||
async def test_async_reconnect_recovers_from_transient_failure(failure):
|
||||
parts, seen = await _run_async(_transient_steps(AsyncListByteStream, failure))
|
||||
assert [p.event for p in parts] == ["values", "values", "end"]
|
||||
assert len(seen) == 3
|
||||
|
||||
|
||||
# Any other 4xx on a reconnect won't go away by retrying, so it raises.
|
||||
|
||||
|
||||
def _client_error_steps(stream_cls):
|
||||
return [
|
||||
_ok(stream_cls, [_event(1)], drop=True, location="/reconnect"),
|
||||
_status(404),
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.usefixtures("no_sleep")
|
||||
def test_sync_reconnect_does_not_retry_a_client_error():
|
||||
with pytest.raises(NotFoundError):
|
||||
_run_sync(_client_error_steps(ListByteStream))
|
||||
|
||||
|
||||
@pytest.mark.usefixtures("no_sleep")
|
||||
async def test_async_reconnect_does_not_retry_a_client_error():
|
||||
with pytest.raises(NotFoundError):
|
||||
await _run_async(_client_error_steps(AsyncListByteStream))
|
||||
|
||||
|
||||
# The first request isn't a reconnect, so it raises even when it names a Location.
|
||||
|
||||
|
||||
def _first_request_503(request):
|
||||
return httpx.Response(
|
||||
503,
|
||||
headers={"Location": "/reconnect"},
|
||||
json={"detail": "upstream"},
|
||||
request=request,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.usefixtures("no_sleep")
|
||||
def test_sync_first_request_failure_is_not_retried():
|
||||
with pytest.raises(InternalServerError):
|
||||
_run_sync([_first_request_503])
|
||||
|
||||
|
||||
@pytest.mark.usefixtures("no_sleep")
|
||||
async def test_async_first_request_failure_is_not_retried():
|
||||
with pytest.raises(InternalServerError):
|
||||
await _run_async([_first_request_503])
|
||||
|
||||
|
||||
# A server that stays down gets five reconnects, backing off 1, 2, 4, 5 and 5s
|
||||
# (each plus up to 25% jitter), so about 20s, before the stream gives up. The
|
||||
# error it gives up with is caused by the last failure, so callers see why.
|
||||
|
||||
GIVE_UP = "Exceeded maximum SSE reconnection attempts"
|
||||
|
||||
|
||||
def _server_down_steps(stream_cls):
|
||||
return [
|
||||
_ok(stream_cls, [_event(1)], drop=True, location="/reconnect"),
|
||||
*[_status(502)] * 5,
|
||||
]
|
||||
|
||||
|
||||
def _assert_backed_off_for_about_20s(sleeps: list[float]) -> None:
|
||||
assert len(sleeps) == 5
|
||||
for slept, delay in zip(sleeps, [1, 2, 4, 5, 5], strict=True):
|
||||
assert delay <= slept <= delay * 1.25
|
||||
|
||||
|
||||
def test_sync_reconnect_gives_up_after_five_backed_off_attempts(no_sleep):
|
||||
with pytest.raises(httpx.TransportError, match=GIVE_UP) as exc_info:
|
||||
_run_sync(_server_down_steps(ListByteStream))
|
||||
assert isinstance(exc_info.value.__cause__, InternalServerError)
|
||||
_assert_backed_off_for_about_20s(no_sleep)
|
||||
|
||||
|
||||
async def test_async_reconnect_gives_up_after_five_backed_off_attempts(no_sleep):
|
||||
with pytest.raises(httpx.TransportError, match=GIVE_UP) as exc_info:
|
||||
await _run_async(_server_down_steps(AsyncListByteStream))
|
||||
assert isinstance(exc_info.value.__cause__, InternalServerError)
|
||||
_assert_backed_off_for_about_20s(no_sleep)
|
||||
|
||||
|
||||
# A long stream that drops 6 times, each reconnect succeeding and making
|
||||
# progress, should not run out of attempts.
|
||||
|
||||
|
||||
@pytest.mark.usefixtures("no_sleep")
|
||||
def test_sync_long_stream_survives_many_spaced_out_drops():
|
||||
parts, _ = _run_sync(_long_stream_steps(ListByteStream, drops=6))
|
||||
assert [p.data for p in parts] == [0, 1, 2, 3, 4, 5, 6, None]
|
||||
|
||||
|
||||
@pytest.mark.usefixtures("no_sleep")
|
||||
async def test_async_long_stream_survives_many_spaced_out_drops():
|
||||
parts, _ = await _run_async(_long_stream_steps(AsyncListByteStream, drops=6))
|
||||
assert [p.data for p in parts] == [0, 1, 2, 3, 4, 5, 6, None]
|
||||
|
||||
|
||||
# Reconnects should back off rather than hammer the server.
|
||||
|
||||
|
||||
def test_sync_reconnect_backs_off(no_sleep):
|
||||
_run_sync(_long_stream_steps(ListByteStream, drops=2))
|
||||
assert len(no_sleep) == 2 and all(d > 0 for d in no_sleep)
|
||||
|
||||
|
||||
async def test_async_reconnect_backs_off(no_sleep):
|
||||
await _run_async(_long_stream_steps(AsyncListByteStream, drops=2))
|
||||
assert len(no_sleep) == 2 and all(d > 0 for d in no_sleep)
|
||||
|
||||
|
||||
# --- _sse_to_v2_dict conversion ---
|
||||
|
||||
|
||||
|
||||
Reference in new issue
Block a user