STT-Assistant/src/stt_daemon/assistant/tools.py

"""Tool registry: typed tools the assistant may invoke.

Arguments arrive from the model as a JSON string; they are parsed and validated
against the tool's schema before anything runs. Unknown tools, malformed JSON,
missing/extra/wrong-typed arguments are all rejected (parse, don't trust).
"""

import json
import logging
from dataclasses import dataclass
from datetime import datetime
from typing import Callable

logger = logging.getLogger(__name__)

READ_ONLY = "read_only"
MUTATING = "mutating"

_TYPE_MAP = {
    "string": str,
    "number": (int, float),
    "integer": int,
    "boolean": bool,
    "array": list,
    "object": dict,
}


@dataclass
class Tool:
    """A callable tool with an OpenAI-style JSON schema."""

    name: str
    description: str
    parameters: dict
    run: Callable[[dict], str]
    side_effect: str = READ_ONLY
    requires_confirmation: bool = False

    def spec(self) -> dict:
        """OpenAI `tools` entry for this tool."""
        return {
            "type": "function",
            "function": {
                "name": self.name,
                "description": self.description,
                "parameters": self.parameters,
            },
        }


def _validate_schema(schema: dict, args: dict) -> str | None:
    """Return an error message, or None when the arguments are acceptable."""
    properties = schema.get("properties", {})
    for key in schema.get("required", []):
        if key not in args:
            return f"missing required argument '{key}'"
    for key, value in args.items():
        if key not in properties:
            return f"unexpected argument '{key}'"
        expected = properties[key].get("type")
        if expected not in _TYPE_MAP:
            continue
        if expected in ("integer", "number") and isinstance(value, bool):
            return f"argument '{key}' must be {expected}"
        if not isinstance(value, _TYPE_MAP[expected]):
            return f"argument '{key}' must be {expected}"
    return None


class ToolRegistry:
    """Holds the tools the assistant may call."""

    def __init__(self) -> None:
        self._tools: dict[str, Tool] = {}

    def register(self, tool: Tool) -> None:
        self._tools[tool.name] = tool

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

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

    def specs(self) -> list[dict]:
        return [tool.spec() for tool in self._tools.values()]

    def validate(self, name: str, arguments) -> tuple[dict | None, str | None]:
        """Parse and validate model-supplied arguments for `name`."""
        tool = self.get(name)
        if tool is None:
            return None, f"unknown tool: {name}"
        if isinstance(arguments, str):
            try:
                parsed = json.loads(arguments) if arguments.strip() else {}
            except json.JSONDecodeError:
                return None, f"invalid JSON arguments for {name}"
        else:
            parsed = dict(arguments or {})
        if not isinstance(parsed, dict):
            return None, f"arguments must be an object for {name}"
        error = _validate_schema(tool.parameters, parsed)
        if error is not None:
            return None, f"{name}: {error}"
        return parsed, None


def _get_datetime(_args: dict) -> str:
    return datetime.now().astimezone().isoformat(timespec="seconds")


def default_registry() -> ToolRegistry:
    """Registry with the built-in tools."""
    registry = ToolRegistry()
    registry.register(
        Tool(
            name="get_datetime",
            description="Get the current local date and time.",
            parameters={"type": "object", "properties": {}, "required": []},
            run=_get_datetime,
        )
    )
    return registry