zai-telegram-bot / memory.py
Ejdjdososs's picture
Upload memory.py with huggingface_hub
831b557 verified
Raw
History Blame Contribute Delete
9.02 kB
"""
الذاكرة الدائمة - SQLite
يحفظ:
- سياق كل مستخدم (آخر N رسالة)
- تفضيلات المستخدم (النموذج المفضل، اللغة)
- إحصائيات الاستخدام
"""
import asyncio
import json
import logging
import os
import time
from typing import List, Dict, Optional
import aiosqlite
from config import config
logger = logging.getLogger(__name__)
class Memory:
"""إدارة ذاكرة البوت الدائمة عبر SQLite"""
def __init__(self, db_path: str):
self.db_path = db_path
# تأكد من وجود المجلد الأب
os.makedirs(os.path.dirname(self.db_path) or ".", exist_ok=True)
self._lock = asyncio.Lock()
async def init(self):
"""تهيئة قاعدة البيانات"""
async with aiosqlite.connect(self.db_path) as db:
await db.execute("""
CREATE TABLE IF NOT EXISTS messages (
id INTEGER PRIMARY KEY AUTOINCREMENT,
user_id INTEGER NOT NULL,
role TEXT NOT NULL,
content TEXT NOT NULL,
model TEXT,
timestamp INTEGER NOT NULL
)
""")
await db.execute("""
CREATE INDEX IF NOT EXISTS idx_messages_user
ON messages(user_id, timestamp DESC)
""")
await db.execute("""
CREATE TABLE IF NOT EXISTS users (
user_id INTEGER PRIMARY KEY,
username TEXT,
first_name TEXT,
language TEXT DEFAULT 'ar',
preferred_model TEXT,
is_owner INTEGER DEFAULT 0,
is_admin INTEGER DEFAULT 0,
first_seen INTEGER NOT NULL,
last_seen INTEGER NOT NULL,
message_count INTEGER DEFAULT 0
)
""")
await db.execute("""
CREATE TABLE IF NOT EXISTS github_cache (
user_id INTEGER NOT NULL,
repo_full_name TEXT NOT NULL,
data TEXT NOT NULL,
timestamp INTEGER NOT NULL,
PRIMARY KEY (user_id, repo_full_name)
)
""")
await db.commit()
logger.info(f"Database initialized at {self.db_path}")
async def register_user(
self,
user_id: int,
username: str = "",
first_name: str = "",
) -> Dict:
"""تسجيل أو تحديث مستخدم"""
now = int(time.time())
is_owner = 1 if config.is_owner(user_id) else 0
is_admin = 1 if config.is_admin(user_id) else 0
async with self._lock, aiosqlite.connect(self.db_path) as db:
# تحقق من وجود المستخدم
cursor = await db.execute(
"SELECT user_id, message_count FROM users WHERE user_id = ?",
(user_id,),
)
row = await cursor.fetchone()
if row is None:
await db.execute(
"""INSERT INTO users
(user_id, username, first_name, language, is_owner, is_admin,
first_seen, last_seen, message_count)
VALUES (?, ?, ?, ?, ?, ?, ?, ?, 0)""",
(user_id, username, first_name, "ar", is_owner, is_admin, now, now),
)
await db.commit()
return {
"user_id": user_id, "username": username,
"first_name": first_name, "is_new": True,
"message_count": 0,
}
else:
await db.execute(
"""UPDATE users SET
username = ?, first_name = ?, is_owner = ?, is_admin = ?,
last_seen = ?
WHERE user_id = ?""",
(username, first_name, is_owner, is_admin, now, user_id),
)
await db.commit()
return {
"user_id": user_id, "username": username,
"first_name": first_name, "is_new": False,
"message_count": row[1],
}
async def increment_message_count(self, user_id: int):
async with self._lock, aiosqlite.connect(self.db_path) as db:
await db.execute(
"UPDATE users SET message_count = message_count + 1 WHERE user_id = ?",
(user_id,),
)
await db.commit()
async def add_message(
self,
user_id: int,
role: str,
content: str,
model: str = "",
):
"""إضافة رسالة للذاكرة"""
# استخدام time.time() بدلاً من int(time.time()) لدقة أعلى
now = time.time()
async with self._lock, aiosqlite.connect(self.db_path) as db:
await db.execute(
"""INSERT INTO messages (user_id, role, content, model, timestamp)
VALUES (?, ?, ?, ?, ?)""",
(user_id, role, content, model, now),
)
# حذف الرسائل القديمة الزائدة عن الحد
await db.execute(
"""DELETE FROM messages WHERE user_id = ? AND id NOT IN (
SELECT id FROM messages WHERE user_id = ?
ORDER BY timestamp DESC, id DESC LIMIT ?
)""",
(user_id, user_id, config.MAX_HISTORY_PER_USER),
)
await db.commit()
async def get_history(self, user_id: int) -> List[Dict[str, str]]:
"""جلب سياق محادثة المستخدم (مرتب زمنياً، مع id كحاكم)"""
async with aiosqlite.connect(self.db_path) as db:
cursor = await db.execute(
"""SELECT role, content FROM messages
WHERE user_id = ? AND role IN ('user', 'assistant')
ORDER BY timestamp ASC, id ASC""",
(user_id,),
)
rows = await cursor.fetchall()
return [{"role": r[0], "content": r[1]} for r in rows]
async def clear_history(self, user_id: int) -> int:
"""حذف ذاكرة مستخدم - يعيد عدد الرسائل المحذوفة"""
async with self._lock, aiosqlite.connect(self.db_path) as db:
cursor = await db.execute(
"SELECT COUNT(*) FROM messages WHERE user_id = ?",
(user_id,),
)
count = (await cursor.fetchone())[0]
await db.execute("DELETE FROM messages WHERE user_id = ?", (user_id,))
await db.commit()
return count
async def set_preferred_model(self, user_id: int, model: str):
async with self._lock, aiosqlite.connect(self.db_path) as db:
await db.execute(
"UPDATE users SET preferred_model = ? WHERE user_id = ?",
(model, user_id),
)
await db.commit()
async def get_preferred_model(self, user_id: int) -> Optional[str]:
async with aiosqlite.connect(self.db_path) as db:
cursor = await db.execute(
"SELECT preferred_model FROM users WHERE user_id = ?",
(user_id,),
)
row = await cursor.fetchone()
return row[0] if row and row[0] else None
async def get_user_language(self, user_id: int) -> str:
async with aiosqlite.connect(self.db_path) as db:
cursor = await db.execute(
"SELECT language FROM users WHERE user_id = ?",
(user_id,),
)
row = await cursor.fetchone()
return row[0] if row else "ar"
async def set_user_language(self, user_id: int, lang: str):
async with self._lock, aiosqlite.connect(self.db_path) as db:
await db.execute(
"UPDATE users SET language = ? WHERE user_id = ?",
(lang, user_id),
)
await db.commit()
async def get_stats(self) -> Dict:
"""إحصائيات عامة"""
async with aiosqlite.connect(self.db_path) as db:
cursor = await db.execute("SELECT COUNT(*) FROM users")
users_count = (await cursor.fetchone())[0]
cursor = await db.execute("SELECT COUNT(*) FROM messages")
messages_count = (await cursor.fetchone())[0]
cursor = await db.execute(
"SELECT COUNT(*) FROM users WHERE is_owner = 1"
)
owners = (await cursor.fetchone())[0]
return {
"users": users_count,
"messages": messages_count,
"owners_configured": owners,
}
# Singleton
memory = Memory(config.DB_PATH)