keystone-agent-sdk 0.3.0__tar.gz
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- keystone_agent_sdk-0.3.0/PKG-INFO +17 -0
- keystone_agent_sdk-0.3.0/__init__.py +47 -0
- keystone_agent_sdk-0.3.0/bundle.py +183 -0
- keystone_agent_sdk-0.3.0/clients/__init__.py +27 -0
- keystone_agent_sdk-0.3.0/clients/_base.py +56 -0
- keystone_agent_sdk-0.3.0/clients/_errors.py +92 -0
- keystone_agent_sdk-0.3.0/clients/gateway.py +56 -0
- keystone_agent_sdk-0.3.0/clients/rag.py +32 -0
- keystone_agent_sdk-0.3.0/durable.py +156 -0
- keystone_agent_sdk-0.3.0/hitl.py +43 -0
- keystone_agent_sdk-0.3.0/keystone_agent_sdk.egg-info/PKG-INFO +17 -0
- keystone_agent_sdk-0.3.0/keystone_agent_sdk.egg-info/SOURCES.txt +46 -0
- keystone_agent_sdk-0.3.0/keystone_agent_sdk.egg-info/dependency_links.txt +1 -0
- keystone_agent_sdk-0.3.0/keystone_agent_sdk.egg-info/requires.txt +13 -0
- keystone_agent_sdk-0.3.0/keystone_agent_sdk.egg-info/top_level.txt +1 -0
- keystone_agent_sdk-0.3.0/loader.py +62 -0
- keystone_agent_sdk-0.3.0/manifest.py +141 -0
- keystone_agent_sdk-0.3.0/provenance.py +109 -0
- keystone_agent_sdk-0.3.0/py.typed +0 -0
- keystone_agent_sdk-0.3.0/pyproject.toml +68 -0
- keystone_agent_sdk-0.3.0/setup.cfg +4 -0
- keystone_agent_sdk-0.3.0/tests/test_bundle.py +149 -0
- keystone_agent_sdk-0.3.0/tests/test_clients.py +96 -0
- keystone_agent_sdk-0.3.0/tests/test_durable.py +131 -0
- keystone_agent_sdk-0.3.0/tests/test_hitl.py +73 -0
- keystone_agent_sdk-0.3.0/tests/test_loader.py +31 -0
- keystone_agent_sdk-0.3.0/tests/test_manifest.py +131 -0
- keystone_agent_sdk-0.3.0/tests/test_multi_agent_semantics.py +276 -0
- keystone_agent_sdk-0.3.0/tests/test_provenance.py +105 -0
- keystone_agent_sdk-0.3.0/tests/test_tracing.py +167 -0
- keystone_agent_sdk-0.3.0/tests/test_validate.py +40 -0
- keystone_agent_sdk-0.3.0/tracing.py +100 -0
- keystone_agent_sdk-0.3.0/validate.py +69 -0
|
@@ -0,0 +1,17 @@
|
|
|
1
|
+
Metadata-Version: 2.4
|
|
2
|
+
Name: keystone-agent-sdk
|
|
3
|
+
Version: 0.3.0
|
|
4
|
+
Summary: Keystone pro-code agent SDK — LangGraph authoring primitives (StateGraph, @tool), the agent.yaml manifest model, and validate(). FR-04 / FDP-3072 (M1).
|
|
5
|
+
Requires-Python: >=3.12
|
|
6
|
+
Requires-Dist: langgraph<2,>=1
|
|
7
|
+
Requires-Dist: langchain-core<2,>=0.3
|
|
8
|
+
Requires-Dist: pydantic<3,>=2
|
|
9
|
+
Requires-Dist: pyyaml<7,>=6
|
|
10
|
+
Requires-Dist: httpx<1,>=0.27
|
|
11
|
+
Requires-Dist: keystone-request-context<0.4,>=0.3
|
|
12
|
+
Requires-Dist: opentelemetry-api<2,>=1.28
|
|
13
|
+
Provides-Extra: dev
|
|
14
|
+
Requires-Dist: pytest<9,>=8; extra == "dev"
|
|
15
|
+
Requires-Dist: pytest-cov<7,>=6; extra == "dev"
|
|
16
|
+
Requires-Dist: ruff==0.15.15; extra == "dev"
|
|
17
|
+
Requires-Dist: opentelemetry-sdk<2,>=1.28; extra == "dev"
|
|
@@ -0,0 +1,47 @@
|
|
|
1
|
+
"""Keystone AI — pro-code agent SDK (``keystone.agent_sdk``).
|
|
2
|
+
|
|
3
|
+
The developer's interface to the platform for **authoring** agents (FR-04 / FDP-3072).
|
|
4
|
+
M1 (walking skeleton) surface:
|
|
5
|
+
|
|
6
|
+
* **Authoring** — build the graph with **LangGraph directly** (``from langgraph.graph import
|
|
7
|
+
StateGraph``) + tools with ``langchain_core`` (``from langchain_core.tools import tool``). The
|
|
8
|
+
SDK does NOT wrap them, so an existing LangGraph agent runs here **unchanged**.
|
|
9
|
+
* **Manifest** — :class:`AgentManifest` (the ``agent.yaml`` schema, owned by the SDK
|
|
10
|
+
and shared by validate / run / deploy) + :func:`load_manifest`.
|
|
11
|
+
* **Validation** — :func:`validate` (manifest + entrypoint importability).
|
|
12
|
+
* **HITL** — :func:`request_approval` (stub; lands in M3 — FDP-3078).
|
|
13
|
+
|
|
14
|
+
Building-block clients live in the ``keystone.agent_sdk.clients`` subpackage
|
|
15
|
+
(``from keystone.agent_sdk.clients import GatewayClient, RAGClient`` — FDP-3118) — kept OFF the
|
|
16
|
+
top-level import so authoring-only agents don't pull the HTTP / request-context deps. The CLI
|
|
17
|
+
(``keystone agent run/validate``) is the separate ``keystone-cli`` package (FDP-3120).
|
|
18
|
+
|
|
19
|
+
Usage::
|
|
20
|
+
|
|
21
|
+
from langgraph.graph import StateGraph # graph = LangGraph, the SDK doesn't wrap it
|
|
22
|
+
from keystone.agent_sdk import AgentManifest, validate
|
|
23
|
+
from keystone.agent_sdk.clients import GatewayClient, RAGClient
|
|
24
|
+
"""
|
|
25
|
+
|
|
26
|
+
from .durable import durable_task, idempotency_key
|
|
27
|
+
from .hitl import request_approval
|
|
28
|
+
from .loader import GraphLoadError, load_graph
|
|
29
|
+
from .manifest import AgentManifest, Dependencies, HitlSpec, Identity, ResourceSpec, SecretRef, load_manifest
|
|
30
|
+
from .validate import ValidationResult, validate
|
|
31
|
+
|
|
32
|
+
__all__ = [
|
|
33
|
+
"AgentManifest",
|
|
34
|
+
"Dependencies",
|
|
35
|
+
"GraphLoadError",
|
|
36
|
+
"HitlSpec",
|
|
37
|
+
"Identity",
|
|
38
|
+
"ResourceSpec",
|
|
39
|
+
"SecretRef",
|
|
40
|
+
"ValidationResult",
|
|
41
|
+
"durable_task",
|
|
42
|
+
"idempotency_key",
|
|
43
|
+
"load_graph",
|
|
44
|
+
"load_manifest",
|
|
45
|
+
"request_approval",
|
|
46
|
+
"validate",
|
|
47
|
+
]
|
|
@@ -0,0 +1,183 @@
|
|
|
1
|
+
"""Package an agent project into a deterministic ``bundle.tar.gz`` + SHA-256 (deploy, FDP-3169).
|
|
2
|
+
|
|
3
|
+
Backs ``keystone agent deploy``: tar the project (source + agent.yaml + uv.lock) honouring ``.gitignore``,
|
|
4
|
+
skipping VCS/venv/cache junk. The SHA-256 (over the UNCOMPRESSED, normalized tar) is a stable idempotency
|
|
5
|
+
key — the same source produces the same hash, so a re-deploy dedupes to the existing version (cf. RAG file
|
|
6
|
+
upload). stdlib-only (no new dep); NOT re-exported at package top-level (deploy-time, not authoring).
|
|
7
|
+
"""
|
|
8
|
+
|
|
9
|
+
from __future__ import annotations
|
|
10
|
+
|
|
11
|
+
import fnmatch
|
|
12
|
+
import gzip
|
|
13
|
+
import hashlib
|
|
14
|
+
import io
|
|
15
|
+
import re
|
|
16
|
+
import tarfile
|
|
17
|
+
from dataclasses import dataclass
|
|
18
|
+
from dataclasses import field as dataclass_field
|
|
19
|
+
from pathlib import Path
|
|
20
|
+
|
|
21
|
+
# Always excluded regardless of .gitignore — VCS, venvs, caches, editor/OS cruft, build output.
|
|
22
|
+
_ALWAYS_IGNORE_DIRS = frozenset(
|
|
23
|
+
{
|
|
24
|
+
".git",
|
|
25
|
+
".venv",
|
|
26
|
+
"venv",
|
|
27
|
+
"env",
|
|
28
|
+
"__pycache__",
|
|
29
|
+
".mypy_cache",
|
|
30
|
+
".ruff_cache",
|
|
31
|
+
".pytest_cache",
|
|
32
|
+
"node_modules",
|
|
33
|
+
".idea",
|
|
34
|
+
".vscode",
|
|
35
|
+
"dist",
|
|
36
|
+
"build",
|
|
37
|
+
".keystone",
|
|
38
|
+
}
|
|
39
|
+
)
|
|
40
|
+
# `.env*` is excluded UNCONDITIONALLY, not left to .gitignore. It is the file the dev is told to put
|
|
41
|
+
# their platform API key in (`RT_AGENT_API_KEY`, access-policy.md §4.2), and a bundle is not a private
|
|
42
|
+
# thing: it lands in S3 and is baked into the agent's image in ECR — where retire/delete cleanup is
|
|
43
|
+
# still OFF (technical-debt.md §7), so a copy outlives a revoked key. Nothing needs it there either:
|
|
44
|
+
# a pod gets its env from the operator (ConfigMap/Secret), and `.env` only exists because
|
|
45
|
+
# `RuntimeSettings` reads `env_file=".env"` for LOCAL runs. Relying on the project's .gitignore would
|
|
46
|
+
# make this depend on a file the dev edits — and the scaffold's .gitignore didn't list `.env` at all
|
|
47
|
+
# until this change, so every project generated before it would have shipped the key.
|
|
48
|
+
_ALWAYS_IGNORE_GLOBS = ("*.pyc", "*.pyo", "*.egg-info", ".DS_Store", "*.log", ".env", ".env.*")
|
|
49
|
+
|
|
50
|
+
# The platform API key format (`ak_<env>_<short_id>.<secret>`, agent-identity reference / platform
|
|
51
|
+
# `api_keys`). Distinctive enough to match on directly, which is the point: this is a targeted check
|
|
52
|
+
# for OUR credential, not a general-purpose secret scanner (detect-secrets already runs pre-commit).
|
|
53
|
+
# Catches the key wherever it was pasted — `agent.yaml`'s `config:`, source, a README, or a `.env`
|
|
54
|
+
# variant under a name the glob above doesn't cover.
|
|
55
|
+
_API_KEY_RE = re.compile(r"\bak_[a-z0-9]{2,16}_[A-Za-z0-9]{4,}\.[A-Za-z0-9_-]{16,}")
|
|
56
|
+
|
|
57
|
+
# Cap per-file scanning: a bundle may legitimately carry a model or fixture, and reading a 500 MB
|
|
58
|
+
# file to look for a 60-char token is not worth it. Anything that fails to decode as UTF-8 is binary
|
|
59
|
+
# and skipped for the same reason.
|
|
60
|
+
_SCAN_MAX_BYTES = 1 << 20 # 1 MiB
|
|
61
|
+
|
|
62
|
+
|
|
63
|
+
class SecretInBundleError(Exception):
|
|
64
|
+
"""A bundled file contains a platform API key. Carries ``file``/``line`` — never the value.
|
|
65
|
+
|
|
66
|
+
The value is deliberately NOT in the message: this surfaces in terminal scrollback and CI job
|
|
67
|
+
logs, so echoing the credential there would leak it a second time in the act of reporting it.
|
|
68
|
+
"""
|
|
69
|
+
|
|
70
|
+
def __init__(self, file: str, line: int) -> None:
|
|
71
|
+
super().__init__(
|
|
72
|
+
f"{file}:{line} contains what looks like a platform API key (ak_<env>_<id>.<secret>). "
|
|
73
|
+
"Refusing to package it: a bundle is uploaded to S3 and baked into the agent image in "
|
|
74
|
+
"ECR, so the key would outlive being revoked. Put it in .env (never bundled) and let the "
|
|
75
|
+
"operator inject it, or reference it from the portal — do not commit it."
|
|
76
|
+
)
|
|
77
|
+
self.file = file
|
|
78
|
+
self.line = line
|
|
79
|
+
|
|
80
|
+
|
|
81
|
+
@dataclass
|
|
82
|
+
class BundleResult:
|
|
83
|
+
"""A packaged agent project. ``data`` = gzip bytes to upload; ``sha256`` = idempotency key.
|
|
84
|
+
|
|
85
|
+
``skipped_env`` names the env files left out (see ``_ALWAYS_IGNORE_GLOBS``). Reported so the
|
|
86
|
+
caller can SAY it: a dev who put config in ``.env`` and sees it missing at runtime would
|
|
87
|
+
otherwise have no way to know the packager dropped it on purpose.
|
|
88
|
+
"""
|
|
89
|
+
|
|
90
|
+
data: bytes
|
|
91
|
+
sha256: str
|
|
92
|
+
file_count: int
|
|
93
|
+
skipped_env: list[str] = dataclass_field(default_factory=list)
|
|
94
|
+
|
|
95
|
+
|
|
96
|
+
def _gitignore_patterns(root: Path) -> list[str]:
|
|
97
|
+
gi = root / ".gitignore"
|
|
98
|
+
if not gi.is_file():
|
|
99
|
+
return []
|
|
100
|
+
out: list[str] = []
|
|
101
|
+
for raw in gi.read_text(encoding="utf-8", errors="replace").splitlines():
|
|
102
|
+
line = raw.strip()
|
|
103
|
+
if not line or line.startswith("#") or line.startswith("!"): # negation unsupported (v1)
|
|
104
|
+
continue
|
|
105
|
+
out.append(line.rstrip("/"))
|
|
106
|
+
return out
|
|
107
|
+
|
|
108
|
+
|
|
109
|
+
def _is_ignored(rel: Path, patterns: list[str]) -> bool:
|
|
110
|
+
if set(rel.parts) & _ALWAYS_IGNORE_DIRS:
|
|
111
|
+
return True
|
|
112
|
+
name, rel_str = rel.name, str(rel)
|
|
113
|
+
if any(fnmatch.fnmatch(name, g) for g in _ALWAYS_IGNORE_GLOBS):
|
|
114
|
+
return True
|
|
115
|
+
return any(
|
|
116
|
+
fnmatch.fnmatch(name, p) or fnmatch.fnmatch(rel_str, p) or any(fnmatch.fnmatch(seg, p) for seg in rel.parts)
|
|
117
|
+
for p in patterns
|
|
118
|
+
)
|
|
119
|
+
|
|
120
|
+
|
|
121
|
+
def create_bundle(project_dir: str | Path) -> BundleResult:
|
|
122
|
+
"""Tar+gzip ``project_dir`` (deterministic: sorted, mtime/uid/gid/mode normalized) → BundleResult.
|
|
123
|
+
|
|
124
|
+
Raises ``NotADirectoryError`` if the path isn't a directory, ``FileNotFoundError`` if there's no
|
|
125
|
+
``agent.yaml`` (a deploy without a manifest is a mistake), and :class:`SecretInBundleError` if a
|
|
126
|
+
file that WOULD be packaged carries a platform API key."""
|
|
127
|
+
root = Path(project_dir).resolve()
|
|
128
|
+
if not root.is_dir():
|
|
129
|
+
raise NotADirectoryError(f"not a directory: {root}")
|
|
130
|
+
if not (root / "agent.yaml").is_file():
|
|
131
|
+
raise FileNotFoundError(f"no agent.yaml in {root} — not an agent project")
|
|
132
|
+
|
|
133
|
+
patterns = _gitignore_patterns(root)
|
|
134
|
+
tar_buf = io.BytesIO()
|
|
135
|
+
count = 0
|
|
136
|
+
skipped_env: list[str] = []
|
|
137
|
+
# Uncompressed tar, sorted + normalized → the hash is stable across machines/times.
|
|
138
|
+
with tarfile.open(fileobj=tar_buf, mode="w") as tar:
|
|
139
|
+
for path in sorted(p for p in root.rglob("*") if p.is_file()):
|
|
140
|
+
rel = path.relative_to(root)
|
|
141
|
+
if _is_ignored(rel, patterns):
|
|
142
|
+
if _is_env_file(rel):
|
|
143
|
+
skipped_env.append(str(rel))
|
|
144
|
+
continue
|
|
145
|
+
payload = path.read_bytes()
|
|
146
|
+
# Check BEFORE adding: the exception must leave no half-built tar behind, and a bundle
|
|
147
|
+
# that was already assembled is a bundle someone can accidentally upload.
|
|
148
|
+
_reject_api_key(str(rel), payload)
|
|
149
|
+
info = tarfile.TarInfo(name=str(rel))
|
|
150
|
+
info.size = len(payload)
|
|
151
|
+
info.mtime = 0
|
|
152
|
+
info.mode = 0o644
|
|
153
|
+
info.uid = info.gid = 0
|
|
154
|
+
info.uname = info.gname = ""
|
|
155
|
+
tar.addfile(info, io.BytesIO(payload))
|
|
156
|
+
count += 1
|
|
157
|
+
|
|
158
|
+
tar_bytes = tar_buf.getvalue()
|
|
159
|
+
sha256 = hashlib.sha256(tar_bytes).hexdigest()
|
|
160
|
+
data = gzip.compress(tar_bytes, mtime=0) # mtime=0 → deterministic gzip too
|
|
161
|
+
return BundleResult(data=data, sha256=sha256, file_count=count, skipped_env=skipped_env)
|
|
162
|
+
|
|
163
|
+
|
|
164
|
+
def _is_env_file(rel: Path) -> bool:
|
|
165
|
+
"""True for the env-file shapes ``_ALWAYS_IGNORE_GLOBS`` drops (``.env``, ``.env.local``, …)."""
|
|
166
|
+
return rel.name == ".env" or rel.name.startswith(".env.")
|
|
167
|
+
|
|
168
|
+
|
|
169
|
+
def _reject_api_key(rel: str, payload: bytes) -> None:
|
|
170
|
+
"""Raise :class:`SecretInBundleError` if ``payload`` carries a platform API key.
|
|
171
|
+
|
|
172
|
+
Binary and oversized files are skipped rather than scanned — see ``_SCAN_MAX_BYTES``. A key can
|
|
173
|
+
only be *pasted* as text, so decoding failure is a reliable "not it" rather than a blind spot.
|
|
174
|
+
"""
|
|
175
|
+
if len(payload) > _SCAN_MAX_BYTES:
|
|
176
|
+
return
|
|
177
|
+
try:
|
|
178
|
+
text = payload.decode("utf-8")
|
|
179
|
+
except UnicodeDecodeError:
|
|
180
|
+
return
|
|
181
|
+
for lineno, line in enumerate(text.splitlines(), start=1):
|
|
182
|
+
if _API_KEY_RE.search(line):
|
|
183
|
+
raise SecretInBundleError(rel, lineno)
|
|
@@ -0,0 +1,27 @@
|
|
|
1
|
+
"""Building-block clients for pro-code agents (FDP-3118).
|
|
2
|
+
|
|
3
|
+
M1 ships ``GatewayClient`` (chat/embeddings) + ``RAGClient`` (search) — the ``rag-qa`` agent's
|
|
4
|
+
deps. guardrails / memory / skill clients land as their stories ship. All clients auto-attach
|
|
5
|
+
the caller identity (from the request-context ContextVar the runtime sets per run) and decode
|
|
6
|
+
errors to :class:`GatewayError`.
|
|
7
|
+
"""
|
|
8
|
+
|
|
9
|
+
from keystone.agent_sdk.clients._base import BaseClient
|
|
10
|
+
from keystone.agent_sdk.clients._errors import (
|
|
11
|
+
GatewayError,
|
|
12
|
+
GatewayForbiddenError,
|
|
13
|
+
GatewayRateLimitedError,
|
|
14
|
+
GatewayUpstreamError,
|
|
15
|
+
)
|
|
16
|
+
from keystone.agent_sdk.clients.gateway import GatewayClient
|
|
17
|
+
from keystone.agent_sdk.clients.rag import RAGClient
|
|
18
|
+
|
|
19
|
+
__all__ = [
|
|
20
|
+
"BaseClient",
|
|
21
|
+
"GatewayClient",
|
|
22
|
+
"GatewayError",
|
|
23
|
+
"GatewayForbiddenError",
|
|
24
|
+
"GatewayRateLimitedError",
|
|
25
|
+
"GatewayUpstreamError",
|
|
26
|
+
"RAGClient",
|
|
27
|
+
]
|
|
@@ -0,0 +1,56 @@
|
|
|
1
|
+
"""Shared base for the SDK building-block clients (proposal §3.5a: ContextVar identity +
|
|
2
|
+
typed error decode + singleton httpx).
|
|
3
|
+
|
|
4
|
+
The client attaches caller-identity headers from the request-context ContextVar that the agent
|
|
5
|
+
runtime sets per run (``keystone.request_context.gateway_identity_headers``) — the agent author
|
|
6
|
+
never sets headers. One ``httpx.AsyncClient`` per client instance (connection pooling).
|
|
7
|
+
"""
|
|
8
|
+
|
|
9
|
+
from __future__ import annotations
|
|
10
|
+
|
|
11
|
+
from typing import Any
|
|
12
|
+
|
|
13
|
+
import httpx
|
|
14
|
+
from keystone.agent_sdk.clients._errors import decode_gateway_error
|
|
15
|
+
from keystone.agent_sdk.tracing import inject_traceparent
|
|
16
|
+
from keystone.request_context import gateway_identity_headers
|
|
17
|
+
|
|
18
|
+
|
|
19
|
+
class BaseClient:
|
|
20
|
+
"""Base HTTP client — identity injection + error decode. Subclasses add typed methods."""
|
|
21
|
+
|
|
22
|
+
def __init__(
|
|
23
|
+
self,
|
|
24
|
+
*,
|
|
25
|
+
base_url: str,
|
|
26
|
+
auth_token: str = "",
|
|
27
|
+
caller_service: str = "agent", # gateway X-Caller-Service vocab for an agent-runtime caller (Arch §10)
|
|
28
|
+
timeout_s: float = 60.0,
|
|
29
|
+
) -> None:
|
|
30
|
+
self._base_url = base_url.rstrip("/")
|
|
31
|
+
self._auth_token = auth_token
|
|
32
|
+
self._caller_service = caller_service
|
|
33
|
+
self._http = httpx.AsyncClient(
|
|
34
|
+
limits=httpx.Limits(max_keepalive_connections=20, max_connections=40),
|
|
35
|
+
timeout=httpx.Timeout(timeout_s),
|
|
36
|
+
)
|
|
37
|
+
|
|
38
|
+
def _headers(self, operation: str) -> dict[str, str]:
|
|
39
|
+
headers = gateway_identity_headers(
|
|
40
|
+
auth_token=self._auth_token,
|
|
41
|
+
operation=operation,
|
|
42
|
+
caller_service=self._caller_service,
|
|
43
|
+
)
|
|
44
|
+
# FDP-3436: propagate the run's trace context so downstream hops (rag,
|
|
45
|
+
# llm-gateway) can join the trace. No active span (offline run) → no-op.
|
|
46
|
+
return inject_traceparent(headers)
|
|
47
|
+
|
|
48
|
+
async def _post(self, path: str, body: dict[str, Any], *, operation: str) -> dict[str, Any]:
|
|
49
|
+
response = await self._http.post(f"{self._base_url}{path}", json=body, headers=self._headers(operation))
|
|
50
|
+
if response.status_code >= 400:
|
|
51
|
+
raise decode_gateway_error(response)
|
|
52
|
+
return response.json()
|
|
53
|
+
|
|
54
|
+
async def aclose(self) -> None:
|
|
55
|
+
"""Close the connection pool. Idempotent."""
|
|
56
|
+
await self._http.aclose()
|
|
@@ -0,0 +1,92 @@
|
|
|
1
|
+
"""Typed errors for the SDK building-block clients.
|
|
2
|
+
|
|
3
|
+
SDK-LOCAL copy of the gateway error-decode shape (FDP-3118 scope call — `decode_gateway_error`
|
|
4
|
+
is NOT lifted to a shared lib this milestone; promote to ``keystone.errors`` when the
|
|
5
|
+
guardrails/memory/skill clients land). The envelope is the keystone-standard
|
|
6
|
+
``{"error": {"code", "message", "request_id"}}`` — so the same decoder serves rag responses too
|
|
7
|
+
(a rag error decodes to :class:`GatewayError` carrying its ``RAG_*`` code).
|
|
8
|
+
"""
|
|
9
|
+
|
|
10
|
+
from __future__ import annotations
|
|
11
|
+
|
|
12
|
+
from typing import TYPE_CHECKING
|
|
13
|
+
|
|
14
|
+
if TYPE_CHECKING:
|
|
15
|
+
import httpx
|
|
16
|
+
|
|
17
|
+
|
|
18
|
+
class GatewayError(Exception):
|
|
19
|
+
"""A building-block call returned a non-2xx envelope."""
|
|
20
|
+
|
|
21
|
+
def __init__(
|
|
22
|
+
self,
|
|
23
|
+
*,
|
|
24
|
+
code: str,
|
|
25
|
+
message: str,
|
|
26
|
+
request_id: str | None = None,
|
|
27
|
+
retry_after: int | None = None,
|
|
28
|
+
) -> None:
|
|
29
|
+
super().__init__(f"[{code}] {message}")
|
|
30
|
+
self.code = code
|
|
31
|
+
self.message = message
|
|
32
|
+
self.request_id = request_id
|
|
33
|
+
self.retry_after = retry_after
|
|
34
|
+
|
|
35
|
+
|
|
36
|
+
class GatewayForbiddenError(GatewayError):
|
|
37
|
+
"""Caller not permitted (model not allowed / unauthorized / invalid key)."""
|
|
38
|
+
|
|
39
|
+
|
|
40
|
+
class GatewayRateLimitedError(GatewayError):
|
|
41
|
+
"""Provider throttle / budget — carries ``retry_after`` when the server sent one."""
|
|
42
|
+
|
|
43
|
+
|
|
44
|
+
class GatewayUpstreamError(GatewayError):
|
|
45
|
+
"""Upstream/provider 5xx or gateway self-error — retryable after backoff."""
|
|
46
|
+
|
|
47
|
+
|
|
48
|
+
_ERROR_CODE_MAP: dict[str, type[GatewayError]] = {
|
|
49
|
+
"LLM_MODEL_NOT_ALLOWED": GatewayForbiddenError,
|
|
50
|
+
"LLM_MODEL_UNAVAILABLE": GatewayForbiddenError,
|
|
51
|
+
"LLM_UNAUTHORIZED": GatewayForbiddenError,
|
|
52
|
+
"LLM_INVALID_KEY": GatewayForbiddenError,
|
|
53
|
+
"LLM_INVALID_CALLER_IDENTITY": GatewayForbiddenError,
|
|
54
|
+
"LLM_PROVIDER_RATE_LIMITED": GatewayRateLimitedError,
|
|
55
|
+
"LLM_BUDGET_EXCEEDED": GatewayRateLimitedError,
|
|
56
|
+
"LLM_COST_CEILING_EXCEEDED": GatewayRateLimitedError,
|
|
57
|
+
"LLM_PROVIDER_ERROR": GatewayUpstreamError,
|
|
58
|
+
"LLM_UPSTREAM_ERROR": GatewayUpstreamError,
|
|
59
|
+
"LLM_UPSTREAM_TIMEOUT": GatewayUpstreamError,
|
|
60
|
+
"LLM_PROVIDER_UNAVAILABLE": GatewayUpstreamError,
|
|
61
|
+
"LLM_NO_PROVIDERS_AVAILABLE": GatewayUpstreamError,
|
|
62
|
+
"LLM_INTERNAL_ERROR": GatewayUpstreamError,
|
|
63
|
+
}
|
|
64
|
+
|
|
65
|
+
|
|
66
|
+
def _parse_retry_after(response: httpx.Response) -> int | None:
|
|
67
|
+
raw = response.headers.get("Retry-After")
|
|
68
|
+
if not isinstance(raw, str):
|
|
69
|
+
return None
|
|
70
|
+
try:
|
|
71
|
+
secs = int(raw.strip())
|
|
72
|
+
except ValueError:
|
|
73
|
+
return None
|
|
74
|
+
return secs if secs > 0 else None
|
|
75
|
+
|
|
76
|
+
|
|
77
|
+
def decode_gateway_error(response: httpx.Response) -> GatewayError:
|
|
78
|
+
"""Decode a non-2xx response into a typed :class:`GatewayError`.
|
|
79
|
+
|
|
80
|
+
Defensive: a non-JSON / non-envelope body (infra error page) falls back to ``HTTP_<status>``.
|
|
81
|
+
"""
|
|
82
|
+
try:
|
|
83
|
+
envelope = (response.json().get("error")) or {}
|
|
84
|
+
code = envelope.get("code") or f"HTTP_{response.status_code}"
|
|
85
|
+
message = envelope.get("message") or response.text[:500]
|
|
86
|
+
request_id = envelope.get("request_id")
|
|
87
|
+
except (ValueError, AttributeError):
|
|
88
|
+
code = f"HTTP_{response.status_code}"
|
|
89
|
+
message = response.text[:500] or response.reason_phrase
|
|
90
|
+
request_id = None
|
|
91
|
+
exc_cls = _ERROR_CODE_MAP.get(code, GatewayError)
|
|
92
|
+
return exc_cls(code=code, message=message, request_id=request_id, retry_after=_parse_retry_after(response))
|
|
@@ -0,0 +1,56 @@
|
|
|
1
|
+
"""``GatewayClient`` — chat completions + embeddings via llm-gateway (OpenAI-compatible).
|
|
2
|
+
|
|
3
|
+
Endpoints from ``specs/api/llm-gateway.openapi.yaml``. Identity headers + error decode are
|
|
4
|
+
handled by :class:`~keystone.agent_sdk.clients._base.BaseClient`.
|
|
5
|
+
"""
|
|
6
|
+
|
|
7
|
+
from __future__ import annotations
|
|
8
|
+
|
|
9
|
+
from typing import TYPE_CHECKING, Any
|
|
10
|
+
|
|
11
|
+
from keystone.agent_sdk.clients._base import BaseClient
|
|
12
|
+
from keystone.agent_sdk.tracing import genai_span, link_generation, record_usage
|
|
13
|
+
|
|
14
|
+
if TYPE_CHECKING:
|
|
15
|
+
from collections.abc import Sequence
|
|
16
|
+
|
|
17
|
+
_CHAT_PATH = "/api/llm-gateway/v1/chat/completions"
|
|
18
|
+
_EMBEDDINGS_PATH = "/api/llm-gateway/v1/embeddings"
|
|
19
|
+
|
|
20
|
+
|
|
21
|
+
class GatewayClient(BaseClient):
|
|
22
|
+
async def chat(
|
|
23
|
+
self,
|
|
24
|
+
*,
|
|
25
|
+
messages: list[dict[str, str]],
|
|
26
|
+
model: str,
|
|
27
|
+
max_tokens: int = 1024,
|
|
28
|
+
temperature: float = 0.0,
|
|
29
|
+
**extra: Any,
|
|
30
|
+
) -> dict[str, Any]:
|
|
31
|
+
"""OpenAI-compatible chat completion. Returns the raw gateway JSON."""
|
|
32
|
+
body: dict[str, Any] = {
|
|
33
|
+
"model": model,
|
|
34
|
+
"messages": messages,
|
|
35
|
+
"max_tokens": max_tokens,
|
|
36
|
+
"temperature": temperature,
|
|
37
|
+
**extra,
|
|
38
|
+
}
|
|
39
|
+
# FDP-3436: SDK-owned GenAI span (spike Q3) — token usage from the response;
|
|
40
|
+
# no-op when the host process has no tracer (local `agent run`).
|
|
41
|
+
with genai_span("chat", model) as span:
|
|
42
|
+
# FDP-3663: must be INSIDE the span — it stamps that span's id into the body so
|
|
43
|
+
# LiteLLM's generation nests under it instead of dangling at the trace root.
|
|
44
|
+
link_generation(span, body)
|
|
45
|
+
result = await self._post(_CHAT_PATH, body, operation="chat_completion")
|
|
46
|
+
record_usage(span, result)
|
|
47
|
+
return result
|
|
48
|
+
|
|
49
|
+
async def embeddings(self, *, texts: Sequence[str], model: str, **extra: Any) -> dict[str, Any]:
|
|
50
|
+
"""Vector embeddings for ``texts``. Returns the raw gateway JSON."""
|
|
51
|
+
body: dict[str, Any] = {"model": model, "input": list(texts), **extra}
|
|
52
|
+
with genai_span("embeddings", model) as span:
|
|
53
|
+
link_generation(span, body)
|
|
54
|
+
result = await self._post(_EMBEDDINGS_PATH, body, operation="embedding")
|
|
55
|
+
record_usage(span, result)
|
|
56
|
+
return result
|
|
@@ -0,0 +1,32 @@
|
|
|
1
|
+
"""``RAGClient`` — hybrid search via the rag service.
|
|
2
|
+
|
|
3
|
+
Path verified against the ``services/rag`` search router (prefix ``/knowledge-bases``). ``search``
|
|
4
|
+
targets the multi-KB endpoint (``/api/rag/v1/knowledge-bases/search`` with ``kb_ids`` in the body).
|
|
5
|
+
Response = ``{"results": [{"content", "chunk_id", "score", ...}], ...}`` (no ``data`` envelope).
|
|
6
|
+
Identity headers + error decode via :class:`~keystone.agent_sdk.clients._base.BaseClient`.
|
|
7
|
+
"""
|
|
8
|
+
|
|
9
|
+
from __future__ import annotations
|
|
10
|
+
|
|
11
|
+
from typing import TYPE_CHECKING, Any
|
|
12
|
+
|
|
13
|
+
from keystone.agent_sdk.clients._base import BaseClient
|
|
14
|
+
|
|
15
|
+
if TYPE_CHECKING:
|
|
16
|
+
from collections.abc import Sequence
|
|
17
|
+
|
|
18
|
+
_SEARCH_PATH = "/api/rag/v1/knowledge-bases/search"
|
|
19
|
+
|
|
20
|
+
|
|
21
|
+
class RAGClient(BaseClient):
|
|
22
|
+
async def search(
|
|
23
|
+
self,
|
|
24
|
+
*,
|
|
25
|
+
query: str,
|
|
26
|
+
kb_ids: Sequence[str],
|
|
27
|
+
top_k: int = 5,
|
|
28
|
+
**extra: Any,
|
|
29
|
+
) -> dict[str, Any]:
|
|
30
|
+
"""Hybrid retrieval across ``kb_ids``. Returns the raw rag JSON (chunks + citations)."""
|
|
31
|
+
body: dict[str, Any] = {"query": query, "kb_ids": list(kb_ids), "top_k": top_k, **extra}
|
|
32
|
+
return await self._post(_SEARCH_PATH, body, operation="search")
|
|
@@ -0,0 +1,156 @@
|
|
|
1
|
+
"""Durable-task discipline for paid / side-effecting steps (M3.3 — FDP-3334).
|
|
2
|
+
|
|
3
|
+
Spike-verified premise (langgraph 1.1.10 + AsyncPostgresSaver, see
|
|
4
|
+
``services/agent-runtime/tests/integration/test_task_memoization.py``): when a node
|
|
5
|
+
crashes AFTER a ``langgraph.func.task`` call completed, the task's result is already
|
|
6
|
+
persisted with the checkpoint; on resume (``ainvoke(None)`` — the platform worker's
|
|
7
|
+
convention) the node body re-runs but completed tasks **replay from the checkpoint
|
|
8
|
+
instead of re-executing**. Wrap every expensive step (LLM call, paid API, mutation)
|
|
9
|
+
in :func:`durable_task` and a crash-resumed run re-bills only the in-flight step.
|
|
10
|
+
|
|
11
|
+
Two rules make the replay trustworthy, and this module owns both:
|
|
12
|
+
|
|
13
|
+
* **Memo-safe arguments** — replay matches tasks by their deterministic position in
|
|
14
|
+
the node's call sequence, so arguments must be *values*, not live handles or
|
|
15
|
+
freshly-minted timestamps. :func:`durable_task` validates every call's arguments
|
|
16
|
+
and rejects clients / connections / callables (construct them INSIDE the task
|
|
17
|
+
body) and ``datetime`` objects (pass pre-computed ISO strings, or move ``now()``
|
|
18
|
+
inside the task body so it is captured with the memoized result).
|
|
19
|
+
* **Idempotency keys for mutations** — memoization stops the *node-retry* double
|
|
20
|
+
fire, but the crash can still land INSIDE a mutating call (money moved, ticket
|
|
21
|
+
created — then SIGKILL before the result persisted). Downstream dedup is the only
|
|
22
|
+
fix: send :func:`idempotency_key` with every mutating call so the retried call is
|
|
23
|
+
recognised. This gate is review-enforced (``.claude/references/review-checklist.md``),
|
|
24
|
+
not runtime-enforced.
|
|
25
|
+
|
|
26
|
+
Usage::
|
|
27
|
+
|
|
28
|
+
from keystone.agent_sdk import durable_task, idempotency_key
|
|
29
|
+
|
|
30
|
+
@durable_task
|
|
31
|
+
async def draft_reply(ticket: str) -> str: # paid LLM call
|
|
32
|
+
return await llm.ainvoke(ticket) # client built/closed inside
|
|
33
|
+
|
|
34
|
+
async def triage(state: State, config: RunnableConfig) -> dict:
|
|
35
|
+
draft = await draft_reply(state["ticket"]) # replayed, not re-billed, on resume
|
|
36
|
+
key = idempotency_key(
|
|
37
|
+
config["configurable"]["thread_id"], "triage", f"create-ticket:{state['id']}"
|
|
38
|
+
)
|
|
39
|
+
await helpdesk.create(draft, idempotency_key=key)
|
|
40
|
+
return {"draft": draft}
|
|
41
|
+
"""
|
|
42
|
+
|
|
43
|
+
from __future__ import annotations
|
|
44
|
+
|
|
45
|
+
import hashlib
|
|
46
|
+
from collections.abc import Callable
|
|
47
|
+
from typing import TYPE_CHECKING, Any, overload
|
|
48
|
+
|
|
49
|
+
from langgraph.func import task
|
|
50
|
+
from pydantic import BaseModel
|
|
51
|
+
|
|
52
|
+
if TYPE_CHECKING:
|
|
53
|
+
from langgraph.types import RetryPolicy
|
|
54
|
+
|
|
55
|
+
_SCALARS = (type(None), bool, int, float, str, bytes)
|
|
56
|
+
|
|
57
|
+
|
|
58
|
+
def _assert_memo_safe(value: Any, path: str) -> None:
|
|
59
|
+
"""Reject arguments that would make a durable task's replay non-deterministic."""
|
|
60
|
+
if isinstance(value, _SCALARS):
|
|
61
|
+
return
|
|
62
|
+
if isinstance(value, (list, tuple)):
|
|
63
|
+
for i, item in enumerate(value):
|
|
64
|
+
_assert_memo_safe(item, f"{path}[{i}]")
|
|
65
|
+
return
|
|
66
|
+
if isinstance(value, dict):
|
|
67
|
+
for k, item in value.items():
|
|
68
|
+
if not isinstance(k, str):
|
|
69
|
+
raise TypeError(
|
|
70
|
+
f"durable_task argument {path} has a non-string dict key {k!r}; "
|
|
71
|
+
"memo-safe arguments use string keys only"
|
|
72
|
+
)
|
|
73
|
+
_assert_memo_safe(item, f"{path}[{k!r}]")
|
|
74
|
+
return
|
|
75
|
+
if isinstance(value, BaseModel):
|
|
76
|
+
_assert_memo_safe(value.model_dump(), path)
|
|
77
|
+
return
|
|
78
|
+
# datetime/date/time before the generic reject, for a targeted message.
|
|
79
|
+
import datetime as _dt
|
|
80
|
+
|
|
81
|
+
if isinstance(value, (_dt.datetime, _dt.date, _dt.time)):
|
|
82
|
+
raise TypeError(
|
|
83
|
+
f"durable_task argument {path} is a {type(value).__name__} — timestamps make "
|
|
84
|
+
"the replayed call non-deterministic. Pass a pre-computed ISO string, or move "
|
|
85
|
+
"now() inside the task body so it is captured with the memoized result."
|
|
86
|
+
)
|
|
87
|
+
if isinstance(value, (set, frozenset)):
|
|
88
|
+
raise TypeError(
|
|
89
|
+
f"durable_task argument {path} is a {type(value).__name__} — iteration order "
|
|
90
|
+
"is non-deterministic; pass a sorted list instead"
|
|
91
|
+
)
|
|
92
|
+
raise TypeError(
|
|
93
|
+
f"durable_task argument {path} is a live object ({type(value).__module__}."
|
|
94
|
+
f"{type(value).__qualname__}) — clients, connections and other handles must not "
|
|
95
|
+
"cross the durable-task boundary. Construct them INSIDE the task body and pass "
|
|
96
|
+
"plain data (str/int/float/bool/bytes/list/dict/pydantic model) in."
|
|
97
|
+
)
|
|
98
|
+
|
|
99
|
+
|
|
100
|
+
@overload
|
|
101
|
+
def durable_task[F: Callable[..., Any]](fn: F) -> F: ...
|
|
102
|
+
@overload
|
|
103
|
+
def durable_task(
|
|
104
|
+
*, name: str | None = None, retry_policy: RetryPolicy | None = None
|
|
105
|
+
) -> Callable[[Callable[..., Any]], Callable[..., Any]]: ...
|
|
106
|
+
|
|
107
|
+
|
|
108
|
+
def durable_task(
|
|
109
|
+
fn: Callable[..., Any] | None = None,
|
|
110
|
+
*,
|
|
111
|
+
name: str | None = None,
|
|
112
|
+
retry_policy: RetryPolicy | None = None,
|
|
113
|
+
) -> Any:
|
|
114
|
+
"""Mark a paid / side-effecting step durable: memoized across crash-resume.
|
|
115
|
+
|
|
116
|
+
Thin, validated wrapper over ``langgraph.func.task`` — a resumed run replays the
|
|
117
|
+
completed call's persisted result instead of re-executing it. Call it from inside
|
|
118
|
+
a graph node (sync or async; ``await fn(...)`` / ``fn(...).result()``). Arguments
|
|
119
|
+
are validated per call (see :func:`_assert_memo_safe`); the result must be
|
|
120
|
+
checkpoint-serialisable, exactly like node state.
|
|
121
|
+
|
|
122
|
+
:param name: memo identity (defaults to the function name). Renaming a durable
|
|
123
|
+
task orphans in-flight runs' memoized results — treat the name as a contract.
|
|
124
|
+
:param retry_policy: forwarded to ``langgraph.func.task`` (transient in-process
|
|
125
|
+
retries; distinct from the platform's crash-resume).
|
|
126
|
+
"""
|
|
127
|
+
|
|
128
|
+
def decorate(f: Callable[..., Any]) -> Callable[..., Any]:
|
|
129
|
+
inner = task(name=name or f.__name__, retry_policy=retry_policy)(f)
|
|
130
|
+
|
|
131
|
+
def call(*args: Any, **kwargs: Any) -> Any:
|
|
132
|
+
for i, a in enumerate(args):
|
|
133
|
+
_assert_memo_safe(a, f"#{i}")
|
|
134
|
+
for k, a in kwargs.items():
|
|
135
|
+
_assert_memo_safe(a, k)
|
|
136
|
+
return inner(*args, **kwargs)
|
|
137
|
+
|
|
138
|
+
call.__name__ = name or f.__name__
|
|
139
|
+
call.__doc__ = f.__doc__
|
|
140
|
+
call.__wrapped__ = f # type: ignore[attr-defined]
|
|
141
|
+
return call # type: ignore[return-value]
|
|
142
|
+
|
|
143
|
+
return decorate(fn) if fn is not None else decorate
|
|
144
|
+
|
|
145
|
+
|
|
146
|
+
def idempotency_key(thread_id: str, node: str, call: str) -> str:
|
|
147
|
+
"""Deterministic dedup key for a mutating call: stable across crash-resume retries.
|
|
148
|
+
|
|
149
|
+
Same logical call → same key on every replay, so the downstream system can drop
|
|
150
|
+
the duplicate (memoization alone cannot help when the crash lands INSIDE the
|
|
151
|
+
mutation). ``thread_id`` comes from the node's ``config["configurable"]["thread_id"]``;
|
|
152
|
+
``node`` is the node name; ``call`` identifies the call site + its business identity
|
|
153
|
+
(e.g. ``f"refund:{order_id}"``) — include a loop index if the node issues several.
|
|
154
|
+
"""
|
|
155
|
+
joined = "\x1f".join((thread_id, node, call))
|
|
156
|
+
return hashlib.sha256(joined.encode("utf-8")).hexdigest()
|