agent/agent_memory.py

"""
NeuraPrompt Agent — Memory & Skills Module (v1.0)
==================================================

Separate database (neuraprompt_agent) — does NOT mix with chat history or Kype.

Collections:
  neuraprompt_agent.conversations  — last 5 messages per user
  neuraprompt_agent.user_profiles  — user facts (name, timezone, preferences)
  neuraprompt_agent.skills         — learned skills, per-user

Design principles:
  * Memory is capped at 5 messages (not 10) to keep the system prompt small.
  * Skills are NOT loaded into the system prompt automatically. The AI decides
    whether to check skills via a tool-like mechanism (the memory_context
    string includes a SKILLS AVAILABLE indicator, and the AI can request them).
  * Both memory and skills trigger UI indicators:
    - "👾 New skill learned" when a skill is saved
    - "👾 Memory recalled" when memory was used in the response

Env vars:
  MONGO_URI — MongoDB connection string (reuses main.py's if available)
"""

from __future__ import annotations
import os
import re
import time
import logging
from datetime import datetime, timezone
from typing import Optional, Dict, Any, List

log = logging.getLogger("agent.memory")

# ── MONGODB (reuses main.py's connection, separate database) ──────────────────
_db = None

def _get_db():
    """Get the neuraprompt_agent database. Reuses main.py's MongoClient."""
    global _db
    if _db is not None:
        return _db
    try:
        import sys
        main_mod = sys.modules.get("main")
        if main_mod and hasattr(main_mod, "mongo_client"):
            _db = main_mod.mongo_client["neuraprompt_agent"]
            return _db
    except Exception:
        pass
    # Fallback: own connection
    try:
        from pymongo import MongoClient
        from pymongo.server_api import ServerApi
        MONGO_URL = os.environ.get("MONGO_URI", "")
        client = MongoClient(
            MONGO_URL, ssl=True,
            tlsAllowInvalidCertificates=False,
            tlsCAFile="/etc/ssl/certs/ca-certificates.crt",
            server_api=ServerApi("1"),
        )
        _db = client["neuraprompt_agent"]
        return _db
    except Exception as e:
        log.error(f"[AgentMemory] DB connection failed: {e}")
        return None


# ── MEMORY (last 5 messages) ──────────────────────────────────────────────────
MAX_MESSAGES = 5

def load_memory(user_id: str) -> Dict[str, Any]:
    """Load the last 5 messages + user profile for a user.
    Returns {messages: [], profile: {}} — empty if no history."""
    db = _get_db()
    if db is None or not user_id:
        return {"messages": [], "profile": {}}

    result = {"messages": [], "profile": {}}
    try:
        # Load conversation (last 5 messages)
        conv = db.conversations.find_one({"user_id": user_id})
        if conv and conv.get("messages"):
            result["messages"] = conv["messages"][-MAX_MESSAGES:]

        # Load user profile (facts)
        profile = db.user_profiles.find_one({"user_id": user_id})
        if profile:
            result["profile"] = {
                k: v for k, v in profile.items()
                if k not in ("_id", "user_id", "created_at", "last_updated")
            }
    except Exception as e:
        log.warning(f"[AgentMemory] load_memory failed: {e}")

    return result

def save_memory(user_id: str, user_message: str, agent_response: str):
    """Save a user message + agent response to conversation history.
    Keeps only the last 5 messages (trims older ones)."""
    db = _get_db()
    if db is None or not user_id:
        return

    try:
        now = datetime.now(timezone.utc)
        new_msgs = [
            {"role": "user", "content": user_message[:500], "timestamp": now.isoformat()},
            {"role": "assistant", "content": agent_response[:500], "timestamp": now.isoformat()},
        ]

        # Atomically push new messages and trim to last 5
        existing = db.conversations.find_one({"user_id": user_id})
        if existing:
            all_msgs = existing.get("messages", []) + new_msgs
            # Keep last MAX_MESSAGES
            trimmed = all_msgs[-MAX_MESSAGES:]
            db.conversations.update_one(
                {"user_id": user_id},
                {"$set": {"messages": trimmed, "last_updated": now.isoformat()}},
            )
        else:
            db.conversations.insert_one({
                "user_id": user_id,
                "messages": new_msgs,
                "created_at": now.isoformat(),
                "last_updated": now.isoformat(),
            })
    except Exception as e:
        log.warning(f"[AgentMemory] save_memory failed: {e}")

def update_user_profile(user_id: str, facts: Dict[str, Any]):
    """Update user profile with learned facts (name, timezone, preferences)."""
    db = _get_db()
    if db is None or not user_id or not facts:
        return
    try:
        now = datetime.now(timezone.utc)
        db.user_profiles.update_one(
            {"user_id": user_id},
            {"$set": {**facts, "last_updated": now.isoformat()}},
            upsert=True,
        )
    except Exception as e:
        log.warning(f"[AgentMemory] update_user_profile failed: {e}")


