rp/pr/tools/memory.py

import os
from typing import Dict, Any, List
from pr.memory.knowledge_store import KnowledgeStore, KnowledgeEntry
import time
import uuid

def add_knowledge_entry(category: str, content: str, metadata: Dict[str, Any] = None, entry_id: str = None) -> Dict[str, Any]:
    """Add a new entry to the knowledge base."""
    try:
        db_path = os.path.expanduser('~/.assistant_db.sqlite')
        store = KnowledgeStore(db_path)
        
        if entry_id is None:
            entry_id = str(uuid.uuid4())[:16]
        
        entry = KnowledgeEntry(
            entry_id=entry_id,
            category=category,
            content=content,
            metadata=metadata or {},
            created_at=time.time(),
            updated_at=time.time()
        )
        
        store.add_entry(entry)
        return {"status": "success", "entry_id": entry_id}
    except Exception as e:
        return {"status": "error", "error": str(e)}

def get_knowledge_entry(entry_id: str) -> Dict[str, Any]:
    """Retrieve a knowledge entry by ID."""
    try:
        db_path = os.path.expanduser('~/.assistant_db.sqlite')
        store = KnowledgeStore(db_path)
        
        entry = store.get_entry(entry_id)
        if entry:
            return {"status": "success", "entry": entry.to_dict()}
        else:
            return {"status": "not_found", "entry_id": entry_id}
    except Exception as e:
        return {"status": "error", "error": str(e)}

def search_knowledge(query: str, category: str = None, top_k: int = 5) -> Dict[str, Any]:
    """Search the knowledge base semantically."""
    try:
        db_path = os.path.expanduser('~/.assistant_db.sqlite')
        store = KnowledgeStore(db_path)
        
        entries = store.search_entries(query, category, top_k)
        results = [entry.to_dict() for entry in entries]
        return {"status": "success", "results": results}
    except Exception as e:
        return {"status": "error", "error": str(e)}

def get_knowledge_by_category(category: str, limit: int = 20) -> Dict[str, Any]:
    """Get knowledge entries by category."""
    try:
        db_path = os.path.expanduser('~/.assistant_db.sqlite')
        store = KnowledgeStore(db_path)
        
        entries = store.get_by_category(category, limit)
        results = [entry.to_dict() for entry in entries]
        return {"status": "success", "entries": results}
    except Exception as e:
        return {"status": "error", "error": str(e)}

def update_knowledge_importance(entry_id: str, importance_score: float) -> Dict[str, Any]:
    """Update the importance score of a knowledge entry."""
    try:
        db_path = os.path.expanduser('~/.assistant_db.sqlite')
        store = KnowledgeStore(db_path)
        
        store.update_importance(entry_id, importance_score)
        return {"status": "success", "entry_id": entry_id, "importance_score": importance_score}
    except Exception as e:
        return {"status": "error", "error": str(e)}

def delete_knowledge_entry(entry_id: str) -> Dict[str, Any]:
    """Delete a knowledge entry."""
    try:
        db_path = os.path.expanduser('~/.assistant_db.sqlite')
        store = KnowledgeStore(db_path)
        
        success = store.delete_entry(entry_id)
        return {"status": "success" if success else "not_found", "entry_id": entry_id}
    except Exception as e:
        return {"status": "error", "error": str(e)}

def get_knowledge_statistics() -> Dict[str, Any]:
    """Get statistics about the knowledge base."""
    try:
        db_path = os.path.expanduser('~/.assistant_db.sqlite')
        store = KnowledgeStore(db_path)
        
        stats = store.get_statistics()
        return {"status": "success", "statistics": stats}
    except Exception as e:
        return {"status": "error", "error": str(e)}