409 lines
17 KiB
Python
409 lines
17 KiB
Python
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' or 'private'
|
|
user_id: Optional[int] # None for global, chat_id for private
|
|
key: str
|
|
category: str # '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_private(self) -> bool:
|
|
return self.type == "private"
|
|
|
|
|
|
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 if they do not exist."""
|
|
try:
|
|
with self._get_connection() as conn:
|
|
conn.execute("""
|
|
CREATE TABLE IF NOT EXISTS memories (
|
|
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
|
type TEXT NOT NULL CHECK(type IN ('global', 'private')),
|
|
user_id INTEGER,
|
|
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
|
|
);
|
|
""")
|
|
# Partial unique indexes to properly handle NULL user_id for global memories
|
|
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_private_key
|
|
ON memories(type, user_id, key) WHERE type = 'private';
|
|
""")
|
|
conn.execute("""
|
|
CREATE INDEX IF NOT EXISTS idx_memories_user_type
|
|
ON memories(type, user_id);
|
|
""")
|
|
conn.execute("""
|
|
CREATE INDEX IF NOT EXISTS idx_memories_category
|
|
ON memories(category);
|
|
""")
|
|
logger.info("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,
|
|
category: str = "general",
|
|
importance: int = 1,
|
|
created_by: Optional[int] = None,
|
|
) -> Tuple[MemoryItem, bool]:
|
|
"""
|
|
Saves or updates a memory.
|
|
Returns: (MemoryItem, is_created: bool)
|
|
"""
|
|
type_ = type_.lower().strip()
|
|
if type_ not in ("global", "private"):
|
|
raise ValueError("Memory type must be 'global' or 'private'")
|
|
|
|
if type_ == "global":
|
|
user_id = None
|
|
elif user_id is None:
|
|
raise ValueError("Private memory requires a valid user_id")
|
|
|
|
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 exists
|
|
if type_ == "global":
|
|
cur = conn.execute(
|
|
"SELECT id, created_at, created_by FROM memories WHERE type = 'global' AND key = ?",
|
|
(key,),
|
|
)
|
|
else:
|
|
cur = conn.execute(
|
|
"SELECT id, created_at, created_by FROM memories WHERE type = 'private' AND user_id = ? AND key = ?",
|
|
(user_id, 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 category = ?, content = ?, importance = ?, updated_at = ?
|
|
WHERE id = ?
|
|
""",
|
|
(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, key, category, content, importance, created_by, created_at, updated_at)
|
|
VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)
|
|
""",
|
|
(type_, user_id, 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,
|
|
key=key,
|
|
category=category,
|
|
content=content,
|
|
importance=importance,
|
|
created_by=creator,
|
|
created_at=created_at,
|
|
updated_at=now,
|
|
)
|
|
logger.info(
|
|
"%s memory [id=%d, key=%s, type=%s, user=%s]",
|
|
"Created" if is_created else "Updated",
|
|
mem_id,
|
|
key,
|
|
type_,
|
|
user_id,
|
|
)
|
|
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) -> Optional[MemoryItem]:
|
|
type_ = type_.lower().strip()
|
|
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,))
|
|
else:
|
|
cur = conn.execute(
|
|
"SELECT * FROM memories WHERE type = 'private' AND user_id = ? AND key = ?",
|
|
(user_id, 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 = '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 = '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 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 == "private" and mem.user_id != user_id and not is_admin:
|
|
raise PermissionError("Cannot delete another user's private 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, is_admin: bool = False
|
|
) -> bool:
|
|
type_ = type_.lower().strip()
|
|
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,))
|
|
else:
|
|
if user_id is None:
|
|
return False
|
|
cur = conn.execute(
|
|
"DELETE FROM memories WHERE type = 'private' AND user_id = ? AND key = ?",
|
|
(user_id, key),
|
|
)
|
|
deleted = cur.rowcount > 0
|
|
if deleted:
|
|
logger.info("Deleted memory by key [key=%s, type=%s, user=%s]", key, type_, user_id)
|
|
return deleted
|
|
|
|
def clear_memories(self, type_: str, user_id: Optional[int] = None, is_admin: bool = False) -> int:
|
|
type_ = type_.lower().strip()
|
|
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'")
|
|
else:
|
|
if user_id is None:
|
|
return 0
|
|
cur = conn.execute("DELETE FROM memories WHERE type = 'private' AND user_id = ?", (user_id,))
|
|
count = cur.rowcount
|
|
logger.info("Cleared %d %s memories for user %s", count, type_, user_id)
|
|
return count
|
|
|
|
def search_memories(
|
|
self, query: str, user_id: Optional[int] = None, type_: Optional[str] = None, limit: int = 50
|
|
) -> List[MemoryItem]:
|
|
q = f"%{query.strip().lower()}%"
|
|
with self._get_connection() as conn:
|
|
if type_ == "global":
|
|
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_ == "private" and user_id is not None:
|
|
cur = conn.execute(
|
|
"""
|
|
SELECT * FROM memories
|
|
WHERE type = '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 user_id is not None:
|
|
cur = conn.execute(
|
|
"""
|
|
SELECT * FROM memories
|
|
WHERE (type = 'global' OR (type = 'private' AND user_id = ?))
|
|
AND (key LIKE ? OR content LIKE ? OR category LIKE ?)
|
|
ORDER BY type DESC, importance DESC, updated_at DESC LIMIT ?
|
|
""",
|
|
(user_id, q, q, q, limit),
|
|
)
|
|
else:
|
|
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) -> Optional[str]:
|
|
"""
|
|
Formats all applicable memories (Global + User's Private) into a concise, token-efficient
|
|
system instruction block for the AI model.
|
|
Returns None if no memories are found.
|
|
"""
|
|
globals_list = self.get_global_memories(limit=50)
|
|
privates_list = self.get_user_memories(user_id=user_id, limit=50) if user_id else []
|
|
|
|
if not globals_list and not privates_list:
|
|
return None
|
|
|
|
lines = [
|
|
"[SYSTEM INSTRUCTION: AI PERSISTENT MEMORY & CONTINUOUS LEARNING]",
|
|
"You have access to persistent long-term memory containing verified system rules and user context.",
|
|
"Use this knowledge to maintain continuity and customize all your responses accordingly.\n"
|
|
]
|
|
|
|
if globals_list:
|
|
lines.append("🌐 GLOBAL SYSTEM MEMORY (قوانین و اطلاعات عمومی سیستم):")
|
|
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 privates_list:
|
|
lines.append(f"👤 USER PRIVATE MEMORY (خاطرات و ترجیحات اختصاصی این کاربر):")
|
|
for m in privates_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 the user shares important facts, personal preferences, coding conventions, or rules that should persist, output the action tag:",
|
|
' [[SAVE_MEMORY: type="private|global", key="<short_slug>", category="preference|rule|tech_stack|fact", content="<concise memory summary>"]]',
|
|
"2. Note: Non-admin users can ONLY save 'private' memories. Only system administrators can set 'global' memories.",
|
|
"3. If a past memory is superseded or changed, update it with the same key or use [[UPDATE_MEMORY: key=\"<slug>\", content=\"<new>\"]].",
|
|
"4. If the user explicitly asks to forget something, use [[DELETE_MEMORY: key=\"<slug>\"]].",
|
|
"5. Do NOT store temporary or trivial small-talk. Only store enduring and valuable knowledge."
|
|
])
|
|
|
|
return "\n".join(lines)
|
|
|
|
def get_stats(self) -> 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 = 'private'")
|
|
private_count = cur.fetchone()["cnt"]
|
|
|
|
cur = conn.execute("SELECT COUNT(DISTINCT user_id) as cnt FROM memories WHERE type = 'private'")
|
|
user_count = cur.fetchone()["cnt"]
|
|
|
|
return {
|
|
"total_memories": global_count + private_count,
|
|
"global_memories": global_count,
|
|
"private_memories": private_count,
|
|
"users_with_memory": user_count,
|
|
}
|
|
|
|
|
|
memory_manager = MemoryManager()
|