backend/app/tools/__init__.py

"""Agent tool registry.

A deliberately small, trimmed tool set shared by the browser agent. Each tool has
a name, a JSON-schema parameter spec, and a handler bound at runtime. The schema
definitions are part of the cached prompt prefix, so keeping the set minimal is
itself a token optimization (the reference docs note unused tool defs cost tokens
on every call).
"""
from __future__ import annotations

from dataclasses import dataclass
from typing import Any, Callable


@dataclass
class Tool:
    name: str
    description: str
    parameters: dict  # JSON-schema
    handler: Callable[..., Any] | None = None


# Canonical browser tool definitions (the "contract" the agent reasons over).
BROWSER_TOOLS: list[Tool] = [
    Tool(
        name="navigate",
        description="Go to a URL.",
        parameters={"type": "object", "properties": {"url": {"type": "string"}},
                    "required": ["url"]},
    ),
    Tool(
        name="click",
        description="Click an element by CSS selector or accessibility label.",
        parameters={"type": "object", "properties": {"selector": {"type": "string"}},
                    "required": ["selector"]},
    ),
    Tool(
        name="fill",
        description="Type text into a form field identified by selector.",
        parameters={"type": "object", "properties": {
            "selector": {"type": "string"}, "text": {"type": "string"}},
            "required": ["selector", "text"]},
    ),
    Tool(
        name="extract",
        description="Extract the structured data visible on the page as JSON.",
        parameters={"type": "object", "properties": {}},
    ),
    Tool(
        name="screenshot",
        description="Capture the current page for the audit trail.",
        parameters={"type": "object", "properties": {}},
    ),
    Tool(
        name="done",
        description="Finish the task; args.result is the final JSON payload.",
        parameters={"type": "object", "properties": {"result": {}},
                    "required": ["result"]},
    ),
]


class ToolRegistry:
    def __init__(self, tools: list[Tool]) -> None:
        self._tools = {t.name: t for t in tools}

    def names(self) -> list[str]:
        return list(self._tools)

    def get(self, name: str) -> Tool | None:
        return self._tools.get(name)

    def bind(self, name: str, handler: Callable[..., Any]) -> None:
        if name in self._tools:
            self._tools[name].handler = handler

    def definitions(self) -> list[dict]:
        return [
            {"name": t.name, "description": t.description, "parameters": t.parameters}
            for t in self._tools.values()
        ]

    def call(self, name: str, args: dict) -> Any:
        tool = self._tools.get(name)
        if tool is None or tool.handler is None:
            raise ValueError(f"unknown or unbound tool: {name}")
        return tool.handler(**(args or {}))


def browser_registry() -> ToolRegistry:
    return ToolRegistry([Tool(t.name, t.description, t.parameters) for t in BROWSER_TOOLS])