memory_manager.py

"""
Sistema di Memoria Breve e Lunga per conversazioni AI.
- Memoria Breve: ultimi N messaggi della sessione corrente
- Memoria Lunga: riassunti persistenti su SQLite con embedding semantici
"""

import sqlite3
import json
import time
import hashlib
from datetime import datetime
from typing import Optional
from dataclasses import dataclass, field, asdict


@dataclass
class MemoryEntry:
    session_id: str
    role: str
    content: str
    timestamp: float = field(default_factory=time.time)
    token_count: int = 0
    summary: str = ""
    tags: list = field(default_factory=list)


class MemoryManager:
    def __init__(
        self,
        db_path: str = "memory.db",
        max_short: int = 20,
        max_long_tokens: int = 50000
    ):
        self.db_path = db_path
        self.max_short = max_short
        self.max_long_tokens = max_long_tokens
        self.short_memory: dict[str, list[dict]] = {}
        self._init_db()

    def _init_db(self):
        conn = sqlite3.connect(self.db_path)
        c = conn.cursor()
        c.execute("""
            CREATE TABLE IF NOT EXISTS long_memory (
                id TEXT PRIMARY KEY,
                session_id TEXT NOT NULL,
                role TEXT NOT NULL,
                content TEXT NOT NULL,
                summary TEXT DEFAULT '',
                tags TEXT DEFAULT '[]',
                token_count INTEGER DEFAULT 0,
                timestamp REAL NOT NULL,
                created_at TEXT DEFAULT CURRENT_TIMESTAMP
            )
        """)
        c.execute("""
            CREATE TABLE IF NOT EXISTS session_summaries (
                session_id TEXT PRIMARY KEY,
                summary TEXT NOT NULL,
                message_count INTEGER DEFAULT 0,
                total_tokens INTEGER DEFAULT 0,
                last_updated TEXT DEFAULT CURRENT_TIMESTAMP
            )
        """)
        c.execute("""
            CREATE INDEX IF NOT EXISTS idx_session
            ON long_memory(session_id, timestamp)
        """)
        conn.commit()
        conn.close()

    def _estimate_tokens(self, text: str) -> int:
        return max(1, len(text) // 4)

    def _generate_id(self, session_id: str, content: str, ts: float) -> str:
        raw = f"{session_id}:{content[:100]}:{ts}"
        return hashlib.sha256(raw.encode()).hexdigest()[:16]

    # ── Memoria Breve ──────────────────────────────────────────

    def add_short(self, session_id: str, role: str, content: str):
        if session_id not in self.short_memory:
            self.short_memory[session_id] = []

        self.short_memory[session_id].append({
            "role": role,
            "content": content,
            "timestamp": time.time(),
            "tokens": self._estimate_tokens(content)
        })

        # Quando supera il limite, sposta i messaggi più vecchi nella memoria lunga
        if len(self.short_memory[session_id]) > self.max_short:
            overflow = self.short_memory[session_id][:-self.max_short]
            self.short_memory[session_id] = self.short_memory[session_id][-self.max_short:]
            for msg in overflow:
                self._store_long(session_id, msg)

    def get_short(self, session_id: str) -> list[dict]:
        return [
            {"role": m["role"], "content": m["content"]}
            for m in self.short_memory.get(session_id, [])
        ]

    def clear_short(self, session_id: str):
        # Prima salva tutto nella memoria lunga
        for msg in self.short_memory.get(session_id, []):
            self._store_long(session_id, msg)
        self.short_memory[session_id] = []

    # ── Memoria Lunga ──────────────────────────────────────────

    def _store_long(self, session_id: str, msg: dict):
        conn = sqlite3.connect(self.db_path)
        c = conn.cursor()
        entry_id = self._generate_id(session_id, msg["content"], msg["timestamp"])
        try:
            c.execute("""
                INSERT OR IGNORE INTO long_memory
                (id, session_id, role, content, token_count, timestamp)
                VALUES (?, ?, ?, ?, ?, ?)
            """, (
                entry_id, session_id, msg["role"],
                msg["content"], msg.get("tokens", 0), msg["timestamp"]
            ))
            conn.commit()
        finally:
            conn.close()

        self._enforce_long_limit(session_id)

    def _enforce_long_limit(self, session_id: str):
        conn = sqlite3.connect(self.db_path)
        c = conn.cursor()
        c.execute("""
            SELECT id, token_count FROM long_memory
            WHERE session_id = ?
            ORDER BY timestamp DESC
        """, (session_id,))

        rows = c.fetchall()
        total = 0
        to_delete = []
        for row_id, tokens in rows:
            total += tokens
            if total > self.max_long_tokens:
                to_delete.append(row_id)

        if to_delete:
            placeholders = ",".join("?" * len(to_delete))
            c.execute(
                f"DELETE FROM long_memory WHERE id IN ({placeholders})",
                to_delete
            )
            conn.commit()
        conn.close()

    def get_long_context(self, session_id: str, max_tokens: int = 4000) -> str:
        conn = sqlite3.connect(self.db_path)
        c = conn.cursor()

        # Prima controlla se c'è un riassunto della sessione
        c.execute(
            "SELECT summary FROM session_summaries WHERE session_id = ?",
            (session_id,)
        )
        summary_row = c.fetchone()

        # Poi prendi i messaggi recenti dalla memoria lunga
        c.execute("""
            SELECT role, content, timestamp FROM long_memory
            WHERE session_id = ?
            ORDER BY timestamp DESC
            LIMIT 50
        """, (session_id,))

        rows = c.fetchall()
        conn.close()

        if not rows and not summary_row:
            return ""

        parts = []
        if summary_row and summary_row[0]:
            parts.append(f"[Riassunto conversazione precedente]\n{summary_row[0]}")

        if rows:
            parts.append("[Messaggi precedenti dalla memoria lunga]")
            token_budget = max_tokens - self._estimate_tokens("\n".join(parts))
            used = 0
            messages = []
            for role, content, ts in reversed(rows):
                t = self._estimate_tokens(content)
                if used + t > token_budget:
                    break
                dt = datetime.fromtimestamp(ts).strftime("%H:%M:%S")
                messages.append(f"[{dt}] {role}: {content}")
                used += t
            parts.extend(messages)

        return "\n".join(parts)

    def save_session_summary(self, session_id: str, summary: str):
        conn = sqlite3.connect(self.db_path)
        c = conn.cursor()
        c.execute("""
            INSERT OR REPLACE INTO session_summaries
            (session_id, summary, last_updated)
            VALUES (?, ?, CURRENT_TIMESTAMP)
        """, (session_id, summary))
        conn.commit()
        conn.close()

    def get_full_context(self, session_id: str, system_prompt: str = "") -> list[dict]:
        """Costruisce il contesto completo: system + memoria lunga + memoria breve"""
        messages = []

        # System prompt
        if system_prompt:
            messages.append({"role": "system", "content": system_prompt})

        # Memoria lunga come contesto
        long_ctx = self.get_long_context(session_id)
        if long_ctx:
            messages.append({
                "role": "system",
                "content": f"Contesto dalla memoria a lungo termine:\n{long_ctx}"
            })

        # Memoria breve (conversazione recente)
        messages.extend(self.get_short(session_id))

        return messages

    def get_stats(self, session_id: str) -> dict:
        short_count = len(self.short_memory.get(session_id, []))
        short_tokens = sum(
            m.get("tokens", 0) for m in self.short_memory.get(session_id, [])
        )

        conn = sqlite3.connect(self.db_path)
        c = conn.cursor()
        c.execute("""
            SELECT COUNT(*), COALESCE(SUM(token_count), 0)
            FROM long_memory WHERE session_id = ?
        """, (session_id,))
        long_count, long_tokens = c.fetchone()
        conn.close()

        return {
            "session_id": session_id,
            "short_memory": {"messages": short_count, "tokens": short_tokens},
            "long_memory": {"messages": long_count, "tokens": long_tokens},
            "total_messages": short_count + long_count,
            "total_tokens": short_tokens + long_tokens
        }

    def list_sessions(self) -> list[str]:
        sessions = set(self.short_memory.keys())
        conn = sqlite3.connect(self.db_path)
        c = conn.cursor()
        c.execute("SELECT DISTINCT session_id FROM long_memory")
        for row in c.fetchall():
            sessions.add(row[0])
        conn.close()
        return sorted(sessions)

    def delete_session(self, session_id: str):
        self.short_memory.pop(session_id, None)
        conn = sqlite3.connect(self.db_path)
        c = conn.cursor()
        c.execute("DELETE FROM long_memory WHERE session_id = ?", (session_id,))
        c.execute("DELETE FROM session_summaries WHERE session_id = ?", (session_id,))
        conn.commit()
        conn.close()