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.
Files changed (33) hide show
  1. keystone_agent_sdk-0.3.0/PKG-INFO +17 -0
  2. keystone_agent_sdk-0.3.0/__init__.py +47 -0
  3. keystone_agent_sdk-0.3.0/bundle.py +183 -0
  4. keystone_agent_sdk-0.3.0/clients/__init__.py +27 -0
  5. keystone_agent_sdk-0.3.0/clients/_base.py +56 -0
  6. keystone_agent_sdk-0.3.0/clients/_errors.py +92 -0
  7. keystone_agent_sdk-0.3.0/clients/gateway.py +56 -0
  8. keystone_agent_sdk-0.3.0/clients/rag.py +32 -0
  9. keystone_agent_sdk-0.3.0/durable.py +156 -0
  10. keystone_agent_sdk-0.3.0/hitl.py +43 -0
  11. keystone_agent_sdk-0.3.0/keystone_agent_sdk.egg-info/PKG-INFO +17 -0
  12. keystone_agent_sdk-0.3.0/keystone_agent_sdk.egg-info/SOURCES.txt +46 -0
  13. keystone_agent_sdk-0.3.0/keystone_agent_sdk.egg-info/dependency_links.txt +1 -0
  14. keystone_agent_sdk-0.3.0/keystone_agent_sdk.egg-info/requires.txt +13 -0
  15. keystone_agent_sdk-0.3.0/keystone_agent_sdk.egg-info/top_level.txt +1 -0
  16. keystone_agent_sdk-0.3.0/loader.py +62 -0
  17. keystone_agent_sdk-0.3.0/manifest.py +141 -0
  18. keystone_agent_sdk-0.3.0/provenance.py +109 -0
  19. keystone_agent_sdk-0.3.0/py.typed +0 -0
  20. keystone_agent_sdk-0.3.0/pyproject.toml +68 -0
  21. keystone_agent_sdk-0.3.0/setup.cfg +4 -0
  22. keystone_agent_sdk-0.3.0/tests/test_bundle.py +149 -0
  23. keystone_agent_sdk-0.3.0/tests/test_clients.py +96 -0
  24. keystone_agent_sdk-0.3.0/tests/test_durable.py +131 -0
  25. keystone_agent_sdk-0.3.0/tests/test_hitl.py +73 -0
  26. keystone_agent_sdk-0.3.0/tests/test_loader.py +31 -0
  27. keystone_agent_sdk-0.3.0/tests/test_manifest.py +131 -0
  28. keystone_agent_sdk-0.3.0/tests/test_multi_agent_semantics.py +276 -0
  29. keystone_agent_sdk-0.3.0/tests/test_provenance.py +105 -0
  30. keystone_agent_sdk-0.3.0/tests/test_tracing.py +167 -0
  31. keystone_agent_sdk-0.3.0/tests/test_validate.py +40 -0
  32. keystone_agent_sdk-0.3.0/tracing.py +100 -0
  33. 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()