mirror of
https://github.com/langchain-ai/langgraph.git
synced 2026-10-31 03:11:16 +01:00
Compare commits
4
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
5bd231cbfd | ||
|
|
3eaf06c67c | ||
|
|
6fab323753 | ||
|
|
8d5637b2dd |
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" },
|
||||
|
||||
@@ -11,7 +11,6 @@ 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
|
||||
@@ -234,13 +233,6 @@ 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,7 +1,6 @@
|
||||
"""Deploy command and subcommands for the LangGraph CLI."""
|
||||
|
||||
import base64
|
||||
import functools
|
||||
import json as json_mod
|
||||
import os
|
||||
import pathlib
|
||||
@@ -20,7 +19,6 @@ 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
|
||||
@@ -1984,6 +1982,17 @@ 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"
|
||||
)
|
||||
@@ -1991,31 +2000,8 @@ 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,
|
||||
**({"access_token": access_token} if access_token else {}),
|
||||
endpoints.control_plane_url, resolved_api_key, tenant_id=tenant_id
|
||||
)
|
||||
|
||||
|
||||
|
||||
@@ -2,7 +2,7 @@
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Callable, Mapping
|
||||
from collections.abc import Mapping
|
||||
from dataclasses import dataclass
|
||||
from typing import Any, Literal
|
||||
from urllib.parse import urlparse
|
||||
@@ -130,18 +130,17 @@ class HostBackendClient:
|
||||
def __init__(
|
||||
self,
|
||||
base_url: str,
|
||||
api_key: str | None,
|
||||
api_key: str,
|
||||
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")
|
||||
self._access_token = access_token
|
||||
headers: dict[str, str] = {"Accept": "application/json"}
|
||||
if api_key:
|
||||
headers["X-Api-Key"] = api_key
|
||||
headers: dict[str, str] = {
|
||||
"X-Api-Key": api_key,
|
||||
"Accept": "application/json",
|
||||
}
|
||||
if tenant_id:
|
||||
headers["X-Tenant-ID"] = tenant_id
|
||||
self._endpoints = ControlPlaneEndpoints.from_control_plane_url(base_url)
|
||||
@@ -171,15 +170,8 @@ 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, headers=headers
|
||||
)
|
||||
resp = self._client.request(method, path, json=payload, params=params)
|
||||
resp.raise_for_status()
|
||||
except httpx.HTTPStatusError as err:
|
||||
detail = _error_detail(err.response)
|
||||
|
||||
@@ -1,225 +0,0 @@
|
||||
"""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,13 +578,12 @@ 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="credentials"):
|
||||
with pytest.raises(click.ClickException, match="API key"):
|
||||
_create_host_backend_client(
|
||||
host_url="https://api.example.com",
|
||||
api_key=None,
|
||||
|
||||
@@ -1,191 +0,0 @@
|
||||
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={}
|
||||
)
|
||||
@@ -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:
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
@@ -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
|
||||
@@ -203,11 +204,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
|
||||
@@ -2885,12 +2881,14 @@ class Pregel(
|
||||
# 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)
|
||||
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_handler(h)
|
||||
if not _is_stream_mode_handler(h)
|
||||
]
|
||||
if "ls_integration" not in callback_manager.metadata:
|
||||
callback_manager.add_metadata({"ls_integration": "langgraph"})
|
||||
@@ -3300,12 +3298,14 @@ class Pregel(
|
||||
# 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)
|
||||
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_handler(h)
|
||||
if not _is_stream_mode_handler(h)
|
||||
]
|
||||
if "ls_integration" not in callback_manager.metadata:
|
||||
callback_manager.add_metadata({"ls_integration": "langgraph"})
|
||||
@@ -3320,20 +3320,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(
|
||||
@@ -4467,7 +4457,7 @@ def _coerce_checkpoint_values(payload: Any, mapper: Callable[[Any], Any]) -> Non
|
||||
payload["values"] = mapper(payload["values"])
|
||||
|
||||
|
||||
def _is_stream_handler(handler: object) -> bool:
|
||||
def _is_stream_mode_handler(handler: object) -> bool:
|
||||
return isinstance(handler, (StreamMessagesHandler, StreamToolCallHandler))
|
||||
|
||||
|
||||
|
||||
@@ -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"
|
||||
|
||||
@@ -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",
|
||||
|
||||
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