Cortex-Inara/cortex/tools/__init__.py

"""
Orchestrator tool registry.

Each tool has two parts:
  1. A Gemini FunctionDeclaration — tells the model what the tool does and what args it takes
  2. A Python async callable — the actual implementation

To add a new tool:
  1. Implement it in a tools/<domain>.py module
  2. Import it here and add (declaration, callable) to _REGISTRY
  3. Add a FunctionDeclaration below and include it in TOOL_DECLARATIONS

IMPORTANT: These tools are separate from the ae_* MCP tools used by the fleet agents.
           Do not modify the ae_* MCP server to support orchestrator needs.
"""

from google.genai import types
from tools.web import search as _web_search
from tools.ae_knowledge import journal_search as _ae_journal_search
from tools.ae_knowledge import journal_entry_create as _ae_journal_entry_create
from tools.ae_tasks import task_list as _ae_task_list
from tools.files import file_read as _file_read
from tools.system import claude_allow_dir as _claude_allow_dir
from tools.tasks import task_list as _task_list, task_create as _task_create
from tools.tasks import task_update as _task_update, task_complete as _task_complete
from tools.scratch import (
    scratch_read as _scratch_read,
    scratch_write as _scratch_write,
    scratch_append as _scratch_append,
    scratch_clear as _scratch_clear,
)


# ---------------------------------------------------------------------------
# Gemini function declarations
# ---------------------------------------------------------------------------

_web_search_declaration = types.FunctionDeclaration(
    name="web_search",
    description=(
        "Search the web for current information. Use this when you need up-to-date "
        "facts, news, documentation, or anything not in your training data."
    ),
    parameters=types.Schema(
        type=types.Type.OBJECT,
        properties={
            "query": types.Schema(
                type=types.Type.STRING,
                description="The search query string",
            ),
            "max_results": types.Schema(
                type=types.Type.INTEGER,
                description="Number of results to return (default 5, max 10)",
            ),
        },
        required=["query"],
    ),
)

_ae_journal_search_declaration = types.FunctionDeclaration(
    name="ae_journal_search",
    description=(
        "Search the Aether Journals knowledge base by keyword. "
        "Use this to look up notes, documentation, meeting summaries, or any saved knowledge. "
        "Always search before creating a new entry to avoid duplicates."
    ),
    parameters=types.Schema(
        type=types.Type.OBJECT,
        properties={
            "query": types.Schema(
                type=types.Type.STRING,
                description="Keyword or phrase to search for",
            ),
            "journal_id": types.Schema(
                type=types.Type.STRING,
                description=(
                    "Optional: scope search to a specific journal by its id_random. "
                    "Omit to search all journals."
                ),
            ),
            "max_results": types.Schema(
                type=types.Type.INTEGER,
                description="Maximum number of entries to return (default 10)",
            ),
        },
        required=["query"],
    ),
)

_ae_journal_entry_create_declaration = types.FunctionDeclaration(
    name="ae_journal_entry_create",
    description=(
        "Create a new entry in an Aether Journal. "
        "Use this to save notes, summaries, or any content the user wants to store. "
        "Always call ae_journal_search first to check for existing entries on the same topic."
    ),
    parameters=types.Schema(
        type=types.Type.OBJECT,
        properties={
            "journal_id": types.Schema(
                type=types.Type.STRING,
                description=(
                    "The id_random of the target journal. "
                    "Ask the user which journal to write to if not specified."
                ),
            ),
            "title": types.Schema(
                type=types.Type.STRING,
                description="Entry title",
            ),
            "content": types.Schema(
                type=types.Type.STRING,
                description="Full entry content (markdown supported)",
            ),
            "summary": types.Schema(
                type=types.Type.STRING,
                description="Optional short summary (1-2 sentences)",
            ),
            "tags": types.Schema(
                type=types.Type.STRING,
                description="Optional comma-separated tags (e.g. 'wireguard, networking, homelab')",
            ),
        },
        required=["journal_id", "title", "content"],
    ),
)

