terrarium-python 0.1.1__py3-none-any.whl
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.
- terrarium/__init__.py +86 -0
- terrarium/cli.py +178 -0
- terrarium/client.py +893 -0
- terrarium/errors.py +66 -0
- terrarium/memory.py +171 -0
- terrarium/messages.py +137 -0
- terrarium/options.py +212 -0
- terrarium/py.typed +0 -0
- terrarium_python-0.1.1.dist-info/METADATA +397 -0
- terrarium_python-0.1.1.dist-info/RECORD +13 -0
- terrarium_python-0.1.1.dist-info/WHEEL +4 -0
- terrarium_python-0.1.1.dist-info/entry_points.txt +2 -0
- terrarium_python-0.1.1.dist-info/licenses/LICENSE +202 -0
terrarium/errors.py
ADDED
|
@@ -0,0 +1,66 @@
|
|
|
1
|
+
"""Typed exceptions so callers can branch on failure class instead of
|
|
2
|
+
string-matching HTTP status codes."""
|
|
3
|
+
|
|
4
|
+
from __future__ import annotations
|
|
5
|
+
|
|
6
|
+
from typing import Any
|
|
7
|
+
|
|
8
|
+
import httpx
|
|
9
|
+
|
|
10
|
+
|
|
11
|
+
class TerrariumError(Exception):
|
|
12
|
+
"""Base for all SDK errors. Carries the HTTP status + response body when known."""
|
|
13
|
+
|
|
14
|
+
def __init__(self, message: str, *, status: int | None = None, body: Any = None) -> None:
|
|
15
|
+
super().__init__(message)
|
|
16
|
+
self.status = status
|
|
17
|
+
self.body = body
|
|
18
|
+
|
|
19
|
+
|
|
20
|
+
class AuthError(TerrariumError):
|
|
21
|
+
"""401/403 — missing, invalid, or insufficiently-scoped token."""
|
|
22
|
+
|
|
23
|
+
|
|
24
|
+
class NotFoundError(TerrariumError):
|
|
25
|
+
"""404 — session/agent/etc. does not exist."""
|
|
26
|
+
|
|
27
|
+
|
|
28
|
+
class ConflictError(TerrariumError):
|
|
29
|
+
"""409 — conflicting state."""
|
|
30
|
+
|
|
31
|
+
|
|
32
|
+
class RateLimitError(TerrariumError):
|
|
33
|
+
"""429 — slow down."""
|
|
34
|
+
|
|
35
|
+
|
|
36
|
+
class ServerError(TerrariumError):
|
|
37
|
+
"""5xx — orchestrator-side failure (transient: retried)."""
|
|
38
|
+
|
|
39
|
+
|
|
40
|
+
class TransportError(TerrariumError):
|
|
41
|
+
"""Connection/timeout failure reaching the orchestrator (transient: retried)."""
|
|
42
|
+
|
|
43
|
+
|
|
44
|
+
def from_status(e: httpx.HTTPStatusError) -> TerrariumError:
|
|
45
|
+
status = e.response.status_code
|
|
46
|
+
try:
|
|
47
|
+
body = e.response.json()
|
|
48
|
+
detail = body.get("detail") if isinstance(body, dict) else body
|
|
49
|
+
except Exception:
|
|
50
|
+
body = e.response.text
|
|
51
|
+
detail = body
|
|
52
|
+
msg = f"{status} {e.request.method} {e.request.url.path}: {detail}"
|
|
53
|
+
cls = {
|
|
54
|
+
401: AuthError, 403: AuthError, 404: NotFoundError,
|
|
55
|
+
409: ConflictError, 429: RateLimitError,
|
|
56
|
+
}.get(status, ServerError if status >= 500 else TerrariumError)
|
|
57
|
+
return cls(msg, status=status, body=body)
|
|
58
|
+
|
|
59
|
+
|
|
60
|
+
def is_transient(e: BaseException) -> bool:
|
|
61
|
+
"""Worth retrying: connection/timeout, 429, or 5xx."""
|
|
62
|
+
if isinstance(e, httpx.TransportError):
|
|
63
|
+
return True
|
|
64
|
+
if isinstance(e, httpx.HTTPStatusError):
|
|
65
|
+
return e.response.status_code == 429 or e.response.status_code >= 500
|
|
66
|
+
return False
|
terrarium/memory.py
ADDED
|
@@ -0,0 +1,171 @@
|
|
|
1
|
+
"""Structured, retrieval-backed memory for long-running agents.
|
|
2
|
+
|
|
3
|
+
The sandbox's ``/memory`` is an FS volume — no retrieval, and on k8s its RWO PVC can't
|
|
4
|
+
multi-attach, so concurrent sessions of one agent silently fork. This module offers memory as
|
|
5
|
+
**client tools** instead: the agent calls ``memory_search`` / ``memory_write`` / ``memory_get``,
|
|
6
|
+
and the handlers run in YOUR process against a real store. That gives keyword/vector retrieval
|
|
7
|
+
at scale AND sidesteps multi-attach entirely (no shared volume) — concurrent sessions and
|
|
8
|
+
scheduled jobs share one store.
|
|
9
|
+
|
|
10
|
+
from terrarium import TerrariumClient, TerrariumOptions
|
|
11
|
+
from terrarium.memory import SqliteMemory, memory_tools
|
|
12
|
+
|
|
13
|
+
store = SqliteMemory("assistant.db") # durable across sessions/processes
|
|
14
|
+
opts = TerrariumOptions(agent_id=agent_id, tools=memory_tools(store))
|
|
15
|
+
|
|
16
|
+
Bring your own backend (pgvector, Redis, an API) by implementing ``MemoryStore``; the reference
|
|
17
|
+
``SqliteMemory`` needs only the stdlib (SQLite FTS5, with a LIKE fallback).
|
|
18
|
+
"""
|
|
19
|
+
from __future__ import annotations
|
|
20
|
+
|
|
21
|
+
import asyncio
|
|
22
|
+
import json
|
|
23
|
+
import re
|
|
24
|
+
import sqlite3
|
|
25
|
+
import threading
|
|
26
|
+
import time
|
|
27
|
+
from typing import Any, Protocol, runtime_checkable
|
|
28
|
+
|
|
29
|
+
from .options import ClientTool, tool
|
|
30
|
+
|
|
31
|
+
__all__ = ["MemoryStore", "SqliteMemory", "memory_tools"]
|
|
32
|
+
|
|
33
|
+
|
|
34
|
+
def _coerce_tags(tags: Any) -> list[str]:
|
|
35
|
+
"""Models routinely pass `tags` as a comma-joined STRING despite the array schema. Coerce
|
|
36
|
+
to a clean list (split a string on commas/whitespace); ignore non-iterables. Without this,
|
|
37
|
+
`" ".join(a_string)` silently stores it character-by-character."""
|
|
38
|
+
if tags is None:
|
|
39
|
+
return []
|
|
40
|
+
if isinstance(tags, str):
|
|
41
|
+
return [t for t in re.split(r"[,\s]+", tags.strip()) if t]
|
|
42
|
+
if isinstance(tags, (list, tuple)):
|
|
43
|
+
return [str(t) for t in tags if t is not None and str(t)]
|
|
44
|
+
return []
|
|
45
|
+
|
|
46
|
+
|
|
47
|
+
@runtime_checkable
|
|
48
|
+
class MemoryStore(Protocol):
|
|
49
|
+
"""A pluggable memory backend. Implement these three async methods over your own store
|
|
50
|
+
(pgvector, Redis, an internal API) and pass it to :func:`memory_tools`."""
|
|
51
|
+
|
|
52
|
+
async def write(self, content: str, tags: list[str] | None = None) -> str:
|
|
53
|
+
"""Persist a memory; return its id."""
|
|
54
|
+
...
|
|
55
|
+
|
|
56
|
+
async def search(self, query: str, limit: int = 5) -> list[dict[str, Any]]:
|
|
57
|
+
"""Return up to ``limit`` memories matching ``query`` (most relevant first)."""
|
|
58
|
+
...
|
|
59
|
+
|
|
60
|
+
async def get(self, memory_id: str) -> dict[str, Any] | None:
|
|
61
|
+
"""Fetch one memory by id, or ``None``."""
|
|
62
|
+
...
|
|
63
|
+
|
|
64
|
+
|
|
65
|
+
def _now() -> str:
|
|
66
|
+
return time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime())
|
|
67
|
+
|
|
68
|
+
|
|
69
|
+
def _fts_match(query: str) -> str:
|
|
70
|
+
# Build an FTS5 query that RECALLS well. Three things matter:
|
|
71
|
+
# - quote each term so user text can't inject FTS5 operators;
|
|
72
|
+
# - prefix `*` (terms ≥3 chars) + the porter tokenizer so inflection doesn't lose a hit
|
|
73
|
+
# ("language" finds "languages", "cat" finds "cats");
|
|
74
|
+
# - join with OR, not the default implicit AND. An LLM's recall query is a bag of terms
|
|
75
|
+
# spread across MANY memories ("timezone language preference" where each fact lives in a
|
|
76
|
+
# different row); AND requires all terms in ONE row, so a broad query recalls nothing.
|
|
77
|
+
# OR surfaces every partial match and `ORDER BY rank` (in _search) floats the best up.
|
|
78
|
+
terms = [t for t in query.replace('"', " ").split() if t]
|
|
79
|
+
return " OR ".join((f'"{t}"*' if len(t) >= 3 else f'"{t}"') for t in terms)
|
|
80
|
+
|
|
81
|
+
|
|
82
|
+
class SqliteMemory:
|
|
83
|
+
"""Reference :class:`MemoryStore` over SQLite. Uses FTS5 for ranked full-text search when
|
|
84
|
+
available, falling back to ``LIKE``. Thread-safe (one connection + lock); pass a file path
|
|
85
|
+
for durability across processes, or the default ``:memory:`` for ephemeral use."""
|
|
86
|
+
|
|
87
|
+
def __init__(self, path: str = ":memory:") -> None:
|
|
88
|
+
self._c = sqlite3.connect(path, check_same_thread=False)
|
|
89
|
+
self._c.row_factory = sqlite3.Row
|
|
90
|
+
self._lock = threading.Lock()
|
|
91
|
+
self._fts = True
|
|
92
|
+
try:
|
|
93
|
+
# porter stemming so "languages"/"language", "preferences"/"preference", etc. recall
|
|
94
|
+
# each other (an LLM rarely recalls with the inflection it wrote). New tables only —
|
|
95
|
+
# an existing table keeps its tokenizer, but prefix matching (see _fts_match) still helps.
|
|
96
|
+
self._c.execute(
|
|
97
|
+
"CREATE VIRTUAL TABLE IF NOT EXISTS memories "
|
|
98
|
+
"USING fts5(content, tags, created_at UNINDEXED, tokenize='porter unicode61')")
|
|
99
|
+
except sqlite3.OperationalError: # SQLite built without FTS5 — degrade to LIKE
|
|
100
|
+
self._fts = False
|
|
101
|
+
self._c.execute(
|
|
102
|
+
"CREATE TABLE IF NOT EXISTS memories "
|
|
103
|
+
"(id INTEGER PRIMARY KEY, content TEXT, tags TEXT, created_at TEXT)")
|
|
104
|
+
self._c.commit()
|
|
105
|
+
|
|
106
|
+
# --- sync core (run off-loop via asyncio.to_thread) ---
|
|
107
|
+
def _write(self, content: str, tags: list[str] | None) -> str:
|
|
108
|
+
with self._lock:
|
|
109
|
+
cur = self._c.execute(
|
|
110
|
+
"INSERT INTO memories(content, tags, created_at) VALUES(?,?,?)",
|
|
111
|
+
(content, " ".join(_coerce_tags(tags)), _now())) # coerce: a stray str must not char-join
|
|
112
|
+
self._c.commit()
|
|
113
|
+
return str(cur.lastrowid)
|
|
114
|
+
|
|
115
|
+
def _search(self, query: str, limit: int) -> list[dict[str, Any]]:
|
|
116
|
+
with self._lock:
|
|
117
|
+
if self._fts and (m := _fts_match(query)):
|
|
118
|
+
rows = self._c.execute(
|
|
119
|
+
"SELECT rowid AS id, content, tags, created_at FROM memories "
|
|
120
|
+
"WHERE memories MATCH ? ORDER BY rank LIMIT ?", (m, limit)).fetchall()
|
|
121
|
+
else:
|
|
122
|
+
rows = self._c.execute(
|
|
123
|
+
"SELECT rowid AS id, content, tags, created_at FROM memories "
|
|
124
|
+
"WHERE content LIKE ? ORDER BY rowid DESC LIMIT ?",
|
|
125
|
+
(f"%{query}%", limit)).fetchall()
|
|
126
|
+
return [dict(r) for r in rows]
|
|
127
|
+
|
|
128
|
+
def _get(self, memory_id: str) -> dict[str, Any] | None:
|
|
129
|
+
with self._lock:
|
|
130
|
+
row = self._c.execute(
|
|
131
|
+
"SELECT rowid AS id, content, tags, created_at FROM memories WHERE rowid=?",
|
|
132
|
+
(memory_id,)).fetchone()
|
|
133
|
+
return dict(row) if row else None
|
|
134
|
+
|
|
135
|
+
async def write(self, content: str, tags: list[str] | None = None) -> str:
|
|
136
|
+
return await asyncio.to_thread(self._write, content, tags)
|
|
137
|
+
|
|
138
|
+
async def search(self, query: str, limit: int = 5) -> list[dict[str, Any]]:
|
|
139
|
+
return await asyncio.to_thread(self._search, query, limit)
|
|
140
|
+
|
|
141
|
+
async def get(self, memory_id: str) -> dict[str, Any] | None:
|
|
142
|
+
return await asyncio.to_thread(self._get, memory_id)
|
|
143
|
+
|
|
144
|
+
|
|
145
|
+
def memory_tools(store: MemoryStore) -> list[ClientTool]:
|
|
146
|
+
"""Ready-made client tools (``memory_write`` / ``memory_search`` / ``memory_get``) backed by
|
|
147
|
+
``store``. Pass the result to ``TerrariumOptions(tools=...)``; the handlers run in your
|
|
148
|
+
process, so the store and its credentials never enter the sandbox."""
|
|
149
|
+
|
|
150
|
+
@tool("memory_write", "Save a durable memory — a fact, preference, or decision — for later recall.",
|
|
151
|
+
{"content": {"type": "string", "description": "The thing to remember"},
|
|
152
|
+
"tags": {"type": "array", "items": {"type": "string"}, "description": "Optional labels"}})
|
|
153
|
+
async def memory_write(args: dict[str, Any]) -> str:
|
|
154
|
+
# Coerce at the protocol boundary so EVERY MemoryStore backend gets a clean list, even
|
|
155
|
+
# when the model passes tags as a comma-joined string (it routinely does).
|
|
156
|
+
mid = await store.write(str(args["content"]), _coerce_tags(args.get("tags")))
|
|
157
|
+
return f"saved memory {mid}"
|
|
158
|
+
|
|
159
|
+
@tool("memory_search", "Search durable memory by keyword or topic. Call this BEFORE answering "
|
|
160
|
+
"to recall relevant past context.",
|
|
161
|
+
{"query": {"type": "string"}, "limit": {"type": "integer", "description": "max results (default 5)"}})
|
|
162
|
+
async def memory_search(args: dict[str, Any]) -> str:
|
|
163
|
+
hits = await store.search(str(args["query"]), int(args.get("limit") or 5))
|
|
164
|
+
return json.dumps(hits) if hits else "no memories matched"
|
|
165
|
+
|
|
166
|
+
@tool("memory_get", "Fetch one memory by its id.", {"id": {"type": "string"}})
|
|
167
|
+
async def memory_get(args: dict[str, Any]) -> str:
|
|
168
|
+
m = await store.get(str(args["id"]))
|
|
169
|
+
return json.dumps(m) if m else "not found"
|
|
170
|
+
|
|
171
|
+
return [memory_write, memory_search, memory_get]
|
terrarium/messages.py
ADDED
|
@@ -0,0 +1,137 @@
|
|
|
1
|
+
"""Typed messages + content blocks — modelled on the Claude Agent SDK's message types
|
|
2
|
+
(``AssistantMessage``, ``TextBlock``, ``ToolUseBlock``, ``ResultMessage`` …) so iterating a
|
|
3
|
+
Terrarium turn feels like iterating a Claude Agent SDK response.
|
|
4
|
+
|
|
5
|
+
Terrarium streams a flat event log; ``parse_message`` lifts each event into the closest
|
|
6
|
+
Claude-SDK message. Every message keeps the original event under ``.raw`` as an escape hatch.
|
|
7
|
+
Terrarium-only events (human-in-the-loop) become ``QuestionMessage`` / ``PermissionMessage``.
|
|
8
|
+
"""
|
|
9
|
+
|
|
10
|
+
from __future__ import annotations
|
|
11
|
+
|
|
12
|
+
from dataclasses import dataclass, field
|
|
13
|
+
from typing import Any, Union
|
|
14
|
+
|
|
15
|
+
|
|
16
|
+
# --- content blocks (mirror claude_agent_sdk) ---
|
|
17
|
+
@dataclass
|
|
18
|
+
class TextBlock:
|
|
19
|
+
text: str
|
|
20
|
+
|
|
21
|
+
|
|
22
|
+
@dataclass
|
|
23
|
+
class ThinkingBlock:
|
|
24
|
+
thinking: str
|
|
25
|
+
signature: str = ""
|
|
26
|
+
|
|
27
|
+
|
|
28
|
+
@dataclass
|
|
29
|
+
class ToolUseBlock:
|
|
30
|
+
id: str
|
|
31
|
+
name: str
|
|
32
|
+
input: dict[str, Any]
|
|
33
|
+
|
|
34
|
+
|
|
35
|
+
@dataclass
|
|
36
|
+
class ToolResultBlock:
|
|
37
|
+
tool_use_id: str
|
|
38
|
+
content: Any
|
|
39
|
+
is_error: bool = False
|
|
40
|
+
|
|
41
|
+
|
|
42
|
+
ContentBlock = Union[TextBlock, ThinkingBlock, ToolUseBlock, ToolResultBlock]
|
|
43
|
+
|
|
44
|
+
|
|
45
|
+
# --- messages (mirror claude_agent_sdk) ---
|
|
46
|
+
@dataclass
|
|
47
|
+
class AssistantMessage:
|
|
48
|
+
content: list[ContentBlock]
|
|
49
|
+
model: str | None = None
|
|
50
|
+
raw: dict[str, Any] = field(default_factory=dict)
|
|
51
|
+
|
|
52
|
+
@property
|
|
53
|
+
def text(self) -> str:
|
|
54
|
+
"""Concatenated text of any TextBlocks (convenience)."""
|
|
55
|
+
return "".join(b.text for b in self.content if isinstance(b, TextBlock))
|
|
56
|
+
|
|
57
|
+
|
|
58
|
+
@dataclass
|
|
59
|
+
class UserMessage:
|
|
60
|
+
content: list[ContentBlock]
|
|
61
|
+
raw: dict[str, Any] = field(default_factory=dict)
|
|
62
|
+
|
|
63
|
+
|
|
64
|
+
@dataclass
|
|
65
|
+
class SystemMessage:
|
|
66
|
+
subtype: str
|
|
67
|
+
data: dict[str, Any]
|
|
68
|
+
raw: dict[str, Any] = field(default_factory=dict)
|
|
69
|
+
|
|
70
|
+
|
|
71
|
+
@dataclass
|
|
72
|
+
class ResultMessage:
|
|
73
|
+
subtype: str
|
|
74
|
+
total_cost_usd: float | None
|
|
75
|
+
usage: dict[str, Any]
|
|
76
|
+
duration_ms: int | None
|
|
77
|
+
num_turns: int | None
|
|
78
|
+
is_error: bool
|
|
79
|
+
raw: dict[str, Any] = field(default_factory=dict)
|
|
80
|
+
|
|
81
|
+
|
|
82
|
+
# --- Terrarium extensions (human-in-the-loop) ---
|
|
83
|
+
@dataclass
|
|
84
|
+
class QuestionMessage:
|
|
85
|
+
"""An AskUserQuestion prompt. Answer with ``session.answer(question_id, {...})`` or via a
|
|
86
|
+
``can_use_tool`` callback returning ``PermissionResultAllow(updated_input={"answers": ...})``."""
|
|
87
|
+
|
|
88
|
+
question_id: str
|
|
89
|
+
questions: list[dict[str, Any]]
|
|
90
|
+
raw: dict[str, Any] = field(default_factory=dict)
|
|
91
|
+
|
|
92
|
+
|
|
93
|
+
@dataclass
|
|
94
|
+
class PermissionMessage:
|
|
95
|
+
"""A gated tool-permission request (when the harness sets ``approval``). Resolve with
|
|
96
|
+
``session.decide(request_id, "allow"|"always"|"deny")`` or a ``can_use_tool`` callback."""
|
|
97
|
+
|
|
98
|
+
request_id: str
|
|
99
|
+
tool_name: str
|
|
100
|
+
input: dict[str, Any]
|
|
101
|
+
title: str | None = None
|
|
102
|
+
description: str | None = None
|
|
103
|
+
raw: dict[str, Any] = field(default_factory=dict)
|
|
104
|
+
|
|
105
|
+
|
|
106
|
+
Message = Union[AssistantMessage, UserMessage, SystemMessage, ResultMessage, QuestionMessage, PermissionMessage]
|
|
107
|
+
|
|
108
|
+
|
|
109
|
+
def parse_message(ev: dict[str, Any]) -> Message | None:
|
|
110
|
+
"""Convert a raw Terrarium event into a typed message, or ``None`` for events with no
|
|
111
|
+
message equivalent (status, rewind_point, answered, decided, …)."""
|
|
112
|
+
t = ev.get("type")
|
|
113
|
+
# The worker stamps the responding model on assistant events (emit(EV_ASSISTANT_TEXT,
|
|
114
|
+
# …, model=msg.model)); carry it through so AssistantMessage.model matches the real SDK.
|
|
115
|
+
model = ev.get("model")
|
|
116
|
+
if t == "assistant_text":
|
|
117
|
+
return AssistantMessage(content=[TextBlock(str(ev.get("text", "")))], model=model, raw=ev)
|
|
118
|
+
if t == "thinking":
|
|
119
|
+
return AssistantMessage(content=[ThinkingBlock(str(ev.get("text", "")))], model=model, raw=ev)
|
|
120
|
+
if t == "tool_use":
|
|
121
|
+
return AssistantMessage(content=[ToolUseBlock(str(ev.get("id", "")), str(ev.get("name", "")), ev.get("input") or {})], model=model, raw=ev)
|
|
122
|
+
if t == "tool_result":
|
|
123
|
+
return UserMessage(content=[ToolResultBlock(str(ev.get("tool_use_id", "")), ev.get("content"), bool(ev.get("is_error")))], raw=ev)
|
|
124
|
+
if t == "result":
|
|
125
|
+
return ResultMessage(
|
|
126
|
+
subtype=str(ev.get("subtype", "")), total_cost_usd=ev.get("total_cost_usd"),
|
|
127
|
+
usage=ev.get("usage") or {}, duration_ms=ev.get("duration_ms"),
|
|
128
|
+
num_turns=ev.get("num_turns"), is_error=bool(ev.get("is_error")), raw=ev)
|
|
129
|
+
if t in ("system", "ready", "session_start", "error", "session_end", "worker_lost"):
|
|
130
|
+
data = ev.get("data") if isinstance(ev.get("data"), dict) else {k: v for k, v in ev.items() if k not in ("type", "seq", "ts")}
|
|
131
|
+
return SystemMessage(subtype=str(ev.get("subtype", t)), data=data, raw=ev)
|
|
132
|
+
if t == "question":
|
|
133
|
+
return QuestionMessage(str(ev.get("question_id", "")), ev.get("questions") or [], raw=ev)
|
|
134
|
+
if t == "permission":
|
|
135
|
+
return PermissionMessage(str(ev.get("request_id", "")), str(ev.get("tool_name", "")),
|
|
136
|
+
ev.get("input") or {}, ev.get("title"), ev.get("description"), raw=ev)
|
|
137
|
+
return None
|
terrarium/options.py
ADDED
|
@@ -0,0 +1,212 @@
|
|
|
1
|
+
"""Options + permission types — modelled on the Claude Agent SDK (``ClaudeAgentOptions``,
|
|
2
|
+
``can_use_tool``, ``PermissionResultAllow/Deny``) so the two SDKs feel the same, with extra
|
|
3
|
+
fields for Terrarium's own features (personas, approval gating, egress profiles, memory).
|
|
4
|
+
"""
|
|
5
|
+
|
|
6
|
+
from __future__ import annotations
|
|
7
|
+
|
|
8
|
+
from dataclasses import asdict, dataclass, field, is_dataclass
|
|
9
|
+
from typing import Any, Awaitable, Callable, Union
|
|
10
|
+
|
|
11
|
+
|
|
12
|
+
@dataclass
|
|
13
|
+
class AgentDefinition:
|
|
14
|
+
"""A programmatic subagent — field-for-field the Claude Agent SDK's ``AgentDefinition``
|
|
15
|
+
(same names, incl. its camelCase), so ported code drops in unchanged. Plain dicts (with
|
|
16
|
+
snake_case or camelCase keys) are accepted anywhere this is."""
|
|
17
|
+
|
|
18
|
+
description: str
|
|
19
|
+
prompt: str
|
|
20
|
+
tools: list[str] | None = None
|
|
21
|
+
disallowedTools: list[str] | None = None # noqa: N815 — Claude-SDK name
|
|
22
|
+
model: str | None = None
|
|
23
|
+
skills: list[str] | None = None
|
|
24
|
+
memory: str | None = None # "user" | "project" | "local"
|
|
25
|
+
mcpServers: "list[str | dict[str, Any]] | None" = None # noqa: N815
|
|
26
|
+
initialPrompt: str | None = None # noqa: N815
|
|
27
|
+
maxTurns: int | None = None # noqa: N815
|
|
28
|
+
background: bool | None = None
|
|
29
|
+
effort: "str | int | None" = None
|
|
30
|
+
permissionMode: str | None = None # noqa: N815
|
|
31
|
+
|
|
32
|
+
|
|
33
|
+
def _agent_spec(a: "AgentDefinition | dict[str, Any]") -> dict[str, Any]:
|
|
34
|
+
if is_dataclass(a) and not isinstance(a, type):
|
|
35
|
+
return {k: v for k, v in asdict(a).items() if v is not None}
|
|
36
|
+
return dict(a)
|
|
37
|
+
|
|
38
|
+
|
|
39
|
+
@dataclass
|
|
40
|
+
class ToolPermissionContext:
|
|
41
|
+
"""Context passed to a ``can_use_tool`` callback (mirrors the Claude SDK's). For a
|
|
42
|
+
permission request it carries the request id + the agent's title/description; for an
|
|
43
|
+
AskUserQuestion it wraps the question event."""
|
|
44
|
+
|
|
45
|
+
request_id: str = ""
|
|
46
|
+
title: str | None = None
|
|
47
|
+
description: str | None = None
|
|
48
|
+
raw: dict[str, Any] = field(default_factory=dict)
|
|
49
|
+
|
|
50
|
+
|
|
51
|
+
@dataclass
|
|
52
|
+
class PermissionResultAllow:
|
|
53
|
+
"""Allow a tool call. ``always`` maps to Terrarium's "always allow (this session)".
|
|
54
|
+
For AskUserQuestion, put the answers under ``updated_input={"answers": {...}}`` —
|
|
55
|
+
exactly as the Claude Agent SDK expects."""
|
|
56
|
+
|
|
57
|
+
updated_input: dict[str, Any] | None = None
|
|
58
|
+
always: bool = False
|
|
59
|
+
|
|
60
|
+
|
|
61
|
+
@dataclass
|
|
62
|
+
class PermissionResultDeny:
|
|
63
|
+
"""Deny a tool call (optionally with a message the agent sees)."""
|
|
64
|
+
|
|
65
|
+
message: str = ""
|
|
66
|
+
|
|
67
|
+
|
|
68
|
+
@dataclass
|
|
69
|
+
class ClientTool:
|
|
70
|
+
"""A custom tool whose handler runs in YOUR process (with your application context) —
|
|
71
|
+
bridged to the sandboxed agent. The agent's tool input crosses out, your handler runs
|
|
72
|
+
here, and only the result you return crosses back in; your code/state/secrets never
|
|
73
|
+
enter the sandbox. Create via :func:`tool`."""
|
|
74
|
+
|
|
75
|
+
name: str
|
|
76
|
+
description: str
|
|
77
|
+
input_schema: dict[str, Any]
|
|
78
|
+
handler: "Callable[[dict[str, Any]], Any]" # (input) -> str | {"content","is_error"} | awaitable
|
|
79
|
+
|
|
80
|
+
def schema(self) -> dict[str, Any]:
|
|
81
|
+
return {"name": self.name, "description": self.description, "input_schema": self.input_schema}
|
|
82
|
+
|
|
83
|
+
|
|
84
|
+
def tool(name: str, description: str, input_schema: dict[str, Any] | None = None):
|
|
85
|
+
"""Decorator turning an (async or sync) ``handler(input) -> result`` into a :class:`ClientTool`.
|
|
86
|
+
Pass the resulting tools to ``TerrariumOptions(tools=[...])``. Mirrors the Claude Agent SDK's
|
|
87
|
+
``@tool`` — except the handler executes client-side, never in the sandbox.
|
|
88
|
+
|
|
89
|
+
@tool("get_user", "Look up a user by id", {"id": {"type": "string"}})
|
|
90
|
+
async def get_user(args):
|
|
91
|
+
return await db.users.find(args["id"]) # your app context, your process
|
|
92
|
+
"""
|
|
93
|
+
def deco(fn):
|
|
94
|
+
return ClientTool(name=name, description=description, input_schema=input_schema or {}, handler=fn)
|
|
95
|
+
return deco
|
|
96
|
+
|
|
97
|
+
|
|
98
|
+
PermissionResult = Union[PermissionResultAllow, PermissionResultDeny]
|
|
99
|
+
# A permission callback (like the Claude Agent SDK's). May be sync OR async — the client
|
|
100
|
+
# awaits the result if it's awaitable.
|
|
101
|
+
CanUseTool = Callable[[str, dict[str, Any], ToolPermissionContext], Union[PermissionResult, Awaitable[PermissionResult]]]
|
|
102
|
+
|
|
103
|
+
# Harness keys the orchestrator understands (kept in sync with terrarium/harness.py).
|
|
104
|
+
_HARNESS_FIELDS = (
|
|
105
|
+
"model", "system_mode", "custom_prompt", "permission_mode", "allowed_tools", "builtin_tools",
|
|
106
|
+
"thinking", "effort", "fallback_model", "max_thinking_tokens", "betas",
|
|
107
|
+
"max_turns", "max_budget_usd", "mcp_servers", "agents", "skills",
|
|
108
|
+
"interactive", "approval", "setting_sources", "env", "extra_options",
|
|
109
|
+
"environments", "memory_mode",
|
|
110
|
+
)
|
|
111
|
+
|
|
112
|
+
|
|
113
|
+
@dataclass
|
|
114
|
+
class TerrariumOptions:
|
|
115
|
+
"""Configuration for a Terrarium agent/session — a superset of the Claude Agent SDK's
|
|
116
|
+
``ClaudeAgentOptions``.
|
|
117
|
+
|
|
118
|
+
The first block mirrors the Claude SDK field-for-field (so existing knowledge carries
|
|
119
|
+
over); the second block adds Terrarium-only capabilities. ``None`` means "leave to the
|
|
120
|
+
orchestrator default" — only set fields are sent.
|
|
121
|
+
"""
|
|
122
|
+
|
|
123
|
+
# --- Claude Agent SDK-aligned ---
|
|
124
|
+
model: str | None = None # "sonnet" | "opus" | "haiku" | full id
|
|
125
|
+
system_prompt: "str | dict[str, Any] | None" = None # custom prompt (→ custom persona), or the
|
|
126
|
+
# Claude-SDK preset dict {"type":"preset","preset":"claude_code"} (→ claude_code persona)
|
|
127
|
+
allowed_tools: list[str] | None = None # AUTO-APPROVE list (Claude-SDK: skips the prompt;
|
|
128
|
+
# does NOT change which tools the agent has)
|
|
129
|
+
builtin_tools: "list[str] | dict[str, Any] | None" = None # AVAILABILITY allowlist — the base set
|
|
130
|
+
# of built-in tools the agent may use. None=all defaults; ["Read","Grep"]=only those; []=none;
|
|
131
|
+
# {"type":"preset","preset":"claude_code"}=all defaults. This is the real "restrict tools" knob.
|
|
132
|
+
permission_mode: str | None = None # default | acceptEdits | plan | bypassPermissions
|
|
133
|
+
max_turns: int | None = None
|
|
134
|
+
max_budget_usd: float | None = None
|
|
135
|
+
mcp_servers: dict[str, Any] | None = None
|
|
136
|
+
agents: "dict[str, AgentDefinition | dict[str, Any]] | None" = None # programmatic subagents
|
|
137
|
+
# (SDK `agents`): name -> AgentDefinition, or a plain dict ({"description", "prompt",
|
|
138
|
+
# and optionally "tools", "disallowed_tools", "model", "skills", "max_turns", ...})
|
|
139
|
+
thinking: dict[str, Any] | None = None # e.g. {"type": "adaptive"} / {"type": "enabled", "budget_tokens": N}
|
|
140
|
+
effort: str | None = None # low | medium | high | xhigh | max
|
|
141
|
+
fallback_model: str | None = None # model to retry with on overload/refusal
|
|
142
|
+
max_thinking_tokens: int | None = None # hard cap on thinking tokens per turn
|
|
143
|
+
betas: list[str] | None = None # API beta flags to opt into
|
|
144
|
+
setting_sources: list[str] | None = None # ["user", "project", "local"]
|
|
145
|
+
env: dict[str, str] | None = None
|
|
146
|
+
skills: "bool | list[str] | str | None" = None # True=mount+discover; "all"; [names]=only these;
|
|
147
|
+
# []=NO skills at all (hides the CLI's built-in skills too — the "bare harness" setting)
|
|
148
|
+
memory_mode: str | None = None # "volume" (default, durable mount) | "synced"
|
|
149
|
+
# (snapshot in/out — much faster k8s launch, loses writes since the last turn if the pod
|
|
150
|
+
# dies abruptly) | "none" (container-local scratch, discarded on stop)
|
|
151
|
+
can_use_tool: CanUseTool | None = None # (tool, input, ctx) -> Allow | Deny
|
|
152
|
+
tools: list[ClientTool] | None = None # custom tools that run in YOUR process (see `tool`)
|
|
153
|
+
|
|
154
|
+
# --- Terrarium extensions ---
|
|
155
|
+
system_mode: str | None = None # minimal | claude_code | custom | assistant
|
|
156
|
+
custom_prompt: str | None = None # alias for system_prompt (explicit persona text)
|
|
157
|
+
approval: "str | list[str] | None" = None # off | edits | all | [tool names] — human-in-the-loop gating
|
|
158
|
+
interactive: bool | None = None # allow AskUserQuestion / approval prompts to block for an operator
|
|
159
|
+
environments: list[str] | None = None # attach to {secrets, egress} bundles — the sole per-agent
|
|
160
|
+
# egress + secret-scoping mechanism. None/[] = no operator secrets + global egress; a list =
|
|
161
|
+
# ONLY those environments' secrets, and egress merged from their profiles (enforce wins; hosts union).
|
|
162
|
+
memory_scope: str | None = None # share memory with another agent id
|
|
163
|
+
extra_options: dict[str, Any] | None = None # new SDK fields only; Terrarium-managed keys are rejected
|
|
164
|
+
|
|
165
|
+
# --- session attach / routing (not part of the harness) ---
|
|
166
|
+
agent_id: str | None = None # attach the session to an existing agent
|
|
167
|
+
title: str | None = None
|
|
168
|
+
|
|
169
|
+
def to_harness(self) -> dict[str, Any]:
|
|
170
|
+
"""Serialize the harness-relevant fields for the create-agent/create-session API."""
|
|
171
|
+
h: dict[str, Any] = {}
|
|
172
|
+
# `system_prompt` is the Claude-SDK name for a bespoke prompt → Terrarium custom
|
|
173
|
+
# persona. Its preset dict form maps onto the matching system_mode.
|
|
174
|
+
sp = self.system_prompt
|
|
175
|
+
preset_mode: str | None = None
|
|
176
|
+
if isinstance(sp, dict):
|
|
177
|
+
if sp.get("type") == "preset" and sp.get("preset") == "claude_code" and "append" not in sp:
|
|
178
|
+
sp = None
|
|
179
|
+
preset_mode = "claude_code"
|
|
180
|
+
else: # unknown preset / "append" — no Terrarium equivalent; refuse loudly
|
|
181
|
+
raise ValueError(
|
|
182
|
+
f"unsupported system_prompt {sp!r}: only "
|
|
183
|
+
'{"type": "preset", "preset": "claude_code"} (without "append") maps to a '
|
|
184
|
+
"persona; for appended guidance use system_prompt=<full text> instead"
|
|
185
|
+
)
|
|
186
|
+
# A claude_code preset can't be combined with a custom prompt or an explicit
|
|
187
|
+
# system_mode — either would silently clobber the preset. Refuse loudly (matching the
|
|
188
|
+
# unknown-preset/unknown-kwarg behavior) rather than picking one.
|
|
189
|
+
if preset_mode and (self.custom_prompt is not None or self.system_mode is not None):
|
|
190
|
+
raise ValueError(
|
|
191
|
+
'conflicting persona: system_prompt={"preset": "claude_code"} cannot be '
|
|
192
|
+
"combined with custom_prompt or an explicit system_mode — set only one"
|
|
193
|
+
)
|
|
194
|
+
custom = self.custom_prompt or sp
|
|
195
|
+
if custom is not None:
|
|
196
|
+
h["custom_prompt"] = custom
|
|
197
|
+
h["system_mode"] = self.system_mode or "custom"
|
|
198
|
+
elif self.system_mode is not None:
|
|
199
|
+
h["system_mode"] = self.system_mode
|
|
200
|
+
elif preset_mode is not None:
|
|
201
|
+
h["system_mode"] = preset_mode
|
|
202
|
+
for k in _HARNESS_FIELDS:
|
|
203
|
+
if k in ("custom_prompt", "system_mode"):
|
|
204
|
+
continue
|
|
205
|
+
v = getattr(self, k, None)
|
|
206
|
+
if v is not None:
|
|
207
|
+
h[k] = v
|
|
208
|
+
if self.agents: # AgentDefinition dataclasses → plain dicts for the wire
|
|
209
|
+
h["agents"] = {name: _agent_spec(a) for name, a in self.agents.items()}
|
|
210
|
+
if self.tools: # send only the SCHEMAS; the handlers stay client-side
|
|
211
|
+
h["client_tools"] = [t.schema() for t in self.tools]
|
|
212
|
+
return h
|
terrarium/py.typed
ADDED
|
File without changes
|