Merge pull request #655 from tobi/feat/local-qmd-index-bench

feat: support project-local QMD indexes
This commit is contained in:
Tobias Lütke 2026-05-16 14:31:05 -04:00 committed by GitHub
commit 17aa34d753
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
7 changed files with 326 additions and 30 deletions

View File

@ -22,6 +22,7 @@ import {
type QMDStore,
type SearchResult,
type HybridQueryResult,
type ExpandedQuery,
} from "../index.js";
import { scoreResults } from "./score.js";
import type {
@ -34,35 +35,130 @@ import type {
type Backend = {
name: string;
run: (store: QMDStore, query: string, limit: number, collection?: string) => Promise<string[]>;
run: (store: QMDStore, query: BenchmarkQuery, limit: number, collection?: string) => Promise<string[]>;
};
type ParsedStructuredQuery = {
searches: ExpandedQuery[];
intent?: string;
};
function parseStructuredQuery(query: string): ParsedStructuredQuery | undefined {
const lines = query.split("\n").map((line, idx) => ({
trimmed: line.trim(),
number: idx + 1,
})).filter(line => line.trimmed.length > 0);
if (lines.length === 0) return undefined;
const prefixRe = /^(lex|vec|hyde):\s*/i;
const intentRe = /^intent:\s*/i;
const searches: ExpandedQuery[] = [];
let intent: string | undefined;
for (const line of lines) {
if (intentRe.test(line.trimmed)) {
if (intent !== undefined) {
throw new Error(`Line ${line.number}: only one intent: line is allowed per benchmark query.`);
}
intent = line.trimmed.replace(intentRe, "").trim();
if (!intent) {
throw new Error(`Line ${line.number}: intent: must include text.`);
}
continue;
}
const match = line.trimmed.match(prefixRe);
if (match) {
const type = match[1]!.toLowerCase() as "lex" | "vec" | "hyde";
const text = line.trimmed.slice(match[0].length).trim();
if (!text) {
throw new Error(`Line ${line.number} (${type}:) must include text.`);
}
searches.push({ type, query: text, line: line.number });
continue;
}
if (lines.length === 1) {
return undefined;
}
throw new Error(`Line ${line.number} is missing a lex:/vec:/hyde:/intent: prefix.`);
}
if (intent && searches.length === 0) {
throw new Error("intent: cannot appear alone. Add at least one lex:, vec:, or hyde: line.");
}
return searches.length > 0 ? { searches, intent } : undefined;
}
function uniqueFiles(files: string[], limit: number): string[] {
const seen = new Set<string>();
const out: string[] = [];
for (const file of files) {
if (seen.has(file)) continue;
seen.add(file);
out.push(file);
if (out.length >= limit) break;
}
return out;
}
const BACKENDS: Backend[] = [
{
name: "bm25",
run: async (store, query, limit, collection) => {
const results = await store.searchLex(query, { limit, collection });
const structured = parseStructuredQuery(query.query);
const lexQueries = structured?.searches.filter(q => q.type === "lex");
if (structured) {
const files: string[] = [];
for (const lex of lexQueries ?? []) {
const results = await store.searchLex(lex.query, { limit, collection });
files.push(...results.map((r: SearchResult) => r.filepath));
}
return uniqueFiles(files, limit);
}
const results = await store.searchLex(query.query, { limit, collection });
return results.map((r: SearchResult) => r.filepath);
},
},
{
name: "vector",
run: async (store, query, limit, collection) => {
const results = await store.searchVector(query, { limit, collection });
const structured = parseStructuredQuery(query.query);
const vectorQueries = structured?.searches.filter(q => q.type === "vec" || q.type === "hyde");
if (structured) {
const files: string[] = [];
for (const vectorQuery of vectorQueries ?? []) {
const results = await store.searchVector(vectorQuery.query, { limit, collection });
files.push(...results.map((r: SearchResult) => r.filepath));
}
return uniqueFiles(files, limit);
}
const results = await store.searchVector(query.query, { limit, collection });
return results.map((r: SearchResult) => r.filepath);
},
},
{
name: "hybrid",
run: async (store, query, limit, collection) => {
const results = await store.search({ query, limit, collection, rerank: false });
const structured = parseStructuredQuery(query.query);
const results = structured
? await store.search({ queries: structured.searches, intent: structured.intent, limit, collection, rerank: false })
: await store.search({ query: query.query, limit, collection, rerank: false });
return results.map((r: HybridQueryResult) => r.file);
},
},
{
name: "full",
run: async (store, query, limit, collection) => {
const results = await store.search({ query, limit, collection, rerank: true });
const structured = parseStructuredQuery(query.query);
const results = structured
? await store.search({ queries: structured.searches, intent: structured.intent, limit, collection, rerank: true })
: await store.search({ query: query.query, limit, collection, rerank: true });
return results.map((r: HybridQueryResult) => r.file);
},
},
@ -79,18 +175,23 @@ async function runQuery(
let resultFiles: string[];
try {
resultFiles = await backend.run(store, query.query, limit, collection);
resultFiles = await backend.run(store, query, limit, collection);
} catch (err: any) {
// Backend may not be available (e.g., no embeddings for vector search)
return {
precision_at_k: 0,
recall: 0,
recall_at_1: 0,
recall_at_3: 0,
recall_at_5: 0,
mrr: 0,
f1: 0,
hits_at_k: 0,
total_expected: query.expected_files.length,
latency_ms: Date.now() - start,
top_files: [],
matched_files: [],
unmatched_expected_files: query.expected_files,
};
}
@ -111,14 +212,14 @@ function formatTable(results: QueryResult[]): string {
const num = (n: number) => n.toFixed(2).padStart(5);
lines.push(
`${pad("Query", 25)} ${pad("Backend", 8)} ${pad("P@k", 6)} ${pad("Recall", 7)} ${pad("MRR", 6)} ${pad("F1", 6)} ${pad("ms", 8)}`
`${pad("Query", 25)} ${pad("Backend", 8)} ${pad("P@k", 6)} ${pad("R@1", 6)} ${pad("R@3", 6)} ${pad("R@5", 6)} ${pad("MRR", 6)} ${pad("F1", 6)} ${pad("ms", 8)}`
);
lines.push("-".repeat(70));
lines.push("-".repeat(88));
for (const r of results) {
for (const [backend, br] of Object.entries(r.backends)) {
lines.push(
`${pad(r.id, 25)} ${pad(backend, 8)} ${num(br.precision_at_k)} ${num(br.recall)} ${num(br.mrr)} ${num(br.f1)} ${String(Math.round(br.latency_ms)).padStart(7)}ms`
`${pad(r.id, 25)} ${pad(backend, 8)} ${num(br.precision_at_k)} ${num(br.recall_at_1)} ${num(br.recall_at_3)} ${num(br.recall_at_5)} ${num(br.mrr)} ${num(br.f1)} ${String(Math.round(br.latency_ms)).padStart(7)}ms`
);
}
lines.push("");
@ -138,13 +239,16 @@ function computeSummary(results: QueryResult[]): BenchmarkResult["summary"] {
}
}
for (const name of backendNames) {
let totalP = 0, totalR = 0, totalMrr = 0, totalF1 = 0, totalLat = 0, count = 0;
for (const name of Array.from(backendNames)) {
let totalP = 0, totalR = 0, totalR1 = 0, totalR3 = 0, totalR5 = 0, totalMrr = 0, totalF1 = 0, totalLat = 0, count = 0;
for (const r of results) {
const br = r.backends[name];
if (!br) continue;
totalP += br.precision_at_k;
totalR += br.recall;
totalR1 += br.recall_at_1;
totalR3 += br.recall_at_3;
totalR5 += br.recall_at_5;
totalMrr += br.mrr;
totalF1 += br.f1;
totalLat += br.latency_ms;
@ -154,6 +258,9 @@ function computeSummary(results: QueryResult[]): BenchmarkResult["summary"] {
summary[name] = {
avg_precision: totalP / count,
avg_recall: totalR / count,
avg_recall_at_1: totalR1 / count,
avg_recall_at_3: totalR3 / count,
avg_recall_at_5: totalR5 / count,
avg_mrr: totalMrr / count,
avg_f1: totalF1 / count,
avg_latency_ms: totalLat / count,
@ -166,7 +273,7 @@ function computeSummary(results: QueryResult[]): BenchmarkResult["summary"] {
export async function runBenchmark(
fixturePath: string,
options: { json?: boolean; collection?: string; backends?: string[] } = {},
options: { json?: boolean; collection?: string; backends?: string[]; dbPath?: string; configPath?: string } = {},
): Promise<BenchmarkResult> {
// Load fixture
const raw = readFileSync(resolve(fixturePath), "utf-8");
@ -177,7 +284,10 @@ export async function runBenchmark(
}
// Open store
const store = await createStore({ dbPath: getDefaultDbPath() });
const store = await createStore({
dbPath: options.dbPath ?? getDefaultDbPath(),
...(options.configPath ? { configPath: options.configPath } : {}),
});
// Filter backends if requested
const activeBackends = options.backends
@ -232,7 +342,7 @@ export async function runBenchmark(
const num = (n: number) => n.toFixed(3).padStart(6);
for (const [name, s] of Object.entries(summary)) {
console.log(
` ${pad(name, 8)} P@k=${num(s.avg_precision)} Recall=${num(s.avg_recall)} MRR=${num(s.avg_mrr)} F1=${num(s.avg_f1)} Avg=${Math.round(s.avg_latency_ms)}ms`
` ${pad(name, 8)} P@k=${num(s.avg_precision)} R@1=${num(s.avg_recall_at_1)} R@3=${num(s.avg_recall_at_3)} R@5=${num(s.avg_recall_at_5)} MRR=${num(s.avg_mrr)} F1=${num(s.avg_f1)} Avg=${Math.round(s.avg_latency_ms)}ms`
);
}
}

View File

@ -11,7 +11,7 @@
*/
export function normalizePath(p: string): string {
if (p.startsWith("qmd://")) {
// qmd://collection/path/to/file → path/to/file
// qmd://collection/docs/readme.md → docs/readme.md
const withoutScheme = p.slice("qmd://".length);
const slashIdx = withoutScheme.indexOf("/");
p = slashIdx >= 0 ? withoutScheme.slice(slashIdx + 1) : withoutScheme;
@ -31,6 +31,30 @@ export function pathsMatch(result: string, expected: string): boolean {
return false;
}
type ScoreMetrics = {
precision_at_k: number;
recall: number;
recall_at_1: number;
recall_at_3: number;
recall_at_5: number;
mrr: number;
f1: number;
hits_at_k: number;
matched_files: string[];
unmatched_expected_files: string[];
};
function hitsWithin(resultFiles: string[], expectedFiles: string[], k: number): number {
const topKResults = resultFiles.slice(0, k);
let hits = 0;
for (const expected of expectedFiles) {
if (topKResults.some(r => pathsMatch(r, expected))) {
hits++;
}
}
return hits;
}
/**
* Score a set of search results against expected files.
*/
@ -38,21 +62,18 @@ export function scoreResults(
resultFiles: string[],
expectedFiles: string[],
topK: number,
): { precision_at_k: number; recall: number; mrr: number; f1: number; hits_at_k: number } {
): ScoreMetrics {
// Count hits in top-k
const topKResults = resultFiles.slice(0, topK);
let hitsAtK = 0;
for (const expected of expectedFiles) {
if (topKResults.some(r => pathsMatch(r, expected))) {
hitsAtK++;
}
}
const hitsAtK = hitsWithin(resultFiles, expectedFiles, topK);
const matchedFiles: string[] = [];
const unmatchedExpectedFiles: string[] = [];
// Count total hits anywhere
let totalHits = 0;
for (const expected of expectedFiles) {
if (resultFiles.some(r => pathsMatch(r, expected))) {
totalHits++;
matchedFiles.push(expected);
} else {
unmatchedExpectedFiles.push(expected);
}
}
@ -67,10 +88,24 @@ export function scoreResults(
const denominator = Math.min(topK, expectedFiles.length);
const precision_at_k = denominator > 0 ? hitsAtK / denominator : 0;
const recall = expectedFiles.length > 0 ? totalHits / expectedFiles.length : 0;
const recall = expectedFiles.length > 0 ? matchedFiles.length / expectedFiles.length : 0;
const recall_at_1 = expectedFiles.length > 0 ? hitsWithin(resultFiles, expectedFiles, 1) / expectedFiles.length : 0;
const recall_at_3 = expectedFiles.length > 0 ? hitsWithin(resultFiles, expectedFiles, 3) / expectedFiles.length : 0;
const recall_at_5 = expectedFiles.length > 0 ? hitsWithin(resultFiles, expectedFiles, 5) / expectedFiles.length : 0;
const f1 = precision_at_k + recall > 0
? 2 * (precision_at_k * recall) / (precision_at_k + recall)
: 0;
return { precision_at_k, recall, mrr, f1, hits_at_k: hitsAtK };
return {
precision_at_k,
recall,
recall_at_1,
recall_at_3,
recall_at_5,
mrr,
f1,
hits_at_k: hitsAtK,
matched_files: matchedFiles,
unmatched_expected_files: unmatchedExpectedFiles,
};
}

View File

@ -37,6 +37,12 @@ export interface BackendResult {
precision_at_k: number;
/** Fraction of expected files found anywhere in results */
recall: number;
/** Fraction of expected files found in the first result */
recall_at_1: number;
/** Fraction of expected files found in the top 3 results */
recall_at_3: number;
/** Fraction of expected files found in the top 5 results */
recall_at_5: number;
/** Reciprocal rank of first relevant result (1/rank, 0 if not found) */
mrr: number;
/** Harmonic mean of precision_at_k and recall */
@ -49,6 +55,10 @@ export interface BackendResult {
latency_ms: number;
/** Top result file paths (for inspection) */
top_files: string[];
/** Expected files that were found anywhere in the returned result set */
matched_files: string[];
/** Expected files missing from the returned result set */
unmatched_expected_files: string[];
}
export interface QueryResult {
@ -65,6 +75,9 @@ export interface BenchmarkResult {
summary: Record<string, {
avg_precision: number;
avg_recall: number;
avg_recall_at_1: number;
avg_recall_at_3: number;
avg_recall_at_5: number;
avg_mrr: number;
avg_f1: number;
avg_latency_ms: number;

View File

@ -98,6 +98,11 @@ import {
listAllContexts,
setConfigIndexName,
loadConfig,
setConfigSource,
findLocalConfigPath,
getLocalDbPath,
getConfigPath,
configExists,
} from "../collections.js";
import { getEmbeddedQmdSkillContent, getEmbeddedQmdSkillFiles } from "../embedded-skills.js";
@ -2652,11 +2657,22 @@ function parseCLI() {
process.env.QMD_FORCE_CPU = "1";
}
// Select index name (default: "index")
// Select index name (default: "index"). If no explicit --index is supplied,
// a project-local .qmd/index.yaml overrides the global config/cache paths.
const indexName = values.index as string | undefined;
if (indexName) {
setIndexName(indexName);
setConfigIndexName(indexName);
setConfigSource();
} else {
const localConfigPath = findLocalConfigPath();
if (localConfigPath) {
setConfigSource({ configPath: localConfigPath });
storeDbPathOverride = getLocalDbPath(localConfigPath);
closeDb();
} else {
setConfigSource();
}
}
// Determine output format
@ -3321,8 +3337,10 @@ if (isMain) {
const { runBenchmark } = await import("../bench/bench.js");
const benchCollection = cli.opts.collection;
await runBenchmark(fixturePath, {
json: !!(cli.opts as { json?: boolean }).json,
json: !!cli.values.json,
collection: Array.isArray(benchCollection) ? benchCollection[0] : benchCollection,
dbPath: getDbPath(),
configPath: configExists() ? getConfigPath() : undefined,
});
break;
}

View File

@ -125,6 +125,34 @@ function getConfigFilePath(): string {
return join(getConfigDir(), `${currentIndexName}.yml`);
}
/**
* Find a project-local QMD config by walking upward from startDir.
* The local config lives at .qmd/index.yaml or .qmd/index.yml and,
* when used by the CLI, keeps both config and index DB writes inside
* the project instead of the global ~/.config / ~/.cache locations.
*/
export function findLocalConfigPath(startDir: string = process.cwd()): string | undefined {
let dir = resolve(startDir);
while (true) {
const qmdDir = join(dir, ".qmd");
const yamlPath = join(qmdDir, "index.yaml");
if (existsSync(yamlPath)) return yamlPath;
const ymlPath = join(qmdDir, "index.yml");
if (existsSync(ymlPath)) return ymlPath;
const parent = dirname(dir);
if (parent === dir) return undefined;
dir = parent;
}
}
/** Return the local SQLite index path paired with a local .qmd/index.yaml file. */
export function getLocalDbPath(configPath: string): string {
return join(dirname(configPath), "index.sqlite");
}
/**
* Ensure config directory exists
*/

View File

@ -99,6 +99,20 @@ describe("scoreResults", () => {
expect(result.mrr).toBeCloseTo(0.5); // 1/2
});
test("reports recall@1/3/5 and matched documents", () => {
const result = scoreResults(
["x.md", "qmd://concepts/a.md", "docs/b.md", "docs/c.md", "docs/d.md"],
["concepts/a.md", "b.md", "missing.md"],
3,
);
expect(result.recall_at_1).toBe(0);
expect(result.recall_at_3).toBeCloseTo(2 / 3);
expect(result.recall_at_5).toBeCloseTo(2 / 3);
expect(result.matched_files).toEqual(["concepts/a.md", "b.md"]);
expect(result.unmatched_expected_files).toEqual(["missing.md"]);
});
test("empty results", () => {
const result = scoreResults([], ["a.md"], 1);
expect(result.precision_at_k).toBe(0);

78
test/local-config.test.ts Normal file
View File

@ -0,0 +1,78 @@
import { existsSync, mkdtempSync, mkdirSync, writeFileSync, rmSync } from "node:fs";
import { execFileSync } from "node:child_process";
import { join } from "node:path";
import { tmpdir } from "node:os";
import { afterEach, describe, expect, test } from "vitest";
import { findLocalConfigPath, getLocalDbPath } from "../src/collections.js";
const roots: string[] = [];
function tempProject(): string {
const root = mkdtempSync(join(tmpdir(), "qmd-local-config-"));
roots.push(root);
return root;
}
afterEach(() => {
for (const root of roots.splice(0)) {
rmSync(root, { recursive: true, force: true });
}
});
describe("local .qmd project config", () => {
test("finds .qmd/index.yaml from nested working directories", () => {
const root = tempProject();
const configPath = join(root, ".qmd", "index.yaml");
mkdirSync(join(root, ".qmd"), { recursive: true });
writeFileSync(configPath, "collections: {}\n");
const nested = join(root, "wiki", "Shopify");
mkdirSync(nested, { recursive: true });
expect(findLocalConfigPath(nested)).toBe(configPath);
});
test("prefers index.yaml over index.yml when both exist", () => {
const root = tempProject();
mkdirSync(join(root, ".qmd"), { recursive: true });
const yaml = join(root, ".qmd", "index.yaml");
const yml = join(root, ".qmd", "index.yml");
writeFileSync(yaml, "collections: {}\n");
writeFileSync(yml, "collections: {}\n");
expect(findLocalConfigPath(root)).toBe(yaml);
});
test("uses .qmd/index.sqlite next to the local config", () => {
const root = tempProject();
mkdirSync(join(root, ".qmd"), { recursive: true });
const configPath = join(root, ".qmd", "index.yaml");
writeFileSync(configPath, "collections: {}\n");
expect(getLocalDbPath(configPath)).toBe(join(root, ".qmd", "index.sqlite"));
});
test("CLI uses local .qmd config and index instead of global cache", () => {
const root = tempProject();
mkdirSync(join(root, ".qmd"), { recursive: true });
mkdirSync(join(root, "docs"), { recursive: true });
writeFileSync(join(root, "docs", "a.md"), "# A\n\nLocal test document.\n");
writeFileSync(join(root, ".qmd", "index.yaml"), `collections:\n docs:\n path: ${JSON.stringify(join(root, "docs"))}\n pattern: "**/*.md"\n context:\n /: Local test docs\n`);
const home = join(root, "home");
const output = execFileSync("bun", [join(process.cwd(), "src/cli/qmd.ts"), "status"], {
cwd: root,
encoding: "utf-8",
env: {
...process.env,
HOME: home,
XDG_CONFIG_HOME: join(home, ".config"),
XDG_CACHE_HOME: join(home, ".cache"),
},
});
expect(output).toContain(`Index: ${join(root, ".qmd", "index.sqlite")}`);
expect(output).toContain("docs (qmd://docs/)");
expect(existsSync(join(root, ".qmd", "index.sqlite"))).toBe(true);
expect(existsSync(join(home, ".cache", "qmd", "index.sqlite"))).toBe(false);
});
});