""" Memory layer: Mem0-style + Letta 3-level + CMA consolidation. - core : user profile, project facts (always in system prompt) - recall : recent conversation turns + auto-summary - archival : vector search over all past facts """ import json, math from datetime import datetime, timezone from uuid import UUID import asyncpg from db import get_pool from ollama_client import embed, chat # ── helpers ────────────────────────────────────────────────────────────────── def _vec_str(v: list[float]) -> str: return "[" + ",".join(f"{x:.6f}" for x in v) + "]" # ── archival memory ─────────────────────────────────────────────────────────── async def store_memory(content: str, user_id: str = "default", memory_type: str = "fact", importance: float = 0.7, metadata: dict = None): pool = await get_pool() vec = await embed(content) async with pool.acquire() as conn: await conn.execute(""" INSERT INTO memories (user_id, content, embedding, importance, memory_type, metadata) VALUES ($1, $2, $3::vector, $4, $5, $6) """, user_id, content, _vec_str(vec), importance, memory_type, json.dumps(metadata or {})) async def search_memories(query: str, user_id: str = "default", limit: int = 5) -> list[dict]: pool = await get_pool() vec = await embed(query) async with pool.acquire() as conn: rows = await conn.fetch(""" SELECT id, content, importance, memory_type, metadata, 1 - (embedding <=> $1::vector) AS score FROM memories WHERE user_id = $2 ORDER BY embedding <=> $1::vector LIMIT $3 """, _vec_str(vec), user_id, limit) # bump access stats ids = [str(r["id"]) for r in rows] if ids: await conn.execute(""" UPDATE memories SET access_count = access_count+1, last_accessed = NOW() WHERE id = ANY($1::uuid[]) """, [UUID(i) for i in ids]) return [dict(r) for r in rows] # ── recall (conversation history) ──────────────────────────────────────────── async def save_turn(session_id: str, role: str, content: str): pool = await get_pool() vec = await embed(content[:2000]) # embed first 2k chars only async with pool.acquire() as conn: await conn.execute(""" INSERT INTO conversations (session_id, role, content, embedding) VALUES ($1, $2, $3, $4::vector) """, session_id, role, content, _vec_str(vec)) async def get_recent_turns(session_id: str, limit: int = 20) -> list[dict]: pool = await get_pool() async with pool.acquire() as conn: rows = await conn.fetch(""" SELECT role, content, created_at FROM conversations WHERE session_id = $1 ORDER BY created_at DESC LIMIT $2 """, session_id, limit) return [dict(r) for r in reversed(rows)] async def count_turns(session_id: str) -> int: pool = await get_pool() async with pool.acquire() as conn: return await conn.fetchval( "SELECT COUNT(*) FROM conversations WHERE session_id=$1", session_id) async def summarize_old_turns(session_id: str, keep_last: int = 20): """Summarize turns older than keep_last into a single summary row.""" pool = await get_pool() async with pool.acquire() as conn: total = await conn.fetchval( "SELECT COUNT(*) FROM conversations WHERE session_id=$1", session_id) if total <= keep_last + 10: return old_rows = await conn.fetch(""" SELECT role, content FROM conversations WHERE session_id=$1 ORDER BY created_at ASC LIMIT $2 """, session_id, total - keep_last) text = "\n".join(f"{r['role'].upper()}: {r['content'][:300]}" for r in old_rows) summary = await chat([ {"role": "system", "content": "Summarize this conversation concisely in 3-5 sentences, keeping key technical facts."}, {"role": "user", "content": text} ]) # Store summary, delete old turns await conn.execute(""" INSERT INTO conversation_summaries (session_id, summary, turn_start, turn_end) VALUES ($1, $2, 0, $3) """, session_id, summary, total - keep_last) # Delete summarised rows (oldest ones) await conn.execute(""" DELETE FROM conversations WHERE id IN ( SELECT id FROM conversations WHERE session_id=$1 ORDER BY created_at ASC LIMIT $2 ) """, session_id, total - keep_last) async def get_summary(session_id: str) -> str | None: pool = await get_pool() async with pool.acquire() as conn: row = await conn.fetchrow(""" SELECT summary FROM conversation_summaries WHERE session_id=$1 ORDER BY created_at DESC LIMIT 1 """, session_id) return row["summary"] if row else None # ── extract & store facts after each turn ──────────────────────────────────── EXTRACT_PROMPT = """Extract important facts, preferences, or technical details from this message. Return a JSON array of strings (facts). Max 5 items. If nothing important, return []. Only technical or personal facts worth remembering long-term. Message: {msg}""" async def extract_and_store(content: str, user_id: str = "default"): try: raw = await chat([ {"role": "user", "content": EXTRACT_PROMPT.format(msg=content[:1500])} ]) # find JSON array in response start = raw.find("[") end = raw.rfind("]") + 1 if start == -1 or end == 0: return facts = json.loads(raw[start:end]) for fact in facts[:5]: if isinstance(fact, str) and len(fact) > 10: await store_memory(fact, user_id=user_id, memory_type="fact") except Exception: pass # never crash the main flow # ── CMA: consolidation & forgetting ────────────────────────────────────────── async def consolidate_memories(user_id: str = "default"): """ 1. Apply Ebbinghaus decay to importance. 2. Find near-duplicate memories (cosine > 0.92) and merge them. """ pool = await get_pool() async with pool.acquire() as conn: # decay await conn.execute(""" UPDATE memories SET importance = importance * EXP( -EXTRACT(EPOCH FROM (NOW() - last_accessed)) / 86400.0 / 30.0 ) WHERE user_id = $1 AND importance > 0.05 """, user_id) # archive very stale await conn.execute(""" UPDATE memories SET metadata = metadata || '{"archived":true}'::jsonb WHERE user_id=$1 AND importance < 0.05 AND (metadata->>'archived') IS NULL """, user_id) # find pairs to merge (cosine similarity > 0.92) pairs = await conn.fetch(""" SELECT a.id AS aid, b.id AS bid, a.content AS ac, b.content AS bc, 1-(a.embedding<=>b.embedding) AS sim FROM memories a JOIN memories b ON b.id > a.id WHERE a.user_id=$1 AND b.user_id=$1 AND (a.metadata->>'archived') IS NULL AND (b.metadata->>'archived') IS NULL AND 1-(a.embedding<=>b.embedding) > 0.92 LIMIT 20 """, user_id) for p in pairs: merged = await chat([ {"role": "system", "content": "Merge these two facts into one concise sentence:"}, {"role": "user", "content": f"1. {p['ac']}\n2. {p['bc']}"} ]) merged = merged.strip() if merged: new_imp = min(1.0, p.get("importance_a", 0.7) + 0.1) vec = await embed(merged) await conn.execute(""" INSERT INTO memories (user_id,content,embedding,importance,memory_type) VALUES ($1,$2,$3::vector,$4,'consolidated') """, user_id, merged, _vec_str(vec), new_imp) await conn.execute("DELETE FROM memories WHERE id IN ($1,$2)", p["aid"], p["bid"])