Beta/ia/memoire.py

528 lines
17 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})")
async def summarize_messages(messages: List[Dict[str, Any]]) -> str:
"""
Fonction de résumé améliorée. Utilise une extraction de mots-clés et de phrases clés
pour un résumé plus intelligent. Placeholder pour intégration IA future.
Cette fonction est `async` pour faciliter un remplacement par un appel réseau/IA.
"""
if not messages:
return ""
# Extraire tous les textes utilisateur et bot
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 simple 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 (premières phrases des messages importants)
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] # première phrase
key_phrases.extend(sentences)
if bm:
sentences = re.split(r'[.!?]', bm)[:1]
key_phrases.extend(sentences)
phrases_str = " | ".join(key_phrases[:5]) # limiter à 5 phrases
# Construire le résumé
summary = f"Mots-clés: {keywords_str}\nPhrases clés: {phrases_str}"
# Tronquer si trop long
if len(summary) > 1500:
summary = summary[:1497] + "..."
await asyncio.sleep(0) # point d'await pour remplacer par un vrai appel
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 ~6000 caractères (ajustable)
joined = "\n".join(parts)
if len(joined) > 6000:
return joined[-6000:]
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) -> None:
path = _user_filepath(user_id)
try:
os.remove(path)
except FileNotFoundError:
pass
# 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()