Beta/brain/memoire.py

600 lines
20 KiB
Python

"""
Module de gestion de la mémoire utilisateur pour le bot Superviseur.
- Fichiers JSON par utilisateur (dossier configurable)
- Historique récent (liste) + résumé global
- Écriture atomique pour éviter la corruption
- Compression automatique de l'historique via une fonction de "summarize" (placeholder)
Usage minimal :
from ia.memoire import add_interaction, build_context_for_model
# Avant de générer la réponse IA :
context_prompt = await build_context_for_model(user_id, max_recent=12)
# inclure `context_prompt` dans le prompt envoyé à ton modèle
# Après avoir généré la réponse IA :
await add_interaction(user_id, channel_id, user_message, bot_response)
Ce module est asynchrone-friendly mais fonctionne aussi en sync pour l'accès disque.
"""
from __future__ import annotations
import os
import json
import tempfile
import asyncio
import aiofiles
import time
try:
import redis.asyncio as aioredis
except Exception: # optional dependency
aioredis = None
import logging
import re
from datetime import datetime, timezone
from typing import Any, Dict, List, Optional
# Config
MEMORY_DIR = os.environ.get("SUPERVISEUR_MEM_DIR", "memoires")
MAX_RECENT = int(os.environ.get("SUPERVISEUR_MAX_RECENT", "20")) # garder X messages avant résumé
KEEP_AFTER_SUMMARIZE = int(os.environ.get("SUPERVISEUR_KEEP_RECENT_AFTER_SUMMARY", "8"))
# Validation des configs
if MAX_RECENT <= 0 or KEEP_AFTER_SUMMARIZE < 0 or KEEP_AFTER_SUMMARIZE >= MAX_RECENT:
raise ValueError("Configurations invalides: vérifiez SUPERVISEUR_MAX_RECENT et SUPERVISEUR_KEEP_RECENT_AFTER_SUMMARY")
# Logger
logger = logging.getLogger(__name__)
# Métriques simples
metrics = {
"total_interactions": 0,
"total_summaries": 0,
"memory_files": 0
}
# Extended metrics
metrics.setdefault("cache_hits", 0)
metrics.setdefault("cache_misses", 0)
metrics.setdefault("redis_reads", 0)
metrics.setdefault("redis_writes", 0)
metrics.setdefault("file_reads", 0)
metrics.setdefault("file_writes", 0)
# Optional redis client
_redis_client = None
# Format d'une entrée d'historique
# {
# "timestamp": "2025-11-25T22:00:00Z",
# "channel_id": 123456789,
# "user_message": "...",
# "bot_message": "..."
# }
def _ensure_memory_dir() -> None:
os.makedirs(MEMORY_DIR, exist_ok=True)
logger.debug(f"Répertoire mémoire assuré: {MEMORY_DIR}")
def _user_filepath(user_identifier: int | str, guild_id: str | None = None) -> str:
_ensure_memory_dir()
base_dir = MEMORY_DIR
if guild_id:
# Create per-guild subdirectory
base_dir = os.path.join(MEMORY_DIR, str(guild_id))
os.makedirs(base_dir, exist_ok=True)
if isinstance(user_identifier, str):
# Sanitize for filename: replace invalid chars and spaces with _
sanitized = re.sub(r'[\/\\:*?"<>|\s]', '_', user_identifier)
return os.path.join(base_dir, f"{sanitized}.json")
else:
return os.path.join(base_dir, f"{user_identifier}.json")
def _redis_key(user_id: int | str, guild_id: str | None = None) -> str:
guild_part = f"_{guild_id}" if guild_id else ""
return f"mem{guild_part}:{user_id}"
def _use_redis() -> bool:
return bool(os.environ.get("SUPERVISEUR_REDIS_URL")) and aioredis is not None
def _ensure_redis():
global _redis_client
if not _use_redis():
return None
if _redis_client is None:
url = os.environ.get("SUPERVISEUR_REDIS_URL")
_redis_client = aioredis.from_url(url)
return _redis_client
async def _atomic_write_async(path: str, data: Dict[str, Any]) -> None:
# Écriture atomique via fichier temporaire + replace en async
dirpath = os.path.dirname(path)
fd, tmp_path = tempfile.mkstemp(dir=dirpath)
try:
# Since aiofiles doesn't support os.fdopen directly, write via path
async with aiofiles.open(tmp_path, "w", encoding="utf-8") as f:
await f.write(json.dumps(data, ensure_ascii=False, indent=2))
await f.flush()
# Ensure durability at OS level
try:
with open(tmp_path, "rb") as f:
os.fsync(f.fileno())
except Exception:
pass
os.replace(tmp_path, path)
finally:
if os.path.exists(tmp_path):
try:
os.remove(tmp_path)
except Exception:
pass
# In-memory write-back cache to avoid disk I/O on every interaction
memory_cache: Dict[str, Dict[str, Any]] = {}
# Dirty flags for cached items
_dirty: Dict[str, bool] = {}
# Per-user locks to avoid concurrent writes
_locks: Dict[str, asyncio.Lock] = {}
async def _get_lock_for(user_id: str) -> asyncio.Lock:
key = str(user_id)
if key not in _locks:
_locks[key] = asyncio.Lock()
return _locks[key]
async def _background_persist(user_id: str, guild_id: str | None) -> None:
"""Persist a cached user memory to disk if dirty."""
key = f"{guild_id or ''}_{user_id}"
lock = await _get_lock_for(key)
async with lock:
if not _dirty.get(key):
return
data = memory_cache.get(key)
if data is None:
return
path = _user_filepath(user_id, guild_id)
try:
if _use_redis():
client = _ensure_redis()
await client.set(_redis_key(user_id, guild_id), json.dumps(data))
metrics["redis_writes"] += 1
_dirty[key] = False
logger.debug(f"Persisted memory for {key} to redis")
else:
await _atomic_write_async(path, data)
_dirty[key] = False
metrics["file_writes"] += 1
logger.debug(f"Persisted memory for {key} to disk")
except Exception as e:
logger.exception(f"Erreur lors de la sauvegarde asynchrone pour {key}: {e}")
async def load_user_memory_async(user_id: int | str, guild_id: str | None = None) -> Dict[str, Any]:
"""Charge (ou crée) la mémoire d'un utilisateur.
Structure renvoyée :
{
"resume_global": str,
"historique_recent": [ { ... }, ... ]
}
"""
key = f"{guild_id or ''}_{user_id}"
# Check in-memory cache first
if key in memory_cache:
metrics["cache_hits"] += 1
return memory_cache[key]
path = _user_filepath(user_id, guild_id)
# Try redis if enabled
if _use_redis():
client = _ensure_redis()
try:
raw = await client.get(_redis_key(user_id, guild_id))
if raw:
metrics["redis_reads"] += 1
data = json.loads(raw)
memory_cache[key] = data
_dirty[key] = False
metrics["cache_hits"] += 1
return data
except Exception as e:
logger.warning(f"Redis load error for {key}: {e}")
if not os.path.exists(path):
# Prime cache
memory_cache[key] = {"resume_global": "", "historique_recent": []}
_dirty[key] = False
metrics["cache_misses"] += 1
return memory_cache[key]
try:
async with aiofiles.open(path, "r", encoding="utf-8") as f:
content = await f.read()
data = json.loads(content)
memory_cache[key] = data
_dirty[key] = False
metrics["file_reads"] += 1
return data
except Exception as e:
logger.warning(f"Erreur lors du chargement de la mémoire pour {user_id} (guild {guild_id}): {e}")
# Si fichier corrompu, on sauvegarde une copie et on recrée
backup = path + ".corrompu." + datetime.now(timezone.utc).isoformat()
try:
os.rename(path, backup)
except Exception:
pass
return {"resume_global": "", "historique_recent": []}
async def save_user_memory_async(user_id: int | str, data: Dict[str, Any], guild_id: str | None = None, persist_immediately: bool = False) -> None:
key = f"{guild_id or ''}_{user_id}"
memory_cache[key] = data
_dirty[key] = True
# If immediate persist requested, write now, else schedule background persist
if persist_immediately:
await _background_persist(user_id, guild_id)
else:
# Schedule a background persist with small delay to group writes
async def _delayed():
await asyncio.sleep(0.3)
await _background_persist(user_id, guild_id)
asyncio.create_task(_delayed())
logger.debug(f"Mémoire marquée dirty pour utilisateur {user_id} (guild {guild_id})")
# Référence globale au LLM manager (set par le bot au démarrage)
_llm_manager = None
def set_llm_manager(llm_manager) -> None:
"""Enregistre le LLM manager pour les résumés."""
global _llm_manager
_llm_manager = llm_manager
async def summarize_messages(messages: List[Dict[str, Any]]) -> str:
"""
Résume une liste de messages via le LLM.
Fallback sur l'extraction de mots-clés si le LLM n'est pas disponible.
"""
if not messages:
return ""
# Construire le transcript à résumer
transcript_lines = []
for m in messages:
um = m.get("user_message", "").strip()
bm = m.get("bot_message", "").strip()
if um:
transcript_lines.append(f"Utilisateur: {um}")
if bm:
transcript_lines.append(f"Bot: {bm}")
transcript = "\n".join(transcript_lines)
# Tronquer le transcript pour éviter les appels trop longs
if len(transcript) > 4000:
transcript = transcript[-4000:]
# Essayer d'utiliser le LLM pour un vrai résumé
if _llm_manager is not None:
try:
prompt = (
"Tu es un assistant spécialisé dans la synthèse de conversations.\n"
"Résume les échanges suivants de manière concise et factuelle.\n"
"Extrais les sujets abordés, les décisions prises, et les informations importantes.\n"
"Le résumé ne doit PAS dépasser 500 caractères.\n\n"
f"CONVERSATION À RÉSUMER :\n{transcript}\n\n"
"RÉSUMÉ :"
)
payload = _llm_manager.build_payload(
prompt=prompt,
system_prompt="Tu es un assistant de synthèse. Sois concis et factuel.",
force_json=False
)
summary = await _llm_manager.call_llama(payload)
if summary and len(summary.strip()) > 10:
# Nettoyer le résumé
summary = summary.strip()
if len(summary) > 1500:
summary = summary[:1497] + "..."
return summary
except Exception as e:
logger.warning(f"Erreur résumé LLM, fallback extractif: {e}")
# Fallback : extraction de mots-clés et phrases clés
return _extractive_summary(messages)
def _extractive_summary(messages: List[Dict[str, Any]]) -> str:
"""Résumé extractif basé sur les mots-clés et phrases clés (fallback)."""
all_texts = []
for m in messages:
um = m.get("user_message", "")
bm = m.get("bot_message", "")
all_texts.append(um + " " + bm)
combined_text = " ".join(all_texts)
# Extraction de mots-clés (mots fréquents, >3 chars)
words = re.findall(r'\b\w{4,}\b', combined_text.lower())
word_freq = {}
for word in words:
word_freq[word] = word_freq.get(word, 0) + 1
# Top 10 mots-clés
top_keywords = sorted(word_freq.items(), key=lambda x: x[1], reverse=True)[:10]
keywords_str = ", ".join([kw for kw, _ in top_keywords])
# Extraire phrases clés
key_phrases = []
for m in messages:
um = m.get("user_message", "").strip()
bm = m.get("bot_message", "").strip()
if um:
sentences = re.split(r'[.!?]', um)[:1]
key_phrases.extend(sentences)
if bm:
sentences = re.split(r'[.!?]', bm)[:1]
key_phrases.extend(sentences)
phrases_str = " | ".join(key_phrases[:5])
summary = f"Mots-clés: {keywords_str}\nPhrases clés: {phrases_str}"
if len(summary) > 1500:
summary = summary[:1497] + "..."
return summary
async def _maybe_summarize(user_id: int | str, data: Dict[str, Any]) -> Dict[str, Any]:
"""Si historique_recent dépasse MAX_RECENT, on en résume une partie dans resume_global."""
hist: List[Dict[str, Any]] = data.get("historique_recent", [])
if len(hist) <= MAX_RECENT:
return data
# On prend les plus anciennes entrées à résumer
to_summarize = hist[: len(hist) - KEEP_AFTER_SUMMARIZE]
remaining = hist[len(hist) - KEEP_AFTER_SUMMARIZE :]
summary_text = await summarize_messages(to_summarize)
# Concatène proprement le résumé au resume_global
existing = data.get("resume_global", "")
if existing:
new_resume = existing + "\n\n" + summary_text
else:
new_resume = summary_text
# tronquer si trop long
if len(new_resume) > 20000:
new_resume = new_resume[-20000:]
data["resume_global"] = new_resume
data["historique_recent"] = remaining
metrics["total_summaries"] += 1
logger.info(f"Résumé effectué pour utilisateur {user_id}: {len(to_summarize)} messages résumés")
return data
async def add_interaction(
user_id: int | str,
channel_id: int | str,
user_message: str,
bot_message: str,
guild_id: str | None = None,
timestamp: Optional[str] = None,
) -> None:
"""Ajoute une interaction à la mémoire utilisateur et sauvegarde.
timestamp : isoformat UTC string ; si non fourni, généré automatiquement.
"""
if timestamp is None:
timestamp = datetime.now(timezone.utc).isoformat()
data = await load_user_memory_async(user_id, guild_id)
hist = data.get("historique_recent", [])
hist.append(
{
"timestamp": timestamp,
"channel_id": int(channel_id) if channel_id is not None else None,
"user_message": user_message,
"bot_message": bot_message,
}
)
data["historique_recent"] = hist
# éventuellement résumer en arrière-plan
data = await _maybe_summarize(user_id, data)
# Sauvegarde asynchrone (mise en cache + flush en arrière-plan)
await save_user_memory_async(user_id, data, guild_id)
async def build_context_for_model(
user_id: int | str,
max_recent: int = 12,
include_resume: bool = True,
guild_id: str | None = None,
) -> str:
"""Construit la portion de prompt à passer au modèle, en respectant la mémoire.
- resume_global (optionnel) en tête
- derniers messages récents (jusqu'à max_recent)
Retour : chaîne prête à concaténer avec le prompt courant.
"""
data = await load_user_memory_async(user_id, guild_id)
parts: List[str] = []
if include_resume:
resume = data.get("resume_global", "")
if resume:
parts.append("Contexte (résumé des conversations précédentes) :\n" + resume)
recent = data.get("historique_recent", [])
# On prend les N derniers
recent_tail = recent[-max_recent:] if recent else []
if recent_tail:
parts.append("Historique de la conversation :")
for m in recent_tail:
um = m.get("user_message", "").replace("\n", " ")
bm = m.get("bot_message", "").replace("\n", " ")
parts.append(f"Utilisateur: {um}")
parts.append(f"Bot: {bm}")
# Contrainte : ne pas renvoyer un prompt de plus de ~3000 caractères (ajustable)
joined = "\n".join(parts)
if len(joined) > 3000:
return joined[-3000:]
return joined
async def purge_old_memories(days: int = 30) -> int:
"""
Supprime les fichiers de mémoire plus vieux que X jours pour conformité RGPD.
Retourne le nombre de fichiers supprimés.
"""
_ensure_memory_dir()
now = time.time()
cutoff = now - (days * 24 * 3600)
purged_count = 0
# Parcourir récursivement si guilds, sinon juste MEMORY_DIR
for root, dirs, files in os.walk(MEMORY_DIR):
for name in files:
if name.endswith(".json"):
path = os.path.join(root, name)
try:
mtime = os.path.getmtime(path)
if mtime < cutoff:
os.remove(path)
purged_count += 1
logger.info(f"RGPD: Fichier expiré supprimé: {path}")
except Exception as e:
logger.error(f"Erreur purge RGPD {path}: {e}")
return purged_count
# Utilitaires optionnels
def list_memory_files() -> List[str]:
_ensure_memory_dir()
files = [f for f in os.listdir(MEMORY_DIR) if f.endswith(".json")]
metrics["memory_files"] = len(files)
return files
def delete_user_memory(user_id: int | str, guild_id: str | None = None) -> None:
"""Supprime la mémoire d'un utilisateur, optionnellement pour un guild spécifique."""
path = _user_filepath(user_id, guild_id)
try:
os.remove(path)
except FileNotFoundError:
pass
# Si aucun guild_id spécifique, supprimer aussi tous les fichiers de mémoire pour cet utilisateur
if guild_id is None:
_ensure_memory_dir()
user_str = str(user_id)
for root, dirs, files in os.walk(MEMORY_DIR):
for name in files:
if name == f"{user_str}.json" or name == f"{user_id}.json":
full_path = os.path.join(root, name)
if full_path != path: # Éviter double suppression
try:
os.remove(full_path)
except FileNotFoundError:
pass
# Nettoyer le cache
key_no_guild = f"_{user_id}" if guild_id is None else f"{guild_id}_{user_id}"
memory_cache.pop(key_no_guild, None)
_dirty.pop(key_no_guild, None)
# Fonction pour obtenir les métriques
def get_metrics() -> Dict[str, Any]:
return metrics.copy()
async def flush_all() -> None:
"""Persist all dirty cached memories to disk (used at shutdown)."""
keys = [k for k, v in _dirty.items() if v]
if not keys:
return
tasks = []
for key in keys:
# Parse key: "guild_user" or "_user" for None guild
if '_' in key:
guild_str, user_str = key.split('_', 1)
guild_id = guild_str if guild_str else None
else:
guild_id = None
user_str = key
tasks.append(_background_persist(user_str, guild_id))
await asyncio.gather(*tasks, return_exceptions=True)
# Tests unitaires simples
def run_tests():
"""Exécute des tests basiques pour vérifier le fonctionnement."""
import tempfile
import shutil
# Créer un répertoire temporaire pour les tests
test_dir = tempfile.mkdtemp()
global MEMORY_DIR
original_dir = MEMORY_DIR
MEMORY_DIR = test_dir
try:
# Test load/save (async)
uid = "test_user"
data = {"resume_global": "test", "historique_recent": []}
import asyncio as _asyncio
_asyncio.run(save_user_memory_async(uid, data, persist_immediately=True))
loaded = _asyncio.run(load_user_memory_async(uid))
assert loaded == data, "Load/save failed"
# Test add_interaction
import asyncio as _asyncio
_asyncio.run(add_interaction(uid, 123, "Hello", "Hi"))
loaded = _asyncio.run(load_user_memory_async(uid))
assert len(loaded["historique_recent"]) == 1, "Add interaction failed"
# Test build_context
ctx = _asyncio.run(build_context_for_model(uid))
assert "Hello" in ctx, "Build context failed"
print("Tous les tests passent !")
finally:
# Nettoyer
shutil.rmtree(test_dir)
MEMORY_DIR = original_dir
# Si tu veux tester vite :
if __name__ == "__main__":
import asyncio
async def demo():
uid = 123456
await add_interaction(uid, 111, "Salut, tu peux modérer ?", "Oui, quel canal ?")
await add_interaction(uid, 111, "Le canal #raids", "Ok je surveille")
ctx = await build_context_for_model(uid)
print("Contexte construit:")
print(ctx)
print("\nMétriques:", get_metrics())
asyncio.run(demo())
run_tests()