src/copilot/agents/memory.py

"""
MemoryAgent - Conversation context management
Stores and retrieves conversation history
Resolves coreferences (this account, that alert, etc)
"""

import json
import re
import sqlite3
from typing import Dict, List
from pathlib import Path
from .base import Agent, AgentConfig, AgentResult
import time


class MemoryAgent(Agent):
    """
    Manages conversation memory using SQLite
    Maintains context across turns within a conversation
    """

    DB_PATH = Path(__file__).parent.parent.parent.parent / "data" / "copilot_memory.db"

    def __init__(self, api_pool):
        config = AgentConfig(
            name="MemoryAgent",
            model="llama-3.1-8b-instant",
            temperature=0.0,
            max_tokens=500,
            timeout_ms=5000,
        )
        super().__init__(config, api_pool)
        self._init_db()

    def _init_db(self):
        """Initialize SQLite database for conversation memory"""
        self.DB_PATH.parent.mkdir(parents=True, exist_ok=True)

        conn = sqlite3.connect(str(self.DB_PATH))
        cursor = conn.cursor()

        cursor.execute("""
            CREATE TABLE IF NOT EXISTS conversation_turns (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                conversation_id TEXT NOT NULL,
                turn_number INTEGER NOT NULL,
                user_message TEXT NOT NULL,
                assistant_response TEXT,
                intent TEXT,
                extracted_entities TEXT,
                timestamp TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                UNIQUE(conversation_id, turn_number)
            )
        """)

        cursor.execute("""
            CREATE INDEX IF NOT EXISTS idx_conversation_id
            ON conversation_turns(conversation_id)
        """)

        conn.commit()
        conn.close()

    def _build_prompt(self, **inputs) -> str:
        return ""

    def _parse_response(self, response_text: str) -> Dict:
        return {}

    async def invoke(
        self,
        conversation_id: str,
        user_message: str = None,
        limit: int = 5,
        **kwargs,
    ) -> AgentResult:
        """
        Retrieve conversation history and resolve coreferences
        """
        start_time = time.time()

        try:
            # Get conversation history
            history = self._get_conversation_history(conversation_id, limit=limit)

            # Extract current context from history
            current_context = self._extract_current_context(history)

            # Resolve coreferences in current message
            resolved_message = user_message
            if user_message and current_context:
                resolved_message = self._resolve_coreferences(user_message, current_context)

            result_data = {
                "conversation_id": conversation_id,
                "history": history,
                "current_context": current_context,
                "resolved_message": resolved_message,
                "turn_count": len(history),
            }

            self.logger.info(
                f"[OK] {self.config.name}: Retrieved {len(history)} turns "
                f"for {conversation_id}"
            )

            return await self._create_result(
                success=True,
                data=result_data,
                tokens_input=0,
                tokens_output=self._estimate_tokens(str(result_data)),
                start_time=start_time,
            )

        except Exception as e:
            self.logger.error(f"[FAIL] {self.config.name}: {e}")
            return await self._create_result(
                success=False,
                data={"history": [], "current_context": {}},
                error=str(e),
                start_time=start_time,
            )

    def _get_conversation_history(self, conversation_id: str, limit: int = 5) -> List[Dict]:
        """Fetch recent turns for a conversation"""
        conn = sqlite3.connect(str(self.DB_PATH))
        conn.row_factory = sqlite3.Row
        cursor = conn.cursor()

        cursor.execute("""
            SELECT * FROM conversation_turns
            WHERE conversation_id = ?
            ORDER BY turn_number DESC
            LIMIT ?
        """, (conversation_id, limit))

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

        turns = []
        for row in rows:
            turn = dict(row)
            # Parse extracted entities JSON
            if turn.get("extracted_entities"):
                try:
                    turn["extracted_entities"] = json.loads(turn["extracted_entities"])
                except json.JSONDecodeError:
                    turn["extracted_entities"] = {}
            turns.append(turn)

        # Reverse to chronological order
        turns.reverse()
        return turns

    def _extract_current_context(self, history: List[Dict]) -> Dict:
        """Extract current context from conversation history"""
        context = {
            "current_account_id": None,
            "current_alert_id": None,
            "current_typology": None,
            "last_intent": None,
            "topics_discussed": [],
        }

        for turn in history:
            entities = turn.get("extracted_entities", {})
            if isinstance(entities, dict):
                # Update with most recent values
                if entities.get("account_id"):
                    context["current_account_id"] = entities["account_id"]
                if entities.get("typologies"):
                    context["current_typology"] = entities["typologies"][0] if entities["typologies"] else None

                context["last_intent"] = turn.get("intent")

                # Track topics
                if turn.get("intent"):
                    context["topics_discussed"].append(turn["intent"])

        return context

    def _resolve_coreferences(self, message: str, context: Dict) -> str:
        """Resolve pronouns and references in message"""
        if not message:
            return message

        resolved = message
        message_lower = message.lower()

        # Coreferences to resolve
        replacements = {
            "this account": context.get("current_account_id"),
            "that account": context.get("current_account_id"),
            "the account": context.get("current_account_id"),
            "this alert": context.get("current_alert_id"),
            "that alert": context.get("current_alert_id"),
            "this typology": context.get("current_typology"),
            "that typology": context.get("current_typology"),
        }

        for pronoun, entity in replacements.items():
            if entity and pronoun in message_lower:
                # Case-insensitive replacement preserving original
                pattern = re.compile(re.escape(pronoun), re.IGNORECASE)
                resolved = pattern.sub(f"{pronoun} ({entity})", resolved)

        return resolved

    def store_turn(
        self,
        conversation_id: str,
        user_message: str,
        assistant_response: str = "",
        intent: str = "GENERAL",
        extracted_entities: Dict = None,
    ) -> int:
        """Store a conversation turn in the database"""
        conn = sqlite3.connect(str(self.DB_PATH))
        cursor = conn.cursor()

        # Get next turn number
        cursor.execute(
            "SELECT COALESCE(MAX(turn_number), 0) + 1 FROM conversation_turns WHERE conversation_id = ?",
            (conversation_id,)
        )
        turn_number = cursor.fetchone()[0]

        # Insert turn
        cursor.execute("""
            INSERT INTO conversation_turns
            (conversation_id, turn_number, user_message, assistant_response, intent, extracted_entities)
            VALUES (?, ?, ?, ?, ?, ?)
        """, (
            conversation_id,
            turn_number,
            user_message,
            assistant_response,
            intent,
            json.dumps(extracted_entities or {}),
        ))

        conn.commit()
        conn.close()

        return turn_number