Add memory management: chat commands + REST API
Chat commands (intercept before LLM):
"запомни что X" / "remember X" — saves with importance 0.9
"забудь про X" / "forget X" — deletes by vector similarity
REST API:
POST /v1/memories — add memory manually
PATCH /v1/memories/{id} — update content or importance
DELETE /v1/memories/{id} — delete by id (existing)
Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
This commit is contained in:
80
main.py
80
main.py
@@ -122,11 +122,52 @@ async def build_context(request: ChatRequest) -> list[dict]:
|
|||||||
|
|
||||||
# ── chat endpoint ─────────────────────────────────────────────────────────────
|
# ── chat endpoint ─────────────────────────────────────────────────────────────
|
||||||
|
|
||||||
|
FORGET_PREFIXES = ("забудь про ", "забудь ", "forget ")
|
||||||
|
REMEMBER_PREFIXES = ("запомни что ", "запомни ", "remember ")
|
||||||
|
|
||||||
|
async def handle_memory_command(user_content: str, user_id: str) -> str | None:
|
||||||
|
low = user_content.lower().strip()
|
||||||
|
for prefix in FORGET_PREFIXES:
|
||||||
|
if low.startswith(prefix):
|
||||||
|
topic = user_content[len(prefix):].strip()
|
||||||
|
pool = await _db.get_pool()
|
||||||
|
vec = await oll.embed(topic)
|
||||||
|
from memory import _vec_str
|
||||||
|
async with pool.acquire() as conn:
|
||||||
|
deleted = await conn.fetchval("""
|
||||||
|
DELETE FROM memories
|
||||||
|
WHERE user_id=$1
|
||||||
|
AND 1-(embedding<=>$2::vector) > 0.60
|
||||||
|
RETURNING content
|
||||||
|
""", user_id, _vec_str(vec))
|
||||||
|
if deleted:
|
||||||
|
return f"Забыл: «{deleted}»"
|
||||||
|
return f"Не нашёл ничего похожего на «{topic}»"
|
||||||
|
for prefix in REMEMBER_PREFIXES:
|
||||||
|
if low.startswith(prefix):
|
||||||
|
fact = user_content[len(prefix):].strip()
|
||||||
|
await mem.store_memory(fact, user_id=user_id, memory_type="manual", importance=0.9)
|
||||||
|
return f"Запомнил: «{fact}»"
|
||||||
|
return None
|
||||||
|
|
||||||
@app.post("/v1/chat/completions")
|
@app.post("/v1/chat/completions")
|
||||||
async def chat_completions(req: ChatRequest, bg: BackgroundTasks):
|
async def chat_completions(req: ChatRequest, bg: BackgroundTasks):
|
||||||
messages = await build_context(req)
|
|
||||||
user_content = next((m.content for m in reversed(req.messages)
|
user_content = next((m.content for m in reversed(req.messages)
|
||||||
if m.role == "user"), "")
|
if m.role == "user"), "")
|
||||||
|
|
||||||
|
# intercept memory commands before building context
|
||||||
|
cmd_response = await handle_memory_command(user_content, req.user_id)
|
||||||
|
if cmd_response:
|
||||||
|
return {
|
||||||
|
"id": "cmpl-cmd",
|
||||||
|
"object": "chat.completion",
|
||||||
|
"model": req.model or os.getenv("CHAT_MODEL"),
|
||||||
|
"choices": [{"index": 0,
|
||||||
|
"message": {"role": "assistant", "content": cmd_response},
|
||||||
|
"finish_reason": "stop"}]
|
||||||
|
}
|
||||||
|
|
||||||
|
messages = await build_context(req)
|
||||||
# use real Ollama model, not our proxy name
|
# use real Ollama model, not our proxy name
|
||||||
ollama_model = None if req.model == "code-memory" else req.model
|
ollama_model = None if req.model == "code-memory" else req.model
|
||||||
|
|
||||||
@@ -200,6 +241,43 @@ async def list_memories(user_id: str = "default", limit: int = 50):
|
|||||||
""", user_id, limit)
|
""", user_id, limit)
|
||||||
return [dict(r) for r in rows]
|
return [dict(r) for r in rows]
|
||||||
|
|
||||||
|
@app.post("/v1/memories")
|
||||||
|
async def create_memory(body: dict, user_id: str = "default"):
|
||||||
|
content = body.get("content", "").strip()
|
||||||
|
if not content:
|
||||||
|
raise HTTPException(400, "content required")
|
||||||
|
importance = float(body.get("importance", 0.9))
|
||||||
|
memory_type = body.get("memory_type", "manual")
|
||||||
|
await mem.store_memory(content, user_id=user_id,
|
||||||
|
memory_type=memory_type, importance=importance)
|
||||||
|
return {"status": "stored", "content": content}
|
||||||
|
|
||||||
|
@app.patch("/v1/memories/{memory_id}")
|
||||||
|
async def update_memory(memory_id: str, body: dict):
|
||||||
|
import uuid as _uuid
|
||||||
|
pool = await _db.get_pool()
|
||||||
|
uid = _uuid.UUID(memory_id)
|
||||||
|
async with pool.acquire() as conn:
|
||||||
|
row = await conn.fetchrow("SELECT * FROM memories WHERE id=$1", uid)
|
||||||
|
if not row:
|
||||||
|
raise HTTPException(404, "memory not found")
|
||||||
|
new_content = body.get("content", row["content"])
|
||||||
|
new_importance = float(body.get("importance", row["importance"]))
|
||||||
|
if new_content != row["content"]:
|
||||||
|
vec = await oll.embed(new_content)
|
||||||
|
from memory import _vec_str
|
||||||
|
await conn.execute("""
|
||||||
|
UPDATE memories SET content=$1, embedding=$2::vector,
|
||||||
|
importance=$3, last_accessed=NOW()
|
||||||
|
WHERE id=$4
|
||||||
|
""", new_content, _vec_str(vec), new_importance, uid)
|
||||||
|
else:
|
||||||
|
await conn.execute("""
|
||||||
|
UPDATE memories SET importance=$1, last_accessed=NOW()
|
||||||
|
WHERE id=$2
|
||||||
|
""", new_importance, uid)
|
||||||
|
return {"updated": memory_id}
|
||||||
|
|
||||||
@app.delete("/v1/memories/{memory_id}")
|
@app.delete("/v1/memories/{memory_id}")
|
||||||
async def delete_memory(memory_id: str):
|
async def delete_memory(memory_id: str):
|
||||||
pool = await _db.get_pool()
|
pool = await _db.get_pool()
|
||||||
|
|||||||
Reference in New Issue
Block a user