qmd.ts 72 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374757677787980818283848586878889909192939495969798991001011021031041051061071081091101111121131141151161171181191201211221231241251261271281291301311321331341351361371381391401411421431441451461471481491501511521531541551561571581591601611621631641651661671681691701711721731741751761771781791801811821831841851861871881891901911921931941951961971981992002012022032042052062072082092102112122132142152162172182192202212222232242252262272282292302312322332342352362372382392402412422432442452462472482492502512522532542552562572582592602612622632642652662672682692702712722732742752762772782792802812822832842852862872882892902912922932942952962972982993003013023033043053063073083093103113123133143153163173183193203213223233243253263273283293303313323333343353363373383393403413423433443453463473483493503513523533543553563573583593603613623633643653663673683693703713723733743753763773783793803813823833843853863873883893903913923933943953963973983994004014024034044054064074084094104114124134144154164174184194204214224234244254264274284294304314324334344354364374384394404414424434444454464474484494504514524534544554564574584594604614624634644654664674684694704714724734744754764774784794804814824834844854864874884894904914924934944954964974984995005015025035045055065075085095105115125135145155165175185195205215225235245255265275285295305315325335345355365375385395405415425435445455465475485495505515525535545555565575585595605615625635645655665675685695705715725735745755765775785795805815825835845855865875885895905915925935945955965975985996006016026036046056066076086096106116126136146156166176186196206216226236246256266276286296306316326336346356366376386396406416426436446456466476486496506516526536546556566576586596606616626636646656666676686696706716726736746756766776786796806816826836846856866876886896906916926936946956966976986997007017027037047057067077087097107117127137147157167177187197207217227237247257267277287297307317327337347357367377387397407417427437447457467477487497507517527537547557567577587597607617627637647657667677687697707717727737747757767777787797807817827837847857867877887897907917927937947957967977987998008018028038048058068078088098108118128138148158168178188198208218228238248258268278288298308318328338348358368378388398408418428438448458468478488498508518528538548558568578588598608618628638648658668678688698708718728738748758768778788798808818828838848858868878888898908918928938948958968978988999009019029039049059069079089099109119129139149159169179189199209219229239249259269279289299309319329339349359369379389399409419429439449459469479489499509519529539549559569579589599609619629639649659669679689699709719729739749759769779789799809819829839849859869879889899909919929939949959969979989991000100110021003100410051006100710081009101010111012101310141015101610171018101910201021102210231024102510261027102810291030103110321033103410351036103710381039104010411042104310441045104610471048104910501051105210531054105510561057105810591060106110621063106410651066106710681069107010711072107310741075107610771078107910801081108210831084108510861087108810891090109110921093109410951096109710981099110011011102110311041105110611071108110911101111111211131114111511161117111811191120112111221123112411251126112711281129113011311132113311341135113611371138113911401141114211431144114511461147114811491150115111521153115411551156115711581159116011611162116311641165116611671168116911701171117211731174117511761177117811791180118111821183118411851186118711881189119011911192119311941195119611971198119912001201120212031204120512061207120812091210121112121213121412151216121712181219122012211222122312241225122612271228122912301231123212331234123512361237123812391240124112421243124412451246124712481249125012511252125312541255125612571258125912601261126212631264126512661267126812691270127112721273127412751276127712781279128012811282128312841285128612871288128912901291129212931294129512961297129812991300130113021303130413051306130713081309131013111312131313141315131613171318131913201321132213231324132513261327132813291330133113321333133413351336133713381339134013411342134313441345134613471348134913501351135213531354135513561357135813591360136113621363136413651366136713681369137013711372137313741375137613771378137913801381138213831384138513861387138813891390139113921393139413951396139713981399140014011402140314041405140614071408140914101411141214131414141514161417141814191420142114221423142414251426142714281429143014311432143314341435143614371438143914401441144214431444144514461447144814491450145114521453145414551456145714581459146014611462146314641465146614671468146914701471147214731474147514761477147814791480148114821483148414851486148714881489149014911492149314941495149614971498149915001501150215031504150515061507150815091510151115121513151415151516151715181519152015211522152315241525152615271528152915301531153215331534153515361537153815391540154115421543154415451546154715481549155015511552155315541555155615571558155915601561156215631564156515661567156815691570157115721573157415751576157715781579158015811582158315841585158615871588158915901591159215931594159515961597159815991600160116021603160416051606160716081609161016111612161316141615161616171618161916201621162216231624162516261627162816291630163116321633163416351636163716381639164016411642164316441645164616471648164916501651165216531654165516561657165816591660166116621663166416651666166716681669167016711672167316741675167616771678167916801681168216831684168516861687168816891690169116921693169416951696169716981699170017011702170317041705170617071708170917101711171217131714171517161717171817191720172117221723172417251726172717281729173017311732173317341735173617371738173917401741174217431744174517461747174817491750175117521753175417551756175717581759176017611762176317641765176617671768176917701771177217731774177517761777177817791780178117821783178417851786178717881789179017911792179317941795179617971798179918001801180218031804180518061807180818091810181118121813181418151816181718181819182018211822182318241825182618271828182918301831183218331834183518361837183818391840184118421843184418451846184718481849185018511852185318541855185618571858185918601861186218631864186518661867186818691870187118721873187418751876187718781879188018811882188318841885188618871888188918901891189218931894189518961897189818991900190119021903190419051906190719081909191019111912191319141915191619171918191919201921192219231924192519261927192819291930193119321933193419351936193719381939194019411942194319441945194619471948194919501951195219531954195519561957195819591960196119621963196419651966196719681969197019711972197319741975197619771978197919801981198219831984198519861987198819891990199119921993199419951996199719981999
  1. #!/usr/bin/env bun
  2. import { Database } from "bun:sqlite";
  3. import { Glob, $ } from "bun";
  4. import { parseArgs } from "util";
  5. import * as sqliteVec from "sqlite-vec";
  6. import {
  7. getDb,
  8. closeDb,
  9. getDbPath,
  10. getPwd,
  11. getRealPath,
  12. homedir,
  13. resolve,
  14. setCustomIndexName,
  15. searchFTS,
  16. searchVec,
  17. reciprocalRankFusion,
  18. extractSnippet,
  19. getContextForFile,
  20. getCollectionIdByName,
  21. findSimilarFiles,
  22. matchFilesByGlob,
  23. getHashesNeedingEmbedding,
  24. getDocument as storeGetDocument,
  25. getMultipleDocuments as storeMultiGetDocuments,
  26. getStatus,
  27. hashContent,
  28. extractTitle,
  29. formatDocForEmbedding,
  30. formatQueryForEmbedding,
  31. chunkDocument,
  32. ensureVecTable,
  33. clearCache,
  34. getCacheKey,
  35. getCachedResult,
  36. setCachedResult,
  37. getIndexHealth,
  38. OLLAMA_URL,
  39. DEFAULT_EMBED_MODEL,
  40. DEFAULT_QUERY_MODEL,
  41. DEFAULT_RERANK_MODEL,
  42. DEFAULT_GLOB,
  43. DEFAULT_MULTI_GET_MAX_BYTES,
  44. } from "./store.js";
  45. import type { SearchResult, RankedResult } from "./store.js";
  46. import {
  47. formatSearchResults,
  48. formatDocuments,
  49. escapeXml,
  50. escapeCSV,
  51. type OutputFormat,
  52. } from "./formatter.js";
  53. // Chunking: ~2000 tokens per chunk, ~3 bytes/token = 6KB
  54. const CHUNK_BYTE_SIZE = 6 * 1024;
  55. // Terminal colors (respects NO_COLOR env)
  56. const useColor = !process.env.NO_COLOR && process.stdout.isTTY;
  57. const c = {
  58. reset: useColor ? "\x1b[0m" : "",
  59. dim: useColor ? "\x1b[2m" : "",
  60. bold: useColor ? "\x1b[1m" : "",
  61. cyan: useColor ? "\x1b[36m" : "",
  62. yellow: useColor ? "\x1b[33m" : "",
  63. green: useColor ? "\x1b[32m" : "",
  64. magenta: useColor ? "\x1b[35m" : "",
  65. blue: useColor ? "\x1b[34m" : "",
  66. };
  67. // Terminal cursor control
  68. const cursor = {
  69. hide() { process.stderr.write('\x1b[?25l'); },
  70. show() { process.stderr.write('\x1b[?25h'); },
  71. };
  72. // Ensure cursor is restored on exit
  73. process.on('SIGINT', () => { cursor.show(); process.exit(130); });
  74. process.on('SIGTERM', () => { cursor.show(); process.exit(143); });
  75. // Terminal progress bar using OSC 9;4 escape sequence
  76. const progress = {
  77. set(percent: number) {
  78. process.stderr.write(`\x1b]9;4;1;${Math.round(percent)}\x07`);
  79. },
  80. clear() {
  81. process.stderr.write(`\x1b]9;4;0\x07`);
  82. },
  83. indeterminate() {
  84. process.stderr.write(`\x1b]9;4;3\x07`);
  85. },
  86. error() {
  87. process.stderr.write(`\x1b]9;4;2\x07`);
  88. },
  89. };
  90. // Format seconds into human-readable ETA
  91. function formatETA(seconds: number): string {
  92. if (seconds < 60) return `${Math.round(seconds)}s`;
  93. if (seconds < 3600) return `${Math.floor(seconds / 60)}m ${Math.round(seconds % 60)}s`;
  94. return `${Math.floor(seconds / 3600)}h ${Math.floor((seconds % 3600) / 60)}m`;
  95. }
  96. // Check index health and print warnings/tips
  97. function checkIndexHealth(db: Database): void {
  98. const { needsEmbedding, totalDocs, daysStale } = getIndexHealth(db);
  99. // Warn if many docs need embedding
  100. if (needsEmbedding > 0) {
  101. const pct = Math.round((needsEmbedding / totalDocs) * 100);
  102. if (pct >= 10) {
  103. process.stderr.write(`${c.yellow}Warning: ${needsEmbedding} documents (${pct}%) need embeddings. Run 'qmd embed' for better results.${c.reset}\n`);
  104. } else {
  105. process.stderr.write(`${c.dim}Tip: ${needsEmbedding} documents need embeddings. Run 'qmd embed' to index them.${c.reset}\n`);
  106. }
  107. }
  108. // Check if most recent document update is older than 2 weeks
  109. if (daysStale !== null && daysStale >= 14) {
  110. process.stderr.write(`${c.dim}Tip: Index last updated ${daysStale} days ago. Run 'qmd update' to refresh.${c.reset}\n`);
  111. }
  112. }
  113. // Compute unique display path for a document
  114. // Always include at least parent folder + filename, add more parent dirs until unique
  115. function computeDisplayPath(
  116. filepath: string,
  117. collectionPath: string,
  118. existingPaths: Set<string>
  119. ): string {
  120. // Get path relative to collection (include collection dir name)
  121. const collectionDir = collectionPath.replace(/\/$/, '');
  122. const collectionName = collectionDir.split('/').pop() || '';
  123. let relativePath: string;
  124. if (filepath.startsWith(collectionDir + '/')) {
  125. // filepath is under collection: use collection name + relative path
  126. relativePath = collectionName + filepath.slice(collectionDir.length);
  127. } else {
  128. // Fallback: just use the filepath
  129. relativePath = filepath;
  130. }
  131. const parts = relativePath.split('/').filter(p => p.length > 0);
  132. // Always include at least parent folder + filename (minimum 2 parts if available)
  133. // Then add more parent dirs until unique
  134. const minParts = Math.min(2, parts.length);
  135. for (let i = parts.length - minParts; i >= 0; i--) {
  136. const candidate = parts.slice(i).join('/');
  137. if (!existingPaths.has(candidate)) {
  138. return candidate;
  139. }
  140. }
  141. // Absolute fallback: use full path (should be unique)
  142. return filepath;
  143. }
  144. // Auto-pull model if not found
  145. async function ensureModelAvailable(model: string): Promise<void> {
  146. try {
  147. const response = await fetch(`${OLLAMA_URL}/api/show`, {
  148. method: "POST",
  149. headers: { "Content-Type": "application/json" },
  150. body: JSON.stringify({ name: model }),
  151. });
  152. if (response.ok) return;
  153. } catch {
  154. // Continue to pull attempt
  155. }
  156. console.log(`Model ${model} not found. Pulling...`);
  157. progress.indeterminate();
  158. const pullResponse = await fetch(`${OLLAMA_URL}/api/pull`, {
  159. method: "POST",
  160. headers: { "Content-Type": "application/json" },
  161. body: JSON.stringify({ name: model, stream: false }),
  162. });
  163. if (!pullResponse.ok) {
  164. progress.error();
  165. throw new Error(`Failed to pull model ${model}: ${pullResponse.status} - ${await pullResponse.text()}`);
  166. }
  167. progress.clear();
  168. console.log(`Model ${model} pulled successfully.`);
  169. }
  170. async function getEmbedding(text: string, model: string, isQuery: boolean = false, title?: string, retried: boolean = false): Promise<number[]> {
  171. const input = isQuery ? formatQueryForEmbedding(text) : formatDocForEmbedding(text, title);
  172. const response = await fetch(`${OLLAMA_URL}/api/embed`, {
  173. method: "POST",
  174. headers: { "Content-Type": "application/json" },
  175. body: JSON.stringify({ model, input }),
  176. });
  177. if (!response.ok) {
  178. const errorText = await response.text();
  179. if (!retried && (errorText.includes("not found") || errorText.includes("does not exist"))) {
  180. await ensureModelAvailable(model);
  181. return getEmbedding(text, model, isQuery, title, true);
  182. }
  183. throw new Error(`Ollama API error: ${response.status} - ${errorText}`);
  184. }
  185. const data = await response.json() as { embeddings: number[][] };
  186. return data.embeddings[0];
  187. }
  188. // Qwen3-Reranker prompt format (trained for yes/no relevance classification)
  189. const RERANK_SYSTEM = `Judge whether the Document meets the requirements based on the Query and the Instruct provided. Note that the answer can only be "yes" or "no".`;
  190. function formatRerankPrompt(query: string, title: string, doc: string): string {
  191. return `<Instruct>: Determine if this document from a Shopify knowledge base is relevant to the search query. The query may reference specific Shopify programs, competitions, features, or named concepts (e.g., "Build a Business" competition, "Shop Pay", "Polaris"). Match documents that discuss the queried topic, even if phrasing differs.
  192. <Query>: ${query}
  193. <Document Title>: ${title}
  194. <Document>: ${doc}`;
  195. }
  196. type LogProb = { token: string; logprob: number };
  197. type RerankResponse = {
  198. response: string;
  199. logprobs?: LogProb[];
  200. };
  201. function parseRerankResponse(data: RerankResponse): number {
  202. if (!data.logprobs || data.logprobs.length === 0) {
  203. throw new Error("Reranker response missing logprobs");
  204. }
  205. const firstToken = data.logprobs[0];
  206. const token = firstToken.token.toLowerCase().trim();
  207. const confidence = Math.exp(firstToken.logprob);
  208. if (token === "yes") {
  209. return confidence;
  210. }
  211. if (token === "no") {
  212. return (1 - confidence) * 0.3;
  213. }
  214. throw new Error(`Unexpected reranker token: "${token}"`);
  215. }
  216. async function rerankSingle(prompt: string, model: string, db?: Database, retried: boolean = false): Promise<number> {
  217. // Use generate with raw template for qwen3-reranker format
  218. // Include empty <think> tags as per HuggingFace reference implementation
  219. const fullPrompt = `<|im_start|>system
  220. ${RERANK_SYSTEM}<|im_end|>
  221. <|im_start|>user
  222. ${prompt}<|im_end|>
  223. <|im_start|>assistant
  224. <think>
  225. </think>
  226. `;
  227. const requestBody = {
  228. model,
  229. prompt: fullPrompt,
  230. raw: true,
  231. stream: false,
  232. logprobs: true,
  233. options: { num_predict: 1 },
  234. };
  235. // Check cache
  236. const cacheKey = db ? getCacheKey(`${OLLAMA_URL}/api/generate`, requestBody) : "";
  237. if (db) {
  238. const cached = getCachedResult(db, cacheKey);
  239. if (cached) {
  240. const data = JSON.parse(cached) as RerankResponse;
  241. return parseRerankResponse(data);
  242. }
  243. }
  244. const response = await fetch(`${OLLAMA_URL}/api/generate`, {
  245. method: "POST",
  246. headers: { "Content-Type": "application/json" },
  247. body: JSON.stringify(requestBody),
  248. });
  249. if (!response.ok) {
  250. const errorText = await response.text();
  251. if (!retried && (errorText.includes("not found") || errorText.includes("does not exist"))) {
  252. await ensureModelAvailable(model);
  253. return rerankSingle(prompt, model, db, true);
  254. }
  255. throw new Error(`Ollama API error: ${response.status} - ${errorText}`);
  256. }
  257. const data = await response.json() as RerankResponse;
  258. // Cache the result
  259. if (db) {
  260. setCachedResult(db, cacheKey, JSON.stringify(data));
  261. }
  262. return parseRerankResponse(data);
  263. }
  264. async function rerank(query: string, documents: { file: string; text: string }[], model: string = DEFAULT_RERANK_MODEL, db?: Database): Promise<{ file: string; score: number }[]> {
  265. const results: { file: string; score: number }[] = [];
  266. const total = documents.length;
  267. const PARALLEL = 5;
  268. process.stderr.write(`Reranking ${total} documents with ${model} (parallel: ${PARALLEL})...\n`);
  269. progress.indeterminate();
  270. // Process in parallel batches
  271. for (let i = 0; i < documents.length; i += PARALLEL) {
  272. const batch = documents.slice(i, i + PARALLEL);
  273. const batchResults = await Promise.all(
  274. batch.map(async (doc) => {
  275. try {
  276. // Extract title from filename for reranker context
  277. const title = doc.file.split('/').pop()?.replace(/\.md$/, '') || doc.file;
  278. const prompt = formatRerankPrompt(query, title, doc.text.slice(0, 4000));
  279. const score = await rerankSingle(prompt, model, db);
  280. return { file: doc.file, score };
  281. } catch (err) {
  282. return { file: doc.file, score: 0 };
  283. }
  284. })
  285. );
  286. results.push(...batchResults);
  287. const processed = Math.min(i + PARALLEL, total);
  288. progress.set((processed / total) * 100);
  289. process.stderr.write(`\rReranking: ${processed}/${total}`);
  290. }
  291. progress.clear();
  292. process.stderr.write("\n");
  293. return results.sort((a, b) => b.score - a.score);
  294. }
  295. function getOrCreateCollection(db: Database, pwd: string, globPattern: string): number {
  296. const now = new Date().toISOString();
  297. // Use INSERT OR IGNORE to handle race conditions, then SELECT
  298. db.prepare(`INSERT OR IGNORE INTO collections (pwd, glob_pattern, created_at) VALUES (?, ?, ?)`).run(pwd, globPattern, now);
  299. const existing = db.prepare(`SELECT id FROM collections WHERE pwd = ? AND glob_pattern = ?`).get(pwd, globPattern) as { id: number };
  300. return existing.id;
  301. }
  302. function cleanupDuplicateCollections(db: Database): void {
  303. // Remove duplicate collections keeping the oldest one
  304. db.exec(`
  305. DELETE FROM collections WHERE id NOT IN (
  306. SELECT MIN(id) FROM collections GROUP BY pwd, glob_pattern
  307. )
  308. `);
  309. // Remove bogus "." glob pattern entries (from earlier bug)
  310. db.exec(`DELETE FROM collections WHERE glob_pattern = '.'`);
  311. }
  312. function formatTimeAgo(date: Date): string {
  313. const seconds = Math.floor((Date.now() - date.getTime()) / 1000);
  314. if (seconds < 60) return `${seconds}s ago`;
  315. const minutes = Math.floor(seconds / 60);
  316. if (minutes < 60) return `${minutes}m ago`;
  317. const hours = Math.floor(minutes / 60);
  318. if (hours < 24) return `${hours}h ago`;
  319. const days = Math.floor(hours / 24);
  320. return `${days}d ago`;
  321. }
  322. function formatBytes(bytes: number): string {
  323. if (bytes < 1024) return `${bytes} B`;
  324. if (bytes < 1024 * 1024) return `${(bytes / 1024).toFixed(1)} KB`;
  325. if (bytes < 1024 * 1024 * 1024) return `${(bytes / (1024 * 1024)).toFixed(1)} MB`;
  326. return `${(bytes / (1024 * 1024 * 1024)).toFixed(1)} GB`;
  327. }
  328. function showStatus(): void {
  329. const dbPath = getDbPath();
  330. const db = getDb();
  331. // Cleanup any duplicate collections
  332. cleanupDuplicateCollections(db);
  333. // Index size
  334. let indexSize = 0;
  335. try {
  336. const stat = Bun.file(dbPath).size;
  337. indexSize = stat;
  338. } catch {}
  339. // Collections info
  340. const collections = db.prepare(`
  341. SELECT c.id, c.pwd, c.glob_pattern, c.created_at,
  342. COUNT(d.id) as doc_count,
  343. SUM(CASE WHEN d.active = 1 THEN 1 ELSE 0 END) as active_count,
  344. MAX(d.modified_at) as last_modified
  345. FROM collections c
  346. LEFT JOIN documents d ON d.collection_id = c.id
  347. GROUP BY c.id
  348. ORDER BY c.created_at DESC
  349. `).all() as { id: number; pwd: string; glob_pattern: string; created_at: string; doc_count: number; active_count: number; last_modified: string | null }[];
  350. // Overall stats
  351. const totalDocs = db.prepare(`SELECT COUNT(*) as count FROM documents WHERE active = 1`).get() as { count: number };
  352. const vectorCount = db.prepare(`SELECT COUNT(*) as count FROM content_vectors`).get() as { count: number };
  353. const needsEmbedding = getHashesNeedingEmbedding(db);
  354. // Most recent update across all collections
  355. const mostRecent = db.prepare(`SELECT MAX(modified_at) as latest FROM documents WHERE active = 1`).get() as { latest: string | null };
  356. console.log(`${c.bold}QMD Status${c.reset}\n`);
  357. console.log(`Index: ${dbPath}`);
  358. console.log(`Size: ${formatBytes(indexSize)}\n`);
  359. console.log(`${c.bold}Documents${c.reset}`);
  360. console.log(` Total: ${totalDocs.count} files indexed`);
  361. console.log(` Vectors: ${vectorCount.count} embedded`);
  362. if (needsEmbedding > 0) {
  363. console.log(` ${c.yellow}Pending: ${needsEmbedding} need embedding${c.reset} (run 'qmd embed')`);
  364. }
  365. if (mostRecent.latest) {
  366. const lastUpdate = new Date(mostRecent.latest);
  367. console.log(` Updated: ${formatTimeAgo(lastUpdate)}`);
  368. }
  369. // Get all path contexts
  370. const pathContexts = db.prepare(`SELECT path_prefix, context FROM path_contexts ORDER BY path_prefix`).all() as { path_prefix: string; context: string }[];
  371. if (collections.length > 0) {
  372. console.log(`\n${c.bold}Collections${c.reset}`);
  373. for (const col of collections) {
  374. const lastMod = col.last_modified ? formatTimeAgo(new Date(col.last_modified)) : "never";
  375. console.log(` ${c.cyan}${col.pwd}${c.reset}`);
  376. console.log(` ${col.glob_pattern} → ${col.active_count} docs (updated ${lastMod})`);
  377. // Show contexts that match this collection's path
  378. const matchingContexts = pathContexts.filter(ctx =>
  379. ctx.path_prefix.startsWith(col.pwd) || col.pwd.startsWith(ctx.path_prefix)
  380. );
  381. for (const ctx of matchingContexts) {
  382. const displayPath = shortPath(ctx.path_prefix);
  383. console.log(` ${c.dim}context: ${displayPath} → "${ctx.context}"${c.reset}`);
  384. }
  385. }
  386. } else {
  387. console.log(`\n${c.dim}No collections. Run 'qmd add .' to index markdown files.${c.reset}`);
  388. }
  389. closeDb();
  390. }
  391. // Update display_paths for all documents that have empty display_path
  392. function updateDisplayPaths(db: Database): number {
  393. // Get all docs with empty display_path, grouped by collection
  394. const emptyDocs = db.prepare(`
  395. SELECT d.id, d.filepath, c.pwd
  396. FROM documents d
  397. JOIN collections c ON d.collection_id = c.id
  398. WHERE d.active = 1 AND (d.display_path IS NULL OR d.display_path = '')
  399. `).all() as { id: number; filepath: string; pwd: string }[];
  400. if (emptyDocs.length === 0) return 0;
  401. // Collect existing display_paths
  402. const existingPaths = new Set<string>(
  403. (db.prepare(`SELECT display_path FROM documents WHERE active = 1 AND display_path != ''`).all() as { display_path: string }[])
  404. .map(r => r.display_path)
  405. );
  406. const updateStmt = db.prepare(`UPDATE documents SET display_path = ? WHERE id = ?`);
  407. let updated = 0;
  408. for (const doc of emptyDocs) {
  409. const displayPath = computeDisplayPath(doc.filepath, doc.pwd, existingPaths);
  410. updateStmt.run(displayPath, doc.id);
  411. existingPaths.add(displayPath);
  412. updated++;
  413. }
  414. return updated;
  415. }
  416. async function updateCollections(): Promise<void> {
  417. const db = getDb();
  418. cleanupDuplicateCollections(db);
  419. // Clear Ollama cache on update
  420. clearCache(db);
  421. const collections = db.prepare(`SELECT id, pwd, glob_pattern FROM collections`).all() as { id: number; pwd: string; glob_pattern: string }[];
  422. if (collections.length === 0) {
  423. console.log(`${c.dim}No collections found. Run 'qmd add .' to index markdown files.${c.reset}`);
  424. closeDb();
  425. return;
  426. }
  427. // Update display_paths for any documents missing them (migration)
  428. const pathsUpdated = updateDisplayPaths(db);
  429. if (pathsUpdated > 0) {
  430. console.log(`${c.green}✓${c.reset} Updated ${pathsUpdated} display paths`);
  431. }
  432. // Don't close db here - indexFiles will reuse it and close at the end
  433. console.log(`${c.bold}Updating ${collections.length} collection(s)...${c.reset}\n`);
  434. for (let i = 0; i < collections.length; i++) {
  435. const col = collections[i];
  436. console.log(`${c.cyan}[${i + 1}/${collections.length}]${c.reset} ${c.bold}${col.pwd}${c.reset}`);
  437. console.log(`${c.dim} Pattern: ${col.glob_pattern}${c.reset}`);
  438. // Temporarily set PWD for indexing
  439. const originalPwd = process.env.PWD;
  440. process.env.PWD = col.pwd;
  441. await indexFiles(col.glob_pattern);
  442. process.env.PWD = originalPwd;
  443. console.log("");
  444. }
  445. console.log(`${c.green}✓ All collections updated.${c.reset}`);
  446. }
  447. async function addContext(pathArg: string, contextText: string): Promise<void> {
  448. const db = getDb();
  449. const now = new Date().toISOString();
  450. // Resolve path - could be relative, absolute, or use ~
  451. let pathPrefix = pathArg;
  452. if (pathPrefix === '.' || pathPrefix === './') {
  453. pathPrefix = getPwd();
  454. } else if (pathPrefix.startsWith('~/')) {
  455. pathPrefix = homedir() + pathPrefix.slice(1);
  456. } else if (!pathPrefix.startsWith('/')) {
  457. pathPrefix = resolve(getPwd(), pathPrefix);
  458. }
  459. // Get realpath and normalize: remove trailing slash
  460. pathPrefix = getRealPath(pathPrefix).replace(/\/$/, '');
  461. // Insert or update
  462. db.prepare(`INSERT INTO path_contexts (path_prefix, context, created_at) VALUES (?, ?, ?)
  463. ON CONFLICT(path_prefix) DO UPDATE SET context = excluded.context`).run(pathPrefix, contextText, now);
  464. console.log(`${c.green}✓${c.reset} Added context for: ${shortPath(pathPrefix)}`);
  465. console.log(`${c.dim}Context: ${contextText}${c.reset}`);
  466. closeDb();
  467. }
  468. function getDocument(filename: string, fromLine?: number, maxLines?: number): void {
  469. const db = getDb();
  470. // Parse :linenum suffix from filename (e.g., "file.md:100")
  471. let filepath = filename;
  472. const colonMatch = filepath.match(/:(\d+)$/);
  473. if (colonMatch && !fromLine) {
  474. fromLine = parseInt(colonMatch[1], 10);
  475. filepath = filepath.slice(0, -colonMatch[0].length);
  476. }
  477. // Expand ~ to home directory
  478. if (filepath.startsWith('~/')) {
  479. filepath = homedir() + filepath.slice(1);
  480. }
  481. // Try exact match on filepath first
  482. let doc = db.prepare(`SELECT filepath, body FROM documents WHERE filepath = ? AND active = 1`).get(filepath) as { filepath: string; body: string } | null;
  483. // Try exact match on display_path
  484. if (!doc) {
  485. doc = db.prepare(`SELECT filepath, body FROM documents WHERE display_path = ? AND active = 1`).get(filepath) as { filepath: string; body: string } | null;
  486. }
  487. // Try matching by filename ending (allows partial paths)
  488. if (!doc) {
  489. doc = db.prepare(`SELECT filepath, body FROM documents WHERE filepath LIKE ? AND active = 1 LIMIT 1`).get(`%${filepath}`) as { filepath: string; body: string } | null;
  490. }
  491. // Try matching by display_path ending
  492. if (!doc) {
  493. doc = db.prepare(`SELECT filepath, body FROM documents WHERE display_path LIKE ? AND active = 1 LIMIT 1`).get(`%${filepath}`) as { filepath: string; body: string } | null;
  494. }
  495. if (!doc) {
  496. // Suggest similar files using Levenshtein distance
  497. const similar = findSimilarFiles(db, filepath, 5, 5);
  498. console.error(`Document not found: ${filename}`);
  499. if (similar.length > 0) {
  500. console.error(`\nDid you mean one of these?`);
  501. for (const s of similar) {
  502. console.error(` ${s}`);
  503. }
  504. }
  505. closeDb();
  506. process.exit(1);
  507. }
  508. // Get context for this file
  509. const context = getContextForFile(db, doc.filepath);
  510. let output = doc.body;
  511. // Apply line filtering if specified
  512. if (fromLine !== undefined || maxLines !== undefined) {
  513. const lines = output.split('\n');
  514. const start = (fromLine || 1) - 1; // Convert to 0-indexed
  515. const end = maxLines !== undefined ? start + maxLines : lines.length;
  516. output = lines.slice(start, end).join('\n');
  517. }
  518. // Output context header if exists
  519. if (context) {
  520. console.log(`Folder Context: ${context}\n---\n`);
  521. }
  522. console.log(output);
  523. closeDb();
  524. }
  525. // Multi-get: fetch multiple documents by glob pattern or comma-separated list
  526. function multiGet(pattern: string, maxLines?: number, maxBytes: number = DEFAULT_MULTI_GET_MAX_BYTES, format: OutputFormat = "cli"): void {
  527. const db = getDb();
  528. // Check if it's a comma-separated list or a glob pattern
  529. const isCommaSeparated = pattern.includes(',') && !pattern.includes('*') && !pattern.includes('?');
  530. let files: { filepath: string; displayPath: string; bodyLength: number }[];
  531. if (isCommaSeparated) {
  532. // Comma-separated list of files
  533. const names = pattern.split(',').map(s => s.trim()).filter(Boolean);
  534. files = [];
  535. for (const name of names) {
  536. // Try exact match on display_path first
  537. let doc = db.prepare(`SELECT filepath, display_path, LENGTH(body) as body_length FROM documents WHERE display_path = ? AND active = 1`).get(name) as { filepath: string; display_path: string; body_length: number } | null;
  538. // Try suffix match
  539. if (!doc) {
  540. doc = db.prepare(`SELECT filepath, display_path, LENGTH(body) as body_length FROM documents WHERE display_path LIKE ? AND active = 1 LIMIT 1`).get(`%${name}`) as { filepath: string; display_path: string; body_length: number } | null;
  541. }
  542. if (doc) {
  543. files.push({ filepath: doc.filepath, displayPath: doc.display_path, bodyLength: doc.body_length });
  544. } else {
  545. // Suggest similar files
  546. const similar = findSimilarFiles(db, name, 5, 3);
  547. console.error(`File not found: ${name}`);
  548. if (similar.length > 0) {
  549. console.error(` Did you mean: ${similar.join(', ')}`);
  550. }
  551. }
  552. }
  553. } else {
  554. // Glob pattern on display_path
  555. files = matchFilesByGlob(db, pattern);
  556. if (files.length === 0) {
  557. console.error(`No files matched pattern: ${pattern}`);
  558. closeDb();
  559. process.exit(1);
  560. }
  561. }
  562. // Collect results for structured output
  563. const results: { file: string; displayPath: string; title: string; body: string; context: string | null; skipped: boolean; skipReason?: string }[] = [];
  564. for (const file of files) {
  565. const context = getContextForFile(db, file.filepath);
  566. // Check size limit
  567. if (file.bodyLength > maxBytes) {
  568. results.push({
  569. file: file.filepath,
  570. displayPath: file.displayPath,
  571. title: file.displayPath.split('/').pop() || file.displayPath,
  572. body: "",
  573. context,
  574. skipped: true,
  575. skipReason: `File too large (${Math.round(file.bodyLength / 1024)}KB > ${Math.round(maxBytes / 1024)}KB). Use 'qmd get ${file.displayPath}' to retrieve.`,
  576. });
  577. continue;
  578. }
  579. const doc = db.prepare(`SELECT body, title FROM documents WHERE filepath = ? AND active = 1`).get(file.filepath) as { body: string; title: string } | null;
  580. if (!doc) continue;
  581. let body = doc.body;
  582. // Apply line limit if specified
  583. if (maxLines !== undefined) {
  584. const lines = body.split('\n');
  585. body = lines.slice(0, maxLines).join('\n');
  586. if (lines.length > maxLines) {
  587. body += `\n\n[... truncated ${lines.length - maxLines} more lines]`;
  588. }
  589. }
  590. results.push({
  591. file: file.filepath,
  592. displayPath: file.displayPath,
  593. title: doc.title || file.displayPath.split('/').pop() || file.displayPath,
  594. body,
  595. context,
  596. skipped: false,
  597. });
  598. }
  599. closeDb();
  600. // Output based on format
  601. if (format === "json") {
  602. const output = results.map(r => ({
  603. file: r.displayPath,
  604. title: r.title,
  605. ...(r.context && { context: r.context }),
  606. ...(r.skipped ? { skipped: true, reason: r.skipReason } : { body: r.body }),
  607. }));
  608. console.log(JSON.stringify(output, null, 2));
  609. } else if (format === "csv") {
  610. const escapeField = (val: string | null): string => {
  611. if (val === null || val === undefined) return "";
  612. const str = String(val);
  613. if (str.includes(",") || str.includes('"') || str.includes("\n")) {
  614. return `"${str.replace(/"/g, '""')}"`;
  615. }
  616. return str;
  617. };
  618. console.log("file,title,context,skipped,body");
  619. for (const r of results) {
  620. console.log([r.displayPath, r.title, r.context || "", r.skipped ? "true" : "false", r.skipped ? r.skipReason : r.body].map(escapeField).join(","));
  621. }
  622. } else if (format === "files") {
  623. for (const r of results) {
  624. const ctx = r.context ? `,"${r.context.replace(/"/g, '""')}"` : "";
  625. const status = r.skipped ? "[SKIPPED]" : "";
  626. console.log(`${r.displayPath}${ctx}${status ? `,${status}` : ""}`);
  627. }
  628. } else if (format === "md") {
  629. for (const r of results) {
  630. console.log(`## ${r.displayPath}\n`);
  631. if (r.title && r.title !== r.displayPath) console.log(`**Title:** ${r.title}\n`);
  632. if (r.context) console.log(`**Context:** ${r.context}\n`);
  633. if (r.skipped) {
  634. console.log(`> ${r.skipReason}\n`);
  635. } else {
  636. console.log("```");
  637. console.log(r.body);
  638. console.log("```\n");
  639. }
  640. }
  641. } else if (format === "xml") {
  642. console.log('<?xml version="1.0" encoding="UTF-8"?>');
  643. console.log("<documents>");
  644. for (const r of results) {
  645. console.log(" <document>");
  646. console.log(` <file>${escapeXml(r.displayPath)}</file>`);
  647. console.log(` <title>${escapeXml(r.title)}</title>`);
  648. if (r.context) console.log(` <context>${escapeXml(r.context)}</context>`);
  649. if (r.skipped) {
  650. console.log(` <skipped>true</skipped>`);
  651. console.log(` <reason>${escapeXml(r.skipReason || "")}</reason>`);
  652. } else {
  653. console.log(` <body>${escapeXml(r.body)}</body>`);
  654. }
  655. console.log(" </document>");
  656. }
  657. console.log("</documents>");
  658. } else {
  659. // CLI format (default)
  660. for (const r of results) {
  661. console.log(`\n${'='.repeat(60)}`);
  662. console.log(`File: ${r.displayPath}`);
  663. console.log(`${'='.repeat(60)}\n`);
  664. if (r.skipped) {
  665. console.log(`[SKIPPED: ${r.skipReason}]`);
  666. continue;
  667. }
  668. if (r.context) {
  669. console.log(`Folder Context: ${r.context}\n---\n`);
  670. }
  671. console.log(r.body);
  672. }
  673. }
  674. }
  675. // Get context for a filepath (finds most specific matching path prefix)
  676. function getContextForFile(db: Database, filepath: string): string | null {
  677. // Find all matching prefixes and return the longest (most specific) one
  678. const result = db.prepare(`
  679. SELECT context FROM path_contexts
  680. WHERE ? LIKE path_prefix || '%'
  681. ORDER BY LENGTH(path_prefix) DESC
  682. LIMIT 1
  683. `).get(filepath) as { context: string } | null;
  684. return result?.context || null;
  685. }
  686. async function dropCollection(globPattern: string): Promise<void> {
  687. const db = getDb();
  688. const pwd = getPwd();
  689. const collection = db.prepare(`SELECT id FROM collections WHERE pwd = ? AND glob_pattern = ?`).get(pwd, globPattern) as { id: number } | null;
  690. if (!collection) {
  691. // No collection to drop - this is fine, we'll create one during indexing
  692. return;
  693. }
  694. // Delete documents in this collection
  695. const deleted = db.prepare(`DELETE FROM documents WHERE collection_id = ?`).run(collection.id);
  696. // Delete the collection
  697. db.prepare(`DELETE FROM collections WHERE id = ?`).run(collection.id);
  698. console.log(`Dropped collection: ${pwd} (${globPattern})`);
  699. console.log(`Removed ${deleted.changes} documents`);
  700. console.log(`(Vectors kept for potential reuse)`);
  701. // Don't close db - indexFiles will use it and close at the end
  702. }
  703. async function indexFiles(globPattern: string = DEFAULT_GLOB): Promise<void> {
  704. const db = getDb();
  705. const pwd = getPwd();
  706. const now = new Date().toISOString();
  707. const excludeDirs = ["node_modules", ".git", ".cache", "vendor", "dist", "build"];
  708. // Clear Ollama cache on index
  709. clearCache(db);
  710. // Get or create collection for this (pwd, glob)
  711. const collectionId = getOrCreateCollection(db, pwd, globPattern);
  712. console.log(`Collection: ${pwd} (${globPattern})`);
  713. progress.indeterminate();
  714. const glob = new Glob(globPattern);
  715. const files: string[] = [];
  716. for await (const file of glob.scan({ cwd: pwd, onlyFiles: true, followSymlinks: true })) {
  717. // Skip node_modules, hidden folders (.*), and other common excludes
  718. const parts = file.split("/");
  719. const shouldSkip = parts.some(part =>
  720. part === "node_modules" ||
  721. part.startsWith(".") ||
  722. excludeDirs.includes(part)
  723. );
  724. if (!shouldSkip) {
  725. files.push(file);
  726. }
  727. }
  728. const total = files.length;
  729. if (total === 0) {
  730. progress.clear();
  731. console.log("No files found matching pattern.");
  732. closeDb();
  733. return;
  734. }
  735. const insertStmt = db.prepare(`INSERT INTO documents (collection_id, name, title, hash, filepath, display_path, body, created_at, modified_at, active) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, 1)`);
  736. const deactivateStmt = db.prepare(`UPDATE documents SET active = 0 WHERE collection_id = ? AND filepath = ? AND active = 1`);
  737. const findActiveStmt = db.prepare(`SELECT id, hash, title, display_path FROM documents WHERE collection_id = ? AND filepath = ? AND active = 1`);
  738. const findActiveAnyCollectionStmt = db.prepare(`SELECT id, collection_id, hash, title, display_path FROM documents WHERE filepath = ? AND active = 1`);
  739. const updateTitleStmt = db.prepare(`UPDATE documents SET title = ?, modified_at = ? WHERE id = ?`);
  740. const updateDisplayPathStmt = db.prepare(`UPDATE documents SET display_path = ? WHERE id = ?`);
  741. // Collect all existing display_paths for uniqueness check
  742. const existingDisplayPaths = new Set<string>(
  743. (db.prepare(`SELECT display_path FROM documents WHERE active = 1 AND display_path != ''`).all() as { display_path: string }[])
  744. .map(r => r.display_path)
  745. );
  746. let indexed = 0, updated = 0, unchanged = 0, processed = 0;
  747. const seenFiles = new Set<string>();
  748. const startTime = Date.now();
  749. for (const relativeFile of files) {
  750. const filepath = getRealPath(resolve(pwd, relativeFile));
  751. seenFiles.add(filepath);
  752. const content = await Bun.file(filepath).text();
  753. const hash = await hashContent(content);
  754. const name = relativeFile.replace(/\.md$/, "").split("/").pop() || relativeFile;
  755. const title = extractTitle(content, relativeFile);
  756. // First check if file exists in THIS collection
  757. const existing = findActiveStmt.get(collectionId, filepath) as { id: number; hash: string; title: string; display_path: string } | null;
  758. if (existing) {
  759. if (existing.hash === hash) {
  760. // Hash unchanged, but check if title needs updating
  761. if (existing.title !== title) {
  762. updateTitleStmt.run(title, now, existing.id);
  763. updated++;
  764. } else {
  765. unchanged++;
  766. }
  767. // Update display_path if empty
  768. if (!existing.display_path) {
  769. const displayPath = computeDisplayPath(filepath, pwd, existingDisplayPaths);
  770. updateDisplayPathStmt.run(displayPath, existing.id);
  771. existingDisplayPaths.add(displayPath);
  772. }
  773. } else {
  774. // Content changed - deactivate old, insert new
  775. existingDisplayPaths.delete(existing.display_path);
  776. deactivateStmt.run(collectionId, filepath);
  777. updated++;
  778. const stat = await Bun.file(filepath).stat();
  779. const displayPath = computeDisplayPath(filepath, pwd, existingDisplayPaths);
  780. insertStmt.run(collectionId, name, title, hash, filepath, displayPath, content, stat ? new Date(stat.birthtime).toISOString() : now, stat ? new Date(stat.mtime).toISOString() : now);
  781. existingDisplayPaths.add(displayPath);
  782. }
  783. } else {
  784. // Check if file exists in ANY collection (would violate unique constraint)
  785. const existingAnywhere = findActiveAnyCollectionStmt.get(filepath) as { id: number; collection_id: number; hash: string; title: string; display_path: string } | null;
  786. if (existingAnywhere) {
  787. // File already indexed in another collection - skip it
  788. unchanged++;
  789. } else {
  790. indexed++;
  791. const stat = await Bun.file(filepath).stat();
  792. const displayPath = computeDisplayPath(filepath, pwd, existingDisplayPaths);
  793. insertStmt.run(collectionId, name, title, hash, filepath, displayPath, content, stat ? new Date(stat.birthtime).toISOString() : now, stat ? new Date(stat.mtime).toISOString() : now);
  794. existingDisplayPaths.add(displayPath);
  795. }
  796. }
  797. processed++;
  798. progress.set((processed / total) * 100);
  799. const elapsed = (Date.now() - startTime) / 1000;
  800. const rate = processed / elapsed;
  801. const remaining = (total - processed) / rate;
  802. const eta = processed > 2 ? ` ETA: ${formatETA(remaining)}` : "";
  803. process.stderr.write(`\rIndexing: ${processed}/${total}${eta} `);
  804. }
  805. // Deactivate documents in this collection that no longer exist
  806. const allActive = db.prepare(`SELECT filepath FROM documents WHERE collection_id = ? AND active = 1`).all(collectionId) as { filepath: string }[];
  807. let removed = 0;
  808. for (const row of allActive) {
  809. if (!seenFiles.has(row.filepath)) {
  810. deactivateStmt.run(collectionId, row.filepath);
  811. removed++;
  812. }
  813. }
  814. // Check if vector index needs updating
  815. const needsEmbedding = getHashesNeedingEmbedding(db);
  816. progress.clear();
  817. console.log(`\nIndexed: ${indexed} new, ${updated} updated, ${unchanged} unchanged, ${removed} removed`);
  818. if (needsEmbedding > 0) {
  819. console.log(`\nRun 'qmd embed' to update embeddings (${needsEmbedding} unique hashes need vectors)`);
  820. }
  821. closeDb();
  822. }
  823. function renderProgressBar(percent: number, width: number = 30): string {
  824. const filled = Math.round((percent / 100) * width);
  825. const empty = width - filled;
  826. const bar = "█".repeat(filled) + "░".repeat(empty);
  827. return bar;
  828. }
  829. async function vectorIndex(model: string = DEFAULT_EMBED_MODEL, force: boolean = false): Promise<void> {
  830. const db = getDb();
  831. const now = new Date().toISOString();
  832. // If force, clear all vectors
  833. if (force) {
  834. console.log(`${c.yellow}Force re-indexing: clearing all vectors...${c.reset}`);
  835. db.exec(`DELETE FROM content_vectors`);
  836. db.exec(`DROP TABLE IF EXISTS vectors_vec`);
  837. }
  838. // Find unique hashes that need embedding (from active documents)
  839. // Use MIN(filepath) to get one representative filepath per hash
  840. const hashesToEmbed = db.prepare(`
  841. SELECT d.hash, d.body, MIN(d.filepath) as filepath, MIN(d.display_path) as display_path
  842. FROM documents d
  843. LEFT JOIN content_vectors v ON d.hash = v.hash AND v.seq = 0
  844. WHERE d.active = 1 AND v.hash IS NULL
  845. GROUP BY d.hash
  846. `).all() as { hash: string; body: string; filepath: string; display_path: string }[];
  847. if (hashesToEmbed.length === 0) {
  848. console.log(`${c.green}✓ All content hashes already have embeddings.${c.reset}`);
  849. closeDb();
  850. return;
  851. }
  852. // Prepare documents with chunks
  853. type ChunkItem = { hash: string; title: string; text: string; seq: number; pos: number; bytes: number; displayName: string };
  854. const allChunks: ChunkItem[] = [];
  855. let multiChunkDocs = 0;
  856. for (const item of hashesToEmbed) {
  857. const encoder = new TextEncoder();
  858. const bodyBytes = encoder.encode(item.body).length;
  859. if (bodyBytes === 0) continue; // Skip empty
  860. const title = extractTitle(item.body, item.filepath);
  861. const displayName = item.display_path || item.filepath;
  862. const chunks = chunkDocument(item.body, CHUNK_BYTE_SIZE);
  863. if (chunks.length > 1) multiChunkDocs++;
  864. for (let seq = 0; seq < chunks.length; seq++) {
  865. allChunks.push({
  866. hash: item.hash,
  867. title,
  868. text: chunks[seq].text,
  869. seq,
  870. pos: chunks[seq].pos,
  871. bytes: encoder.encode(chunks[seq].text).length,
  872. displayName,
  873. });
  874. }
  875. }
  876. if (allChunks.length === 0) {
  877. console.log(`${c.green}✓ No non-empty documents to embed.${c.reset}`);
  878. closeDb();
  879. return;
  880. }
  881. const totalBytes = allChunks.reduce((sum, c) => sum + c.bytes, 0);
  882. const totalChunks = allChunks.length;
  883. const totalDocs = hashesToEmbed.length;
  884. console.log(`${c.bold}Embedding ${totalDocs} documents${c.reset} ${c.dim}(${totalChunks} chunks, ${formatBytes(totalBytes)})${c.reset}`);
  885. if (multiChunkDocs > 0) {
  886. console.log(`${c.dim}${multiChunkDocs} documents split into multiple chunks${c.reset}`);
  887. }
  888. console.log(`${c.dim}Model: ${model}${c.reset}\n`);
  889. // Hide cursor during embedding
  890. cursor.hide();
  891. // Get embedding dimensions from first chunk
  892. progress.indeterminate();
  893. const firstEmbedding = await getEmbedding(allChunks[0].text, model, false, allChunks[0].title);
  894. ensureVecTable(db, firstEmbedding.length);
  895. const insertVecStmt = db.prepare(`INSERT OR REPLACE INTO vectors_vec (hash_seq, embedding) VALUES (?, ?)`);
  896. const insertContentVectorStmt = db.prepare(`INSERT OR REPLACE INTO content_vectors (hash, seq, pos, model, embedded_at) VALUES (?, ?, ?, ?, ?)`);
  897. let chunksEmbedded = 0, errors = 0, bytesProcessed = 0;
  898. const startTime = Date.now();
  899. // Insert first chunk
  900. const firstHashSeq = `${allChunks[0].hash}_${allChunks[0].seq}`;
  901. insertVecStmt.run(firstHashSeq, new Float32Array(firstEmbedding));
  902. insertContentVectorStmt.run(allChunks[0].hash, allChunks[0].seq, allChunks[0].pos, model, now);
  903. chunksEmbedded++;
  904. bytesProcessed += allChunks[0].bytes;
  905. for (let i = 1; i < allChunks.length; i++) {
  906. const chunk = allChunks[i];
  907. try {
  908. const embedding = await getEmbedding(chunk.text, model, false, chunk.title);
  909. const hashSeq = `${chunk.hash}_${chunk.seq}`;
  910. insertVecStmt.run(hashSeq, new Float32Array(embedding));
  911. insertContentVectorStmt.run(chunk.hash, chunk.seq, chunk.pos, model, now);
  912. chunksEmbedded++;
  913. bytesProcessed += chunk.bytes;
  914. } catch (err) {
  915. errors++;
  916. bytesProcessed += chunk.bytes;
  917. progress.error();
  918. console.error(`\n${c.yellow}⚠ Error embedding "${chunk.displayName}" chunk ${chunk.seq}: ${err}${c.reset}`);
  919. }
  920. const percent = (bytesProcessed / totalBytes) * 100;
  921. progress.set(percent);
  922. const elapsed = (Date.now() - startTime) / 1000;
  923. const bytesPerSec = bytesProcessed / elapsed;
  924. const remainingBytes = totalBytes - bytesProcessed;
  925. const etaSec = remainingBytes / bytesPerSec;
  926. const bar = renderProgressBar(percent);
  927. const percentStr = percent.toFixed(0).padStart(3);
  928. const throughput = `${formatBytes(bytesPerSec)}/s`;
  929. const eta = elapsed > 2 ? formatETA(etaSec) : "...";
  930. const errStr = errors > 0 ? ` ${c.yellow}${errors} err${c.reset}` : "";
  931. process.stderr.write(`\r${c.cyan}${bar}${c.reset} ${c.bold}${percentStr}%${c.reset} ${c.dim}${chunksEmbedded}/${totalChunks}${c.reset}${errStr} ${c.dim}${throughput} ETA ${eta}${c.reset} `);
  932. }
  933. progress.clear();
  934. cursor.show();
  935. const totalTimeSec = (Date.now() - startTime) / 1000;
  936. const avgThroughput = formatBytes(totalBytes / totalTimeSec);
  937. console.log(`\r${c.green}${renderProgressBar(100)}${c.reset} ${c.bold}100%${c.reset} `);
  938. console.log(`\n${c.green}✓ Done!${c.reset} Embedded ${c.bold}${chunksEmbedded}${c.reset} chunks from ${c.bold}${totalDocs}${c.reset} documents in ${c.bold}${formatETA(totalTimeSec)}${c.reset} ${c.dim}(${avgThroughput}/s)${c.reset}`);
  939. if (errors > 0) {
  940. console.log(`${c.yellow}⚠ ${errors} chunks failed${c.reset}`);
  941. }
  942. closeDb();
  943. }
  944. // Sanitize a term for FTS5: remove punctuation except apostrophes
  945. function sanitizeFTS5Term(term: string): string {
  946. // Remove all non-alphanumeric except apostrophes (for contractions like "don't")
  947. return term.replace(/[^\w']/g, '').trim();
  948. }
  949. // Build FTS5 query: phrase-aware with fallback to individual terms
  950. function buildFTS5Query(query: string): string {
  951. // Sanitize the full query for phrase matching
  952. const sanitizedQuery = query.replace(/[^\w\s']/g, '').trim();
  953. const terms = query
  954. .split(/\s+/)
  955. .map(sanitizeFTS5Term)
  956. .filter(term => term.length >= 2); // Skip single chars and empty
  957. if (terms.length === 0) return "";
  958. if (terms.length === 1) return `"${terms[0].replace(/"/g, '""')}"`;
  959. // Strategy: exact phrase OR proximity match OR individual terms
  960. // Exact phrase matches rank highest, then close proximity, then any term
  961. const phrase = `"${sanitizedQuery.replace(/"/g, '""')}"`;
  962. const quotedTerms = terms.map(t => `"${t.replace(/"/g, '""')}"`);
  963. // FTS5 NEAR syntax: NEAR(term1 term2, distance)
  964. const nearPhrase = `NEAR(${quotedTerms.join(' ')}, 10)`;
  965. const orTerms = quotedTerms.join(' OR ');
  966. // Exact phrase > proximity > any term
  967. return `(${phrase}) OR (${nearPhrase}) OR (${orTerms})`;
  968. }
  969. // Normalize BM25 score to 0-1 range using sigmoid
  970. function normalizeBM25(score: number): number {
  971. // BM25 scores are negative in SQLite (lower = better)
  972. // Typical range: -15 (excellent) to -2 (weak match)
  973. // Map to 0-1 where higher is better
  974. const absScore = Math.abs(score);
  975. // Sigmoid-ish normalization: maps ~2-15 range to ~0.1-0.95
  976. return 1 / (1 + Math.exp(-(absScore - 5) / 3));
  977. }
  978. // Get collection ID by name (matches pwd or glob_pattern suffix)
  979. function getCollectionIdByName(db: Database, name: string): number | null {
  980. // Search both pwd and glob_pattern columns for the name
  981. const result = db.prepare(`
  982. SELECT id FROM collections
  983. WHERE pwd LIKE ? OR glob_pattern LIKE ?
  984. ORDER BY LENGTH(pwd) DESC
  985. LIMIT 1
  986. `).get(`%${name}%`, `%${name}%`) as { id: number } | null;
  987. return result?.id || null;
  988. }
  989. function searchFTS(db: Database, query: string, limit: number = 20, collectionId?: number): SearchResult[] {
  990. const ftsQuery = buildFTS5Query(query);
  991. if (!ftsQuery) return [];
  992. // BM25 weights: name=10, body=1 (title matches ranked higher)
  993. let sql = `
  994. SELECT d.filepath, d.display_path, d.title, d.body, bm25(documents_fts, 10.0, 1.0) as score
  995. FROM documents_fts f
  996. JOIN documents d ON d.id = f.rowid
  997. WHERE documents_fts MATCH ? AND d.active = 1
  998. `;
  999. const params: (string | number)[] = [ftsQuery];
  1000. if (collectionId !== undefined) {
  1001. sql += ` AND d.collection_id = ?`;
  1002. params.push(collectionId);
  1003. }
  1004. sql += ` ORDER BY score LIMIT ?`;
  1005. params.push(limit);
  1006. const stmt = db.prepare(sql);
  1007. const results = stmt.all(...params) as { filepath: string; display_path: string; title: string; body: string; score: number }[];
  1008. return results.map(r => ({
  1009. file: r.filepath,
  1010. displayPath: r.display_path,
  1011. title: r.title,
  1012. body: r.body,
  1013. score: normalizeBM25(r.score),
  1014. source: "fts" as const,
  1015. }));
  1016. }
  1017. async function searchVec(db: Database, query: string, model: string, limit: number = 20, collectionId?: number): Promise<SearchResult[]> {
  1018. const tableExists = db.prepare(`SELECT name FROM sqlite_master WHERE type='table' AND name='vectors_vec'`).get();
  1019. if (!tableExists) return [];
  1020. const queryEmbedding = await getEmbedding(query, model, true);
  1021. const queryVec = new Float32Array(queryEmbedding);
  1022. // Join: vectors_vec -> content_vectors -> documents
  1023. // Over-retrieve to handle multiple chunks per document, then dedupe
  1024. let sql = `
  1025. SELECT d.filepath, d.display_path, d.title, d.body, vec.distance, cv.pos
  1026. FROM vectors_vec vec
  1027. JOIN content_vectors cv ON vec.hash_seq = cv.hash || '_' || cv.seq
  1028. JOIN documents d ON d.hash = cv.hash AND d.active = 1
  1029. WHERE vec.embedding MATCH ? AND k = ?
  1030. `;
  1031. if (collectionId !== undefined) {
  1032. sql += ` AND d.collection_id = ${collectionId}`;
  1033. }
  1034. sql += ` ORDER BY vec.distance`;
  1035. const stmt = db.prepare(sql);
  1036. const rawResults = stmt.all(queryVec, limit * 3) as { filepath: string; display_path: string; title: string; body: string; distance: number; pos: number }[];
  1037. // Aggregate chunks per document: max score + small bonus for additional matches
  1038. const byFile = new Map<string, { filepath: string; displayPath: string; title: string; body: string; chunkCount: number; bestPos: number; bestDist: number }>();
  1039. for (const r of rawResults) {
  1040. const existing = byFile.get(r.filepath);
  1041. if (!existing) {
  1042. byFile.set(r.filepath, { filepath: r.filepath, displayPath: r.display_path, title: r.title, body: r.body, chunkCount: 1, bestPos: r.pos, bestDist: r.distance });
  1043. } else {
  1044. existing.chunkCount++;
  1045. if (r.distance < existing.bestDist) {
  1046. existing.bestDist = r.distance;
  1047. existing.bestPos = r.pos;
  1048. }
  1049. }
  1050. }
  1051. // Score = max chunk score + 0.02 bonus per additional chunk (capped at +0.1)
  1052. return Array.from(byFile.values())
  1053. .map(r => {
  1054. const maxScore = 1 / (1 + r.bestDist);
  1055. const bonusChunks = Math.min(r.chunkCount - 1, 5);
  1056. const bonus = bonusChunks * 0.02;
  1057. return {
  1058. file: r.filepath,
  1059. displayPath: r.displayPath,
  1060. title: r.title,
  1061. body: r.body,
  1062. score: maxScore + bonus,
  1063. source: "vec" as const,
  1064. chunkPos: r.bestPos,
  1065. };
  1066. })
  1067. .sort((a, b) => b.score - a.score)
  1068. .slice(0, limit);
  1069. }
  1070. function normalizeScores(results: SearchResult[]): SearchResult[] {
  1071. if (results.length === 0) return results;
  1072. const maxScore = Math.max(...results.map(r => r.score));
  1073. const minScore = Math.min(...results.map(r => r.score));
  1074. const range = maxScore - minScore || 1;
  1075. return results.map(r => ({ ...r, score: (r.score - minScore) / range }));
  1076. }
  1077. // Reciprocal Rank Fusion: combines multiple ranked lists
  1078. // RRF score = sum(1 / (k + rank)) across all lists where doc appears
  1079. // k=60 is standard, provides good balance between top and lower ranks
  1080. export type RankedResult = { file: string; displayPath: string; title: string; body: string; score: number };
  1081. function reciprocalRankFusion(
  1082. resultLists: RankedResult[][],
  1083. weights: number[] = [], // Weight per result list (default 1.0)
  1084. k: number = 60
  1085. ): RankedResult[] {
  1086. const scores = new Map<string, { score: number; displayPath: string; title: string; body: string; bestRank: number }>();
  1087. for (let listIdx = 0; listIdx < resultLists.length; listIdx++) {
  1088. const results = resultLists[listIdx];
  1089. const weight = weights[listIdx] ?? 1.0;
  1090. for (let rank = 0; rank < results.length; rank++) {
  1091. const doc = results[rank];
  1092. const rrfScore = weight / (k + rank + 1);
  1093. const existing = scores.get(doc.file);
  1094. if (existing) {
  1095. existing.score += rrfScore;
  1096. existing.bestRank = Math.min(existing.bestRank, rank);
  1097. } else {
  1098. scores.set(doc.file, { score: rrfScore, displayPath: doc.displayPath, title: doc.title, body: doc.body, bestRank: rank });
  1099. }
  1100. }
  1101. }
  1102. // Add bonus for best rank: documents that ranked #1-3 in any list get a boost
  1103. // This prevents dilution of exact matches by expansion queries
  1104. return Array.from(scores.entries())
  1105. .map(([file, { score, displayPath, title, body, bestRank }]) => {
  1106. let bonus = 0;
  1107. if (bestRank === 0) bonus = 0.05; // Ranked #1 somewhere
  1108. else if (bestRank <= 2) bonus = 0.02; // Ranked top-3 somewhere
  1109. return { file, displayPath, title, body, score: score + bonus };
  1110. })
  1111. .sort((a, b) => b.score - a.score);
  1112. }
  1113. type OutputOptions = {
  1114. format: OutputFormat;
  1115. full: boolean;
  1116. limit: number;
  1117. minScore: number;
  1118. all?: boolean;
  1119. collection?: string; // Filter by collection name (pwd suffix match)
  1120. };
  1121. // Extract snippet with more context lines for CLI display
  1122. function extractSnippetWithContext(body: string, query: string, contextLines = 3, chunkPos?: number): { line: number; snippet: string; hasMatch: boolean } {
  1123. // If chunkPos provided, focus search on that area
  1124. let lineOffset = 0;
  1125. let searchBody = body;
  1126. if (chunkPos && chunkPos > 0) {
  1127. const contextStart = Math.max(0, chunkPos - 200);
  1128. searchBody = body.slice(contextStart);
  1129. if (contextStart > 0) {
  1130. lineOffset = body.slice(0, contextStart).split('\n').length - 1;
  1131. }
  1132. }
  1133. const lines = searchBody.split('\n');
  1134. const queryTerms = query.toLowerCase().split(/\s+/).filter(t => t.length > 0);
  1135. let bestLine = 0, bestScore = -1;
  1136. for (let i = 0; i < lines.length; i++) {
  1137. const lineLower = lines[i].toLowerCase();
  1138. let score = 0;
  1139. for (const term of queryTerms) {
  1140. if (lineLower.includes(term)) score++;
  1141. }
  1142. if (score > bestScore) {
  1143. bestScore = score;
  1144. bestLine = i;
  1145. }
  1146. }
  1147. // No query match found - return beginning of chunk area or file
  1148. if (bestScore <= 0) {
  1149. const preview = lines.slice(0, contextLines * 2).join('\n').trim();
  1150. return { line: lineOffset + 1, snippet: preview, hasMatch: false };
  1151. }
  1152. const startLine = Math.max(0, bestLine - contextLines);
  1153. const endLine = Math.min(lines.length, bestLine + contextLines + 1);
  1154. const snippet = lines.slice(startLine, endLine).join('\n').trim();
  1155. return { line: lineOffset + bestLine + 1, snippet, hasMatch: true };
  1156. }
  1157. // Highlight query terms in text (skip short words < 3 chars)
  1158. function highlightTerms(text: string, query: string): string {
  1159. if (!useColor) return text;
  1160. const terms = query.toLowerCase().split(/\s+/).filter(t => t.length >= 3);
  1161. let result = text;
  1162. for (const term of terms) {
  1163. const regex = new RegExp(`(${term.replace(/[.*+?^${}()|[\]\\]/g, '\\$&')})`, 'gi');
  1164. result = result.replace(regex, `${c.yellow}${c.bold}$1${c.reset}`);
  1165. }
  1166. return result;
  1167. }
  1168. // Format score with color based on value
  1169. function formatScore(score: number): string {
  1170. const pct = (score * 100).toFixed(0).padStart(3);
  1171. if (!useColor) return `${pct}%`;
  1172. if (score >= 0.7) return `${c.green}${pct}%${c.reset}`;
  1173. if (score >= 0.4) return `${c.yellow}${pct}%${c.reset}`;
  1174. return `${c.dim}${pct}%${c.reset}`;
  1175. }
  1176. // Shorten directory path for display - relative to $HOME (used for context paths, not documents)
  1177. function shortPath(dirpath: string): string {
  1178. const home = homedir();
  1179. if (dirpath.startsWith(home)) {
  1180. return '~' + dirpath.slice(home.length);
  1181. }
  1182. return dirpath;
  1183. }
  1184. function outputResults(results: { file: string; displayPath: string; title: string; body: string; score: number; context?: string | null; chunkPos?: number }[], query: string, opts: OutputOptions): void {
  1185. const filtered = results.filter(r => r.score >= opts.minScore).slice(0, opts.limit);
  1186. if (filtered.length === 0) {
  1187. console.log("No results found above minimum score threshold.");
  1188. return;
  1189. }
  1190. if (opts.format === "json") {
  1191. // JSON output for LLM consumption
  1192. const output = filtered.map(row => ({
  1193. score: Math.round(row.score * 100) / 100,
  1194. file: row.displayPath,
  1195. title: row.title,
  1196. ...(row.context && { context: row.context }),
  1197. ...(opts.full && { body: row.body }),
  1198. ...(!opts.full && { snippet: extractSnippet(row.body, query, 300, row.chunkPos).snippet }),
  1199. }));
  1200. console.log(JSON.stringify(output, null, 2));
  1201. } else if (opts.format === "files") {
  1202. // Simple score,filepath,context output
  1203. for (const row of filtered) {
  1204. const ctx = row.context ? `,"${row.context.replace(/"/g, '""')}"` : "";
  1205. console.log(`${row.score.toFixed(2)},${row.displayPath}${ctx}`);
  1206. }
  1207. } else if (opts.format === "cli") {
  1208. for (let i = 0; i < filtered.length; i++) {
  1209. const row = filtered[i];
  1210. const { line, snippet, hasMatch } = extractSnippetWithContext(row.body, query, 2, row.chunkPos);
  1211. // Line 1: filepath
  1212. const path = row.displayPath;
  1213. const lineInfo = hasMatch ? `:${line}` : "";
  1214. console.log(`${c.cyan}${path}${c.dim}${lineInfo}${c.reset}`);
  1215. // Line 2: Title (if available)
  1216. if (row.title) {
  1217. console.log(`${c.bold}Title: ${row.title}${c.reset}`);
  1218. }
  1219. // Line 3: Context (if available)
  1220. if (row.context) {
  1221. console.log(`${c.dim}Context: ${row.context}${c.reset}`);
  1222. }
  1223. // Line 4: Score
  1224. const score = formatScore(row.score);
  1225. console.log(`Score: ${c.bold}${score}${c.reset}`);
  1226. console.log();
  1227. // Snippet with highlighting (no leading | chars for better word wrap)
  1228. const highlighted = highlightTerms(snippet, query);
  1229. console.log(highlighted);
  1230. // Double empty line between results
  1231. if (i < filtered.length - 1) console.log('\n');
  1232. }
  1233. } else if (opts.format === "md") {
  1234. for (const row of filtered) {
  1235. const heading = row.title || row.displayPath;
  1236. if (opts.full) {
  1237. console.log(`---\n# ${heading}\n\n${row.body}\n`);
  1238. } else {
  1239. const { snippet } = extractSnippet(row.body, query, 500, row.chunkPos);
  1240. console.log(`---\n# ${heading}\n\n${snippet}\n`);
  1241. }
  1242. }
  1243. } else if (opts.format === "xml") {
  1244. for (const row of filtered) {
  1245. const titleAttr = row.title ? ` title="${row.title.replace(/"/g, '&quot;')}"` : "";
  1246. if (opts.full) {
  1247. console.log(`<file name="${row.displayPath}"${titleAttr}>\n${row.body}\n</file>\n`);
  1248. } else {
  1249. const { snippet } = extractSnippet(row.body, query, 500, row.chunkPos);
  1250. console.log(`<file name="${row.displayPath}"${titleAttr}>\n${snippet}\n</file>\n`);
  1251. }
  1252. }
  1253. } else {
  1254. // CSV format
  1255. console.log("score,file,title,line,snippet");
  1256. for (const row of filtered) {
  1257. const { line, snippet } = extractSnippet(row.body, query, 500, row.chunkPos);
  1258. const content = opts.full ? row.body : snippet;
  1259. console.log(`${row.score.toFixed(4)},${escapeCSV(row.displayPath)},${escapeCSV(row.title)},${line},${escapeCSV(content)}`);
  1260. }
  1261. }
  1262. }
  1263. function search(query: string, opts: OutputOptions): void {
  1264. const db = getDb();
  1265. // Resolve collection filter if specified
  1266. let collectionId: number | undefined;
  1267. if (opts.collection) {
  1268. collectionId = getCollectionIdByName(db, opts.collection) ?? undefined;
  1269. if (collectionId === undefined) {
  1270. console.error(`Collection not found: ${opts.collection}`);
  1271. closeDb();
  1272. process.exit(1);
  1273. }
  1274. }
  1275. // Use large limit for --all, otherwise fetch more than needed and let outputResults filter
  1276. const fetchLimit = opts.all ? 100000 : Math.max(50, opts.limit * 2);
  1277. const results = searchFTS(db, query, fetchLimit, collectionId);
  1278. // Add context to results
  1279. const resultsWithContext = results.map(r => ({
  1280. ...r,
  1281. context: getContextForFile(db, r.file),
  1282. }));
  1283. closeDb();
  1284. if (resultsWithContext.length === 0) {
  1285. console.log("No results found.");
  1286. return;
  1287. }
  1288. outputResults(resultsWithContext, query, opts);
  1289. }
  1290. async function vectorSearch(query: string, opts: OutputOptions, model: string = DEFAULT_EMBED_MODEL): Promise<void> {
  1291. const db = getDb();
  1292. // Resolve collection filter if specified
  1293. let collectionId: number | undefined;
  1294. if (opts.collection) {
  1295. collectionId = getCollectionIdByName(db, opts.collection) ?? undefined;
  1296. if (collectionId === undefined) {
  1297. console.error(`Collection not found: ${opts.collection}`);
  1298. closeDb();
  1299. process.exit(1);
  1300. }
  1301. }
  1302. const tableExists = db.prepare(`SELECT name FROM sqlite_master WHERE type='table' AND name='vectors_vec'`).get();
  1303. if (!tableExists) {
  1304. console.error("Vector index not found. Run 'qmd embed' first to create embeddings.");
  1305. closeDb();
  1306. return;
  1307. }
  1308. // Check index health and warn about issues
  1309. checkIndexHealth(db);
  1310. // Expand query to multiple variations (with caching)
  1311. const queries = await expandQuery(query, DEFAULT_QUERY_MODEL, db);
  1312. process.stderr.write(`Searching with ${queries.length} query variations...\n`);
  1313. // Collect results from all query variations
  1314. // For --all, fetch more results per query
  1315. const perQueryLimit = opts.all ? 500 : 20;
  1316. const allResults = new Map<string, { file: string; displayPath: string; title: string; body: string; score: number }>();
  1317. for (const q of queries) {
  1318. const vecResults = await searchVec(db, q, model, perQueryLimit, collectionId);
  1319. for (const r of vecResults) {
  1320. const existing = allResults.get(r.file);
  1321. if (!existing || r.score > existing.score) {
  1322. allResults.set(r.file, { file: r.file, displayPath: r.displayPath, title: r.title, body: r.body, score: r.score });
  1323. }
  1324. }
  1325. }
  1326. // Sort by max score and limit to requested count
  1327. const results = Array.from(allResults.values())
  1328. .sort((a, b) => b.score - a.score)
  1329. .slice(0, opts.limit)
  1330. .map(r => ({ ...r, context: getContextForFile(db, r.file) }));
  1331. closeDb();
  1332. if (results.length === 0) {
  1333. console.log("No results found.");
  1334. return;
  1335. }
  1336. outputResults(results, query, { ...opts, limit: results.length }); // Already limited
  1337. }
  1338. async function expandQuery(query: string, model: string = DEFAULT_QUERY_MODEL, db?: Database): Promise<string[]> {
  1339. process.stderr.write("Generating query variations...\n");
  1340. const prompt = `You are a search query expander. Given a search query, generate 2 alternative queries that would help find relevant documents.
  1341. Rules:
  1342. - Use synonyms and related terminology (e.g., "craft" → "craftsmanship", "quality", "excellence")
  1343. - Rephrase to capture different angles (e.g., "engineering culture" → "technical excellence", "developer practices")
  1344. - Keep proper nouns and named concepts exactly as written (e.g., "Build a Business", "Stripe", "Shopify")
  1345. - Each variation should be 3-8 words, natural search terms
  1346. - Do NOT just append words like "search" or "find" or "documents"
  1347. Query: "${query}"
  1348. Output exactly 2 variations, one per line, no numbering or bullets:`;
  1349. const requestBody = {
  1350. model,
  1351. prompt,
  1352. stream: false,
  1353. think: false,
  1354. options: { num_predict: 150 },
  1355. };
  1356. // Check cache
  1357. const cacheDb = db || getDb();
  1358. const cacheKey = getCacheKey(`${OLLAMA_URL}/api/generate`, requestBody);
  1359. const cached = getCachedResult(cacheDb, cacheKey);
  1360. let responseText: string;
  1361. if (cached) {
  1362. responseText = cached;
  1363. } else {
  1364. const response = await fetch(`${OLLAMA_URL}/api/generate`, {
  1365. method: "POST",
  1366. headers: { "Content-Type": "application/json" },
  1367. body: JSON.stringify(requestBody),
  1368. });
  1369. if (!response.ok) {
  1370. const errorText = await response.text();
  1371. if (errorText.includes("not found") || errorText.includes("does not exist")) {
  1372. await ensureModelAvailable(model);
  1373. if (!db) cacheDb.close();
  1374. return expandQuery(query, model, db);
  1375. }
  1376. if (!db) cacheDb.close();
  1377. return [query];
  1378. }
  1379. const data = await response.json() as { response: string };
  1380. responseText = data.response;
  1381. setCachedResult(cacheDb, cacheKey, responseText);
  1382. }
  1383. if (!db) cacheDb.close();
  1384. const lines = responseText.trim().split('\n')
  1385. .map(l => l.replace(/^[\d\.\-\*\"\s]+/, '').replace(/["\s]+$/, '').trim())
  1386. .filter(l => l.length > 2 && l.length < 100 && !l.startsWith('<') && !l.toLowerCase().includes('variation'))
  1387. .slice(0, 2);
  1388. const allQueries = [query, ...lines];
  1389. process.stderr.write(`${c.dim}Queries: ${allQueries.join(' | ')}${c.reset}\n`);
  1390. return allQueries;
  1391. }
  1392. async function querySearch(query: string, opts: OutputOptions, embedModel: string = DEFAULT_EMBED_MODEL, rerankModel: string = DEFAULT_RERANK_MODEL): Promise<void> {
  1393. const db = getDb();
  1394. // Resolve collection filter if specified
  1395. let collectionId: number | undefined;
  1396. if (opts.collection) {
  1397. collectionId = getCollectionIdByName(db, opts.collection) ?? undefined;
  1398. if (collectionId === undefined) {
  1399. console.error(`Collection not found: ${opts.collection}`);
  1400. closeDb();
  1401. process.exit(1);
  1402. }
  1403. }
  1404. // Check index health and warn about issues
  1405. checkIndexHealth(db);
  1406. // Expand query to multiple variations (with caching)
  1407. const queries = await expandQuery(query, DEFAULT_QUERY_MODEL, db);
  1408. process.stderr.write(`Searching with ${queries.length} query variations...\n`);
  1409. // Collect ranked result lists for RRF fusion
  1410. const rankedLists: RankedResult[][] = [];
  1411. const hasVectors = !!db.prepare(`SELECT name FROM sqlite_master WHERE type='table' AND name='vectors_vec'`).get();
  1412. for (const q of queries) {
  1413. // FTS search - get ranked results
  1414. const ftsResults = searchFTS(db, q, 20, collectionId);
  1415. if (ftsResults.length > 0) {
  1416. rankedLists.push(ftsResults.map(r => ({ file: r.file, displayPath: r.displayPath, title: r.title, body: r.body, score: r.score })));
  1417. }
  1418. // Vector search - get ranked results
  1419. if (hasVectors) {
  1420. const vecResults = await searchVec(db, q, embedModel, 20, collectionId);
  1421. if (vecResults.length > 0) {
  1422. rankedLists.push(vecResults.map(r => ({ file: r.file, displayPath: r.displayPath, title: r.title, body: r.body, score: r.score })));
  1423. }
  1424. }
  1425. }
  1426. // Apply Reciprocal Rank Fusion to combine all ranked lists
  1427. // Give 2x weight to original query results (first 2 lists: FTS + vector)
  1428. const weights = rankedLists.map((_, i) => i < 2 ? 2.0 : 1.0);
  1429. const fused = reciprocalRankFusion(rankedLists, weights);
  1430. const candidates = fused.slice(0, 30); // Over-retrieve for reranking
  1431. if (candidates.length === 0) {
  1432. console.log("No results found.");
  1433. closeDb();
  1434. return;
  1435. }
  1436. // Rerank with the original query (with caching)
  1437. const reranked = await rerank(
  1438. query,
  1439. candidates.map(c => ({ file: c.file, text: c.body })),
  1440. rerankModel,
  1441. db
  1442. );
  1443. // Blend RRF position score with reranker score using position-aware weights
  1444. // Top retrieval results get more protection from reranker disagreement
  1445. const candidateMap = new Map(candidates.map(c => [c.file, { displayPath: c.displayPath, title: c.title, body: c.body }]));
  1446. const rrfRankMap = new Map(candidates.map((c, i) => [c.file, i + 1])); // 1-indexed rank
  1447. const finalResults = reranked.map(r => {
  1448. const rrfRank = rrfRankMap.get(r.file) || 30;
  1449. // Position-aware blending: top retrieval results preserved more
  1450. // Rank 1-3: 75% RRF, 25% reranker (trust retrieval for exact matches)
  1451. // Rank 4-10: 60% RRF, 40% reranker
  1452. // Rank 11+: 40% RRF, 60% reranker (trust reranker for lower-ranked)
  1453. let rrfWeight: number;
  1454. if (rrfRank <= 3) {
  1455. rrfWeight = 0.75;
  1456. } else if (rrfRank <= 10) {
  1457. rrfWeight = 0.60;
  1458. } else {
  1459. rrfWeight = 0.40;
  1460. }
  1461. const rrfScore = 1 / rrfRank; // Position-based: 1, 0.5, 0.33...
  1462. const blendedScore = rrfWeight * rrfScore + (1 - rrfWeight) * r.score;
  1463. const candidate = candidateMap.get(r.file);
  1464. return {
  1465. file: r.file,
  1466. displayPath: candidate?.displayPath || "",
  1467. title: candidate?.title || "",
  1468. body: candidate?.body || "",
  1469. score: blendedScore,
  1470. context: getContextForFile(db, r.file),
  1471. };
  1472. }).sort((a, b) => b.score - a.score);
  1473. closeDb();
  1474. outputResults(finalResults, query, opts);
  1475. }
  1476. // Parse CLI arguments using util.parseArgs
  1477. function parseCLI() {
  1478. const { values, positionals } = parseArgs({
  1479. args: Bun.argv.slice(2), // Skip bun and script path
  1480. options: {
  1481. // Global options
  1482. index: { type: "string" },
  1483. help: { type: "boolean", short: "h" },
  1484. // Search options
  1485. n: { type: "string" },
  1486. "min-score": { type: "string" },
  1487. all: { type: "boolean" },
  1488. full: { type: "boolean" },
  1489. csv: { type: "boolean" },
  1490. md: { type: "boolean" },
  1491. xml: { type: "boolean" },
  1492. files: { type: "boolean" },
  1493. json: { type: "boolean" },
  1494. collection: { type: "string", short: "c" }, // Filter by collection
  1495. // Add options
  1496. drop: { type: "boolean" },
  1497. // Embed options
  1498. force: { type: "boolean", short: "f" },
  1499. // Get options
  1500. l: { type: "string" }, // max lines
  1501. from: { type: "string" }, // start line
  1502. "max-bytes": { type: "string" }, // max bytes for multi-get
  1503. },
  1504. allowPositionals: true,
  1505. strict: false, // Allow unknown options to pass through
  1506. });
  1507. // Set global index name in store
  1508. if (values.index) {
  1509. setCustomIndexName(values.index);
  1510. }
  1511. // Determine output format
  1512. let format: OutputFormat = "cli";
  1513. if (values.csv) format = "csv";
  1514. else if (values.md) format = "md";
  1515. else if (values.xml) format = "xml";
  1516. else if (values.files) format = "files";
  1517. else if (values.json) format = "json";
  1518. // Default limit: 20 for --files/--json, 5 otherwise
  1519. // --all means return all results (use very large limit)
  1520. const defaultLimit = (format === "files" || format === "json") ? 20 : 5;
  1521. const isAll = values.all || false;
  1522. const opts: OutputOptions = {
  1523. format,
  1524. full: values.full || false,
  1525. limit: isAll ? 100000 : (values.n ? parseInt(values.n, 10) || defaultLimit : defaultLimit),
  1526. minScore: values["min-score"] ? parseFloat(values["min-score"]) || 0 : 0,
  1527. all: isAll,
  1528. collection: values.collection as string | undefined,
  1529. };
  1530. return {
  1531. command: positionals[0] || "",
  1532. args: positionals.slice(1),
  1533. query: positionals.slice(1).join(" "),
  1534. opts,
  1535. values,
  1536. };
  1537. }
  1538. function showHelp(): void {
  1539. console.log("Usage:");
  1540. console.log(" qmd add [--drop] [glob] - Add/update collection from $PWD (default: **/*.md)");
  1541. console.log(" qmd add-context <path> <text> - Add context description for files under path");
  1542. console.log(" qmd get <file>[:line] [-l N] [--from N] - Get document (optionally from line, max N lines)");
  1543. console.log(" qmd multi-get <pattern> [-l N] [--max-bytes N] - Get multiple docs by glob or comma-separated list");
  1544. console.log(" qmd status - Show index status and collections");
  1545. console.log(" qmd update - Re-index all collections");
  1546. console.log(" qmd embed [-f] - Create vector embeddings (chunks ~6KB each)");
  1547. console.log(" qmd cleanup - Remove cache and orphaned data, vacuum DB");
  1548. console.log(" qmd search <query> - Full-text search (BM25)");
  1549. console.log(" qmd vsearch <query> - Vector similarity search");
  1550. console.log(" qmd query <query> - Combined search with query expansion + reranking");
  1551. console.log(" qmd mcp - Start MCP server (for AI agent integration)");
  1552. console.log("");
  1553. console.log("Global options:");
  1554. console.log(" --index <name> - Use custom index name (default: index)");
  1555. console.log("");
  1556. console.log("Search options:");
  1557. console.log(" -n <num> - Number of results (default: 5, or 20 for --files)");
  1558. console.log(" --all - Return all matches (use with --min-score to filter)");
  1559. console.log(" --min-score <num> - Minimum similarity score");
  1560. console.log(" --full - Output full document instead of snippet");
  1561. console.log(" --files - Output score,filepath,context (default: 20 results)");
  1562. console.log(" --json - JSON output with snippets (default: 20 results)");
  1563. console.log(" --csv - CSV output with snippets");
  1564. console.log(" --md - Markdown output");
  1565. console.log(" --xml - XML output");
  1566. console.log(" -c, --collection <name> - Filter results to a specific collection");
  1567. console.log("");
  1568. console.log("Multi-get options:");
  1569. console.log(" -l <num> - Maximum lines per file");
  1570. console.log(" --max-bytes <num> - Skip files larger than N bytes (default: 10240)");
  1571. console.log(" --json/--csv/--md/--xml/--files - Output format (same as search)");
  1572. console.log("");
  1573. console.log("Environment:");
  1574. console.log(" OLLAMA_URL - Ollama server URL (default: http://localhost:11434)");
  1575. console.log("");
  1576. console.log("Models:");
  1577. console.log(` Embedding: ${DEFAULT_EMBED_MODEL}`);
  1578. console.log(` Reranking: ${DEFAULT_RERANK_MODEL}`);
  1579. console.log("");
  1580. console.log(`Index: ${getDbPath()}`);
  1581. }
  1582. // Main CLI - only run if this is the main module
  1583. if (import.meta.main) {
  1584. const cli = parseCLI();
  1585. if (!cli.command || cli.values.help) {
  1586. showHelp();
  1587. process.exit(cli.values.help ? 0 : 1);
  1588. }
  1589. switch (cli.command) {
  1590. case "add": {
  1591. const globArg = cli.args[0];
  1592. // Treat "." as "use default glob in current directory"
  1593. const globPattern = (!globArg || globArg === ".") ? DEFAULT_GLOB : globArg;
  1594. if (cli.values.drop) {
  1595. await dropCollection(globPattern);
  1596. }
  1597. await indexFiles(globPattern);
  1598. break;
  1599. }
  1600. case "add-context": {
  1601. // qmd add-context <path> <context> OR qmd add-context <context> (uses .)
  1602. if (cli.args.length === 0) {
  1603. console.error("Usage: qmd add-context <path> <context>");
  1604. console.error(" qmd add-context . \"Description of files in current directory\"");
  1605. process.exit(1);
  1606. }
  1607. let pathArg: string;
  1608. let contextText: string;
  1609. if (cli.args.length === 1) {
  1610. // Single arg = context for current directory
  1611. pathArg = ".";
  1612. contextText = cli.args[0];
  1613. } else {
  1614. pathArg = cli.args[0];
  1615. contextText = cli.args.slice(1).join(" ");
  1616. }
  1617. await addContext(pathArg, contextText);
  1618. break;
  1619. }
  1620. case "get": {
  1621. if (!cli.args[0]) {
  1622. console.error("Usage: qmd get <filepath>[:line] [--from <line>] [-l <lines>]");
  1623. process.exit(1);
  1624. }
  1625. const fromLine = cli.values.from ? parseInt(cli.values.from as string, 10) : undefined;
  1626. const maxLines = cli.values.l ? parseInt(cli.values.l as string, 10) : undefined;
  1627. getDocument(cli.args[0], fromLine, maxLines);
  1628. break;
  1629. }
  1630. case "multi-get": {
  1631. if (!cli.args[0]) {
  1632. console.error("Usage: qmd multi-get <pattern> [-l <lines>] [--max-bytes <bytes>] [--json|--csv|--md|--xml|--files]");
  1633. console.error(" pattern: glob (e.g., 'journals/2025-05*.md') or comma-separated list");
  1634. process.exit(1);
  1635. }
  1636. const maxLinesMulti = cli.values.l ? parseInt(cli.values.l as string, 10) : undefined;
  1637. const maxBytes = cli.values["max-bytes"] ? parseInt(cli.values["max-bytes"] as string, 10) : DEFAULT_MULTI_GET_MAX_BYTES;
  1638. multiGet(cli.args[0], maxLinesMulti, maxBytes, cli.opts.format);
  1639. break;
  1640. }
  1641. case "status":
  1642. showStatus();
  1643. break;
  1644. case "update":
  1645. await updateCollections();
  1646. break;
  1647. case "embed":
  1648. await vectorIndex(DEFAULT_EMBED_MODEL, cli.values.force || false);
  1649. break;
  1650. case "search":
  1651. if (!cli.query) {
  1652. console.error("Usage: qmd search [options] <query>");
  1653. process.exit(1);
  1654. }
  1655. search(cli.query, cli.opts);
  1656. break;
  1657. case "vsearch":
  1658. if (!cli.query) {
  1659. console.error("Usage: qmd vsearch [options] <query>");
  1660. process.exit(1);
  1661. }
  1662. // Default min-score for vector search is 0.3
  1663. if (!cli.values["min-score"]) {
  1664. cli.opts.minScore = 0.3;
  1665. }
  1666. await vectorSearch(cli.query, cli.opts);
  1667. break;
  1668. case "query":
  1669. if (!cli.query) {
  1670. console.error("Usage: qmd query [options] <query>");
  1671. process.exit(1);
  1672. }
  1673. await querySearch(cli.query, cli.opts);
  1674. break;
  1675. case "mcp": {
  1676. const { startMcpServer } = await import("./mcp.js");
  1677. await startMcpServer();
  1678. break;
  1679. }
  1680. case "cleanup": {
  1681. const db = getDb();
  1682. // 1. Clear ollama_cache
  1683. const cacheCount = db.prepare(`SELECT COUNT(*) as c FROM ollama_cache`).get() as { c: number };
  1684. db.exec(`DELETE FROM ollama_cache`);
  1685. console.log(`${c.green}✓${c.reset} Cleared ${cacheCount.c} cached API responses`);
  1686. // 2. Remove orphaned vectors (no active document with that hash)
  1687. const orphanedVecs = db.prepare(`
  1688. SELECT COUNT(*) as c FROM content_vectors cv
  1689. WHERE NOT EXISTS (
  1690. SELECT 1 FROM documents d WHERE d.hash = cv.hash AND d.active = 1
  1691. )
  1692. `).get() as { c: number };
  1693. if (orphanedVecs.c > 0) {
  1694. db.exec(`
  1695. DELETE FROM vectors_vec WHERE hash_seq IN (
  1696. SELECT cv.hash || '_' || cv.seq FROM content_vectors cv
  1697. WHERE NOT EXISTS (
  1698. SELECT 1 FROM documents d WHERE d.hash = cv.hash AND d.active = 1
  1699. )
  1700. )
  1701. `);
  1702. db.exec(`
  1703. DELETE FROM content_vectors WHERE hash NOT IN (
  1704. SELECT hash FROM documents WHERE active = 1
  1705. )
  1706. `);
  1707. console.log(`${c.green}✓${c.reset} Removed ${orphanedVecs.c} orphaned embedding chunks`);
  1708. } else {
  1709. console.log(`${c.dim}No orphaned embeddings to remove${c.reset}`);
  1710. }
  1711. // 3. Count inactive documents
  1712. const inactiveDocs = db.prepare(`SELECT COUNT(*) as c FROM documents WHERE active = 0`).get() as { c: number };
  1713. if (inactiveDocs.c > 0) {
  1714. db.exec(`DELETE FROM documents WHERE active = 0`);
  1715. console.log(`${c.green}✓${c.reset} Removed ${inactiveDocs.c} inactive document records`);
  1716. }
  1717. // 4. Vacuum to reclaim space
  1718. db.exec(`VACUUM`);
  1719. console.log(`${c.green}✓${c.reset} Database vacuumed`);
  1720. closeDb();
  1721. break;
  1722. }
  1723. default:
  1724. console.error(`Unknown command: ${cli.command}`);
  1725. console.error("Run 'qmd --help' for usage.");
  1726. process.exit(1);
  1727. }
  1728. } // end if (import.meta.main)