From 9974ce8dd98b757e1ab7fc048cd4caf2c8ff3700 Mon Sep 17 00:00:00 2001
From: jze9 <jze9programer@gmail.com>
Date: Thu, 28 May 2026 20:41:53 +0500
Subject: [PATCH] Add memory management: chat commands + REST API
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

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>
---
 main.py | 80 ++++++++++++++++++++++++++++++++++++++++++++++++++++++++-
 1 file changed, 79 insertions(+), 1 deletion(-)

diff --git a/main.py b/main.py
index d481649..b16e5ee 100644
--- a/main.py
+++ b/main.py
@@ -122,11 +122,52 @@ async def build_context(request: ChatRequest) -> list[dict]:
 
 # ── 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")
 async def chat_completions(req: ChatRequest, bg: BackgroundTasks):
-    messages = await build_context(req)
     user_content = next((m.content for m in reversed(req.messages)
                          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
     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)
         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}")
 async def delete_memory(memory_id: str):
     pool = await _db.get_pool()