path: root/tools.py

from abc import ABC, abstractmethod
from dataclasses import dataclass
from datetime import UTC, datetime, timedelta, timezone
from pathlib import Path

import requests

from memory import WorkingMemory


class ToolError(Exception):
    pass


@dataclass
class ToolResult:
    content: str


class Tool(ABC):
    @property
    @abstractmethod
    def name(self) -> str: ...

    @property
    @abstractmethod
    def description(self) -> str: ...

    @property
    @abstractmethod
    def schema(self) -> dict[str, object]: ...

    @abstractmethod
    def execute(self, arguments: dict[str, object]) -> ToolResult: ...


class ToolRegistry:
    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 all(self) -> list[Tool]:
        return list(self.tools.values())

    def schemas(self) -> list[dict[str, object]]:
        return [
            {
                "type": "function",
                "function": {
                    "name": tool.name,
                    "description": tool.description,
                    "parameters": tool.schema,
                },
            }
            for tool in self.tools.values()
        ]


class TaskStartTool(Tool):
    def __init__(self, memory: WorkingMemory) -> None:
        self.memory = memory

    @property
    def name(self) -> str:
        return "task_start"

    @property
    def description(self) -> str:
        return (
            "Record the user's ongoing project or goal as the current task. "
            "Use this when the user intends to work on something across multiple "
            "turns or over a period of time. Do not use it for one-off requests "
            "or temporary side conversations."
        )

    @property
    def schema(self) -> dict[str, object]:
        return {
            "type": "object",
            "properties": {
                "objective": {
                    "type": "string",
                    "description": (
                        "A concise description of the ongoing objective the user wants "
                        "to pursue. For example: 'Build an image tagging API' or "
                        "'Design and implement a Rust CLI that converts directory trees "
                        "into static HTML.'"
                    ),
                },
            },
            "required": ["objective"],
        }

    def execute(self, arguments: dict[str, object]) -> ToolResult:
        objective = arguments.get("objective")

        if not isinstance(objective, str):
            raise ToolError("objective must be a string")

        if not objective.strip():
            raise ToolError("objective must not be empty")

        self.memory.start_task(objective)

        return ToolResult(content=f"Started task: {objective}")


class MemoryAddFactTool(Tool):
    def __init__(self, memory: WorkingMemory) -> None:
        self.memory = memory

    @property
    def name(self) -> str:
        return "memory_add_fact"

    @property
    def description(self) -> str:
        return "Add a useful fact to the current working memory."

    @property
    def schema(self) -> dict[str, object]:
        return {
            "type": "object",
            "properties": {
                "fact": {
                    "type": "string",
                    "description": "A fact worth remembering during the current task.",
                },
            },
            "required": ["fact"],
        }

    def execute(self, arguments: dict[str, object]) -> ToolResult:
        fact = arguments.get("fact")

        if not isinstance(fact, str):
            raise ToolError("fact must be a string")

        self.memory.add_fact(fact)

        return ToolResult(content=f"Added working-memory fact: {fact}")


class MemoryAddNoteTool(Tool):
    def __init__(self, memory: WorkingMemory) -> None:
        self.memory = memory

    @property
    def name(self) -> str:
        return "memory_add_note"

    @property
    def description(self) -> str:
        return "Add a useful note to the current working memory."

    @property
    def schema(self) -> dict[str, object]:
        return {
            "type": "object",
            "properties": {
                "note": {
                    "type": "string",
                    "description": "A useful observation, reminder, or note for the current task.",
                },
            },
            "required": ["note"],
        }

    def execute(self, arguments: dict[str, object]) -> ToolResult:
        note = arguments.get("note")

        if not isinstance(note, str):
            raise ToolError("note must be a string")

        self.memory.add_note(note)

        return ToolResult(content=f"Added working-memory note: {note}")


class MemoryRemoveFactTool(Tool):
    def __init__(self, memory: WorkingMemory) -> None:
        self.memory = memory

    @property
    def name(self) -> str:
        return "memory_remove_fact"

    @property
    def description(self) -> str:
        return "Remove a fact from the working memory when it is no longer true or relevant."

    @property
    def schema(self) -> dict[str, object]:
        return {
            "type": "object",
            "properties": {
                "fact": {
                    "type": "string",
                    "description": "The exact fact to remove from working memory.",
                },
            },
            "required": ["fact"],
        }

    def execute(self, arguments: dict[str, object]) -> ToolResult:
        fact = arguments.get("fact")

        if not isinstance(fact, str):
            raise ToolError("fact must be a string")

        self.memory.add_fact(fact)

        return ToolResult(content=f"Removed fact: {fact}")