# ── SKILLS (per-user, AI-decided loading) ─────────────────────────────────────
MAX_SKILLS = 20

def get_skills_list(user_id: str) -> List[str]:
    """Get a list of skill NAMES (not content) for the system prompt.
    The AI sees the names and can decide whether to request details."""
    db = _get_db()
    if db is None or not user_id:
        return []
    try:
        skills = list(db.skills.find({"user_id": user_id}, {"skill_name": 1, "_id": 0}))
        return [s["skill_name"] for s in skills[:MAX_SKILLS]]
    except Exception as e:
        log.warning(f"[AgentMemory] get_skills_list failed: {e}")
        return []

def get_skill_detail(user_id: str, skill_name: str) -> Optional[str]:
    """Get the full body/description of a specific skill (when AI requests it).
    Prefers skill_body, then when_to_use+description, then skill_description."""
    db = _get_db()
    if db is None or not user_id:
        return None
    try:
        skill = db.skills.find_one({"user_id": user_id, "skill_name": skill_name})
        if not skill:
            return None
        body = (skill.get("skill_body") or "").strip()
        if body:
            when = (skill.get("when_to_use") or "").strip()
            if when:
                return f"**When to use:** {when}\n\n{body}"
            return body
        desc = skill.get("skill_description", "") or ""
        when = (skill.get("when_to_use") or "").strip()
        if when and desc:
            return f"**When to use:** {when}\n\n{desc}"
        return desc or when or None
    except Exception as e:
        log.warning(f"[AgentMemory] get_skill_detail failed: {e}")
    return None

def save_skill(
    user_id: str,
    skill_name: str,
    skill_description: str,
    when_to_use: str = "",
    skill_body: str = "",
) -> bool:
    """Save a new skill. Returns True if it's actually new (not a duplicate).
    Supports richer fields: when_to_use + skill_body (falls back to skill_description)."""
    db = _get_db()
    if db is None or not user_id or not skill_name:
        return False

    try:
        now = datetime.now(timezone.utc).isoformat()
        fields = {
            "skill_description": skill_description,
            "when_to_use": when_to_use or "",
            "skill_body": skill_body or skill_description or "",
            "updated_at": now,
        }
        # Check if this skill already exists (by name)
        existing = db.skills.find_one({"user_id": user_id, "skill_name": skill_name})
        if existing:
            changed = any(existing.get(k) != v for k, v in fields.items() if k != "updated_at")
            if changed:
                db.skills.update_one(
                    {"user_id": user_id, "skill_name": skill_name},
                    {"$set": fields},
                )
            return False  # Not new

        # Check skill count — if at max, remove oldest
        count = db.skills.count_documents({"user_id": user_id})
        if count >= MAX_SKILLS:
            oldest = db.skills.find_one({"user_id": user_id}, sort=[("learned_at", 1)])
            if oldest:
                db.skills.delete_one({"_id": oldest["_id"]})

        # Insert new skill
        db.skills.insert_one({
            "user_id": user_id,
            "skill_name": skill_name,
            "skill_description": skill_description,
            "when_to_use": when_to_use or "",
            "skill_body": skill_body or skill_description or "",
            "learned_at": now,
        })
        log.info(f"[AgentMemory] New skill saved for {user_id}: {skill_name}")
        return True  # New skill
    except Exception as e:
        log.warning(f"[AgentMemory] save_skill failed: {e}")
        return False


def match_skills_for_query(user_id: str, query: str, limit: int = 2) -> List[Dict[str, Any]]:
    """Match builtins + user skills against a query.
    Returns list of dicts: {name, when_to_use, body, source}."""
    results: List[Dict[str, Any]] = []
    try:
        from .skills_catalog import match_builtins
    except ImportError:
        try:
            from skills_catalog import match_builtins  # type: ignore
        except ImportError:
            match_builtins = lambda q, limit=2: []  # type: ignore

    for sk in match_builtins(query, limit=limit):
        results.append({
            "name": sk["name"],
            "when_to_use": sk.get("when_to_use", ""),
            "body": sk.get("body", ""),
            "source": "builtin",
        })

    if user_id and query:
        db = _get_db()
        if db is not None:
            try:
                q_lower = query.lower()
                q_tokens = set(re.findall(r"[a-z0-9_]{3,}", q_lower))
                scored = []
                for doc in db.skills.find({"user_id": user_id}):
                    name = (doc.get("skill_name") or "").lower()
                    when = (doc.get("when_to_use") or doc.get("skill_description") or "").lower()
                    score = 0
                    if name and (name in q_lower or name.replace("_", " ") in q_lower):
                        score += 6
                    when_tokens = set(re.findall(r"[a-z0-9_]{3,}", when))
                    score += min(8, len(when_tokens & q_tokens) * 2)
                    if score >= 4:
                        body = doc.get("skill_body") or doc.get("skill_description") or ""
                        scored.append((score, {
                            "name": doc.get("skill_name", ""),
                            "when_to_use": doc.get("when_to_use") or "",
                            "body": body,
                            "source": "user",
                        }))
                scored.sort(key=lambda x: -x[0])
                for _, item in scored[:limit]:
                    # Avoid duplicate names already filled by builtins
                    if not any(r["name"] == item["name"] for r in results):
                        results.append(item)
            except Exception as e:
                log.warning(f"[AgentMemory] match_skills_for_query failed: {e}")

    return results[: max(limit, 2)]


