mirror of
https://github.com/langchain-ai/langgraph.git
synced 2026-11-03 04:41:21 +01:00
Compare commits
9
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
d9abeabfcd | ||
|
|
5bd231cbfd | ||
|
|
3eaf06c67c | ||
|
|
6fab323753 | ||
|
|
8d5637b2dd | ||
|
|
1f0cf816d5 | ||
|
|
6f2df053f4 | ||
|
|
ffc3336367 | ||
|
|
5f4b28a9f7 |
No files matched your search
@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
|
||||
|
||||
[project]
|
||||
name = "langgraph-checkpoint-sqlite"
|
||||
version = "3.1.2"
|
||||
version = "3.1.3"
|
||||
description = "Library with a SQLite implementation of LangGraph checkpoint saver."
|
||||
authors = []
|
||||
requires-python = ">=3.10"
|
||||
@@ -12,7 +12,7 @@ readme = "README.md"
|
||||
license = "MIT"
|
||||
license-files = ['LICENSE']
|
||||
dependencies = [
|
||||
"langgraph-checkpoint>=4.3.0,<5.0.0",
|
||||
"langgraph-checkpoint>=4.3.1,<5.0.0",
|
||||
"aiosqlite>=0.20",
|
||||
"sqlite-vec>=0.1.6",
|
||||
]
|
||||
|
||||
Generated
+1
-1
@@ -360,7 +360,7 @@ test = [
|
||||
|
||||
[[package]]
|
||||
name = "langgraph-checkpoint-sqlite"
|
||||
version = "3.1.2"
|
||||
version = "3.1.3"
|
||||
source = { editable = "." }
|
||||
dependencies = [
|
||||
{ name = "aiosqlite" },
|
||||
|
||||
@@ -86,47 +86,6 @@ 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:
|
||||
|
||||
@@ -13,7 +13,6 @@ 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
|
||||
@@ -2636,109 +2635,6 @@ 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,18 +362,6 @@ 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:
|
||||
|
||||
@@ -1,130 +0,0 @@
|
||||
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
|
||||
@@ -1,6 +1,7 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import copy
|
||||
import enum
|
||||
import inspect
|
||||
import logging
|
||||
@@ -23,9 +24,16 @@ from typing import (
|
||||
Optional,
|
||||
Protocol,
|
||||
TypeGuard,
|
||||
TypeVar,
|
||||
cast,
|
||||
)
|
||||
|
||||
from langchain_core.callbacks.base import BaseCallbackHandler, BaseCallbackManager
|
||||
from langchain_core.callbacks.manager import (
|
||||
AsyncCallbackManagerForChainRun,
|
||||
BaseRunManager,
|
||||
CallbackManagerForChainRun,
|
||||
)
|
||||
from langchain_core.runnables.base import (
|
||||
Runnable,
|
||||
RunnableConfig,
|
||||
@@ -41,6 +49,7 @@ from langchain_core.runnables.config import (
|
||||
var_child_runnable_config,
|
||||
)
|
||||
from langchain_core.runnables.utils import Input, Output
|
||||
from langchain_core.tracers.base import BaseTracer
|
||||
from langchain_core.tracers.langchain import LangChainTracer
|
||||
from langgraph.store.base import BaseStore
|
||||
|
||||
@@ -57,7 +66,7 @@ from langgraph._internal._constants import (
|
||||
)
|
||||
from langgraph._internal._typing import MISSING
|
||||
from langgraph.errors import NodeError
|
||||
from langgraph.types import StreamWriter
|
||||
from langgraph.types import StreamWriter, TracePolicy
|
||||
|
||||
try:
|
||||
from langchain_core.tracers._streaming import _StreamingCallbackHandler
|
||||
@@ -85,6 +94,29 @@ def _trace_payload(value: Any, transform: Callable[[Any], Any] | None) -> Any:
|
||||
return value
|
||||
|
||||
|
||||
def _is_stream_handler(handler: BaseCallbackHandler) -> bool:
|
||||
# Inheritance, not isinstance on the runtime-checkable protocol, which also
|
||||
# matches any exporter that merely has the tap methods: stream handlers get a
|
||||
# node's real payload and its streamed output.
|
||||
return _StreamingCallbackHandler in type(handler).__mro__
|
||||
|
||||
|
||||
_M = TypeVar("_M", bound=BaseCallbackManager | BaseRunManager)
|
||||
|
||||
|
||||
def _split_handlers(manager: _M) -> tuple[_M, _M]:
|
||||
# The recorded payload exists to keep tracer ingest small, so only tracers get it.
|
||||
# `astream_log`'s handler is a tracer too, but it feeds the caller's own stream.
|
||||
real, tracers = copy.copy(manager), copy.copy(manager)
|
||||
tracers.handlers = [
|
||||
h
|
||||
for h in manager.handlers
|
||||
if isinstance(h, BaseTracer) and not _is_stream_handler(h)
|
||||
]
|
||||
real.handlers = [h for h in manager.handlers if h not in tracers.handlers]
|
||||
return real, tracers
|
||||
|
||||
|
||||
def _set_config_context(
|
||||
config: RunnableConfig, run: Any = None
|
||||
) -> Token[RunnableConfig | None]:
|
||||
@@ -593,13 +625,15 @@ class RunnableSeq(Runnable):
|
||||
*steps: RunnableLike,
|
||||
name: str | None = None,
|
||||
trace_inputs: Callable[[Any], Any] | None = None,
|
||||
trace_outputs: Callable[[Any], Any] | None = None,
|
||||
trace_policy: TracePolicy | None = None,
|
||||
) -> None:
|
||||
"""Create a new RunnableSeq.
|
||||
|
||||
Args:
|
||||
steps: The steps to include in the sequence.
|
||||
name: The name of the `Runnable`.
|
||||
trace_inputs: Transforms the input every callback handler receives.
|
||||
trace_policy: Processes what tracers record, after `trace_inputs`.
|
||||
|
||||
Raises:
|
||||
ValueError: If the sequence has less than 2 steps.
|
||||
@@ -619,7 +653,67 @@ class RunnableSeq(Runnable):
|
||||
self.steps = steps_flat
|
||||
self.name = name
|
||||
self.trace_inputs = trace_inputs
|
||||
self.trace_outputs = trace_outputs
|
||||
self.trace_policy = trace_policy or TracePolicy()
|
||||
|
||||
def _start(self, input: Any, config: RunnableConfig) -> CallbackManagerForChainRun:
|
||||
callback_manager = get_callback_manager_for_config(config)
|
||||
payload = _trace_payload(input, self.trace_inputs)
|
||||
recorded = _trace_payload(payload, self.trace_policy.process_inputs)
|
||||
name = config.get("run_name") or self.get_name()
|
||||
run_id = config.pop("run_id", None)
|
||||
if recorded is payload:
|
||||
return callback_manager.on_chain_start(
|
||||
None, payload, name=name, run_id=run_id
|
||||
)
|
||||
real, tracers = _split_handlers(callback_manager)
|
||||
run_manager = real.on_chain_start(None, payload, name=name, run_id=run_id)
|
||||
tracers.on_chain_start(None, recorded, name=name, run_id=run_manager.run_id)
|
||||
run_manager.handlers = callback_manager.handlers
|
||||
return run_manager
|
||||
|
||||
async def _astart(
|
||||
self, input: Any, config: RunnableConfig
|
||||
) -> AsyncCallbackManagerForChainRun:
|
||||
callback_manager = get_async_callback_manager_for_config(config)
|
||||
payload = _trace_payload(input, self.trace_inputs)
|
||||
recorded = _trace_payload(payload, self.trace_policy.process_inputs)
|
||||
name = config.get("run_name") or self.get_name()
|
||||
run_id = config.pop("run_id", None)
|
||||
if recorded is payload:
|
||||
return await callback_manager.on_chain_start(
|
||||
None, payload, name=name, run_id=run_id
|
||||
)
|
||||
real, tracers = _split_handlers(callback_manager)
|
||||
# Tracers start last, so a cancellation before then leaves no tracer run open.
|
||||
run_manager = await real.on_chain_start(None, payload, name=name, run_id=run_id)
|
||||
await tracers.on_chain_start(
|
||||
None, recorded, name=name, run_id=run_manager.run_id
|
||||
)
|
||||
run_manager.handlers = callback_manager.handlers
|
||||
return run_manager
|
||||
|
||||
def _end(self, run_manager: CallbackManagerForChainRun, output: Any) -> None:
|
||||
recorded = _trace_payload(output, self.trace_policy.process_outputs)
|
||||
if recorded is output:
|
||||
run_manager.on_chain_end(output)
|
||||
else:
|
||||
real, tracers = _split_handlers(run_manager)
|
||||
tracers.on_chain_end(recorded)
|
||||
real.on_chain_end(output)
|
||||
|
||||
async def _aend(
|
||||
self, run_manager: AsyncCallbackManagerForChainRun, output: Any
|
||||
) -> None:
|
||||
recorded = _trace_payload(output, self.trace_policy.process_outputs)
|
||||
if recorded is output:
|
||||
await run_manager.on_chain_end(output)
|
||||
else:
|
||||
real, tracers = _split_handlers(run_manager)
|
||||
# Tracers end first: on_chain_end is shielded once it starts, so a
|
||||
# cancellation can only skip the other handlers' end, never leave a
|
||||
# tracer run open.
|
||||
await tracers.on_chain_end(recorded)
|
||||
await real.on_chain_end(output)
|
||||
|
||||
def __or__(
|
||||
self,
|
||||
@@ -676,15 +770,8 @@ class RunnableSeq(Runnable):
|
||||
) -> Any:
|
||||
if config is None:
|
||||
config = ensure_config()
|
||||
# setup callbacks and context
|
||||
callback_manager = get_callback_manager_for_config(config)
|
||||
# start the root run
|
||||
run_manager = callback_manager.on_chain_start(
|
||||
None,
|
||||
_trace_payload(input, self.trace_inputs),
|
||||
name=config.get("run_name") or self.get_name(),
|
||||
run_id=config.pop("run_id", None),
|
||||
)
|
||||
run_manager = self._start(input, config)
|
||||
# invoke all steps in sequence
|
||||
try:
|
||||
for i, step in enumerate(self.steps):
|
||||
@@ -712,7 +799,7 @@ class RunnableSeq(Runnable):
|
||||
run_manager.on_chain_error(e)
|
||||
raise
|
||||
else:
|
||||
run_manager.on_chain_end(_trace_payload(input, self.trace_outputs))
|
||||
self._end(run_manager, input)
|
||||
return input
|
||||
|
||||
async def ainvoke(
|
||||
@@ -723,15 +810,8 @@ class RunnableSeq(Runnable):
|
||||
) -> Any:
|
||||
if config is None:
|
||||
config = ensure_config()
|
||||
# setup callbacks
|
||||
callback_manager = get_async_callback_manager_for_config(config)
|
||||
# start the root run
|
||||
run_manager = await callback_manager.on_chain_start(
|
||||
None,
|
||||
_trace_payload(input, self.trace_inputs),
|
||||
name=config.get("run_name") or self.get_name(),
|
||||
run_id=config.pop("run_id", None),
|
||||
)
|
||||
run_manager = await self._astart(input, config)
|
||||
|
||||
# invoke all steps in sequence
|
||||
try:
|
||||
@@ -765,7 +845,7 @@ class RunnableSeq(Runnable):
|
||||
await run_manager.on_chain_error(e)
|
||||
raise
|
||||
else:
|
||||
await run_manager.on_chain_end(_trace_payload(input, self.trace_outputs))
|
||||
await self._aend(run_manager, input)
|
||||
return input
|
||||
|
||||
def stream(
|
||||
@@ -776,15 +856,8 @@ class RunnableSeq(Runnable):
|
||||
) -> Iterator[Any]:
|
||||
if config is None:
|
||||
config = ensure_config()
|
||||
# setup callbacks
|
||||
callback_manager = get_callback_manager_for_config(config)
|
||||
# start the root run
|
||||
run_manager = callback_manager.on_chain_start(
|
||||
None,
|
||||
_trace_payload(input, self.trace_inputs),
|
||||
name=config.get("run_name") or self.get_name(),
|
||||
run_id=config.pop("run_id", None),
|
||||
)
|
||||
run_manager = self._start(input, config)
|
||||
# get the run object
|
||||
for h in run_manager.handlers:
|
||||
if isinstance(h, LangChainTracer):
|
||||
@@ -814,10 +887,9 @@ class RunnableSeq(Runnable):
|
||||
)
|
||||
iterator = step.transform(iterator, config)
|
||||
# populates streamed_output in astream_log() output if needed
|
||||
if _StreamingCallbackHandler is not None:
|
||||
for h in run_manager.handlers:
|
||||
if isinstance(h, _StreamingCallbackHandler):
|
||||
iterator = h.tap_output_iter(run_manager.run_id, iterator)
|
||||
for h in run_manager.handlers:
|
||||
if _is_stream_handler(h):
|
||||
iterator = h.tap_output_iter(run_manager.run_id, iterator)
|
||||
# consume into final output
|
||||
output = context.run(_consume_iter, iterator)
|
||||
# sequence doesn't emit output, yield to mark as generator
|
||||
@@ -826,7 +898,7 @@ class RunnableSeq(Runnable):
|
||||
run_manager.on_chain_error(e)
|
||||
raise
|
||||
else:
|
||||
run_manager.on_chain_end(_trace_payload(output, self.trace_outputs))
|
||||
self._end(run_manager, output)
|
||||
|
||||
async def astream(
|
||||
self,
|
||||
@@ -836,15 +908,8 @@ class RunnableSeq(Runnable):
|
||||
) -> AsyncIterator[Any]:
|
||||
if config is None:
|
||||
config = ensure_config()
|
||||
# setup callbacks
|
||||
callback_manager = get_async_callback_manager_for_config(config)
|
||||
# start the root run
|
||||
run_manager = await callback_manager.on_chain_start(
|
||||
None,
|
||||
_trace_payload(input, self.trace_inputs),
|
||||
name=config.get("run_name") or self.get_name(),
|
||||
run_id=config.pop("run_id", None),
|
||||
)
|
||||
run_manager = await self._astart(input, config)
|
||||
# stream the last steps
|
||||
# transform the input stream of each step with the next
|
||||
# steps that don't natively support transforming an input stream will
|
||||
@@ -880,12 +945,11 @@ class RunnableSeq(Runnable):
|
||||
if hasattr(aiterator, "aclose"):
|
||||
stack.push_async_callback(aiterator.aclose)
|
||||
# populates streamed_output in astream_log() output if needed
|
||||
if _StreamingCallbackHandler is not None:
|
||||
for h in run_manager.handlers:
|
||||
if isinstance(h, _StreamingCallbackHandler):
|
||||
aiterator = h.tap_output_aiter(
|
||||
run_manager.run_id, aiterator
|
||||
)
|
||||
for h in run_manager.handlers:
|
||||
if _is_stream_handler(h):
|
||||
aiterator = h.tap_output_aiter(
|
||||
run_manager.run_id, aiterator
|
||||
)
|
||||
# consume into final output
|
||||
output = await asyncio.create_task(
|
||||
_consume_aiter(aiterator), context=context
|
||||
@@ -896,9 +960,7 @@ class RunnableSeq(Runnable):
|
||||
await run_manager.on_chain_error(e)
|
||||
raise
|
||||
else:
|
||||
await run_manager.on_chain_end(
|
||||
_trace_payload(output, self.trace_outputs)
|
||||
)
|
||||
await self._aend(run_manager, output)
|
||||
else:
|
||||
try:
|
||||
async with AsyncExitStack() as stack:
|
||||
@@ -914,12 +976,11 @@ class RunnableSeq(Runnable):
|
||||
if hasattr(aiterator, "aclose"):
|
||||
stack.push_async_callback(aiterator.aclose)
|
||||
# populates streamed_output in astream_log() output if needed
|
||||
if _StreamingCallbackHandler is not None:
|
||||
for h in run_manager.handlers:
|
||||
if isinstance(h, _StreamingCallbackHandler):
|
||||
aiterator = h.tap_output_aiter(
|
||||
run_manager.run_id, aiterator
|
||||
)
|
||||
for h in run_manager.handlers:
|
||||
if _is_stream_handler(h):
|
||||
aiterator = h.tap_output_aiter(
|
||||
run_manager.run_id, aiterator
|
||||
)
|
||||
# consume into final output
|
||||
output = await _consume_aiter(aiterator)
|
||||
# sequence doesn't emit output, yield to mark as generator
|
||||
@@ -928,9 +989,7 @@ class RunnableSeq(Runnable):
|
||||
await run_manager.on_chain_error(e)
|
||||
raise
|
||||
else:
|
||||
await run_manager.on_chain_end(
|
||||
_trace_payload(output, self.trace_outputs)
|
||||
)
|
||||
await self._aend(run_manager, output)
|
||||
|
||||
|
||||
def _consume_iter(it: Iterator[Any]) -> Any:
|
||||
|
||||
@@ -22,7 +22,7 @@ from langchain_core.messages import (
|
||||
)
|
||||
from typing_extensions import TypedDict, deprecated
|
||||
|
||||
from langgraph._internal._constants import CONF, CONFIG_KEY_SEND, NS_SEP
|
||||
from langgraph._internal._constants import CONF, CONFIG_KEY_SEND
|
||||
from langgraph.graph.state import StateGraph
|
||||
from langgraph.pregel._messages import (
|
||||
_messages_delta_reducer as _messages_delta_reducer,
|
||||
@@ -343,7 +343,7 @@ def push_message(
|
||||
)
|
||||
|
||||
from langgraph.config import get_config
|
||||
from langgraph.pregel._messages import StreamMessagesHandler
|
||||
from langgraph.pregel._messages import StreamMessagesHandler, graph_namespace
|
||||
|
||||
config = get_config()
|
||||
message = next(x for x in convert_to_messages([message]))
|
||||
@@ -359,15 +359,12 @@ def push_message(
|
||||
):
|
||||
handlers = config["callbacks"]
|
||||
|
||||
if stream_handler := next(
|
||||
(x for x in handlers if isinstance(x, StreamMessagesHandler)), None
|
||||
):
|
||||
metadata = config["metadata"]
|
||||
message_meta = (
|
||||
tuple(cast(str, metadata["langgraph_checkpoint_ns"]).split(NS_SEP)),
|
||||
metadata,
|
||||
)
|
||||
stream_handler._emit(message_meta, message, dedupe=False)
|
||||
metadata = config["metadata"]
|
||||
for handler in handlers:
|
||||
if isinstance(handler, StreamMessagesHandler):
|
||||
ns = graph_namespace(metadata)
|
||||
if handler.subgraphs or not ns or ns == handler.parent_ns:
|
||||
handler._emit((ns, metadata), message, dedupe=False)
|
||||
|
||||
if state_key:
|
||||
config[CONF][CONFIG_KEY_SEND]([(state_key, message)])
|
||||
|
||||
@@ -132,7 +132,7 @@ from langgraph.pregel.debug import (
|
||||
map_debug_tasks,
|
||||
)
|
||||
from langgraph.pregel.protocol import StreamChunk, StreamProtocol
|
||||
from langgraph.runtime import RunControl, Runtime
|
||||
from langgraph.runtime import RunControl, Runtime, _RunExits
|
||||
from langgraph.types import (
|
||||
All,
|
||||
CachePolicy,
|
||||
@@ -1497,6 +1497,10 @@ class PregelLoop:
|
||||
if self._delta_write_futs is not None:
|
||||
self._delta_write_futs.append(fut)
|
||||
|
||||
def _run_exits(self) -> _RunExits | None:
|
||||
runtime = self.config[CONF].get(CONFIG_KEY_RUNTIME)
|
||||
return runtime._run_exits if isinstance(runtime, Runtime) else None
|
||||
|
||||
def _suppress_interrupt(
|
||||
self,
|
||||
exc_type: type[BaseException] | None,
|
||||
@@ -1874,6 +1878,10 @@ class SyncPregelLoop(PregelLoop, AbstractContextManager):
|
||||
self._exit_delta_writes = (
|
||||
[] if self.durability == "exit" and self.checkpointer is not None else None
|
||||
)
|
||||
# Bottom of the stack, so run-exit callbacks run last: after the
|
||||
# executor has stopped every task and the final checkpoint is saved.
|
||||
if (exits := self._run_exits()) is not None and exits.bind(is_async=False):
|
||||
self.stack.push(exits.exit)
|
||||
self.submit = self.stack.enter_context(BackgroundExecutor(self.config))
|
||||
self.channels, self.managed = channels_from_checkpoint(
|
||||
self.specs,
|
||||
@@ -2130,6 +2138,10 @@ class AsyncPregelLoop(PregelLoop, AbstractAsyncContextManager):
|
||||
self._exit_delta_writes = (
|
||||
[] if self.durability == "exit" and self.checkpointer is not None else None
|
||||
)
|
||||
# Bottom of the stack, so run-exit callbacks run last: after the
|
||||
# executor has stopped every task and the final checkpoint is saved.
|
||||
if (exits := self._run_exits()) is not None and exits.bind(is_async=True):
|
||||
self.stack.push_async_exit(exits.aexit)
|
||||
self.submit = await self.stack.enter_async_context(
|
||||
AsyncBackgroundExecutor(self.config)
|
||||
)
|
||||
|
||||
@@ -23,7 +23,7 @@ from langchain_core.outputs import ChatGeneration, ChatGenerationChunk, LLMResul
|
||||
from pydantic import BaseModel
|
||||
|
||||
from langgraph._internal._config import filter_to_user_tags
|
||||
from langgraph._internal._constants import NS_SEP
|
||||
from langgraph._internal._constants import NS_SEP, PUSH
|
||||
from langgraph.channels.binop import _get_overwrite
|
||||
from langgraph.channels.delta import DeltaChannel
|
||||
from langgraph.constants import TAG_HIDDEN, TAG_NOSTREAM
|
||||
@@ -55,6 +55,23 @@ def _state_values(obj: Any) -> Sequence[Any]:
|
||||
return ()
|
||||
|
||||
|
||||
def graph_namespace(metadata: Mapping[str, Any]) -> tuple[str, ...]:
|
||||
"""Namespace of the graph a task runs in, from the task's run metadata.
|
||||
|
||||
Drops the task's own segment, plus one per enclosing `@task` call: a call
|
||||
runs in its caller's graph but nests its segment under the caller's. A
|
||||
call's `langgraph_path` is `(PUSH, <caller's path>, <index>, True)`, while
|
||||
a `Send` task's ends in `False`.
|
||||
"""
|
||||
ns = tuple(cast(str, metadata["langgraph_checkpoint_ns"]).split(NS_SEP))
|
||||
path = metadata.get("langgraph_path", ())
|
||||
drop = 1
|
||||
while len(path) == 4 and path[0] == PUSH and path[-1] is True:
|
||||
drop += 1
|
||||
path = path[1]
|
||||
return ns[:-drop]
|
||||
|
||||
|
||||
class StreamMessagesHandler(BaseCallbackHandler, _StreamingCallbackHandler):
|
||||
"""A callback handler that implements stream_mode=messages.
|
||||
|
||||
@@ -148,9 +165,7 @@ class StreamMessagesHandler(BaseCallbackHandler, _StreamingCallbackHandler):
|
||||
**kwargs: Any,
|
||||
) -> Any:
|
||||
if metadata and (not tags or (TAG_NOSTREAM not in tags)):
|
||||
ns = tuple(cast(str, metadata["langgraph_checkpoint_ns"]).split(NS_SEP))[
|
||||
:-1
|
||||
]
|
||||
ns = graph_namespace(metadata)
|
||||
if not self.subgraphs and len(ns) > 0 and ns != self.parent_ns:
|
||||
return
|
||||
if (filtered_tags := filter_to_user_tags(tags)) is not None:
|
||||
@@ -213,9 +228,7 @@ class StreamMessagesHandler(BaseCallbackHandler, _StreamingCallbackHandler):
|
||||
and kwargs.get("name") == metadata.get("langgraph_node")
|
||||
and (not tags or TAG_HIDDEN not in tags)
|
||||
):
|
||||
ns = tuple(cast(str, metadata["langgraph_checkpoint_ns"]).split(NS_SEP))[
|
||||
:-1
|
||||
]
|
||||
ns = graph_namespace(metadata)
|
||||
if not self.subgraphs and len(ns) > 0 and ns != self.parent_ns:
|
||||
return
|
||||
self.metadata[run_id] = (ns, metadata)
|
||||
|
||||
@@ -227,23 +227,14 @@ class PregelNode:
|
||||
def node(self) -> Runnable[Any, Any] | None:
|
||||
"""Get a runnable that combines `bound` and `writers`."""
|
||||
writers = self.flat_writers
|
||||
trace_inputs = self.trace_policy.process_inputs if self.trace_policy else None
|
||||
trace_outputs = self.trace_policy.process_outputs if self.trace_policy else None
|
||||
if self.bound is DEFAULT_BOUND and not writers:
|
||||
return None
|
||||
elif self.bound is DEFAULT_BOUND and len(writers) == 1:
|
||||
return writers[0]
|
||||
elif self.bound is DEFAULT_BOUND:
|
||||
return RunnableSeq(
|
||||
*writers, trace_inputs=trace_inputs, trace_outputs=trace_outputs
|
||||
)
|
||||
return RunnableSeq(*writers, trace_policy=self.trace_policy)
|
||||
elif writers:
|
||||
return RunnableSeq(
|
||||
self.bound,
|
||||
*writers,
|
||||
trace_inputs=trace_inputs,
|
||||
trace_outputs=trace_outputs,
|
||||
)
|
||||
return RunnableSeq(self.bound, *writers, trace_policy=self.trace_policy)
|
||||
else:
|
||||
return self.bound
|
||||
|
||||
|
||||
@@ -7,9 +7,9 @@ from uuid import UUID
|
||||
|
||||
from langchain_core.callbacks import BaseCallbackHandler
|
||||
|
||||
from langgraph._internal._constants import NS_SEP
|
||||
from langgraph.constants import TAG_NOSTREAM
|
||||
from langgraph.errors import GraphDrained, GraphInterrupt
|
||||
from langgraph.pregel._messages import graph_namespace
|
||||
from langgraph.pregel.protocol import StreamChunk
|
||||
|
||||
try:
|
||||
@@ -92,12 +92,9 @@ class StreamToolCallHandler(BaseCallbackHandler, _StreamingCallbackHandler):
|
||||
) -> tuple[str, ...] | None:
|
||||
"""Resolve the namespace this tool call should emit at, or `None` to skip.
|
||||
|
||||
Mirrors `StreamMessagesHandler.on_chat_model_start`'s namespace
|
||||
derivation: parses `langgraph_checkpoint_ns` (which ends with
|
||||
the `node_name:task_id` of the calling node), drops that
|
||||
trailing segment, and returns the containing subgraph's own
|
||||
namespace. Returns `None` when the call should be silently
|
||||
suppressed:
|
||||
Mirrors `StreamMessagesHandler.on_chat_model_start`: the namespace
|
||||
is that of the graph the tool runs in (`graph_namespace`). Returns
|
||||
`None` when the call should be silently suppressed:
|
||||
|
||||
- `metadata` is missing — handler is attached to a context
|
||||
without Pregel routing info.
|
||||
@@ -110,11 +107,9 @@ class StreamToolCallHandler(BaseCallbackHandler, _StreamingCallbackHandler):
|
||||
return None
|
||||
if tags and TAG_NOSTREAM in tags:
|
||||
return None
|
||||
nskey = metadata.get("langgraph_checkpoint_ns")
|
||||
if not nskey:
|
||||
ns: tuple[str, ...] = ()
|
||||
else:
|
||||
ns = tuple(cast(str, nskey).split(NS_SEP))[:-1]
|
||||
ns = (
|
||||
graph_namespace(metadata) if metadata.get("langgraph_checkpoint_ns") else ()
|
||||
)
|
||||
if not self.subgraphs and len(ns) > 0 and ns != self.parent_ns:
|
||||
return None
|
||||
return ns
|
||||
|
||||
@@ -97,6 +97,7 @@ from langgraph._internal._runnable import (
|
||||
Runnable,
|
||||
RunnableLike,
|
||||
RunnableSeq,
|
||||
_is_stream_handler,
|
||||
coerce_to_runnable,
|
||||
)
|
||||
from langgraph._internal._timeout import coerce_timeout_policy
|
||||
@@ -174,6 +175,7 @@ from langgraph.runtime import (
|
||||
RunControl,
|
||||
Runtime,
|
||||
ServerInfo,
|
||||
_RunExits,
|
||||
)
|
||||
from langgraph.stream._mux import StreamMux
|
||||
from langgraph.stream._types import StreamTransformer
|
||||
@@ -203,11 +205,6 @@ from langgraph.types import (
|
||||
from langgraph.typing import ContextT, InputT, OutputT, StateT
|
||||
from langgraph.warnings import LangGraphDeprecatedSinceV10
|
||||
|
||||
try:
|
||||
from langchain_core.tracers._streaming import _StreamingCallbackHandler
|
||||
except ImportError:
|
||||
_StreamingCallbackHandler = None # type: ignore
|
||||
|
||||
__all__ = ("NodeBuilder", "Pregel")
|
||||
|
||||
_WriteValue = Callable[[Input], Output] | Any
|
||||
@@ -2880,6 +2877,20 @@ class Pregel(
|
||||
for h in callback_manager.inheritable_handlers
|
||||
if not isinstance(h, StreamMessagesHandlerV2)
|
||||
]
|
||||
if CONFIG_KEY_TASK_ID not in config[CONF]:
|
||||
# Not run as a task (a child invoked with its own thread, say),
|
||||
# this graph's namespace starts at the root, so a caller's
|
||||
# handler would take this graph's events for its own graph's.
|
||||
callback_manager.handlers = [
|
||||
h
|
||||
for h in callback_manager.handlers
|
||||
if not _is_stream_mode_handler(h)
|
||||
]
|
||||
callback_manager.inheritable_handlers = [
|
||||
h
|
||||
for h in callback_manager.inheritable_handlers
|
||||
if not _is_stream_mode_handler(h)
|
||||
]
|
||||
if "ls_integration" not in callback_manager.metadata:
|
||||
callback_manager.add_metadata({"ls_integration": "langgraph"})
|
||||
run_manager = callback_manager.on_chain_start(
|
||||
@@ -2898,9 +2909,9 @@ class Pregel(
|
||||
)
|
||||
# set up subgraph checkpointing
|
||||
config = self._own_checkpoint_config(config)
|
||||
stream_handlers: list[StreamMessagesHandler | StreamToolCallHandler] = []
|
||||
# set up messages stream mode
|
||||
if "messages" in stream_modes:
|
||||
ns_ = cast(str | None, config[CONF].get(CONFIG_KEY_CHECKPOINT_NS))
|
||||
use_stream_messages_v2 = bool(
|
||||
version == "v2" and config[CONF].get(CONFIG_KEY_STREAM_MESSAGES_V2)
|
||||
)
|
||||
@@ -2909,24 +2920,12 @@ class Pregel(
|
||||
if use_stream_messages_v2
|
||||
else StreamMessagesHandler
|
||||
)
|
||||
run_manager.inheritable_handlers.append(
|
||||
messages_handler_cls(
|
||||
stream.put,
|
||||
subgraphs,
|
||||
parent_ns=tuple(ns_.split(NS_SEP)) if ns_ else None,
|
||||
)
|
||||
)
|
||||
stream_handlers.append(messages_handler_cls(stream.put, subgraphs))
|
||||
|
||||
# set up tools stream mode
|
||||
if "tools" in stream_modes:
|
||||
ns_tools = cast(str | None, config[CONF].get(CONFIG_KEY_CHECKPOINT_NS))
|
||||
run_manager.inheritable_handlers.append(
|
||||
StreamToolCallHandler(
|
||||
stream.put,
|
||||
subgraphs,
|
||||
parent_ns=tuple(ns_tools.split(NS_SEP)) if ns_tools else None,
|
||||
)
|
||||
)
|
||||
stream_handlers.append(StreamToolCallHandler(stream.put, subgraphs))
|
||||
run_manager.inheritable_handlers.extend(stream_handlers)
|
||||
|
||||
# set up custom stream mode
|
||||
if "custom" in stream_modes:
|
||||
@@ -2968,6 +2967,8 @@ class Pregel(
|
||||
execution_info=None,
|
||||
server_info=server_info,
|
||||
control=control or parent_runtime.control or RunControl(),
|
||||
# a subgraph shares the root run's exit callbacks
|
||||
_run_exits=parent_runtime._run_exits or _RunExits(),
|
||||
)
|
||||
runtime = parent_runtime.merge(runtime)
|
||||
config[CONF][CONFIG_KEY_RUNTIME] = runtime
|
||||
@@ -3009,6 +3010,10 @@ class Pregel(
|
||||
cache_policy=self.cache_policy,
|
||||
has_graph_lifecycle_callbacks=bool(graph_callback_manager.handlers),
|
||||
) as loop:
|
||||
# Known only now: the loop numbers the namespace of every
|
||||
# subgraph a task runs after its first.
|
||||
for handler in stream_handlers:
|
||||
handler.parent_ns = loop.checkpoint_ns
|
||||
emit_graph_lifecycle_events(loop)
|
||||
# create runner
|
||||
runner = PregelRunner(
|
||||
@@ -3291,6 +3296,20 @@ class Pregel(
|
||||
for h in callback_manager.inheritable_handlers
|
||||
if not isinstance(h, StreamMessagesHandlerV2)
|
||||
]
|
||||
if CONFIG_KEY_TASK_ID not in config[CONF]:
|
||||
# Not run as a task (a child invoked with its own thread, say),
|
||||
# this graph's namespace starts at the root, so a caller's
|
||||
# handler would take this graph's events for its own graph's.
|
||||
callback_manager.handlers = [
|
||||
h
|
||||
for h in callback_manager.handlers
|
||||
if not _is_stream_mode_handler(h)
|
||||
]
|
||||
callback_manager.inheritable_handlers = [
|
||||
h
|
||||
for h in callback_manager.inheritable_handlers
|
||||
if not _is_stream_mode_handler(h)
|
||||
]
|
||||
if "ls_integration" not in callback_manager.metadata:
|
||||
callback_manager.add_metadata({"ls_integration": "langgraph"})
|
||||
run_manager = await callback_manager.on_chain_start(
|
||||
@@ -3304,20 +3323,10 @@ class Pregel(
|
||||
run_id=run_manager.run_id,
|
||||
)
|
||||
# if running from astream_log() run each proc with streaming
|
||||
do_stream = (
|
||||
next(
|
||||
(
|
||||
True
|
||||
for h in run_manager.handlers
|
||||
if isinstance(h, _StreamingCallbackHandler)
|
||||
and not isinstance(
|
||||
h, (StreamMessagesHandler, StreamToolCallHandler)
|
||||
)
|
||||
),
|
||||
False,
|
||||
)
|
||||
if _StreamingCallbackHandler is not None
|
||||
else False
|
||||
do_stream = any(
|
||||
_is_stream_handler(h)
|
||||
and not isinstance(h, (StreamMessagesHandler, StreamToolCallHandler))
|
||||
for h in run_manager.handlers
|
||||
)
|
||||
if checkpointer is None and durability is not None:
|
||||
warnings.warn(
|
||||
@@ -3325,10 +3334,9 @@ class Pregel(
|
||||
)
|
||||
# set up subgraph checkpointing
|
||||
config = self._own_checkpoint_config(config)
|
||||
stream_handlers: list[StreamMessagesHandler | StreamToolCallHandler] = []
|
||||
# set up messages stream mode
|
||||
if "messages" in stream_modes:
|
||||
# namespace can be None in a root level graph?
|
||||
ns_ = cast(str | None, config[CONF].get(CONFIG_KEY_CHECKPOINT_NS))
|
||||
use_stream_messages_v2 = bool(
|
||||
version == "v2" and config[CONF].get(CONFIG_KEY_STREAM_MESSAGES_V2)
|
||||
)
|
||||
@@ -3337,24 +3345,12 @@ class Pregel(
|
||||
if use_stream_messages_v2
|
||||
else StreamMessagesHandler
|
||||
)
|
||||
run_manager.inheritable_handlers.append(
|
||||
messages_handler_cls(
|
||||
stream_put,
|
||||
subgraphs,
|
||||
parent_ns=tuple(ns_.split(NS_SEP)) if ns_ else None,
|
||||
)
|
||||
)
|
||||
stream_handlers.append(messages_handler_cls(stream_put, subgraphs))
|
||||
|
||||
# set up tools stream mode
|
||||
if "tools" in stream_modes:
|
||||
ns_tools = cast(str | None, config[CONF].get(CONFIG_KEY_CHECKPOINT_NS))
|
||||
run_manager.inheritable_handlers.append(
|
||||
StreamToolCallHandler(
|
||||
stream_put,
|
||||
subgraphs,
|
||||
parent_ns=tuple(ns_tools.split(NS_SEP)) if ns_tools else None,
|
||||
)
|
||||
)
|
||||
stream_handlers.append(StreamToolCallHandler(stream_put, subgraphs))
|
||||
run_manager.inheritable_handlers.extend(stream_handlers)
|
||||
|
||||
# set up custom stream mode
|
||||
def stream_writer(c: Any) -> None:
|
||||
@@ -3411,6 +3407,8 @@ class Pregel(
|
||||
execution_info=None,
|
||||
server_info=server_info,
|
||||
control=control or parent_runtime.control or RunControl(),
|
||||
# a subgraph shares the root run's exit callbacks
|
||||
_run_exits=parent_runtime._run_exits or _RunExits(),
|
||||
)
|
||||
runtime = parent_runtime.merge(runtime)
|
||||
config[CONF][CONFIG_KEY_RUNTIME] = runtime
|
||||
@@ -3463,6 +3461,10 @@ class Pregel(
|
||||
cache_policy=self.cache_policy,
|
||||
has_graph_lifecycle_callbacks=bool(graph_callback_manager.handlers),
|
||||
) as loop:
|
||||
# Known only now: the loop numbers the namespace of every
|
||||
# subgraph a task runs after its first.
|
||||
for handler in stream_handlers:
|
||||
handler.parent_ns = loop.checkpoint_ns
|
||||
await aemit_graph_lifecycle_events(loop)
|
||||
# create runner
|
||||
runner = PregelRunner(
|
||||
@@ -4460,6 +4462,10 @@ def _coerce_checkpoint_values(payload: Any, mapper: Callable[[Any], Any]) -> Non
|
||||
payload["values"] = mapper(payload["values"])
|
||||
|
||||
|
||||
def _is_stream_mode_handler(handler: object) -> bool:
|
||||
return isinstance(handler, (StreamMessagesHandler, StreamToolCallHandler))
|
||||
|
||||
|
||||
def _resolve_parent_ns(
|
||||
graph_config: RunnableConfig | None, call_config: RunnableConfig | None
|
||||
) -> tuple[str, ...]:
|
||||
|
||||
@@ -1,7 +1,11 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Callable
|
||||
import inspect
|
||||
import logging
|
||||
import threading
|
||||
from collections.abc import Awaitable, Callable
|
||||
from dataclasses import dataclass, field, replace
|
||||
from types import TracebackType
|
||||
from typing import Any, Generic, cast
|
||||
|
||||
from langgraph.store.base import BaseStore
|
||||
@@ -104,6 +108,129 @@ class RunControl:
|
||||
return self._drain_reason
|
||||
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
ExitCallback = Callable[[], Awaitable[None] | None]
|
||||
"""A cleanup callback registered with `Runtime.on_exit`."""
|
||||
|
||||
|
||||
class _RunExits:
|
||||
"""Cleanup callbacks for one root graph invocation.
|
||||
|
||||
Created by the root `stream`/`astream` call and shared by every subgraph
|
||||
it runs, through `Runtime`. The root Pregel loop binds it at the bottom of
|
||||
its exit stack, so callbacks run after the loop has stopped every task,
|
||||
persisted its final checkpoint and handled any interrupt, whatever the
|
||||
reason the run ended.
|
||||
"""
|
||||
|
||||
__slots__ = ("_callbacks", "_is_async", "_lock", "_state")
|
||||
|
||||
def __init__(self) -> None:
|
||||
self._callbacks: list[ExitCallback] = []
|
||||
self._lock = threading.Lock()
|
||||
self._is_async = False
|
||||
# "new" -> "bound" (owned by a loop) -> "closed" (callbacks ran)
|
||||
self._state = "new"
|
||||
|
||||
def bind(self, *, is_async: bool) -> bool:
|
||||
"""Claim these exits for a loop. Only the first (root) loop wins."""
|
||||
with self._lock:
|
||||
if self._state != "new":
|
||||
return False
|
||||
self._state = "bound"
|
||||
self._is_async = is_async
|
||||
return True
|
||||
|
||||
def register(self, callback: ExitCallback) -> None:
|
||||
if not callable(callback):
|
||||
msg = "on_exit() expects a callable"
|
||||
raise TypeError(msg)
|
||||
with self._lock:
|
||||
if self._state == "closed":
|
||||
msg = "on_exit() called after the graph run has already exited"
|
||||
raise RuntimeError(msg)
|
||||
if not self._is_async and _is_async_callable(callback):
|
||||
msg = (
|
||||
"on_exit() got an async callback in a synchronous graph run "
|
||||
"(invoke/stream). Register a sync callback, or run the graph "
|
||||
"with ainvoke/astream."
|
||||
)
|
||||
raise TypeError(msg)
|
||||
self._callbacks.append(callback)
|
||||
|
||||
def _drain(self) -> list[ExitCallback]:
|
||||
with self._lock:
|
||||
self._state = "closed"
|
||||
callbacks = self._callbacks[::-1] # LIFO
|
||||
self._callbacks = []
|
||||
return callbacks
|
||||
|
||||
def exit(
|
||||
self,
|
||||
exc_type: type[BaseException] | None,
|
||||
exc_value: BaseException | None,
|
||||
traceback: TracebackType | None,
|
||||
) -> bool:
|
||||
errors: list[Exception] = []
|
||||
for callback in self._drain():
|
||||
try:
|
||||
result = callback()
|
||||
if inspect.isawaitable(result):
|
||||
if inspect.iscoroutine(result):
|
||||
result.close()
|
||||
msg = (
|
||||
f"on_exit() callback {callback!r} returned an awaitable "
|
||||
"in a synchronous graph run; it was not awaited."
|
||||
)
|
||||
raise TypeError(msg)
|
||||
except Exception as e:
|
||||
errors.append(e)
|
||||
return self._finish(errors, exc_value)
|
||||
|
||||
async def aexit(
|
||||
self,
|
||||
exc_type: type[BaseException] | None,
|
||||
exc_value: BaseException | None,
|
||||
traceback: TracebackType | None,
|
||||
) -> bool:
|
||||
errors: list[Exception] = []
|
||||
for callback in self._drain():
|
||||
try:
|
||||
result = callback()
|
||||
if inspect.isawaitable(result):
|
||||
await result
|
||||
except Exception as e:
|
||||
errors.append(e)
|
||||
return self._finish(errors, exc_value)
|
||||
|
||||
@staticmethod
|
||||
def _finish(errors: list[Exception], exc_value: BaseException | None) -> bool:
|
||||
# Every callback runs even if an earlier one failed. A failing callback
|
||||
# never replaces the run's own exception (error, cancellation, stream
|
||||
# close): it is logged instead. If the run otherwise ended cleanly, the
|
||||
# first callback error is raised so cleanup bugs aren't silent.
|
||||
if not errors:
|
||||
return False
|
||||
if exc_value is None:
|
||||
for e in errors[1:]:
|
||||
logger.error("on_exit() callback failed", exc_info=e)
|
||||
raise errors[0]
|
||||
for e in errors:
|
||||
logger.error(
|
||||
"on_exit() callback failed while the graph run was exiting with %r",
|
||||
exc_value,
|
||||
exc_info=e,
|
||||
)
|
||||
return False
|
||||
|
||||
|
||||
def _is_async_callable(callback: Any) -> bool:
|
||||
return inspect.iscoroutinefunction(callback) or inspect.iscoroutinefunction(
|
||||
getattr(callback, "__call__", None)
|
||||
)
|
||||
|
||||
|
||||
def _no_op_stream_writer(_: Any) -> None: ...
|
||||
|
||||
|
||||
@@ -237,6 +364,40 @@ class Runtime(Generic[ContextT]):
|
||||
graph runtime.
|
||||
"""
|
||||
|
||||
_run_exits: _RunExits | None = field(default=None, repr=False, compare=False)
|
||||
"""Exit callbacks of the root graph run. Use `on_exit` to register one."""
|
||||
|
||||
def on_exit(self, callback: ExitCallback) -> None:
|
||||
"""Register cleanup to run when the current graph run ends.
|
||||
|
||||
The callback runs once the root invocation (the `invoke`/`stream`/
|
||||
`ainvoke`/`astream` call the caller made) exits for any reason: it
|
||||
completed, raised, was cancelled, paused on an interrupt, or the caller
|
||||
closed the stream or broke out of it. Calls from subgraphs register on
|
||||
that same root run.
|
||||
|
||||
Callbacks take no arguments and run in reverse registration order,
|
||||
after every task of the run has stopped and the final checkpoint has
|
||||
been written. In `ainvoke`/`astream` a callback may be sync or async; in
|
||||
`invoke`/`stream` it must be sync. A failing callback never hides the
|
||||
run's own exception: it is logged, and raised only if the run otherwise
|
||||
ended without one.
|
||||
|
||||
A process that is killed does not run its callbacks, so use them to
|
||||
free resources early, not as the only guarantee they are freed.
|
||||
|
||||
Args:
|
||||
callback: A zero-argument sync or async callable.
|
||||
|
||||
Raises:
|
||||
RuntimeError: When called outside a graph run, or after it exited.
|
||||
TypeError: When given an async callback in a synchronous run.
|
||||
"""
|
||||
if self._run_exits is None:
|
||||
msg = "on_exit() can only be called during a graph run"
|
||||
raise RuntimeError(msg)
|
||||
self._run_exits.register(callback)
|
||||
|
||||
def merge(self, other: Runtime[ContextT]) -> Runtime[ContextT]:
|
||||
"""Merge two runtimes together.
|
||||
|
||||
@@ -255,6 +416,7 @@ class Runtime(Generic[ContextT]):
|
||||
execution_info=other.execution_info or self.execution_info,
|
||||
server_info=other.server_info or self.server_info,
|
||||
control=other.control or self.control,
|
||||
_run_exits=other._run_exits or self._run_exits,
|
||||
)
|
||||
|
||||
def override(
|
||||
|
||||
@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
|
||||
|
||||
[project]
|
||||
name = "langgraph"
|
||||
version = "1.2.18"
|
||||
version = "1.2.19"
|
||||
description = "Building stateful, multi-actor applications with LLMs"
|
||||
authors = []
|
||||
requires-python = ">=3.10"
|
||||
|
||||
@@ -0,0 +1,339 @@
|
||||
"""Tests for `Runtime.on_exit`: cleanup that runs however a graph run ends."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import sys
|
||||
|
||||
import pytest
|
||||
from langgraph.checkpoint.memory import InMemorySaver
|
||||
from typing_extensions import TypedDict
|
||||
|
||||
from langgraph.graph import END, START, StateGraph
|
||||
from langgraph.runtime import Runtime, get_runtime
|
||||
from langgraph.types import Command, interrupt
|
||||
|
||||
pytestmark = pytest.mark.anyio
|
||||
|
||||
NEEDS_CONTEXTVARS = pytest.mark.skipif(
|
||||
sys.version_info < (3, 11),
|
||||
reason="Python 3.11+ is required for async contextvars support",
|
||||
)
|
||||
|
||||
|
||||
class State(TypedDict, total=False):
|
||||
value: str
|
||||
|
||||
|
||||
def _two_node_graph(first, second=None, *, checkpointer=None):
|
||||
builder = StateGraph(State)
|
||||
builder.add_node("first", first)
|
||||
builder.add_edge(START, "first")
|
||||
if second is None:
|
||||
builder.add_edge("first", END)
|
||||
else:
|
||||
builder.add_node("second", second)
|
||||
builder.add_edge("first", "second")
|
||||
builder.add_edge("second", END)
|
||||
return builder.compile(checkpointer=checkpointer)
|
||||
|
||||
|
||||
# sync
|
||||
|
||||
|
||||
def test_runs_on_success_in_lifo_order() -> None:
|
||||
events: list[str] = []
|
||||
|
||||
def first(state: State, runtime: Runtime) -> State:
|
||||
runtime.on_exit(lambda: events.append("cleanup-a"))
|
||||
runtime.on_exit(lambda: events.append("cleanup-b"))
|
||||
events.append("first")
|
||||
return {"value": "x"}
|
||||
|
||||
def second(state: State) -> State:
|
||||
events.append("second")
|
||||
return {}
|
||||
|
||||
graph = _two_node_graph(first, second)
|
||||
assert graph.invoke({}) == {"value": "x"}
|
||||
assert events == ["first", "second", "cleanup-b", "cleanup-a"]
|
||||
|
||||
|
||||
def test_runs_on_error_and_keeps_the_run_exception() -> None:
|
||||
events: list[str] = []
|
||||
|
||||
def failing_cleanup() -> None:
|
||||
events.append("failing")
|
||||
raise KeyError("cleanup failed")
|
||||
|
||||
def first(state: State) -> State:
|
||||
runtime = get_runtime()
|
||||
runtime.on_exit(lambda: events.append("cleanup"))
|
||||
runtime.on_exit(failing_cleanup)
|
||||
raise ValueError("node failed")
|
||||
|
||||
graph = _two_node_graph(first)
|
||||
with pytest.raises(ValueError, match="node failed"):
|
||||
graph.invoke({})
|
||||
# the failing callback did not stop the earlier one from running
|
||||
assert events == ["failing", "cleanup"]
|
||||
|
||||
|
||||
def test_callback_error_is_raised_when_the_run_succeeded() -> None:
|
||||
def failing_cleanup() -> None:
|
||||
raise KeyError("cleanup failed")
|
||||
|
||||
def first(state: State, runtime: Runtime) -> State:
|
||||
runtime.on_exit(failing_cleanup)
|
||||
return {}
|
||||
|
||||
graph = _two_node_graph(first)
|
||||
with pytest.raises(KeyError, match="cleanup failed"):
|
||||
graph.invoke({})
|
||||
|
||||
|
||||
def test_runs_on_interrupt() -> None:
|
||||
events: list[str] = []
|
||||
|
||||
def first(state: State, runtime: Runtime) -> State:
|
||||
runtime.on_exit(lambda: events.append("cleanup"))
|
||||
return {"value": interrupt("approve?")}
|
||||
|
||||
graph = _two_node_graph(first, checkpointer=InMemorySaver())
|
||||
config = {"configurable": {"thread_id": "1"}}
|
||||
result = graph.invoke({}, config)
|
||||
assert "__interrupt__" in result
|
||||
assert events == ["cleanup"]
|
||||
|
||||
# each resumed run is its own invocation with its own callbacks
|
||||
graph.invoke(Command(resume="ok"), config)
|
||||
assert events == ["cleanup", "cleanup"]
|
||||
|
||||
|
||||
def test_runs_when_the_stream_is_closed_early() -> None:
|
||||
events: list[str] = []
|
||||
|
||||
def first(state: State, runtime: Runtime) -> State:
|
||||
runtime.on_exit(lambda: events.append("cleanup"))
|
||||
return {"value": "x"}
|
||||
|
||||
def second(state: State) -> State:
|
||||
events.append("second")
|
||||
return {}
|
||||
|
||||
graph = _two_node_graph(first, second)
|
||||
stream = graph.stream({}, stream_mode="updates")
|
||||
assert next(stream) == {"first": {"value": "x"}}
|
||||
assert events == []
|
||||
stream.close()
|
||||
assert events == ["cleanup"]
|
||||
|
||||
|
||||
def test_runs_when_breaking_out_of_stream() -> None:
|
||||
events: list[str] = []
|
||||
|
||||
def first(state: State, runtime: Runtime) -> State:
|
||||
runtime.on_exit(lambda: events.append("cleanup"))
|
||||
return {"value": "x"}
|
||||
|
||||
graph = _two_node_graph(first, lambda s: {})
|
||||
for _ in graph.stream({}, stream_mode="updates"):
|
||||
break
|
||||
# the dropped generator is closed by refcounting on CPython
|
||||
assert events == ["cleanup"]
|
||||
|
||||
|
||||
def test_subgraph_registers_on_the_root_run() -> None:
|
||||
events: list[str] = []
|
||||
|
||||
def inner(state: State, runtime: Runtime) -> State:
|
||||
runtime.on_exit(lambda: events.append("cleanup"))
|
||||
events.append("inner")
|
||||
return {}
|
||||
|
||||
sub = StateGraph(State)
|
||||
sub.add_node("inner", inner)
|
||||
sub.add_edge(START, "inner")
|
||||
subgraph = sub.compile()
|
||||
|
||||
def call_subgraph(state: State) -> State:
|
||||
subgraph.invoke(state)
|
||||
events.append("subgraph returned")
|
||||
return {}
|
||||
|
||||
def after(state: State) -> State:
|
||||
events.append("after")
|
||||
return {}
|
||||
|
||||
builder = StateGraph(State)
|
||||
builder.add_node("as_node", subgraph)
|
||||
builder.add_node("call_subgraph", call_subgraph)
|
||||
builder.add_node("after", after)
|
||||
builder.add_edge(START, "as_node")
|
||||
builder.add_edge("as_node", "call_subgraph")
|
||||
builder.add_edge("call_subgraph", "after")
|
||||
graph = builder.compile()
|
||||
|
||||
graph.invoke({})
|
||||
# nothing runs when either subgraph ends; both run when the root ends
|
||||
assert events == [
|
||||
"inner",
|
||||
"inner",
|
||||
"subgraph returned",
|
||||
"after",
|
||||
"cleanup",
|
||||
"cleanup",
|
||||
]
|
||||
|
||||
|
||||
def test_async_callback_rejected_in_sync_run() -> None:
|
||||
async def acleanup() -> None:
|
||||
pass
|
||||
|
||||
def first(state: State, runtime: Runtime) -> State:
|
||||
runtime.on_exit(acleanup)
|
||||
return {}
|
||||
|
||||
graph = _two_node_graph(first)
|
||||
with pytest.raises(TypeError, match="async callback in a synchronous"):
|
||||
graph.invoke({})
|
||||
|
||||
|
||||
def test_outside_a_run_raises() -> None:
|
||||
with pytest.raises(RuntimeError, match="during a graph run"):
|
||||
Runtime().on_exit(lambda: None)
|
||||
|
||||
|
||||
def test_each_invocation_gets_its_own_callbacks() -> None:
|
||||
calls: list[int] = []
|
||||
|
||||
def first(state: State, runtime: Runtime) -> State:
|
||||
runtime.on_exit(lambda: calls.append(1))
|
||||
return {}
|
||||
|
||||
graph = _two_node_graph(first)
|
||||
graph.invoke({})
|
||||
graph.invoke({})
|
||||
assert calls == [1, 1]
|
||||
|
||||
|
||||
# async
|
||||
|
||||
|
||||
async def test_async_runs_on_success_sync_and_async_callbacks() -> None:
|
||||
events: list[str] = []
|
||||
|
||||
async def acleanup() -> None:
|
||||
await asyncio.sleep(0)
|
||||
events.append("async-cleanup")
|
||||
|
||||
async def first(state: State, runtime: Runtime) -> State:
|
||||
runtime.on_exit(lambda: events.append("sync-cleanup"))
|
||||
runtime.on_exit(acleanup)
|
||||
events.append("first")
|
||||
return {"value": "x"}
|
||||
|
||||
graph = _two_node_graph(first)
|
||||
assert await graph.ainvoke({}) == {"value": "x"}
|
||||
assert events == ["first", "async-cleanup", "sync-cleanup"]
|
||||
|
||||
|
||||
async def test_async_runs_on_error_and_keeps_the_run_exception() -> None:
|
||||
events: list[str] = []
|
||||
|
||||
async def failing_cleanup() -> None:
|
||||
raise KeyError("cleanup failed")
|
||||
|
||||
async def first(state: State, runtime: Runtime) -> State:
|
||||
runtime.on_exit(lambda: events.append("cleanup"))
|
||||
runtime.on_exit(failing_cleanup)
|
||||
raise ValueError("node failed")
|
||||
|
||||
graph = _two_node_graph(first)
|
||||
with pytest.raises(ValueError, match="node failed"):
|
||||
await graph.ainvoke({})
|
||||
assert events == ["cleanup"]
|
||||
|
||||
|
||||
async def test_async_runs_on_interrupt() -> None:
|
||||
events: list[str] = []
|
||||
|
||||
async def first(state: State, runtime: Runtime) -> State:
|
||||
runtime.on_exit(lambda: events.append("cleanup"))
|
||||
return {"value": interrupt("approve?")}
|
||||
|
||||
graph = _two_node_graph(first, checkpointer=InMemorySaver())
|
||||
result = await graph.ainvoke({}, {"configurable": {"thread_id": "1"}})
|
||||
assert "__interrupt__" in result
|
||||
assert events == ["cleanup"]
|
||||
|
||||
|
||||
async def test_async_runs_on_cancel_after_the_node_unwinds() -> None:
|
||||
events: list[str] = []
|
||||
started = asyncio.Event()
|
||||
|
||||
async def acleanup() -> None:
|
||||
events.append("cleanup")
|
||||
|
||||
async def first(state: State, runtime: Runtime) -> State:
|
||||
runtime.on_exit(acleanup)
|
||||
started.set()
|
||||
try:
|
||||
await asyncio.sleep(60)
|
||||
finally:
|
||||
events.append("node finally")
|
||||
return {}
|
||||
|
||||
graph = _two_node_graph(first)
|
||||
task = asyncio.create_task(graph.ainvoke({}))
|
||||
await started.wait()
|
||||
task.cancel()
|
||||
with pytest.raises(asyncio.CancelledError):
|
||||
await task
|
||||
# cleanup ran after the node had stopped
|
||||
assert events == ["node finally", "cleanup"]
|
||||
|
||||
|
||||
async def test_async_runs_when_the_stream_is_closed_early() -> None:
|
||||
events: list[str] = []
|
||||
|
||||
async def acleanup() -> None:
|
||||
events.append("cleanup")
|
||||
|
||||
async def first(state: State, runtime: Runtime) -> State:
|
||||
runtime.on_exit(acleanup)
|
||||
return {"value": "x"}
|
||||
|
||||
async def second(state: State) -> State:
|
||||
events.append("second")
|
||||
return {}
|
||||
|
||||
graph = _two_node_graph(first, second)
|
||||
stream = graph.astream({}, stream_mode="updates")
|
||||
assert await stream.__anext__() == {"first": {"value": "x"}}
|
||||
assert events == []
|
||||
await stream.aclose()
|
||||
assert events == ["cleanup"]
|
||||
|
||||
|
||||
@NEEDS_CONTEXTVARS
|
||||
async def test_async_subgraph_registers_on_the_root_run() -> None:
|
||||
events: list[str] = []
|
||||
|
||||
async def inner(state: State) -> State:
|
||||
get_runtime().on_exit(lambda: events.append("cleanup"))
|
||||
return {}
|
||||
|
||||
sub = StateGraph(State)
|
||||
sub.add_node("inner", inner)
|
||||
sub.add_edge(START, "inner")
|
||||
subgraph = sub.compile()
|
||||
|
||||
async def call_subgraph(state: State) -> State:
|
||||
await subgraph.ainvoke(state)
|
||||
events.append("subgraph returned")
|
||||
return {}
|
||||
|
||||
graph = _two_node_graph(call_subgraph, lambda s: events.append("after") or {})
|
||||
await graph.ainvoke({})
|
||||
assert events == ["subgraph returned", "after", "cleanup"]
|
||||
@@ -0,0 +1,308 @@
|
||||
import sys
|
||||
from typing import Any
|
||||
|
||||
import pytest
|
||||
from langchain_core.language_models.fake_chat_models import GenericFakeChatModel
|
||||
from langchain_core.messages import AIMessage
|
||||
from langchain_core.tools import tool
|
||||
|
||||
from langgraph.func import entrypoint, task
|
||||
from langgraph.graph import START, MessagesState, StateGraph
|
||||
from langgraph.graph.message import push_message
|
||||
from langgraph.graph.state import CompiledStateGraph
|
||||
|
||||
pytestmark = pytest.mark.anyio
|
||||
|
||||
NEEDS_CONTEXTVARS = pytest.mark.skipif(
|
||||
sys.version_info < (3, 11),
|
||||
reason="Python 3.11+ is required for async contextvars support",
|
||||
)
|
||||
|
||||
|
||||
@tool
|
||||
def echo(text: str) -> str:
|
||||
"""Echo the text."""
|
||||
return text
|
||||
|
||||
|
||||
def _one_node_graph(node: Any) -> CompiledStateGraph:
|
||||
return (
|
||||
StateGraph(MessagesState)
|
||||
.add_node(node)
|
||||
.add_edge(START, node.__name__)
|
||||
.compile()
|
||||
)
|
||||
|
||||
|
||||
def _chat_and_echo(state: MessagesState) -> dict:
|
||||
echo.invoke({"text": "done"})
|
||||
return {"messages": [GenericFakeChatModel(messages=iter(["one two"])).invoke("hi")]}
|
||||
|
||||
|
||||
async def _achat_and_echo(state: MessagesState) -> dict:
|
||||
await echo.ainvoke({"text": "done"})
|
||||
model = GenericFakeChatModel(messages=iter(["one two"]))
|
||||
return {"messages": [await model.ainvoke("hi")]}
|
||||
|
||||
|
||||
def _segment_names(ns: tuple[str, ...]) -> tuple[str, ...]:
|
||||
return tuple(segment.split(":")[0] for segment in ns)
|
||||
|
||||
|
||||
def _summary(part: dict) -> tuple[tuple[str, ...], str]:
|
||||
if part["type"] == "messages":
|
||||
return _segment_names(part["ns"]), part["data"][0].content
|
||||
return _segment_names(part["ns"]), part["data"]["event"]
|
||||
|
||||
|
||||
_CHAT_AND_ECHO_EVENTS = [
|
||||
((), "tool-started"),
|
||||
((), "tool-finished"),
|
||||
((), "one"),
|
||||
((), " "),
|
||||
((), "two"),
|
||||
]
|
||||
|
||||
_CHILD_CONFIGS = pytest.mark.parametrize(
|
||||
"child_config",
|
||||
[
|
||||
None,
|
||||
{"configurable": {"thread_id": "child"}},
|
||||
{"configurable": {"thread_id": "parent"}},
|
||||
],
|
||||
ids=["no_config", "own_thread", "caller_thread"],
|
||||
)
|
||||
|
||||
|
||||
_TASK_CALLERS = pytest.mark.parametrize(
|
||||
"nested", [False, True], ids=["from_entrypoint", "from_task_in_task"]
|
||||
)
|
||||
|
||||
|
||||
@_TASK_CALLERS
|
||||
@pytest.mark.parametrize("subgraphs", [False, True])
|
||||
def test_task_events_stream_under_the_graph_namespace(
|
||||
nested: bool, subgraphs: bool
|
||||
) -> None:
|
||||
@task
|
||||
def chat_and_echo() -> None:
|
||||
_chat_and_echo({"messages": []})
|
||||
|
||||
@task
|
||||
def outer() -> None:
|
||||
chat_and_echo().result()
|
||||
|
||||
@entrypoint()
|
||||
def workflow(inputs: dict) -> None:
|
||||
(outer() if nested else chat_and_echo()).result()
|
||||
|
||||
parts = workflow.stream(
|
||||
{}, stream_mode=["messages", "tools"], subgraphs=subgraphs, version="v2"
|
||||
)
|
||||
|
||||
assert [_summary(part) for part in parts] == _CHAT_AND_ECHO_EVENTS
|
||||
|
||||
|
||||
@NEEDS_CONTEXTVARS
|
||||
@_TASK_CALLERS
|
||||
@pytest.mark.parametrize("subgraphs", [False, True])
|
||||
async def test_task_events_stream_under_the_graph_namespace_async(
|
||||
nested: bool, subgraphs: bool
|
||||
) -> None:
|
||||
@task
|
||||
async def chat_and_echo() -> None:
|
||||
await _achat_and_echo({"messages": []})
|
||||
|
||||
@task
|
||||
async def outer() -> None:
|
||||
await chat_and_echo()
|
||||
|
||||
@entrypoint()
|
||||
async def workflow(inputs: dict) -> None:
|
||||
await (outer() if nested else chat_and_echo())
|
||||
|
||||
parts = workflow.astream(
|
||||
{}, stream_mode=["messages", "tools"], subgraphs=subgraphs, version="v2"
|
||||
)
|
||||
|
||||
assert [_summary(part) async for part in parts] == _CHAT_AND_ECHO_EVENTS
|
||||
|
||||
|
||||
def test_task_events_stream_under_the_namespace_of_their_subgraph() -> None:
|
||||
@task
|
||||
def chat_and_echo() -> None:
|
||||
_chat_and_echo({"messages": []})
|
||||
|
||||
def node(state: MessagesState) -> dict:
|
||||
chat_and_echo().result()
|
||||
return {}
|
||||
|
||||
sub = _one_node_graph(node)
|
||||
parent = StateGraph(MessagesState).add_node("sub", sub).add_edge(START, "sub")
|
||||
|
||||
parts = parent.compile().stream(
|
||||
{"messages": []},
|
||||
stream_mode=["messages", "tools"],
|
||||
subgraphs=True,
|
||||
version="v2",
|
||||
)
|
||||
|
||||
assert [_summary(part) for part in parts] == [
|
||||
(("sub",), event) for _, event in _CHAT_AND_ECHO_EVENTS
|
||||
]
|
||||
|
||||
|
||||
@_CHILD_CONFIGS
|
||||
def test_child_graph_streams_only_to_the_node_streaming_it(
|
||||
child_config: dict | None,
|
||||
) -> None:
|
||||
child = _one_node_graph(_chat_and_echo)
|
||||
received = []
|
||||
|
||||
def node(state: MessagesState) -> dict:
|
||||
received.extend(
|
||||
child.stream(
|
||||
{"messages": []}, child_config, stream_mode=["messages", "tools"]
|
||||
)
|
||||
)
|
||||
return {}
|
||||
|
||||
parts = _one_node_graph(node).stream(
|
||||
{"messages": []},
|
||||
{"configurable": {"thread_id": "parent"}},
|
||||
stream_mode=["messages", "tools"],
|
||||
)
|
||||
|
||||
assert list(parts) == []
|
||||
assert len(received) == len(_CHAT_AND_ECHO_EVENTS)
|
||||
|
||||
|
||||
@NEEDS_CONTEXTVARS
|
||||
@_CHILD_CONFIGS
|
||||
async def test_child_graph_streams_only_to_the_node_streaming_it_async(
|
||||
child_config: dict | None,
|
||||
) -> None:
|
||||
child = _one_node_graph(_achat_and_echo)
|
||||
received = []
|
||||
|
||||
async def node(state: MessagesState) -> dict:
|
||||
stream = child.astream(
|
||||
{"messages": []}, child_config, stream_mode=["messages", "tools"]
|
||||
)
|
||||
received.extend([part async for part in stream])
|
||||
return {}
|
||||
|
||||
parts = _one_node_graph(node).astream(
|
||||
{"messages": []},
|
||||
{"configurable": {"thread_id": "parent"}},
|
||||
stream_mode=["messages", "tools"],
|
||||
)
|
||||
|
||||
assert [part async for part in parts] == []
|
||||
assert len(received) == len(_CHAT_AND_ECHO_EVENTS)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"child_config",
|
||||
[
|
||||
{"configurable": {"thread_id": "child"}},
|
||||
{"configurable": {"thread_id": "parent"}},
|
||||
],
|
||||
ids=["own_thread", "caller_thread"],
|
||||
)
|
||||
def test_child_graph_on_its_own_thread_skips_a_caller_streaming_subgraphs(
|
||||
child_config: dict,
|
||||
) -> None:
|
||||
child = _one_node_graph(_chat_and_echo)
|
||||
|
||||
def node(state: MessagesState) -> dict:
|
||||
child.invoke({"messages": []}, child_config)
|
||||
return {}
|
||||
|
||||
parts = _one_node_graph(node).stream(
|
||||
{"messages": []},
|
||||
{"configurable": {"thread_id": "parent"}},
|
||||
stream_mode=["messages", "tools"],
|
||||
subgraphs=True,
|
||||
)
|
||||
|
||||
assert list(parts) == []
|
||||
|
||||
|
||||
def test_node_gets_the_events_of_every_subgraph_it_streams() -> None:
|
||||
sub = _one_node_graph(_chat_and_echo)
|
||||
counts = []
|
||||
|
||||
def node(state: MessagesState) -> dict:
|
||||
for _ in range(2):
|
||||
stream = sub.stream({"messages": []}, stream_mode=["messages", "tools"])
|
||||
counts.append(len(list(stream)))
|
||||
return {}
|
||||
|
||||
_one_node_graph(node).invoke({"messages": []})
|
||||
|
||||
assert counts == [len(_CHAT_AND_ECHO_EVENTS)] * 2
|
||||
|
||||
|
||||
@NEEDS_CONTEXTVARS
|
||||
async def test_node_gets_the_events_of_every_subgraph_it_streams_async() -> None:
|
||||
sub = _one_node_graph(_achat_and_echo)
|
||||
counts = []
|
||||
|
||||
async def node(state: MessagesState) -> dict:
|
||||
for _ in range(2):
|
||||
stream = sub.astream({"messages": []}, stream_mode=["messages", "tools"])
|
||||
counts.append(len([part async for part in stream]))
|
||||
return {}
|
||||
|
||||
await _one_node_graph(node).ainvoke({"messages": []})
|
||||
|
||||
assert counts == [len(_CHAT_AND_ECHO_EVENTS)] * 2
|
||||
|
||||
|
||||
def _push(state: MessagesState) -> dict:
|
||||
push_message(AIMessage(id="pushed", content="hi"), state_key=None)
|
||||
return {}
|
||||
|
||||
|
||||
def _namespaces(graph: CompiledStateGraph, subgraphs: bool) -> list[tuple[str, ...]]:
|
||||
parts = graph.stream(
|
||||
{"messages": []}, stream_mode="messages", subgraphs=subgraphs, version="v2"
|
||||
)
|
||||
return [_segment_names(part["ns"]) for part in parts]
|
||||
|
||||
|
||||
def _streams_pushing_graph(received: list) -> Any:
|
||||
pushing = _one_node_graph(_push)
|
||||
|
||||
def node(state: MessagesState) -> dict:
|
||||
for message, _ in pushing.stream({"messages": []}, stream_mode="messages"):
|
||||
received.append(message.id)
|
||||
return {}
|
||||
|
||||
return node
|
||||
|
||||
|
||||
def test_pushed_message_streams_under_the_namespace_of_its_graph() -> None:
|
||||
pushing = _one_node_graph(_push)
|
||||
parent = StateGraph(MessagesState).add_node("sub", pushing).add_edge(START, "sub")
|
||||
|
||||
assert _namespaces(pushing, subgraphs=True) == [()]
|
||||
assert _namespaces(parent.compile(), subgraphs=False) == []
|
||||
assert _namespaces(parent.compile(), subgraphs=True) == [("sub",)]
|
||||
|
||||
|
||||
def test_pushed_message_streams_to_the_node_that_streams_its_graph() -> None:
|
||||
received: list[str] = []
|
||||
graph = _one_node_graph(_streams_pushing_graph(received))
|
||||
|
||||
assert _namespaces(graph, subgraphs=False) == []
|
||||
assert received == ["pushed"]
|
||||
|
||||
|
||||
def test_pushed_message_streams_to_every_stream_that_takes_it() -> None:
|
||||
received: list[str] = []
|
||||
graph = _one_node_graph(_streams_pushing_graph(received))
|
||||
|
||||
assert _namespaces(graph, subgraphs=True) == [("node",)]
|
||||
assert received == ["pushed"]
|
||||
@@ -1,11 +1,24 @@
|
||||
"""End-to-end tests for `TracePolicy` input processing on node trace runs."""
|
||||
|
||||
import asyncio
|
||||
from collections.abc import AsyncIterator, Iterator
|
||||
from typing import Any
|
||||
from uuid import UUID
|
||||
|
||||
import pytest
|
||||
from langchain_core.callbacks import BaseCallbackHandler
|
||||
from langchain_core.language_models.fake_chat_models import FakeMessagesListChatModel
|
||||
from langchain_core.messages import AIMessage, HumanMessage
|
||||
from langchain_core.runnables import RunnableConfig
|
||||
from typing_extensions import TypedDict
|
||||
|
||||
from langgraph.graph import END, START, StateGraph
|
||||
from langgraph._internal._runnable import (
|
||||
ASYNCIO_ACCEPTS_CONTEXT,
|
||||
RunnableCallable,
|
||||
RunnableSeq,
|
||||
)
|
||||
from langgraph.graph import END, START, MessagesState, StateGraph
|
||||
from langgraph.graph.state import CompiledStateGraph
|
||||
from langgraph.types import TracePolicy
|
||||
from tests.fake_tracer import FakeTracer, Run
|
||||
|
||||
@@ -147,3 +160,179 @@ async def test_trace_policy_transforms_recorded_inputs_async() -> None:
|
||||
run = _node_run(tracer, "n")
|
||||
assert run.inputs == {"scrubbed_in": True}
|
||||
assert run.outputs == {"value": 6}
|
||||
|
||||
|
||||
_RECORDING_POLICY = TracePolicy(
|
||||
process_inputs=lambda _: {"recorded": "inputs"},
|
||||
process_outputs=lambda _: {"recorded": "outputs"},
|
||||
)
|
||||
|
||||
|
||||
def _chat_graph() -> CompiledStateGraph:
|
||||
model = FakeMessagesListChatModel(
|
||||
responses=[AIMessage("model reply", id="ai-model")]
|
||||
)
|
||||
|
||||
def chat(state: MessagesState) -> dict[str, Any]:
|
||||
reply = model.invoke(state["messages"])
|
||||
return {
|
||||
"messages": [
|
||||
*state["messages"],
|
||||
reply,
|
||||
AIMessage("from node", id="ai-node"),
|
||||
]
|
||||
}
|
||||
|
||||
return (
|
||||
StateGraph(MessagesState)
|
||||
.add_node("chat", chat, trace_policy=_RECORDING_POLICY)
|
||||
.add_edge(START, "chat")
|
||||
.compile()
|
||||
)
|
||||
|
||||
|
||||
def _assert_scrubbed_node_run(tracer: FakeTracer, *, model_is_child: bool) -> None:
|
||||
run = _node_run(tracer, "chat")
|
||||
assert run.inputs == {"recorded": "inputs"}
|
||||
assert run.outputs == {"recorded": "outputs"}
|
||||
if model_is_child:
|
||||
assert [c.name for c in run.child_runs] == ["FakeMessagesListChatModel"]
|
||||
|
||||
|
||||
def test_trace_policy_scrubs_tracer_not_messages_stream() -> None:
|
||||
tracer = FakeTracer()
|
||||
streamed = [
|
||||
message.content
|
||||
for message, _ in _chat_graph().stream(
|
||||
{"messages": [HumanMessage("hello", id="h-1")]},
|
||||
{"callbacks": [tracer]},
|
||||
stream_mode="messages",
|
||||
)
|
||||
]
|
||||
|
||||
assert streamed == ["model reply", "from node"]
|
||||
_assert_scrubbed_node_run(tracer, model_is_child=True)
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_trace_policy_scrubs_tracer_not_messages_stream_async() -> None:
|
||||
tracer = FakeTracer()
|
||||
streamed = [
|
||||
message.content
|
||||
async for message, _ in _chat_graph().astream(
|
||||
{"messages": [HumanMessage("hello", id="h-1")]},
|
||||
{"callbacks": [tracer]},
|
||||
stream_mode="messages",
|
||||
)
|
||||
]
|
||||
|
||||
assert streamed == ["model reply", "from node"]
|
||||
_assert_scrubbed_node_run(tracer, model_is_child=ASYNCIO_ACCEPTS_CONTEXT)
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_trace_policy_scrubs_tracer_not_astream_events() -> None:
|
||||
tracer = FakeTracer()
|
||||
chat_events = [
|
||||
event
|
||||
async for event in _chat_graph().astream_events(
|
||||
{"messages": [HumanMessage("hello", id="h-1")]},
|
||||
{"callbacks": [tracer]},
|
||||
version="v2",
|
||||
)
|
||||
if event["name"] == "chat"
|
||||
and event["event"] in ("on_chain_start", "on_chain_end")
|
||||
]
|
||||
|
||||
assert [e["event"] for e in chat_events] == ["on_chain_start", "on_chain_end"]
|
||||
start, end = (e["data"] for e in chat_events)
|
||||
assert [m.content for m in start["input"]["messages"]] == ["hello"]
|
||||
assert [m.content for m in end["output"]["messages"]] == [
|
||||
"hello",
|
||||
"model reply",
|
||||
"from node",
|
||||
]
|
||||
_assert_scrubbed_node_run(tracer, model_is_child=ASYNCIO_ACCEPTS_CONTEXT)
|
||||
|
||||
|
||||
class _HandlerWithTapMethods(BaseCallbackHandler):
|
||||
def __init__(self) -> None:
|
||||
self.node_payloads: list[Any] = []
|
||||
self.node_run_ids: set[UUID] = set()
|
||||
|
||||
def _record(self, run_id: UUID, payload: Any) -> None:
|
||||
if run_id in self.node_run_ids:
|
||||
self.node_payloads.append(payload)
|
||||
|
||||
def tap_output_iter(self, run_id: UUID, output: Iterator[Any]) -> Iterator[Any]:
|
||||
for chunk in output:
|
||||
self._record(run_id, chunk)
|
||||
yield chunk
|
||||
|
||||
async def tap_output_aiter(
|
||||
self, run_id: UUID, output: AsyncIterator[Any]
|
||||
) -> AsyncIterator[Any]:
|
||||
async for chunk in output:
|
||||
self._record(run_id, chunk)
|
||||
yield chunk
|
||||
|
||||
def on_chain_start(
|
||||
self, serialized: Any, inputs: Any, *, run_id: UUID, **kwargs: Any
|
||||
) -> None:
|
||||
if kwargs.get("name") == "n":
|
||||
self.node_run_ids.add(run_id)
|
||||
self._record(run_id, inputs)
|
||||
|
||||
def on_chain_end(self, outputs: Any, *, run_id: UUID, **kwargs: Any) -> None:
|
||||
self._record(run_id, outputs)
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
@pytest.mark.parametrize("via_events", [False, True], ids=["ainvoke", "astream_events"])
|
||||
async def test_trace_policy_reaches_tracers_but_not_custom_handlers(
|
||||
via_events: bool,
|
||||
) -> None:
|
||||
graph = (
|
||||
StateGraph(State)
|
||||
.add_node("n", _incr, trace_policy=_RECORDING_POLICY)
|
||||
.add_edge(START, "n")
|
||||
.add_edge("n", END)
|
||||
.compile()
|
||||
)
|
||||
handler = _HandlerWithTapMethods()
|
||||
tracer = FakeTracer()
|
||||
config: RunnableConfig = {"callbacks": [handler, tracer]}
|
||||
|
||||
if via_events:
|
||||
async for _ in graph.astream_events({"value": 1}, config, version="v2"):
|
||||
pass
|
||||
else:
|
||||
assert await graph.ainvoke({"value": 1}, config) == {"value": 2}
|
||||
assert handler.node_payloads == [{"value": 1}, {"value": 2}]
|
||||
run = _node_run(tracer, "n")
|
||||
assert (run.inputs, run.outputs) == (
|
||||
{"recorded": "inputs"},
|
||||
{"recorded": "outputs"},
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_trace_policy_cancel_at_node_end_still_closes_tracer_run() -> None:
|
||||
def cancel_own_task(value: Any) -> Any:
|
||||
asyncio.current_task().cancel()
|
||||
return value
|
||||
|
||||
tracer = FakeTracer()
|
||||
seq = RunnableSeq(
|
||||
RunnableCallable(lambda x: x, trace=False),
|
||||
RunnableCallable(cancel_own_task, trace=False),
|
||||
trace_policy=_RECORDING_POLICY,
|
||||
)
|
||||
with pytest.raises(asyncio.CancelledError):
|
||||
await asyncio.create_task(seq.ainvoke({}, {"callbacks": [tracer]}))
|
||||
for _ in range(100):
|
||||
if tracer.runs:
|
||||
break
|
||||
await asyncio.sleep(0.01)
|
||||
|
||||
assert [r.outputs for r in tracer.runs] == [{"recorded": "outputs"}]
|
||||
Generated
+2
-2
@@ -1439,7 +1439,7 @@ wheels = [
|
||||
|
||||
[[package]]
|
||||
name = "langgraph"
|
||||
version = "1.2.18"
|
||||
version = "1.2.19"
|
||||
source = { editable = "." }
|
||||
dependencies = [
|
||||
{ name = "langchain-core" },
|
||||
@@ -1722,7 +1722,7 @@ test = [
|
||||
|
||||
[[package]]
|
||||
name = "langgraph-checkpoint-sqlite"
|
||||
version = "3.1.2"
|
||||
version = "3.1.3"
|
||||
source = { editable = "../checkpoint-sqlite" }
|
||||
dependencies = [
|
||||
{ name = "aiosqlite" },
|
||||
|
||||
Generated
+2
-2
@@ -286,7 +286,7 @@ wheels = [
|
||||
|
||||
[[package]]
|
||||
name = "langgraph"
|
||||
version = "1.2.18"
|
||||
version = "1.2.19"
|
||||
source = { editable = "../langgraph" }
|
||||
dependencies = [
|
||||
{ name = "langchain-core" },
|
||||
@@ -467,7 +467,7 @@ test = [
|
||||
|
||||
[[package]]
|
||||
name = "langgraph-checkpoint-sqlite"
|
||||
version = "3.1.2"
|
||||
version = "3.1.3"
|
||||
source = { editable = "../checkpoint-sqlite" }
|
||||
dependencies = [
|
||||
{ name = "aiosqlite" },
|
||||
|
||||
@@ -5,7 +5,7 @@ from langgraph_sdk.encryption.types import DecryptResult, EncryptionContext
|
||||
from langgraph_sdk.priority import Priority
|
||||
from langgraph_sdk.priority.types import PriorityContext, PrioritySettings
|
||||
|
||||
__version__ = "0.4.10"
|
||||
__version__ = "0.4.11"
|
||||
|
||||
__all__ = [
|
||||
"Auth",
|
||||
|
||||
@@ -151,7 +151,8 @@ class HttpClient:
|
||||
on_response: Callable[[httpx.Response], None] | None = None,
|
||||
reconnect_limit: int = 5,
|
||||
) -> Any:
|
||||
"""Send a request that automatically reconnects to Location header."""
|
||||
"""Send a request that reconnects to its Location header if the response is
|
||||
cut off, with backoff between attempts."""
|
||||
request_headers, content = await _aencode_json(json)
|
||||
if headers:
|
||||
request_headers.update(headers)
|
||||
@@ -181,13 +182,27 @@ class HttpClient:
|
||||
stacklevel=2,
|
||||
)
|
||||
await r.aclose()
|
||||
return await self.request_reconnect(
|
||||
loc,
|
||||
"GET",
|
||||
headers=headers,
|
||||
# don't pass on_response so it's only called once
|
||||
reconnect_limit=reconnect_limit - 1,
|
||||
)
|
||||
for attempt in range(reconnect_limit):
|
||||
await asyncio.sleep(_reconnect_delay(attempt, base=1.0, cap=5.0))
|
||||
# A restart or a deploy looks like a refused connection, a
|
||||
# dropped body, a 429 or a 5xx; other errors won't go away.
|
||||
try:
|
||||
return await self.request_reconnect(
|
||||
loc,
|
||||
"GET",
|
||||
headers=headers,
|
||||
# don't pass on_response so it's only called once
|
||||
reconnect_limit=0,
|
||||
)
|
||||
except httpx.HTTPStatusError as err:
|
||||
status = err.response.status_code
|
||||
if attempt == reconnect_limit - 1 or (
|
||||
status != 429 and status < 500
|
||||
):
|
||||
raise
|
||||
except httpx.TransportError:
|
||||
if attempt == reconnect_limit - 1:
|
||||
raise
|
||||
|
||||
async def stream(
|
||||
self,
|
||||
@@ -230,6 +245,17 @@ class HttpClient:
|
||||
}
|
||||
if last_event_id:
|
||||
current_headers["Last-Event-ID"] = last_event_id
|
||||
if reconnect_path is not None and not any(
|
||||
key.lower() == "last-event-id" and value
|
||||
for key, value in current_headers.items()
|
||||
):
|
||||
warnings.warn(
|
||||
"Reconnecting without a Last-Event-ID, so events sent while "
|
||||
"disconnected are not replayed. Only runs created with "
|
||||
"stream_resumable=True send the event IDs a reconnect resumes "
|
||||
"from.",
|
||||
stacklevel=2,
|
||||
)
|
||||
|
||||
current_method = method if reconnect_path is None else "GET"
|
||||
current_content = content if reconnect_path is None else None
|
||||
|
||||
@@ -152,7 +152,8 @@ class SyncHttpClient:
|
||||
on_response: Callable[[httpx.Response], None] | None = None,
|
||||
reconnect_limit: int = 5,
|
||||
) -> Any:
|
||||
"""Send a request that automatically reconnects to Location header."""
|
||||
"""Send a request that reconnects to its Location header if the response is
|
||||
cut off, with backoff between attempts."""
|
||||
request_headers, content = _encode_json(json)
|
||||
if headers:
|
||||
request_headers.update(headers)
|
||||
@@ -182,13 +183,27 @@ class SyncHttpClient:
|
||||
stacklevel=2,
|
||||
)
|
||||
r.close()
|
||||
return self.request_reconnect(
|
||||
loc,
|
||||
"GET",
|
||||
headers=headers,
|
||||
# don't pass on_response so it's only called once
|
||||
reconnect_limit=reconnect_limit - 1,
|
||||
)
|
||||
for attempt in range(reconnect_limit):
|
||||
time.sleep(_reconnect_delay(attempt, base=1.0, cap=5.0))
|
||||
# A restart or a deploy looks like a refused connection, a
|
||||
# dropped body, a 429 or a 5xx; other errors won't go away.
|
||||
try:
|
||||
return self.request_reconnect(
|
||||
loc,
|
||||
"GET",
|
||||
headers=headers,
|
||||
# don't pass on_response so it's only called once
|
||||
reconnect_limit=0,
|
||||
)
|
||||
except httpx.HTTPStatusError as err:
|
||||
status = err.response.status_code
|
||||
if attempt == reconnect_limit - 1 or (
|
||||
status != 429 and status < 500
|
||||
):
|
||||
raise
|
||||
except httpx.TransportError:
|
||||
if attempt == reconnect_limit - 1:
|
||||
raise
|
||||
|
||||
def stream(
|
||||
self,
|
||||
@@ -233,6 +248,17 @@ class SyncHttpClient:
|
||||
}
|
||||
if last_event_id:
|
||||
current_headers["Last-Event-ID"] = last_event_id
|
||||
if reconnect_path is not None and not any(
|
||||
key.lower() == "last-event-id" and value
|
||||
for key, value in current_headers.items()
|
||||
):
|
||||
warnings.warn(
|
||||
"Reconnecting without a Last-Event-ID, so events sent while "
|
||||
"disconnected are not replayed. Only runs created with "
|
||||
"stream_resumable=True send the event IDs a reconnect resumes "
|
||||
"from.",
|
||||
stacklevel=2,
|
||||
)
|
||||
|
||||
current_method = method if reconnect_path is None else "GET"
|
||||
current_content = content if reconnect_path is None else None
|
||||
|
||||
@@ -0,0 +1,17 @@
|
||||
import asyncio
|
||||
import time
|
||||
|
||||
import pytest
|
||||
|
||||
|
||||
@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
|
||||
@@ -1,7 +1,5 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import time
|
||||
from collections.abc import Callable, Iterator, Sequence
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
@@ -70,19 +68,6 @@ 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:
|
||||
@@ -501,9 +486,10 @@ _CALLER_CURSORS = pytest.mark.parametrize(
|
||||
def test_sync_stream_reconnects_without_a_cursor_after_an_empty_id(
|
||||
empty_id: bytes, headers: dict[str, str] | None
|
||||
):
|
||||
reconnect = _sync_reconnect(
|
||||
[b"id: 1\nevent: values\ndata: 1\n\n", empty_id], headers
|
||||
)
|
||||
with pytest.warns(UserWarning, match="Last-Event-ID"):
|
||||
reconnect = _sync_reconnect(
|
||||
[b"id: 1\nevent: values\ndata: 1\n\n", empty_id], headers
|
||||
)
|
||||
|
||||
assert "last-event-id" not in reconnect.headers
|
||||
|
||||
@@ -514,9 +500,10 @@ def test_sync_stream_reconnects_without_a_cursor_after_an_empty_id(
|
||||
async def test_async_stream_reconnects_without_a_cursor_after_an_empty_id(
|
||||
empty_id: bytes, headers: dict[str, str] | None
|
||||
):
|
||||
reconnect = await _async_reconnect(
|
||||
[b"id: 1\nevent: values\ndata: 1\n\n", empty_id], headers
|
||||
)
|
||||
with pytest.warns(UserWarning, match="Last-Event-ID"):
|
||||
reconnect = await _async_reconnect(
|
||||
[b"id: 1\nevent: values\ndata: 1\n\n", empty_id], headers
|
||||
)
|
||||
|
||||
assert "last-event-id" not in reconnect.headers
|
||||
|
||||
@@ -555,6 +542,45 @@ async def test_async_stream_keeps_the_caller_cursor_until_the_stream_sends_an_id
|
||||
assert reconnect.headers.get_list("last-event-id") == ["5"]
|
||||
|
||||
|
||||
# A reconnect that sends no Last-Event-ID can't replay what was sent while it was
|
||||
# disconnected, so it warns.
|
||||
_RECONNECT_CURSORS = pytest.mark.parametrize(
|
||||
("chunks", "headers", "warns"),
|
||||
[
|
||||
([b"event: values\ndata: 1\n\n"], None, True),
|
||||
([b"id: 1\nevent: values\ndata: 1\n\n"], None, False),
|
||||
([b"event: values\ndata: 1\n\n"], {"Last-Event-ID": "5"}, False),
|
||||
],
|
||||
ids=["no-cursor", "stream-cursor", "caller-cursor"],
|
||||
)
|
||||
|
||||
|
||||
@_RECONNECT_CURSORS
|
||||
@pytest.mark.usefixtures("no_sleep")
|
||||
def test_sync_reconnect_warns_only_without_a_cursor(
|
||||
chunks: list[bytes],
|
||||
headers: dict[str, str] | None,
|
||||
warns: bool,
|
||||
recwarn: pytest.WarningsRecorder,
|
||||
):
|
||||
_sync_reconnect(chunks, headers)
|
||||
|
||||
assert any("Last-Event-ID" in str(w.message) for w in recwarn) is warns
|
||||
|
||||
|
||||
@_RECONNECT_CURSORS
|
||||
@pytest.mark.usefixtures("no_sleep")
|
||||
async def test_async_reconnect_warns_only_without_a_cursor(
|
||||
chunks: list[bytes],
|
||||
headers: dict[str, str] | None,
|
||||
warns: bool,
|
||||
recwarn: pytest.WarningsRecorder,
|
||||
):
|
||||
await _async_reconnect(chunks, headers)
|
||||
|
||||
assert any("Last-Event-ID" in str(w.message) for w in recwarn) is warns
|
||||
|
||||
|
||||
# --- reconnects survive a server restart or a deploy ---
|
||||
|
||||
SSE = {"Content-Type": "text/event-stream"}
|
||||
@@ -772,6 +798,39 @@ async def test_async_long_stream_survives_many_spaced_out_drops():
|
||||
assert [p.data for p in parts] == [0, 1, 2, 3, 4, 5, 6, None]
|
||||
|
||||
|
||||
# Each reconnect resumes after the newest event the stream got, so a reconnect
|
||||
# loop on a resumable run replays only what it missed. A failed reconnect keeps
|
||||
# the cursor for the next one.
|
||||
RECONNECT_LOOPS = {
|
||||
"progress_between_drops": (
|
||||
lambda stream_cls: _long_stream_steps(stream_cls, drops=3),
|
||||
[None, "0", "1", "2"],
|
||||
),
|
||||
"failed_reconnect": (
|
||||
lambda stream_cls: _transient_steps(stream_cls, _status(502)),
|
||||
[None, "1", "1"],
|
||||
),
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.usefixtures("no_sleep")
|
||||
@pytest.mark.parametrize(
|
||||
("steps", "cursors"), RECONNECT_LOOPS.values(), ids=RECONNECT_LOOPS.keys()
|
||||
)
|
||||
def test_sync_reconnect_loop_resumes_after_the_newest_event(steps, cursors):
|
||||
_, seen = _run_sync(steps(ListByteStream))
|
||||
assert [r.headers.get("last-event-id") for r in seen] == cursors
|
||||
|
||||
|
||||
@pytest.mark.usefixtures("no_sleep")
|
||||
@pytest.mark.parametrize(
|
||||
("steps", "cursors"), RECONNECT_LOOPS.values(), ids=RECONNECT_LOOPS.keys()
|
||||
)
|
||||
async def test_async_reconnect_loop_resumes_after_the_newest_event(steps, cursors):
|
||||
_, seen = await _run_async(steps(AsyncListByteStream))
|
||||
assert [r.headers.get("last-event-id") for r in seen] == cursors
|
||||
|
||||
|
||||
# Reconnects should back off rather than hammer the server.
|
||||
|
||||
|
||||
|
||||
@@ -1,7 +1,8 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
from collections.abc import AsyncIterator, Iterator
|
||||
from collections.abc import AsyncIterator, Callable, Iterator, Sequence
|
||||
from typing import Any
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
@@ -24,7 +25,16 @@ class _BodyDropsMidRead(httpx.SyncByteStream, httpx.AsyncByteStream):
|
||||
raise httpx.ReadError("connection dropped")
|
||||
|
||||
|
||||
def _wait_handler(requests: list[httpx.Request]):
|
||||
def _done(_request: httpx.Request) -> httpx.Response:
|
||||
return httpx.Response(200, json={"done": True})
|
||||
|
||||
|
||||
def _wait_handler(
|
||||
requests: list[httpx.Request],
|
||||
reconnects: Sequence[Callable[[httpx.Request], httpx.Response]] = (_done,),
|
||||
):
|
||||
"""POST drops mid-body; the reconnects answer from `reconnects` in order."""
|
||||
|
||||
def handler(request: httpx.Request) -> httpx.Response:
|
||||
requests.append(request)
|
||||
if request.method == "POST":
|
||||
@@ -33,7 +43,7 @@ def _wait_handler(requests: list[httpx.Request]):
|
||||
headers={"Location": "/threads/t1/runs/r1/join"},
|
||||
stream=_BodyDropsMidRead(),
|
||||
)
|
||||
return httpx.Response(200, json={"done": True})
|
||||
return reconnects[len(requests) - 2](request)
|
||||
|
||||
return handler
|
||||
|
||||
@@ -46,6 +56,7 @@ def _assert_bodyless_get(request: httpx.Request) -> None:
|
||||
assert request.read() == b""
|
||||
|
||||
|
||||
@pytest.mark.usefixtures("no_sleep")
|
||||
def test_sync_wait_reconnects_with_a_bodyless_get():
|
||||
requests: list[httpx.Request] = []
|
||||
transport = httpx.MockTransport(_wait_handler(requests))
|
||||
@@ -58,6 +69,7 @@ def test_sync_wait_reconnects_with_a_bodyless_get():
|
||||
_assert_bodyless_get(requests[1])
|
||||
|
||||
|
||||
@pytest.mark.usefixtures("no_sleep")
|
||||
async def test_async_wait_reconnects_with_a_bodyless_get():
|
||||
requests: list[httpx.Request] = []
|
||||
transport = httpx.MockTransport(_wait_handler(requests))
|
||||
@@ -72,6 +84,89 @@ async def test_async_wait_reconnects_with_a_bodyless_get():
|
||||
_assert_bodyless_get(requests[1])
|
||||
|
||||
|
||||
def _connect_error(request: httpx.Request) -> httpx.Response:
|
||||
raise httpx.ConnectError("[Errno 61] Connection refused", request=request)
|
||||
|
||||
|
||||
async def _wait(
|
||||
handler: Callable[[httpx.Request], httpx.Response], *, asynchronous: bool
|
||||
) -> Any:
|
||||
transport = httpx.MockTransport(handler)
|
||||
if asynchronous:
|
||||
async with httpx.AsyncClient(
|
||||
transport=transport, base_url="https://example.com"
|
||||
) as client:
|
||||
return await RunsClient(HttpClient(client)).wait("t1", "agent", input={})
|
||||
with httpx.Client(transport=transport, base_url="https://example.com") as client:
|
||||
return SyncRunsClient(SyncHttpClient(client)).wait("t1", "agent", input={})
|
||||
|
||||
|
||||
# What a server restart or a deploy looks like to the reconnect GET: each one is
|
||||
# retried after a backoff.
|
||||
RECONNECT_FAILURES = {
|
||||
"429": lambda request: httpx.Response(429, request=request),
|
||||
"503": lambda request: httpx.Response(503, request=request),
|
||||
"connect_error": _connect_error,
|
||||
"body_drops": lambda _request: httpx.Response(
|
||||
200,
|
||||
headers={"Location": "/threads/t1/runs/r1/join"},
|
||||
stream=_BodyDropsMidRead(),
|
||||
),
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.parametrize("asynchronous", [False, True])
|
||||
@pytest.mark.parametrize(
|
||||
"failure", RECONNECT_FAILURES.values(), ids=RECONNECT_FAILURES.keys()
|
||||
)
|
||||
async def test_wait_retries_a_failed_reconnect_after_a_backoff(
|
||||
failure, asynchronous, no_sleep
|
||||
):
|
||||
requests: list[httpx.Request] = []
|
||||
with pytest.warns(UserWarning, match="attempting reconnect"):
|
||||
result = await _wait(
|
||||
_wait_handler(requests, [failure, _done]), asynchronous=asynchronous
|
||||
)
|
||||
|
||||
assert result == {"done": True}
|
||||
assert len(requests) == 3
|
||||
assert len(no_sleep) == 2 and min(no_sleep) >= 1, no_sleep
|
||||
|
||||
|
||||
@pytest.mark.usefixtures("no_sleep")
|
||||
@pytest.mark.parametrize("asynchronous", [False, True])
|
||||
async def test_wait_raises_after_five_failed_reconnects(asynchronous):
|
||||
requests: list[httpx.Request] = []
|
||||
with (
|
||||
pytest.warns(UserWarning, match="attempting reconnect"),
|
||||
pytest.raises(httpx.HTTPStatusError, match="503"),
|
||||
):
|
||||
await _wait(
|
||||
_wait_handler(requests, [RECONNECT_FAILURES["503"]] * 5),
|
||||
asynchronous=asynchronous,
|
||||
)
|
||||
|
||||
assert len(requests) == 6
|
||||
|
||||
|
||||
@pytest.mark.usefixtures("no_sleep")
|
||||
@pytest.mark.parametrize("asynchronous", [False, True])
|
||||
async def test_wait_does_not_retry_a_client_error(asynchronous):
|
||||
requests: list[httpx.Request] = []
|
||||
with (
|
||||
pytest.warns(UserWarning, match="attempting reconnect"),
|
||||
pytest.raises(httpx.HTTPStatusError, match="404"),
|
||||
):
|
||||
await _wait(
|
||||
_wait_handler(
|
||||
requests, [lambda request: httpx.Response(404, request=request)]
|
||||
),
|
||||
asynchronous=asynchronous,
|
||||
)
|
||||
|
||||
assert len(requests) == 2
|
||||
|
||||
|
||||
def test_sync_wait_raises_the_run_error():
|
||||
requests: list[httpx.Request] = []
|
||||
|
||||
|
||||
Generated
+1
-1
@@ -301,7 +301,7 @@ wheels = [
|
||||
|
||||
[[package]]
|
||||
name = "langgraph"
|
||||
version = "1.2.18"
|
||||
version = "1.2.19"
|
||||
source = { editable = "../langgraph" }
|
||||
dependencies = [
|
||||
{ name = "langchain-core" },
|
||||
|
||||
Reference in new issue
Block a user