added README and RAG

2025-09-10 17:34:22 -04:00
parent 2a6bed386c
commit 96938f4b18
3 changed files with 485 additions and 2 deletions
@@ -0,0 +1,289 @@
+import { serve } from "bun";
+import fs from "node:fs";
+import path from "node:path";
+
+// types
+interface Chunk {
+	id: string;
+	text: string;
+	metadata?: Record<string, unknown>;
+	vector: number[];
+}
+
+interface Collection {
+	name: string;
+	chunks: Chunk[];
+}
+
+interface OllamaChatMessage {
+	role: "system" | "user" | "assistant";
+	content: string;
+}
+
+interface OllamaChatRequest {
+	model?: string;
+	messages: OllamaChatMessage[];
+	stream?: boolean;
+}
+
+interface OllamaChatResponse {
+	message?: OllamaChatMessage;
+	[k: string]: unknown;
+}
+
+interface UpsertInputItem {
+	text: string;
+	metadata?: Record<string, unknown>;
+}
+
+interface OpenAPIObject {
+	openapi: string;
+	info: { title: string; version: string };
+	paths: Record<string, unknown>;
+}
+
+// env
+const PORT: number = Number(process.env.PORT || 8788),
+	HOST: string = process.env.HOST || "0.0.0.0",
+	OLLAMA_BASE: string = process.env.OLLAMA_BASE || "http://localhost:11434",
+	OLLAMA_CHAT_MODEL: string = process.env.OLLAMA_CHAT_MODEL || "llama3.1",
+	OLLAMA_EMBED_MODEL: string = process.env.OLLAMA_EMBED_MODEL || "nomic-embed-text",
+	DATA_DIR: string = process.env.DATA_DIR || path.resolve("./data"),
+	SNAPSHOT: string = path.join(DATA_DIR, "rag.json");
+
+// in-memory db
+const db: Map<string, Collection> = new Map();
+
+// util: smol json persistence
+function ensureDirs(): void {
+	if (!fs.existsSync(DATA_DIR)) fs.mkdirSync(DATA_DIR, { recursive: true });
+}
+
+// you can probably guess
+function loadSnapshot(): void {
+	try {
+		ensureDirs();
+		if (fs.existsSync(SNAPSHOT)) {
+			const raw = fs.readFileSync(SNAPSHOT, "utf8");
+			const obj = JSON.parse(raw || "{}") as Record<string, Collection>;
+			for (const [name, value] of Object.entries(obj)) db.set(name, value);
+		}
+	} catch (e) {
+		console.warn("failed to load snapshot:", e);
+	}
+}
+
+// you can probably guess 2
+function saveSnapshot(): void {
+	try {
+		ensureDirs();
+		const obj = Object.fromEntries(db.entries());
+		fs.writeFileSync(SNAPSHOT, JSON.stringify(obj, null, 2));
+	} catch (e) {
+		console.warn("failed to save snapshot:", e);
+	}
+}
+
+loadSnapshot();
+
+// basic text splitter (recursive by punctuation, then by length)
+function chunkText(text: string, maxLen = 800): string[] {
+	const parts = text
+		.split(/\n{2,}/g)
+		.flatMap(p => p.split(/(?<=[.!?])\s+/g))
+		.flatMap(s => s.length > maxLen ? s.match(new RegExp(`.{1,${maxLen}}`, "g")) || [] : [s])
+		.map(s => s.trim())
+		.filter(Boolean);
+	return parts;
+}
+
+// cosine similarity
+function dot(a: number[], b: number[]): number { let s = 0; for (let i = 0; i < a.length; i++) s += (a[i] || 0) * (b[i] || 0); return s; }
+function norm(a: number[]): number { return Math.sqrt(dot(a, a)); }
+function cosineSim(a: number[], b: number[]): number { const d = dot(a, b), n = norm(a) * norm(b) || 1; return d / n; }
+
+// call ollama embeddings
+async function embedAll(texts: string[]): Promise<number[][]> {
+	const primary = await fetch(`${OLLAMA_BASE}/api/embed`, {
+		method: "POST",
+		headers: { "content-type": "application/json" },
+		body: JSON.stringify({ model: OLLAMA_EMBED_MODEL, input: texts })
+	});
+
+	if (primary.ok) {
+		const j: { embeddings: number[][] } = await primary.json();
+		return j.embeddings;
+	}
+
+	const results: number[][] = [];
+	for (const t of texts) {
+		const r = await fetch(`${OLLAMA_BASE}/api/embeddings`, {
+			method: "POST",
+			headers: { "content-type": "application/json" },
+			body: JSON.stringify({ model: OLLAMA_EMBED_MODEL, prompt: t })
+		});
+
+		if (!r.ok) throw new Error(`embed failed: ${r.status}`);
+
+		const j: { embedding: number[] } = await r.json();
+		results.push(j.embedding);
+	}
+	return results;
+}
+
+// call ollama chat/generate with retrieved context
+async function ollamaChat(req: OllamaChatRequest): Promise<OllamaChatResponse> {
+	const res = await fetch(`${OLLAMA_BASE}/api/chat`, {
+		method: "POST",
+		headers: { "content-type": "application/json" },
+		body: JSON.stringify({ model: req.model || OLLAMA_CHAT_MODEL, messages: req.messages, stream: req.stream })
+	});
+
+	if (!res.ok) throw new Error(`ollama chat failed: ${res.status}`);
+	const j: OllamaChatResponse = await res.json();
+
+	return j;
+}
+
+// openapi for open webui tool integration
+const OPENAPI: OpenAPIObject = {
+	openapi: "3.1.0",
+	info: { title: "RAG Server (Ollama)", version: "1.0.0" },
+	paths: {
+		"/collections": {
+			get: { operationId: "listCollections" },
+			post: { operationId: "createCollection" }
+		},
+		"/upsert": { post: { operationId: "upsert" } },
+		"/query": { post: { operationId: "query" } },
+		"/chat": { post: { operationId: "chat" } }
+	}
+};
+
+// tiny router
+async function json<T = any>(req: Request): Promise<T> { try { return await req.json() as T; } catch { return {} as T; } }
+function sendJson(_res: unknown, status: number, obj: unknown): Response {
+	return new Response(JSON.stringify(obj), { status, headers: { "content-type": "application/json; charset=utf-8" } });
+}
+
+async function handleCollections(req: Request): Promise<Response> {
+	if (req.method === "GET") {
+		return sendJson(null, 200, { collections: Array.from(db.keys()) });
+	}
+
+	if (req.method === "POST") {
+		const body = await json<{ name?: string }>(req),
+			name = String(body?.name || "").trim();
+
+		if (!name) return sendJson(null, 400, { error: "name required" });
+		if (!db.has(name)) db.set(name, { name, chunks: [] });
+
+		saveSnapshot();
+		return sendJson(null, 200, { ok: true });
+	}
+
+	return new Response("not found", { status: 404 });
+}
+
+async function handleUpsert(req: Request): Promise<Response> {
+	const body = await json<{ collection?: string; items?: UpsertInputItem[] }>(req),
+		collection = String(body?.collection || "").trim(),
+		items: UpsertInputItem[] = Array.isArray(body?.items) ? body.items : [];
+
+	if (!collection) return sendJson(null, 400, { error: "collection required" });
+	if (!db.has(collection)) db.set(collection, { name: collection, chunks: [] });
+
+	const col = db.get(collection)!,
+		chunksToIndex: { text: string; metadata?: Record<string, unknown>; _id: string }[] = [];
+
+	for (const it of items) {
+		const parts = chunkText(String(it.text || ""));
+		for (const p of parts) chunksToIndex.push({ text: p, metadata: it.metadata || {}, _id: crypto.randomUUID() });
+	}
+
+	const vecs = await embedAll(chunksToIndex.map(x => x.text));
+	for (let i = 0; i < chunksToIndex.length; i++) {
+		const item = chunksToIndex[i],
+			doc: Chunk = { id: item._id, text: item.text, metadata: item.metadata, vector: vecs[i] };
+
+		col.chunks.push(doc);
+	}
+
+	saveSnapshot();
+	return sendJson(null, 200, { ok: true, indexed: chunksToIndex.length });
+}
+
+async function handleQuery(req: Request): Promise<Response> {
+	const body = await json<{ collection?: string; query?: string; topK?: number }>(req),
+		collection = String(body?.collection || "").trim(),
+		query = String(body?.query || "").trim(),
+		topK = Number(body?.topK || 5);
+
+	if (!collection || !query) return sendJson(null, 400, { error: "collection and query required" });
+
+	const col = db.get(collection);
+	if (!col) return sendJson(null, 404, { error: "collection not found" });
+
+	const [qvec] = await embedAll([query]),
+		scored = col.chunks.map((c) => ({ c, score: cosineSim(qvec, c.vector) }))
+			.sort((a, b) => b.score - a.score)
+			.slice(0, topK)
+			.map(x => ({ id: x.c.id, text: x.c.text, metadata: x.c.metadata, score: x.score }));
+	return sendJson(null, 200, { matches: scored });
+}
+
+async function handleChat(req: Request): Promise<Response> {
+	const body = await json<{ collection?: string; query?: string; topK?: number; model?: string }>(req),
+		collection = String(body?.collection || "").trim(),
+		query = String(body?.query || "").trim(),
+		topK = Number(body?.topK || 5),
+		model = body?.model || OLLAMA_CHAT_MODEL;
+
+	if (!collection || !query) return sendJson(null, 400, { error: "collection and query required" });
+
+	const col = db.get(collection);
+	if (!col) return sendJson(null, 404, { error: "collection not found" });
+
+	const [qvec] = await embedAll([query]),
+		matches = col.chunks.map((c) => ({ c, score: cosineSim(qvec, c.vector) }))
+			.sort((a, b) => b.score - a.score)
+			.slice(0, topK);
+
+	const context = matches.map((m, i) => `[[doc ${i + 1} score=${m.score.toFixed(3)}]]\n${m.c.text}`).join("\n\n"),
+		system: string = `you are a helpful assistant. use ONLY the provided context to answer. if the answer isn't in the context, say you don't know. cite as [doc N].`,
+		user: string = `question: ${query}\n\ncontext:\n${context}`;
+
+	const out = await ollamaChat({ model, messages: [{ role: "system", content: system }, { role: "user", content: user }], stream: false });
+	return sendJson(null, 200, {
+		answer: out?.message?.content || "",
+		citations: matches.map((m, i) => ({ id: m.c.id, score: m.score, text: m.c.text }))
+	});
+}
+
+const pickFunc = (pathname: string) => {
+	switch (pathname) {
+		case "/collections":
+			return handleCollections;
+		case "/upsert":
+			return handleUpsert;
+		case "/query":
+			return handleQuery;
+		case "/chat":
+			return handleChat;
+		default:
+			return undefined;
+	}
+}
+
+const server = serve({
+	port: PORT,
+	hostname: HOST,
+	fetch: async (req: Request): Promise<Response> => {
+		const u = new URL(req.url);
+		if (req.method === "GET" && u.pathname === "/") return new Response("ok");
+		if (req.method === "GET" && u.pathname === "/openapi.json") return sendJson(null, 200, OPENAPI);
+		return pickFunc(u.pathname)?.call(req) || new Response("not found", { status: 404 });
+	}
+});
+
+console.log(`[rag] listening on http://${HOST}:${PORT}`);