fa66a8098c
Extraction: system+user split with explicit SAVE/DON'T SAVE rules, few-shot examples, stricter JSON-only format. Min fact length 15 chars. System prompt: model now knows about its 3-level memory, how to use recalled facts, and that user can control memory via chat commands. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
236 lines
11 KiB
Python
236 lines
11 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_SYSTEM = """Ты — система извлечения фактов из переписки. Твоя задача: найти факты о ПОЛЬЗОВАТЕЛЕ, достойные долгосрочного хранения.
|
||
|
||
СОХРАНЯЙ:
|
||
- Имя, роль, профессия пользователя ("я разработчик", "меня зовут Иван")
|
||
- Язык программирования, фреймворк, стек ("пишу на Go", "использую FastAPI")
|
||
- ОС, редактор, инструменты ("сижу на Arch", "использую Neovim")
|
||
- Характеристики проекта ("проект на Django", "БД PostgreSQL", "деплой на Proxmox")
|
||
- Явные предпочтения ("предпочитаю tabs", "не люблю TypeScript")
|
||
- Технические ограничения ("сервер 4GB RAM", "Python 3.11")
|
||
- Личные факты которые пользователь сообщил о себе
|
||
|
||
НЕ СОХРАНЯЙ:
|
||
- Вопросы (строки со знаком ?)
|
||
- Приветствия ("привет", "hello", "окей")
|
||
- Одиночные слова или короткие реплики без факта
|
||
- То что сказал ассистент — только слова пользователя
|
||
- Запросы без фактов ("помоги мне", "объясни как")
|
||
- Мусор, шутки, эмоции без информации
|
||
|
||
ФОРМАТ ОТВЕТА: только JSON массив строк, без пояснений, без markdown.
|
||
Правильно: ["пользователь пишет на Python 3.12", "использует Neovim"]
|
||
Неправильно: ```json\n["факт"]\n``` или "Вот факты: [...]"
|
||
|
||
Если фактов нет — верни пустой массив: []"""
|
||
|
||
EXTRACT_PROMPT = """Извлеки факты из сообщения пользователя. Максимум 5 фактов.
|
||
|
||
Сообщение пользователя:
|
||
{msg}
|
||
|
||
JSON массив:"""
|
||
|
||
async def extract_and_store(content: str, user_id: str = "default"):
|
||
try:
|
||
raw = await chat([
|
||
{"role": "system", "content": EXTRACT_SYSTEM},
|
||
{"role": "user", "content": EXTRACT_PROMPT.format(msg=content[:1500])}
|
||
])
|
||
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) > 15:
|
||
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"])
|