Initial commit: code-memory service
OpenAI-compatible FastAPI proxy with GraphRAG + persistent memory. Includes 3-level Letta memory, CMA consolidation, AST-based code indexing. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
This commit is contained in:
jze9
198
memory.py
Normal file
198
memory.py
Normal file
@@ -0,0 +1,198 @@
|
||||
"""
|
||||
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"])
|
||||
Reference in New Issue
Block a user