backend/app/services/memory_manager.py

"""
Project Jarvis — Memory Manager
Handles long-term cognitive consolidation and automated memory extraction.
"""

import logging
import json
from datetime import datetime
from typing import Optional, List, Dict

from app.services import llm
from app.core import context

logger = logging.getLogger("friday.memory")

SUMMARIZATION_PROMPT = """
Summarize the following conversation history between FRIDAY and the user.
Extract:
1. Key facts about the user (preferences, names mentioned).
2. Status of ongoing tasks.
3. Summary of the most important context.

Output format: JSON with "summary", "facts", and "tasks".
"""

def consolidate_session(session_id: str):
    """
    Summarizes the session and saves key facts to the persistent memory.
    This should be called periodically or at the end of a long interaction.
    """
    logger.info(f"Starting cognitive consolidation for session: {session_id}")
    
    turns = context.get_conversation_messages(session_id)
    if not turns or len(turns) < 4:
        logger.info("Session too short for consolidation.")
        return

    # 1. Generate Summary & Extract Facts
    history_str = "\n".join([f"{t['role']}: {t['content']}" for t in turns])
    prompt = f"{SUMMARIZATION_PROMPT}\n\n[HISTORY]\n{history_str}"
    
    try:
        raw_result = llm.chat(
            messages=[{"role": "user", "content": prompt}],
            system_prompt="You are a data extraction engine. Respond ONLY with valid JSON.",
            use_tools=False
        )
        
        # Parse JSON from LLM response
        result = _parse_json_flexible(raw_result)
        
        # 2. Save facts to persistent memory
        if result.get("facts"):
            for fact in result["facts"]:
                # In a real app, this would call the memory service (Supabase/Chroma)
                # For now, we store in session metadata
                entities = context.get_last_entities(session_id)
                entities.update(fact if isinstance(fact, dict) else {"extracted_fact": fact})
                context.set_metadata(session_id, "entities", entities)
                logger.info(f"✓ Extracted and saved fact: {fact}")

        logger.info(f"✓ Consolidation complete. Summary: {result.get('summary', 'No summary generated.')}")

    except Exception as e:
        logger.error(f"Memory consolidation failed: {e}")

def _parse_json_flexible(text: str) -> dict:
    """Safely parse JSON from LLM output even if wrapped in markdown blocks."""
    text = text.strip()
    if text.startswith("```"):
        lines = text.split("\n")
        if lines[0].startswith("```"):
            lines = lines[1:]
        if lines[-1].strip() == "```":
            lines = lines[:-1]
        text = "\n".join(lines).strip()
    
    try:
        return json.loads(text)
    except json.JSONDecodeError:
        logger.warning("Failed to parse JSON exactly, attempting fallback.")
        return {"summary": text, "facts": [], "tasks": []}