mirror of
https://github.com/langchain-ai/langgraph.git
synced 2026-11-03 21:01:24 +01:00
Compare commits
5
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
6e144fb893 | ||
|
|
1f0cf816d5 | ||
|
|
6f2df053f4 | ||
|
|
ffc3336367 | ||
|
|
5f4b28a9f7 |
No files matched your search
@@ -11,6 +11,7 @@ import click.exceptions
|
||||
|
||||
import langgraph_cli.config
|
||||
import langgraph_cli.docker
|
||||
from langgraph_cli import langsmith_auth
|
||||
from langgraph_cli.analytics import log_command
|
||||
from langgraph_cli.config import Config
|
||||
from langgraph_cli.constants import DEFAULT_CONFIG, DEFAULT_PORT
|
||||
@@ -233,6 +234,13 @@ def cli():
|
||||
pass
|
||||
|
||||
|
||||
@cli.command(help="Log in to LangSmith in the browser (no API key needed).")
|
||||
@click.option("--no-browser", is_flag=True, help="Print the URL instead of opening it.")
|
||||
def login(no_browser: bool) -> None:
|
||||
api_url = langsmith_auth.resolve_api_url(os.environ.get("LANGSMITH_ENDPOINT"))
|
||||
langsmith_auth.login(api_url, open_browser=not no_browser)
|
||||
|
||||
|
||||
# Wire the deploy group (defined in deploy.py) into the top-level CLI.
|
||||
cli.add_command(deploy)
|
||||
|
||||
|
||||
@@ -1,6 +1,7 @@
|
||||
"""Deploy command and subcommands for the LangGraph CLI."""
|
||||
|
||||
import base64
|
||||
import functools
|
||||
import json as json_mod
|
||||
import os
|
||||
import pathlib
|
||||
@@ -19,6 +20,7 @@ import click.exceptions
|
||||
from dotenv import dotenv_values, set_key
|
||||
|
||||
import langgraph_cli.config
|
||||
from langgraph_cli import langsmith_auth
|
||||
from langgraph_cli.analytics import log_command
|
||||
from langgraph_cli.config import Config
|
||||
from langgraph_cli.constants import DEFAULT_CONFIG
|
||||
@@ -1982,17 +1984,6 @@ def _create_host_backend_client(
|
||||
if val:
|
||||
resolved_api_key = val
|
||||
break
|
||||
if not resolved_api_key:
|
||||
if effective_no_input:
|
||||
raise click.ClickException(
|
||||
"No LangSmith API key found. Set LANGSMITH_API_KEY in the "
|
||||
"environment or .env file."
|
||||
)
|
||||
click.secho(
|
||||
"No LangSmith API key found. Create one at Settings > API Keys in LangSmith.",
|
||||
fg="yellow",
|
||||
)
|
||||
resolved_api_key = click.prompt("Enter LangSmith API key", hide_input=True)
|
||||
tenant_id = env_vars.get("LANGSMITH_TENANT_ID") or os.environ.get(
|
||||
"LANGSMITH_TENANT_ID"
|
||||
)
|
||||
@@ -2000,8 +1991,31 @@ def _create_host_backend_client(
|
||||
"LANGSMITH_ENDPOINT"
|
||||
)
|
||||
endpoints = ControlPlaneEndpoints.resolve(host_url, langsmith_endpoint)
|
||||
|
||||
access_token: Callable[[], str] | None = None
|
||||
if not resolved_api_key:
|
||||
profile = langsmith_auth.load_profile()
|
||||
api_url = langsmith_auth.resolve_api_url(langsmith_endpoint)
|
||||
if profile.get("api_key"):
|
||||
resolved_api_key = profile["api_key"]
|
||||
else:
|
||||
if not any((profile.get("oauth") or {}).values()):
|
||||
if effective_no_input:
|
||||
raise click.ClickException(
|
||||
"No LangSmith credentials found. Set LANGSMITH_API_KEY "
|
||||
"in the environment or .env file (an API key), or run "
|
||||
"`langgraph login`."
|
||||
)
|
||||
click.secho("You are not logged in to LangSmith.", fg="yellow")
|
||||
langsmith_auth.login(api_url)
|
||||
profile = langsmith_auth.load_profile()
|
||||
access_token = functools.partial(langsmith_auth.get_access_token, api_url)
|
||||
tenant_id = tenant_id or profile.get("workspace_id")
|
||||
return HostBackendClient(
|
||||
endpoints.control_plane_url, resolved_api_key, tenant_id=tenant_id
|
||||
endpoints.control_plane_url,
|
||||
resolved_api_key,
|
||||
tenant_id=tenant_id,
|
||||
**({"access_token": access_token} if access_token else {}),
|
||||
)
|
||||
|
||||
|
||||
|
||||
@@ -2,7 +2,7 @@
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Mapping
|
||||
from collections.abc import Callable, Mapping
|
||||
from dataclasses import dataclass
|
||||
from typing import Any, Literal
|
||||
from urllib.parse import urlparse
|
||||
@@ -130,17 +130,18 @@ class HostBackendClient:
|
||||
def __init__(
|
||||
self,
|
||||
base_url: str,
|
||||
api_key: str,
|
||||
api_key: str | None,
|
||||
tenant_id: str | None = None,
|
||||
*,
|
||||
access_token: Callable[[], str] | None = None,
|
||||
transport: httpx.BaseTransport | None = None,
|
||||
):
|
||||
if not base_url:
|
||||
raise click.UsageError("Host backend URL is required")
|
||||
headers: dict[str, str] = {
|
||||
"X-Api-Key": api_key,
|
||||
"Accept": "application/json",
|
||||
}
|
||||
self._access_token = access_token
|
||||
headers: dict[str, str] = {"Accept": "application/json"}
|
||||
if api_key:
|
||||
headers["X-Api-Key"] = api_key
|
||||
if tenant_id:
|
||||
headers["X-Tenant-ID"] = tenant_id
|
||||
self._endpoints = ControlPlaneEndpoints.from_control_plane_url(base_url)
|
||||
@@ -170,8 +171,15 @@ class HostBackendClient:
|
||||
payload: dict[str, Any] | None = None,
|
||||
params: dict[str, Any] | None = None,
|
||||
) -> Any:
|
||||
headers = (
|
||||
{"Authorization": f"Bearer {self._access_token()}"}
|
||||
if self._access_token
|
||||
else None
|
||||
)
|
||||
try:
|
||||
resp = self._client.request(method, path, json=payload, params=params)
|
||||
resp = self._client.request(
|
||||
method, path, json=payload, params=params, headers=headers
|
||||
)
|
||||
resp.raise_for_status()
|
||||
except httpx.HTTPStatusError as err:
|
||||
detail = _error_detail(err.response)
|
||||
|
||||
@@ -0,0 +1,225 @@
|
||||
"""LangSmith OAuth login, shared with the ``langsmith`` CLI via ``~/.langsmith/config.json``."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import datetime
|
||||
import json
|
||||
import os
|
||||
import pathlib
|
||||
import time
|
||||
import webbrowser
|
||||
from typing import Any
|
||||
from urllib.parse import urlparse
|
||||
|
||||
import click
|
||||
import httpx
|
||||
|
||||
CLIENT_ID = "langsmith-cli"
|
||||
DEFAULT_API_URL = "https://api.smith.langchain.com"
|
||||
_WELL_KNOWN = "/.well-known/oauth-authorization-server"
|
||||
_REFRESH_LEEWAY = datetime.timedelta(seconds=60)
|
||||
_DEVICE_GRANT = "urn:ietf:params:oauth:grant-type:device_code"
|
||||
|
||||
|
||||
def _config_path() -> pathlib.Path:
|
||||
override = os.environ.get("LANGSMITH_CONFIG_FILE")
|
||||
if override:
|
||||
return pathlib.Path(override)
|
||||
return pathlib.Path.home() / ".langsmith" / "config.json"
|
||||
|
||||
|
||||
def _read_config() -> dict[str, Any]:
|
||||
try:
|
||||
config = json.loads(_config_path().read_text(encoding="utf-8"))
|
||||
except (OSError, ValueError):
|
||||
return {}
|
||||
return config if isinstance(config, dict) else {}
|
||||
|
||||
|
||||
def _write_config(config: dict[str, Any]) -> None:
|
||||
path = _config_path()
|
||||
path.parent.mkdir(mode=0o700, parents=True, exist_ok=True)
|
||||
tmp = path.with_name(path.name + ".tmp")
|
||||
tmp.write_text(json.dumps(config, indent=2) + "\n", encoding="utf-8")
|
||||
os.chmod(tmp, 0o600)
|
||||
os.replace(tmp, path)
|
||||
|
||||
|
||||
def _profile_name(config: dict[str, Any]) -> str:
|
||||
return (
|
||||
os.environ.get("LANGSMITH_PROFILE")
|
||||
or config.get("current_profile")
|
||||
or "default"
|
||||
)
|
||||
|
||||
|
||||
def load_profile() -> dict[str, Any]:
|
||||
"""Return the active profile from the LangSmith config file, or ``{}``."""
|
||||
config = _read_config()
|
||||
profile = (config.get("profiles") or {}).get(_profile_name(config))
|
||||
return profile if isinstance(profile, dict) else {}
|
||||
|
||||
|
||||
def resolve_api_url(endpoint: str | None = None) -> str:
|
||||
"""Pick the LangSmith API URL: explicit endpoint, then profile, then cloud."""
|
||||
return endpoint or load_profile().get("api_url") or DEFAULT_API_URL
|
||||
|
||||
|
||||
def _save_tokens(token: dict[str, Any], api_url: str | None = None) -> None:
|
||||
config = _read_config()
|
||||
name = _profile_name(config)
|
||||
profile = config.setdefault("profiles", {}).setdefault(name, {})
|
||||
if api_url:
|
||||
profile["api_url"] = api_url
|
||||
oauth = profile.setdefault("oauth", {})
|
||||
oauth["access_token"] = token["access_token"]
|
||||
if token.get("refresh_token"):
|
||||
oauth["refresh_token"] = token["refresh_token"]
|
||||
if token.get("expires_in"):
|
||||
expires_at = datetime.datetime.now(datetime.timezone.utc) + (
|
||||
datetime.timedelta(seconds=int(token["expires_in"]))
|
||||
)
|
||||
oauth["expires_at"] = expires_at.strftime("%Y-%m-%dT%H:%M:%SZ")
|
||||
config.setdefault("current_profile", name)
|
||||
_write_config(config)
|
||||
|
||||
|
||||
def _oauth_endpoints(api_url: str) -> dict[str, str]:
|
||||
"""Discover the OAuth endpoints, falling back to ``<root>/oauth/*``."""
|
||||
root = api_url.rstrip("/")
|
||||
for suffix in ("/api/v1", "/api"):
|
||||
if root.endswith(suffix):
|
||||
root = root[: -len(suffix)]
|
||||
break
|
||||
for base in (root, f"{root}/api"):
|
||||
try:
|
||||
resp = httpx.get(base + _WELL_KNOWN, timeout=10)
|
||||
doc = resp.json() if resp.is_success else None
|
||||
except (httpx.HTTPError, ValueError):
|
||||
continue
|
||||
if (
|
||||
isinstance(doc, dict)
|
||||
and str(doc.get("issuer", "")).rstrip("/") == base
|
||||
and all(
|
||||
urlparse(str(doc.get(key, ""))).netloc == urlparse(base).netloc
|
||||
for key in ("token_endpoint", "device_authorization_endpoint")
|
||||
)
|
||||
):
|
||||
return {
|
||||
"issuer": base,
|
||||
"token": doc["token_endpoint"],
|
||||
"device": doc["device_authorization_endpoint"],
|
||||
}
|
||||
return {
|
||||
"issuer": root,
|
||||
"token": f"{root}/oauth/token",
|
||||
"device": f"{root}/oauth/device/code",
|
||||
}
|
||||
|
||||
|
||||
def _expiring(oauth: dict[str, Any]) -> bool:
|
||||
if not oauth.get("access_token"):
|
||||
return True
|
||||
expires_at = oauth.get("expires_at")
|
||||
if not expires_at:
|
||||
return False
|
||||
try:
|
||||
expiry = datetime.datetime.fromisoformat(expires_at.replace("Z", "+00:00"))
|
||||
except ValueError:
|
||||
return True
|
||||
if expiry.tzinfo is None:
|
||||
expiry = expiry.replace(tzinfo=datetime.timezone.utc)
|
||||
return expiry <= datetime.datetime.now(datetime.timezone.utc) + _REFRESH_LEEWAY
|
||||
|
||||
|
||||
def get_access_token(api_url: str) -> str:
|
||||
"""Return a valid access token from the profile, refreshing it if needed."""
|
||||
oauth = load_profile().get("oauth") or {}
|
||||
if not _expiring(oauth):
|
||||
return oauth["access_token"]
|
||||
refresh_token = oauth.get("refresh_token")
|
||||
if not refresh_token:
|
||||
raise click.ClickException(
|
||||
"LangSmith session expired. Run `langgraph login` to sign in again."
|
||||
)
|
||||
endpoints = _oauth_endpoints(api_url)
|
||||
try:
|
||||
resp = httpx.post(
|
||||
endpoints["token"],
|
||||
data={
|
||||
"grant_type": "refresh_token",
|
||||
"client_id": CLIENT_ID,
|
||||
"refresh_token": refresh_token,
|
||||
},
|
||||
timeout=10,
|
||||
)
|
||||
resp.raise_for_status()
|
||||
token = resp.json()
|
||||
except (httpx.HTTPError, ValueError):
|
||||
raise click.ClickException(
|
||||
"Could not refresh your LangSmith session. "
|
||||
"Run `langgraph login` to sign in again."
|
||||
) from None
|
||||
if not isinstance(token, dict) or not token.get("access_token"):
|
||||
raise click.ClickException(
|
||||
"Could not refresh your LangSmith session. "
|
||||
"Run `langgraph login` to sign in again."
|
||||
)
|
||||
_save_tokens(token)
|
||||
return token["access_token"]
|
||||
|
||||
|
||||
def login(api_url: str, *, open_browser: bool = True) -> None:
|
||||
"""Sign in with the OAuth device flow and store tokens in the profile."""
|
||||
endpoints = _oauth_endpoints(api_url)
|
||||
try:
|
||||
resp = httpx.post(
|
||||
endpoints["device"],
|
||||
data={"client_id": CLIENT_ID, "resource": endpoints["issuer"]},
|
||||
timeout=10,
|
||||
)
|
||||
resp.raise_for_status()
|
||||
device = resp.json()
|
||||
device_code, user_code = device["device_code"], device["user_code"]
|
||||
verification_url = (
|
||||
device.get("verification_uri_complete") or device["verification_uri"]
|
||||
)
|
||||
except (httpx.HTTPError, ValueError, KeyError):
|
||||
raise click.ClickException(
|
||||
f"Could not start LangSmith login against {api_url}."
|
||||
) from None
|
||||
|
||||
click.echo(f"Open this URL to authorize the CLI:\n{verification_url}\n")
|
||||
click.echo(f"Code: {user_code}\n")
|
||||
if open_browser:
|
||||
webbrowser.open(verification_url)
|
||||
click.echo("Waiting for authorization...")
|
||||
|
||||
interval = max(int(device.get("interval") or 5), 5)
|
||||
deadline = time.monotonic() + int(device.get("expires_in") or 600)
|
||||
while time.monotonic() < deadline:
|
||||
time.sleep(interval)
|
||||
try:
|
||||
resp = httpx.post(
|
||||
endpoints["token"],
|
||||
data={
|
||||
"grant_type": _DEVICE_GRANT,
|
||||
"client_id": CLIENT_ID,
|
||||
"device_code": device_code,
|
||||
"resource": endpoints["issuer"],
|
||||
},
|
||||
timeout=10,
|
||||
)
|
||||
body = resp.json()
|
||||
except (httpx.HTTPError, ValueError):
|
||||
continue
|
||||
if resp.is_success and body.get("access_token"):
|
||||
_save_tokens(body, api_url)
|
||||
click.secho("Logged in to LangSmith.", fg="green")
|
||||
return
|
||||
error = body.get("error") if isinstance(body, dict) else None
|
||||
if error == "slow_down":
|
||||
interval += 5
|
||||
elif error != "authorization_pending":
|
||||
raise click.ClickException(f"LangSmith login failed: {error or 'error'}.")
|
||||
raise click.ClickException("Timed out waiting for LangSmith login.")
|
||||
@@ -578,12 +578,13 @@ def test_a_deployed_result_prints_the_deployment_url(
|
||||
class TestCreateHostBackendClientNoInput:
|
||||
def test_raises_when_no_api_key_and_no_input(self, monkeypatch, tmp_path):
|
||||
|
||||
monkeypatch.setenv("LANGSMITH_CONFIG_FILE", str(tmp_path / "missing.json"))
|
||||
monkeypatch.setattr(deploy_mod, "_no_input", True)
|
||||
monkeypatch.delenv("LANGSMITH_API_KEY", raising=False)
|
||||
monkeypatch.delenv("LANGCHAIN_API_KEY", raising=False)
|
||||
monkeypatch.delenv("LANGGRAPH_HOST_API_KEY", raising=False)
|
||||
|
||||
with pytest.raises(click.ClickException, match="API key"):
|
||||
with pytest.raises(click.ClickException, match="credentials"):
|
||||
_create_host_backend_client(
|
||||
host_url="https://api.example.com",
|
||||
api_key=None,
|
||||
|
||||
@@ -0,0 +1,191 @@
|
||||
import json
|
||||
from datetime import datetime, timedelta, timezone
|
||||
|
||||
import click
|
||||
import httpx
|
||||
import pytest
|
||||
|
||||
from langgraph_cli import deploy as deploy_mod
|
||||
from langgraph_cli import langsmith_auth
|
||||
from langgraph_cli.deploy import _create_host_backend_client
|
||||
from langgraph_cli.host_backend import HostBackendClient
|
||||
|
||||
API_URL = "https://api.smith.langchain.com"
|
||||
|
||||
|
||||
def _iso(delta: timedelta) -> str:
|
||||
return (datetime.now(timezone.utc) + delta).strftime("%Y-%m-%dT%H:%M:%SZ")
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def config_file(tmp_path, monkeypatch):
|
||||
path = tmp_path / "config.json"
|
||||
monkeypatch.setenv("LANGSMITH_CONFIG_FILE", str(path))
|
||||
monkeypatch.delenv("LANGSMITH_PROFILE", raising=False)
|
||||
for name in ("LANGGRAPH_HOST_API_KEY", "LANGSMITH_API_KEY", "LANGCHAIN_API_KEY"):
|
||||
monkeypatch.delenv(name, raising=False)
|
||||
return path
|
||||
|
||||
|
||||
def _write_profile(path, **oauth):
|
||||
path.write_text(
|
||||
json.dumps(
|
||||
{
|
||||
"current_profile": "default",
|
||||
"profiles": {
|
||||
"default": {
|
||||
"api_url": API_URL,
|
||||
"workspace_id": "ws-1",
|
||||
"oauth": oauth,
|
||||
}
|
||||
},
|
||||
}
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
def _no_discovery(monkeypatch):
|
||||
def get(*_a, **_k):
|
||||
raise httpx.ConnectError("no discovery")
|
||||
|
||||
monkeypatch.setattr(httpx, "get", get)
|
||||
|
||||
|
||||
def test_valid_token_is_returned_without_network(config_file, monkeypatch):
|
||||
_write_profile(
|
||||
config_file, access_token="tok", expires_at=_iso(timedelta(minutes=5))
|
||||
)
|
||||
|
||||
def boom(*_a, **_k):
|
||||
raise AssertionError("no network expected")
|
||||
|
||||
monkeypatch.setattr(httpx, "get", boom)
|
||||
monkeypatch.setattr(httpx, "post", boom)
|
||||
|
||||
assert langsmith_auth.get_access_token(API_URL) == "tok"
|
||||
|
||||
|
||||
def test_expired_token_is_refreshed_and_saved(config_file, monkeypatch):
|
||||
_write_profile(
|
||||
config_file,
|
||||
access_token="old",
|
||||
refresh_token="refresh-1",
|
||||
expires_at=_iso(timedelta(seconds=-5)),
|
||||
)
|
||||
_no_discovery(monkeypatch)
|
||||
seen = {}
|
||||
|
||||
def post(url, data, timeout):
|
||||
seen.update(url=url, data=data)
|
||||
return httpx.Response(
|
||||
200,
|
||||
json={
|
||||
"access_token": "new",
|
||||
"refresh_token": "refresh-2",
|
||||
"expires_in": 300,
|
||||
},
|
||||
request=httpx.Request("POST", url),
|
||||
)
|
||||
|
||||
monkeypatch.setattr(httpx, "post", post)
|
||||
|
||||
assert langsmith_auth.get_access_token(API_URL) == "new"
|
||||
assert seen["url"] == f"{API_URL}/oauth/token"
|
||||
assert seen["data"]["refresh_token"] == "refresh-1"
|
||||
saved = json.loads(config_file.read_text())["profiles"]["default"]["oauth"]
|
||||
assert saved["access_token"] == "new"
|
||||
assert saved["refresh_token"] == "refresh-2"
|
||||
|
||||
|
||||
def test_expired_token_without_refresh_token_asks_to_log_in(config_file):
|
||||
_write_profile(
|
||||
config_file, access_token="old", expires_at=_iso(timedelta(seconds=-5))
|
||||
)
|
||||
with pytest.raises(click.ClickException, match="langgraph login"):
|
||||
langsmith_auth.get_access_token(API_URL)
|
||||
|
||||
|
||||
def test_login_stores_tokens_after_pending(config_file, monkeypatch):
|
||||
_no_discovery(monkeypatch)
|
||||
monkeypatch.setattr(langsmith_auth.time, "sleep", lambda _s: None)
|
||||
monkeypatch.setattr(langsmith_auth.webbrowser, "open", lambda _u: True)
|
||||
token_responses = iter(
|
||||
[
|
||||
{"error": "authorization_pending"},
|
||||
{"access_token": "tok", "refresh_token": "r", "expires_in": 300},
|
||||
]
|
||||
)
|
||||
|
||||
def post(url, data, timeout):
|
||||
request = httpx.Request("POST", url)
|
||||
if url.endswith("/oauth/device/code"):
|
||||
return httpx.Response(
|
||||
200,
|
||||
json={
|
||||
"device_code": "dc",
|
||||
"user_code": "ABCD",
|
||||
"verification_uri": "https://example.com/activate",
|
||||
"expires_in": 600,
|
||||
"interval": 5,
|
||||
},
|
||||
request=request,
|
||||
)
|
||||
body = next(token_responses)
|
||||
return httpx.Response(
|
||||
400 if "error" in body else 200, json=body, request=request
|
||||
)
|
||||
|
||||
monkeypatch.setattr(httpx, "post", post)
|
||||
|
||||
langsmith_auth.login(API_URL)
|
||||
|
||||
profile = json.loads(config_file.read_text())["profiles"]["default"]
|
||||
assert profile["api_url"] == API_URL
|
||||
assert profile["oauth"]["access_token"] == "tok"
|
||||
|
||||
|
||||
def test_client_sends_bearer_token_per_request():
|
||||
seen = []
|
||||
|
||||
def handler(req: httpx.Request) -> httpx.Response:
|
||||
seen.append(dict(req.headers))
|
||||
return httpx.Response(200, json={"resources": []})
|
||||
|
||||
tokens = iter(["t1", "t2"])
|
||||
client = HostBackendClient(
|
||||
"https://api.example.com",
|
||||
None,
|
||||
access_token=lambda: next(tokens),
|
||||
transport=httpx.MockTransport(handler),
|
||||
)
|
||||
client.list_deployments()
|
||||
client.list_deployments()
|
||||
|
||||
assert [h["authorization"] for h in seen] == ["Bearer t1", "Bearer t2"]
|
||||
assert "x-api-key" not in seen[0]
|
||||
|
||||
|
||||
def test_factory_uses_logged_in_profile(config_file, monkeypatch):
|
||||
_write_profile(
|
||||
config_file, access_token="tok", expires_at=_iso(timedelta(minutes=5))
|
||||
)
|
||||
monkeypatch.setattr(deploy_mod, "_no_input", True)
|
||||
monkeypatch.delenv("LANGSMITH_TENANT_ID", raising=False)
|
||||
monkeypatch.delenv("LANGSMITH_ENDPOINT", raising=False)
|
||||
|
||||
client = _create_host_backend_client(
|
||||
host_url="https://api.example.com", api_key=None, env_vars={}
|
||||
)
|
||||
|
||||
assert client._access_token is not None
|
||||
assert client._access_token() == "tok"
|
||||
assert client._client.headers["x-tenant-id"] == "ws-1"
|
||||
assert "x-api-key" not in client._client.headers
|
||||
|
||||
|
||||
def test_factory_without_credentials_and_no_input_fails(config_file, monkeypatch):
|
||||
monkeypatch.setattr(deploy_mod, "_no_input", True)
|
||||
with pytest.raises(click.ClickException, match="langgraph login"):
|
||||
_create_host_backend_client(
|
||||
host_url="https://api.example.com", api_key=None, env_vars={}
|
||||
)
|
||||
@@ -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)])
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -2880,6 +2880,18 @@ 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_handler(h)
|
||||
]
|
||||
callback_manager.inheritable_handlers = [
|
||||
h
|
||||
for h in callback_manager.inheritable_handlers
|
||||
if not _is_stream_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 +2910,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 +2921,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:
|
||||
@@ -3009,6 +3009,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 +3295,18 @@ 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_handler(h)
|
||||
]
|
||||
callback_manager.inheritable_handlers = [
|
||||
h
|
||||
for h in callback_manager.inheritable_handlers
|
||||
if not _is_stream_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(
|
||||
@@ -3325,10 +3341,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 +3352,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:
|
||||
@@ -3463,6 +3466,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 +4467,10 @@ def _coerce_checkpoint_values(payload: Any, mapper: Callable[[Any], Any]) -> Non
|
||||
payload["values"] = mapper(payload["values"])
|
||||
|
||||
|
||||
def _is_stream_handler(handler: object) -> bool:
|
||||
return isinstance(handler, (StreamMessagesHandler, StreamToolCallHandler))
|
||||
|
||||
|
||||
def _resolve_parent_ns(
|
||||
graph_config: RunnableConfig | None, call_config: RunnableConfig | None
|
||||
) -> tuple[str, ...]:
|
||||
|
||||
@@ -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"]
|
||||
@@ -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] = []
|
||||
|
||||
|
||||
Reference in new issue
Block a user