commit 3aa6ec8b3af702a74307f63e977d423e8ce7c559 Author: jze9 Date: Thu May 28 20:27:16 2026 +0500 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 diff --git a/.env.example b/.env.example new file mode 100644 index 0000000..1a7382a --- /dev/null +++ b/.env.example @@ -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 diff --git a/db.py b/db.py new file mode 100644 index 0000000..50d743c --- /dev/null +++ b/db.py @@ -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); + """) diff --git a/graph.py b/graph.py new file mode 100644 index 0000000..d45700f --- /dev/null +++ b/graph.py @@ -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]) diff --git a/main.py b/main.py new file mode 100644 index 0000000..c32bfae --- /dev/null +++ b/main.py @@ -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"} diff --git a/memory.py b/memory.py new file mode 100644 index 0000000..acd1267 --- /dev/null +++ b/memory.py @@ -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"]) diff --git a/ollama_client.py b/ollama_client.py new file mode 100644 index 0000000..f145237 --- /dev/null +++ b/ollama_client.py @@ -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", "") diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000..cc688f5 --- /dev/null +++ b/requirements.txt @@ -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