class MemoryRemoveNoteTool(Tool):
    def __init__(self, memory: WorkingMemory) -> None:
        self.memory = memory

    @property
    def name(self) -> str:
        return "memory_remove_note"

    @property
    def description(self) -> str:
        return "Remove a note from the working memory when it is no longer true or relevant."

    @property
    def schema(self) -> dict[str, object]:
        return {
            "type": "object",
            "properties": {
                "note": {
                    "type": "string",
                    "description": "The exact note to remove from working memory.",
                },
            },
            "required": ["note"],
        }

    def execute(self, arguments: dict[str, object]) -> ToolResult:
        note = arguments.get("note")

        if not isinstance(note, str):
            raise ToolError("note must be a string")

        self.memory.add_note(note)

        return ToolResult(content=f"Removed note: {note}")


class HelloTool(Tool):
    @property
    def name(self) -> str:
        return "hello"

    @property
    def description(self) -> str:
        return "Say hello to someone."

    @property
    def schema(self) -> dict[str, object]:
        return {
            "type": "object",
            "properties": {
                "name": {
                    "type": "string",
                    "description": "The person to greet.",
                },
            },
            "required": ["name"],
        }

    def execute(self, arguments: dict[str, object]) -> ToolResult:
        name = arguments.get("name", "stranger")
        return ToolResult(content=f"Hello, {name}!")


class CalculatorTool(Tool):
    @property
    def name(self) -> str:
        return "calculator"

    @property
    def description(self) -> str:
        return "Perform basic arithemtic"

    @property
    def schema(self) -> dict[str, object]:
        return {
            "type": "object",
            "properties": {
                "a": {
                    "type": "number",
                    "description": "The first number.",
                },
                "b": {
                    "type": "number",
                    "description": "The second number.",
                },
                "operation": {
                    "type": "string",
                    "enum": ["add", "subtract", "multiply", "divide"],
                    "description": "The arithemtic operation to perform",
                },
            },
            "required": ["a", "b", "operation"],
        }

    def execute(self, arguments: dict[str, object]) -> ToolResult:
        a = arguments["a"]
        b = arguments["b"]
        operation = arguments["operation"]

        if not isinstance(a, (int, float)):
            raise ToolError("a must be a number")

        if not isinstance(b, (int, float)):
            raise ToolError("b must be a number")

        if not isinstance(operation, str):
            raise ToolError("operation must be a string")

        if operation == "add":
            result = a + b
        elif operation == "subtract":
            result = a - b
        elif operation == "multiply":
            result = a * b
        elif operation == "divide":
            if b == 0:
                raise ToolError("cannot divide by zero")
            result = a / b
        else:
            raise ToolError(f"unknown operation: {operation}")

        return ToolResult(content=str(result))


class WeatherTool(Tool):
    @property
    def name(self) -> str:
        return "weather"

    @property
    def description(self) -> str:
        return "Get current weather conditions for a location"

    @property
    def schema(self) -> dict[str, object]:
        return {
            "type": "object",
            "properties": {
                "location": {
                    "type": "string",
                    "description": ("The city or location to get weather for."),
                },
            },
            "required": ["location"],
        }

    def execute(self, arguments: dict[str, object]) -> ToolResult:
        location = arguments.get("location")

        if not isinstance(location, str):
            raise ToolError("location must be a string")

        geocoding_response = requests.get(
            "https://geocoding-api.open-meteo.com/v1/search",
            params={
                "name": location,
                "count": 1,
                "language": "en",
                "format": "json",
            },
            timeout=10,
        )
        geocoding_response.raise_for_status()

        geocoding_data = geocoding_response.json()
        results = geocoding_data.get("results")

        if not results:
            raise ToolError(f"location not found: {location}")

        place = results[0]

        latitude = place["latitude"]
        longitude = place["longitude"]
        timezone_name = place["timezone"]

        forecast_response = requests.get(
            "https://api.open-meteo.com/v1/forecast",
            params={
                "latitude": latitude,
                "longitude": longitude,
                "hourly": (
                    "temperature_2m,"
                    "precipitation_probability,"
                    "precipitation,"
                    "weather_code,"
                    "wind_speed_10m"
                ),
                "timezone": timezone_name,
                "forecast_days": 2,
            },
            timeout=10,
        )
        forecast_response.raise_for_status()

        forecast = forecast_response.json()
        hourly = forecast["hourly"]

        tomorrow = (
            (datetime.now(UTC) + timedelta(days=1))
            .astimezone(timezone(timedelta(seconds=forecast["utc_offset_seconds"])))
            .date()
        )

        temperatures = []
        precipitation_probabilities = []
        precipitation = []
        wind_speeds = []
        weather_codes = []

        for index, timestamp in enumerate(hourly["time"]):
            forecast_date = datetime.fromisoformat(timestamp).date()

            if forecast_date != tomorrow:
                continue

            temperatures.append(hourly["temperature_2m"][index])
            precipitation_probabilities.append(
                hourly["precipitation_probability"][index]
            )
            precipitation.append(hourly["precipitation"][index])
            wind_speeds.append(hourly["wind_speed_10m"][index])
            weather_codes.append(hourly["weather_code"][index])

        if not temperatures:
            raise ToolError("no forecast data available for tomorrow")
        return ToolResult(
            content=(
                f"Location: {place['name']}, {place['country']}\n"
                f"Forecast date: {tomorrow}\n"
                f"Temperature: {min(temperatures):.1f}–"
                f"{max(temperatures):.1f} °C\n"
                f"Maximum precipitation probability: "
                f"{max(precipitation_probabilities)}%\n"
                f"Total precipitation: {sum(precipitation):.1f} mm\n"
                f"Maximum wind speed: {max(wind_speeds):.1f} km/h\n"
                f"Weather codes: {sorted(set(weather_codes))}"
            )
        )


