ml-knowledge-platform/knowledge_platform/tools/base.py

"""Tool base classes for Crystal's agentic tool framework.

Defines the foundational abstractions: Tool (base class for all tools),
ToolResult (execution outcome), and ToolParameter (schema definition).
All tools produce Anthropic-compatible JSON schemas for parameter validation.
"""

from __future__ import annotations

import asyncio
from abc import ABC, abstractmethod
from enum import Enum
from typing import Any, ClassVar

from pydantic import BaseModel, Field


class ToolParameter(BaseModel):
    """Schema definition for a single tool parameter.

    Maps directly to JSON Schema properties for Anthropic tool calling.
    """

    name: str = Field(description="Parameter name")
    type: str = Field(description="JSON Schema type (string, integer, boolean, array, object)")
    description: str = Field(description="Human-readable description for LLM consumption")
    required: bool = Field(default=True, description="Whether the parameter is required")
    default: Any = Field(default=None, description="Default value when not provided")
    enum: list[str] | None = Field(default=None, description="Allowed values for enum types")
    items: dict[str, Any] | None = Field(
        default=None, description="Schema for array item types"
    )


class ToolResultStatus(str, Enum):
    """Outcome status of a tool execution."""

    SUCCESS = "success"
    ERROR = "error"
    TIMEOUT = "timeout"


class ToolResult(BaseModel):
    """Outcome of a tool execution.

    Encapsulates success/failure state, output data, and error information.
    Designed for direct serialization back to the LLM.
    """

    status: ToolResultStatus = Field(description="Execution outcome status")
    output: Any = Field(default=None, description="Tool output data (any JSON-serializable value)")
    error: str | None = Field(default=None, description="Error message if execution failed")
    metadata: dict[str, Any] = Field(
        default_factory=dict,
        description="Additional execution metadata (timing, tool name, etc.)",
    )

    @classmethod
    def success(cls, output: Any, **metadata: Any) -> ToolResult:
        """Create a successful result."""
        return cls(status=ToolResultStatus.SUCCESS, output=output, metadata=metadata)

    @classmethod
    def fail(cls, error: str, **metadata: Any) -> ToolResult:
        """Create a failed result."""
        return cls(status=ToolResultStatus.ERROR, error=error, metadata=metadata)

    @classmethod
    def timed_out(cls, timeout_seconds: float, **metadata: Any) -> ToolResult:
        """Create a timeout result."""
        return cls(
            status=ToolResultStatus.TIMEOUT,
            error=f"Execution timed out after {timeout_seconds}s",
            metadata=metadata,
        )

    @property
    def is_success(self) -> bool:
        return self.status == ToolResultStatus.SUCCESS

    @property
    def is_error(self) -> bool:
        return self.status in (ToolResultStatus.ERROR, ToolResultStatus.TIMEOUT)


class Tool(ABC):
    """Base class for all Crystal tools.

    Subclasses must define class-level `name`, `description`, and `parameters`,
    then implement the async `execute()` method.

    The `to_anthropic_schema()` method produces a tool definition compatible
    with the Anthropic Messages API tool_use format.

    Example::

        class ReadFileTool(Tool):
            name = "read_file"
            description = "Read contents of a file"
            parameters = [
                ToolParameter(
                    name="file_path",
                    type="string",
                    description="Absolute path to the file",
                ),
            ]

            async def execute(self, **kwargs: Any) -> ToolResult:
                path = kwargs["file_path"]
                content = Path(path).read_text()
                return ToolResult.success(content)
    """

    name: ClassVar[str]
    description: ClassVar[str]
    parameters: ClassVar[list[ToolParameter]]

    @abstractmethod
    async def execute(self, **kwargs: Any) -> ToolResult:
        """Execute the tool with validated parameters.

        Args:
            **kwargs: Tool-specific parameters, pre-validated against the schema.

        Returns:
            ToolResult with execution outcome.
        """

    def get_input_schema(self) -> dict[str, Any]:
        """Generate JSON Schema for this tool's parameters.

        Returns a JSON Schema object suitable for Anthropic tool calling.
        """
        properties: dict[str, Any] = {}
        required: list[str] = []

        for param in self.parameters:
            prop: dict[str, Any] = {
                "type": param.type,
                "description": param.description,
            }
            if param.enum is not None:
                prop["enum"] = param.enum
            if param.items is not None:
                prop["items"] = param.items
            if param.default is not None:
                prop["default"] = param.default

            properties[param.name] = prop

            if param.required:
                required.append(param.name)

        return {
            "type": "object",
            "properties": properties,
            "required": required,
        }

    def to_anthropic_schema(self) -> dict[str, Any]:
        """Produce Anthropic Messages API tool definition.

        Returns:
            Dict with ``name``, ``description``, and ``input_schema`` keys,
            ready for the ``tools`` parameter of the Messages API.
        """
        return {
            "name": self.name,
            "description": self.description,
            "input_schema": self.get_input_schema(),
        }

    def validate_parameters(self, params: dict[str, Any]) -> list[str]:
        """Validate parameters against the schema.

        Returns a list of validation error messages (empty if valid).
        """
        errors: list[str] = []
        schema = self.get_input_schema()
        required_params = schema.get("required", [])
        properties = schema.get("properties", {})

        for req in required_params:
            if req not in params:
                errors.append(f"Missing required parameter: {req}")

        for key in params:
            if key not in properties:
                errors.append(f"Unknown parameter: {key}")

        return errors

    def __repr__(self) -> str:
        return f"<Tool {self.name}>"