""" 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