_ae_task_list_declaration = types.FunctionDeclaration(
    name="ae_task_list",
    description=(
        "List tasks from the agents_sync Kanban board (todo and in-progress). "
        "Use this when asked about current work, pending tasks, or project status."
    ),
    parameters=types.Schema(
        type=types.Type.OBJECT,
        properties={
            "include_done": types.Schema(
                type=types.Type.BOOLEAN,
                description="If true, also include completed tasks (default false)",
            ),
        },
    ),
)

_file_read_declaration = types.FunctionDeclaration(
    name="file_read",
    description=(
        "Read a local file and return its contents. "
        "Allowed directories: ~/agents_sync/, ~/OSIT_dev/, ~/DgrZone_Nextcloud/, ~/OSIT_Nextcloud/. "
        "Use this to read documentation, notes, CLAUDE.md files, or config references. "
        "If given a directory path, returns a directory listing instead."
    ),
    parameters=types.Schema(
        type=types.Type.OBJECT,
        properties={
            "path": types.Schema(
                type=types.Type.STRING,
                description=(
                    "Absolute or home-relative path to the file "
                    "(e.g. ~/agents_sync/CLAUDE.md or /home/scott/agents_sync/tasks/01_todo/)"
                ),
            ),
            "max_lines": types.Schema(
                type=types.Type.INTEGER,
                description="Optional line limit (default 500)",
            ),
        },
        required=["path"],
    ),
)


# ---------------------------------------------------------------------------
# Registry: maps tool name → async callable
# ---------------------------------------------------------------------------

_CALLABLES: dict[str, callable] = {
    "web_search": _web_search,
    "ae_journal_search": _ae_journal_search,
    "ae_journal_entry_create": _ae_journal_entry_create,
    "ae_task_list": _ae_task_list,
    "file_read": _file_read,
    "claude_allow_dir": _claude_allow_dir,
    "task_list": _task_list,
    "task_create": _task_create,
    "task_update": _task_update,
    "task_complete": _task_complete,
    "scratch_read": _scratch_read,
    "scratch_write": _scratch_write,
    "scratch_append": _scratch_append,
    "scratch_clear": _scratch_clear,
}

_claude_allow_dir_declaration = types.FunctionDeclaration(
    name="claude_allow_dir",
    description=(
        "Add a directory to Claude Code's auto-allow list so Claude can read or write "
        "files there without prompting. Edits ~/.claude/settings.json on the local machine. "
        "Use this when Claude is silently hanging or being blocked from accessing a directory. "
        "Changes take effect in the next Claude Code session."
    ),
    parameters=types.Schema(
        type=types.Type.OBJECT,
        properties={
            "path": types.Schema(
                type=types.Type.STRING,
                description=(
                    "Absolute or home-relative path to the directory "
                    "(e.g. ~/OSIT_dev/aether_api_fastapi or /home/scott/agents_sync)"
                ),
            ),
            "mode": types.Schema(
                type=types.Type.STRING,
                description="Permission mode: 'r' (read-only), 'w' (write-only), or 'rw' (both). Default: rw",
            ),
        },
        required=["path"],
    ),
)

_task_list_declaration = types.FunctionDeclaration(
    name="task_list",
    description=(
        "List personal tasks from Inara's task list. "
        "Use this to check what's on the list, review pending work, or find a task ID. "
        "Optionally filter by status: 'todo', 'in_progress', or 'done'."
    ),
    parameters=types.Schema(
        type=types.Type.OBJECT,
        properties={
            "status": types.Schema(
                type=types.Type.STRING,
                description="Filter by status: 'todo', 'in_progress', or 'done'. Omit to list all.",
            ),
        },
    ),
)

_task_create_declaration = types.FunctionDeclaration(
    name="task_create",
    description=(
        "Add a new task to Inara's personal task list. "
        "Use this when the user asks to remember something, add a to-do, or track a follow-up."
    ),
    parameters=types.Schema(
        type=types.Type.OBJECT,
        properties={
            "title": types.Schema(
                type=types.Type.STRING,
                description="Short task title",
            ),
            "description": types.Schema(
                type=types.Type.STRING,
                description="Optional longer description or context",
            ),
            "priority": types.Schema(
                type=types.Type.STRING,
                description="Priority: 'low', 'normal', or 'high'. Default: normal.",
            ),
        },
        required=["title"],
    ),
)

