diff --git a/src/qmd.ts b/src/qmd.ts index b6678d2..244578f 100755 --- a/src/qmd.ts +++ b/src/qmd.ts @@ -1748,7 +1748,7 @@ type OutputOptions = { limit: number; minScore: number; all?: boolean; - collection?: string; // Filter by collection name (pwd suffix match) + collection?: string | string[]; // Filter by collection name(s) lineNumbers?: boolean; // Add line numbers to output context?: string; // Optional context for query expansion }; @@ -1902,24 +1902,47 @@ function outputResults(results: { file: string; displayPath: string; title: stri } } -function search(query: string, opts: OutputOptions): void { - const db = getDb(); - - // Validate collection filter if specified - let collectionName: string | undefined; - if (opts.collection) { - const coll = getCollectionFromYaml(opts.collection); +// Resolve -c collection filter: supports single string, array, or undefined. +// Returns validated collection names (exits on unknown collection). +function resolveCollectionFilter(raw: string | string[] | undefined): string[] { + if (!raw) return []; + const names = Array.isArray(raw) ? raw : [raw]; + const validated: string[] = []; + for (const name of names) { + const coll = getCollectionFromYaml(name); if (!coll) { - console.error(`Collection not found: ${opts.collection}`); + console.error(`Collection not found: ${name}`); closeDb(); process.exit(1); } - collectionName = opts.collection; + validated.push(name); } + return validated; +} + +// Post-filter results to only include files from specified collections. +function filterByCollections(results: T[], collectionNames: string[]): T[] { + if (collectionNames.length <= 1) return results; + const prefixes = collectionNames.map(n => `qmd://${n}/`); + return results.filter(r => { + const path = r.filepath || r.file || ''; + return prefixes.some(p => path.startsWith(p)); + }); +} + +function search(query: string, opts: OutputOptions): void { + const db = getDb(); + + // Validate collection filter (supports multiple -c flags) + const collectionNames = resolveCollectionFilter(opts.collection); + const singleCollection = collectionNames.length === 1 ? collectionNames[0] : undefined; // Use large limit for --all, otherwise fetch more than needed and let outputResults filter const fetchLimit = opts.all ? 100000 : Math.max(50, opts.limit * 2); - const results = searchFTS(db, query, fetchLimit, collectionName); + const results = filterByCollections( + searchFTS(db, query, fetchLimit, singleCollection), + collectionNames + ); // Add context to results const resultsWithContext = results.map(r => ({ @@ -1960,20 +1983,15 @@ function logExpansionTree(originalQuery: string, expanded: ExpandedQuery[]): voi async function vectorSearch(query: string, opts: OutputOptions, _model: string = DEFAULT_EMBED_MODEL): Promise { const store = getStore(); - if (opts.collection) { - const coll = getCollectionFromYaml(opts.collection); - if (!coll) { - console.error(`Collection not found: ${opts.collection}`); - closeDb(); - process.exit(1); - } - } + // Validate collection filter (supports multiple -c flags) + const collectionNames = resolveCollectionFilter(opts.collection); + const singleCollection = collectionNames.length === 1 ? collectionNames[0] : undefined; checkIndexHealth(store.db); await withLLMSession(async () => { - const results = await vectorSearchQuery(store, query, { - collection: opts.collection, + let results = await vectorSearchQuery(store, query, { + collection: singleCollection, limit: opts.all ? 500 : (opts.limit || 10), minScore: opts.minScore || 0.3, hooks: { @@ -1984,6 +2002,14 @@ async function vectorSearch(query: string, opts: OutputOptions, _model: string = }, }); + // Post-filter for multi-collection + if (collectionNames.length > 1) { + results = results.filter(r => { + const prefixes = collectionNames.map(n => `qmd://${n}/`); + return prefixes.some(p => r.file.startsWith(p)); + }); + } + closeDb(); if (results.length === 0) { @@ -2006,20 +2032,15 @@ async function vectorSearch(query: string, opts: OutputOptions, _model: string = async function querySearch(query: string, opts: OutputOptions, _embedModel: string = DEFAULT_EMBED_MODEL, _rerankModel: string = DEFAULT_RERANK_MODEL): Promise { const store = getStore(); - if (opts.collection) { - const coll = getCollectionFromYaml(opts.collection); - if (!coll) { - console.error(`Collection not found: ${opts.collection}`); - closeDb(); - process.exit(1); - } - } + // Validate collection filter (supports multiple -c flags) + const collectionNames = resolveCollectionFilter(opts.collection); + const singleCollection = collectionNames.length === 1 ? collectionNames[0] : undefined; checkIndexHealth(store.db); await withLLMSession(async () => { - const results = await hybridQuery(store, query, { - collection: opts.collection, + let results = await hybridQuery(store, query, { + collection: singleCollection, limit: opts.all ? 500 : (opts.limit || 10), minScore: opts.minScore || 0, hooks: { @@ -2040,6 +2061,14 @@ async function querySearch(query: string, opts: OutputOptions, _embedModel: stri }, }); + // Post-filter for multi-collection + if (collectionNames.length > 1) { + results = results.filter(r => { + const prefixes = collectionNames.map(n => `qmd://${n}/`); + return prefixes.some(p => r.file.startsWith(p)); + }); + } + closeDb(); if (results.length === 0) { @@ -2088,7 +2117,7 @@ function parseCLI() { xml: { type: "boolean" }, files: { type: "boolean" }, json: { type: "boolean" }, - collection: { type: "string", short: "c" }, // Filter by collection + collection: { type: "string", short: "c", multiple: true }, // Filter by collection(s) // Collection options name: { type: "string" }, // collection name mask: { type: "string" }, // glob pattern @@ -2137,7 +2166,7 @@ function parseCLI() { limit: isAll ? 100000 : (values.n ? parseInt(String(values.n), 10) || defaultLimit : defaultLimit), minScore: values["min-score"] ? parseFloat(String(values["min-score"])) || 0 : 0, all: isAll, - collection: values.collection as string | undefined, + collection: values.collection as string[] | undefined, lineNumbers: !!values["line-numbers"], }; diff --git a/test/multi-collection-filter.test.ts b/test/multi-collection-filter.test.ts new file mode 100644 index 0000000..31500b1 --- /dev/null +++ b/test/multi-collection-filter.test.ts @@ -0,0 +1,145 @@ +/** + * Unit tests for multi-collection filter logic (PR #191). + * + * Tests the filterByCollections post-filter and the resolveCollectionFilter + * behavior for single-collection vs multi-collection search. + */ + +import { describe, test, expect } from "vitest"; + +// Reproduce the filterByCollections logic from qmd.ts for testing +// (the function is private in qmd.ts) +function filterByCollections( + results: T[], + collectionNames: string[], +): T[] { + if (collectionNames.length <= 1) return results; + const prefixes = collectionNames.map((n) => `qmd://${n}/`); + return results.filter((r) => { + const path = r.filepath || r.file || ""; + return prefixes.some((p) => path.startsWith(p)); + }); +} + +describe("filterByCollections", () => { + const results = [ + { filepath: "qmd://docs/readme.md", file: "qmd://docs/readme.md" }, + { filepath: "qmd://notes/todo.md", file: "qmd://notes/todo.md" }, + { filepath: "qmd://journals/2024/jan.md", file: "qmd://journals/2024/jan.md" }, + { filepath: "qmd://docs/api.md", file: "qmd://docs/api.md" }, + ]; + + test("returns all results when no collections specified", () => { + expect(filterByCollections(results, [])).toEqual(results); + }); + + test("returns all results for single collection (no-op, handled by SQL filter)", () => { + expect(filterByCollections(results, ["docs"])).toEqual(results); + }); + + test("filters to matching collections when multiple specified", () => { + const filtered = filterByCollections(results, ["docs", "journals"]); + expect(filtered).toHaveLength(3); + expect(filtered.map((r) => r.filepath)).toEqual([ + "qmd://docs/readme.md", + "qmd://journals/2024/jan.md", + "qmd://docs/api.md", + ]); + }); + + test("filters correctly with two collections", () => { + const filtered = filterByCollections(results, ["notes", "journals"]); + expect(filtered).toHaveLength(2); + expect(filtered.map((r) => r.filepath)).toEqual([ + "qmd://notes/todo.md", + "qmd://journals/2024/jan.md", + ]); + }); + + test("returns empty when no results match collections", () => { + const filtered = filterByCollections(results, ["archive", "trash"]); + expect(filtered).toHaveLength(0); + }); + + test("uses file field when filepath is missing", () => { + const fileOnlyResults = [ + { file: "qmd://docs/readme.md" }, + { file: "qmd://notes/todo.md" }, + ]; + const filtered = filterByCollections(fileOnlyResults, ["docs", "notes"]); + expect(filtered).toHaveLength(2); + }); + + test("uses filepath over file when both present", () => { + const mixedResults = [ + { filepath: "qmd://docs/readme.md", file: "qmd://notes/todo.md" }, + ]; + const filtered = filterByCollections(mixedResults, ["docs", "notes"]); + expect(filtered).toHaveLength(1); + // Should match via filepath (docs), not file (notes) + expect(filtered[0].filepath).toBe("qmd://docs/readme.md"); + }); +}); + +describe("resolveCollectionFilter input normalization", () => { + // Test the array normalization logic without the DB dependency + function normalizeCollectionInput(raw: string | string[] | undefined): string[] { + if (!raw) return []; + return Array.isArray(raw) ? raw : [raw]; + } + + test("undefined returns empty array", () => { + expect(normalizeCollectionInput(undefined)).toEqual([]); + }); + + test("single string returns single-element array", () => { + expect(normalizeCollectionInput("docs")).toEqual(["docs"]); + }); + + test("array passes through", () => { + expect(normalizeCollectionInput(["docs", "notes"])).toEqual(["docs", "notes"]); + }); + + test("empty string returns single-element array", () => { + expect(normalizeCollectionInput("")).toEqual([]); + }); +}); + +describe("collection option type from parseArgs", () => { + // Verify that parseArgs with `multiple: true` produces string[] + test("parseArgs multiple:true produces array for repeated flags", () => { + const { parseArgs } = require("node:util"); + const { values } = parseArgs({ + args: ["-c", "docs", "-c", "notes"], + options: { + collection: { type: "string", short: "c", multiple: true }, + }, + strict: true, + }); + expect(values.collection).toEqual(["docs", "notes"]); + }); + + test("parseArgs multiple:true produces array for single flag", () => { + const { parseArgs } = require("node:util"); + const { values } = parseArgs({ + args: ["-c", "docs"], + options: { + collection: { type: "string", short: "c", multiple: true }, + }, + strict: true, + }); + expect(values.collection).toEqual(["docs"]); + }); + + test("parseArgs multiple:true produces undefined when flag absent", () => { + const { parseArgs } = require("node:util"); + const { values } = parseArgs({ + args: [], + options: { + collection: { type: "string", short: "c", multiple: true }, + }, + strict: true, + }); + expect(values.collection).toBeUndefined(); + }); +});