path: root/personal_agent/tools/memory.py

from uuid import UUID

from personal_agent.memory.models import LongTermMemory, MemoryCategory, WorkingMemory
from personal_agent.memory.store import LongTermMemoryStore
from personal_agent.tools.base import Tool, ToolError, ToolResult


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 durable goal the user wants to pursue. "
                        "Describe what they are trying to accomplish, not mutable implementation "
                        "details, preferences, or constraints. 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.remove_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.remove_note(note)

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


class LongTermMemorySaveTool(Tool):
    def __init__(self, store: LongTermMemoryStore) -> None:
        self.store = store

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

    @property
    def description(self) -> str:
        return (
            "Save durable information for use across future conversations. "
            "Use semantic for stable facts or preferences, and episodic for "
            "meaningful past events."
        )

    @property
    def schema(self) -> dict[str, object]:
        return {
            "type": "object",
            "properties": {
                "content": {
                    "type": "string",
                    "description": "The durable information to remember.",
                },
                "category": {
                    "type": "string",
                    "enum": ["semantic", "episodic"],
                    "description": (
                        "semantic for stable facts/preferences; "
                        "episodic for meaningful past events."
                    ),
                },
            },
            "required": ["content", "category"],
        }

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

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

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

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

        try:
            memory_category = MemoryCategory(category)
        except ValueError:
            raise ToolError("category must be 'semantic' or 'episodic'")

        memory = LongTermMemory.create(content, memory_category)
        self.store.add(memory)

        return ToolResult(content=f"Saved memory: {memory.id}")


class LongTermMemorySearchTool(Tool):
    def __init__(self, store: LongTermMemoryStore) -> None:
        self.store = store

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

    @property
    def description(self) -> str:
        return (
            "Search long-term memory for information relevant to the current "
            "conversation. Use semantic for stable facts or preferences, and "
            "episodic for meaningful past events."
        )

    @property
    def schema(self) -> dict[str, object]:
        return {
            "type": "object",
            "properties": {
                "query": {
                    "type": "string",
                    "description": "The information to search for.",
                },
                "category": {
                    "type": "string",
                    "enum": ["semantic", "episodic"],
                    "description": "Optionally restrict results to one category.",
                },
                "limit": {
                    "type": "integer",
                    "minimum": 1,
                    "description": "Maximum number of memories to return.",
                },
            },
            "required": ["query"],
        }

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

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

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

        if category is not None:
            if not isinstance(category, str):
                raise ToolError("category must be a string")

            try:
                memory_category = MemoryCategory(category)
            except ValueError:
                raise ToolError("category must be 'semantic' or 'episodic'")
        else:
            memory_category = None

        if not isinstance(limit, int) or isinstance(limit, bool):
            raise ToolError("limit must be an integer")

        if limit <= 0:
            raise ToolError("limit must be positive")

        memories = self.store.search(
            query=query,
            category=memory_category,
            limit=limit,
        )

        if not memories:
            return ToolResult(content="No matching memories found.")

        lines = [
            f"- [{memory.category.value}] {memory.id}: {memory.content}"
            for memory in memories
        ]

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


class LongTermMemoryDeleteTool(Tool):
    def __init__(self, store: LongTermMemoryStore) -> None:
        self.store = store

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

    @property
    def description(self) -> str:
        return (
            "Delete a long-term memory by its exact ID. Only delete a memory "
            "when its ID is known from a previous memory search or save result."
        )

    @property
    def schema(self) -> dict[str, object]:
        return {
            "type": "object",
            "properties": {
                "id": {
                    "type": "string",
                    "description": "The exact UUID of the memory to delete.",
                },
            },
            "required": ["id"],
        }

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

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

        try:
            parsed_id = UUID(memory_id)
        except ValueError:
            raise ToolError("id must be a valid UUID")

        self.store.delete(parsed_id)

        return ToolResult(content=f"Deleted memory: {parsed_id}")