commit 39193ea252b1572f5d2d4c86a46ffaa184899df3 Author: Tobi Lutke Date: Sun Dec 7 19:16:16 2025 -0500 Initial commit: QMD - Quick Markdown Search A CLI tool for searching markdown knowledge bases using hybrid retrieval: - BM25 full-text search via SQLite FTS5 - Vector semantic search via sqlite-vec + Ollama embeddings - LLM re-ranking with qwen3-reranker (logprobs-based scoring) - Reciprocal Rank Fusion with weighted queries and position-aware blending Features: - `qmd add .` - Index markdown files in current directory - `qmd embed` - Generate vector embeddings - `qmd search` - BM25 full-text search - `qmd vsearch` - Vector similarity search - `qmd query` - Hybrid search with query expansion + reranking πŸ€– Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude Opus 4.5 diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..9f2e2b7 --- /dev/null +++ b/.gitignore @@ -0,0 +1,9 @@ +node_modules/ +*.sqlite +.DS_Store +archive/ +texts/ +.cursor/ +*.md +!README.md +!CLAUDE.md diff --git a/CLAUDE.md b/CLAUDE.md new file mode 100644 index 0000000..b8100b7 --- /dev/null +++ b/CLAUDE.md @@ -0,0 +1,111 @@ +--- +description: Use Bun instead of Node.js, npm, pnpm, or vite. +globs: "*.ts, *.tsx, *.html, *.css, *.js, *.jsx, package.json" +alwaysApply: false +--- + +Default to using Bun instead of Node.js. + +- Use `bun ` instead of `node ` or `ts-node ` +- Use `bun test` instead of `jest` or `vitest` +- Use `bun build ` instead of `webpack` or `esbuild` +- Use `bun install` instead of `npm install` or `yarn install` or `pnpm install` +- Use `bun run + + +``` + +With the following `frontend.tsx`: + +```tsx#frontend.tsx +import React from "react"; + +// import .css files directly and it works +import './index.css'; + +import { createRoot } from "react-dom/client"; + +const root = createRoot(document.body); + +export default function Frontend() { + return

Hello, world!

; +} + +root.render(); +``` + +Then, run index.ts + +```sh +bun --hot ./index.ts +``` + +For more information, read the Bun API docs in `node_modules/bun-types/docs/**.md`. diff --git a/README.md b/README.md new file mode 100644 index 0000000..fb751d8 --- /dev/null +++ b/README.md @@ -0,0 +1,333 @@ +# QMD - Quick Markdown Search + +A CLI tool for searching markdown knowledge bases using hybrid retrieval: combining BM25 full-text search, vector semantic search, and LLM re-ranking. + +## Architecture + +``` +β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β” +β”‚ QMD Search Pipeline β”‚ +β””β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜ + + β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β” + β”‚ User Query β”‚ + β””β”€β”€β”€β”€β”€β”€β”€β”€β”¬β”€β”€β”€β”€β”€β”€β”€β”€β”˜ + β”‚ + β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”΄β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β” + β–Ό β–Ό + β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β” β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β” + β”‚ Query Expansionβ”‚ β”‚ Direct Query β”‚ + β”‚ (qwen3:0.6b) β”‚ β”‚ (Γ—2 weight) β”‚ + β””β”€β”€β”€β”€β”€β”€β”€β”¬β”€β”€β”€β”€β”€β”€β”€β”€β”˜ β””β”€β”€β”€β”€β”€β”€β”€β”¬β”€β”€β”€β”€β”€β”€β”€β”€β”˜ + β”‚ β”‚ + β”‚ 1 alternative query β”‚ + β””β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”¬β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜ + β”‚ + β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”Όβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β” + β–Ό β–Ό β–Ό + β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β” β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β” β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β” + β”‚ FTS Search β”‚ β”‚ FTS Search β”‚ β”‚ FTS Search β”‚ + β”‚ (BM25) β”‚ β”‚ (BM25) β”‚ β”‚ (BM25) β”‚ + β””β”€β”€β”€β”€β”€β”€β”€β”¬β”€β”€β”€β”€β”€β”€β”€β”€β”˜ β””β”€β”€β”€β”€β”€β”€β”€β”¬β”€β”€β”€β”€β”€β”€β”€β”€β”˜ β””β”€β”€β”€β”€β”€β”€β”€β”¬β”€β”€β”€β”€β”€β”€β”€β”€β”˜ + β”‚ β”‚ β”‚ + β”Œβ”€β”€β”€β”€β”€β”€β”€β”΄β”€β”€β”€β”€β”€β”€β”€β”€β” β”Œβ”€β”€β”€β”€β”€β”€β”€β”΄β”€β”€β”€β”€β”€β”€β”€β”€β” β”Œβ”€β”€β”€β”€β”€β”€β”€β”΄β”€β”€β”€β”€β”€β”€β”€β”€β” + β”‚ Vector Search β”‚ β”‚ Vector Search β”‚ β”‚ Vector Search β”‚ + β”‚(embeddinggemma)β”‚ β”‚(embeddinggemma)β”‚ β”‚(embeddinggemma)β”‚ + β””β”€β”€β”€β”€β”€β”€β”€β”¬β”€β”€β”€β”€β”€β”€β”€β”€β”˜ β””β”€β”€β”€β”€β”€β”€β”€β”¬β”€β”€β”€β”€β”€β”€β”€β”€β”˜ β””β”€β”€β”€β”€β”€β”€β”€β”¬β”€β”€β”€β”€β”€β”€β”€β”€β”˜ + β”‚ β”‚ β”‚ + β””β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”Όβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜ + β”‚ + β–Ό + β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β” + β”‚ RRF Fusion + Bonus β”‚ + β”‚ (Top-rank preserved) β”‚ + β”‚ Top 30 Kept β”‚ + β””β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”¬β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜ + β”‚ + β–Ό + β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β” + β”‚ LLM Re-ranking β”‚ + β”‚ (qwen3-reranker) β”‚ + β”‚ Yes/No + logprobs β”‚ + β””β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”¬β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜ + β”‚ + β–Ό + β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β” + β”‚ Position-Aware Blend β”‚ + β”‚ (RRF + Reranker) β”‚ + β””β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜ +``` + +## Score Normalization & Fusion + +### Search Backends + +| Backend | Raw Score | Conversion | Range | +|---------|-----------|------------|-------| +| **FTS (BM25)** | SQLite FTS5 BM25 | `Math.abs(score)` | 0 to ~25+ | +| **Vector** | Cosine distance | `1 / (1 + distance)` | 0.0 to 1.0 | +| **Reranker** | LLM 0-10 rating | `score / 10` | 0.0 to 1.0 | + +### Fusion Strategy + +The `query` command uses **Reciprocal Rank Fusion (RRF)** with position-aware blending: + +1. **Query Expansion**: Original query (Γ—2 for weighting) + 1 LLM variation +2. **Parallel Retrieval**: Each query searches both FTS and vector indexes +3. **RRF Fusion**: Combine all result lists using `score = Ξ£(1/(k+rank+1))` where k=60 +4. **Top-Rank Bonus**: Documents ranking #1 in any list get +0.05, #2-3 get +0.02 +5. **Top-K Selection**: Take top 30 candidates for reranking +6. **Re-ranking**: LLM scores each document (yes/no with logprobs confidence) +7. **Position-Aware Blending**: + - RRF rank 1-3: 75% retrieval, 25% reranker (preserves exact matches) + - RRF rank 4-10: 60% retrieval, 40% reranker + - RRF rank 11+: 40% retrieval, 60% reranker (trust reranker more) + +**Why this approach**: Pure RRF can dilute exact matches when expanded queries don't match. The top-rank bonus preserves documents that score #1 for the original query. Position-aware blending prevents the reranker from destroying high-confidence retrieval results. + +### Score Interpretation + +| Score | Meaning | +|-------|---------| +| 0.8 - 1.0 | Highly relevant | +| 0.5 - 0.8 | Moderately relevant | +| 0.2 - 0.5 | Somewhat relevant | +| 0.0 - 0.2 | Low relevance | + +## Requirements + +### System Requirements + +- **Bun** >= 1.0.0 +- **macOS**: Homebrew SQLite (for extension support) + ```sh + brew install sqlite + ``` +- **Ollama** running locally (default: `http://localhost:11434`) + +### Ollama Models + +QMD uses three models (auto-pulled if missing): + +| Model | Purpose | Size | +|-------|---------|------| +| `embeddinggemma` | Vector embeddings | ~1.6GB | +| `ExpedientFalcon/qwen3-reranker:0.6b-q8_0` | Re-ranking (trained) | ~640MB | +| `qwen3:0.6b` | Query expansion | ~400MB | + +```sh +# Pre-pull models (optional) +ollama pull embeddinggemma +ollama pull ExpedientFalcon/qwen3-reranker:0.6b-q8_0 +ollama pull qwen3:0.6b +``` + +## Installation + +```sh +bun install +``` + +## Usage + +### Index Markdown Files + +```sh +# Index all .md files in current directory +qmd index + +# Index with custom glob pattern +qmd index "**/*.md" + +# Index specific directory +qmd index "docs/**/*.md" +``` + +### Generate Vector Embeddings + +```sh +# Embed all indexed documents +qmd embed +``` + +### Search Commands + +``` +β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β” +β”‚ Search Modes β”‚ +β”œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”¬β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€ +β”‚ search β”‚ BM25 full-text search only β”‚ +β”‚ vsearch β”‚ Vector semantic search only β”‚ +β”‚ query β”‚ Hybrid: FTS + Vector + Query Expansion + Re-ranking β”‚ +β””β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”΄β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜ +``` + +```sh +# Full-text search (fast, keyword-based) +qmd search "authentication flow" + +# Vector search (semantic similarity) +qmd vsearch "how to login" + +# Hybrid search with re-ranking (best quality) +qmd query "user authentication" +``` + +### Options + +```sh +-n # Number of results (default: 5) +--min-score # Minimum score threshold (default: 0) +--full # Show full document content +-csv # CSV output (for piping/scripting) +-md # Output as markdown +-xml # Output as XML +--index # Use named index +``` + +### Output Format + +Default output is colorized CLI format (respects `NO_COLOR` env): + +``` + 93% docs/guide.md:42 + β”‚ This section covers the **craftsmanship** of building + β”‚ quality software with attention to detail. + β”‚ See also: engineering principles + + 67% notes/meeting.md:15 + β”‚ Discussion about code quality and craftsmanship + β”‚ in the development process. +``` + +- **Score**: Color-coded (green >70%, yellow >40%, dim otherwise) +- **Path**: Shortened relative to current directory +- **Line**: Line number where match was found (omitted for vector-only results) +- **Snippet**: Context around match with query terms highlighted + +### Examples + +```sh +# Get 10 results with minimum score 0.3 +qmd query -n 10 --min-score 0.3 "API design patterns" + +# Output as markdown for LLM context +qmd search -md --full "error handling" + +# Use separate index for different knowledge base +qmd --index work search "quarterly reports" +``` + +### Manage Collections + +```sh +# List all indexed collections +qmd list + +# Show database statistics +qmd stats + +# Forget a collection +qmd forget +``` + +## Data Storage + +Index stored in: `~/.cache/qmd/index.sqlite` + +### Schema + +```sql +collections -- Indexed directories and glob patterns +documents -- Markdown content with metadata +documents_fts -- FTS5 full-text index +content_vectors -- Embedding cache (by content hash) +vectors_vec -- sqlite-vec vector index +``` + +## Environment Variables + +| Variable | Default | Description | +|----------|---------|-------------| +| `OLLAMA_URL` | `http://localhost:11434` | Ollama API endpoint | +| `XDG_CACHE_HOME` | `~/.cache` | Cache directory location | + +## How It Works + +### Indexing Flow + +``` +Markdown Files ──► Parse Title ──► Hash Content ──► Store in SQLite + β”‚ β”‚ + └─► FTS5 Index β—„β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜ +``` + +### Embedding Flow + +``` +Document ──► Format for EmbeddingGemma ──► Ollama API ──► Store Vector + "title: X | text: Y" /api/embed +``` + +### Query Flow (Hybrid) + +``` +Query ──► Expand (3 variations) ──► FTS + Vector (per variation) + β”‚ + β–Ό + Merge (max score) + β”‚ + β–Ό + Top 25 candidates + β”‚ + β–Ό + LLM Re-rank (0-10) + β”‚ + β–Ό + Final ranked results +``` + +## Model Configuration + +Models are configured as constants in `qmd.ts`: + +```typescript +const DEFAULT_EMBED_MODEL = "embeddinggemma"; +const DEFAULT_RERANK_MODEL = "ExpedientFalcon/qwen3-reranker:0.6b-q8_0"; +const DEFAULT_QUERY_MODEL = "qwen3:0.6b"; +``` + +### EmbeddingGemma Prompt Format + +``` +// For queries +"task: search result | query: {query}" + +// For documents +"title: {title} | text: {content}" +``` + +### Qwen3-Reranker + +A dedicated reranker model trained on relevance classification: + +``` +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". + +User: : Given a search query, determine if the document is relevant... + : {query} + : {doc} +``` + +- Uses `logprobs: true` to extract token probabilities +- Outputs yes/no with confidence score (0.0 - 1.0) +- `num_predict: 1` - Only need the yes/no token + +### Qwen3 (Query Expansion) + +- `num_predict: 150` - For generating query variations + +## License + +MIT diff --git a/bun.lock b/bun.lock new file mode 100644 index 0000000..0c5cc4e --- /dev/null +++ b/bun.lock @@ -0,0 +1,47 @@ +{ + "lockfileVersion": 1, + "configVersion": 1, + "workspaces": { + "": { + "name": "2025-12-07-bm25-q", + "dependencies": { + "sqlite-vec": "^0.1.7-alpha.2", + }, + "devDependencies": { + "@types/bun": "latest", + }, + "optionalDependencies": { + "sqlite-vec-darwin-arm64": "^0.1.7-alpha.2", + "sqlite-vec-darwin-x64": "^0.1.7-alpha.2", + "sqlite-vec-linux-x64": "^0.1.7-alpha.2", + "sqlite-vec-win32-x64": "^0.1.7-alpha.2", + }, + "peerDependencies": { + "typescript": "^5", + }, + }, + }, + "packages": { + "@types/bun": ["@types/bun@1.3.3", "", { "dependencies": { "bun-types": "1.3.3" } }, "sha512-ogrKbJ2X5N0kWLLFKeytG0eHDleBYtngtlbu9cyBKFtNL3cnpDZkNdQj8flVf6WTZUX5ulI9AY1oa7ljhSrp+g=="], + + "@types/node": ["@types/node@24.10.1", "", { "dependencies": { "undici-types": "~7.16.0" } }, "sha512-GNWcUTRBgIRJD5zj+Tq0fKOJ5XZajIiBroOF0yvj2bSU1WvNdYS/dn9UxwsujGW4JX06dnHyjV2y9rRaybH0iQ=="], + + "bun-types": ["bun-types@1.3.3", "", { "dependencies": { "@types/node": "*" } }, "sha512-z3Xwlg7j2l9JY27x5Qn3Wlyos8YAp0kKRlrePAOjgjMGS5IG6E7Jnlx736vH9UVI4wUICwwhC9anYL++XeOgTQ=="], + + "sqlite-vec": ["sqlite-vec@0.1.7-alpha.2", "", { "optionalDependencies": { "sqlite-vec-darwin-arm64": "0.1.7-alpha.2", "sqlite-vec-darwin-x64": "0.1.7-alpha.2", "sqlite-vec-linux-arm64": "0.1.7-alpha.2", "sqlite-vec-linux-x64": "0.1.7-alpha.2", "sqlite-vec-windows-x64": "0.1.7-alpha.2" } }, "sha512-rNgRCv+4V4Ed3yc33Qr+nNmjhtrMnnHzXfLVPeGb28Dx5mmDL3Ngw/Wk8vhCGjj76+oC6gnkmMG8y73BZWGBwQ=="], + + "sqlite-vec-darwin-arm64": ["sqlite-vec-darwin-arm64@0.1.7-alpha.2", "", { "os": "darwin", "cpu": "arm64" }, "sha512-raIATOqFYkeCHhb/t3r7W7Cf2lVYdf4J3ogJ6GFc8PQEgHCPEsi+bYnm2JT84MzLfTlSTIdxr4/NKv+zF7oLPw=="], + + "sqlite-vec-darwin-x64": ["sqlite-vec-darwin-x64@0.1.7-alpha.2", "", { "os": "darwin", "cpu": "x64" }, "sha512-jeZEELsQjjRsVojsvU5iKxOvkaVuE+JYC8Y4Ma8U45aAERrDYmqZoHvgSG7cg1PXL3bMlumFTAmHynf1y4pOzA=="], + + "sqlite-vec-linux-arm64": ["sqlite-vec-linux-arm64@0.1.7-alpha.2", "", { "os": "linux", "cpu": "arm64" }, "sha512-6Spj4Nfi7tG13jsUG+W7jnT0bCTWbyPImu2M8nWp20fNrd1SZ4g3CSlDAK8GBdavX7wRlbBHCZ+BDa++rbDewA=="], + + "sqlite-vec-linux-x64": ["sqlite-vec-linux-x64@0.1.7-alpha.2", "", { "os": "linux", "cpu": "x64" }, "sha512-IcgrbHaDccTVhXDf8Orwdc2+hgDLAFORl6OBUhcvlmwswwBP1hqBTSEhovClG4NItwTOBNgpwOoQ7Qp3VDPWLg=="], + + "sqlite-vec-windows-x64": ["sqlite-vec-windows-x64@0.1.7-alpha.2", "", { "os": "win32", "cpu": "x64" }, "sha512-TRP6hTjAcwvQ6xpCZvjP00pdlda8J38ArFy1lMYhtQWXiIBmWnhMaMbq4kaeCYwvTTddfidatRS+TJrwIKB/oQ=="], + + "typescript": ["typescript@5.9.3", "", { "bin": { "tsc": "bin/tsc", "tsserver": "bin/tsserver" } }, "sha512-jl1vZzPDinLr9eUt3J/t7V6FgNEw9QjvBPdysz9KfQDD41fQrC2Y4vKQdiaUpFT4bXlb1RHhLpp8wtm6M5TgSw=="], + + "undici-types": ["undici-types@7.16.0", "", {}, "sha512-Zz+aZWSj8LE6zoxD+xrjh4VfkIG8Ya6LvYkZqtUQGJPZjYl53ypCaUwWqo7eI0x66KBGeRo+mlBEkMSeSZ38Nw=="], + } +} diff --git a/package.json b/package.json new file mode 100644 index 0000000..38e57d8 --- /dev/null +++ b/package.json @@ -0,0 +1,47 @@ +{ + "name": "qmd", + "version": "1.0.0", + "description": "Quick Markdown Search - Full-text and vector search for markdown files", + "type": "module", + "bin": { + "qmd": "./qmd" + }, + "scripts": { + "qmd": "bun qmd.ts", + "index": "bun qmd.ts index", + "vector": "bun qmd.ts vector", + "search": "bun qmd.ts search", + "vsearch": "bun qmd.ts vsearch", + "rerank": "bun qmd.ts rerank", + "link": "bun link" + }, + "dependencies": { + "sqlite-vec": "^0.1.7-alpha.2" + }, + "optionalDependencies": { + "sqlite-vec-darwin-arm64": "^0.1.7-alpha.2", + "sqlite-vec-darwin-x64": "^0.1.7-alpha.2", + "sqlite-vec-linux-x64": "^0.1.7-alpha.2", + "sqlite-vec-win32-x64": "^0.1.7-alpha.2" + }, + "devDependencies": { + "@types/bun": "latest" + }, + "peerDependencies": { + "typescript": "^5" + }, + "engines": { + "bun": ">=1.0.0" + }, + "keywords": [ + "markdown", + "search", + "fts", + "vector", + "sqlite", + "bm25", + "embeddings", + "ollama" + ], + "license": "MIT" +} diff --git a/qmd b/qmd new file mode 100755 index 0000000..7db8a3c --- /dev/null +++ b/qmd @@ -0,0 +1,14 @@ +#!/bin/bash +# qmd - Quick Markdown Search +# Run with: ./qmd or symlink to PATH + +# Resolve symlinks to find actual script location +SOURCE="${BASH_SOURCE[0]}" +while [ -L "$SOURCE" ]; do + DIR="$(cd -P "$(dirname "$SOURCE")" && pwd)" + SOURCE="$(readlink "$SOURCE")" + [[ $SOURCE != /* ]] && SOURCE="$DIR/$SOURCE" +done +SCRIPT_DIR="$(cd -P "$(dirname "$SOURCE")" && pwd)" + +exec bun "$SCRIPT_DIR/qmd.ts" "$@" diff --git a/qmd.ts b/qmd.ts new file mode 100755 index 0000000..6cb7039 --- /dev/null +++ b/qmd.ts @@ -0,0 +1,1250 @@ +#!/usr/bin/env bun +import { Database } from "bun:sqlite"; +import { Glob } from "bun"; +import { mkdirSync, existsSync } from "node:fs"; +import { homedir } from "node:os"; +import { resolve } from "node:path"; +import * as sqliteVec from "sqlite-vec"; + +// On macOS, use Homebrew's SQLite which supports extensions +if (process.platform === "darwin") { + const homebrewSqlitePath = "/opt/homebrew/opt/sqlite/lib/libsqlite3.dylib"; + if (existsSync(homebrewSqlitePath)) { + Database.setCustomSQLite(homebrewSqlitePath); + } +} + +const DEFAULT_EMBED_MODEL = "embeddinggemma"; +const DEFAULT_RERANK_MODEL = "ExpedientFalcon/qwen3-reranker:0.6b-q8_0"; +const DEFAULT_QUERY_MODEL = "qwen3:0.6b"; +const DEFAULT_GLOB = "**/*.md"; +const OLLAMA_URL = process.env.OLLAMA_URL || "http://localhost:11434"; + +// Terminal colors (respects NO_COLOR env) +const useColor = !process.env.NO_COLOR && process.stdout.isTTY; +const c = { + reset: useColor ? "\x1b[0m" : "", + dim: useColor ? "\x1b[2m" : "", + bold: useColor ? "\x1b[1m" : "", + cyan: useColor ? "\x1b[36m" : "", + yellow: useColor ? "\x1b[33m" : "", + green: useColor ? "\x1b[32m" : "", + magenta: useColor ? "\x1b[35m" : "", + blue: useColor ? "\x1b[34m" : "", +}; + +// Global state for --index option +let customIndexName: string | null = null; + +// Terminal progress bar using OSC 9;4 escape sequence +const progress = { + set(percent: number) { + process.stderr.write(`\x1b]9;4;1;${Math.round(percent)}\x07`); + }, + clear() { + process.stderr.write(`\x1b]9;4;0\x07`); + }, + indeterminate() { + process.stderr.write(`\x1b]9;4;3\x07`); + }, + error() { + process.stderr.write(`\x1b]9;4;2\x07`); + }, +}; + +// Format seconds into human-readable ETA +function formatETA(seconds: number): string { + if (seconds < 60) return `${Math.round(seconds)}s`; + if (seconds < 3600) return `${Math.floor(seconds / 60)}m ${Math.round(seconds % 60)}s`; + return `${Math.floor(seconds / 3600)}h ${Math.floor((seconds % 3600) / 60)}m`; +} + +function getDbPath(): string { + const cacheDir = process.env.XDG_CACHE_HOME || resolve(homedir(), ".cache"); + const qmdCacheDir = resolve(cacheDir, "qmd"); + mkdirSync(qmdCacheDir, { recursive: true }); + const dbName = customIndexName || "index"; + return resolve(qmdCacheDir, `${dbName}.sqlite`); +} + +function getPwd(): string { + return process.env.PWD || process.cwd(); +} + +/* +Schema: + +CREATE TABLE collections ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + pwd TEXT NOT NULL, + glob_pattern TEXT NOT NULL, + created_at TEXT NOT NULL, + UNIQUE(pwd, glob_pattern) +); + +CREATE TABLE documents ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + collection_id INTEGER NOT NULL, + name TEXT NOT NULL, + title TEXT NOT NULL, + hash TEXT NOT NULL, + filepath TEXT NOT NULL, + body TEXT NOT NULL, + created_at TEXT NOT NULL, + modified_at TEXT NOT NULL, + active INTEGER NOT NULL DEFAULT 1, + FOREIGN KEY (collection_id) REFERENCES collections(id) +); + +CREATE TABLE content_vectors ( + hash TEXT PRIMARY KEY, + embedding BLOB NOT NULL, + model TEXT NOT NULL, + embedded_at TEXT NOT NULL +); + +CREATE VIRTUAL TABLE documents_fts USING fts5(...); +*/ + +function getDb(): Database { + const db = new Database(getDbPath()); + sqliteVec.load(db); + db.exec("PRAGMA journal_mode = WAL"); + + // Collections table + db.exec(` + CREATE TABLE IF NOT EXISTS collections ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + pwd TEXT NOT NULL, + glob_pattern TEXT NOT NULL, + created_at TEXT NOT NULL, + UNIQUE(pwd, glob_pattern) + ) + `); + + // Documents table with collection_id and full filepath + db.exec(` + CREATE TABLE IF NOT EXISTS documents ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + collection_id INTEGER NOT NULL, + name TEXT NOT NULL, + title TEXT NOT NULL, + hash TEXT NOT NULL, + filepath TEXT NOT NULL, + body TEXT NOT NULL, + created_at TEXT NOT NULL, + modified_at TEXT NOT NULL, + active INTEGER NOT NULL DEFAULT 1, + FOREIGN KEY (collection_id) REFERENCES collections(id) + ) + `); + + // Content vectors keyed by hash (UNIQUE) + db.exec(` + CREATE TABLE IF NOT EXISTS content_vectors ( + hash TEXT PRIMARY KEY, + model TEXT NOT NULL, + embedded_at TEXT NOT NULL + ) + `); + + // FTS on documents + db.exec(` + CREATE VIRTUAL TABLE IF NOT EXISTS documents_fts USING fts5( + name, body, + content='documents', + content_rowid='id', + tokenize='porter unicode61' + ) + `); + + db.exec(` + CREATE TRIGGER IF NOT EXISTS documents_ai AFTER INSERT ON documents BEGIN + INSERT INTO documents_fts(rowid, name, body) VALUES (new.id, new.name, new.body); + END + `); + + db.exec(` + CREATE TRIGGER IF NOT EXISTS documents_ad AFTER DELETE ON documents BEGIN + INSERT INTO documents_fts(documents_fts, rowid, name, body) VALUES('delete', old.id, old.name, old.body); + END + `); + + db.exec(` + CREATE TRIGGER IF NOT EXISTS documents_au AFTER UPDATE ON documents BEGIN + INSERT INTO documents_fts(documents_fts, rowid, name, body) VALUES('delete', old.id, old.name, old.body); + INSERT INTO documents_fts(rowid, name, body) VALUES (new.id, new.name, new.body); + END + `); + + db.exec(`CREATE INDEX IF NOT EXISTS idx_documents_collection ON documents(collection_id, active)`); + db.exec(`CREATE INDEX IF NOT EXISTS idx_documents_hash ON documents(hash)`); + db.exec(`CREATE INDEX IF NOT EXISTS idx_documents_filepath ON documents(filepath, active)`); + + return db; +} + +function ensureVecTable(db: Database, dimensions: number): void { + const tableInfo = db.prepare(`SELECT sql FROM sqlite_master WHERE type='table' AND name='vectors_vec'`).get() as { sql: string } | null; + if (tableInfo) { + const match = tableInfo.sql.match(/float\[(\d+)\]/); + if (match && parseInt(match[1]) === dimensions) return; + db.exec("DROP TABLE IF EXISTS vectors_vec"); + } + db.exec(`CREATE VIRTUAL TABLE vectors_vec USING vec0(hash TEXT PRIMARY KEY, embedding float[${dimensions}])`); +} + +function getHashesNeedingEmbedding(db: Database): number { + const result = db.prepare(` + SELECT COUNT(DISTINCT d.hash) as count + FROM documents d + LEFT JOIN content_vectors v ON d.hash = v.hash + WHERE d.active = 1 AND v.hash IS NULL + `).get() as { count: number }; + return result.count; +} + +async function hashContent(content: string): Promise { + const hash = new Bun.CryptoHasher("sha256"); + hash.update(content); + return hash.digest("hex"); +} + +// Extract title from first markdown headline, or use filename as fallback +function extractTitle(content: string, filename: string): string { + const match = content.match(/^##?\s+(.+)$/m); + if (match) return match[1].trim(); + return filename.replace(/\.md$/, "").split("/").pop() || filename; +} + +// Format text for EmbeddingGemma +function formatQueryForEmbedding(query: string): string { + return `task: search result | query: ${query}`; +} + +function formatDocForEmbedding(text: string, title?: string): string { + return `title: ${title || "none"} | text: ${text}`; +} + +// Auto-pull model if not found +async function ensureModelAvailable(model: string): Promise { + try { + const response = await fetch(`${OLLAMA_URL}/api/show`, { + method: "POST", + headers: { "Content-Type": "application/json" }, + body: JSON.stringify({ name: model }), + }); + if (response.ok) return; + } catch { + // Continue to pull attempt + } + + console.log(`Model ${model} not found. Pulling...`); + progress.indeterminate(); + + const pullResponse = await fetch(`${OLLAMA_URL}/api/pull`, { + method: "POST", + headers: { "Content-Type": "application/json" }, + body: JSON.stringify({ name: model, stream: false }), + }); + + if (!pullResponse.ok) { + progress.error(); + throw new Error(`Failed to pull model ${model}: ${pullResponse.status} - ${await pullResponse.text()}`); + } + + progress.clear(); + console.log(`Model ${model} pulled successfully.`); +} + +async function getEmbedding(text: string, model: string, isQuery: boolean = false, title?: string, retried: boolean = false): Promise { + const input = isQuery ? formatQueryForEmbedding(text) : formatDocForEmbedding(text, title); + + const response = await fetch(`${OLLAMA_URL}/api/embed`, { + method: "POST", + headers: { "Content-Type": "application/json" }, + body: JSON.stringify({ model, input }), + }); + if (!response.ok) { + const errorText = await response.text(); + if (!retried && (errorText.includes("not found") || errorText.includes("does not exist"))) { + await ensureModelAvailable(model); + return getEmbedding(text, model, isQuery, title, true); + } + throw new Error(`Ollama API error: ${response.status} - ${errorText}`); + } + const data = await response.json() as { embeddings: number[][] }; + return data.embeddings[0]; +} + +// Qwen3-Reranker prompt format (trained for yes/no relevance classification) +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".`; + +function formatRerankPrompt(query: string, title: string, doc: string): string { + return `: 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. +: ${query} +: ${title} +: ${doc}`; +} + +type LogProb = { token: string; logprob: number }; +type RerankResponse = { + response: string; + logprobs?: LogProb[]; +}; + +async function rerankSingle(prompt: string, model: string, retried: boolean = false): Promise { + // Use generate with raw template for qwen3-reranker format + // Include empty tags as per HuggingFace reference implementation + const fullPrompt = `<|im_start|>system +${RERANK_SYSTEM}<|im_end|> +<|im_start|>user +${prompt}<|im_end|> +<|im_start|>assistant + + + + +`; + + const response = await fetch(`${OLLAMA_URL}/api/generate`, { + method: "POST", + headers: { "Content-Type": "application/json" }, + body: JSON.stringify({ + model, + prompt: fullPrompt, + raw: true, + stream: false, + logprobs: true, + options: { num_predict: 1 }, + }), + }); + + if (!response.ok) { + const errorText = await response.text(); + if (!retried && (errorText.includes("not found") || errorText.includes("does not exist"))) { + await ensureModelAvailable(model); + return rerankSingle(prompt, model, true); + } + throw new Error(`Ollama API error: ${response.status} - ${errorText}`); + } + + const data = await response.json() as RerankResponse; + + // Extract score from logprobs - required for proper reranking + if (!data.logprobs || data.logprobs.length === 0) { + throw new Error("Reranker response missing logprobs - ensure Ollama supports logprobs"); + } + + const firstToken = data.logprobs[0]; + const token = firstToken.token.toLowerCase().trim(); + const confidence = Math.exp(firstToken.logprob); // 0-1, higher = more confident + + if (token === "yes") { + // Relevant: return confidence (e.g., 0.93 for high confidence yes) + return confidence; + } + if (token === "no") { + // Not relevant: return low score, scaled by inverse confidence + // High confidence "no" β†’ very low score + return (1 - confidence) * 0.3; // Cap at 0.3 for uncertain "no" + } + + throw new Error(`Unexpected reranker token: "${token}" (expected "yes" or "no")`); +} + +async function rerank(query: string, documents: { file: string; text: string }[], model: string = DEFAULT_RERANK_MODEL): Promise<{ file: string; score: number }[]> { + const results: { file: string; score: number }[] = []; + const total = documents.length; + const PARALLEL = 5; + + process.stderr.write(`Reranking ${total} documents with ${model} (parallel: ${PARALLEL})...\n`); + progress.indeterminate(); + + // Process in parallel batches + for (let i = 0; i < documents.length; i += PARALLEL) { + const batch = documents.slice(i, i + PARALLEL); + const batchResults = await Promise.all( + batch.map(async (doc) => { + try { + // Extract title from filename for reranker context + const title = doc.file.split('/').pop()?.replace(/\.md$/, '') || doc.file; + const prompt = formatRerankPrompt(query, title, doc.text.slice(0, 4000)); + const score = await rerankSingle(prompt, model); + return { file: doc.file, score }; + } catch (err) { + return { file: doc.file, score: 0 }; + } + }) + ); + results.push(...batchResults); + + const processed = Math.min(i + PARALLEL, total); + progress.set((processed / total) * 100); + process.stderr.write(`\rReranking: ${processed}/${total}`); + } + + progress.clear(); + process.stderr.write("\n"); + + return results.sort((a, b) => b.score - a.score); +} + +function getOrCreateCollection(db: Database, pwd: string, globPattern: string): number { + const now = new Date().toISOString(); + const existing = db.prepare(`SELECT id FROM collections WHERE pwd = ? AND glob_pattern = ?`).get(pwd, globPattern) as { id: number } | null; + if (existing) return existing.id; + + db.prepare(`INSERT INTO collections (pwd, glob_pattern, created_at) VALUES (?, ?, ?)`).run(pwd, globPattern, now); + return (db.prepare(`SELECT last_insert_rowid() as id`).get() as { id: number }).id; +} + +function listCollections(): void { + const db = getDb(); + const collections = db.prepare(` + SELECT c.id, c.pwd, c.glob_pattern, c.created_at, + COUNT(d.id) as doc_count, + SUM(CASE WHEN d.active = 1 THEN 1 ELSE 0 END) as active_count + FROM collections c + LEFT JOIN documents d ON d.collection_id = c.id + GROUP BY c.id + ORDER BY c.created_at DESC + `).all() as { id: number; pwd: string; glob_pattern: string; created_at: string; doc_count: number; active_count: number }[]; + + if (collections.length === 0) { + console.log("No collections found."); + db.close(); + return; + } + + console.log("Collections:\n"); + for (const c of collections) { + console.log(` ${c.pwd}`); + console.log(` Pattern: ${c.glob_pattern}`); + console.log(` Documents: ${c.active_count} active (${c.doc_count} total)`); + console.log(` Created: ${c.created_at}\n`); + } + + const hashCount = db.prepare(`SELECT COUNT(*) as count FROM content_vectors`).get() as { count: number }; + console.log(`Vectors: ${hashCount.count} unique content hashes embedded`); + + db.close(); +} + +async function updateAllCollections(): Promise { + const db = getDb(); + const collections = db.prepare(`SELECT id, pwd, glob_pattern FROM collections`).all() as { id: number; pwd: string; glob_pattern: string }[]; + + if (collections.length === 0) { + console.log("No collections found."); + db.close(); + return; + } + + db.close(); + + console.log(`Updating ${collections.length} collection(s)...\n`); + + for (const c of collections) { + console.log(`\n--- ${c.pwd} (${c.glob_pattern}) ---`); + // Temporarily set PWD for indexing + const originalPwd = process.env.PWD; + process.env.PWD = c.pwd; + await indexFiles(c.glob_pattern); + process.env.PWD = originalPwd; + } + + console.log("\nAll collections updated."); +} + +async function dropCollection(globPattern: string): Promise { + const db = getDb(); + const pwd = getPwd(); + + const collection = db.prepare(`SELECT id FROM collections WHERE pwd = ? AND glob_pattern = ?`).get(pwd, globPattern) as { id: number } | null; + + if (!collection) { + console.log(`No collection found for ${pwd} with pattern ${globPattern}`); + db.close(); + return; + } + + // Delete documents in this collection + const deleted = db.prepare(`DELETE FROM documents WHERE collection_id = ?`).run(collection.id); + + // Delete the collection + db.prepare(`DELETE FROM collections WHERE id = ?`).run(collection.id); + + console.log(`Dropped collection: ${pwd} (${globPattern})`); + console.log(`Removed ${deleted.changes} documents`); + console.log(`(Vectors kept for potential reuse)`); + + db.close(); +} + +async function indexFiles(globPattern: string = DEFAULT_GLOB): Promise { + const db = getDb(); + const pwd = getPwd(); + const now = new Date().toISOString(); + const excludeDirs = ["node_modules", ".git", ".cache", "vendor", "dist", "build"]; + + // Get or create collection for this (pwd, glob) + const collectionId = getOrCreateCollection(db, pwd, globPattern); + console.log(`Collection: ${pwd} (${globPattern})`); + + progress.indeterminate(); + const glob = new Glob(globPattern); + const files: string[] = []; + for await (const file of glob.scan({ cwd: pwd, onlyFiles: true, followSymlinks: true })) { + // Skip node_modules, hidden folders (.*), and other common excludes + const parts = file.split("/"); + const shouldSkip = parts.some(part => + part === "node_modules" || + part.startsWith(".") || + excludeDirs.includes(part) + ); + if (!shouldSkip) { + files.push(file); + } + } + + const total = files.length; + if (total === 0) { + progress.clear(); + console.log("No files found matching pattern."); + db.close(); + return; + } + + const insertStmt = db.prepare(`INSERT INTO documents (collection_id, name, title, hash, filepath, body, created_at, modified_at, active) VALUES (?, ?, ?, ?, ?, ?, ?, ?, 1)`); + const deactivateStmt = db.prepare(`UPDATE documents SET active = 0 WHERE collection_id = ? AND filepath = ? AND active = 1`); + const findActiveStmt = db.prepare(`SELECT id, hash FROM documents WHERE collection_id = ? AND filepath = ? AND active = 1`); + + let indexed = 0, updated = 0, unchanged = 0, processed = 0; + const seenFiles = new Set(); + const startTime = Date.now(); + + for (const relativeFile of files) { + const filepath = resolve(pwd, relativeFile); + seenFiles.add(filepath); + + const content = await Bun.file(filepath).text(); + const hash = await hashContent(content); + const name = relativeFile.replace(/\.md$/, "").split("/").pop() || relativeFile; + const title = extractTitle(content, relativeFile); + const existing = findActiveStmt.get(collectionId, filepath) as { id: number; hash: string } | null; + + if (existing) { + if (existing.hash === hash) { + unchanged++; + } else { + deactivateStmt.run(collectionId, filepath); + updated++; + const stat = await Bun.file(filepath).stat(); + insertStmt.run(collectionId, name, title, hash, filepath, content, stat ? new Date(stat.birthtime).toISOString() : now, stat ? new Date(stat.mtime).toISOString() : now); + } + } else { + indexed++; + const stat = await Bun.file(filepath).stat(); + insertStmt.run(collectionId, name, title, hash, filepath, content, stat ? new Date(stat.birthtime).toISOString() : now, stat ? new Date(stat.mtime).toISOString() : now); + } + + processed++; + progress.set((processed / total) * 100); + const elapsed = (Date.now() - startTime) / 1000; + const rate = processed / elapsed; + const remaining = (total - processed) / rate; + const eta = processed > 2 ? ` ETA: ${formatETA(remaining)}` : ""; + process.stderr.write(`\rIndexing: ${processed}/${total}${eta} `); + } + + // Deactivate documents in this collection that no longer exist + const allActive = db.prepare(`SELECT filepath FROM documents WHERE collection_id = ? AND active = 1`).all(collectionId) as { filepath: string }[]; + let removed = 0; + for (const row of allActive) { + if (!seenFiles.has(row.filepath)) { + deactivateStmt.run(collectionId, row.filepath); + removed++; + } + } + + // Check if vector index needs updating + const needsEmbedding = getHashesNeedingEmbedding(db); + + progress.clear(); + console.log(`\nIndexed: ${indexed} new, ${updated} updated, ${unchanged} unchanged, ${removed} removed`); + + if (needsEmbedding > 0) { + console.log(`\nRun 'qmd vector' to update embeddings (${needsEmbedding} unique hashes need vectors)`); + } + + db.close(); +} + +async function vectorIndex(model: string = DEFAULT_EMBED_MODEL, force: boolean = false): Promise { + const db = getDb(); + const now = new Date().toISOString(); + + // If force, clear all vectors + if (force) { + console.log("Force re-indexing: clearing all vectors..."); + db.exec(`DELETE FROM content_vectors`); + db.exec(`DROP TABLE IF EXISTS vectors_vec`); + } + + // Find unique hashes that need embedding (from active documents) + const hashesToEmbed = db.prepare(` + SELECT DISTINCT d.hash, d.title, d.body + FROM documents d + LEFT JOIN content_vectors v ON d.hash = v.hash + WHERE d.active = 1 AND v.hash IS NULL + `).all() as { hash: string; title: string; body: string }[]; + + if (hashesToEmbed.length === 0) { + console.log("All content hashes already have embeddings."); + db.close(); + return; + } + + const total = hashesToEmbed.length; + console.log(`Embedding ${total} unique content hashes with ${model}...`); + + progress.indeterminate(); + const firstEmbedding = await getEmbedding(hashesToEmbed[0].body, model, false, hashesToEmbed[0].title); + console.log(`Embedding dimensions: ${firstEmbedding.length}`); + ensureVecTable(db, firstEmbedding.length); + + const insertVecStmt = db.prepare(`INSERT INTO vectors_vec (hash, embedding) VALUES (?, ?)`); + const insertContentVectorStmt = db.prepare(`INSERT OR REPLACE INTO content_vectors (hash, model, embedded_at) VALUES (?, ?, ?)`); + + let embedded = 0, errors = 0; + const startTime = Date.now(); + + // Insert first + insertVecStmt.run(hashesToEmbed[0].hash, new Float32Array(firstEmbedding)); + insertContentVectorStmt.run(hashesToEmbed[0].hash, model, now); + embedded++; + progress.set((embedded / total) * 100); + process.stderr.write(`\rEmbedding: ${embedded}/${total}`); + + for (let i = 1; i < hashesToEmbed.length; i++) { + const item = hashesToEmbed[i]; + try { + const embedding = await getEmbedding(item.body, model, false, item.title); + insertVecStmt.run(item.hash, new Float32Array(embedding)); + insertContentVectorStmt.run(item.hash, model, now); + embedded++; + } catch (err) { + errors++; + progress.error(); + console.error(`\nError embedding hash ${item.hash.slice(0, 8)}...: ${err}`); + } + const processed = embedded + errors; + progress.set((processed / total) * 100); + const elapsed = (Date.now() - startTime) / 1000; + const rate = processed / elapsed; + const remaining = (total - processed) / rate; + const eta = processed > 2 ? ` ETA: ${formatETA(remaining)}` : ""; + process.stderr.write(`\rEmbedding: ${embedded}/${total}${errors > 0 ? ` (${errors} errors)` : ""}${eta} `); + } + + progress.clear(); + const totalTime = ((Date.now() - startTime) / 1000).toFixed(1); + console.log(`\nDone! Embedded ${embedded} hashes${errors > 0 ? `, ${errors} errors` : ""} in ${totalTime}s.`); + db.close(); +} + +function escapeCSV(value: string): string { + if (value.includes('"') || value.includes(',') || value.includes('\n')) { + return `"${value.replace(/"/g, '""')}"`; + } + return value; +} + +function extractSnippet(body: string, query: string, maxLen = 500): { line: number; snippet: string } { + const lines = body.split('\n'); + const queryTerms = query.toLowerCase().split(/\s+/).filter(t => t.length > 0); + let bestLine = 0, bestScore = -1; + + for (let i = 0; i < lines.length; i++) { + const lineLower = lines[i].toLowerCase(); + let score = 0; + for (const term of queryTerms) { + if (lineLower.includes(term)) score++; + } + if (score > bestScore) { + bestScore = score; + bestLine = i; + } + } + + const startLine = Math.max(0, bestLine - 1); + const endLine = Math.min(lines.length, bestLine + 2); + let snippet = lines.slice(startLine, endLine).join('\n'); + if (snippet.length > maxLen) snippet = snippet.substring(0, maxLen - 3) + "..."; + return { line: bestLine + 1, snippet }; +} + +type SearchResult = { file: string; body: string; score: number; source: "fts" | "vec" }; + +// Build FTS5 query: phrase-aware with fallback to individual terms +function buildFTS5Query(query: string): string { + const terms = query + .split(/\s+/) + .filter(term => term.length >= 2); // Skip single chars + + if (terms.length === 0) return ""; + if (terms.length === 1) return `"${terms[0].replace(/"/g, '""')}"`; + + // Strategy: exact phrase OR proximity match OR individual terms + // Exact phrase matches rank highest, then close proximity, then any term + const phrase = `"${query.replace(/"/g, '""')}"`; + const quotedTerms = terms.map(t => `"${t.replace(/"/g, '""')}"`); + + // FTS5 NEAR syntax: NEAR(term1 term2, distance) + const nearPhrase = `NEAR(${quotedTerms.join(' ')}, 10)`; + const orTerms = quotedTerms.join(' OR '); + + // Exact phrase > proximity > any term + return `(${phrase}) OR (${nearPhrase}) OR (${orTerms})`; +} + +// Normalize BM25 score to 0-1 range using sigmoid +function normalizeBM25(score: number): number { + // BM25 scores are negative in SQLite (lower = better) + // Typical range: -15 (excellent) to -2 (weak match) + // Map to 0-1 where higher is better + const absScore = Math.abs(score); + // Sigmoid-ish normalization: maps ~2-15 range to ~0.1-0.95 + return 1 / (1 + Math.exp(-(absScore - 5) / 3)); +} + +function searchFTS(db: Database, query: string, limit: number = 20): SearchResult[] { + const ftsQuery = buildFTS5Query(query); + if (!ftsQuery) return []; + + // BM25 weights: name=10, body=1 (title matches ranked higher) + const stmt = db.prepare(` + SELECT d.filepath, d.body, bm25(documents_fts, 10.0, 1.0) as score + FROM documents_fts f + JOIN documents d ON d.id = f.rowid + WHERE documents_fts MATCH ? AND d.active = 1 + ORDER BY score + LIMIT ? + `); + const results = stmt.all(ftsQuery, limit) as { filepath: string; body: string; score: number }[]; + return results.map(r => ({ + file: r.filepath, + body: r.body, + score: normalizeBM25(r.score), + source: "fts" as const, + })); +} + +async function searchVec(db: Database, query: string, model: string, limit: number = 20): Promise { + const tableExists = db.prepare(`SELECT name FROM sqlite_master WHERE type='table' AND name='vectors_vec'`).get(); + if (!tableExists) return []; + + const queryEmbedding = await getEmbedding(query, model, true); + const queryVec = new Float32Array(queryEmbedding); + + // Join: documents -> content_vectors -> vectors_vec + const stmt = db.prepare(` + SELECT d.filepath, d.body, vec.distance + FROM vectors_vec vec + JOIN documents d ON d.hash = vec.hash + WHERE vec.embedding MATCH ? AND k = ? AND d.active = 1 + ORDER BY vec.distance + `); + const results = stmt.all(queryVec, limit) as { filepath: string; body: string; distance: number }[]; + return results.map(r => ({ + file: r.filepath, + body: r.body, + score: 1 / (1 + r.distance), + source: "vec" as const, + })); +} + +function normalizeScores(results: SearchResult[]): SearchResult[] { + if (results.length === 0) return results; + const maxScore = Math.max(...results.map(r => r.score)); + const minScore = Math.min(...results.map(r => r.score)); + const range = maxScore - minScore || 1; + return results.map(r => ({ ...r, score: (r.score - minScore) / range })); +} + +// Reciprocal Rank Fusion: combines multiple ranked lists +// RRF score = sum(1 / (k + rank)) across all lists where doc appears +// k=60 is standard, provides good balance between top and lower ranks +type RankedResult = { file: string; body: string; score: number }; + +function reciprocalRankFusion( + resultLists: RankedResult[][], + weights: number[] = [], // Weight per result list (default 1.0) + k: number = 60 +): RankedResult[] { + const scores = new Map(); + + for (let listIdx = 0; listIdx < resultLists.length; listIdx++) { + const results = resultLists[listIdx]; + const weight = weights[listIdx] ?? 1.0; + for (let rank = 0; rank < results.length; rank++) { + const doc = results[rank]; + const rrfScore = weight / (k + rank + 1); + const existing = scores.get(doc.file); + if (existing) { + existing.score += rrfScore; + existing.bestRank = Math.min(existing.bestRank, rank); + } else { + scores.set(doc.file, { score: rrfScore, body: doc.body, bestRank: rank }); + } + } + } + + // Add bonus for best rank: documents that ranked #1-3 in any list get a boost + // This prevents dilution of exact matches by expansion queries + return Array.from(scores.entries()) + .map(([file, { score, body, bestRank }]) => { + let bonus = 0; + if (bestRank === 0) bonus = 0.05; // Ranked #1 somewhere + else if (bestRank <= 2) bonus = 0.02; // Ranked top-3 somewhere + return { file, body, score: score + bonus }; + }) + .sort((a, b) => b.score - a.score); +} + +type OutputFormat = "cli" | "csv" | "md" | "xml"; +type OutputOptions = { + format: OutputFormat; + full: boolean; + limit: number; + minScore: number; +}; + +// Extract snippet with more context lines for CLI display +function extractSnippetWithContext(body: string, query: string, contextLines = 3): { line: number; snippet: string; hasMatch: boolean } { + const lines = body.split('\n'); + const queryTerms = query.toLowerCase().split(/\s+/).filter(t => t.length > 0); + let bestLine = 0, bestScore = -1; + + for (let i = 0; i < lines.length; i++) { + const lineLower = lines[i].toLowerCase(); + let score = 0; + for (const term of queryTerms) { + if (lineLower.includes(term)) score++; + } + if (score > bestScore) { + bestScore = score; + bestLine = i; + } + } + + // No query match found - return beginning of file + if (bestScore <= 0) { + const preview = lines.slice(0, contextLines * 2).join('\n').trim(); + return { line: 1, snippet: preview, hasMatch: false }; + } + + const startLine = Math.max(0, bestLine - contextLines); + const endLine = Math.min(lines.length, bestLine + contextLines + 1); + const snippet = lines.slice(startLine, endLine).join('\n').trim(); + return { line: bestLine + 1, snippet, hasMatch: true }; +} + +// Highlight query terms in text (skip short words < 3 chars) +function highlightTerms(text: string, query: string): string { + if (!useColor) return text; + const terms = query.toLowerCase().split(/\s+/).filter(t => t.length >= 3); + let result = text; + for (const term of terms) { + const regex = new RegExp(`(${term.replace(/[.*+?^${}()|[\]\\]/g, '\\$&')})`, 'gi'); + result = result.replace(regex, `${c.yellow}${c.bold}$1${c.reset}`); + } + return result; +} + +// Format score with color based on value +function formatScore(score: number): string { + const pct = (score * 100).toFixed(0).padStart(3); + if (!useColor) return `${pct}%`; + if (score >= 0.7) return `${c.green}${pct}%${c.reset}`; + if (score >= 0.4) return `${c.yellow}${pct}%${c.reset}`; + return `${c.dim}${pct}%${c.reset}`; +} + +// Shorten filepath for display +function shortPath(filepath: string): string { + const cwd = getPwd(); + if (filepath.startsWith(cwd)) { + return filepath.slice(cwd.length + 1); + } + // Show last 2 path components + const parts = filepath.split('/'); + if (parts.length > 2) { + return '.../' + parts.slice(-2).join('/'); + } + return filepath; +} + +function outputResults(results: { file: string; body: string; score: number }[], query: string, opts: OutputOptions): void { + const filtered = results.filter(r => r.score >= opts.minScore).slice(0, opts.limit); + + if (filtered.length === 0) { + console.log("No results found above minimum score threshold."); + return; + } + + if (opts.format === "cli") { + for (let i = 0; i < filtered.length; i++) { + const row = filtered[i]; + const { line, snippet, hasMatch } = extractSnippetWithContext(row.body, query, 2); + + // Header: score and filename + const score = formatScore(row.score); + const path = shortPath(row.file); + const lineInfo = hasMatch ? `:${line}` : ""; + console.log(`${c.bold}${score}${c.reset} ${c.cyan}${path}${c.dim}${lineInfo}${c.reset}`); + + // Snippet with highlighting + const highlighted = highlightTerms(snippet, query); + const indented = highlighted.split('\n').map(l => ` ${c.dim}β”‚${c.reset} ${l}`).join('\n'); + console.log(indented); + + if (i < filtered.length - 1) console.log(); + } + } else if (opts.format === "md") { + for (const row of filtered) { + if (opts.full) { + console.log(`---\n# ${row.file}\n\n${row.body}\n`); + } else { + const { snippet } = extractSnippet(row.body, query); + console.log(`---\n# ${row.file}\n\n${snippet}\n`); + } + } + } else if (opts.format === "xml") { + for (const row of filtered) { + if (opts.full) { + console.log(`\n${row.body}\n\n`); + } else { + const { snippet } = extractSnippet(row.body, query); + console.log(`\n${snippet}\n\n`); + } + } + } else { + // CSV format + console.log("score,file,line,snippet"); + for (const row of filtered) { + const { line, snippet } = extractSnippet(row.body, query); + const content = opts.full ? row.body : snippet; + console.log(`${row.score.toFixed(4)},${escapeCSV(row.file)},${line},${escapeCSV(content)}`); + } + } +} + +function search(query: string, opts: OutputOptions): void { + const db = getDb(); + const results = searchFTS(db, query, 50); + db.close(); + + if (results.length === 0) { + console.log("No results found."); + return; + } + outputResults(results, query, opts); +} + +async function vectorSearch(query: string, opts: OutputOptions, model: string = DEFAULT_EMBED_MODEL): Promise { + const db = getDb(); + + const tableExists = db.prepare(`SELECT name FROM sqlite_master WHERE type='table' AND name='vectors_vec'`).get(); + if (!tableExists) { + console.error("Vector index not found. Run 'qmd vector' first to create embeddings."); + db.close(); + return; + } + + // Expand query to multiple variations + const queries = await expandQuery(query); + process.stderr.write(`Searching with ${queries.length} query variations...\n`); + + // Collect results from all query variations + const allResults = new Map(); + + for (const q of queries) { + const vecResults = await searchVec(db, q, model, 20); + for (const r of vecResults) { + const existing = allResults.get(r.file); + if (!existing || r.score > existing.score) { + allResults.set(r.file, { file: r.file, body: r.body, score: r.score }); + } + } + } + + db.close(); + + // Sort by max score and limit to requested count + const results = Array.from(allResults.values()) + .sort((a, b) => b.score - a.score) + .slice(0, opts.limit); + + if (results.length === 0) { + console.log("No results found."); + return; + } + outputResults(results, query, { ...opts, limit: results.length }); // Already limited +} + +async function expandQuery(query: string, model: string = DEFAULT_QUERY_MODEL): Promise { + process.stderr.write("Generating query variations...\n"); + + const prompt = `Generate 3 search query variations to find documents about this topic. + +IMPORTANT: Keep multi-word phrases intact if they look like names (e.g., "Build a Business" should stay as "Build a Business", not "create a company"). + +Query: "${query}" + +Output 3 variations, one per line:`; + + const response = await fetch(`${OLLAMA_URL}/api/generate`, { + method: "POST", + headers: { "Content-Type": "application/json" }, + body: JSON.stringify({ + model, + prompt, + stream: false, + think: false, // Disable thinking mode for qwen3 models + options: { num_predict: 150 }, + }), + }); + + if (!response.ok) { + const errorText = await response.text(); + if (errorText.includes("not found") || errorText.includes("does not exist")) { + await ensureModelAvailable(model); + return expandQuery(query, model); + } + // Fall back to original query if expansion fails + return [query]; + } + + const data = await response.json() as { response: string }; + const lines = data.response.trim().split('\n') + .map(l => l.replace(/^[\d\.\-\*\"\s]+/, '').replace(/["\s]+$/, '').trim()) + .filter(l => l.length > 0 && !l.startsWith('<')) + .slice(0, 1); // Only 1 expanded query to preserve original query signal + + // Original query + expansions (original gets 2x weight in RRF) + const allQueries = [query, ...lines]; + process.stderr.write(`Queries:\n - ${allQueries.join('\n - ')}\n`); + return allQueries; +} + +async function querySearch(query: string, opts: OutputOptions, embedModel: string = DEFAULT_EMBED_MODEL, rerankModel: string = DEFAULT_RERANK_MODEL): Promise { + const db = getDb(); + + // Expand query to multiple variations + const queries = await expandQuery(query); + process.stderr.write(`Searching with ${queries.length} query variations...\n`); + + // Collect ranked result lists for RRF fusion + const rankedLists: RankedResult[][] = []; + const hasVectors = !!db.prepare(`SELECT name FROM sqlite_master WHERE type='table' AND name='vectors_vec'`).get(); + + for (const q of queries) { + // FTS search - get ranked results + const ftsResults = searchFTS(db, q, 20); + if (ftsResults.length > 0) { + rankedLists.push(ftsResults.map(r => ({ file: r.file, body: r.body, score: r.score }))); + } + + // Vector search - get ranked results + if (hasVectors) { + const vecResults = await searchVec(db, q, embedModel, 20); + if (vecResults.length > 0) { + rankedLists.push(vecResults.map(r => ({ file: r.file, body: r.body, score: r.score }))); + } + } + } + + // Apply Reciprocal Rank Fusion to combine all ranked lists + // Give 2x weight to original query results (first 2 lists: FTS + vector) + const weights = rankedLists.map((_, i) => i < 2 ? 2.0 : 1.0); + const fused = reciprocalRankFusion(rankedLists, weights); + const candidates = fused.slice(0, 30); // Over-retrieve for reranking + + if (candidates.length === 0) { + console.log("No results found."); + db.close(); + return; + } + + // Rerank with the original query + const reranked = await rerank( + query, + candidates.map(c => ({ file: c.file, text: c.body })), + rerankModel + ); + + db.close(); + + // Blend RRF position score with reranker score using position-aware weights + // Top retrieval results get more protection from reranker disagreement + const bodyMap = new Map(candidates.map(c => [c.file, c.body])); + const rrfRankMap = new Map(candidates.map((c, i) => [c.file, i + 1])); // 1-indexed rank + + const finalResults = reranked.map(r => { + const rrfRank = rrfRankMap.get(r.file) || 30; + // Position-aware blending: top retrieval results preserved more + // Rank 1-3: 75% RRF, 25% reranker (trust retrieval for exact matches) + // Rank 4-10: 60% RRF, 40% reranker + // Rank 11+: 40% RRF, 60% reranker (trust reranker for lower-ranked) + let rrfWeight: number; + if (rrfRank <= 3) { + rrfWeight = 0.75; + } else if (rrfRank <= 10) { + rrfWeight = 0.60; + } else { + rrfWeight = 0.40; + } + const rrfScore = 1 / rrfRank; // Position-based: 1, 0.5, 0.33... + const blendedScore = rrfWeight * rrfScore + (1 - rrfWeight) * r.score; + return { + file: r.file, + body: bodyMap.get(r.file) || "", + score: blendedScore, + }; + }).sort((a, b) => b.score - a.score); + + outputResults(finalResults, query, opts); +} + +// Parse CLI options +function parseOptions(args: string[], defaultMinScore: number = 0): { opts: OutputOptions; query: string } { + let format: OutputFormat = "cli"; + let full = false; + let limit = 5; + let minScore = defaultMinScore; + const queryParts: string[] = []; + + for (let i = 0; i < args.length; i++) { + const arg = args[i]; + if (arg === "-n" && i + 1 < args.length) { + limit = parseInt(args[++i], 10) || 5; + } else if (arg === "--min-score" && i + 1 < args.length) { + minScore = parseFloat(args[++i]) || defaultMinScore; + } else if (arg === "--full") { + full = true; + } else if (arg === "-csv" || arg === "--csv") { + format = "csv"; + } else if (arg === "-md" || arg === "--md") { + format = "md"; + } else if (arg === "-xml" || arg === "--xml") { + format = "xml"; + } else if (!arg.startsWith("-")) { + queryParts.push(arg); + } + } + + return { + opts: { format, full, limit, minScore }, + query: queryParts.join(" "), + }; +} + +// Parse global options and extract remaining args +function parseGlobalOptions(args: string[]): string[] { + const remaining: string[] = []; + for (let i = 0; i < args.length; i++) { + if (args[i] === "--index" && i + 1 < args.length) { + customIndexName = args[++i]; + } else { + remaining.push(args[i]); + } + } + return remaining; +} + +// Main CLI +const rawArgs = process.argv.slice(2); +const args = parseGlobalOptions(rawArgs); + +if (args.length === 0) { + console.log("Usage:"); + console.log(" qmd add [--drop] [glob] - Add/update collection from $PWD (default: **/*.md)"); + console.log(" qmd collections - List all collections"); + console.log(" qmd update-all - Re-index all collections"); + console.log(" qmd embed [-f] - Create vector embeddings for all content"); + console.log(" qmd search - Full-text search (BM25)"); + console.log(" qmd vsearch - Vector similarity search"); + console.log(" qmd query - Combined search with query expansion + reranking"); + console.log(""); + console.log("Global options:"); + console.log(" --index - Use custom index name (default: index)"); + console.log(""); + console.log("Search options:"); + console.log(" -n - Number of results (default: 5)"); + console.log(" --min-score - Minimum similarity score"); + console.log(" --full - Output full document instead of snippet"); + console.log(" -csv - CSV output (default is colorized CLI)"); + console.log(" -md - Markdown output"); + console.log(" -xml - XML output"); + console.log(""); + console.log("Environment:"); + console.log(" OLLAMA_URL - Ollama server URL (default: http://localhost:11434)"); + console.log(""); + console.log("Models:"); + console.log(` Embedding: ${DEFAULT_EMBED_MODEL}`); + console.log(` Reranking: ${DEFAULT_RERANK_MODEL}`); + console.log(""); + console.log(`Index: ${getDbPath()}`); + process.exit(1); +} + +const cmd = args[0]; + +if (cmd === "add") { + const addArgs = args.slice(1); + const drop = addArgs.includes("--drop"); + const globArg = addArgs.find(a => !a.startsWith("-")); + // Treat "." as "use default glob in current directory" + const globPattern = (!globArg || globArg === ".") ? DEFAULT_GLOB : globArg; + + if (drop) { + await dropCollection(globPattern); + } else { + await indexFiles(globPattern); + } +} else if (cmd === "collections") { + listCollections(); +} else if (cmd === "update-all") { + await updateAllCollections(); +} else if (cmd === "embed") { + const embedArgs = args.slice(1); + const force = embedArgs.includes("-f") || embedArgs.includes("--force"); + await vectorIndex(DEFAULT_EMBED_MODEL, force); +} else if (cmd === "search") { + const { opts, query } = parseOptions(args.slice(1), 0); + if (!query) { + console.error("Usage: qmd search [-n num] [--min-score num] [--full] [-csv|-md|-xml] "); + process.exit(1); + } + search(query, opts); +} else if (cmd === "vsearch") { + const { opts, query } = parseOptions(args.slice(1), 0.3); + if (!query) { + console.error("Usage: qmd vsearch [-n num] [--min-score num] [--full] [-csv|-md|-xml] "); + process.exit(1); + } + await vectorSearch(query, opts); +} else if (cmd === "query") { + const { opts, query } = parseOptions(args.slice(1), 0); + if (!query) { + console.error("Usage: qmd query [-n num] [--min-score num] [--full] [-csv|-md|-xml] "); + process.exit(1); + } + await querySearch(query, opts); +} else { + console.error(`Unknown command: ${cmd}`); + console.error("Run 'qmd' without arguments for usage."); + process.exit(1); +} diff --git a/tsconfig.json b/tsconfig.json new file mode 100644 index 0000000..bfa0fea --- /dev/null +++ b/tsconfig.json @@ -0,0 +1,29 @@ +{ + "compilerOptions": { + // Environment setup & latest features + "lib": ["ESNext"], + "target": "ESNext", + "module": "Preserve", + "moduleDetection": "force", + "jsx": "react-jsx", + "allowJs": true, + + // Bundler mode + "moduleResolution": "bundler", + "allowImportingTsExtensions": true, + "verbatimModuleSyntax": true, + "noEmit": true, + + // Best practices + "strict": true, + "skipLibCheck": true, + "noFallthroughCasesInSwitch": true, + "noUncheckedIndexedAccess": true, + "noImplicitOverride": true, + + // Some stricter flags (disabled by default) + "noUnusedLocals": false, + "noUnusedParameters": false, + "noPropertyAccessFromIndexSignature": false + } +}