forked from Rockachopa/Timmy-time-dashboard
This commit is contained in:
564
src/timmy/tools/_registry.py
Normal file
564
src/timmy/tools/_registry.py
Normal file
@@ -0,0 +1,564 @@
|
||||
"""Tool registry, full toolkit construction, and tool catalog.
|
||||
|
||||
Provides:
|
||||
- Internal _register_* helpers for wiring tools into toolkits
|
||||
- create_full_toolkit (orchestrator toolkit)
|
||||
- create_experiment_tools (Lab agent toolkit)
|
||||
- AGENT_TOOLKITS / get_tools_for_agent registry
|
||||
- get_all_available_tools catalog
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
from collections.abc import Callable
|
||||
from pathlib import Path
|
||||
|
||||
from timmy.tools._base import (
|
||||
_AGNO_TOOLS_AVAILABLE,
|
||||
_ImportError,
|
||||
FileTools,
|
||||
PythonTools,
|
||||
ShellTools,
|
||||
Toolkit,
|
||||
)
|
||||
from timmy.tools.file_tools import (
|
||||
_make_smart_read_file,
|
||||
create_data_tools,
|
||||
create_research_tools,
|
||||
create_writing_tools,
|
||||
)
|
||||
from timmy.tools.system_tools import (
|
||||
calculator,
|
||||
consult_grok,
|
||||
create_code_tools,
|
||||
create_devops_tools,
|
||||
create_security_tools,
|
||||
web_fetch,
|
||||
)
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Internal _register_* helpers
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _register_web_fetch_tool(toolkit: Toolkit) -> None:
|
||||
"""Register the web_fetch tool for full-page content extraction."""
|
||||
try:
|
||||
toolkit.register(web_fetch, name="web_fetch")
|
||||
except Exception as exc:
|
||||
logger.error("Failed to register web_fetch tool: %s", exc)
|
||||
raise
|
||||
|
||||
|
||||
def _register_core_tools(toolkit: Toolkit, base_path: Path) -> None:
|
||||
"""Register core execution and file tools."""
|
||||
# Python execution
|
||||
python_tools = PythonTools()
|
||||
toolkit.register(python_tools.run_python_code, name="python")
|
||||
|
||||
# Shell commands
|
||||
shell_tools = ShellTools()
|
||||
toolkit.register(shell_tools.run_shell_command, name="shell")
|
||||
|
||||
# File operations
|
||||
file_tools = FileTools(base_dir=base_path)
|
||||
toolkit.register(_make_smart_read_file(file_tools), name="read_file")
|
||||
toolkit.register(file_tools.save_file, name="write_file")
|
||||
toolkit.register(file_tools.list_files, name="list_files")
|
||||
|
||||
# Calculator — exact arithmetic (never let the LLM guess)
|
||||
toolkit.register(calculator, name="calculator")
|
||||
|
||||
|
||||
def _register_grok_tool(toolkit: Toolkit) -> None:
|
||||
"""Register Grok consultation tool if available."""
|
||||
try:
|
||||
from timmy.backends import grok_available
|
||||
|
||||
if grok_available():
|
||||
toolkit.register(consult_grok, name="consult_grok")
|
||||
logger.info("Grok consultation tool registered")
|
||||
except (ImportError, AttributeError) as exc:
|
||||
logger.error("Failed to register Grok tool: %s", exc)
|
||||
raise
|
||||
|
||||
|
||||
def _register_memory_tools(toolkit: Toolkit) -> None:
|
||||
"""Register memory search, write, and forget tools."""
|
||||
try:
|
||||
from timmy.memory_system import memory_forget, memory_read, memory_search, memory_write
|
||||
|
||||
toolkit.register(memory_search, name="memory_search")
|
||||
toolkit.register(memory_write, name="memory_write")
|
||||
toolkit.register(memory_read, name="memory_read")
|
||||
toolkit.register(memory_forget, name="memory_forget")
|
||||
except (ImportError, AttributeError) as exc:
|
||||
logger.error("Failed to register Memory tools: %s", exc)
|
||||
raise
|
||||
|
||||
|
||||
def _register_agentic_loop_tool(toolkit: Toolkit) -> None:
|
||||
"""Register agentic loop tool for background multi-step task execution."""
|
||||
try:
|
||||
from timmy.agentic_loop import run_agentic_loop
|
||||
|
||||
def plan_and_execute(task: str) -> str:
|
||||
"""Execute a complex multi-step task in the background with progress tracking.
|
||||
|
||||
Use this when a task requires 3 or more sequential tool calls that may
|
||||
take significant time. The task will run in the background and stream
|
||||
progress updates to the user via WebSocket.
|
||||
|
||||
Args:
|
||||
task: Full description of the multi-step task to execute.
|
||||
|
||||
Returns:
|
||||
Task ID and confirmation that background execution has started.
|
||||
"""
|
||||
import asyncio
|
||||
|
||||
task_id = None
|
||||
|
||||
async def _launch():
|
||||
nonlocal task_id
|
||||
result = await run_agentic_loop(task)
|
||||
return result
|
||||
|
||||
# Spawn as a background task on the running event loop
|
||||
try:
|
||||
asyncio.get_running_loop()
|
||||
future = asyncio.ensure_future(_launch())
|
||||
task_id = id(future)
|
||||
logger.info("Agentic loop started (task=%s)", task[:80])
|
||||
except RuntimeError:
|
||||
# No running loop — run synchronously (shouldn't happen in prod)
|
||||
result = asyncio.run(_launch())
|
||||
return f"Task completed: {result.summary}"
|
||||
|
||||
return (
|
||||
"Background task started. I'll execute this step-by-step "
|
||||
"and stream progress updates. You can monitor via the dashboard."
|
||||
)
|
||||
|
||||
toolkit.register(plan_and_execute, name="plan_and_execute")
|
||||
except (ImportError, AttributeError) as exc:
|
||||
logger.error("Failed to register plan_and_execute tool: %s", exc)
|
||||
raise
|
||||
|
||||
|
||||
def _register_introspection_tools(toolkit: Toolkit) -> None:
|
||||
"""Register system introspection tools for runtime environment queries."""
|
||||
try:
|
||||
from timmy.tools_intro import (
|
||||
check_ollama_health,
|
||||
get_memory_status,
|
||||
get_system_info,
|
||||
run_self_tests,
|
||||
)
|
||||
|
||||
toolkit.register(get_system_info, name="get_system_info")
|
||||
toolkit.register(check_ollama_health, name="check_ollama_health")
|
||||
toolkit.register(get_memory_status, name="get_memory_status")
|
||||
toolkit.register(run_self_tests, name="run_self_tests")
|
||||
except (ImportError, AttributeError) as exc:
|
||||
logger.error("Failed to register Introspection tools: %s", exc)
|
||||
raise
|
||||
|
||||
try:
|
||||
from timmy.mcp_tools import update_gitea_avatar
|
||||
|
||||
toolkit.register(update_gitea_avatar, name="update_gitea_avatar")
|
||||
except (ImportError, AttributeError) as exc:
|
||||
logger.error("Failed to register update_gitea_avatar tool: %s", exc)
|
||||
raise
|
||||
|
||||
try:
|
||||
from timmy.session_logger import self_reflect, session_history
|
||||
|
||||
toolkit.register(session_history, name="session_history")
|
||||
toolkit.register(self_reflect, name="self_reflect")
|
||||
except (ImportError, AttributeError) as exc:
|
||||
logger.error("Failed to register session_history tool: %s", exc)
|
||||
raise
|
||||
|
||||
|
||||
def _register_delegation_tools(toolkit: Toolkit) -> None:
|
||||
"""Register inter-agent delegation tools."""
|
||||
try:
|
||||
from timmy.tools_delegation import delegate_task, delegate_to_kimi, list_swarm_agents
|
||||
|
||||
toolkit.register(delegate_task, name="delegate_task")
|
||||
toolkit.register(delegate_to_kimi, name="delegate_to_kimi")
|
||||
toolkit.register(list_swarm_agents, name="list_swarm_agents")
|
||||
except Exception as exc:
|
||||
logger.error("Failed to register Delegation tools: %s", exc)
|
||||
raise
|
||||
|
||||
|
||||
def _register_gematria_tool(toolkit: Toolkit) -> None:
|
||||
"""Register the gematria computation tool."""
|
||||
try:
|
||||
from timmy.gematria import gematria
|
||||
|
||||
toolkit.register(gematria, name="gematria")
|
||||
except (ImportError, AttributeError) as exc:
|
||||
logger.error("Failed to register Gematria tool: %s", exc)
|
||||
raise
|
||||
|
||||
|
||||
def _register_artifact_tools(toolkit: Toolkit) -> None:
|
||||
"""Register artifact tools — notes and decision logging."""
|
||||
try:
|
||||
from timmy.memory_system import jot_note, log_decision
|
||||
|
||||
toolkit.register(jot_note, name="jot_note")
|
||||
toolkit.register(log_decision, name="log_decision")
|
||||
except (ImportError, AttributeError) as exc:
|
||||
logger.error("Failed to register Artifact tools: %s", exc)
|
||||
raise
|
||||
|
||||
|
||||
def _register_thinking_tools(toolkit: Toolkit) -> None:
|
||||
"""Register thinking/introspection tools for self-reflection."""
|
||||
try:
|
||||
from timmy.thinking import search_thoughts
|
||||
|
||||
toolkit.register(search_thoughts, name="thought_search")
|
||||
except (ImportError, AttributeError) as exc:
|
||||
logger.error("Failed to register Thinking tools: %s", exc)
|
||||
raise
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Full toolkit factories
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def create_full_toolkit(base_dir: str | Path | None = None):
|
||||
"""Create a full toolkit with all available tools (for the orchestrator).
|
||||
|
||||
Includes: web search, file read/write, shell commands, python execution,
|
||||
memory search for contextual recall, and Grok consultation.
|
||||
"""
|
||||
if not _AGNO_TOOLS_AVAILABLE:
|
||||
# Return None when tools aren't available (tests)
|
||||
return None
|
||||
|
||||
from config import settings
|
||||
from timmy.tool_safety import DANGEROUS_TOOLS
|
||||
|
||||
toolkit = Toolkit(name="full")
|
||||
# Set requires_confirmation_tools AFTER construction (avoids agno WARNING
|
||||
# about tools not yet registered) but BEFORE register() calls (so each
|
||||
# Function gets requires_confirmation=True). Fixes #79.
|
||||
toolkit.requires_confirmation_tools = list(DANGEROUS_TOOLS)
|
||||
|
||||
base_path = Path(base_dir) if base_dir else Path(settings.repo_root)
|
||||
|
||||
_register_core_tools(toolkit, base_path)
|
||||
_register_web_fetch_tool(toolkit)
|
||||
_register_grok_tool(toolkit)
|
||||
_register_memory_tools(toolkit)
|
||||
_register_agentic_loop_tool(toolkit)
|
||||
_register_introspection_tools(toolkit)
|
||||
_register_delegation_tools(toolkit)
|
||||
_register_gematria_tool(toolkit)
|
||||
_register_artifact_tools(toolkit)
|
||||
_register_thinking_tools(toolkit)
|
||||
|
||||
# Gitea issue management is now provided by the gitea-mcp server
|
||||
# (wired in as MCPTools in agent.py, not registered here)
|
||||
|
||||
return toolkit
|
||||
|
||||
|
||||
def create_experiment_tools(base_dir: str | Path | None = None):
|
||||
"""Create tools for the experiment agent (Lab).
|
||||
|
||||
Includes: prepare_experiment, run_experiment, evaluate_result,
|
||||
plus shell + file ops for editing training code.
|
||||
"""
|
||||
if not _AGNO_TOOLS_AVAILABLE:
|
||||
raise ImportError(f"Agno tools not available: {_ImportError}")
|
||||
|
||||
from config import settings
|
||||
|
||||
toolkit = Toolkit(name="experiment")
|
||||
|
||||
from timmy.autoresearch import evaluate_result, prepare_experiment, run_experiment
|
||||
|
||||
workspace = (
|
||||
Path(base_dir) if base_dir else Path(settings.repo_root) / settings.autoresearch_workspace
|
||||
)
|
||||
|
||||
def _prepare(repo_url: str = "https://github.com/karpathy/autoresearch.git") -> str:
|
||||
"""Clone and prepare an autoresearch experiment workspace."""
|
||||
return prepare_experiment(workspace, repo_url)
|
||||
|
||||
def _run(timeout: int = 0) -> str:
|
||||
"""Run a single training experiment with wall-clock timeout."""
|
||||
t = timeout or settings.autoresearch_time_budget
|
||||
result = run_experiment(workspace, timeout=t, metric_name=settings.autoresearch_metric)
|
||||
if result["success"] and result["metric"] is not None:
|
||||
return (
|
||||
f"{settings.autoresearch_metric}: {result['metric']:.4f} ({result['duration_s']}s)"
|
||||
)
|
||||
return result.get("error") or "Experiment failed"
|
||||
|
||||
def _evaluate(current: float, baseline: float) -> str:
|
||||
"""Compare current metric against baseline."""
|
||||
return evaluate_result(current, baseline, metric_name=settings.autoresearch_metric)
|
||||
|
||||
toolkit.register(_prepare, name="prepare_experiment")
|
||||
toolkit.register(_run, name="run_experiment")
|
||||
toolkit.register(_evaluate, name="evaluate_result")
|
||||
|
||||
# Also give Lab access to file + shell tools for editing train.py
|
||||
shell_tools = ShellTools()
|
||||
toolkit.register(shell_tools.run_shell_command, name="shell")
|
||||
|
||||
base_path = Path(base_dir) if base_dir else Path(settings.repo_root)
|
||||
file_tools = FileTools(base_dir=base_path)
|
||||
toolkit.register(_make_smart_read_file(file_tools), name="read_file")
|
||||
toolkit.register(file_tools.save_file, name="write_file")
|
||||
toolkit.register(file_tools.list_files, name="list_files")
|
||||
|
||||
return toolkit
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Agent toolkit registry
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _create_stub_toolkit(name: str):
|
||||
"""Create a minimal Agno toolkit for creative agents.
|
||||
|
||||
Creative agents use their own dedicated tool modules rather than
|
||||
Agno-wrapped functions. This stub ensures AGENT_TOOLKITS has an
|
||||
entry so ToolExecutor doesn't fall back to the full toolkit.
|
||||
"""
|
||||
if not _AGNO_TOOLS_AVAILABLE:
|
||||
return None
|
||||
toolkit = Toolkit(name=name)
|
||||
return toolkit
|
||||
|
||||
|
||||
# Mapping of agent IDs to their toolkits
|
||||
AGENT_TOOLKITS: dict[str, Callable[[], Toolkit]] = {
|
||||
"echo": create_research_tools,
|
||||
"mace": create_security_tools,
|
||||
"helm": create_devops_tools,
|
||||
"seer": create_data_tools,
|
||||
"forge": create_code_tools,
|
||||
"quill": create_writing_tools,
|
||||
"lab": create_experiment_tools,
|
||||
"pixel": lambda base_dir=None: _create_stub_toolkit("pixel"),
|
||||
"lyra": lambda base_dir=None: _create_stub_toolkit("lyra"),
|
||||
"reel": lambda base_dir=None: _create_stub_toolkit("reel"),
|
||||
}
|
||||
|
||||
|
||||
def get_tools_for_agent(agent_id: str, base_dir: str | Path | None = None) -> "Toolkit | None":
|
||||
"""Get the appropriate toolkit for an agent.
|
||||
|
||||
Args:
|
||||
agent_id: The agent ID (echo, mace, helm, seer, forge, quill)
|
||||
base_dir: Optional base directory for file operations
|
||||
|
||||
Returns:
|
||||
A Toolkit instance or None if agent_id is not recognized
|
||||
"""
|
||||
factory = AGENT_TOOLKITS.get(agent_id)
|
||||
if factory:
|
||||
return factory(base_dir)
|
||||
return None
|
||||
|
||||
|
||||
# Backward-compat aliases
|
||||
get_tools_for_persona = get_tools_for_agent
|
||||
PERSONA_TOOLKITS = AGENT_TOOLKITS
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Tool catalog
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _core_tool_catalog() -> dict:
|
||||
"""Return core file and execution tools catalog entries."""
|
||||
return {
|
||||
"shell": {
|
||||
"name": "Shell Commands",
|
||||
"description": "Execute shell commands (sandboxed)",
|
||||
"available_in": ["forge", "mace", "helm", "orchestrator"],
|
||||
},
|
||||
"python": {
|
||||
"name": "Python Execution",
|
||||
"description": "Execute Python code for analysis and scripting",
|
||||
"available_in": ["forge", "seer", "orchestrator"],
|
||||
},
|
||||
"read_file": {
|
||||
"name": "Read File",
|
||||
"description": "Read contents of local files",
|
||||
"available_in": ["echo", "seer", "forge", "quill", "mace", "helm", "orchestrator"],
|
||||
},
|
||||
"write_file": {
|
||||
"name": "Write File",
|
||||
"description": "Write content to local files",
|
||||
"available_in": ["forge", "quill", "helm", "orchestrator"],
|
||||
},
|
||||
"list_files": {
|
||||
"name": "List Files",
|
||||
"description": "List files in a directory",
|
||||
"available_in": ["echo", "seer", "forge", "quill", "mace", "helm", "orchestrator"],
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
def _analysis_tool_catalog() -> dict:
|
||||
"""Return analysis and calculation tools catalog entries."""
|
||||
return {
|
||||
"calculator": {
|
||||
"name": "Calculator",
|
||||
"description": "Evaluate mathematical expressions with exact results",
|
||||
"available_in": ["orchestrator"],
|
||||
},
|
||||
"web_fetch": {
|
||||
"name": "Web Fetch",
|
||||
"description": "Fetch a web page and extract clean readable text (trafilatura)",
|
||||
"available_in": ["orchestrator"],
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
def _ai_tool_catalog() -> dict:
|
||||
"""Return AI assistant and frontier reasoning tools catalog entries."""
|
||||
return {
|
||||
"consult_grok": {
|
||||
"name": "Consult Grok",
|
||||
"description": "Premium frontier reasoning via xAI Grok (opt-in, Lightning-payable)",
|
||||
"available_in": ["orchestrator"],
|
||||
},
|
||||
"aider": {
|
||||
"name": "Aider AI Assistant",
|
||||
"description": "Local AI coding assistant using Ollama (qwen3:30b or deepseek-coder)",
|
||||
"available_in": ["forge", "orchestrator"],
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
def _introspection_tool_catalog() -> dict:
|
||||
"""Return system introspection tools catalog entries."""
|
||||
return {
|
||||
"get_system_info": {
|
||||
"name": "System Info",
|
||||
"description": "Introspect runtime environment - discover model, Python version, config",
|
||||
"available_in": ["orchestrator"],
|
||||
},
|
||||
"check_ollama_health": {
|
||||
"name": "Ollama Health",
|
||||
"description": "Check if Ollama is accessible and what models are available",
|
||||
"available_in": ["orchestrator"],
|
||||
},
|
||||
"get_memory_status": {
|
||||
"name": "Memory Status",
|
||||
"description": "Check status of memory tiers (hot memory, vault)",
|
||||
"available_in": ["orchestrator"],
|
||||
},
|
||||
"session_history": {
|
||||
"name": "Session History",
|
||||
"description": "Search past conversation logs for messages, tool calls, errors, and decisions",
|
||||
"available_in": ["orchestrator"],
|
||||
},
|
||||
"thought_search": {
|
||||
"name": "Thought Search",
|
||||
"description": "Query Timmy's own thought history for past reflections and insights",
|
||||
"available_in": ["orchestrator"],
|
||||
},
|
||||
"self_reflect": {
|
||||
"name": "Self-Reflect",
|
||||
"description": "Review recent conversations to spot patterns, low-confidence answers, and errors",
|
||||
"available_in": ["orchestrator"],
|
||||
},
|
||||
"update_gitea_avatar": {
|
||||
"name": "Update Gitea Avatar",
|
||||
"description": "Generate and upload a wizard-themed avatar to Timmy's Gitea profile",
|
||||
"available_in": ["orchestrator"],
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
def _experiment_tool_catalog() -> dict:
|
||||
"""Return ML experiment tools catalog entries."""
|
||||
return {
|
||||
"prepare_experiment": {
|
||||
"name": "Prepare Experiment",
|
||||
"description": "Clone autoresearch repo and run data preparation for ML experiments",
|
||||
"available_in": ["lab", "orchestrator"],
|
||||
},
|
||||
"run_experiment": {
|
||||
"name": "Run Experiment",
|
||||
"description": "Execute a time-boxed ML training experiment and capture metrics",
|
||||
"available_in": ["lab", "orchestrator"],
|
||||
},
|
||||
"evaluate_result": {
|
||||
"name": "Evaluate Result",
|
||||
"description": "Compare experiment metric against baseline to assess improvement",
|
||||
"available_in": ["lab", "orchestrator"],
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
_CREATIVE_CATALOG_SOURCES: list[tuple[str, str, list[str]]] = [
|
||||
("creative.tools.git_tools", "GIT_TOOL_CATALOG", ["forge", "helm", "orchestrator"]),
|
||||
("creative.tools.image_tools", "IMAGE_TOOL_CATALOG", ["pixel", "orchestrator"]),
|
||||
("creative.tools.music_tools", "MUSIC_TOOL_CATALOG", ["lyra", "orchestrator"]),
|
||||
("creative.tools.video_tools", "VIDEO_TOOL_CATALOG", ["reel", "orchestrator"]),
|
||||
("creative.director", "DIRECTOR_TOOL_CATALOG", ["orchestrator"]),
|
||||
("creative.assembler", "ASSEMBLER_TOOL_CATALOG", ["reel", "orchestrator"]),
|
||||
]
|
||||
|
||||
|
||||
def _import_creative_catalogs(catalog: dict) -> None:
|
||||
"""Import and merge creative tool catalogs from creative module."""
|
||||
for module_path, attr_name, available_in in _CREATIVE_CATALOG_SOURCES:
|
||||
_merge_catalog(catalog, module_path, attr_name, available_in)
|
||||
|
||||
|
||||
def _merge_catalog(
|
||||
catalog: dict, module_path: str, attr_name: str, available_in: list[str]
|
||||
) -> None:
|
||||
"""Import a single creative catalog and merge its entries."""
|
||||
try:
|
||||
from importlib import import_module
|
||||
|
||||
source_catalog = getattr(import_module(module_path), attr_name)
|
||||
for tool_id, info in source_catalog.items():
|
||||
catalog[tool_id] = {
|
||||
"name": info["name"],
|
||||
"description": info["description"],
|
||||
"available_in": available_in,
|
||||
}
|
||||
except ImportError:
|
||||
logger.debug("Optional catalog %s.%s not available", module_path, attr_name)
|
||||
|
||||
|
||||
def get_all_available_tools() -> dict[str, dict]:
|
||||
"""Get a catalog of all available tools and their descriptions.
|
||||
|
||||
Returns:
|
||||
Dict mapping tool categories to their tools and descriptions.
|
||||
"""
|
||||
catalog = {}
|
||||
catalog.update(_core_tool_catalog())
|
||||
catalog.update(_analysis_tool_catalog())
|
||||
catalog.update(_ai_tool_catalog())
|
||||
catalog.update(_introspection_tool_catalog())
|
||||
catalog.update(_experiment_tool_catalog())
|
||||
_import_creative_catalogs(catalog)
|
||||
return catalog
|
||||
Reference in New Issue
Block a user