# ── SKILL EXTRACTION (lightweight, post-response) ─────────────────────────────
# Patterns that indicate a learnable preference
SKILL_PATTERNS = [
    # Language/framework preferences
    (r'(?:i prefer|i like|i always use|i use)\s+(python|javascript|typescript|react|vue|angular|django|flask|fastapi|tailwind|bootstrap|node\.?js|rust|go\b|java\b|c\+\+)',
     "prefers_{0}", "User prefers using {0}"),
    # Code style
    (r'(?:i prefer|i like).{0,20}(?:dark mode|light mode|tabs|spaces|camelcase|snake_case)',
     "code_style", "User mentioned code style preference"),
    # Timezone
    (r'(?:my timezone|i.?m in|my time zone).{0,30}(africa|europe|america|asia|[A-Z]{2,4}[-+]?\d{0,2})',
     "timezone", "User timezone info"),
    # Project context
    (r'(?:my project|my app|my website|i.?m building|i.?m working on).{0,50}',
     "project_context", "User mentioned a project they're working on"),
]

def extract_skills(user_message: str, agent_response: str, user_id: str) -> List[Dict[str, str]]:
    """Extract potential skills from the conversation. Returns list of {name, description}.
    Only extracts from the USER's message, not the agent's response."""
    skills_found = []
    for pattern, name_template, desc_template in SKILL_PATTERNS:
        matches = re.findall(pattern, user_message, re.IGNORECASE)
        for match in matches:
            match_str = match if isinstance(match, str) else match[0] if match else ""
            if match_str:
                skill_name = name_template.format(match_str.lower().replace(" ", "_").replace(".", ""))
                skill_desc = desc_template.format(match_str)
                is_new = save_skill(user_id, skill_name, skill_desc)
                if is_new:
                    skills_found.append({"name": skill_name, "description": skill_desc})
    return skills_found


# ── CONTEXT BUILDER (builds the memory_context string for the system prompt) ──
def build_memory_context(user_id: str, query: str = "") -> str:
    """Build a compact memory context string for the system prompt.
    Includes last 5 messages (short) + skill names.
    If query is provided, also injects up to 2 auto-matched Active skill bodies."""
    mem = load_memory(user_id)
    parts = []

    # User profile (very compact)
    profile = mem.get("profile", {})
    if profile:
        facts = [f"{k}: {v}" for k, v in profile.items() if k != "messages"]
        if facts:
            parts.append("User profile: " + ", ".join(facts[:5]))

    # Recent messages (very compact — just role + first 100 chars)
    messages = mem.get("messages", [])
    if messages:
        recent = []
        for msg in messages[-MAX_MESSAGES:]:
            role = msg.get("role", "?")
            content = msg.get("content", "")[:100]
            recent.append(f"{role}: {content}")
        parts.append("Recent conversation:\n" + "\n".join(recent))

    # Skills (names only — AI decides if it needs details)
    skills = get_skills_list(user_id)
    # Also surface builtin names so the model knows they exist
    try:
        from .skills_catalog import BUILTIN_SKILLS
    except ImportError:
        try:
            from skills_catalog import BUILTIN_SKILLS  # type: ignore
        except ImportError:
            BUILTIN_SKILLS = []
    builtin_names = [s["name"] for s in BUILTIN_SKILLS]
    name_bits = []
    if builtin_names:
        name_bits.append("builtins: " + ", ".join(builtin_names))
    if skills:
        name_bits.append("learned: " + ", ".join(skills))
    if name_bits:
        parts.append("Learned skills (available if needed): " + " | ".join(name_bits))

    # Active skills — auto-inject up to 2 matched full bodies for the current query
    if query:
        matched = match_skills_for_query(user_id, query, limit=2)
        if matched:
            active = ["Active skills (follow these):"]
            for m in matched[:2]:
                active.append(f"### {m['name']}\n{m.get('body') or m.get('when_to_use') or ''}")
            parts.append("\n".join(active))

    if not parts:
        return ""

    return "\n".join(parts)