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
5
.env.example
Normal file
5
.env.example
Normal file
@@ -0,0 +1,5 @@
|
||||
DATABASE_URL=postgresql://codeai:codeai_secret_2026@192.168.20.17:5432/codeai_db
|
||||
OLLAMA_URL=http://192.168.20.47:11434
|
||||
EMBED_MODEL=nomic-embed-text
|
||||
CHAT_MODEL=qwen2.5-coder:7b-instruct-q4_K_M
|
||||
EMBED_DIM=768
|
||||
73
db.py
Normal file
73
db.py
Normal file
@@ -0,0 +1,73 @@
|
||||
import asyncpg
|
||||
import os
|
||||
from typing import Optional
|
||||
|
||||
pool: Optional[asyncpg.Pool] = None
|
||||
|
||||
async def init_db():
|
||||
global pool
|
||||
pool = await asyncpg.create_pool(os.getenv("DATABASE_URL"), min_size=2, max_size=10)
|
||||
await _create_tables()
|
||||
|
||||
async def get_pool() -> asyncpg.Pool:
|
||||
return pool
|
||||
|
||||
async def _create_tables():
|
||||
async with pool.acquire() as conn:
|
||||
await conn.execute("""
|
||||
CREATE TABLE IF NOT EXISTS memories (
|
||||
id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
|
||||
user_id TEXT DEFAULT 'default',
|
||||
content TEXT NOT NULL,
|
||||
embedding vector(768),
|
||||
importance FLOAT DEFAULT 0.7,
|
||||
access_count INT DEFAULT 0,
|
||||
last_accessed TIMESTAMPTZ DEFAULT NOW(),
|
||||
created_at TIMESTAMPTZ DEFAULT NOW(),
|
||||
memory_type TEXT DEFAULT 'fact',
|
||||
metadata JSONB DEFAULT '{}'
|
||||
);
|
||||
CREATE TABLE IF NOT EXISTS code_nodes (
|
||||
id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
|
||||
project_id TEXT NOT NULL,
|
||||
node_type TEXT NOT NULL,
|
||||
name TEXT NOT NULL,
|
||||
full_name TEXT,
|
||||
content TEXT,
|
||||
embedding vector(768),
|
||||
file_path TEXT,
|
||||
line_start INT,
|
||||
line_end INT,
|
||||
metadata JSONB DEFAULT '{}'
|
||||
);
|
||||
CREATE TABLE IF NOT EXISTS code_edges (
|
||||
id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
|
||||
project_id TEXT NOT NULL,
|
||||
source_id UUID REFERENCES code_nodes(id) ON DELETE CASCADE,
|
||||
target_id UUID REFERENCES code_nodes(id) ON DELETE CASCADE,
|
||||
relation_type TEXT
|
||||
);
|
||||
CREATE TABLE IF NOT EXISTS conversations (
|
||||
id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
|
||||
session_id TEXT NOT NULL,
|
||||
role TEXT NOT NULL,
|
||||
content TEXT NOT NULL,
|
||||
embedding vector(768),
|
||||
created_at TIMESTAMPTZ DEFAULT NOW()
|
||||
);
|
||||
CREATE TABLE IF NOT EXISTS conversation_summaries (
|
||||
id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
|
||||
session_id TEXT NOT NULL,
|
||||
summary TEXT NOT NULL,
|
||||
turn_start INT DEFAULT 0,
|
||||
turn_end INT DEFAULT 0,
|
||||
created_at TIMESTAMPTZ DEFAULT NOW()
|
||||
);
|
||||
CREATE INDEX IF NOT EXISTS idx_mem_embed ON memories USING hnsw (embedding vector_cosine_ops);
|
||||
CREATE INDEX IF NOT EXISTS idx_node_embed ON code_nodes USING hnsw (embedding vector_cosine_ops);
|
||||
CREATE INDEX IF NOT EXISTS idx_conv_sess ON conversations (session_id, created_at DESC);
|
||||
CREATE INDEX IF NOT EXISTS idx_mem_user ON memories (user_id, last_accessed DESC);
|
||||
CREATE INDEX IF NOT EXISTS idx_node_proj ON code_nodes (project_id, node_type);
|
||||
CREATE INDEX IF NOT EXISTS idx_edge_src ON code_edges (source_id);
|
||||
CREATE INDEX IF NOT EXISTS idx_edge_dst ON code_edges (target_id);
|
||||
""")
|
||||
208
graph.py
Normal file
208
graph.py
Normal file
@@ -0,0 +1,208 @@
|
||||
"""
|
||||
GraphRAG: AST-based code graph stored in PostgreSQL.
|
||||
Supports Python (ast module). Multi-hop traversal for architecture queries.
|
||||
"""
|
||||
import ast, os, json
|
||||
from pathlib import Path
|
||||
from uuid import UUID
|
||||
|
||||
from db import get_pool
|
||||
from ollama_client import embed
|
||||
|
||||
def _vec_str(v: list[float]) -> str:
|
||||
return "[" + ",".join(f"{x:.6f}" for x in v) + "]"
|
||||
|
||||
# ── AST Parser ────────────────────────────────────────────────────────────────
|
||||
|
||||
class PythonGraphBuilder(ast.NodeVisitor):
|
||||
def __init__(self, file_path: str, source: str):
|
||||
self.file_path = file_path
|
||||
self.source_lines = source.splitlines()
|
||||
self.nodes: list[dict] = []
|
||||
self.edges: list[dict] = []
|
||||
self._class_stack: list[str] = []
|
||||
self._func_stack: list[str] = []
|
||||
|
||||
def _full_name(self, name: str) -> str:
|
||||
parts = self._class_stack + self._func_stack + [name]
|
||||
return ".".join(parts)
|
||||
|
||||
def _src(self, node) -> str:
|
||||
try:
|
||||
return ast.get_source_segment("\n".join(self.source_lines), node) or ""
|
||||
except Exception:
|
||||
return ""
|
||||
|
||||
def visit_ClassDef(self, node):
|
||||
full = self._full_name(node.name)
|
||||
self.nodes.append({"type": "class", "name": node.name, "full_name": full,
|
||||
"file": self.file_path, "line_start": node.lineno,
|
||||
"line_end": node.end_lineno, "content": self._src(node)})
|
||||
self._class_stack.append(node.name)
|
||||
self.generic_visit(node)
|
||||
self._class_stack.pop()
|
||||
|
||||
def visit_FunctionDef(self, node):
|
||||
full = self._full_name(node.name)
|
||||
parent = ".".join(self._class_stack + self._func_stack) or None
|
||||
self.nodes.append({"type": "function", "name": node.name, "full_name": full,
|
||||
"file": self.file_path, "line_start": node.lineno,
|
||||
"line_end": node.end_lineno, "content": self._src(node)})
|
||||
if parent:
|
||||
self.edges.append({"source": parent, "target": full, "rel": "has_method"})
|
||||
self._func_stack.append(node.name)
|
||||
self.generic_visit(node)
|
||||
self._func_stack.pop()
|
||||
|
||||
visit_AsyncFunctionDef = visit_FunctionDef
|
||||
|
||||
def visit_Call(self, node):
|
||||
caller = ".".join(self._class_stack + self._func_stack)
|
||||
if not caller:
|
||||
self.generic_visit(node)
|
||||
return
|
||||
if isinstance(node.func, ast.Attribute) and hasattr(node.func.value, "id"):
|
||||
callee = f"{node.func.value.id}.{node.func.attr}"
|
||||
elif isinstance(node.func, ast.Name):
|
||||
callee = node.func.id
|
||||
else:
|
||||
self.generic_visit(node)
|
||||
return
|
||||
self.edges.append({"source": caller, "target": callee, "rel": "calls"})
|
||||
self.generic_visit(node)
|
||||
|
||||
def visit_Import(self, node):
|
||||
file_node = self.file_path
|
||||
for alias in node.names:
|
||||
self.edges.append({"source": file_node, "target": alias.name, "rel": "imports"})
|
||||
self.generic_visit(node)
|
||||
|
||||
def visit_ImportFrom(self, node):
|
||||
if node.module:
|
||||
file_node = self.file_path
|
||||
self.edges.append({"source": file_node, "target": node.module, "rel": "imports"})
|
||||
self.generic_visit(node)
|
||||
|
||||
|
||||
def parse_python_file(path: str) -> tuple[list[dict], list[dict]]:
|
||||
try:
|
||||
source = Path(path).read_text(encoding="utf-8", errors="ignore")
|
||||
tree = ast.parse(source)
|
||||
builder = PythonGraphBuilder(path, source)
|
||||
builder.visit(tree)
|
||||
# add file-level node
|
||||
builder.nodes.insert(0, {"type": "file", "name": path, "full_name": path,
|
||||
"file": path, "line_start": 1,
|
||||
"line_end": len(source.splitlines()),
|
||||
"content": source[:500]})
|
||||
return builder.nodes, builder.edges
|
||||
except SyntaxError:
|
||||
return [], []
|
||||
|
||||
|
||||
# ── Indexer ───────────────────────────────────────────────────────────────────
|
||||
|
||||
async def index_project(project_id: str, root_dir: str,
|
||||
extensions: list[str] = None):
|
||||
if extensions is None:
|
||||
extensions = [".py"]
|
||||
pool = await get_pool()
|
||||
# clear old data
|
||||
async with pool.acquire() as conn:
|
||||
await conn.execute("DELETE FROM code_edges WHERE project_id=$1", project_id)
|
||||
await conn.execute("DELETE FROM code_nodes WHERE project_id=$1", project_id)
|
||||
|
||||
node_name_to_id: dict[str, UUID] = {}
|
||||
|
||||
for ext in extensions:
|
||||
for path in Path(root_dir).rglob(f"*{ext}"):
|
||||
if ".git" in str(path) or "venv" in str(path) or "__pycache__" in str(path):
|
||||
continue
|
||||
nodes, edges = parse_python_file(str(path))
|
||||
for n in nodes:
|
||||
content_for_embed = f"{n['type']} {n['full_name']}\n{n['content'][:800]}"
|
||||
vec = await embed(content_for_embed)
|
||||
async with pool.acquire() as conn:
|
||||
uid = await conn.fetchval("""
|
||||
INSERT INTO code_nodes
|
||||
(project_id, node_type, name, full_name, content,
|
||||
embedding, file_path, line_start, line_end)
|
||||
VALUES ($1,$2,$3,$4,$5,$6::vector,$7,$8,$9)
|
||||
RETURNING id
|
||||
""", project_id, n["type"], n["name"], n["full_name"],
|
||||
n["content"][:2000], _vec_str(vec),
|
||||
n["file"], n.get("line_start"), n.get("line_end"))
|
||||
node_name_to_id[n["full_name"]] = uid
|
||||
|
||||
# store edges (resolve names to IDs later in a second pass)
|
||||
for e in edges:
|
||||
async with pool.acquire() as conn:
|
||||
src_id = node_name_to_id.get(e["source"])
|
||||
tgt_id = node_name_to_id.get(e["target"])
|
||||
if src_id and tgt_id:
|
||||
await conn.execute("""
|
||||
INSERT INTO code_edges (project_id, source_id, target_id, relation_type)
|
||||
VALUES ($1,$2,$3,$4) ON CONFLICT DO NOTHING
|
||||
""", project_id, src_id, tgt_id, e["rel"])
|
||||
|
||||
return {"indexed_nodes": len(node_name_to_id)}
|
||||
|
||||
|
||||
# ── Graph search ──────────────────────────────────────────────────────────────
|
||||
|
||||
async def search_code(query: str, project_id: str, limit: int = 5) -> list[dict]:
|
||||
"""Vector search over code nodes."""
|
||||
pool = await get_pool()
|
||||
vec = await embed(query)
|
||||
async with pool.acquire() as conn:
|
||||
rows = await conn.fetch("""
|
||||
SELECT id, node_type, full_name, content, file_path, line_start,
|
||||
1-(embedding<=>$1::vector) AS score
|
||||
FROM code_nodes WHERE project_id=$2
|
||||
ORDER BY embedding<=>$1::vector LIMIT $3
|
||||
""", _vec_str(vec), project_id, limit)
|
||||
return [dict(r) for r in rows]
|
||||
|
||||
|
||||
async def get_neighbors(node_id: UUID, project_id: str, hops: int = 2) -> list[dict]:
|
||||
"""Multi-hop graph traversal starting from node_id."""
|
||||
pool = await get_pool()
|
||||
async with pool.acquire() as conn:
|
||||
# recursive CTE for graph traversal
|
||||
rows = await conn.fetch("""
|
||||
WITH RECURSIVE graph AS (
|
||||
SELECT id, node_type, full_name, content, file_path, 0 AS depth
|
||||
FROM code_nodes WHERE id=$1
|
||||
UNION
|
||||
SELECT n.id, n.node_type, n.full_name, n.content, n.file_path,
|
||||
g.depth+1
|
||||
FROM code_nodes n
|
||||
JOIN code_edges e ON (e.target_id=n.id AND e.source_id=g.id)
|
||||
OR (e.source_id=n.id AND e.target_id=g.id)
|
||||
JOIN graph g ON TRUE
|
||||
WHERE g.depth < $2 AND n.project_id=$3
|
||||
)
|
||||
SELECT DISTINCT id, node_type, full_name, content, file_path, depth
|
||||
FROM graph ORDER BY depth
|
||||
LIMIT 30
|
||||
""", node_id, hops, project_id)
|
||||
return [dict(r) for r in rows]
|
||||
|
||||
|
||||
async def graphrag_context(query: str, project_id: str) -> str:
|
||||
"""Full GraphRAG: vector search + multi-hop expansion."""
|
||||
seed_nodes = await search_code(query, project_id, limit=3)
|
||||
if not seed_nodes:
|
||||
return ""
|
||||
all_nodes: dict[str, dict] = {str(n["id"]): n for n in seed_nodes}
|
||||
for n in seed_nodes:
|
||||
neighbors = await get_neighbors(n["id"], project_id, hops=2)
|
||||
for nb in neighbors:
|
||||
all_nodes[str(nb["id"])] = nb
|
||||
parts = []
|
||||
for n in sorted(all_nodes.values(), key=lambda x: x.get("depth", 0)):
|
||||
parts.append(
|
||||
f"[{n['node_type']}] {n['full_name']} ({n.get('file_path','')}"
|
||||
f":{n.get('line_start','')})\n{(n.get('content') or '')[:400]}"
|
||||
)
|
||||
return "\n\n".join(parts[:10])
|
||||
194
main.py
Normal file
194
main.py
Normal file
@@ -0,0 +1,194 @@
|
||||
"""
|
||||
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"]
|
||||
|
||||
# ── 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 = (
|
||||
"You are an expert coding assistant with deep knowledge of the project architecture. "
|
||||
"Use the provided memory and code context to give precise, informed answers. "
|
||||
"Always reason about the full architecture, not just isolated snippets."
|
||||
)
|
||||
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 ─────────────────────────────────────────────────────────────
|
||||
|
||||
@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"), "")
|
||||
|
||||
if req.stream:
|
||||
async def gen() -> AsyncIterator[bytes]:
|
||||
full_response = []
|
||||
async for chunk in oll.chat_stream(messages, model=req.model):
|
||||
full_response.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"
|
||||
# save after stream
|
||||
response_text = "".join(full_response)
|
||||
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)
|
||||
|
||||
return StreamingResponse(gen(), media_type="text/event-stream")
|
||||
|
||||
response_text = await oll.chat(messages, model=req.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.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"}
|
||||
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"])
|
||||
36
ollama_client.py
Normal file
36
ollama_client.py
Normal file
@@ -0,0 +1,36 @@
|
||||
import httpx
|
||||
import os
|
||||
from typing import AsyncIterator
|
||||
|
||||
OLLAMA = os.getenv("OLLAMA_URL", "http://192.168.20.47:11434")
|
||||
EMBED_MODEL = os.getenv("EMBED_MODEL", "nomic-embed-text")
|
||||
CHAT_MODEL = os.getenv("CHAT_MODEL", "qwen2.5-coder:7b-instruct-q4_K_M")
|
||||
|
||||
async def embed(text: str) -> list[float]:
|
||||
async with httpx.AsyncClient(timeout=60) as c:
|
||||
r = await c.post(f"{OLLAMA}/api/embeddings",
|
||||
json={"model": EMBED_MODEL, "prompt": text})
|
||||
r.raise_for_status()
|
||||
return r.json()["embedding"]
|
||||
|
||||
async def chat(messages: list[dict], model: str = None, stream: bool = False):
|
||||
async with httpx.AsyncClient(timeout=300) as c:
|
||||
r = await c.post(f"{OLLAMA}/api/chat",
|
||||
json={"model": model or CHAT_MODEL,
|
||||
"messages": messages,
|
||||
"stream": False})
|
||||
r.raise_for_status()
|
||||
return r.json()["message"]["content"]
|
||||
|
||||
async def chat_stream(messages: list[dict], model: str = None) -> AsyncIterator[str]:
|
||||
async with httpx.AsyncClient(timeout=300) as c:
|
||||
async with c.stream("POST", f"{OLLAMA}/api/chat",
|
||||
json={"model": model or CHAT_MODEL,
|
||||
"messages": messages,
|
||||
"stream": True}) as r:
|
||||
async for line in r.aiter_lines():
|
||||
if line:
|
||||
import json
|
||||
data = json.loads(line)
|
||||
if not data.get("done"):
|
||||
yield data.get("message", {}).get("content", "")
|
||||
9
requirements.txt
Normal file
9
requirements.txt
Normal file
@@ -0,0 +1,9 @@
|
||||
fastapi>=0.115.0
|
||||
uvicorn[standard]>=0.30.0
|
||||
httpx>=0.27.0
|
||||
asyncpg>=0.29.0
|
||||
pgvector>=0.3.5
|
||||
networkx>=3.3
|
||||
apscheduler>=3.10.4
|
||||
python-dotenv>=1.0.0
|
||||
pydantic>=2.7.0
|
||||
Reference in New Issue
Block a user