"""Core MemorySystem - Orchestrates all storage, search, and lifecycle operations""" import logging from typing import Optional, List, Dict, Any from config import Config from embedding_service import EmbeddingService from storage.mysql_store import MySQLStore from storage.redis_cache import RedisCache from storage.vector_search import embedding_to_bytes, search_similar, find_duplicates from lifecycle.compressor import MemoryCompressor from llm_client import LLMClient from lifecycle.pipeline import Pipeline from lifecycle.aggregator import ScenarioAggregator from lifecycle.persona_generator import PersonaGenerator from lifecycle.fact_extractor import FactExtractor from lifecycle.cleaner import MemoryCleaner logger = logging.getLogger(__name__) class MemorySystem: """Multi-agent memory system with team isolation.""" def __init__(self, config: Config = None): cfg = config or Config() self.mysql = MySQLStore( host=cfg.MYSQL_HOST, port=cfg.MYSQL_PORT, user=cfg.MYSQL_USER, password=cfg.MYSQL_PASSWORD, database=cfg.MYSQL_DATABASE, unix_socket=cfg.MYSQL_UNIX_SOCKET, ) self.redis = RedisCache( host=cfg.REDIS_HOST, port=cfg.REDIS_PORT, password=cfg.REDIS_PASSWORD, db=cfg.REDIS_DB, ) self.embedder = EmbeddingService(base_url=cfg.EMBEDDING_SERVICE_URL) self.compressor = MemoryCompressor(self.mysql, self.redis, self.embedder) self.cleaner = MemoryCleaner(self.mysql) self.default_ttl = cfg.WORKING_MEMORY_TTL self.llm = LLMClient(providers_config=cfg.llm_providers) self.aggregator = ScenarioAggregator(self.mysql, self.llm) self.persona_gen = PersonaGenerator(self.mysql, self.llm) self.fact_extractor = FactExtractor(self.mysql, self.redis, self.llm, self.embedder) self.pipeline = Pipeline(self.mysql, self.redis, self.compressor, self.cleaner) # ── Team management ────────────────────────────────────────── def create_team(self, team_id: str, name: str, description: str = "", config: dict = None) -> dict: return self.mysql.create_team(team_id, name, description, config) def delete_team(self, team_id: str) -> bool: return self.mysql.delete_team(team_id) def get_team(self, team_id: str) -> Optional[dict]: return self.mysql.get_team(team_id) # ── Agent management ───────────────────────────────────────── def create_agent(self, agent_id: str, team_id: str, name: str, role: str = "") -> dict: # Verify team exists if not self.mysql.get_team(team_id): raise ValueError(f"Team {team_id} not found") return self.mysql.create_agent(agent_id, team_id, name, role) def delete_agent(self, agent_id: str) -> bool: return self.mysql.delete_agent(agent_id) def get_agent(self, agent_id: str) -> Optional[dict]: return self.mysql.get_agent(agent_id) def get_personal_memory(self, memory_id: str) -> Optional[dict]: return self.mysql.get_personal_memory(memory_id) def get_team_memory(self, memory_id: str) -> Optional[dict]: return self.mysql.get_team_memory(memory_id) def get_agents_by_team(self, team_id: str) -> list: return self.mysql.get_agents_by_team(team_id) # ── Personal memories ──────────────────────────────────────── def add_personal_memory(self, agent_id: str, content: str, importance: float = 0.5, metadata: dict = None, enable_dedup: bool = True, dedup_threshold: float = 0.85) -> dict: agent = self.mysql.get_agent(agent_id) if not agent: raise ValueError(f"Agent {agent_id} not found") # Generate embedding embedding = self.embedder.embed(content) emb_bytes = embedding_to_bytes(embedding) # Dedup check: compare against existing memories if enable_dedup: stored = self.mysql.get_personal_memories_with_embeddings(agent_id) if stored: dupes = find_duplicates(embedding, stored, threshold=dedup_threshold) if dupes: best_id, best_score, _ = dupes[0] logger.info(f"Dedup: memory {best_id} already covers this (score={best_score:.4f}), skipping") existing = self.mysql.get_personal_memory(best_id) if existing: existing["dedup_skipped"] = True existing["dedup_score"] = round(best_score, 4) return existing return self.mysql.add_personal_memory( agent_id=agent_id, team_id=agent["team_id"], content=content, embedding=emb_bytes, importance=importance, metadata=metadata, ) def search_personal_memories(self, agent_id: str, query: str, limit: int = 10, min_score: float = 0.0, max_chars_per_memory: int = 0, max_total_chars: int = 0) -> list: # Generate query embedding query_embedding = self.embedder.embed(query) # Get all memories with embeddings for this agent stored = self.mysql.get_personal_memories_with_embeddings(agent_id) if not stored: return [] # Vector search results = search_similar(query_embedding, stored, top_k=limit, min_score=min_score) # Format output with char limits output = [] total_chars = 0 for memory_id, score, item in results: self.mysql.touch_personal_memory(memory_id) content = item["content"] # Truncate single memory if needed if max_chars_per_memory > 0 and len(content) > max_chars_per_memory: content = content[:max_chars_per_memory] + "..." # Check total budget if max_total_chars > 0: if total_chars + len(content) > max_total_chars: remaining = max_total_chars - total_chars if remaining > 50: # only include if meaningful content = content[:remaining] + "..." else: break total_chars += len(content) output.append({ "id": memory_id, "content": content, "score": round(score, 4), "importance": item.get("importance", 0.5), "metadata": item.get("metadata"), }) return output def get_recent_personal_memories(self, agent_id: str, limit: int = 20) -> list: rows = self.mysql.get_personal_memories_by_agent(agent_id, limit) output = [ { "id": r["id"], "content": r["content"], "importance": r.get("importance", 0.5), "metadata": r.get("metadata"), "created_at": str(r.get("created_at", "")), } for r in rows ] for item in output: self.mysql.touch_personal_memory(item["id"]) return output def update_personal_memory(self, memory_id: str, content: str = None, importance: float = None, metadata: dict = None) -> bool: embedding = None if content is not None: emb = self.embedder.embed(content) embedding = embedding_to_bytes(emb) result = self.mysql.update_personal_memory( memory_id, content=content, importance=importance, metadata=metadata, embedding=embedding, ) if result: self.mysql.touch_personal_memory(memory_id) return result def delete_personal_memory(self, memory_id: str) -> bool: return self.mysql.delete_personal_memory(memory_id) # ── Working memory ─────────────────────────────────────────── def add_working_memory(self, agent_id: str, content: str, ttl: int = None) -> dict: ttl = ttl or self.default_ttl return self.redis.add_working_memory(agent_id, content, ttl=ttl) def get_working_memories(self, agent_id: str, limit: int = 20) -> list: return self.redis.get_working_memories(agent_id, limit) def clear_working_memory(self, agent_id: str) -> bool: return self.redis.clear_working_memory(agent_id) # ── Team memories ──────────────────────────────────────────── def add_team_memory(self, team_id: str, content: str, importance: float = 0.5, category: str = "general", metadata: dict = None, enable_dedup: bool = True, dedup_threshold: float = 0.85) -> dict: if not self.mysql.get_team(team_id): raise ValueError(f"Team {team_id} not found") embedding = self.embedder.embed(content) emb_bytes = embedding_to_bytes(embedding) # Dedup check if enable_dedup: stored = self.mysql.get_team_memories_with_embeddings(team_id) if stored: dupes = find_duplicates(embedding, stored, threshold=dedup_threshold) if dupes: best_id, best_score, _ = dupes[0] logger.info(f"Dedup: team memory {best_id} already covers this (score={best_score:.4f}), skipping") existing = self.mysql.get_team_memory(best_id) if existing: existing["dedup_skipped"] = True existing["dedup_score"] = round(best_score, 4) return existing return self.mysql.add_team_memory( team_id=team_id, content=content, embedding=emb_bytes, importance=importance, category=category, metadata=metadata, ) def search_team_memories(self, team_id: str, query: str, limit: int = 10, min_score: float = 0.0, max_chars_per_memory: int = 0, max_total_chars: int = 0) -> list: query_embedding = self.embedder.embed(query) stored = self.mysql.get_team_memories_with_embeddings(team_id) if not stored: return [] results = search_similar(query_embedding, stored, top_k=limit, min_score=min_score) output = [] total_chars = 0 for memory_id, score, item in results: self.mysql.touch_team_memory(memory_id) content = item["content"] if max_chars_per_memory > 0 and len(content) > max_chars_per_memory: content = content[:max_chars_per_memory] + "..." if max_total_chars > 0: if total_chars + len(content) > max_total_chars: remaining = max_total_chars - total_chars if remaining > 50: content = content[:remaining] + "..." else: break total_chars += len(content) output.append({ "id": memory_id, "content": content, "score": round(score, 4), "importance": item.get("importance", 0.5), "metadata": item.get("metadata"), }) return output def get_recent_team_memories(self, team_id: str, limit: int = 20) -> list: rows = self.mysql.get_team_memories_by_team(team_id, limit) output = [ { "id": r["id"], "content": r["content"], "importance": r.get("importance", 0.5), "category": r.get("category", "general"), "created_at": str(r.get("created_at", "")), } for r in rows ] for item in output: self.mysql.touch_team_memory(item["id"]) return output def update_team_memory(self, memory_id: str, content: str = None, importance: float = None) -> bool: embedding = None if content is not None: emb = self.embedder.embed(content) embedding = embedding_to_bytes(emb) result = self.mysql.update_team_memory( memory_id, content=content, importance=importance, embedding=embedding, ) if result: self.mysql.touch_team_memory(memory_id) return result def delete_team_memory(self, memory_id: str) -> bool: return self.mysql.delete_team_memory(memory_id) # ── Lifecycle ──────────────────────────────────────────────── def compress_working_memories(self, agent_id: str, max_items: int = 20, summary_callback=None) -> str: return self.compressor.compress(agent_id, max_items, summary_callback) # -- Pipeline automation -- # -- Scenario aggregation -- # -- Persona generation -- # -- Fact extraction -- def extract_facts(self, agent_id: str, team_id: str, max_memories: int = 20, delete_after: bool = False): return self.fact_extractor.extract(agent_id, team_id, max_memories, delete_after) def generate_persona(self, agent_id: str, team_id: str, max_items: int = 30): return self.persona_gen.generate(agent_id, team_id, max_items) def get_persona(self, agent_id: str): return self.persona_gen.get_persona(agent_id) def delete_persona(self, persona_id: str, team_id: str): return self.persona_gen.delete_persona(persona_id, team_id) def aggregate_scenarios(self, agent_id: str, team_id: str, max_memories: int = 30): return self.aggregator.aggregate(agent_id, team_id, max_memories) def get_scenarios(self, agent_id: str): return self.aggregator.get_scenarios(agent_id) def delete_scenario(self, scenario_id: str, team_id: str): return self.aggregator.delete_scenario(scenario_id, team_id) def get_pipeline_config(self, team_id: str): return self.pipeline.get_config(team_id) def set_pipeline_config(self, team_id: str, **kwargs): return self.pipeline.set_config(team_id, **kwargs) def check_auto_compress(self, agent_id: str): return self.pipeline.check_after_wm_write(agent_id) def run_auto_cleanup(self, team_id: str = None): return self.pipeline.run_cleanup(team_id) def cleanup_memories(self, team_id: str = None, max_age_days: int = 90, min_importance: float = 0.2) -> dict: return self.cleaner.cleanup(team_id, max_age_days, min_importance) def rebuild_vector_index(self, team_id: str = None) -> dict: """Re-embed all memories (useful after model update).""" # Get all personal memories conn = self.mysql._get_conn() with conn.cursor() as cur: if team_id: cur.execute("SELECT id, content FROM personal_memories WHERE team_id = %s", (team_id,)) else: cur.execute("SELECT id, content FROM personal_memories") personals = cur.fetchall() if team_id: cur.execute("SELECT id, content FROM team_memories WHERE team_id = %s", (team_id,)) else: cur.execute("SELECT id, content FROM team_memories") teams = cur.fetchall() rebuilt = 0 for row in personals: try: emb = self.embedder.embed(row["content"]) self.mysql.update_personal_memory(row["id"], embedding=embedding_to_bytes(emb)) rebuilt += 1 except Exception as e: logger.warning(f"Failed to rebuild personal memory {row['id']}: {e}") for row in teams: try: emb = self.embedder.embed(row["content"]) self.mysql.update_team_memory(row["id"], embedding=embedding_to_bytes(emb)) rebuilt += 1 except Exception as e: logger.warning(f"Failed to rebuild team memory {row['id']}: {e}") return {"rebuilt": rebuilt, "total": len(personals) + len(teams)} # ── Stats ──────────────────────────────────────────────────── def get_stats(self, team_id: str = None) -> dict: stats = self.mysql.get_stats(team_id) stats["redis_ok"] = self.redis.health_check() stats["embedding_ok"] = self.embedder.health_check() return stats