class FileSystemListTool(Tool):
    @property
    def name(self) -> str:
        return "filesystem_list"

    @property
    def description(self) -> str:
        return "List files and directories inside the agent's sandbox."

    @property
    def schema(self) -> dict[str, object]:
        return {
            "type": "object",
            "properties": {
                "path": {
                    "type": "string",
                    "description": (
                        "Relative path inside the filesystem sandbox. "
                        "Use '.' for the sandbox root."
                    ),
                },
            },
            "required": ["path"],
        }

    def execute(self, arguments: dict[str, object]) -> ToolResult:
        path = arguments.get("path")

        if not isinstance(path, str):
            raise ToolError("path must be a string")

        sandbox = Path.home() / "Projects" / "personal-agent" / "data"
        target = (sandbox / path).resolve()

        if not target.is_relative_to(sandbox.resolve()):
            raise ToolError("path is outside the filesystem sandbox")

        if not target.exists():
            raise ToolError(f"path does not exist: {path}")

        if not target.is_dir():
            raise ToolError(f"path is not a directory: {path}")

        entries = []

        for entry in sorted(target.iterdir()):
            kind = "directory" if entry.is_dir() else "file"
            entries.append(f"{kind}: {entry.name}")

        if not entries:
            return ToolResult(content="Directory is empty.")

        return ToolResult(content="\n".join(entries))


class FileSystemReadTool(Tool):
    MAX_FILE_SIZE = 1024 * 1024  # 1 MiB

    @property
    def name(self) -> str:
        return "filesystem_read"

    @property
    def description(self) -> str:
        return "Read the contents of a file inside the agent's sandbox."

    @property
    def schema(self) -> dict[str, object]:
        return {
            "type": "object",
            "properties": {
                "path": {
                    "type": "string",
                    "description": ("Relative path inside the filesystem sandbox."),
                },
            },
            "required": ["path"],
        }

    def execute(self, arguments: dict[str, object]) -> ToolResult:
        path = arguments.get("path")

        if not isinstance(path, str):
            raise ToolError("path must be a string")

        sandbox = Path.home() / "Projects" / "personal-agent" / "data"
        sandbox = sandbox.resolve()
        target = (sandbox / path).resolve()

        if not target.is_relative_to(sandbox):
            raise ToolError("path is outside the filesystem sandbox")

        if not target.exists():
            raise ToolError(f"path does not exist: {path}")

        if not target.is_file():
            raise ToolError(f"path is not a file: {path}")

        if target.stat().st_size > self.MAX_FILE_SIZE:
            raise ToolError("file is too large to read")

        try:
            content = target.read_text(encoding="utf-8")
        except ToolError as error:
            raise ToolError("file is not valid UTF-8 text") from error

        return ToolResult(content=content)


if __name__ == "__main__":
    registry = ToolRegistry()

    registry.register(HelloTool())
    print(registry.schemas())
    tool = registry.get("hello")

    if tool is not None:
        result = tool.execute({"name": "Joseph"})
        print(result.content)