import os import sys import time import sqlite3 import logging from pathlib import Path from typing import Optional, Dict, Any, List, Tuple from dataclasses import dataclass, field, asdict from config import settings logger = logging.getLogger("AGYMemoryManager") DATA_DIR = Path("/root/telegram-agy-bot/data") DB_PATH = DATA_DIR / "memory.db" @dataclass class MemoryItem: id: int type: str # 'global', 'user', 'project' user_id: Optional[int] = None # None for global, chat_id for user/project project_name: Optional[str] = None # None for global/user, project_name for project key: str = "" category: str = "general" # 'rule', 'preference', 'tech_stack', 'fact', 'workflow', 'general' content: str = "" importance: int = 1 created_by: Optional[int] = None created_at: float = field(default_factory=time.time) updated_at: float = field(default_factory=time.time) def to_dict(self) -> Dict[str, Any]: return asdict(self) @property def is_global(self) -> bool: return self.type == "global" @property def is_user(self) -> bool: return self.type == "user" @property def is_project(self) -> bool: return self.type == "project" class MemoryManager: def __init__(self, db_path: Path = DB_PATH): self.db_path = db_path self._init_db() def _get_connection(self) -> sqlite3.Connection: self.db_path.parent.mkdir(parents=True, exist_ok=True) conn = sqlite3.connect(str(self.db_path), timeout=15.0) conn.row_factory = sqlite3.Row conn.execute("PRAGMA journal_mode=WAL;") conn.execute("PRAGMA synchronous=NORMAL;") return conn def _init_db(self): """Initializes tables and indexes, migrating from old schema if needed.""" try: with self._get_connection() as conn: # Check if old memories table exists and needs migration cur = conn.execute("SELECT name FROM sqlite_master WHERE type='table' AND name='memories';") table_exists = cur.fetchone() is not None if table_exists: # Check columns in existing table cur = conn.execute("PRAGMA table_info(memories);") cols = [row["name"] for row in cur.fetchall()] if "project_name" not in cols: logger.info("Migrating memories table to support 3-tier hierarchy (adding project_name)...") # Migrate table conn.execute(""" CREATE TABLE memories_new ( id INTEGER PRIMARY KEY AUTOINCREMENT, type TEXT NOT NULL CHECK(type IN ('global', 'user', 'project', 'private')), user_id INTEGER, project_name TEXT, key TEXT NOT NULL, category TEXT NOT NULL DEFAULT 'general', content TEXT NOT NULL, importance INTEGER NOT NULL DEFAULT 1, created_by INTEGER, created_at REAL NOT NULL, updated_at REAL NOT NULL ); """) conn.execute(""" INSERT INTO memories_new (id, type, user_id, project_name, key, category, content, importance, created_by, created_at, updated_at) SELECT id, CASE WHEN type='private' THEN 'user' ELSE type END, user_id, NULL, key, category, content, importance, created_by, created_at, updated_at FROM memories; """) conn.execute("DROP TABLE memories;") conn.execute("ALTER TABLE memories_new RENAME TO memories;") else: conn.execute(""" CREATE TABLE memories ( id INTEGER PRIMARY KEY AUTOINCREMENT, type TEXT NOT NULL CHECK(type IN ('global', 'user', 'project', 'private')), user_id INTEGER, project_name TEXT, key TEXT NOT NULL, category TEXT NOT NULL DEFAULT 'general', content TEXT NOT NULL, importance INTEGER NOT NULL DEFAULT 1, created_by INTEGER, created_at REAL NOT NULL, updated_at REAL NOT NULL ); """) # Create partial unique indexes for each tier conn.execute(""" CREATE UNIQUE INDEX IF NOT EXISTS idx_memories_global_key ON memories(type, key) WHERE type = 'global'; """) conn.execute(""" CREATE UNIQUE INDEX IF NOT EXISTS idx_memories_user_key ON memories(type, user_id, key) WHERE type IN ('user', 'private'); """) conn.execute(""" CREATE UNIQUE INDEX IF NOT EXISTS idx_memories_project_key ON memories(type, user_id, project_name, key) WHERE type = 'project'; """) conn.execute(""" CREATE INDEX IF NOT EXISTS idx_memories_lookup ON memories(type, user_id, project_name); """) conn.execute(""" CREATE INDEX IF NOT EXISTS idx_memories_category ON memories(category); """) logger.info("3-Tier Memory database initialized successfully at %s", self.db_path) except Exception as e: logger.error("Failed to initialize memory database: %s", e) raise def save_or_update( self, type_: str, key: str, content: str, user_id: Optional[int] = None, project_name: Optional[str] = None, category: str = "general", importance: int = 1, created_by: Optional[int] = None, ) -> Tuple[MemoryItem, bool]: """ Saves or updates a memory in one of the 3 tiers (global, user, project). Returns: (MemoryItem, is_created: bool) """ type_ = type_.lower().strip() if type_ == "private": type_ = "user" if type_ not in ("global", "user", "project"): raise ValueError("Memory type must be 'global', 'user', or 'project'") if type_ == "global": user_id = None project_name = None elif type_ == "user": if user_id is None: raise ValueError("User memory requires a valid user_id") project_name = None elif type_ == "project": if user_id is None or not project_name: raise ValueError("Project memory requires both user_id and project_name") project_name = project_name.strip() key = key.strip().lower() if not key: raise ValueError("Memory key cannot be empty") content = content.strip() if not content: raise ValueError("Memory content cannot be empty") category = category.strip().lower() or "general" importance = max(1, min(5, int(importance))) now = time.time() with self._get_connection() as conn: # Check if existing item exists if type_ == "global": cur = conn.execute( "SELECT id, created_at, created_by FROM memories WHERE type = 'global' AND key = ?", (key,), ) elif type_ == "user": cur = conn.execute( "SELECT id, created_at, created_by FROM memories WHERE type IN ('user', 'private') AND user_id = ? AND key = ?", (user_id, key), ) else: # project cur = conn.execute( "SELECT id, created_at, created_by FROM memories WHERE type = 'project' AND user_id = ? AND project_name = ? AND key = ?", (user_id, project_name, key), ) row = cur.fetchone() if row: mem_id = row["id"] orig_created_at = row["created_at"] orig_created_by = row["created_by"] or created_by conn.execute( """ UPDATE memories SET type = ?, category = ?, content = ?, importance = ?, updated_at = ? WHERE id = ? """, (type_, category, content, importance, now, mem_id), ) is_created = False created_at = orig_created_at creator = orig_created_by else: cur = conn.execute( """ INSERT INTO memories (type, user_id, project_name, key, category, content, importance, created_by, created_at, updated_at) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?) """, (type_, user_id, project_name, key, category, content, importance, created_by, now, now), ) mem_id = cur.lastrowid is_created = True created_at = now creator = created_by item = MemoryItem( id=mem_id, type=type_, user_id=user_id, project_name=project_name, key=key, category=category, content=content, importance=importance, created_by=creator, created_at=created_at, updated_at=now, ) logger.info( "%s memory [id=%d, type=%s, key=%s, user=%s, proj=%s]", "Created" if is_created else "Updated", mem_id, type_, key, user_id, project_name, ) return item, is_created def get_by_id(self, memory_id: int) -> Optional[MemoryItem]: with self._get_connection() as conn: cur = conn.execute("SELECT * FROM memories WHERE id = ?", (memory_id,)) row = cur.fetchone() if row: return MemoryItem(**dict(row)) return None def get_by_key( self, type_: str, key: str, user_id: Optional[int] = None, project_name: Optional[str] = None ) -> Optional[MemoryItem]: type_ = type_.lower().strip() if type_ == "private": type_ = "user" key = key.lower().strip() with self._get_connection() as conn: if type_ == "global": cur = conn.execute("SELECT * FROM memories WHERE type = 'global' AND key = ?", (key,)) elif type_ == "user": cur = conn.execute( "SELECT * FROM memories WHERE type IN ('user', 'private') AND user_id = ? AND key = ?", (user_id, key), ) else: # project cur = conn.execute( "SELECT * FROM memories WHERE type = 'project' AND user_id = ? AND project_name = ? AND key = ?", (user_id, project_name, key), ) row = cur.fetchone() if row: return MemoryItem(**dict(row)) return None def get_global_memories(self, category: Optional[str] = None, limit: int = 100) -> List[MemoryItem]: with self._get_connection() as conn: if category: cur = conn.execute( "SELECT * FROM memories WHERE type = 'global' AND category = ? ORDER BY importance DESC, updated_at DESC LIMIT ?", (category.lower().strip(), limit), ) else: cur = conn.execute( "SELECT * FROM memories WHERE type = 'global' ORDER BY importance DESC, updated_at DESC LIMIT ?", (limit,), ) return [MemoryItem(**dict(r)) for r in cur.fetchall()] def get_user_memories( self, user_id: int, category: Optional[str] = None, limit: int = 100 ) -> List[MemoryItem]: with self._get_connection() as conn: if category: cur = conn.execute( "SELECT * FROM memories WHERE type IN ('user', 'private') AND user_id = ? AND category = ? ORDER BY importance DESC, updated_at DESC LIMIT ?", (user_id, category.lower().strip(), limit), ) else: cur = conn.execute( "SELECT * FROM memories WHERE type IN ('user', 'private') AND user_id = ? ORDER BY importance DESC, updated_at DESC LIMIT ?", (user_id, limit), ) return [MemoryItem(**dict(r)) for r in cur.fetchall()] def get_project_memories( self, user_id: int, project_name: str, category: Optional[str] = None, limit: int = 100 ) -> List[MemoryItem]: with self._get_connection() as conn: if category: cur = conn.execute( "SELECT * FROM memories WHERE type = 'project' AND user_id = ? AND project_name = ? AND category = ? ORDER BY importance DESC, updated_at DESC LIMIT ?", (user_id, project_name, category.lower().strip(), limit), ) else: cur = conn.execute( "SELECT * FROM memories WHERE type = 'project' AND user_id = ? AND project_name = ? ORDER BY importance DESC, updated_at DESC LIMIT ?", (user_id, project_name, limit), ) return [MemoryItem(**dict(r)) for r in cur.fetchall()] def delete_by_id(self, memory_id: int, user_id: Optional[int] = None, is_admin: bool = False) -> bool: with self._get_connection() as conn: cur = conn.execute("SELECT * FROM memories WHERE id = ?", (memory_id,)) row = cur.fetchone() if not row: return False mem = MemoryItem(**dict(row)) if mem.type == "global" and not is_admin: raise PermissionError("Only administrators can delete global memories") if mem.type in ("user", "project", "private") and mem.user_id != user_id and not is_admin: raise PermissionError("Cannot delete another user's memory") conn.execute("DELETE FROM memories WHERE id = ?", (memory_id,)) logger.info("Deleted memory [id=%d, key=%s, type=%s]", memory_id, mem.key, mem.type) return True def delete_by_key( self, type_: str, key: str, user_id: Optional[int] = None, project_name: Optional[str] = None, is_admin: bool = False, ) -> bool: type_ = type_.lower().strip() if type_ == "private": type_ = "user" key = key.lower().strip() with self._get_connection() as conn: if type_ == "global": if not is_admin: raise PermissionError("Only administrators can delete global memories") cur = conn.execute("DELETE FROM memories WHERE type = 'global' AND key = ?", (key,)) elif type_ == "user": if user_id is None: return False cur = conn.execute( "DELETE FROM memories WHERE type IN ('user', 'private') AND user_id = ? AND key = ?", (user_id, key), ) else: # project if user_id is None or not project_name: return False cur = conn.execute( "DELETE FROM memories WHERE type = 'project' AND user_id = ? AND project_name = ? AND key = ?", (user_id, project_name, key), ) deleted = cur.rowcount > 0 if deleted: logger.info("Deleted memory by key [key=%s, type=%s, user=%s, proj=%s]", key, type_, user_id, project_name) return deleted def clear_memories( self, type_: str, user_id: Optional[int] = None, project_name: Optional[str] = None, is_admin: bool = False, ) -> int: type_ = type_.lower().strip() if type_ == "private": type_ = "user" with self._get_connection() as conn: if type_ == "global": if not is_admin: raise PermissionError("Only administrators can clear global memories") cur = conn.execute("DELETE FROM memories WHERE type = 'global'") elif type_ == "user": if user_id is None: return 0 cur = conn.execute("DELETE FROM memories WHERE type IN ('user', 'private') AND user_id = ?", (user_id,)) else: # project if user_id is None or not project_name: return 0 cur = conn.execute( "DELETE FROM memories WHERE type = 'project' AND user_id = ? AND project_name = ?", (user_id, project_name), ) count = cur.rowcount logger.info("Cleared %d %s memories for user %s (proj: %s)", count, type_, user_id, project_name) return count def search_memories( self, query: str, user_id: Optional[int] = None, project_name: Optional[str] = None, type_: Optional[str] = None, limit: int = 50, is_admin: bool = False, ) -> List[MemoryItem]: q = f"%{query.strip().lower()}%" with self._get_connection() as conn: if type_ == "global": if not is_admin: return [] cur = conn.execute( """ SELECT * FROM memories WHERE type = 'global' AND (key LIKE ? OR content LIKE ? OR category LIKE ?) ORDER BY importance DESC, updated_at DESC LIMIT ? """, (q, q, q, limit), ) elif type_ == "user" and user_id is not None: cur = conn.execute( """ SELECT * FROM memories WHERE type IN ('user', 'private') AND user_id = ? AND (key LIKE ? OR content LIKE ? OR category LIKE ?) ORDER BY importance DESC, updated_at DESC LIMIT ? """, (user_id, q, q, q, limit), ) elif type_ == "project" and user_id is not None and project_name: cur = conn.execute( """ SELECT * FROM memories WHERE type = 'project' AND user_id = ? AND project_name = ? AND (key LIKE ? OR content LIKE ? OR category LIKE ?) ORDER BY importance DESC, updated_at DESC LIMIT ? """, (user_id, project_name, q, q, q, limit), ) elif user_id is not None: if is_admin: cur = conn.execute( """ SELECT * FROM memories WHERE (type = 'global' OR (type IN ('user', 'private') AND user_id = ?) OR (type = 'project' AND user_id = ? AND project_name = ?)) AND (key LIKE ? OR content LIKE ? OR category LIKE ?) ORDER BY importance DESC, updated_at DESC LIMIT ? """, (user_id, user_id, project_name or "", q, q, q, limit), ) else: cur = conn.execute( """ SELECT * FROM memories WHERE ((type IN ('user', 'private') AND user_id = ?) OR (type = 'project' AND user_id = ? AND project_name = ?)) AND (key LIKE ? OR content LIKE ? OR category LIKE ?) ORDER BY importance DESC, updated_at DESC LIMIT ? """, (user_id, user_id, project_name or "", q, q, q, limit), ) else: if not is_admin: return [] cur = conn.execute( """ SELECT * FROM memories WHERE type = 'global' AND (key LIKE ? OR content LIKE ? OR category LIKE ?) ORDER BY importance DESC, updated_at DESC LIMIT ? """, (q, q, q, limit), ) return [MemoryItem(**dict(r)) for r in cur.fetchall()] def format_memories_for_prompt( self, user_id: Optional[int] = None, project_name: Optional[str] = None ) -> Optional[str]: """ Formats all 3 tiers of memories into a structured, concise instruction block for AI: 1. 🌐 Global System Rules 2. 👤 User Personal Preferences & Profile 3. 📁 Active Project Specific Knowledge """ globals_list = self.get_global_memories(limit=50) users_list = self.get_user_memories(user_id=user_id, limit=50) if user_id else [] projects_list = ( self.get_project_memories(user_id=user_id, project_name=project_name, limit=50) if user_id and project_name else [] ) if not globals_list and not users_list and not projects_list: return None lines = [ "[SYSTEM INSTRUCTION: AI PERSISTENT HIERARCHICAL MEMORY & CONTINUOUS LEARNING]", "You have access to persistent 3-tier long-term memory containing verified system rules, project context, and user profile.", "Use this knowledge to maintain continuity and customize all your responses accordingly.\n", "⚖️ STRICT PRECEDENCE HIERARCHY (سلسله‌مراتب قطعی اولویت و حل تعارض):", "1. 🥇 Tier 1 (Highest Authority): Global System Rules ALWAYS strictly supersede all Project and User instructions.", "2. 🥈 Tier 2 (Intermediate Authority): Active Project Knowledge & Rules ALWAYS supersede general User preferences inside this project.", "3. 🥉 Tier 3 (Baseline Authority): User Personal Preferences apply across projects as defaults when not overridden by Project or Global rules.\n" ] if globals_list: lines.append("🌐 1. GLOBAL SYSTEM RULES & KNOWLEDGE (قوانین عمومی سیستم - بالاترین اولویت):") for m in globals_list: cat = f"[{m.category}]" if m.category and m.category != "general" else "" lines.append(f"• [key: {m.key}] {cat} {m.content}") lines.append("") if projects_list: lines.append(f"📁 2. ACTIVE PROJECT KNOWLEDGE [پروژه: {project_name}] (دانش و تصمیمات اختصاصی این پروژه - اولویت دوم):") for m in projects_list: cat = f"[{m.category}]" if m.category and m.category != "general" else "" lines.append(f"• [key: {m.key}] {cat} {m.content}") lines.append("") if users_list: lines.append("👤 3. USER PERSONAL PROFILE & PREFERENCES (ترجیحات و مشخصات شخصی کاربر - اولویت پایه):") for m in users_list: cat = f"[{m.category}]" if m.category and m.category != "general" else "" lines.append(f"• [key: {m.key}] {cat} {m.content}") lines.append("") lines.extend([ "🧠 AUTONOMOUS MEMORY MANAGEMENT RULES:", "1. When user shares durable facts, choose the appropriate tier:", ' • User-level (developer habits, preferred tools, personal info): type="user"', ' • Project-level (architecture, database schema, libraries, API designs, fixed bugs in this project): type="project"', ' Action Tag: [[SAVE_MEMORY: type="user|project", key="", category="preference|rule|tech_stack|fact", content=""]]', "2. Note: Global memory is read-only for AI and managed exclusively by system administrators via the Telegram Bot UI.", "3. If past memory is superseded, update it: [[UPDATE_MEMORY: type=\"user|project\", key=\"\", content=\"\"]].", "4. If user asks to forget something: [[DELETE_MEMORY: type=\"user|project\", key=\"\"]].", "5. Do NOT store temporary or trivial small-talk. Only store enduring and valuable knowledge." ]) return "\n".join(lines) def get_stats(self, user_id: Optional[int] = None, project_name: Optional[str] = None) -> Dict[str, Any]: with self._get_connection() as conn: cur = conn.execute("SELECT COUNT(*) as cnt FROM memories WHERE type = 'global'") global_count = cur.fetchone()["cnt"] cur = conn.execute("SELECT COUNT(*) as cnt FROM memories WHERE type IN ('user', 'private')") user_count = cur.fetchone()["cnt"] cur = conn.execute("SELECT COUNT(*) as cnt FROM memories WHERE type = 'project'") project_count = cur.fetchone()["cnt"] cur = conn.execute("SELECT COUNT(DISTINCT user_id) as cnt FROM memories WHERE type IN ('user', 'project', 'private')") users_with_mem = cur.fetchone()["cnt"] this_user_count = 0 this_project_count = 0 if user_id: cur = conn.execute("SELECT COUNT(*) as cnt FROM memories WHERE type IN ('user', 'private') AND user_id = ?", (user_id,)) this_user_count = cur.fetchone()["cnt"] if project_name: cur = conn.execute("SELECT COUNT(*) as cnt FROM memories WHERE type = 'project' AND user_id = ? AND project_name = ?", (user_id, project_name)) this_project_count = cur.fetchone()["cnt"] return { "total_memories": global_count + user_count + project_count, "global_memories": global_count, "user_memories": user_count, "project_memories": project_count, "users_with_memory": users_with_mem, "current_user_memories": this_user_count, "current_project_memories": this_project_count, } def ingest_and_cleanup_agents_md( self, workspace_path: str, user_id: Optional[int] = None, project_name: Optional[str] = None, is_admin: bool = False, ) -> List[MemoryItem]: """ Detects if AGENTS.md (or agents.md / GEMINI.md) exists in the workspace. Parses sections and rules, extracts and saves them into memory.db, and finally deletes the AGENTS.md file as per system rule. """ if not workspace_path: return [] ws = Path(workspace_path) if not ws.is_dir(): return [] target_file = None for candidate in (ws / "AGENTS.md", ws / "agents.md", ws / "GEMINI.md", ws / "gemini.md"): if candidate.is_file(): target_file = candidate break if not target_file: return [] try: content_text = target_file.read_text(encoding="utf-8").strip() if not content_text: try: target_file.unlink() except Exception: pass return [] # Determine target memory tier if project_name and project_name != "default" and user_id: mem_type = "project" elif project_name == "default" and is_admin: mem_type = "global" elif user_id and project_name: mem_type = "project" elif is_admin: mem_type = "global" elif user_id: mem_type = "user" else: mem_type = "global" # Parse markdown sections import re sections = re.split(r'(?m)^#{1,4}\s+', content_text) extracted_items = [] for section in sections: section = section.strip() if not section: continue lines = section.splitlines() header = lines[0].strip() body = "\n".join(lines[1:]).strip() if len(lines) > 1 else header if not body: body = header # Generate a clean slug key clean_title = re.sub(r'[\(\)\[\]\{\}\/\\:,\.\?\!_#\*\`]', ' ', header) clean_title = re.sub(r'\s+', '_', clean_title.strip().lower()) clean_title = clean_title.strip('_')[:50] or "rule" if not clean_title or clean_title.isdigit(): clean_title = f"rule_{clean_title}" # Determine category sec_lower = (header + " " + body).lower() if any(w in sec_lower for w in ("tech", "stack", "php", "python", "node", "db", "database", "پشته", "فناوری")): cat = "tech_stack" elif any(w in sec_lower for w in ("workflow", "process", "فرآیند", "جریان", "هاب", "پایپلاین")): cat = "workflow" elif any(w in sec_lower for w in ("preference", "ترجیح", "قالب", "فارسی", "زبان")): cat = "preference" else: cat = "rule" extracted_items.append((clean_title, cat, body)) if not extracted_items: extracted_items.append((f"{project_name or 'project'}_instructions", "rule", content_text)) saved_items = [] for slug_key, category, body_text in extracted_items: if len(body_text.strip()) < 3: continue item, _ = self.save_or_update( type_=mem_type, key=slug_key, content=body_text, user_id=user_id, project_name=project_name if mem_type == "project" else None, category=category, importance=2, ) saved_items.append(item) logger.info( f"Successfully ingested {len(saved_items)} memories from {target_file} into tier '{mem_type}' (user={user_id}, proj={project_name})." ) # Delete file after ingestion try: target_file.unlink() logger.info(f"Deleted ingested file: {target_file}") except Exception as de: logger.warning(f"Could not delete file {target_file} after ingestion: {de}") return saved_items except Exception as ex: logger.error(f"Error ingesting {target_file}: {ex}", exc_info=True) return [] memory_manager = MemoryManager()