tools/memory_tools.py

"""Memory and Knowledge Management Tools"""
import os
import json
import time
from typing import Dict, List, Optional, Any
from config import Config
from utils import Logger

class MemoryTools:
    """Persistent memory storage and retrieval"""
    
    def __init__(self, memory_file: str = None):
        self.memory_file = memory_file or Config.MEMORY_FILE
        self.logger = Logger()
        self._ensure_memory()
    
    def _ensure_memory(self):
        """Ensure memory file exists"""
        if not os.path.exists(self.memory_file):
            with open(self.memory_file, "w") as f:
                json.dump({
                    "facts": [],
                    "conversations": [],
                    "tasks": [],
                    "preferences": {},
                    "created_at": time.time()
                }, f)
    
    def _load(self) -> Dict:
        """Load memory from disk"""
        try:
            with open(self.memory_file, "r") as f:
                return json.load(f)
        except:
            return {"facts": [], "conversations": [], "tasks": [], "preferences": {}}
    
    def _save(self, data: Dict):
        """Save memory to disk"""
        with open(self.memory_file, "w") as f:
            json.dump(data, f, indent=2, ensure_ascii=False)
    
    def memory_save(self, category: str, key: str, value: Any,
                    tags: List[str] = None, ttl: int = 0) -> Dict:
        """Save information to memory"""
        try:
            memory = self._load()
            
            entry = {
                "key": key,
                "value": value,
                "tags": tags or [],
                "created": time.time(),
                "ttl": ttl,
                "accessed": 0
            }
            
            if category == "fact":
                memory["facts"] = [f for f in memory["facts"] if f["key"] != key]
                memory["facts"].append(entry)
            elif category == "task":
                memory["tasks"].append(entry)
            elif category == "preference":
                memory["preferences"][key] = value
            elif category == "conversation":
                memory["conversations"].append({
                    **entry,
                    "role": value.get("role", "unknown"),
                    "content": value.get("content", "")
                })
            
            self._save(memory)
            self.logger.log("memory_save", {"category": category, "key": key})
            
            return {
                "success": True,
                "category": category,
                "key": key,
                "stored": True
            }
            
        except Exception as e:
            return {"success": False, "error": str(e)}
    
    def memory_search(self, query: str = "", category: str = "",
                      tags: List[str] = None, limit: int = 10,
                      time_range: str = "") -> Dict:
        """Search memory"""
        try:
            memory = self._load()
            results = []
            
            cutoff = 0
            if time_range == "today":
                cutoff = time.time() - 86400
            elif time_range == "week":
                cutoff = time.time() - 604800
            elif time_range == "month":
                cutoff = time.time() - 2592000
            
            categories = [category] if category else ["facts", "tasks", "conversations"]
            
            for cat in categories:
                items = memory.get(cat, [])
                for item in items:
                    if cutoff and item.get("created", 0) < cutoff:
                        continue
                    
                    if tags and not any(t in item.get("tags", []) for t in tags):
                        continue
                    
                    score = 0
                    if query:
                        query_lower = query.lower()
                        if query_lower in item.get("key", "").lower():
                            score += 10
                        if query_lower in str(item.get("value", "")).lower():
                            score += 5
                        for tag in item.get("tags", []):
                            if query_lower in tag.lower():
                                score += 3
                    
                    if not query or score > 0:
                        results.append({
                            "category": cat,
                            "score": score,
                            **item
                        })
            
            if query:
                results.sort(key=lambda x: x["score"], reverse=True)
            else:
                results.sort(key=lambda x: x.get("created", 0), reverse=True)
            
            for r in results[:limit]:
                r["accessed"] = r.get("accessed", 0) + 1
            
            self._save(memory)
            self.logger.log("memory_search", {"query": query, "results": len(results)})
            
            return {
                "success": True,
                "query": query,
                "results": results[:limit],
                "total": len(results)
            }
            
        except Exception as e:
            return {"success": False, "error": str(e)}
    
    def memory_delete(self, key: str = "", category: str = "",
                      older_than: int = 0, confirm: bool = True) -> Dict:
        """Delete memory entries"""
        try:
            if confirm and not key.startswith("_auto_"):
                from utils import UI
                if not UI.confirm(f"Delete memory entries matching '{key}'?"):
                    return {"success": False, "error": "User cancelled"}
            
            memory = self._load()
            deleted = 0
            
            if key:
                for cat in memory:
                    if isinstance(memory[cat], list):
                        original = len(memory[cat])
                        memory[cat] = [i for i in memory[cat] if i.get("key") != key]
                        deleted += original - len(memory[cat])
            
            if older_than:
                cutoff = time.time() - older_than
                for cat in memory:
                    if isinstance(memory[cat], list):
                        original = len(memory[cat])
                        memory[cat] = [i for i in memory[cat] if i.get("created", 0) > cutoff]
                        deleted += original - len(memory[cat])
            
            self._save(memory)
            
            return {
                "success": True,
                "deleted": deleted,
                "criteria": {"key": key, "category": category, "older_than": older_than}
            }
            
        except Exception as e:
            return {"success": False, "error": str(e)}
    
    def memory_stats(self) -> Dict:
        """Get memory statistics"""
        try:
            memory = self._load()
            return {
                "success": True,
                "facts": len(memory.get("facts", [])),
                "tasks": len(memory.get("tasks", [])),
                "conversations": len(memory.get("conversations", [])),
                "preferences": len(memory.get("preferences", {})),
                "file_size": os.path.getsize(self.memory_file) if os.path.exists(self.memory_file) else 0
            }
        except Exception as e:
            return {"success": False, "error": str(e)}