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>
304 lines
13 KiB
Python
304 lines
13 KiB
Python
"""
|
||
code-memory: OpenAI-compatible proxy with GraphRAG + persistent memory.
|
||
Endpoints:
|
||
POST /v1/chat/completions — drop-in Ollama replacement with memory
|
||
POST /v1/index — index a code repository
|
||
GET /v1/memories — list memories
|
||
DELETE /v1/memories/{id} — delete a memory
|
||
GET /health — health check
|
||
"""
|
||
import json, os, asyncio
|
||
from contextlib import asynccontextmanager
|
||
from typing import AsyncIterator
|
||
|
||
import asyncpg
|
||
from fastapi import FastAPI, HTTPException, BackgroundTasks
|
||
from fastapi.responses import StreamingResponse
|
||
from pydantic import BaseModel, Field
|
||
from dotenv import load_dotenv
|
||
|
||
load_dotenv()
|
||
|
||
import db as _db
|
||
import memory as mem
|
||
import graph as gph
|
||
import ollama_client as oll
|
||
from apscheduler.schedulers.asyncio import AsyncIOScheduler
|
||
|
||
# ── lifespan ──────────────────────────────────────────────────────────────────
|
||
|
||
scheduler = AsyncIOScheduler()
|
||
|
||
@asynccontextmanager
|
||
async def lifespan(app: FastAPI):
|
||
await _db.init_db()
|
||
# CMA consolidation runs every 6 hours
|
||
scheduler.add_job(mem.consolidate_memories, "interval", hours=6,
|
||
kwargs={"user_id": "default"})
|
||
scheduler.start()
|
||
yield
|
||
scheduler.shutdown()
|
||
|
||
app = FastAPI(title="code-memory", lifespan=lifespan)
|
||
|
||
# ── models ────────────────────────────────────────────────────────────────────
|
||
|
||
class Message(BaseModel):
|
||
role: str
|
||
content: str
|
||
|
||
class ChatRequest(BaseModel):
|
||
model: str = Field(default=None)
|
||
messages: list[Message]
|
||
stream: bool = False
|
||
session_id: str = "default"
|
||
user_id: str = "default"
|
||
project_id: str | None = None
|
||
|
||
class IndexRequest(BaseModel):
|
||
project_id: str
|
||
root_dir: str
|
||
extensions: list[str] = [".py"]
|
||
|
||
# ── models list endpoint ──────────────────────────────────────────────────────
|
||
|
||
@app.get("/v1/models")
|
||
async def list_models():
|
||
return {
|
||
"object": "list",
|
||
"data": [{
|
||
"id": "code-memory",
|
||
"object": "model",
|
||
"owned_by": "code-memory",
|
||
"created": 1700000000,
|
||
}]
|
||
}
|
||
|
||
# ── context assembly ──────────────────────────────────────────────────────────
|
||
|
||
async def build_context(request: ChatRequest) -> list[dict]:
|
||
user_msg = next((m.content for m in reversed(request.messages)
|
||
if m.role == "user"), "")
|
||
|
||
# 1. Relevant archival memories
|
||
memories = await mem.search_memories(user_msg, user_id=request.user_id, limit=6)
|
||
mem_block = ""
|
||
if memories:
|
||
mem_block = "## Relevant memories\n" + "\n".join(
|
||
f"- [{m['memory_type']}] {m['content']}" for m in memories
|
||
if m.get("score", 0) > 0.4
|
||
)
|
||
|
||
# 2. GraphRAG code context
|
||
code_block = ""
|
||
if request.project_id:
|
||
ctx = await gph.graphrag_context(user_msg, request.project_id)
|
||
if ctx:
|
||
code_block = f"## Code architecture context\n{ctx}"
|
||
|
||
# 3. Conversation summary (if exists)
|
||
summary = await mem.get_summary(request.session_id)
|
||
summary_block = f"## Earlier conversation summary\n{summary}" if summary else ""
|
||
|
||
# 4. Assemble system prompt
|
||
base_system = """Ты — AI-ассистент по программированию с постоянной памятью.
|
||
|
||
У тебя есть три уровня памяти:
|
||
1. Долгосрочная память — факты о пользователе и проектах (показаны ниже в "Relevant memories")
|
||
2. История диалога — последние 20 сообщений текущей сессии
|
||
3. Контекст кода — граф архитектуры проекта (если проиндексирован)
|
||
|
||
Используй память чтобы давать точные ответы с учётом контекста пользователя. \
|
||
Если в памяти есть релевантный факт — опирайся на него, не переспрашивай то что уже известно.
|
||
|
||
Пользователь может управлять твоей памятью прямо в чате:
|
||
- "запомни что X" — ты сохранишь X с высоким приоритетом
|
||
- "забудь про X" — ты удалишь похожие воспоминания
|
||
- "что ты помнишь?" — покажи список активных воспоминаний через GET /v1/memories"""
|
||
context_parts = [p for p in [base_system, summary_block, mem_block, code_block] if p]
|
||
system_content = "\n\n".join(context_parts)
|
||
|
||
# 5. Build messages: system + recent history + current
|
||
recent = await mem.get_recent_turns(request.session_id, limit=20)
|
||
messages = [{"role": "system", "content": system_content}]
|
||
messages += [{"role": r["role"], "content": r["content"]} for r in recent]
|
||
# add current user messages (skip if already in history)
|
||
for m in request.messages:
|
||
if not any(h["role"] == m.role and h["content"] == m.content for h in recent[-2:]):
|
||
messages.append({"role": m.role, "content": m.content})
|
||
|
||
return messages
|
||
|
||
# ── 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):
|
||
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
|
||
|
||
if req.stream:
|
||
collected: list[str] = []
|
||
|
||
async def gen() -> AsyncIterator[bytes]:
|
||
try:
|
||
async for chunk in oll.chat_stream(messages, model=ollama_model):
|
||
collected.append(chunk)
|
||
data = json.dumps({"choices": [{"delta": {"content": chunk},
|
||
"finish_reason": None}]})
|
||
yield f"data: {data}\n\n".encode()
|
||
yield b"data: [DONE]\n\n"
|
||
except Exception as e:
|
||
yield f"data: {json.dumps({'error': str(e)})}\n\n".encode()
|
||
|
||
async def _save_stream():
|
||
response_text = "".join(collected)
|
||
await mem.save_turn(req.session_id, "user", user_content)
|
||
await mem.save_turn(req.session_id, "assistant", response_text)
|
||
await mem.extract_and_store(user_content, user_id=req.user_id)
|
||
await mem.extract_and_store(response_text, user_id=req.user_id)
|
||
await mem.summarize_old_turns(req.session_id)
|
||
|
||
bg.add_task(_save_stream)
|
||
return StreamingResponse(gen(), media_type="text/event-stream")
|
||
|
||
response_text = await oll.chat(messages, model=ollama_model)
|
||
|
||
# background: save turns + extract facts
|
||
async def _save():
|
||
await mem.save_turn(req.session_id, "user", user_content)
|
||
await mem.save_turn(req.session_id, "assistant", response_text)
|
||
await mem.extract_and_store(user_content, user_id=req.user_id)
|
||
await mem.extract_and_store(response_text, user_id=req.user_id)
|
||
await mem.summarize_old_turns(req.session_id)
|
||
|
||
bg.add_task(_save)
|
||
|
||
return {
|
||
"id": "cmpl-1",
|
||
"object": "chat.completion",
|
||
"model": req.model or os.getenv("CHAT_MODEL"),
|
||
"choices": [{"index": 0,
|
||
"message": {"role": "assistant", "content": response_text},
|
||
"finish_reason": "stop"}]
|
||
}
|
||
|
||
# ── index endpoint ────────────────────────────────────────────────────────────
|
||
|
||
@app.post("/v1/index")
|
||
async def index_code(req: IndexRequest, bg: BackgroundTasks):
|
||
if not os.path.isdir(req.root_dir):
|
||
raise HTTPException(404, f"Directory not found: {req.root_dir}")
|
||
bg.add_task(gph.index_project, req.project_id, req.root_dir, req.extensions)
|
||
return {"status": "indexing started", "project_id": req.project_id}
|
||
|
||
# ── memory endpoints ──────────────────────────────────────────────────────────
|
||
|
||
@app.get("/v1/memories")
|
||
async def list_memories(user_id: str = "default", limit: int = 50):
|
||
pool = await _db.get_pool()
|
||
async with pool.acquire() as conn:
|
||
rows = await conn.fetch("""
|
||
SELECT id, content, importance, memory_type, access_count,
|
||
last_accessed, created_at
|
||
FROM memories WHERE user_id=$1
|
||
AND (metadata->>'archived') IS NULL
|
||
ORDER BY importance DESC, last_accessed DESC LIMIT $2
|
||
""", 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()
|
||
async with pool.acquire() as conn:
|
||
await conn.execute("DELETE FROM memories WHERE id=$1",
|
||
__import__("uuid").UUID(memory_id))
|
||
return {"deleted": memory_id}
|
||
|
||
@app.get("/health")
|
||
async def health():
|
||
pool = await _db.get_pool()
|
||
async with pool.acquire() as conn:
|
||
await conn.fetchval("SELECT 1")
|
||
return {"status": "ok", "db": "connected"}
|