tools/memory.py

"""
Memory & Knowledge tools for persistent assistant learning.
"""

import time
import logging
from datetime import datetime, timezone
from typing import Optional
from pydantic import BaseModel, Field

from .base import register_tool
from .utils import format_created_at
from ..redis_client import redis_client

logger = logging.getLogger(__name__)

@register_tool(
    name="save_memory",
    description="Ghi nhớ một thông tin, sự thật hoặc ghi chú quan trọng về người dùng hoặc dự án để sử dụng sau này.",
    parameters=[
        {"name": "content", "type": "string", "description": "Nội dung kiến thức/thông tin cần ghi nhớ.", "required": True},
        {"name": "category", "type": "string", "description": "Phân loại thông tin (VD: 'personal', 'work', 'preference').", "required": False}
    ]
)
def tool_save_memory(content: str, category: str = "general") -> dict:
    """
    Saves a memory to Redis.
    """
    try:
        mem_id = f"mem_{int(time.time() * 1000)}"
        memory_data = {
            "id": mem_id,
            "content": content,
            "category": category,
            "timestamp": int(time.time() * 1000)
        }
        
        if redis_client.save_memory(memory_data):
            return {
                "status": "success",
                "message": f"Đã ghi nhớ: '{content}' vào danh mục '{category}'",
                "memory_id": mem_id
            }
        else:
            return {"status": "error", "message": "Không thể lưu vào Redis."}
            
    except Exception as e:
        logger.error(f"Error in save_memory: {e}")
        return {"status": "error", "message": str(e)}

@register_tool(
    name="get_memories",
    description="Tìm kiếm và truy xuất các thông tin đã ghi nhớ trước đây dựa trên từ khóa.",
    parameters=[
        {"name": "query", "type": "string", "description": "TỪ KHÓA DUY NHẤT (VD: 'azure', 'ghét'). Để trống nếu muốn lấy toàn bộ 100 ghi nhớ mới nhất.", "required": False},
        {"name": "limit", "type": "integer", "description": "Số lượng kết quả tối đa.", "required": False}
    ]
)
def tool_get_memories(query: str = None, limit: int = 100) -> dict:
    """
    Retrieves memories from Redis.
    """
    try:
        memories = redis_client.list_memories(query=query, limit=limit)
        for mem in memories:
            if "timestamp" in mem:
                try:
                    iso = datetime.fromtimestamp(int(mem["timestamp"]) / 1000, tz=timezone.utc).isoformat()
                except (ValueError, TypeError):
                    iso = str(mem["timestamp"])
                mem["timestamp"] = format_created_at(iso)
        return {
            "status": "success",
            "count": len(memories),
            "memories": memories
        }
    except Exception as e:
        logger.error(f"Error in get_memories: {e}")
        return {"status": "error", "message": str(e)}