rp/pr/tools/memory.py

import os
import time
import uuid
from typing import Any, Dict

from pr.memory.knowledge_store import KnowledgeEntry, KnowledgeStore


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