|
| 1 | +// SQLite-backed Vectorize adapter for self-host RAG (#979). Implements the Cloudflare `Vectorize` binding |
| 2 | +// surface (upsert / query / deleteByIds) that gittensory's RAG (reviewVectorAdapter) wraps, backed by a |
| 3 | +// SQLite table with brute-force cosine similarity. For a repo's worth of chunks (hundreds–few-thousand |
| 4 | +// vectors per namespace) this is fast enough; namespaces (one per repo) keep each query's candidate set |
| 5 | +// small. Embeddings come from the OpenAI-compatible AI adapter's /embeddings path (e.g. Ollama bge-m3, 1024-d). |
| 6 | +import type { SqliteDriver } from "./d1-adapter"; |
| 7 | + |
| 8 | +const TABLE = "_selfhost_vectors"; |
| 9 | +const DDL = ` |
| 10 | +CREATE TABLE IF NOT EXISTS ${TABLE} ( |
| 11 | + id TEXT PRIMARY KEY, |
| 12 | + namespace TEXT NOT NULL DEFAULT '', |
| 13 | + embedding TEXT NOT NULL, |
| 14 | + metadata TEXT |
| 15 | +); |
| 16 | +CREATE INDEX IF NOT EXISTS ${TABLE}_ns ON ${TABLE}(namespace);`; |
| 17 | + |
| 18 | +interface VectorRecord { |
| 19 | + id: string; |
| 20 | + values: number[]; |
| 21 | + namespace?: string; |
| 22 | + metadata?: Record<string, unknown>; |
| 23 | +} |
| 24 | +interface QueryOptions { |
| 25 | + topK?: number; |
| 26 | + namespace?: string; |
| 27 | + returnMetadata?: string; |
| 28 | +} |
| 29 | +interface Match { |
| 30 | + id: string; |
| 31 | + score: number; |
| 32 | + metadata?: Record<string, unknown>; |
| 33 | +} |
| 34 | + |
| 35 | +export function cosineSimilarity(a: number[], b: number[]): number { |
| 36 | + let dot = 0; |
| 37 | + let na = 0; |
| 38 | + let nb = 0; |
| 39 | + const n = Math.min(a.length, b.length); |
| 40 | + for (let i = 0; i < n; i += 1) { |
| 41 | + const x = a[i] as number; |
| 42 | + const y = b[i] as number; |
| 43 | + dot += x * y; |
| 44 | + na += x * x; |
| 45 | + nb += y * y; |
| 46 | + } |
| 47 | + const denom = Math.sqrt(na) * Math.sqrt(nb); |
| 48 | + return denom === 0 ? 0 : dot / denom; |
| 49 | +} |
| 50 | + |
| 51 | +export function createSqliteVectorize(driver: SqliteDriver): Vectorize { |
| 52 | + driver.exec(DDL); |
| 53 | + const adapter = { |
| 54 | + async upsert(vectors: VectorRecord[]): Promise<{ count: number; ids: string[] }> { |
| 55 | + for (const v of vectors) { |
| 56 | + driver.query( |
| 57 | + `INSERT INTO ${TABLE} (id, namespace, embedding, metadata) VALUES (?,?,?,?) |
| 58 | + ON CONFLICT(id) DO UPDATE SET namespace=excluded.namespace, embedding=excluded.embedding, metadata=excluded.metadata`, |
| 59 | + [v.id, v.namespace ?? "", JSON.stringify(v.values), v.metadata ? JSON.stringify(v.metadata) : null], |
| 60 | + ); |
| 61 | + } |
| 62 | + return { count: vectors.length, ids: vectors.map((v) => v.id) }; |
| 63 | + }, |
| 64 | + async query(vector: number[], opts: QueryOptions): Promise<{ matches: Match[] }> { |
| 65 | + const { rows } = opts.namespace |
| 66 | + ? driver.query(`SELECT id, embedding, metadata FROM ${TABLE} WHERE namespace=?`, [opts.namespace]) |
| 67 | + : driver.query(`SELECT id, embedding, metadata FROM ${TABLE}`, []); |
| 68 | + const scored: Match[] = rows.map((r) => { |
| 69 | + const values = JSON.parse(r.embedding as string) as number[]; |
| 70 | + const metadata = r.metadata ? (JSON.parse(r.metadata as string) as Record<string, unknown>) : undefined; |
| 71 | + const score = cosineSimilarity(vector, values); |
| 72 | + return metadata === undefined ? { id: r.id as string, score } : { id: r.id as string, score, metadata }; |
| 73 | + }); |
| 74 | + scored.sort((a, b) => b.score - a.score); |
| 75 | + return { matches: scored.slice(0, opts.topK ?? 12) }; |
| 76 | + }, |
| 77 | + async deleteByIds(ids: string[]): Promise<{ count: number }> { |
| 78 | + for (let i = 0; i < ids.length; i += 90) { |
| 79 | + const batch = ids.slice(i, i + 90); |
| 80 | + driver.query(`DELETE FROM ${TABLE} WHERE id IN (${batch.map(() => "?").join(",")})`, batch); |
| 81 | + } |
| 82 | + return { count: ids.length }; |
| 83 | + }, |
| 84 | + }; |
| 85 | + return adapter as unknown as Vectorize; |
| 86 | +} |
0 commit comments