Add local AI and Qdrant vector backends

2026-05-02 09:57:57 +10:00
parent 00f0afb8f4
commit 782ace9a34
37 changed files with 1440 additions and 321 deletions
@@ -0,0 +1,109 @@
+import { getSQLiteClient } from '@/services/database/sqlite-client';
+import { getEmbeddingProviderInfo } from '@/services/embedding/provider';
+import type { Chunk } from '@/types/database';
+import type { RankedChunk, RankedNode, VectorBackend, VectorBackendHealth } from './index';
+
+function vectorLiteral(embedding: number[]): string {
+  return `[${embedding.join(',')}]`;
+}
+
+export class SqliteVecBackend implements VectorBackend {
+  async upsertChunk(chunkId: number, _nodeId: number, _chunkIndex: number, _text: string, embedding: number[]): Promise<void> {
+    const sqlite = getSQLiteClient();
+    sqlite.prepare('DELETE FROM vec_chunks WHERE chunk_id = ?').run(BigInt(chunkId));
+    sqlite.prepare('INSERT OR REPLACE INTO vec_chunks (chunk_id, embedding) VALUES (?, ?)').run(BigInt(chunkId), vectorLiteral(embedding));
+  }
+
+  async upsertNode(nodeId: number, _text: string, embedding: number[]): Promise<void> {
+    const sqlite = getSQLiteClient();
+    sqlite.prepare('DELETE FROM vec_nodes WHERE node_id = ?').run(BigInt(nodeId));
+    sqlite.prepare('INSERT OR REPLACE INTO vec_nodes (node_id, embedding) VALUES (?, ?)').run(BigInt(nodeId), vectorLiteral(embedding));
+  }
+
+  async searchChunks(
+    queryEmbedding: number[],
+    similarityThreshold: number,
+    limit: number,
+    nodeIds?: number[]
+  ): Promise<RankedChunk[]> {
+    const sqlite = getSQLiteClient();
+    const vectorLimit = Math.max(limit * 10, 50);
+
+    if (nodeIds && nodeIds.length > 0) {
+      const result = sqlite.query<RankedChunk>(`
+        SELECT c.*, (1.0 / (1.0 + v.distance)) AS similarity
+        FROM (
+          SELECT chunk_id, distance
+          FROM vec_chunks
+          WHERE embedding MATCH ?
+          ORDER BY distance
+          LIMIT ?
+        ) v
+        JOIN chunks c ON c.id = v.chunk_id
+        WHERE c.node_id IN (${nodeIds.map(() => '?').join(',')})
+          AND (1.0 / (1.0 + v.distance)) >= ?
+        ORDER BY similarity DESC
+        LIMIT ?
+      `, [vectorLiteral(queryEmbedding), vectorLimit, ...nodeIds, similarityThreshold, limit]);
+      return result.rows;
+    }
+
+    const result = sqlite.query<RankedChunk>(`
+      WITH vector_results AS (
+        SELECT chunk_id, distance
+        FROM vec_chunks
+        WHERE embedding MATCH ?
+        ORDER BY distance
+        LIMIT ?
+      )
+      SELECT c.*, (1.0 / (1.0 + vr.distance)) AS similarity
+      FROM vector_results vr
+      JOIN chunks c ON c.id = vr.chunk_id
+      WHERE (1.0 / (1.0 + vr.distance)) >= ?
+      ORDER BY similarity DESC
+      LIMIT ?
+    `, [vectorLiteral(queryEmbedding), vectorLimit, similarityThreshold, limit]);
+    return result.rows;
+  }
+
+  async searchNodes(queryEmbedding: number[], limit: number): Promise<RankedNode[]> {
+    const sqlite = getSQLiteClient();
+    const result = sqlite.query<{ node_id: number; distance: number }>(`
+      SELECT node_id, distance
+      FROM vec_nodes
+      WHERE embedding MATCH ?
+      ORDER BY distance
+      LIMIT ?
+    `, [vectorLiteral(queryEmbedding), Math.max(limit * 2, 50)]);
+
+    return result.rows.map((row) => ({
+      nodeId: Number(row.node_id),
+      score: 1.0 / (1.0 + Number(row.distance)),
+    }));
+  }
+
+  async deleteChunksByNode(nodeId: number): Promise<void> {
+    const sqlite = getSQLiteClient();
+    const chunks = sqlite.query<Pick<Chunk, 'id'>>('SELECT id FROM chunks WHERE node_id = ?', [nodeId]).rows;
+    for (const chunk of chunks) {
+      sqlite.prepare('DELETE FROM vec_chunks WHERE chunk_id = ?').run(BigInt(chunk.id));
+    }
+  }
+
+  async deleteNode(nodeId: number): Promise<void> {
+    const sqlite = getSQLiteClient();
+    await this.deleteChunksByNode(nodeId);
+    sqlite.prepare('DELETE FROM vec_nodes WHERE node_id = ?').run(BigInt(nodeId));
+  }
+
+  async healthCheck(): Promise<VectorBackendHealth> {
+    const sqlite = getSQLiteClient();
+    const ok = await sqlite.checkVectorExtension();
+    return {
+      ok,
+      backend: 'sqlite-vec',
+      dimensions: getEmbeddingProviderInfo().dimensions,
+      detail: ok ? 'sqlite-vec extension loaded' : 'sqlite-vec extension unavailable',
+    };
+  }
+}