656 lines
23 KiB
Python
656 lines
23 KiB
Python
"""
|
||
tools.py — AIPA Tool Registry
|
||
Provides a lightweight @tool() decorator, schema export, and tool execution
|
||
for agent tool-calling loops.
|
||
|
||
Agent-management tools allow Evelyn (and any tool-enabled agent) to read and
|
||
modify agent definitions without restarting the session.
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import inspect
|
||
import json
|
||
from pathlib import Path
|
||
from typing import Callable, Any
|
||
|
||
import config as _config
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Registry
|
||
# ---------------------------------------------------------------------------
|
||
|
||
_REGISTRY: dict[str, dict] = {} # name → {fn, schema}
|
||
|
||
|
||
def tool(description: str, parameters: dict):
|
||
"""
|
||
Decorator that registers a callable as a named tool.
|
||
|
||
Parameters
|
||
----------
|
||
description : str
|
||
Plain-language description shown to the LLM.
|
||
parameters : dict
|
||
JSON Schema object for the tool's parameters (OpenAI function-calling
|
||
format — {"type": "object", "properties": {...}, "required": [...]}).
|
||
"""
|
||
def decorator(fn: Callable) -> Callable:
|
||
name = fn.__name__
|
||
_REGISTRY[name] = {
|
||
"fn": fn,
|
||
"schema": {
|
||
"type": "function",
|
||
"function": {
|
||
"name": name,
|
||
"description": description,
|
||
"parameters": parameters,
|
||
},
|
||
},
|
||
}
|
||
return fn
|
||
return decorator
|
||
|
||
|
||
def get_schemas(names: list[str]) -> list[dict]:
|
||
"""Return OpenAI-compatible tool schemas for the named tools."""
|
||
result = []
|
||
for name in names:
|
||
entry = _REGISTRY.get(name)
|
||
if entry is None:
|
||
raise ValueError(
|
||
f"Tool {name!r} is not registered. "
|
||
f"Available tools: {list(_REGISTRY)}"
|
||
)
|
||
result.append(entry["schema"])
|
||
return result
|
||
|
||
|
||
def call(name: str, arguments_json: str) -> str:
|
||
"""
|
||
Execute a registered tool by name with a JSON-encoded arguments string.
|
||
Returns a JSON-encoded result string (always a string for message history).
|
||
"""
|
||
entry = _REGISTRY.get(name)
|
||
if entry is None:
|
||
return json.dumps({"error": f"Unknown tool: {name!r}"})
|
||
try:
|
||
kwargs = json.loads(arguments_json) if arguments_json else {}
|
||
result = entry["fn"](**kwargs)
|
||
return json.dumps(result, ensure_ascii=False)
|
||
except Exception as exc:
|
||
return json.dumps({"error": str(exc)})
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# YAML helpers (ruamel for comment-preserving round-trips)
|
||
# ---------------------------------------------------------------------------
|
||
|
||
def _load_agents_yaml() -> tuple[Any, Any]:
|
||
"""
|
||
Load agents.yaml using ruamel.yaml for comment preservation.
|
||
Returns (ruamel_instance, data_dict).
|
||
"""
|
||
from ruamel.yaml import YAML
|
||
ryaml = YAML()
|
||
ryaml.preserve_quotes = True
|
||
path = _config._AGENTS_CONFIG_PATH
|
||
with open(path, encoding="utf-8") as fh:
|
||
data = ryaml.load(fh)
|
||
return ryaml, data
|
||
|
||
|
||
def _save_agents_yaml(ryaml: Any, data: Any) -> None:
|
||
"""Write agents.yaml back with comment preservation."""
|
||
path = _config._AGENTS_CONFIG_PATH
|
||
with open(path, "w", encoding="utf-8") as fh:
|
||
ryaml.dump(data, fh)
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Agent-management tools
|
||
# ---------------------------------------------------------------------------
|
||
|
||
@tool(
|
||
description="List all agents currently defined in agents.yaml, including their title, provider, model, temperature, max_tokens, stateful flag, and tools list.",
|
||
parameters={
|
||
"type": "object",
|
||
"properties": {},
|
||
"required": [],
|
||
},
|
||
)
|
||
def list_agents() -> list[dict]:
|
||
"""Return all agent configs as a list of dicts."""
|
||
result = []
|
||
for name, cfg in _config.AGENT_CONFIGS.items():
|
||
result.append({
|
||
"name": name,
|
||
"title": _config.agent_title(name),
|
||
"color": cfg.get("color", ""),
|
||
"prompt_file": cfg.get("prompt_file", ""),
|
||
"provider": _config.agent_provider_name(name),
|
||
"provider_type": _config.agent_provider_type(name),
|
||
"model": _config.agent_model(name),
|
||
"temperature": cfg.get("temperature"),
|
||
"max_tokens": cfg.get("max_tokens"),
|
||
"stateful": cfg.get("stateful"),
|
||
"tools": cfg.get("tools", []),
|
||
})
|
||
return result
|
||
|
||
|
||
@tool(
|
||
description="Read the full configuration for a specific agent by name, including provider type, thinking mode, model, and tools.",
|
||
parameters={
|
||
"type": "object",
|
||
"properties": {
|
||
"agent_name": {
|
||
"type": "string",
|
||
"description": "The agent's short name (e.g. 'evelyn', 'atlas').",
|
||
}
|
||
},
|
||
"required": ["agent_name"],
|
||
},
|
||
)
|
||
def read_agent_config(agent_name: str) -> dict:
|
||
"""Return the full resolved config for one agent."""
|
||
name = agent_name.lower()
|
||
if name not in _config.AGENT_CONFIGS:
|
||
return {"error": f"No agent named {agent_name!r}. Known: {list(_config.AGENT_CONFIGS)}"}
|
||
cfg = _config.AGENT_CONFIGS[name]
|
||
provider_name = _config.agent_provider_name(name)
|
||
provider_def = _config._providers.get(provider_name, {})
|
||
return {
|
||
"name": name,
|
||
"title": _config.agent_title(name),
|
||
"color": cfg.get("color", ""),
|
||
"prompt_file": cfg.get("prompt_file"),
|
||
"provider": provider_name,
|
||
"provider_type": _config.agent_provider_type(name),
|
||
"provider_base_url": provider_def.get("base_url", ""),
|
||
"provider_default_model": provider_def.get("default_model", ""),
|
||
"model": _config.agent_model(name),
|
||
"temperature": cfg.get("temperature"),
|
||
"max_tokens": cfg.get("max_tokens"),
|
||
"stateful": cfg.get("stateful"),
|
||
"tools": cfg.get("tools", []),
|
||
}
|
||
|
||
|
||
@tool(
|
||
description=(
|
||
"Create or update an agent definition in agents.yaml. "
|
||
"Pass only the fields you want to set or change — omitted fields are left unchanged on update, "
|
||
"or get sensible defaults on create. "
|
||
"After calling this tool, the session must be reloaded (/reload) for changes to take effect."
|
||
),
|
||
parameters={
|
||
"type": "object",
|
||
"properties": {
|
||
"agent_name": {
|
||
"type": "string",
|
||
"description": "Short name for the agent (e.g. 'evelyn'). Used as the YAML key.",
|
||
},
|
||
"title": {
|
||
"type": "string",
|
||
"description": "Human-readable title (e.g. 'Director of Personnel & Systems').",
|
||
},
|
||
"color": {
|
||
"type": "string",
|
||
"description": "Rich terminal color name (e.g. 'white', 'cyan', 'green').",
|
||
},
|
||
"prompt_file": {
|
||
"type": "string",
|
||
"description": "Filename (not path) in agents/prompts/ (e.g. 'evelyn_director_of_personnel.md').",
|
||
},
|
||
"provider": {
|
||
"type": "string",
|
||
"description": "Provider name from agents.yaml providers section (e.g. 'vastblueai_thinking').",
|
||
},
|
||
"model": {
|
||
"type": "string",
|
||
"description": "Model name as the provider expects it. Empty string means use provider default.",
|
||
},
|
||
"temperature": {
|
||
"type": "number",
|
||
"description": "Sampling temperature 0.0–1.0. Use 0.6 for thinking, 0.7 for non-thinking.",
|
||
},
|
||
"max_tokens": {
|
||
"type": "integer",
|
||
"description": "Maximum response tokens.",
|
||
},
|
||
"stateful": {
|
||
"type": "boolean",
|
||
"description": "Whether the agent maintains conversation history across calls.",
|
||
},
|
||
"tools": {
|
||
"type": "array",
|
||
"items": {"type": "string"},
|
||
"description": "List of tool names from the tools registry that this agent can use.",
|
||
},
|
||
},
|
||
"required": ["agent_name"],
|
||
},
|
||
)
|
||
def upsert_agent_definition(
|
||
agent_name: str,
|
||
title: str | None = None,
|
||
color: str | None = None,
|
||
prompt_file: str | None = None,
|
||
provider: str | None = None,
|
||
model: str | None = None,
|
||
temperature: float | None = None,
|
||
max_tokens: int | None = None,
|
||
stateful: bool | None = None,
|
||
tools: list[str] | None = None,
|
||
) -> dict:
|
||
"""Create or update an agent in agents.yaml."""
|
||
name = agent_name.lower()
|
||
|
||
# Validate provider if given
|
||
if provider is not None and provider != "none":
|
||
valid_providers = list(_config._providers.keys())
|
||
if provider not in valid_providers:
|
||
return {
|
||
"error": f"Unknown provider {provider!r}. Valid providers: {valid_providers}"
|
||
}
|
||
|
||
# Validate tools if given
|
||
if tools is not None:
|
||
unknown = [t for t in tools if t not in _REGISTRY]
|
||
if unknown:
|
||
return {
|
||
"error": f"Unknown tools: {unknown}. Registered tools: {list(_REGISTRY)}"
|
||
}
|
||
|
||
ryaml, data = _load_agents_yaml()
|
||
agents_section = data.get("agents", {})
|
||
|
||
creating = name not in agents_section
|
||
|
||
if creating:
|
||
# Supply required defaults for new agents
|
||
if not prompt_file:
|
||
return {"error": "prompt_file is required when creating a new agent."}
|
||
agents_section[name] = {
|
||
"title": title or name.capitalize(),
|
||
"color": color or "white",
|
||
"prompt_file": prompt_file,
|
||
"provider": provider or "",
|
||
"model": model if model is not None else "",
|
||
"temperature": temperature if temperature is not None else 0.7,
|
||
"max_tokens": max_tokens or 16384,
|
||
"stateful": stateful if stateful is not None else True,
|
||
}
|
||
if tools:
|
||
agents_section[name]["tools"] = tools
|
||
else:
|
||
entry = agents_section[name]
|
||
if title is not None: entry["title"] = title
|
||
if color is not None: entry["color"] = color
|
||
if prompt_file is not None: entry["prompt_file"] = prompt_file
|
||
if provider is not None: entry["provider"] = provider
|
||
if model is not None: entry["model"] = model
|
||
if temperature is not None: entry["temperature"] = temperature
|
||
if max_tokens is not None: entry["max_tokens"] = max_tokens
|
||
if stateful is not None: entry["stateful"] = stateful
|
||
if tools is not None: entry["tools"] = tools
|
||
|
||
_save_agents_yaml(ryaml, data)
|
||
action = "created" if creating else "updated"
|
||
return {"ok": True, "action": action, "agent": name, "note": "Run /reload to apply changes."}
|
||
|
||
|
||
@tool(
|
||
description="Read the system prompt file for an agent. Returns the raw Markdown content of the prompt file.",
|
||
parameters={
|
||
"type": "object",
|
||
"properties": {
|
||
"agent_name": {
|
||
"type": "string",
|
||
"description": "The agent's short name (e.g. 'evelyn').",
|
||
}
|
||
},
|
||
"required": ["agent_name"],
|
||
},
|
||
)
|
||
def read_agent_prompt(agent_name: str) -> dict:
|
||
"""Read and return the raw prompt file for an agent."""
|
||
name = agent_name.lower()
|
||
cfg = _config.AGENT_CONFIGS.get(name)
|
||
if cfg is None:
|
||
return {"error": f"No agent named {agent_name!r}."}
|
||
path = _config.PROMPTS_DIR / cfg["prompt_file"]
|
||
if not path.exists():
|
||
return {"error": f"Prompt file not found: {path}"}
|
||
return {"prompt_file": cfg["prompt_file"], "content": path.read_text(encoding="utf-8")}
|
||
|
||
|
||
@tool(
|
||
description=(
|
||
"Write or overwrite a prompt file in the agents/prompts/ directory. "
|
||
"Use the prompt_file name exactly (e.g. 'evelyn_director_of_personnel.md'). "
|
||
"The file will be created if it does not exist."
|
||
),
|
||
parameters={
|
||
"type": "object",
|
||
"properties": {
|
||
"prompt_file": {
|
||
"type": "string",
|
||
"description": "Filename (not path) in agents/prompts/. Must end with .md.",
|
||
},
|
||
"content": {
|
||
"type": "string",
|
||
"description": "Full Markdown content to write to the file.",
|
||
},
|
||
},
|
||
"required": ["prompt_file", "content"],
|
||
},
|
||
)
|
||
def write_agent_prompt(prompt_file: str, content: str) -> dict:
|
||
"""Write (or overwrite) a prompt file."""
|
||
if not prompt_file.endswith(".md"):
|
||
return {"error": "prompt_file must end with .md"}
|
||
# Disallow path traversal
|
||
if "/" in prompt_file or "\\" in prompt_file:
|
||
return {"error": "prompt_file must be a filename only, not a path."}
|
||
path = _config.PROMPTS_DIR / prompt_file
|
||
path.write_text(content, encoding="utf-8")
|
||
return {"ok": True, "wrote": str(path), "bytes": len(content.encode())}
|
||
|
||
|
||
@tool(
|
||
description="List all provider definitions from agents.yaml, showing type, base_url, default_model, and extra_body settings.",
|
||
parameters={
|
||
"type": "object",
|
||
"properties": {},
|
||
"required": [],
|
||
},
|
||
)
|
||
def list_providers() -> list[dict]:
|
||
"""Return all provider definitions."""
|
||
result = []
|
||
for name, pdef in _config._providers.items():
|
||
entry = {"name": name}
|
||
entry["type"] = pdef.get("type", "")
|
||
entry["base_url"] = pdef.get("base_url", "")
|
||
entry["default_model"] = pdef.get("default_model", "")
|
||
entry["extra_body"] = pdef.get("extra_body", {})
|
||
result.append(entry)
|
||
return result
|
||
|
||
|
||
@tool(
|
||
description="List all tools registered in the tool registry, with their names and descriptions.",
|
||
parameters={
|
||
"type": "object",
|
||
"properties": {},
|
||
"required": [],
|
||
},
|
||
)
|
||
def list_tools() -> list[dict]:
|
||
"""Return all registered tools with name and description."""
|
||
return [
|
||
{
|
||
"name": name,
|
||
"description": entry["schema"]["function"]["description"],
|
||
}
|
||
for name, entry in _REGISTRY.items()
|
||
]
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Task dispatch (used by Miranda — orchestrator drains this queue after her turn)
|
||
# ---------------------------------------------------------------------------
|
||
|
||
_pending_tasks: list[dict] = []
|
||
|
||
|
||
def get_pending_tasks() -> list[dict]:
|
||
"""
|
||
Drain and return all task specs queued by dispatch_task tool calls.
|
||
Call once per directive, after the orchestrator's tool loop completes.
|
||
"""
|
||
tasks = _pending_tasks.copy()
|
||
_pending_tasks.clear()
|
||
return tasks
|
||
|
||
|
||
@tool(
|
||
description=(
|
||
"Dispatch a task to a Lead agent. "
|
||
"Call this tool once per Lead you want to task — multiple calls dispatch multiple tasks in parallel. "
|
||
"Do NOT respond with TASK BRIEF text; use this tool instead. "
|
||
"Results will be returned to you when all Leads complete so you can synthesize."
|
||
),
|
||
parameters={
|
||
"type": "object",
|
||
"properties": {
|
||
"to": {
|
||
"type": "string",
|
||
"description": "Lead agent name (lowercase): atlas, cole, clio, iris, or evelyn.",
|
||
},
|
||
"directive": {
|
||
"type": "string",
|
||
"description": "What the Principal wants — the full task objective.",
|
||
},
|
||
"scope": {
|
||
"type": "string",
|
||
"description": "Exactly what this Lead is responsible for delivering.",
|
||
},
|
||
"constraints": {
|
||
"type": "string",
|
||
"description": "Deadlines, depth limits, format requirements, or assumptions to hold. Use 'NONE' if none.",
|
||
},
|
||
"dependencies": {
|
||
"type": "string",
|
||
"description": "Other tasks this depends on, or 'NONE'.",
|
||
},
|
||
"return_format": {
|
||
"type": "string",
|
||
"description": "Specific structure or emphasis you need back from this Lead beyond the standard STATUS / SUMMARY / FINDINGS / OPEN ITEMS format.",
|
||
},
|
||
},
|
||
"required": ["to", "directive", "scope"],
|
||
},
|
||
)
|
||
def dispatch_task(
|
||
to: str,
|
||
directive: str,
|
||
scope: str,
|
||
constraints: str = "NONE",
|
||
dependencies: str = "NONE",
|
||
return_format: str = "",
|
||
) -> dict:
|
||
"""Queue a task for dispatch to a Lead agent after Miranda's turn completes."""
|
||
lead = to.lower().strip()
|
||
if lead not in _config.LEAD_AGENT_NAMES:
|
||
return {
|
||
"error": f"Unknown lead {to!r}. Valid leads: {sorted(_config.LEAD_AGENT_NAMES)}"
|
||
}
|
||
_pending_tasks.append({
|
||
"to": lead,
|
||
"directive": directive,
|
||
"scope": scope,
|
||
"constraints": constraints,
|
||
"dependencies": dependencies,
|
||
"return_format": return_format,
|
||
})
|
||
return {
|
||
"queued": True,
|
||
"to": lead,
|
||
"note": "Task queued. You will receive lead outputs to synthesize when all tasks complete.",
|
||
}
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Knowledge base tools
|
||
# ---------------------------------------------------------------------------
|
||
|
||
@tool(
|
||
description=(
|
||
"Write or overwrite a Markdown document in the knowledge base. "
|
||
"The document is stored as a .md file on the filesystem. "
|
||
"Existing changelogs are preserved and a new entry is appended automatically. "
|
||
"Use Markdown with a top-level '# Title' heading as the first line."
|
||
),
|
||
parameters={
|
||
"type": "object",
|
||
"properties": {
|
||
"filename": {
|
||
"type": "string",
|
||
"description": "Filename for the document (e.g. 'project_goals.md'). Must end with .md.",
|
||
},
|
||
"content": {
|
||
"type": "string",
|
||
"description": "Full Markdown content to write (excluding changelog — that is managed automatically).",
|
||
},
|
||
"collection": {
|
||
"type": "string",
|
||
"description": "Collection (subdirectory) to store the document in. Defaults to the configured default collection.",
|
||
},
|
||
"agent": {
|
||
"type": "string",
|
||
"description": "Name of the agent writing this document — recorded in the changelog.",
|
||
},
|
||
"changelog_entry": {
|
||
"type": "string",
|
||
"description": "Short description of what changed, for the changelog (e.g. 'Added Q2 goals'). Auto-generated if omitted.",
|
||
},
|
||
},
|
||
"required": ["filename", "content"],
|
||
},
|
||
)
|
||
def kb_write(
|
||
filename: str,
|
||
content: str,
|
||
collection: str = "",
|
||
agent: str = "",
|
||
changelog_entry: str = "",
|
||
) -> dict:
|
||
"""Write or overwrite a KB Markdown document."""
|
||
import knowledge_base as _kb
|
||
return _kb.kb_write(
|
||
filename=filename,
|
||
content=content,
|
||
collection=collection,
|
||
agent=agent,
|
||
changelog_entry=changelog_entry,
|
||
)
|
||
|
||
|
||
@tool(
|
||
description=(
|
||
"Read a Markdown document from the knowledge base. "
|
||
"Returns the document body and its changelog history separately."
|
||
),
|
||
parameters={
|
||
"type": "object",
|
||
"properties": {
|
||
"filename": {
|
||
"type": "string",
|
||
"description": "Filename of the document to read (e.g. 'project_goals.md').",
|
||
},
|
||
"collection": {
|
||
"type": "string",
|
||
"description": "Collection containing the document. Defaults to the configured default collection.",
|
||
},
|
||
},
|
||
"required": ["filename"],
|
||
},
|
||
)
|
||
def kb_read(filename: str, collection: str = "") -> dict:
|
||
"""Read a KB document, returning body and changelog separately."""
|
||
import knowledge_base as _kb
|
||
return _kb.kb_read(filename=filename, collection=collection)
|
||
|
||
|
||
@tool(
|
||
description=(
|
||
"Search the knowledge base by semantic similarity. "
|
||
"Syncs the embedding index from the filesystem before querying. "
|
||
"Returns matching documents ordered by relevance (lower distance = more similar)."
|
||
),
|
||
parameters={
|
||
"type": "object",
|
||
"properties": {
|
||
"query": {
|
||
"type": "string",
|
||
"description": "Natural-language search query.",
|
||
},
|
||
"n_results": {
|
||
"type": "integer",
|
||
"description": "Maximum number of results to return. Defaults to 5.",
|
||
},
|
||
"collection": {
|
||
"type": "string",
|
||
"description": "Collection to search. Defaults to the configured default collection.",
|
||
},
|
||
},
|
||
"required": ["query"],
|
||
},
|
||
)
|
||
def kb_search(query: str, n_results: int = 5, collection: str = "") -> list[dict]:
|
||
"""Semantic search in the knowledge base."""
|
||
import knowledge_base as _kb
|
||
return _kb.kb_search(query=query, n_results=n_results, collection=collection)
|
||
|
||
|
||
@tool(
|
||
description="List all Markdown documents in a knowledge base collection with their titles and modification times.",
|
||
parameters={
|
||
"type": "object",
|
||
"properties": {
|
||
"collection": {
|
||
"type": "string",
|
||
"description": "Collection to list. Defaults to the configured default collection.",
|
||
},
|
||
},
|
||
"required": [],
|
||
},
|
||
)
|
||
def kb_list(collection: str = "") -> list[dict]:
|
||
"""List all documents in a KB collection."""
|
||
import knowledge_base as _kb
|
||
return _kb.kb_list(collection=collection)
|
||
|
||
|
||
@tool(
|
||
description=(
|
||
"Delete a Markdown document from the knowledge base. "
|
||
"The file is removed from disk; the embedding index is cleaned up on next search."
|
||
),
|
||
parameters={
|
||
"type": "object",
|
||
"properties": {
|
||
"filename": {
|
||
"type": "string",
|
||
"description": "Filename of the document to delete (e.g. 'old_notes.md').",
|
||
},
|
||
"collection": {
|
||
"type": "string",
|
||
"description": "Collection containing the document. Defaults to the configured default collection.",
|
||
},
|
||
},
|
||
"required": ["filename"],
|
||
},
|
||
)
|
||
def kb_delete(filename: str, collection: str = "") -> dict:
|
||
"""Delete a KB document from the filesystem."""
|
||
import knowledge_base as _kb
|
||
return _kb.kb_delete(filename=filename, collection=collection)
|
||
|
||
|
||
@tool(
|
||
description="List all collections in the knowledge base with their document counts.",
|
||
parameters={
|
||
"type": "object",
|
||
"properties": {},
|
||
"required": [],
|
||
},
|
||
)
|
||
def kb_list_collections() -> list[dict]:
|
||
"""List all knowledge base collections."""
|
||
import knowledge_base as _kb
|
||
return _kb.kb_list_collections()
|