path: root/personal_agent/tools/memory.py

from datetime import UTC, datetime
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 "working_memory_add_fact"

    @property
    def description(self) -> str:
        return (
            "Add a fact to the current task's working memory. Working memory is "
            "temporary and is used to track information needed during the current "
            "task or conversation. Do not use this for durable user preferences "
            "or facts that should survive the conversation."
        )

    @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 "working_memory_add_note"

    @property
    def description(self) -> str:
        return (
            "Add a temporary note to the current task's working memory. Use this "
            "for information that may help with the current task or conversation. "
            "Do not use this for durable information that should survive the "
            "conversation."
        )

    @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 "working_memory_remove_fact"

    @property
    def description(self) -> str:
        return (
            "Remove a fact from the current task's working memory. This only "
            "changes temporary working memory; it does not delete persistent "
            "long-term memory."
        )

    @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 "working_memory_remove_note"

    @property
    def description(self) -> str:
        return (
            "Remove a note from the current task's working memory. This only "
            "changes temporary working memory; it does not delete persistent "
            "long-term memory."
        )

    @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 "long_term_memory_save"

    @property
    def description(self) -> str:
        return (
            "Save persistent information for use in future conversations. Use "
            "this for durable user preferences, facts, decisions, goals, "
            "relationships, or meaningful past experiences. This is long-term "
            "memory, not temporary working memory."
        )

    @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."
                    ),
                },
                "importance": {
                    "type": "integer",
                    "enum": [1, 2, 3],
                    "description": (
                        "How useful this memory is likely to be in future "
                        "conversations: 1 = useful but minor, "
                        "2 = normally useful, "
                        "3 = especially important for future behavior. "
                        "Importance measure future usefullness, not emotional intensity."
                    ),
                },
            },
            "required": ["content", "category", "importance"],
        }

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

        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")

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

        if importance not in (1, 2, 3):
            raise ToolError("importance must be between 1 and 3")

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

        memory = LongTermMemory.create(content, memory_category, importance)
        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 "long_term_memory_search"

    @property
    def description(self) -> str:
        return (
            "Search persistent long-term memories from previous conversations. "
            "Use this when information from past conversations may be relevant, "
            "when the user asks what you remember, or when checking whether a "
            "durable fact or preference is already stored. This does not search "
            "working memory."
        )

    @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.")

        now = datetime.now(UTC)

        lines = [
            (
                f"- [{memory.category.value}] {memory.id}\n"
                f"  content: {memory.content}\n"
                f"  created_at: {memory.created_at.isoformat()}\n"
                f"  importance: {memory.importance}"
                f"  stale: {(now - memory.created_at).days >= 365}"
            )
            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 "long_term_memory_delete"

    @property
    def description(self) -> str:
        return (
            "Delete a specific persistent long-term memory by its ID. This "
            "deletes long-term memory, not working memory. When replacing or "
            "forgetting an existing memory, search for it first to obtain its ID."
        )

    @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}")