store.ts 56 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273747576777879808182838485868788899091929394959697989910010110210310410510610710810911011111211311411511611711811912012112212312412512612712812913013113213313413513613713813914014114214314414514614714814915015115215315415515615715815916016116216316416516616716816917017117217317417517617717817918018118218318418518618718818919019119219319419519619719819920020120220320420520620720820921021121221321421521621721821922022122222322422522622722822923023123223323423523623723823924024124224324424524624724824925025125225325425525625725825926026126226326426526626726826927027127227327427527627727827928028128228328428528628728828929029129229329429529629729829930030130230330430530630730830931031131231331431531631731831932032132232332432532632732832933033133233333433533633733833934034134234334434534634734834935035135235335435535635735835936036136236336436536636736836937037137237337437537637737837938038138238338438538638738838939039139239339439539639739839940040140240340440540640740840941041141241341441541641741841942042142242342442542642742842943043143243343443543643743843944044144244344444544644744844945045145245345445545645745845946046146246346446546646746846947047147247347447547647747847948048148248348448548648748848949049149249349449549649749849950050150250350450550650750850951051151251351451551651751851952052152252352452552652752852953053153253353453553653753853954054154254354454554654754854955055155255355455555655755855956056156256356456556656756856957057157257357457557657757857958058158258358458558658758858959059159259359459559659759859960060160260360460560660760860961061161261361461561661761861962062162262362462562662762862963063163263363463563663763863964064164264364464564664764864965065165265365465565665765865966066166266366466566666766866967067167267367467567667767867968068168268368468568668768868969069169269369469569669769869970070170270370470570670770870971071171271371471571671771871972072172272372472572672772872973073173273373473573673773873974074174274374474574674774874975075175275375475575675775875976076176276376476576676776876977077177277377477577677777877978078178278378478578678778878979079179279379479579679779879980080180280380480580680780880981081181281381481581681781881982082182282382482582682782882983083183283383483583683783883984084184284384484584684784884985085185285385485585685785885986086186286386486586686786886987087187287387487587687787887988088188288388488588688788888989089189289389489589689789889990090190290390490590690790890991091191291391491591691791891992092192292392492592692792892993093193293393493593693793893994094194294394494594694794894995095195295395495595695795895996096196296396496596696796896997097197297397497597697797897998098198298398498598698798898999099199299399499599699799899910001001100210031004100510061007100810091010101110121013101410151016101710181019102010211022102310241025102610271028102910301031103210331034103510361037103810391040104110421043104410451046104710481049105010511052105310541055105610571058105910601061106210631064106510661067106810691070107110721073107410751076107710781079108010811082108310841085108610871088108910901091109210931094109510961097109810991100110111021103110411051106110711081109111011111112111311141115111611171118111911201121112211231124112511261127112811291130113111321133113411351136113711381139114011411142114311441145114611471148114911501151115211531154115511561157115811591160116111621163116411651166116711681169117011711172117311741175117611771178117911801181118211831184118511861187118811891190119111921193119411951196119711981199120012011202120312041205120612071208120912101211121212131214121512161217121812191220122112221223122412251226122712281229123012311232123312341235123612371238123912401241124212431244124512461247124812491250125112521253125412551256125712581259126012611262126312641265126612671268126912701271127212731274127512761277127812791280128112821283128412851286128712881289129012911292129312941295129612971298129913001301130213031304130513061307130813091310131113121313131413151316131713181319132013211322132313241325132613271328132913301331133213331334133513361337133813391340134113421343134413451346134713481349135013511352135313541355135613571358135913601361136213631364136513661367136813691370137113721373137413751376137713781379138013811382138313841385138613871388138913901391139213931394139513961397139813991400140114021403140414051406140714081409141014111412141314141415141614171418141914201421142214231424142514261427142814291430143114321433143414351436143714381439144014411442144314441445144614471448144914501451145214531454145514561457145814591460146114621463146414651466146714681469147014711472147314741475147614771478147914801481148214831484148514861487148814891490149114921493149414951496149714981499150015011502150315041505150615071508150915101511151215131514151515161517151815191520152115221523152415251526152715281529153015311532153315341535153615371538153915401541154215431544154515461547154815491550155115521553155415551556155715581559156015611562156315641565156615671568156915701571157215731574157515761577157815791580158115821583158415851586158715881589159015911592159315941595159615971598159916001601160216031604160516061607160816091610161116121613161416151616161716181619162016211622162316241625162616271628162916301631163216331634163516361637163816391640164116421643164416451646164716481649165016511652165316541655165616571658165916601661166216631664166516661667
  1. /**
  2. * QMD Store - Core data access and retrieval functions
  3. *
  4. * This module provides all database operations, search functions, and document
  5. * retrieval for QMD. It returns raw data structures that can be formatted by
  6. * CLI or MCP consumers.
  7. *
  8. * Usage:
  9. * const store = createStore("/path/to/db.sqlite");
  10. * // or use default path:
  11. * const store = createStore();
  12. */
  13. import { Database } from "bun:sqlite";
  14. import { Glob } from "bun";
  15. import * as sqliteVec from "sqlite-vec";
  16. import {
  17. Ollama,
  18. getDefaultOllama,
  19. formatQueryForEmbedding,
  20. formatDocForEmbedding,
  21. type RerankDocument,
  22. } from "./llm";
  23. // =============================================================================
  24. // Configuration
  25. // =============================================================================
  26. const HOME = Bun.env.HOME || "/tmp";
  27. export const DEFAULT_EMBED_MODEL = "embeddinggemma";
  28. export const DEFAULT_RERANK_MODEL = "ExpedientFalcon/qwen3-reranker:0.6b-q8_0";
  29. export const DEFAULT_QUERY_MODEL = "qwen3:0.6b";
  30. export const DEFAULT_GLOB = "**/*.md";
  31. export const DEFAULT_MULTI_GET_MAX_BYTES = 10 * 1024; // 10KB
  32. // Re-export OLLAMA_URL for backwards compatibility
  33. export const OLLAMA_URL = getDefaultOllama().getBaseUrl();
  34. // Chunking: ~2000 tokens per chunk, ~3 bytes/token = 6KB
  35. const CHUNK_BYTE_SIZE = 6 * 1024;
  36. // =============================================================================
  37. // Path utilities
  38. // =============================================================================
  39. export function homedir(): string {
  40. return HOME;
  41. }
  42. export function resolve(...paths: string[]): string {
  43. let result = paths[0].startsWith('/') ? '' : Bun.env.PWD || process.cwd();
  44. for (const p of paths) {
  45. if (p.startsWith('/')) {
  46. result = p;
  47. } else {
  48. result = result + '/' + p;
  49. }
  50. }
  51. const parts = result.split('/').filter(Boolean);
  52. const normalized: string[] = [];
  53. for (const part of parts) {
  54. if (part === '..') normalized.pop();
  55. else if (part !== '.') normalized.push(part);
  56. }
  57. return '/' + normalized.join('/');
  58. }
  59. export function getDefaultDbPath(indexName: string = "index"): string {
  60. // Allow override via INDEX_PATH for testing
  61. if (Bun.env.INDEX_PATH) {
  62. return Bun.env.INDEX_PATH;
  63. }
  64. const cacheDir = Bun.env.XDG_CACHE_HOME || resolve(homedir(), ".cache");
  65. const qmdCacheDir = resolve(cacheDir, "qmd");
  66. try { Bun.spawnSync(["mkdir", "-p", qmdCacheDir]); } catch {}
  67. return resolve(qmdCacheDir, `${indexName}.sqlite`);
  68. }
  69. export function getPwd(): string {
  70. return process.env.PWD || process.cwd();
  71. }
  72. export function getRealPath(path: string): string {
  73. try {
  74. const result = Bun.spawnSync(["realpath", path]);
  75. if (result.success) {
  76. return result.stdout.toString().trim();
  77. }
  78. } catch {}
  79. return resolve(path);
  80. }
  81. // =============================================================================
  82. // Virtual Path Utilities (qmd://)
  83. // =============================================================================
  84. export type VirtualPath = {
  85. collectionName: string;
  86. path: string; // relative path within collection
  87. };
  88. /**
  89. * Parse a virtual path like "qmd://collection-name/path/to/file.md"
  90. * into its components.
  91. */
  92. export function parseVirtualPath(virtualPath: string): VirtualPath | null {
  93. const match = virtualPath.match(/^qmd:\/\/([^\/]+)\/(.+)$/);
  94. if (!match) return null;
  95. return {
  96. collectionName: match[1],
  97. path: match[2],
  98. };
  99. }
  100. /**
  101. * Build a virtual path from collection name and relative path.
  102. */
  103. export function buildVirtualPath(collectionName: string, path: string): string {
  104. return `qmd://${collectionName}/${path}`;
  105. }
  106. /**
  107. * Check if a path is a virtual path (starts with qmd://).
  108. */
  109. export function isVirtualPath(path: string): boolean {
  110. return path.startsWith('qmd://');
  111. }
  112. /**
  113. * Resolve a virtual path to absolute filesystem path.
  114. */
  115. export function resolveVirtualPath(db: Database, virtualPath: string): string | null {
  116. const parsed = parseVirtualPath(virtualPath);
  117. if (!parsed) return null;
  118. const coll = getCollectionByName(db, parsed.collectionName);
  119. if (!coll) return null;
  120. return resolve(coll.pwd, parsed.path);
  121. }
  122. /**
  123. * Convert an absolute filesystem path to a virtual path.
  124. * Returns null if the file is not in any indexed collection.
  125. */
  126. export function toVirtualPath(db: Database, absolutePath: string): string | null {
  127. const doc = db.prepare(`
  128. SELECT c.name, d.path
  129. FROM documents d
  130. JOIN collections c ON c.id = d.collection_id
  131. WHERE c.pwd || '/' || d.path = ? AND d.active = 1
  132. LIMIT 1
  133. `).get(absolutePath) as { name: string; path: string } | null;
  134. if (!doc) return null;
  135. return buildVirtualPath(doc.name, doc.path);
  136. }
  137. // =============================================================================
  138. // Database initialization
  139. // =============================================================================
  140. // On macOS, use Homebrew's SQLite which supports extensions
  141. if (process.platform === "darwin") {
  142. const homebrewSqlitePath = "/opt/homebrew/opt/sqlite/lib/libsqlite3.dylib";
  143. try {
  144. if (Bun.file(homebrewSqlitePath).size > 0) {
  145. Database.setCustomSQLite(homebrewSqlitePath);
  146. }
  147. } catch {}
  148. }
  149. function initializeDatabase(db: Database): void {
  150. sqliteVec.load(db);
  151. db.exec("PRAGMA journal_mode = WAL");
  152. db.exec("PRAGMA foreign_keys = ON");
  153. // Check if we need to migrate from old schema
  154. const tables = db.prepare(`SELECT name FROM sqlite_master WHERE type='table'`).all() as { name: string }[];
  155. const tableNames = tables.map(t => t.name);
  156. const needsMigration = tableNames.includes('documents') && !tableNames.includes('content');
  157. if (needsMigration) {
  158. migrateToContentAddressable(db);
  159. return; // Migration will call initializeDatabase again
  160. }
  161. // Content-addressable storage - the source of truth for document content
  162. db.exec(`
  163. CREATE TABLE IF NOT EXISTS content (
  164. hash TEXT PRIMARY KEY,
  165. doc TEXT NOT NULL,
  166. created_at TEXT NOT NULL
  167. )
  168. `);
  169. // Collections table with name field
  170. db.exec(`
  171. CREATE TABLE IF NOT EXISTS collections (
  172. id INTEGER PRIMARY KEY AUTOINCREMENT,
  173. name TEXT NOT NULL UNIQUE,
  174. pwd TEXT NOT NULL,
  175. glob_pattern TEXT NOT NULL,
  176. created_at TEXT NOT NULL,
  177. updated_at TEXT NOT NULL,
  178. UNIQUE(pwd, glob_pattern)
  179. )
  180. `);
  181. // Documents table - file system layer mapping virtual paths to content hashes
  182. db.exec(`
  183. CREATE TABLE IF NOT EXISTS documents (
  184. id INTEGER PRIMARY KEY AUTOINCREMENT,
  185. collection_id INTEGER NOT NULL,
  186. path TEXT NOT NULL,
  187. title TEXT NOT NULL,
  188. hash TEXT NOT NULL,
  189. created_at TEXT NOT NULL,
  190. modified_at TEXT NOT NULL,
  191. active INTEGER NOT NULL DEFAULT 1,
  192. FOREIGN KEY (collection_id) REFERENCES collections(id) ON DELETE CASCADE,
  193. FOREIGN KEY (hash) REFERENCES content(hash) ON DELETE CASCADE,
  194. UNIQUE(collection_id, path)
  195. )
  196. `);
  197. db.exec(`CREATE INDEX IF NOT EXISTS idx_documents_collection ON documents(collection_id, active)`);
  198. db.exec(`CREATE INDEX IF NOT EXISTS idx_documents_hash ON documents(hash)`);
  199. db.exec(`CREATE INDEX IF NOT EXISTS idx_documents_path ON documents(path, active)`);
  200. // Path-based context (collection-scoped, hierarchical)
  201. db.exec(`
  202. CREATE TABLE IF NOT EXISTS path_contexts (
  203. id INTEGER PRIMARY KEY AUTOINCREMENT,
  204. collection_id INTEGER NOT NULL,
  205. path_prefix TEXT NOT NULL,
  206. context TEXT NOT NULL,
  207. created_at TEXT NOT NULL,
  208. FOREIGN KEY (collection_id) REFERENCES collections(id) ON DELETE CASCADE,
  209. UNIQUE(collection_id, path_prefix)
  210. )
  211. `);
  212. db.exec(`CREATE INDEX IF NOT EXISTS idx_path_contexts_collection ON path_contexts(collection_id, path_prefix)`);
  213. // Cache table for Ollama API calls
  214. db.exec(`
  215. CREATE TABLE IF NOT EXISTS ollama_cache (
  216. hash TEXT PRIMARY KEY,
  217. result TEXT NOT NULL,
  218. created_at TEXT NOT NULL
  219. )
  220. `);
  221. // Content vectors
  222. const cvInfo = db.prepare(`PRAGMA table_info(content_vectors)`).all() as { name: string }[];
  223. const hasSeqColumn = cvInfo.some(col => col.name === 'seq');
  224. if (cvInfo.length > 0 && !hasSeqColumn) {
  225. db.exec(`DROP TABLE IF EXISTS content_vectors`);
  226. db.exec(`DROP TABLE IF EXISTS vectors_vec`);
  227. }
  228. db.exec(`
  229. CREATE TABLE IF NOT EXISTS content_vectors (
  230. hash TEXT NOT NULL,
  231. seq INTEGER NOT NULL DEFAULT 0,
  232. pos INTEGER NOT NULL DEFAULT 0,
  233. model TEXT NOT NULL,
  234. embedded_at TEXT NOT NULL,
  235. PRIMARY KEY (hash, seq)
  236. )
  237. `);
  238. // FTS - index path and content (joined from content table)
  239. db.exec(`
  240. CREATE VIRTUAL TABLE IF NOT EXISTS documents_fts USING fts5(
  241. path, body,
  242. tokenize='porter unicode61'
  243. )
  244. `);
  245. // Triggers to keep FTS in sync
  246. db.exec(`
  247. CREATE TRIGGER IF NOT EXISTS documents_ai AFTER INSERT ON documents BEGIN
  248. INSERT INTO documents_fts(rowid, path, body)
  249. SELECT new.id, new.path, c.doc
  250. FROM content c
  251. WHERE c.hash = new.hash;
  252. END
  253. `);
  254. db.exec(`
  255. CREATE TRIGGER IF NOT EXISTS documents_ad AFTER DELETE ON documents BEGIN
  256. DELETE FROM documents_fts WHERE rowid = old.id;
  257. END
  258. `);
  259. db.exec(`
  260. CREATE TRIGGER IF NOT EXISTS documents_au AFTER UPDATE ON documents BEGIN
  261. UPDATE documents_fts
  262. SET path = new.path,
  263. body = (SELECT doc FROM content WHERE hash = new.hash)
  264. WHERE rowid = new.id;
  265. END
  266. `);
  267. }
  268. function migrateToContentAddressable(db: Database): void {
  269. console.log("Migrating database to content-addressable schema...");
  270. // Start transaction
  271. db.exec("BEGIN TRANSACTION");
  272. try {
  273. // Rename old tables
  274. db.exec("ALTER TABLE documents RENAME TO documents_old");
  275. db.exec("ALTER TABLE collections RENAME TO collections_old");
  276. db.exec("ALTER TABLE path_contexts RENAME TO path_contexts_old");
  277. db.exec("DROP TABLE IF EXISTS documents_fts");
  278. db.exec("DROP TRIGGER IF EXISTS documents_ai");
  279. db.exec("DROP TRIGGER IF EXISTS documents_ad");
  280. db.exec("DROP TRIGGER IF EXISTS documents_au");
  281. // Create new schema
  282. db.exec(`
  283. CREATE TABLE content (
  284. hash TEXT PRIMARY KEY,
  285. doc TEXT NOT NULL,
  286. created_at TEXT NOT NULL
  287. )
  288. `);
  289. db.exec(`
  290. CREATE TABLE collections (
  291. id INTEGER PRIMARY KEY AUTOINCREMENT,
  292. name TEXT NOT NULL UNIQUE,
  293. pwd TEXT NOT NULL,
  294. glob_pattern TEXT NOT NULL,
  295. created_at TEXT NOT NULL,
  296. updated_at TEXT NOT NULL,
  297. UNIQUE(pwd, glob_pattern)
  298. )
  299. `);
  300. db.exec(`
  301. CREATE TABLE documents (
  302. id INTEGER PRIMARY KEY AUTOINCREMENT,
  303. collection_id INTEGER NOT NULL,
  304. path TEXT NOT NULL,
  305. title TEXT NOT NULL,
  306. hash TEXT NOT NULL,
  307. created_at TEXT NOT NULL,
  308. modified_at TEXT NOT NULL,
  309. active INTEGER NOT NULL DEFAULT 1,
  310. FOREIGN KEY (collection_id) REFERENCES collections(id) ON DELETE CASCADE,
  311. FOREIGN KEY (hash) REFERENCES content(hash) ON DELETE CASCADE,
  312. UNIQUE(collection_id, path)
  313. )
  314. `);
  315. db.exec(`
  316. CREATE TABLE path_contexts (
  317. id INTEGER PRIMARY KEY AUTOINCREMENT,
  318. collection_id INTEGER NOT NULL,
  319. path_prefix TEXT NOT NULL,
  320. context TEXT NOT NULL,
  321. created_at TEXT NOT NULL,
  322. FOREIGN KEY (collection_id) REFERENCES collections(id) ON DELETE CASCADE,
  323. UNIQUE(collection_id, path_prefix)
  324. )
  325. `);
  326. // Migrate data: Extract unique content hashes
  327. console.log("Migrating content...");
  328. db.exec(`
  329. INSERT INTO content (hash, doc, created_at)
  330. SELECT hash, body, MIN(created_at) as created_at
  331. FROM documents_old
  332. WHERE active = 1
  333. GROUP BY hash
  334. `);
  335. // Migrate collections: generate names from pwd basename
  336. console.log("Migrating collections...");
  337. // First insert with pwd as temporary name
  338. db.exec(`
  339. INSERT INTO collections (id, name, pwd, glob_pattern, created_at, updated_at)
  340. SELECT
  341. id,
  342. pwd as name,
  343. pwd,
  344. glob_pattern,
  345. created_at,
  346. created_at as updated_at
  347. FROM collections_old
  348. `);
  349. // Then update names to basenames using application logic
  350. const collections = db.prepare(`SELECT id, pwd FROM collections`).all() as { id: number; pwd: string }[];
  351. for (const coll of collections) {
  352. const parts = coll.pwd.split('/').filter(Boolean);
  353. const name = parts[parts.length - 1] || 'root';
  354. db.prepare(`UPDATE collections SET name = ? WHERE id = ?`).run(name, coll.id);
  355. }
  356. // Handle duplicate collection names by appending collection_id
  357. const duplicates = db.prepare(`
  358. SELECT name, COUNT(*) as cnt
  359. FROM collections
  360. GROUP BY name
  361. HAVING cnt > 1
  362. `).all() as { name: string; cnt: number }[];
  363. for (const dup of duplicates) {
  364. const rows = db.prepare(`SELECT id FROM collections WHERE name = ? ORDER BY id`).all(dup.name) as { id: number }[];
  365. for (let i = 1; i < rows.length; i++) {
  366. db.prepare(`UPDATE collections SET name = ? WHERE id = ?`).run(`${dup.name}-${rows[i].id}`, rows[i].id);
  367. }
  368. }
  369. // Migrate documents: convert filepath to relative path within collection
  370. console.log("Migrating documents...");
  371. const oldDocs = db.prepare(`
  372. SELECT d.id, d.collection_id, d.filepath, d.title, d.hash, d.created_at, d.modified_at, c.pwd
  373. FROM documents_old d
  374. JOIN collections c ON c.id = d.collection_id
  375. WHERE d.active = 1
  376. `).all() as Array<{
  377. id: number;
  378. collection_id: number;
  379. filepath: string;
  380. title: string;
  381. hash: string;
  382. created_at: string;
  383. modified_at: string;
  384. pwd: string;
  385. }>;
  386. const insertDoc = db.prepare(`
  387. INSERT INTO documents (collection_id, path, title, hash, created_at, modified_at, active)
  388. VALUES (?, ?, ?, ?, ?, ?, 1)
  389. `);
  390. for (const doc of oldDocs) {
  391. // Convert absolute filepath to relative path within collection
  392. let path = doc.filepath;
  393. if (path.startsWith(doc.pwd + '/')) {
  394. path = path.slice(doc.pwd.length + 1);
  395. } else if (path.startsWith(doc.pwd)) {
  396. path = path.slice(doc.pwd.length);
  397. }
  398. // Remove leading slash if present
  399. path = path.replace(/^\/+/, '');
  400. try {
  401. insertDoc.run(doc.collection_id, path, doc.title, doc.hash, doc.created_at, doc.modified_at);
  402. } catch (e) {
  403. console.warn(`Skipping duplicate path: ${path} in collection ${doc.collection_id}`);
  404. }
  405. }
  406. // Migrate path_contexts: associate with collections based on path prefix
  407. console.log("Migrating path contexts...");
  408. const oldContexts = db.prepare(`SELECT * FROM path_contexts_old`).all() as Array<{
  409. path_prefix: string;
  410. context: string;
  411. created_at: string;
  412. }>;
  413. const insertContext = db.prepare(`
  414. INSERT INTO path_contexts (collection_id, path_prefix, context, created_at)
  415. VALUES (?, ?, ?, ?)
  416. `);
  417. const allCollections = db.prepare(`SELECT id, pwd FROM collections`).all() as Array<{ id: number; pwd: string }>;
  418. for (const ctx of oldContexts) {
  419. // Find collection(s) that match this path prefix
  420. for (const coll of allCollections) {
  421. if (ctx.path_prefix.startsWith(coll.pwd)) {
  422. // Convert absolute path_prefix to relative within collection
  423. let relPath = ctx.path_prefix;
  424. if (relPath.startsWith(coll.pwd + '/')) {
  425. relPath = relPath.slice(coll.pwd.length + 1);
  426. } else if (relPath.startsWith(coll.pwd)) {
  427. relPath = relPath.slice(coll.pwd.length);
  428. }
  429. relPath = relPath.replace(/^\/+/, '');
  430. try {
  431. insertContext.run(coll.id, relPath, ctx.context, ctx.created_at);
  432. } catch (e) {
  433. // Ignore duplicates
  434. }
  435. }
  436. }
  437. }
  438. // Drop old tables
  439. db.exec("DROP TABLE documents_old");
  440. db.exec("DROP TABLE collections_old");
  441. db.exec("DROP TABLE path_contexts_old");
  442. // Recreate FTS and triggers
  443. db.exec(`
  444. CREATE VIRTUAL TABLE documents_fts USING fts5(
  445. path, body,
  446. tokenize='porter unicode61'
  447. )
  448. `);
  449. db.exec(`
  450. CREATE TRIGGER documents_ai AFTER INSERT ON documents BEGIN
  451. INSERT INTO documents_fts(rowid, path, body)
  452. SELECT new.id, new.path, c.doc
  453. FROM content c
  454. WHERE c.hash = new.hash;
  455. END
  456. `);
  457. db.exec(`
  458. CREATE TRIGGER documents_ad AFTER DELETE ON documents BEGIN
  459. DELETE FROM documents_fts WHERE rowid = old.id;
  460. END
  461. `);
  462. db.exec(`
  463. CREATE TRIGGER documents_au AFTER UPDATE ON documents BEGIN
  464. UPDATE documents_fts
  465. SET path = new.path,
  466. body = (SELECT doc FROM content WHERE hash = new.hash)
  467. WHERE rowid = new.id;
  468. END
  469. `);
  470. // Populate FTS from migrated data
  471. console.log("Rebuilding full-text search index...");
  472. db.exec(`
  473. INSERT INTO documents_fts(rowid, path, body)
  474. SELECT d.id, d.path, c.doc
  475. FROM documents d
  476. JOIN content c ON c.hash = d.hash
  477. WHERE d.active = 1
  478. `);
  479. // Create indexes
  480. db.exec(`CREATE INDEX idx_documents_collection ON documents(collection_id, active)`);
  481. db.exec(`CREATE INDEX idx_documents_hash ON documents(hash)`);
  482. db.exec(`CREATE INDEX idx_documents_path ON documents(path, active)`);
  483. db.exec(`CREATE INDEX idx_path_contexts_collection ON path_contexts(collection_id, path_prefix)`);
  484. db.exec("COMMIT");
  485. console.log("Migration complete!");
  486. } catch (e) {
  487. db.exec("ROLLBACK");
  488. console.error("Migration failed:", e);
  489. throw e;
  490. }
  491. }
  492. function ensureVecTableInternal(db: Database, dimensions: number): void {
  493. const tableInfo = db.prepare(`SELECT sql FROM sqlite_master WHERE type='table' AND name='vectors_vec'`).get() as { sql: string } | null;
  494. if (tableInfo) {
  495. const match = tableInfo.sql.match(/float\[(\d+)\]/);
  496. const hasHashSeq = tableInfo.sql.includes('hash_seq');
  497. if (match && parseInt(match[1]) === dimensions && hasHashSeq) return;
  498. db.exec("DROP TABLE IF EXISTS vectors_vec");
  499. }
  500. db.exec(`CREATE VIRTUAL TABLE vectors_vec USING vec0(hash_seq TEXT PRIMARY KEY, embedding float[${dimensions}])`);
  501. }
  502. // =============================================================================
  503. // Store Factory
  504. // =============================================================================
  505. export type Store = {
  506. db: Database;
  507. dbPath: string;
  508. close: () => void;
  509. ensureVecTable: (dimensions: number) => void;
  510. // Index health
  511. getHashesNeedingEmbedding: () => number;
  512. getIndexHealth: () => IndexHealthInfo;
  513. getStatus: () => IndexStatus;
  514. // Caching
  515. getCacheKey: typeof getCacheKey;
  516. getCachedResult: (cacheKey: string) => string | null;
  517. setCachedResult: (cacheKey: string, result: string) => void;
  518. clearCache: () => void;
  519. // Context
  520. getContextForFile: (filepath: string) => string | null;
  521. getContextForPath: (collectionId: number, path: string) => string | null;
  522. getCollectionIdByName: (name: string) => number | null;
  523. getCollectionByName: (name: string) => { id: number; name: string; pwd: string; glob_pattern: string } | null;
  524. // Virtual paths
  525. parseVirtualPath: typeof parseVirtualPath;
  526. buildVirtualPath: typeof buildVirtualPath;
  527. isVirtualPath: typeof isVirtualPath;
  528. resolveVirtualPath: (virtualPath: string) => string | null;
  529. toVirtualPath: (absolutePath: string) => string | null;
  530. // Search
  531. searchFTS: (query: string, limit?: number, collectionId?: number) => SearchResult[];
  532. searchVec: (query: string, model: string, limit?: number, collectionId?: number) => Promise<SearchResult[]>;
  533. // Query expansion & reranking
  534. expandQuery: (query: string, model?: string) => Promise<string[]>;
  535. rerank: (query: string, documents: { file: string; text: string }[], model?: string) => Promise<{ file: string; score: number }[]>;
  536. // Document retrieval
  537. findDocument: (filename: string, options?: { includeBody?: boolean }) => DocumentResult | DocumentNotFound;
  538. getDocumentBody: (doc: DocumentResult | { filepath: string }, fromLine?: number, maxLines?: number) => string | null;
  539. findDocuments: (pattern: string, options?: { includeBody?: boolean; maxBytes?: number }) => { docs: MultiGetResult[]; errors: string[] };
  540. // Legacy compatibility
  541. getDocument: (filename: string, fromLine?: number, maxLines?: number) => (DocumentResult & { body: string }) | DocumentNotFound;
  542. getMultipleDocuments: (pattern: string, maxLines?: number, maxBytes?: number) => { files: MultiGetFile[]; errors: string[] };
  543. // Fuzzy matching
  544. findSimilarFiles: (query: string, maxDistance?: number, limit?: number) => string[];
  545. matchFilesByGlob: (pattern: string) => { filepath: string; displayPath: string; bodyLength: number }[];
  546. };
  547. /**
  548. * Create a new store instance with the given database path.
  549. * If no path is provided, uses the default path (~/.cache/qmd/index.sqlite).
  550. *
  551. * @param dbPath - Path to the SQLite database file
  552. * @returns Store instance with all methods bound to the database
  553. */
  554. export function createStore(dbPath?: string): Store {
  555. const resolvedPath = dbPath || getDefaultDbPath();
  556. const db = new Database(resolvedPath);
  557. initializeDatabase(db);
  558. return {
  559. db,
  560. dbPath: resolvedPath,
  561. close: () => db.close(),
  562. ensureVecTable: (dimensions: number) => ensureVecTableInternal(db, dimensions),
  563. // Index health
  564. getHashesNeedingEmbedding: () => getHashesNeedingEmbedding(db),
  565. getIndexHealth: () => getIndexHealth(db),
  566. getStatus: () => getStatus(db),
  567. // Caching
  568. getCacheKey,
  569. getCachedResult: (cacheKey: string) => getCachedResult(db, cacheKey),
  570. setCachedResult: (cacheKey: string, result: string) => setCachedResult(db, cacheKey, result),
  571. clearCache: () => clearCache(db),
  572. // Context
  573. getContextForFile: (filepath: string) => getContextForFile(db, filepath),
  574. getContextForPath: (collectionId: number, path: string) => getContextForPath(db, collectionId, path),
  575. getCollectionIdByName: (name: string) => getCollectionIdByName(db, name),
  576. getCollectionByName: (name: string) => getCollectionByName(db, name),
  577. // Virtual paths
  578. parseVirtualPath,
  579. buildVirtualPath,
  580. isVirtualPath,
  581. resolveVirtualPath: (virtualPath: string) => resolveVirtualPath(db, virtualPath),
  582. toVirtualPath: (absolutePath: string) => toVirtualPath(db, absolutePath),
  583. // Search
  584. searchFTS: (query: string, limit?: number, collectionId?: number) => searchFTS(db, query, limit, collectionId),
  585. searchVec: (query: string, model: string, limit?: number, collectionId?: number) => searchVec(db, query, model, limit, collectionId),
  586. // Query expansion & reranking
  587. expandQuery: (query: string, model?: string) => expandQuery(query, model, db),
  588. rerank: (query: string, documents: { file: string; text: string }[], model?: string) => rerank(query, documents, model, db),
  589. // Document retrieval
  590. findDocument: (filename: string, options?: { includeBody?: boolean }) => findDocument(db, filename, options),
  591. getDocumentBody: (doc: DocumentResult | { filepath: string }, fromLine?: number, maxLines?: number) => getDocumentBody(db, doc, fromLine, maxLines),
  592. findDocuments: (pattern: string, options?: { includeBody?: boolean; maxBytes?: number }) => findDocuments(db, pattern, options),
  593. // Legacy compatibility
  594. getDocument: (filename: string, fromLine?: number, maxLines?: number) => getDocument(db, filename, fromLine, maxLines),
  595. getMultipleDocuments: (pattern: string, maxLines?: number, maxBytes?: number) => getMultipleDocuments(db, pattern, maxLines, maxBytes),
  596. // Fuzzy matching
  597. findSimilarFiles: (query: string, maxDistance?: number, limit?: number) => findSimilarFiles(db, query, maxDistance, limit),
  598. matchFilesByGlob: (pattern: string) => matchFilesByGlob(db, pattern),
  599. };
  600. }
  601. // =============================================================================
  602. // Legacy compatibility - will be removed
  603. // =============================================================================
  604. let _legacyDb: Database | null = null;
  605. let _legacyDbPath: string | null = null;
  606. /** @deprecated Use createStore() instead */
  607. export function setCustomIndexName(name: string | null): void {
  608. _legacyDbPath = name ? getDefaultDbPath(name) : null;
  609. _legacyDb = null; // Reset so next getDb() creates new connection
  610. }
  611. /** @deprecated Use createStore() instead */
  612. export function getDbPath(): string {
  613. return _legacyDbPath || getDefaultDbPath();
  614. }
  615. /** @deprecated Use createStore() instead */
  616. export function getDb(): Database {
  617. if (!_legacyDb) {
  618. _legacyDb = new Database(getDbPath());
  619. initializeDatabase(_legacyDb);
  620. }
  621. return _legacyDb;
  622. }
  623. /** @deprecated Use store.db.close() instead. Closes the legacy db and resets singleton. */
  624. export function closeDb(): void {
  625. if (_legacyDb) {
  626. _legacyDb.close();
  627. _legacyDb = null;
  628. }
  629. }
  630. /** @deprecated Use store.ensureVecTable() instead */
  631. export function ensureVecTable(db: Database, dimensions: number): void {
  632. ensureVecTableInternal(db, dimensions);
  633. }
  634. // =============================================================================
  635. // Core Document Type
  636. // =============================================================================
  637. /**
  638. * Unified document result type with all metadata.
  639. * Body is optional - use getDocumentBody() to load it separately if needed.
  640. */
  641. export type DocumentResult = {
  642. filepath: string; // Full filesystem path
  643. displayPath: string; // Short display path (e.g., "docs/readme.md")
  644. title: string; // Document title (from first heading or filename)
  645. context: string | null; // Folder context description if configured
  646. hash: string; // Content hash for caching/change detection
  647. collectionId: number; // Parent collection ID
  648. modifiedAt: string; // Last modification timestamp
  649. bodyLength: number; // Body length in bytes (useful before loading)
  650. body?: string; // Document body (optional, load with getDocumentBody)
  651. };
  652. /**
  653. * Search result extends DocumentResult with score and source info
  654. */
  655. export type SearchResult = DocumentResult & {
  656. score: number; // Relevance score (0-1)
  657. source: "fts" | "vec"; // Search source (full-text or vector)
  658. chunkPos?: number; // Character position of matching chunk (for vector search)
  659. };
  660. /**
  661. * Ranked result for RRF fusion (simplified, used internally)
  662. */
  663. export type RankedResult = {
  664. file: string;
  665. displayPath: string;
  666. title: string;
  667. body: string;
  668. score: number;
  669. };
  670. /**
  671. * Error result when document is not found
  672. */
  673. export type DocumentNotFound = {
  674. error: "not_found";
  675. query: string;
  676. similarFiles: string[];
  677. };
  678. /**
  679. * Result from multi-get operations
  680. */
  681. export type MultiGetResult = {
  682. doc: DocumentResult;
  683. skipped: false;
  684. } | {
  685. doc: Pick<DocumentResult, "filepath" | "displayPath">;
  686. skipped: true;
  687. skipReason: string;
  688. };
  689. export type CollectionInfo = {
  690. id: number;
  691. path: string;
  692. pattern: string;
  693. documents: number;
  694. lastUpdated: string;
  695. };
  696. export type IndexStatus = {
  697. totalDocuments: number;
  698. needsEmbedding: number;
  699. hasVectorIndex: boolean;
  700. collections: CollectionInfo[];
  701. };
  702. // =============================================================================
  703. // Index health
  704. // =============================================================================
  705. export function getHashesNeedingEmbedding(db: Database): number {
  706. const result = db.prepare(`
  707. SELECT COUNT(DISTINCT d.hash) as count
  708. FROM documents d
  709. LEFT JOIN content_vectors v ON d.hash = v.hash AND v.seq = 0
  710. WHERE d.active = 1 AND v.hash IS NULL
  711. `).get() as { count: number };
  712. return result.count;
  713. }
  714. export type IndexHealthInfo = {
  715. needsEmbedding: number;
  716. totalDocs: number;
  717. daysStale: number | null;
  718. };
  719. export function getIndexHealth(db: Database): IndexHealthInfo {
  720. const needsEmbedding = getHashesNeedingEmbedding(db);
  721. const totalDocs = (db.prepare(`SELECT COUNT(*) as count FROM documents WHERE active = 1`).get() as { count: number }).count;
  722. const mostRecent = db.prepare(`SELECT MAX(modified_at) as latest FROM documents WHERE active = 1`).get() as { latest: string | null };
  723. let daysStale: number | null = null;
  724. if (mostRecent?.latest) {
  725. const lastUpdate = new Date(mostRecent.latest);
  726. daysStale = Math.floor((Date.now() - lastUpdate.getTime()) / (24 * 60 * 60 * 1000));
  727. }
  728. return { needsEmbedding, totalDocs, daysStale };
  729. }
  730. // =============================================================================
  731. // Caching
  732. // =============================================================================
  733. export function getCacheKey(url: string, body: object): string {
  734. const hash = new Bun.CryptoHasher("sha256");
  735. hash.update(url);
  736. hash.update(JSON.stringify(body));
  737. return hash.digest("hex");
  738. }
  739. export function getCachedResult(db: Database, cacheKey: string): string | null {
  740. const row = db.prepare(`SELECT result FROM ollama_cache WHERE hash = ?`).get(cacheKey) as { result: string } | null;
  741. return row?.result || null;
  742. }
  743. export function setCachedResult(db: Database, cacheKey: string, result: string): void {
  744. const now = new Date().toISOString();
  745. db.prepare(`INSERT OR REPLACE INTO ollama_cache (hash, result, created_at) VALUES (?, ?, ?)`).run(cacheKey, result, now);
  746. if (Math.random() < 0.01) {
  747. db.exec(`DELETE FROM ollama_cache WHERE hash NOT IN (SELECT hash FROM ollama_cache ORDER BY created_at DESC LIMIT 1000)`);
  748. }
  749. }
  750. export function clearCache(db: Database): void {
  751. db.exec(`DELETE FROM ollama_cache`);
  752. }
  753. // =============================================================================
  754. // Document helpers
  755. // =============================================================================
  756. export async function hashContent(content: string): Promise<string> {
  757. const hash = new Bun.CryptoHasher("sha256");
  758. hash.update(content);
  759. return hash.digest("hex");
  760. }
  761. export function extractTitle(content: string, filename: string): string {
  762. const match = content.match(/^##?\s+(.+)$/m);
  763. if (match) {
  764. const title = match[1].trim();
  765. if (title === "📝 Notes" || title === "Notes") {
  766. const nextMatch = content.match(/^##\s+(.+)$/m);
  767. if (nextMatch) return nextMatch[1].trim();
  768. }
  769. return title;
  770. }
  771. return filename.replace(/\.md$/, "").split("/").pop() || filename;
  772. }
  773. // Re-export from llm.ts for backwards compatibility
  774. export { formatQueryForEmbedding, formatDocForEmbedding };
  775. export function chunkDocument(content: string, maxBytes: number = CHUNK_BYTE_SIZE): { text: string; pos: number }[] {
  776. const encoder = new TextEncoder();
  777. const totalBytes = encoder.encode(content).length;
  778. if (totalBytes <= maxBytes) {
  779. return [{ text: content, pos: 0 }];
  780. }
  781. const chunks: { text: string; pos: number }[] = [];
  782. let charPos = 0;
  783. while (charPos < content.length) {
  784. let endPos = charPos;
  785. let byteCount = 0;
  786. while (endPos < content.length && byteCount < maxBytes) {
  787. const charBytes = encoder.encode(content[endPos]).length;
  788. if (byteCount + charBytes > maxBytes) break;
  789. byteCount += charBytes;
  790. endPos++;
  791. }
  792. if (endPos < content.length && endPos > charPos) {
  793. const slice = content.slice(charPos, endPos);
  794. const paragraphBreak = slice.lastIndexOf('\n\n');
  795. const sentenceEnd = Math.max(
  796. slice.lastIndexOf('. '),
  797. slice.lastIndexOf('.\n'),
  798. slice.lastIndexOf('? '),
  799. slice.lastIndexOf('?\n'),
  800. slice.lastIndexOf('! '),
  801. slice.lastIndexOf('!\n')
  802. );
  803. const lineBreak = slice.lastIndexOf('\n');
  804. const spaceBreak = slice.lastIndexOf(' ');
  805. let breakPoint = -1;
  806. if (paragraphBreak > slice.length * 0.5) {
  807. breakPoint = paragraphBreak + 2;
  808. } else if (sentenceEnd > slice.length * 0.5) {
  809. breakPoint = sentenceEnd + 2;
  810. } else if (lineBreak > slice.length * 0.3) {
  811. breakPoint = lineBreak + 1;
  812. } else if (spaceBreak > slice.length * 0.3) {
  813. breakPoint = spaceBreak + 1;
  814. }
  815. if (breakPoint > 0) {
  816. endPos = charPos + breakPoint;
  817. }
  818. }
  819. if (endPos <= charPos) {
  820. endPos = charPos + 1;
  821. }
  822. chunks.push({ text: content.slice(charPos, endPos), pos: charPos });
  823. charPos = endPos;
  824. }
  825. return chunks;
  826. }
  827. // =============================================================================
  828. // Fuzzy matching
  829. // =============================================================================
  830. function levenshtein(a: string, b: string): number {
  831. const m = a.length, n = b.length;
  832. if (m === 0) return n;
  833. if (n === 0) return m;
  834. const dp: number[][] = Array.from({ length: m + 1 }, (_, i) => [i]);
  835. for (let j = 1; j <= n; j++) dp[0][j] = j;
  836. for (let i = 1; i <= m; i++) {
  837. for (let j = 1; j <= n; j++) {
  838. const cost = a[i - 1] === b[j - 1] ? 0 : 1;
  839. dp[i][j] = Math.min(dp[i - 1][j] + 1, dp[i][j - 1] + 1, dp[i - 1][j - 1] + cost);
  840. }
  841. }
  842. return dp[m][n];
  843. }
  844. export function findSimilarFiles(db: Database, query: string, maxDistance: number = 3, limit: number = 5): string[] {
  845. const allFiles = db.prepare(`SELECT display_path FROM documents WHERE active = 1`).all() as { display_path: string }[];
  846. const queryLower = query.toLowerCase();
  847. const scored = allFiles
  848. .map(f => ({ path: f.display_path, dist: levenshtein(f.display_path.toLowerCase(), queryLower) }))
  849. .filter(f => f.dist <= maxDistance)
  850. .sort((a, b) => a.dist - b.dist)
  851. .slice(0, limit);
  852. return scored.map(f => f.path);
  853. }
  854. export function matchFilesByGlob(db: Database, pattern: string): { filepath: string; displayPath: string; bodyLength: number }[] {
  855. const allFiles = db.prepare(`
  856. SELECT
  857. 'qmd://' || c.name || '/' || d.path as virtual_path,
  858. LENGTH(content.doc) as body_length,
  859. d.collection_id,
  860. d.path
  861. FROM documents d
  862. JOIN collections c ON c.id = d.collection_id
  863. JOIN content ON content.hash = d.hash
  864. WHERE d.active = 1
  865. `).all() as { virtual_path: string; body_length: number; collection_id: number; path: string }[];
  866. const glob = new Glob(pattern);
  867. return allFiles
  868. .filter(f => glob.match(f.virtual_path) || glob.match(f.path))
  869. .map(f => ({
  870. filepath: f.virtual_path, // Use virtual path as filepath
  871. displayPath: f.virtual_path,
  872. bodyLength: f.body_length
  873. }));
  874. }
  875. // =============================================================================
  876. // Context
  877. // =============================================================================
  878. /**
  879. * Get context for a file path using hierarchical inheritance.
  880. * Contexts are collection-scoped and inherit from parent directories.
  881. * For example, context at "/talks" applies to "/talks/2024/keynote.md".
  882. *
  883. * @param db Database instance
  884. * @param collectionId Collection ID
  885. * @param path Relative path within the collection
  886. * @returns Context string or null if no context is defined
  887. */
  888. export function getContextForPath(db: Database, collectionId: number, path: string): string | null {
  889. // Find the most specific (longest) matching path prefix for this collection
  890. const result = db.prepare(`
  891. SELECT context FROM path_contexts
  892. WHERE collection_id = ?
  893. AND (? LIKE path_prefix || '/%' OR ? = path_prefix OR path_prefix = '')
  894. ORDER BY LENGTH(path_prefix) DESC
  895. LIMIT 1
  896. `).get(collectionId, path, path) as { context: string } | null;
  897. return result?.context || null;
  898. }
  899. /**
  900. * Legacy function for backward compatibility - resolves filepath to collection+path first
  901. */
  902. export function getContextForFile(db: Database, filepath: string): string | null {
  903. // Try to find the document to get its collection_id and path
  904. const doc = db.prepare(`
  905. SELECT d.collection_id, d.path
  906. FROM documents d
  907. JOIN collections c ON c.id = d.collection_id
  908. WHERE c.pwd || '/' || d.path = ? AND d.active = 1
  909. LIMIT 1
  910. `).get(filepath) as { collection_id: number; path: string } | null;
  911. if (!doc) return null;
  912. return getContextForPath(db, doc.collection_id, doc.path);
  913. }
  914. /**
  915. * Get collection ID by its name (exact match).
  916. */
  917. export function getCollectionIdByName(db: Database, name: string): number | null {
  918. const result = db.prepare(`
  919. SELECT id FROM collections
  920. WHERE name = ?
  921. LIMIT 1
  922. `).get(name) as { id: number } | null;
  923. return result?.id || null;
  924. }
  925. /**
  926. * Get collection by name.
  927. */
  928. export function getCollectionByName(db: Database, name: string): { id: number; name: string; pwd: string; glob_pattern: string } | null {
  929. const result = db.prepare(`
  930. SELECT id, name, pwd, glob_pattern FROM collections
  931. WHERE name = ?
  932. LIMIT 1
  933. `).get(name) as { id: number; name: string; pwd: string; glob_pattern: string } | null;
  934. return result;
  935. }
  936. // =============================================================================
  937. // FTS Search
  938. // =============================================================================
  939. function sanitizeFTS5Term(term: string): string {
  940. return term.replace(/[^\p{L}\p{N}']/gu, '').toLowerCase();
  941. }
  942. function buildFTS5Query(query: string): string | null {
  943. const terms = query.split(/\s+/)
  944. .map(t => sanitizeFTS5Term(t))
  945. .filter(t => t.length > 0);
  946. if (terms.length === 0) return null;
  947. if (terms.length === 1) return `"${terms[0]}"*`;
  948. return terms.map(t => `"${t}"*`).join(' AND ');
  949. }
  950. export function searchFTS(db: Database, query: string, limit: number = 20, collectionId?: number): SearchResult[] {
  951. const ftsQuery = buildFTS5Query(query);
  952. if (!ftsQuery) return [];
  953. let sql = `
  954. SELECT
  955. 'qmd://' || c.name || '/' || d.path as filepath,
  956. 'qmd://' || c.name || '/' || d.path as display_path,
  957. d.title,
  958. content.doc as body,
  959. bm25(documents_fts, 10.0, 1.0) as score
  960. FROM documents_fts f
  961. JOIN documents d ON d.id = f.rowid
  962. JOIN collections c ON c.id = d.collection_id
  963. JOIN content ON content.hash = d.hash
  964. WHERE documents_fts MATCH ? AND d.active = 1
  965. `;
  966. const params: (string | number)[] = [ftsQuery];
  967. if (collectionId !== undefined) {
  968. sql += ` AND d.collection_id = ?`;
  969. params.push(collectionId);
  970. }
  971. sql += ` ORDER BY score LIMIT ?`;
  972. params.push(limit);
  973. const rows = db.prepare(sql).all(...params) as { filepath: string; display_path: string; title: string; body: string; score: number }[];
  974. const maxScore = rows.length > 0 ? Math.max(...rows.map(r => Math.abs(r.score))) : 1;
  975. return rows.map(row => ({
  976. file: row.filepath,
  977. displayPath: row.display_path,
  978. title: row.title,
  979. body: row.body,
  980. score: Math.abs(row.score) / maxScore,
  981. source: "fts" as const,
  982. }));
  983. }
  984. // =============================================================================
  985. // Vector Search
  986. // =============================================================================
  987. export async function searchVec(db: Database, query: string, model: string, limit: number = 20, collectionId?: number): Promise<SearchResult[]> {
  988. const tableExists = db.prepare(`SELECT name FROM sqlite_master WHERE type='table' AND name='vectors_vec'`).get();
  989. if (!tableExists) return [];
  990. const embedding = await getEmbedding(query, model, true);
  991. if (!embedding) return [];
  992. // sqlite-vec requires "k = ?" for KNN queries
  993. let sql = `
  994. SELECT
  995. v.hash_seq,
  996. v.distance,
  997. 'qmd://' || c.name || '/' || d.path as filepath,
  998. 'qmd://' || c.name || '/' || d.path as display_path,
  999. d.title,
  1000. content.doc as body,
  1001. cv.pos
  1002. FROM vectors_vec v
  1003. JOIN content_vectors cv ON cv.hash || '_' || cv.seq = v.hash_seq
  1004. JOIN documents d ON d.hash = cv.hash AND d.active = 1
  1005. JOIN collections c ON c.id = d.collection_id
  1006. JOIN content ON content.hash = d.hash
  1007. WHERE v.embedding MATCH ? AND k = ?
  1008. `;
  1009. if (collectionId !== undefined) {
  1010. sql += ` AND d.collection_id = ${collectionId}`;
  1011. }
  1012. sql += ` ORDER BY v.distance`;
  1013. const rows = db.prepare(sql).all(new Float32Array(embedding), limit * 3) as { hash_seq: string; distance: number; filepath: string; display_path: string; title: string; body: string; pos: number }[];
  1014. const seen = new Map<string, { row: typeof rows[0]; bestDist: number }>();
  1015. for (const row of rows) {
  1016. const existing = seen.get(row.filepath);
  1017. if (!existing || row.distance < existing.bestDist) {
  1018. seen.set(row.filepath, { row, bestDist: row.distance });
  1019. }
  1020. }
  1021. return Array.from(seen.values())
  1022. .sort((a, b) => a.bestDist - b.bestDist)
  1023. .slice(0, limit)
  1024. .map(({ row }) => ({
  1025. file: row.filepath,
  1026. displayPath: row.display_path,
  1027. title: row.title,
  1028. body: row.body,
  1029. score: 1 / (1 + row.distance),
  1030. source: "vec" as const,
  1031. chunkPos: row.pos,
  1032. }));
  1033. }
  1034. // =============================================================================
  1035. // Embeddings
  1036. // =============================================================================
  1037. async function getEmbedding(text: string, model: string, isQuery: boolean): Promise<number[] | null> {
  1038. const ollama = getDefaultOllama();
  1039. const result = await ollama.embed(text, { model, isQuery });
  1040. return result?.embedding || null;
  1041. }
  1042. // =============================================================================
  1043. // Query expansion
  1044. // =============================================================================
  1045. export async function expandQuery(query: string, model: string = DEFAULT_QUERY_MODEL, db: Database): Promise<string[]> {
  1046. // Check cache first
  1047. const cacheKey = getCacheKey("expandQuery", { query, model });
  1048. const cached = getCachedResult(db, cacheKey);
  1049. if (cached) {
  1050. const lines = cached.split('\n').map(l => l.trim()).filter(l => l.length > 0);
  1051. return [query, ...lines.slice(0, 2)];
  1052. }
  1053. const ollama = getDefaultOllama();
  1054. const results = await ollama.expandQuery(query, model, 2);
  1055. // Cache the expanded queries (excluding original)
  1056. if (results.length > 1) {
  1057. setCachedResult(db, cacheKey, results.slice(1).join('\n'));
  1058. }
  1059. return results;
  1060. }
  1061. // =============================================================================
  1062. // Reranking
  1063. // =============================================================================
  1064. export async function rerank(query: string, documents: { file: string; text: string }[], model: string = DEFAULT_RERANK_MODEL, db: Database): Promise<{ file: string; score: number }[]> {
  1065. const cachedResults: Map<string, number> = new Map();
  1066. const uncachedDocs: RerankDocument[] = [];
  1067. // Check cache for each document
  1068. for (const doc of documents) {
  1069. const cacheKey = getCacheKey("rerank", { query, file: doc.file, model });
  1070. const cached = getCachedResult(db, cacheKey);
  1071. if (cached !== null) {
  1072. cachedResults.set(doc.file, parseFloat(cached));
  1073. } else {
  1074. uncachedDocs.push({ file: doc.file, text: doc.text });
  1075. }
  1076. }
  1077. // Rerank uncached documents using Ollama
  1078. if (uncachedDocs.length > 0) {
  1079. const ollama = getDefaultOllama();
  1080. const rerankResult = await ollama.rerank(query, uncachedDocs, { model });
  1081. // Cache results
  1082. for (const result of rerankResult.results) {
  1083. const cacheKey = getCacheKey("rerank", { query, file: result.file, model });
  1084. setCachedResult(db, cacheKey, result.score.toString());
  1085. cachedResults.set(result.file, result.score);
  1086. }
  1087. }
  1088. // Return all results sorted by score
  1089. return documents
  1090. .map(doc => ({ file: doc.file, score: cachedResults.get(doc.file) || 0 }))
  1091. .sort((a, b) => b.score - a.score);
  1092. }
  1093. // =============================================================================
  1094. // Reciprocal Rank Fusion
  1095. // =============================================================================
  1096. export function reciprocalRankFusion(
  1097. resultLists: RankedResult[][],
  1098. weights: number[] = [],
  1099. k: number = 60
  1100. ): RankedResult[] {
  1101. const scores = new Map<string, { result: RankedResult; rrfScore: number; topRank: number }>();
  1102. for (let listIdx = 0; listIdx < resultLists.length; listIdx++) {
  1103. const list = resultLists[listIdx];
  1104. const weight = weights[listIdx] ?? 1.0;
  1105. for (let rank = 0; rank < list.length; rank++) {
  1106. const result = list[rank];
  1107. const rrfContribution = weight / (k + rank + 1);
  1108. const existing = scores.get(result.file);
  1109. if (existing) {
  1110. existing.rrfScore += rrfContribution;
  1111. existing.topRank = Math.min(existing.topRank, rank);
  1112. } else {
  1113. scores.set(result.file, {
  1114. result,
  1115. rrfScore: rrfContribution,
  1116. topRank: rank,
  1117. });
  1118. }
  1119. }
  1120. }
  1121. // Top-rank bonus
  1122. for (const entry of scores.values()) {
  1123. if (entry.topRank === 0) {
  1124. entry.rrfScore += 0.05;
  1125. } else if (entry.topRank <= 2) {
  1126. entry.rrfScore += 0.02;
  1127. }
  1128. }
  1129. return Array.from(scores.values())
  1130. .sort((a, b) => b.rrfScore - a.rrfScore)
  1131. .map(e => ({ ...e.result, score: e.rrfScore }));
  1132. }
  1133. // =============================================================================
  1134. // Document retrieval
  1135. // =============================================================================
  1136. type DbDocRow = {
  1137. filepath: string;
  1138. display_path: string;
  1139. title: string;
  1140. hash: string;
  1141. collection_id: number;
  1142. modified_at: string;
  1143. body_length: number;
  1144. body?: string;
  1145. };
  1146. /**
  1147. * Find a document by filename/path (with fuzzy matching)
  1148. * Returns document metadata without body by default
  1149. */
  1150. export function findDocument(db: Database, filename: string, options: { includeBody?: boolean } = {}): DocumentResult | DocumentNotFound {
  1151. let filepath = filename;
  1152. const colonMatch = filepath.match(/:(\d+)$/);
  1153. if (colonMatch) {
  1154. filepath = filepath.slice(0, -colonMatch[0].length);
  1155. }
  1156. if (filepath.startsWith('~/')) {
  1157. filepath = homedir() + filepath.slice(1);
  1158. }
  1159. const selectCols = options.includeBody
  1160. ? `filepath, display_path, title, hash, collection_id, modified_at, LENGTH(body) as body_length, body`
  1161. : `filepath, display_path, title, hash, collection_id, modified_at, LENGTH(body) as body_length`;
  1162. // Try various match strategies
  1163. let doc = db.prepare(`SELECT ${selectCols} FROM documents WHERE filepath = ? AND active = 1`).get(filepath) as DbDocRow | null;
  1164. if (!doc) {
  1165. doc = db.prepare(`SELECT ${selectCols} FROM documents WHERE display_path = ? AND active = 1`).get(filepath) as DbDocRow | null;
  1166. }
  1167. if (!doc) {
  1168. doc = db.prepare(`SELECT ${selectCols} FROM documents WHERE filepath LIKE ? AND active = 1 LIMIT 1`).get(`%${filepath}`) as DbDocRow | null;
  1169. }
  1170. if (!doc) {
  1171. doc = db.prepare(`SELECT ${selectCols} FROM documents WHERE display_path LIKE ? AND active = 1 LIMIT 1`).get(`%${filepath}`) as DbDocRow | null;
  1172. }
  1173. if (!doc) {
  1174. const similar = findSimilarFiles(db, filepath, 5, 5);
  1175. return { error: "not_found", query: filename, similarFiles: similar };
  1176. }
  1177. const context = getContextForFile(db, doc.filepath);
  1178. return {
  1179. filepath: doc.filepath,
  1180. displayPath: doc.display_path,
  1181. title: doc.title,
  1182. context,
  1183. hash: doc.hash,
  1184. collectionId: doc.collection_id,
  1185. modifiedAt: doc.modified_at,
  1186. bodyLength: doc.body_length,
  1187. ...(options.includeBody && doc.body !== undefined && { body: doc.body }),
  1188. };
  1189. }
  1190. /**
  1191. * Get the body content for a document
  1192. * Optionally slice by line range
  1193. */
  1194. export function getDocumentBody(db: Database, doc: DocumentResult | { filepath: string }, fromLine?: number, maxLines?: number): string | null {
  1195. const filepath = 'filepath' in doc ? doc.filepath : doc.filepath;
  1196. const row = db.prepare(`SELECT body FROM documents WHERE filepath = ? AND active = 1`).get(filepath) as { body: string } | null;
  1197. if (!row) return null;
  1198. let body = row.body;
  1199. if (fromLine !== undefined || maxLines !== undefined) {
  1200. const lines = body.split('\n');
  1201. const start = (fromLine || 1) - 1;
  1202. const end = maxLines !== undefined ? start + maxLines : lines.length;
  1203. body = lines.slice(start, end).join('\n');
  1204. }
  1205. return body;
  1206. }
  1207. /**
  1208. * Legacy function for backwards compatibility
  1209. * Combines findDocument + getDocumentBody with line slicing
  1210. */
  1211. export function getDocument(db: Database, filename: string, fromLine?: number, maxLines?: number): (DocumentResult & { body: string }) | DocumentNotFound {
  1212. // Parse :line suffix
  1213. let parsedFromLine = fromLine;
  1214. let filepath = filename;
  1215. const colonMatch = filepath.match(/:(\d+)$/);
  1216. if (colonMatch && !parsedFromLine) {
  1217. parsedFromLine = parseInt(colonMatch[1], 10);
  1218. filepath = filepath.slice(0, -colonMatch[0].length);
  1219. }
  1220. const result = findDocument(db, filepath, { includeBody: true });
  1221. if ("error" in result) return result;
  1222. let body = result.body || "";
  1223. if (parsedFromLine !== undefined || maxLines !== undefined) {
  1224. const lines = body.split('\n');
  1225. const start = (parsedFromLine || 1) - 1;
  1226. const end = maxLines !== undefined ? start + maxLines : lines.length;
  1227. body = lines.slice(start, end).join('\n');
  1228. }
  1229. return { ...result, body };
  1230. }
  1231. /**
  1232. * Find multiple documents by glob pattern or comma-separated list
  1233. * Returns documents without body by default (use getDocumentBody to load)
  1234. */
  1235. export function findDocuments(
  1236. db: Database,
  1237. pattern: string,
  1238. options: { includeBody?: boolean; maxBytes?: number } = {}
  1239. ): { docs: MultiGetResult[]; errors: string[] } {
  1240. const isCommaSeparated = pattern.includes(',') && !pattern.includes('*') && !pattern.includes('?');
  1241. const errors: string[] = [];
  1242. const maxBytes = options.maxBytes ?? DEFAULT_MULTI_GET_MAX_BYTES;
  1243. const selectCols = options.includeBody
  1244. ? `filepath, display_path, title, hash, collection_id, modified_at, LENGTH(body) as body_length, body`
  1245. : `filepath, display_path, title, hash, collection_id, modified_at, LENGTH(body) as body_length`;
  1246. let fileRows: DbDocRow[];
  1247. if (isCommaSeparated) {
  1248. const names = pattern.split(',').map(s => s.trim()).filter(Boolean);
  1249. fileRows = [];
  1250. for (const name of names) {
  1251. let doc = db.prepare(`SELECT ${selectCols} FROM documents WHERE display_path = ? AND active = 1`).get(name) as DbDocRow | null;
  1252. if (!doc) {
  1253. doc = db.prepare(`SELECT ${selectCols} FROM documents WHERE display_path LIKE ? AND active = 1 LIMIT 1`).get(`%${name}`) as DbDocRow | null;
  1254. }
  1255. if (doc) {
  1256. fileRows.push(doc);
  1257. } else {
  1258. const similar = findSimilarFiles(db, name, 5, 3);
  1259. let msg = `File not found: ${name}`;
  1260. if (similar.length > 0) {
  1261. msg += ` (did you mean: ${similar.join(', ')}?)`;
  1262. }
  1263. errors.push(msg);
  1264. }
  1265. }
  1266. } else {
  1267. // Glob pattern match
  1268. const matched = matchFilesByGlob(db, pattern);
  1269. if (matched.length === 0) {
  1270. errors.push(`No files matched pattern: ${pattern}`);
  1271. return { docs: [], errors };
  1272. }
  1273. const filepaths = matched.map(m => m.filepath);
  1274. const placeholders = filepaths.map(() => '?').join(',');
  1275. fileRows = db.prepare(`SELECT ${selectCols} FROM documents WHERE filepath IN (${placeholders}) AND active = 1`).all(...filepaths) as DbDocRow[];
  1276. }
  1277. const results: MultiGetResult[] = [];
  1278. for (const row of fileRows) {
  1279. const context = getContextForFile(db, row.filepath);
  1280. if (row.body_length > maxBytes) {
  1281. results.push({
  1282. doc: { filepath: row.filepath, displayPath: row.display_path },
  1283. skipped: true,
  1284. skipReason: `File too large (${Math.round(row.body_length / 1024)}KB > ${Math.round(maxBytes / 1024)}KB)`,
  1285. });
  1286. continue;
  1287. }
  1288. results.push({
  1289. doc: {
  1290. filepath: row.filepath,
  1291. displayPath: row.display_path,
  1292. title: row.title || row.display_path.split('/').pop() || row.display_path,
  1293. context,
  1294. hash: row.hash,
  1295. collectionId: row.collection_id,
  1296. modifiedAt: row.modified_at,
  1297. bodyLength: row.body_length,
  1298. ...(options.includeBody && row.body !== undefined && { body: row.body }),
  1299. },
  1300. skipped: false,
  1301. });
  1302. }
  1303. return { docs: results, errors };
  1304. }
  1305. /**
  1306. * Legacy function for backwards compatibility
  1307. */
  1308. export function getMultipleDocuments(db: Database, pattern: string, maxLines?: number, maxBytes: number = DEFAULT_MULTI_GET_MAX_BYTES): { files: MultiGetFile[]; errors: string[] } {
  1309. const { docs, errors } = findDocuments(db, pattern, { includeBody: true, maxBytes });
  1310. const files: MultiGetFile[] = docs.map(result => {
  1311. if (result.skipped) {
  1312. return {
  1313. filepath: result.doc.filepath,
  1314. displayPath: result.doc.displayPath,
  1315. title: "",
  1316. body: "",
  1317. context: null,
  1318. skipped: true as const,
  1319. skipReason: result.skipReason,
  1320. };
  1321. }
  1322. let body = result.doc.body || "";
  1323. if (maxLines !== undefined) {
  1324. const lines = body.split('\n');
  1325. body = lines.slice(0, maxLines).join('\n');
  1326. if (lines.length > maxLines) {
  1327. body += `\n\n[... truncated ${lines.length - maxLines} more lines]`;
  1328. }
  1329. }
  1330. return {
  1331. filepath: result.doc.filepath,
  1332. displayPath: result.doc.displayPath,
  1333. title: result.doc.title,
  1334. body,
  1335. context: result.doc.context,
  1336. skipped: false as const,
  1337. };
  1338. });
  1339. return { files, errors };
  1340. }
  1341. // Keep the old MultiGetFile type for backwards compatibility
  1342. export type MultiGetFile = {
  1343. filepath: string;
  1344. displayPath: string;
  1345. title: string;
  1346. body: string;
  1347. context: string | null;
  1348. skipped: false;
  1349. } | {
  1350. filepath: string;
  1351. displayPath: string;
  1352. title: string;
  1353. body: string;
  1354. context: string | null;
  1355. skipped: true;
  1356. skipReason: string;
  1357. };
  1358. // =============================================================================
  1359. // Status
  1360. // =============================================================================
  1361. export function getStatus(db: Database): IndexStatus {
  1362. const collections = db.prepare(`
  1363. SELECT c.id, c.pwd, c.glob_pattern, c.created_at,
  1364. COUNT(d.id) as active_count,
  1365. MAX(d.modified_at) as last_doc_update
  1366. FROM collections c
  1367. LEFT JOIN documents d ON d.collection_id = c.id AND d.active = 1
  1368. GROUP BY c.id
  1369. ORDER BY last_doc_update DESC
  1370. `).all() as { id: number; pwd: string; glob_pattern: string; created_at: string; active_count: number; last_doc_update: string | null }[];
  1371. const totalDocs = (db.prepare(`SELECT COUNT(*) as c FROM documents WHERE active = 1`).get() as { c: number }).c;
  1372. const needsEmbedding = getHashesNeedingEmbedding(db);
  1373. const hasVectors = !!db.prepare(`SELECT name FROM sqlite_master WHERE type='table' AND name='vectors_vec'`).get();
  1374. return {
  1375. totalDocuments: totalDocs,
  1376. needsEmbedding,
  1377. hasVectorIndex: hasVectors,
  1378. collections: collections.map(col => ({
  1379. id: col.id,
  1380. path: col.pwd,
  1381. pattern: col.glob_pattern,
  1382. documents: col.active_count,
  1383. lastUpdated: col.last_doc_update || col.created_at,
  1384. })),
  1385. };
  1386. }
  1387. // =============================================================================
  1388. // Snippet extraction
  1389. // =============================================================================
  1390. export type SnippetResult = {
  1391. line: number; // 1-indexed line number of best match
  1392. snippet: string; // The snippet text with diff-style header
  1393. linesBefore: number; // Lines in document before snippet
  1394. linesAfter: number; // Lines in document after snippet
  1395. snippetLines: number; // Number of lines in snippet
  1396. };
  1397. export function extractSnippet(body: string, query: string, maxLen = 500, chunkPos?: number): SnippetResult {
  1398. const totalLines = body.split('\n').length;
  1399. let searchBody = body;
  1400. let lineOffset = 0;
  1401. if (chunkPos && chunkPos > 0) {
  1402. const contextStart = Math.max(0, chunkPos - 100);
  1403. const contextEnd = Math.min(body.length, chunkPos + maxLen + 100);
  1404. searchBody = body.slice(contextStart, contextEnd);
  1405. if (contextStart > 0) {
  1406. lineOffset = body.slice(0, contextStart).split('\n').length - 1;
  1407. }
  1408. }
  1409. const lines = searchBody.split('\n');
  1410. const queryTerms = query.toLowerCase().split(/\s+/).filter(t => t.length > 0);
  1411. let bestLine = 0, bestScore = -1;
  1412. for (let i = 0; i < lines.length; i++) {
  1413. const lineLower = lines[i].toLowerCase();
  1414. let score = 0;
  1415. for (const term of queryTerms) {
  1416. if (lineLower.includes(term)) score++;
  1417. }
  1418. if (score > bestScore) {
  1419. bestScore = score;
  1420. bestLine = i;
  1421. }
  1422. }
  1423. const start = Math.max(0, bestLine - 1);
  1424. const end = Math.min(lines.length, bestLine + 3);
  1425. const snippetLines = lines.slice(start, end);
  1426. let snippetText = snippetLines.join('\n');
  1427. if (snippetText.length > maxLen) snippetText = snippetText.substring(0, maxLen - 3) + "...";
  1428. const absoluteStart = lineOffset + start + 1; // 1-indexed
  1429. const snippetLineCount = snippetLines.length;
  1430. const linesBefore = absoluteStart - 1;
  1431. const linesAfter = totalLines - (absoluteStart + snippetLineCount - 1);
  1432. // Format with diff-style header: @@ -start,count @@ (linesBefore before, linesAfter after)
  1433. const header = `@@ -${absoluteStart},${snippetLineCount} @@ (${linesBefore} before, ${linesAfter} after)`;
  1434. const snippet = `${header}\n${snippetText}`;
  1435. return {
  1436. line: lineOffset + bestLine + 1,
  1437. snippet,
  1438. linesBefore,
  1439. linesAfter,
  1440. snippetLines: snippetLineCount,
  1441. };
  1442. }