diff --git a/src/qmd.ts b/src/qmd.ts index 3d23541..ac89227 100755 --- a/src/qmd.ts +++ b/src/qmd.ts @@ -1972,8 +1972,11 @@ async function vectorSearch(query: string, opts: OutputOptions, model: string = const perQueryLimit = opts.all ? 500 : 20; const allResults = new Map(); - // Use Promise.all for concurrent vector searches - await Promise.all(vectorQueries.map(async (q) => { + // IMPORTANT: Run vector searches sequentially, not with Promise.all. + // node-llama-cpp's embedding context hangs when multiple concurrent embed() calls + // are made. This is a known limitation of the LlamaEmbeddingContext. + // See: https://github.com/tobi/qmd/pull/23 + for (const q of vectorQueries) { const vecResults = await searchVec(db, q, model, perQueryLimit, collectionName as any); for (const r of vecResults) { const existing = allResults.get(r.filepath); @@ -1981,7 +1984,7 @@ async function vectorSearch(query: string, opts: OutputOptions, model: string = allResults.set(r.filepath, { file: r.filepath, displayPath: r.displayPath, title: r.title, body: r.body || "", score: r.score, hash: r.hash }); } } - })); + } // Sort by max score and limit to requested count const results = Array.from(allResults.values()) diff --git a/src/store.test.ts b/src/store.test.ts index e82beb8..1c4546c 100644 --- a/src/store.test.ts +++ b/src/store.test.ts @@ -1855,6 +1855,42 @@ describe("LlamaCpp Integration", () => { await cleanupTestDb(store); }); + // Regression test for https://github.com/tobi/qmd/pull/23 + // sqlite-vec virtual tables hang when combined with JOINs in the same query. + // The fix uses a two-step approach: vector query first, then separate JOINs. + test("searchVec uses two-step query to avoid sqlite-vec JOIN hang", async () => { + const store = await createTestStore(); + const collectionName = await createTestCollection(); + + const hash = "regression_test_hash"; + await insertTestDocument(store.db, collectionName, { + name: "regression-doc", + hash, + body: "Test content for vector search regression", + filepath: "/test/regression.md", + displayPath: "regression.md", + }); + + // Create vector table and insert a test vector + store.ensureVecTable(768); + const embedding = Array(768).fill(0).map(() => Math.random()); + store.db.prepare(`INSERT INTO content_vectors (hash, seq, pos, model, embedded_at) VALUES (?, 0, 0, 'test', ?)`).run(hash, new Date().toISOString()); + store.db.prepare(`INSERT INTO vectors_vec (hash_seq, embedding) VALUES (?, ?)`).run(`${hash}_0`, new Float32Array(embedding)); + + // This should complete quickly (not hang) due to the two-step fix + // The old code with JOINs in the sqlite-vec query would hang indefinitely + const startTime = Date.now(); + const results = await store.searchVec("test content", "embeddinggemma", 5); + const elapsed = Date.now() - startTime; + + // If the query took more than 5 seconds, something is wrong + // (the hang bug would cause it to never return at all) + expect(elapsed).toBeLessThan(5000); + expect(results.length).toBeGreaterThan(0); + + await cleanupTestDb(store); + }); + test("expandQuery returns original plus expanded queries", async () => { const store = await createTestStore(); diff --git a/src/store.ts b/src/store.ts index e14c7ae..8d27fc8 100644 --- a/src/store.ts +++ b/src/store.ts @@ -1679,48 +1679,66 @@ export async function searchVec(db: Database, query: string, model: string, limi const embedding = await getEmbedding(query, model, true); if (!embedding) return []; - // sqlite-vec requires "k = ?" for KNN queries - let sql = ` + // IMPORTANT: We use a two-step query approach here because sqlite-vec virtual tables + // hang indefinitely when combined with JOINs in the same query. Do NOT try to + // "optimize" this by combining into a single query with JOINs - it will break. + // See: https://github.com/tobi/qmd/pull/23 + + // Step 1: Get vector matches from sqlite-vec (no JOINs allowed) + const vecResults = db.prepare(` + SELECT hash_seq, distance + FROM vectors_vec + WHERE embedding MATCH ? AND k = ? + `).all(new Float32Array(embedding), limit * 3) as { hash_seq: string; distance: number }[]; + + if (vecResults.length === 0) return []; + + // Step 2: Get chunk info and document data + const hashSeqs = vecResults.map(r => r.hash_seq); + const distanceMap = new Map(vecResults.map(r => [r.hash_seq, r.distance])); + + // Build query for document lookup + const placeholders = hashSeqs.map(() => '?').join(','); + let docSql = ` SELECT - v.hash_seq, - v.distance, + cv.hash || '_' || cv.seq as hash_seq, + cv.hash, + cv.pos, 'qmd://' || d.collection || '/' || d.path as filepath, d.collection || '/' || d.path as display_path, d.title, - content.doc as body, - cv.hash, - cv.pos - FROM vectors_vec v - JOIN content_vectors cv ON cv.hash || '_' || cv.seq = v.hash_seq + content.doc as body + FROM content_vectors cv JOIN documents d ON d.hash = cv.hash AND d.active = 1 JOIN content ON content.hash = d.hash - WHERE v.embedding MATCH ? AND k = ? + WHERE cv.hash || '_' || cv.seq IN (${placeholders}) `; - - const params: (Float32Array | number | string)[] = [new Float32Array(embedding), limit * 3]; + const params: string[] = [...hashSeqs]; if (collectionId) { - // Filter by collection name - sql += ` AND d.collection = ?`; + docSql += ` AND d.collection = ?`; params.push(String(collectionId)); } - sql += ` ORDER BY v.distance`; + const docRows = db.prepare(docSql).all(...params) as { + hash_seq: string; hash: string; pos: number; filepath: string; + display_path: string; title: string; body: string; + }[]; - const rows = db.prepare(sql).all(...params) as { hash_seq: string; distance: number; filepath: string; display_path: string; title: string; body: string; hash: string; pos: number }[]; - - const seen = new Map(); - for (const row of rows) { + // Combine with distances and dedupe by filepath + const seen = new Map(); + for (const row of docRows) { + const distance = distanceMap.get(row.hash_seq) ?? 1; const existing = seen.get(row.filepath); - if (!existing || row.distance < existing.bestDist) { - seen.set(row.filepath, { row, bestDist: row.distance }); + if (!existing || distance < existing.bestDist) { + seen.set(row.filepath, { row, bestDist: distance }); } } return Array.from(seen.values()) .sort((a, b) => a.bestDist - b.bestDist) .slice(0, limit) - .map(({ row }) => { + .map(({ row, bestDist }) => { const collectionName = row.filepath.split('//')[1]?.split('/')[0] || ""; return { filepath: row.filepath, @@ -1733,7 +1751,7 @@ export async function searchVec(db: Database, query: string, model: string, limi bodyLength: row.body.length, body: row.body, context: getContextForFile(db, row.filepath), - score: 1 - row.distance, // Cosine similarity = 1 - cosine distance + score: 1 - bestDist, // Cosine similarity = 1 - cosine distance source: "vec" as const, chunkPos: row.pos, };