ai-rpg/app/engine/tools/base.py

"""Base types for tool system: Tool, ToolContext, ToolResult, ToolRegistry."""

from __future__ import annotations

import abc
import uuid
from collections.abc import Awaitable, Callable
from dataclasses import dataclass, field
from datetime import datetime, timezone
from typing import Any

from sqlalchemy.ext.asyncio import AsyncSession

from app.core.logging import get_logger
from app.core.state_validator import apply_patch
from app.models import Entity, StepToolCall, World

_logger = get_logger(__name__)


# --------------------------------------------------------------------------- #
# Context & Result
# --------------------------------------------------------------------------- #
@dataclass
class ToolContext:
    """Per-iteration context passed to every tool call."""

    db: AsyncSession
    world: World
    user_id: uuid.UUID | None = None
    step_id: uuid.UUID | None = None
    stage: str = ""
    sse_emitter: Any = None  # callable: async (event, data) -> None
    pending_state_changes: dict[str, Any] = field(default_factory=dict)


@dataclass
class ToolResult:
    ok: bool
    data: dict[str, Any] = field(default_factory=dict)
    message: str = ""
    error_code: str = ""
    error_message: str = ""

    def to_dict(self) -> dict[str, Any]:
        if self.ok:
            return {"ok": True, "data": self.data, "message": self.message}
        return {
            "ok": False,
            "error": {"code": self.error_code, "message": self.error_message},
        }


# --------------------------------------------------------------------------- #
# Tool base class
# --------------------------------------------------------------------------- #
class Tool(abc.ABC):
    """Abstract base for all tools."""

    name: str = ""
    category: str = "game"  # game | interaction | schema
    stages: set[str] = set()  # which stages can use this tool
    description: str = ""
    parameters_schema: dict[str, Any] = {}

    @abc.abstractmethod
    async def execute(self, arguments: dict[str, Any], ctx: ToolContext) -> ToolResult:
        """Run the tool. Must be idempotent within a single transaction."""

    def to_openai_format(self) -> dict[str, Any]:
        """Serialize to OpenAI tools format."""
        return {
            "type": "function",
            "function": {
                "name": self.name,
                "description": self.description,
                "parameters": self.parameters_schema,
            },
        }


# --------------------------------------------------------------------------- #
# Registry
# --------------------------------------------------------------------------- #
class ToolRegistry:
    """Holds all registered tools and dispatches calls."""

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

    def register(self, tool: Tool) -> None:
        if not tool.name:
            raise ValueError("Tool name is required")
        if tool.name in self._tools:
            raise ValueError(f"Tool {tool.name} already registered")
        self._tools[tool.name] = tool

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

    def list_for_stage(self, stage: str) -> list[Tool]:
        """Return all tools available at the given stage."""
        return [t for t in self._tools.values() if stage in t.stages or "*" in t.stages]

    def to_openai_format(self, stage: str) -> list[dict[str, Any]]:
        return [t.to_openai_format() for t in self.list_for_stage(stage)]

    async def execute(
        self,
        name: str,
        arguments: dict[str, Any],
        ctx: ToolContext,
    ) -> ToolResult:
        """Execute a tool by name. Logs to step_tool_calls and emits SSE."""
        from sqlalchemy import select as sa_select

        tool = self.get(name)
        executed_at = datetime.now(timezone.utc)
        if tool is None:
            result = ToolResult(
                ok=False,
                error_code="unknown_tool",
                error_message=f"Tool {name!r} is not registered",
            )
        else:
            try:
                result = await tool.execute(arguments, ctx)
            except Exception as e:  # noqa: BLE001
                _logger.exception("tool_execution_failed", tool=name, error=str(e))
                result = ToolResult(
                    ok=False,
                    error_code="tool_exception",
                    error_message=str(e),
                )

        # Log to step_tool_calls (if we have a step_id)
        if ctx.step_id is not None:
            try:
                ctx.db.add(
                    StepToolCall(
                        step_id=ctx.step_id,
                        tool_name=name,
                        arguments=arguments,
                        result=result.to_dict(),
                        is_success=result.ok,
                        executed_at=executed_at,
                    )
                )
                await ctx.db.flush()
            except Exception as e:  # noqa: BLE001
                _logger.error("tool_log_failed", tool=name, error=str(e))

        # Emit SSE
        if ctx.sse_emitter is not None:
            try:
                await ctx.sse_emitter(
                    "tool_call",
                    {
                        "tool": name,
                        "arguments": arguments,
                        "result": result.to_dict(),
                        "is_success": result.ok,
                    },
                )
            except Exception as e:  # noqa: BLE001
                _logger.warning("sse_tool_call_failed", tool=name, error=str(e))

        return result


# --------------------------------------------------------------------------- #
# Helpers used by entity_* tools
# --------------------------------------------------------------------------- #
async def get_entity_by_query(
    db: AsyncSession, world_id: uuid.UUID, query: Any
) -> Entity | None:
    """Resolve entity by UUID string or by {entity_type, name}."""
    from sqlalchemy import select as sa_select

    if isinstance(query, str):
        try:
            eid = uuid.UUID(query)
        except ValueError:
            return None
        return (
            await db.execute(
                sa_select(Entity).where(
                    Entity.id == eid, Entity.world_id == world_id
                )
            )
        ).scalar_one_or_none()
    elif isinstance(query, dict):
        et = query.get("entity_type")
        nm = query.get("name")
        if not et or not nm:
            return None
        return (
            await db.execute(
                sa_select(Entity).where(
                    Entity.world_id == world_id,
                    Entity.entity_type == et,
                    Entity.name == nm,
                    Entity.deleted_at.is_(None),
                )
            )
        ).scalar_one_or_none()
    return None


def apply_env_patch(environment: dict, patch: dict) -> tuple[dict, list[str]]:
    """Wrapper around state_validator.apply_patch for environment dicts."""
    return apply_patch(environment, patch)


# Singleton registry (instantiated in `app.engine.tools.__init__`)
_registry: ToolRegistry | None = None


def get_registry() -> ToolRegistry:
    global _registry
    if _registry is None:
        from app.engine.tools.register_all import build_default_registry

        _registry = build_default_registry()
    return _registry


def reset_registry() -> None:
    """Reset the cached registry — used in tests."""
    global _registry
    _registry = None