AI Update: پس فاز اول رو پیاده کن
This commit is contained in:
Antigravity Bot
Binary file not shown.
@@ -0,0 +1,408 @@
|
||||
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()
|
||||
File diff suppressed because one or more lines are too long
Reference in New Issue
Block a user