74c3fac2cc
Before inserting, check if a near-identical memory exists (cosine > 0.85). If found, skip insert. Prevents flooding the DB with repeated facts from the same conversation. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
209 lines
9.0 KiB
Python
209 lines
9.0 KiB
Python
"""
|
|
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:
|
|
# skip if near-duplicate already exists (cosine > 0.85)
|
|
existing = await conn.fetchval("""
|
|
SELECT id FROM memories
|
|
WHERE user_id = $1
|
|
AND (metadata->>'archived') IS NULL
|
|
AND 1 - (embedding <=> $2::vector) > 0.85
|
|
LIMIT 1
|
|
""", user_id, _vec_str(vec))
|
|
if existing:
|
|
return
|
|
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"])
|