retoor
8df5cf7a45
tai.py (8643 lines) is now a thin facade over tai_pkg/ (33 feature modules, one class or feature per file). tai-orig.py preserves the pre-split single-file version byte for byte (sha256 verified against snapshot); import tai and ./tai.py behavior unchanged. Equivalence proof, all run on this tree immediately pre-commit: differential battery tai-orig.py vs facade 3718 checks 0 failures (surface, pure fuzz, tool dispatch, store/audit, release ops, fake-HOME install roundtrip, CLI runs); test_tai.py 261 OK; test_seal.py 16 OK; --version/--help byte-identical; dir(tai) zero removals. Fix included: opencode wire-protocol schema echo leaked into replies; opencode_parse now scans with JSONDecoder.raw_decode capped at 200 scans (13 OpencodeParseTests). Test count 261 = 237 HEAD + 8 pre-existing uncommitted ToolArgumentsTests (now committed with this change, disclosed) + 13 OpencodeParseTests + 3 EmbeddedStringsTests. Patch-target-only test migration (24 hunks), zero assertion changes. tai_pkg clean-regenerated after mid-verification tamper was found and reverted.
117 lines
17 KiB
Python
117 lines
17 KiB
Python
"""Agent tool registry, schemas, and lazy selection."""
|
|
|
|
import os
|
|
import re
|
|
|
|
from tai_pkg.consts import Consts
|
|
from tai_pkg.plexer import PlexerControllerFactory
|
|
|
|
__all__ = ["ToolRegistry", "TOOL_SCHEMAS", "TOOL_ENV_GATES"]
|
|
|
|
|
|
class ToolRegistry:
|
|
"""ToolRegistry."""
|
|
|
|
@staticmethod
|
|
def tool_schema(name, description, properties, required):
|
|
return {"type": "function", "function": {"name": name, "description": description, "parameters": {"type": "object", "properties": properties, "required": required}}}
|
|
|
|
@staticmethod
|
|
def tool_env_ok(name):
|
|
gate = TOOL_ENV_GATES.get(name)
|
|
return gate is None or gate()
|
|
|
|
@staticmethod
|
|
def tool_catalog():
|
|
lines = ["", "", "## Tool catalog (core loads always; name a lazy tool or topic to load it)"]
|
|
for schema in TOOL_SCHEMAS:
|
|
name = schema["function"]["name"]
|
|
if not ToolRegistry.tool_env_ok(name):
|
|
continue
|
|
if name in Consts.CORE_TOOLS:
|
|
continue
|
|
first = schema["function"]["description"].split(".")[0][:80]
|
|
lines.append("- %s: %s" % (name, first))
|
|
return "\n".join(lines)
|
|
|
|
@staticmethod
|
|
def select_tools(text):
|
|
words = set(re.findall(r"[a-z0-9]+", (text or "").lower()))
|
|
picked = []
|
|
for schema in TOOL_SCHEMAS:
|
|
name = schema["function"]["name"]
|
|
if not ToolRegistry.tool_env_ok(name):
|
|
continue
|
|
if name in Consts.CORE_TOOLS:
|
|
picked.append(schema)
|
|
continue
|
|
tokens = [item for item in name.split("_") if len(item) >= 4 and item not in Consts.LAZY_NAME_STOP]
|
|
if tokens and set(tokens) & words:
|
|
picked.append(schema)
|
|
continue
|
|
if set(Consts.TOOL_TAGS.get(name, ())) & words:
|
|
picked.append(schema)
|
|
continue
|
|
desc = set(word for word in re.findall(r"[a-z0-9]+", schema["function"]["description"].lower()) if len(word) >= 6)
|
|
if desc & words:
|
|
picked.append(schema)
|
|
return picked
|
|
|
|
|
|
TOOL_SCHEMAS = [
|
|
ToolRegistry.tool_schema("shell", "Run a shell command. Output is truncated. On home, destructive commands ask the user first. On sandbox, commands run isolated in the container and need no approval. Pass secrets to expose named vault secrets as TAI_SECRET_<NAME> variables for that command only; using secrets asks the user first. For local commands only: use web_fetch for HTTP(S), never curl or wget.", {"command": {"type": "string"}, "workdir": {"type": "string"}, "secrets": {"type": "array"}}, ["command"]),
|
|
ToolRegistry.tool_schema("read_file", "Read a text file. Large files are truncated.", {"path": {"type": "string"}}, ["path"]),
|
|
ToolRegistry.tool_schema("write_file", "Write content to a file, creating parent directories. Overwrites existing files.", {"path": {"type": "string"}, "content": {"type": "string"}}, ["path", "content"]),
|
|
ToolRegistry.tool_schema("edit_file", "Replace one exact text match in a file. Fails unless the match is unique.", {"path": {"type": "string"}, "find": {"type": "string"}, "replace": {"type": "string"}}, ["path", "find", "replace"]),
|
|
ToolRegistry.tool_schema("web_search", "Search the web. Set images for image search, content to include fetched page text.", {"query": {"type": "string"}, "images": {"type": "boolean"}, "content": {"type": "boolean"}}, ["query"]),
|
|
ToolRegistry.tool_schema("web_fetch", "HTTP client: fetch a URL with any method and return status, content type, size, and content. GET by default; pass method (POST, PUT, PATCH, DELETE, HEAD, OPTIONS), headers, and a string body for APIs. HTML becomes readable text, JSON stays raw, raw forces the raw body. Use this for all HTTP(S), never curl or wget via shell. Pass auth_secret to send a vault secret as an Authorization header; the value stays sealed and secret use asks the user first.", {"url": {"type": "string"}, "method": {"type": "string"}, "headers": {"type": "object"}, "body": {"type": "string"}, "timeout": {"type": "integer"}, "raw": {"type": "boolean"}, "auth_secret": {"type": "string"}, "auth_header": {"type": "string"}, "auth_prefix": {"type": "string"}}, ["url"]),
|
|
ToolRegistry.tool_schema("speak", "Synthesize text to speech with a free voice, save the MP3, and play it when a player exists.", {"text": {"type": "string"}, "voice": {"type": "string"}}, ["text"]),
|
|
ToolRegistry.tool_schema("listen", "Record from the microphone for some seconds and transcribe it. Needs OS audio tools.", {"seconds": {"type": "integer"}}, []),
|
|
ToolRegistry.tool_schema("remember", "Update your own system message with new knowledge or behavior. The instruction is merged into the existing system message, unrelated parts stay intact. Use it to forget as well by instructing what to remove. Execute by default when behavior, preferences, or durable facts need updating. Never pass passwords, tokens, or secret values here; store those with store_secret instead.", {"instruction": {"type": "string"}}, ["instruction"]),
|
|
ToolRegistry.tool_schema("recall", "Search past session memory of the current profile by keyword, optionally filtered by tags.", {"query": {"type": "string"}, "tags": {"type": "array"}}, ["query"]),
|
|
ToolRegistry.tool_schema("load_skill", "Load a skill by name. Returns full instructions plus bundled file paths. Missing skills with a blueprint are researched and built on demand.", {"name": {"type": "string"}}, ["name"]),
|
|
ToolRegistry.tool_schema("get_current_terminal_content", "Capture visible text of the current tmux pane including scrollback. Works inside tmux or against a running tmux server.", {"lines": {"type": "integer"}}, []),
|
|
ToolRegistry.tool_schema("export_terminal_log", "Only available inside a live tmux session. Exports the full pane scrollback (tmux capture-pane -pS -, the entire history, not just what's visible) to a log file under tai's home dir, then chunks it and indexes the chunks for retrieval with SQLite FTS5 + bm25 (lexical only, no embeddings). From then on, use the search tool (kind chunk) to answer questions about anything that happened in this terminal; it returns the actual chunk text, not just a snippet.", {"tags": {"type": "array"}}, []),
|
|
ToolRegistry.tool_schema("plexer_send", "Only available when a tai-managed terminal session is active (either tai's own native wrapper, or an active tmux pane). Sends literal keystrokes into that session's shell, as if typed. Set enter to also press Enter and actually run it. No output is returned by this call; follow up with plexer_read.", {"text": {"type": "string"}, "enter": {"type": "boolean"}}, ["text"]),
|
|
ToolRegistry.tool_schema("plexer_read", "Only available when a tai-managed terminal session is active. Reads back the last N lines the session's shell has produced (its actual rendered output, not a tool-call summary).", {"lines": {"type": "integer"}}, []),
|
|
ToolRegistry.tool_schema("plexer_type", "Only available when a tai-managed terminal session is active. Types text into the session one character at a time with a delay between keystrokes, for a demo/animated-typing effect, then presses Enter unless enter is set to false.", {"text": {"type": "string"}, "delay": {"type": "number"}, "enter": {"type": "boolean"}}, ["text"]),
|
|
ToolRegistry.tool_schema("plexer_status", "Only available when a tai-managed terminal session is active. Reports which backend is managing the current session (tai's own native pty wrapper, or tmux) and backend-specific details (pid/socket/log, or pane id).", {}, []),
|
|
ToolRegistry.tool_schema("fork", "Spawn a background subagent with its own context that works while you continue. Returns an agent id immediately. Collect its summarized result with poll. Subagents get a smaller step budget and a time limit.", {"task": {"type": "string"}, "timeout": {"type": "integer"}, "profile": {"type": "string"}}, ["task"]),
|
|
ToolRegistry.tool_schema("poll", "Collect a background subagent result by id. Waits up to wait seconds, then reports running or the result.", {"id": {"type": "integer"}, "wait": {"type": "integer"}}, ["id"]),
|
|
ToolRegistry.tool_schema("sysinfo", "Inspect the host machine where tai runs. Runs os, python, venv, root, container, binaries, cpu, and disk checks in parallel and reports each with timing. Pass checks to run a named subset.", {"checks": {"type": "array"}}, []),
|
|
ToolRegistry.tool_schema("create_skill", "Create a new agent skill by name from a brief. Runs a dedicated deep-research worker on this machine (it calls sysinfo first, then researches with web search and shell until it has verified information) and writes SKILL.md plus supporting files into the skill directory. Runs synchronously and can take many minutes. Scope project writes under ./.tai/skills (the default), scope home under ~/.tai/skills.", {"name": {"type": "string"}, "brief": {"type": "string"}, "scope": {"type": "string"}, "timeout": {"type": "integer"}}, ["name", "brief"]),
|
|
ToolRegistry.tool_schema("install_skill", "Install a skill that already exists somewhere else, instead of building one from scratch. source can be a local directory containing SKILL.md (or one subdirectory that does), a bare SKILL.md file, a local .zip, or a git URL (needs git installed; most sources need nothing beyond the standard library). Validates the SKILL.md frontmatter, then lands it under tai's own skill directories: copy vendors a private snapshot, symlink points at the source in place (the default when the source already sits in a recognized skills folder like .claude/skills or .agents/skills, so there is one canonical copy instead of a drifting duplicate). Scope defaults like create_skill: project when the cwd looks like a project, else home. Overwriting an existing skill of the same name asks the user first unless force is set.", {"source": {"type": "string"}, "name": {"type": "string"}, "scope": {"type": "string"}, "mode": {"type": "string"}, "force": {"type": "boolean"}}, ["source"]),
|
|
ToolRegistry.tool_schema("store_secret", "Store a password, token, or secret in the sealed vault under a name. Values are encrypted at rest, never shown back, and only usable by reference; prefer /secret set in the REPL so the value never enters the conversation. Optional username, host, port, notes, expires (ISO datetime), and tags describe it; name and value suffice, never pester the user for more.", {"name": {"type": "string"}, "value": {"type": "string"}, "username": {"type": "string"}, "host": {"type": "string"}, "port": {"type": "integer"}, "notes": {"type": "string"}, "expires": {"type": "string"}, "tags": {"type": "array"}}, ["name", "value"]),
|
|
ToolRegistry.tool_schema("list_secrets", "List vault secret names with metadata and tags. Values are never revealed.", {}, []),
|
|
ToolRegistry.tool_schema("delete_secret", "Delete a vault secret by name. Asks the user first.", {"name": {"type": "string"}}, ["name"]),
|
|
ToolRegistry.tool_schema("schedule", "Schedule a prompt to run at a future time or on a repeating interval. Exactly one of at (ISO datetime, naive means local time) or every (interval seconds, minimum 60) is required. When due, a background subagent runs the prompt without blocking anyone; inspect outcomes with schedules. Asks the user first. Optional tags label it.", {"at": {"type": "string"}, "every": {"type": "integer"}, "prompt": {"type": "string"}, "name": {"type": "string"}, "timeout": {"type": "integer"}, "profile": {"type": "string"}, "tags": {"type": "array"}}, ["prompt"]),
|
|
ToolRegistry.tool_schema("unschedule", "Delete a scheduled prompt by id. Asks the user first.", {"id": {"type": "integer"}}, ["id"]),
|
|
ToolRegistry.tool_schema("schedules", "List scheduled prompts with due times and last outcomes.", {}, []),
|
|
ToolRegistry.tool_schema("record_save", "Save text as a tagged record in the vault and get back a mem:<hash> id. Use for anything too big or too durable for chat: research notes, command output, transcripts. Kind is one of note, output, file, research, transcript. Tag rules: lowercase singular, reuse tags from the tags tool, at most 5; known words in the text attach automatically and the record links to recent same-tag records.", {"title": {"type": "string"}, "content": {"type": "string"}, "kind": {"type": "string"}, "tags": {"type": "array"}}, ["content"]),
|
|
ToolRegistry.tool_schema("record_read", "Read a slice of a record by mem:<hash> id. Offset and limit page through huge records without loading them whole.", {"id": {"type": "string"}, "offset": {"type": "integer"}, "limit": {"type": "integer"}}, ["id"]),
|
|
ToolRegistry.tool_schema("record_search", "Search records by text, kind, and tags. Returns mem:<hash> ids with sizes so you can page in only what you need.", {"query": {"type": "string"}, "kind": {"type": "string"}, "tags": {"type": "array"}, "limit": {"type": "integer"}}, []),
|
|
ToolRegistry.tool_schema("record_delete", "Delete a record by mem:<hash> id. Asks the user first.", {"id": {"type": "string"}}, ["id"]),
|
|
ToolRegistry.tool_schema("graph_link", "Link two vault nodes with a relation: mem:<hash> records, secret:<name> secrets, sched:<id> schedules. Both endpoints must exist.", {"src": {"type": "string"}, "dst": {"type": "string"}, "relation": {"type": "string"}}, ["src", "dst"]),
|
|
ToolRegistry.tool_schema("graph_query", "Show the neighborhood of a vault node (mem:<hash>, secret:<name>, sched:<id>) by breadth-first traversal. Depth caps at 4.", {"node": {"type": "string"}, "depth": {"type": "integer"}, "limit": {"type": "integer"}}, ["node"]),
|
|
ToolRegistry.tool_schema("delete_file", "Delete a file. Requires reading it first; the full pre-image stays in the audit trail so restore can undelete. Asks the user first.", {"path": {"type": "string"}}, ["path"]),
|
|
ToolRegistry.tool_schema("audit", "Show the file audit trail for time travel: every write, edit, delete, shell snapshot, release, and restore with sizes and messages. Filter by exact path or tag, or rank by a full-text query over paths and messages.", {"path": {"type": "string"}, "tag": {"type": "string"}, "query": {"type": "string"}, "limit": {"type": "integer"}}, []),
|
|
ToolRegistry.tool_schema("restore", "Restore a file to an audit row: the post-image for writes and edits, the pre-image for deletes and shell snapshots. Refuses truncated images. Asks the user first.", {"id": {"type": "integer"}}, ["id"]),
|
|
ToolRegistry.tool_schema("release", "Cut a release of the running script: bump VERSION by part (patch for fixes, minor for features, major for breaking changes), keep a versioned backup, log the message in the audit trail. Asks the user first.", {"part": {"type": "string"}, "message": {"type": "string"}}, ["part", "message"]),
|
|
ToolRegistry.tool_schema("tags", "List vault tags with usage counts, most used first. Consult before tagging so new items reuse established tags instead of inventing synonyms.", {"prefix": {"type": "string"}, "limit": {"type": "integer"}}, []),
|
|
ToolRegistry.tool_schema("install", "Manage tai installations: binary, command-not-found hook and plexer auto-wrap hook for bash/zsh/fish, venv, scheduler service, telegram service, container. Status reports what exists. Install adds missing pieces (skipping a shell hook loudly if that shell isn't on the machine), upgrade refreshes in place and restarts services, reinstall rebuilds artifacts, uninstall removes them. Service changes take effect immediately. The vault is never touched by any action. Non-status actions ask the user first.", {"action": {"type": "string"}, "targets": {"type": "array"}}, ["action"]),
|
|
ToolRegistry.tool_schema("create_bot", "Create a bot in the current profile: a name plus its own system message and history. Give at least one of description, rules, or behavior; optional nicknames register as @mention aliases (a short form of the name registers automatically). Names are lowercase.", {"name": {"type": "string"}, "description": {"type": "string"}, "rules": {"type": "string"}, "behavior": {"type": "string"}, "nicknames": {"type": "array"}}, ["name"]),
|
|
ToolRegistry.tool_schema("search", "Search everything at once with ranked full-text matching (SQLite FTS5 + bm25, no vectors): records, episodic events, the audit trail, and RAG chunks. Chunks are the retrieval-grade granular hits (e.g. one slice of an exported terminal log) and come back with their actual text, not just a snippet; use the chunk's parent id with record_read to see more surrounding context. One call replaces paging through each store separately. expand adds one graph hop per record hit.", {"query": {"type": "string"}, "kinds": {"type": "array"}, "limit": {"type": "integer"}, "expand": {"type": "boolean"}}, ["query"]),
|
|
]
|
|
|
|
|
|
# Tools gated on top of the lazy keyword match: a gate must pass before a tool can even be
|
|
# offered to the model, regardless of relevance. Used for tools that only make sense in a
|
|
# specific environment (e.g. a tmux pane to capture).
|
|
TOOL_ENV_GATES = {
|
|
"export_terminal_log": lambda: bool(os.environ.get("TMUX")),
|
|
"plexer_send": lambda: PlexerControllerFactory.detect() is not None,
|
|
"plexer_read": lambda: PlexerControllerFactory.detect() is not None,
|
|
"plexer_type": lambda: PlexerControllerFactory.detect() is not None,
|
|
"plexer_status": lambda: PlexerControllerFactory.detect() is not None,
|
|
}
|