_task_update_declaration = types.FunctionDeclaration(
    name="task_update",
    description=(
        "Update an existing task. Use task_list first to get the task ID. "
        "Can update status, title, description, or priority. "
        "To just mark complete, use task_complete instead."
    ),
    parameters=types.Schema(
        type=types.Type.OBJECT,
        properties={
            "task_id": types.Schema(
                type=types.Type.STRING,
                description="Task ID (e.g. t_abc123) — get from task_list",
            ),
            "status": types.Schema(
                type=types.Type.STRING,
                description="New status: 'todo', 'in_progress', or 'done'",
            ),
            "title": types.Schema(
                type=types.Type.STRING,
                description="Updated title",
            ),
            "description": types.Schema(
                type=types.Type.STRING,
                description="Updated description",
            ),
            "priority": types.Schema(
                type=types.Type.STRING,
                description="Updated priority: 'low', 'normal', or 'high'",
            ),
        },
        required=["task_id"],
    ),
)

_task_complete_declaration = types.FunctionDeclaration(
    name="task_complete",
    description=(
        "Mark a task as done. Use task_list first to get the task ID. "
        "Shorthand for task_update with status='done'."
    ),
    parameters=types.Schema(
        type=types.Type.OBJECT,
        properties={
            "task_id": types.Schema(
                type=types.Type.STRING,
                description="Task ID (e.g. t_abc123) — get from task_list",
            ),
        },
        required=["task_id"],
    ),
)

_scratch_read_declaration = types.FunctionDeclaration(
    name="scratch_read",
    description=(
        "Read the full contents of the scratchpad. "
        "Use this to recall working notes, mid-task context, or anything previously jotted down. "
        "The scratchpad is transient — nothing here is distilled or archived."
    ),
    parameters=types.Schema(type=types.Type.OBJECT, properties={}),
)

_scratch_write_declaration = types.FunctionDeclaration(
    name="scratch_write",
    description=(
        "Replace the entire scratchpad with new content. "
        "Use this to set a clean working note, replacing whatever was there before. "
        "For adding without replacing, use scratch_append instead."
    ),
    parameters=types.Schema(
        type=types.Type.OBJECT,
        properties={
            "content": types.Schema(
                type=types.Type.STRING,
                description="The new scratchpad content (markdown supported)",
            ),
        },
        required=["content"],
    ),
)

_scratch_append_declaration = types.FunctionDeclaration(
    name="scratch_append",
    description=(
        "Add a new section to the bottom of the scratchpad without replacing existing content. "
        "Each section gets a timestamp heading unless you supply one."
    ),
    parameters=types.Schema(
        type=types.Type.OBJECT,
        properties={
            "content": types.Schema(
                type=types.Type.STRING,
                description="The content to append (markdown supported)",
            ),
            "heading": types.Schema(
                type=types.Type.STRING,
                description="Optional section heading. Defaults to current UTC timestamp.",
            ),
        },
        required=["content"],
    ),
)

_scratch_clear_declaration = types.FunctionDeclaration(
    name="scratch_clear",
    description="Erase everything in the scratchpad. Use when the working notes are no longer needed.",
    parameters=types.Schema(type=types.Type.OBJECT, properties={}),
)


# Gemini Tool object — pass this to GenerateContentConfig
TOOL_DECLARATIONS = [
    types.Tool(function_declarations=[
        _web_search_declaration,
        _ae_journal_search_declaration,
        _ae_journal_entry_create_declaration,
        _ae_task_list_declaration,
        _file_read_declaration,
        _claude_allow_dir_declaration,
        _task_list_declaration,
        _task_create_declaration,
        _task_update_declaration,
        _task_complete_declaration,
        _scratch_read_declaration,
        _scratch_write_declaration,
        _scratch_append_declaration,
        _scratch_clear_declaration,
    ])
]


async def call_tool(name: str, args: dict) -> str:
    """Dispatch a tool call by name. Returns result as a string."""
    fn = _CALLABLES.get(name)
    if fn is None:
        return f"Unknown tool: {name}"
    return await fn(**args)