voyageai-cli 1.16.0 → 1.18.0

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
package/package.json CHANGED
@@ -1,6 +1,6 @@
1
1
  {
2
2
  "name": "voyageai-cli",
3
- "version": "1.16.0",
3
+ "version": "1.18.0",
4
4
  "description": "CLI for Voyage AI embeddings, reranking, and MongoDB Atlas Vector Search",
5
5
  "bin": {
6
6
  "vai": "./src/cli.js"
package/src/cli.js CHANGED
@@ -23,6 +23,8 @@ const { registerBenchmark } = require('./commands/benchmark');
23
23
  const { registerEstimate } = require('./commands/estimate');
24
24
  const { registerInit } = require('./commands/init');
25
25
  const { registerChunk } = require('./commands/chunk');
26
+ const { registerQuery } = require('./commands/query');
27
+ const { registerPipeline } = require('./commands/pipeline');
26
28
  const { registerAbout } = require('./commands/about');
27
29
  const { showBanner, showQuickStart, getVersion } = require('./lib/banner');
28
30
 
@@ -51,6 +53,8 @@ registerBenchmark(program);
51
53
  registerEstimate(program);
52
54
  registerInit(program);
53
55
  registerChunk(program);
56
+ registerQuery(program);
57
+ registerPipeline(program);
54
58
  registerAbout(program);
55
59
 
56
60
  // Append disclaimer to all help output
@@ -19,7 +19,7 @@ _vai_completions() {
19
19
  prev="\${COMP_WORDS[COMP_CWORD-1]}"
20
20
 
21
21
  # Top-level commands
22
- commands="embed rerank store search index models ping config demo explain similarity ingest estimate init chunk completions help"
22
+ commands="embed rerank store search index models ping config demo explain similarity ingest estimate init chunk query pipeline completions help"
23
23
 
24
24
  # Subcommands
25
25
  local index_subs="create list delete"
@@ -114,6 +114,14 @@ _vai_completions() {
114
114
  COMPREPLY=( \$(compgen -W "--strategy --chunk-size --overlap --min-size --output --text-field --extensions --ignore --dry-run --stats --json --quiet --help" -- "\$cur") )
115
115
  return 0
116
116
  ;;
117
+ query)
118
+ COMPREPLY=( \$(compgen -W "--db --collection --index --field --model --dimensions --limit --top-k --rerank --no-rerank --rerank-model --text-field --filter --num-candidates --show-vectors --json --quiet --help" -- "\$cur") )
119
+ return 0
120
+ ;;
121
+ pipeline)
122
+ COMPREPLY=( \$(compgen -W "--db --collection --field --index --model --dimensions --strategy --chunk-size --overlap --batch-size --text-field --extensions --ignore --create-index --dry-run --json --quiet --help" -- "\$cur") )
123
+ return 0
124
+ ;;
117
125
  completions)
118
126
  COMPREPLY=( \$(compgen -W "bash zsh --help" -- "\$cur") )
119
127
  return 0
@@ -187,6 +195,8 @@ _vai() {
187
195
  'estimate:Estimate embedding costs — symmetric vs asymmetric'
188
196
  'init:Initialize project with .vai.json'
189
197
  'chunk:Chunk documents for embedding'
198
+ 'query:Search + rerank in one shot'
199
+ 'pipeline:Chunk, embed, and store documents'
190
200
  'completions:Generate shell completion scripts'
191
201
  'help:Display help for command'
192
202
  )
@@ -425,6 +435,46 @@ _vai() {
425
435
  '--json[JSON output]' \\
426
436
  '(-q --quiet)'{-q,--quiet}'[Suppress non-essential output]'
427
437
  ;;
438
+ query)
439
+ _arguments \\
440
+ '1:query text:' \\
441
+ '--db[Database name]:database:' \\
442
+ '--collection[Collection name]:collection:' \\
443
+ '--index[Vector search index]:index:' \\
444
+ '--field[Embedding field]:field:' \\
445
+ '(-m --model)'{-m,--model}'[Embedding model]:model:(\$models)' \\
446
+ '(-d --dimensions)'{-d,--dimensions}'[Output dimensions]:dims:' \\
447
+ '(-l --limit)'{-l,--limit}'[Search candidates]:limit:' \\
448
+ '(-k --top-k)'{-k,--top-k}'[Final results]:k:' \\
449
+ '--rerank[Enable reranking]' \\
450
+ '--no-rerank[Skip reranking]' \\
451
+ '--rerank-model[Reranking model]:model:' \\
452
+ '--text-field[Document text field]:field:' \\
453
+ '--filter[Pre-filter JSON]:json:' \\
454
+ '--json[JSON output]' \\
455
+ '(-q --quiet)'{-q,--quiet}'[Suppress non-essential output]'
456
+ ;;
457
+ pipeline)
458
+ _arguments \\
459
+ '1:input:_files' \\
460
+ '--db[Database name]:database:' \\
461
+ '--collection[Collection name]:collection:' \\
462
+ '--field[Embedding field]:field:' \\
463
+ '--index[Vector search index]:index:' \\
464
+ '(-m --model)'{-m,--model}'[Embedding model]:model:(\$models)' \\
465
+ '(-d --dimensions)'{-d,--dimensions}'[Output dimensions]:dims:' \\
466
+ '(-s --strategy)'{-s,--strategy}'[Chunking strategy]:strategy:(fixed sentence paragraph recursive markdown)' \\
467
+ '(-c --chunk-size)'{-c,--chunk-size}'[Chunk size]:size:' \\
468
+ '--overlap[Chunk overlap]:chars:' \\
469
+ '--batch-size[Texts per API call]:size:' \\
470
+ '--text-field[Text field for JSON]:field:' \\
471
+ '--extensions[File extensions]:exts:' \\
472
+ '--ignore[Dirs to skip]:dirs:' \\
473
+ '--create-index[Auto-create vector index]' \\
474
+ '--dry-run[Preview without executing]' \\
475
+ '--json[JSON output]' \\
476
+ '(-q --quiet)'{-q,--quiet}'[Suppress non-essential output]'
477
+ ;;
428
478
  completions)
429
479
  _arguments \\
430
480
  '1:shell:(bash zsh)'
@@ -0,0 +1,311 @@
1
+ 'use strict';
2
+
3
+ const fs = require('fs');
4
+ const path = require('path');
5
+ const { chunk, estimateTokens, STRATEGIES } = require('../lib/chunker');
6
+ const { readFile, scanDirectory, isSupported, getReaderType } = require('../lib/readers');
7
+ const { loadProject } = require('../lib/project');
8
+ const { getDefaultModel } = require('../lib/catalog');
9
+ const { generateEmbeddings } = require('../lib/api');
10
+ const { getMongoCollection } = require('../lib/mongo');
11
+ const ui = require('../lib/ui');
12
+
13
+ /**
14
+ * Format number with commas.
15
+ */
16
+ function fmtNum(n) {
17
+ return n.toLocaleString('en-US');
18
+ }
19
+
20
+ /**
21
+ * Resolve input path(s) to file list.
22
+ */
23
+ function resolveFiles(input, opts) {
24
+ const resolved = path.resolve(input);
25
+ if (!fs.existsSync(resolved)) {
26
+ throw new Error(`Not found: ${input}`);
27
+ }
28
+
29
+ const stat = fs.statSync(resolved);
30
+ if (stat.isFile()) return [resolved];
31
+
32
+ if (stat.isDirectory()) {
33
+ const scanOpts = {};
34
+ if (opts.extensions) scanOpts.extensions = opts.extensions.split(',').map(e => e.trim());
35
+ if (opts.ignore) scanOpts.ignore = opts.ignore.split(',').map(d => d.trim());
36
+ return scanDirectory(resolved, scanOpts);
37
+ }
38
+
39
+ return [];
40
+ }
41
+
42
+ /**
43
+ * Register the pipeline command on a Commander program.
44
+ * @param {import('commander').Command} program
45
+ */
46
+ function registerPipeline(program) {
47
+ program
48
+ .command('pipeline <input>')
49
+ .description('End-to-end: chunk → embed → store in MongoDB Atlas')
50
+ .option('--db <database>', 'Database name')
51
+ .option('--collection <name>', 'Collection name')
52
+ .option('--field <name>', 'Embedding field name')
53
+ .option('--index <name>', 'Vector search index name')
54
+ .option('-m, --model <model>', 'Embedding model')
55
+ .option('-d, --dimensions <n>', 'Output dimensions', (v) => parseInt(v, 10))
56
+ .option('-s, --strategy <strategy>', 'Chunking strategy')
57
+ .option('-c, --chunk-size <n>', 'Target chunk size in characters', (v) => parseInt(v, 10))
58
+ .option('--overlap <n>', 'Overlap between chunks', (v) => parseInt(v, 10))
59
+ .option('--batch-size <n>', 'Texts per embedding API call', (v) => parseInt(v, 10), 25)
60
+ .option('--text-field <name>', 'Text field for JSON/JSONL input', 'text')
61
+ .option('--extensions <exts>', 'File extensions to include')
62
+ .option('--ignore <dirs>', 'Directory names to skip', 'node_modules,.git,__pycache__')
63
+ .option('--create-index', 'Auto-create vector search index if it doesn\'t exist')
64
+ .option('--dry-run', 'Show what would happen without executing')
65
+ .option('--json', 'Machine-readable JSON output')
66
+ .option('-q, --quiet', 'Suppress non-essential output')
67
+ .action(async (input, opts) => {
68
+ let client;
69
+ try {
70
+ // Merge project config
71
+ const { config: proj } = loadProject();
72
+ const projChunk = proj.chunk || {};
73
+
74
+ const db = opts.db || proj.db;
75
+ const collection = opts.collection || proj.collection;
76
+ const field = opts.field || proj.field || 'embedding';
77
+ const index = opts.index || proj.index || 'vector_index';
78
+ const model = opts.model || proj.model || getDefaultModel();
79
+ const dimensions = opts.dimensions || proj.dimensions;
80
+ const strategy = opts.strategy || projChunk.strategy || 'recursive';
81
+ const chunkSize = opts.chunkSize || projChunk.size || 512;
82
+ const overlap = opts.overlap != null ? opts.overlap : (projChunk.overlap != null ? projChunk.overlap : 50);
83
+ const batchSize = opts.batchSize || 25;
84
+ const textField = opts.textField || 'text';
85
+
86
+ if (!db || !collection) {
87
+ console.error(ui.error('Database and collection required. Use --db/--collection or "vai init".'));
88
+ process.exit(1);
89
+ }
90
+
91
+ if (!STRATEGIES.includes(strategy)) {
92
+ console.error(ui.error(`Unknown strategy: "${strategy}". Available: ${STRATEGIES.join(', ')}`));
93
+ process.exit(1);
94
+ }
95
+
96
+ // Step 1: Resolve files
97
+ const files = resolveFiles(input, opts);
98
+ if (files.length === 0) {
99
+ console.error(ui.error('No supported files found.'));
100
+ process.exit(1);
101
+ }
102
+
103
+ const basePath = fs.statSync(path.resolve(input)).isDirectory()
104
+ ? path.resolve(input)
105
+ : process.cwd();
106
+
107
+ const verbose = !opts.json && !opts.quiet;
108
+
109
+ if (verbose) {
110
+ console.log('');
111
+ console.log(ui.bold('🚀 Pipeline: chunk → embed → store'));
112
+ console.log(ui.dim(` Files: ${files.length} | Strategy: ${strategy} | Model: ${model}`));
113
+ console.log(ui.dim(` Target: ${db}.${collection} (field: ${field})`));
114
+ console.log('');
115
+ }
116
+
117
+ // Step 2: Chunk all files
118
+ if (verbose) console.log(ui.bold('Step 1/3 — Chunking'));
119
+
120
+ const allChunks = [];
121
+ let totalInputChars = 0;
122
+ const fileErrors = [];
123
+
124
+ for (const filePath of files) {
125
+ const relPath = path.relative(basePath, filePath);
126
+ try {
127
+ const content = await readFile(filePath, { textField });
128
+ const texts = typeof content === 'string'
129
+ ? [{ text: content, metadata: {} }]
130
+ : content;
131
+
132
+ for (const item of texts) {
133
+ const useStrategy = (strategy === 'recursive' && filePath.endsWith('.md'))
134
+ ? 'markdown' : strategy;
135
+
136
+ const chunks = chunk(item.text, {
137
+ strategy: useStrategy,
138
+ size: chunkSize,
139
+ overlap,
140
+ });
141
+
142
+ totalInputChars += item.text.length;
143
+
144
+ for (let ci = 0; ci < chunks.length; ci++) {
145
+ allChunks.push({
146
+ text: chunks[ci],
147
+ metadata: {
148
+ ...item.metadata,
149
+ source: relPath,
150
+ chunk_index: ci,
151
+ total_chunks: chunks.length,
152
+ },
153
+ });
154
+ }
155
+ }
156
+
157
+ if (verbose) console.log(` ${ui.green('✓')} ${relPath} → ${allChunks.length} chunks total`);
158
+ } catch (err) {
159
+ fileErrors.push({ file: relPath, error: err.message });
160
+ if (verbose) console.error(` ${ui.red('✗')} ${relPath}: ${err.message}`);
161
+ }
162
+ }
163
+
164
+ if (allChunks.length === 0) {
165
+ console.error(ui.error('No chunks produced. Check your files and chunk settings.'));
166
+ process.exit(1);
167
+ }
168
+
169
+ const totalTokens = allChunks.reduce((sum, c) => sum + estimateTokens(c.text), 0);
170
+
171
+ if (verbose) {
172
+ console.log(ui.dim(` ${fmtNum(allChunks.length)} chunks, ~${fmtNum(totalTokens)} tokens`));
173
+ console.log('');
174
+ }
175
+
176
+ // Dry run — stop here
177
+ if (opts.dryRun) {
178
+ if (opts.json) {
179
+ console.log(JSON.stringify({
180
+ dryRun: true,
181
+ files: files.length,
182
+ chunks: allChunks.length,
183
+ estimatedTokens: totalTokens,
184
+ strategy, chunkSize, overlap, model, db, collection, field,
185
+ }, null, 2));
186
+ } else {
187
+ console.log(ui.success(`Dry run complete: ${fmtNum(allChunks.length)} chunks from ${files.length} files.`));
188
+ const cost = (totalTokens / 1e6) * 0.12;
189
+ console.log(ui.dim(` Estimated embedding cost: ~$${cost.toFixed(4)} with ${model}`));
190
+ }
191
+ return;
192
+ }
193
+
194
+ // Step 3: Embed in batches
195
+ if (verbose) console.log(ui.bold('Step 2/3 — Embedding'));
196
+
197
+ const batches = [];
198
+ for (let i = 0; i < allChunks.length; i += batchSize) {
199
+ batches.push(allChunks.slice(i, i + batchSize));
200
+ }
201
+
202
+ let embeddedCount = 0;
203
+ let totalApiTokens = 0;
204
+ const embeddings = new Array(allChunks.length);
205
+
206
+ for (let bi = 0; bi < batches.length; bi++) {
207
+ const batch = batches[bi];
208
+ const texts = batch.map(c => c.text);
209
+
210
+ if (verbose) {
211
+ const pct = Math.round(((bi + 1) / batches.length) * 100);
212
+ process.stderr.write(`\r Batch ${bi + 1}/${batches.length} (${pct}%)...`);
213
+ }
214
+
215
+ const embedOpts = { model, inputType: 'document' };
216
+ if (dimensions) embedOpts.dimensions = dimensions;
217
+
218
+ const result = await generateEmbeddings(texts, embedOpts);
219
+ totalApiTokens += result.usage?.total_tokens || 0;
220
+
221
+ for (let j = 0; j < result.data.length; j++) {
222
+ embeddings[embeddedCount + j] = result.data[j].embedding;
223
+ }
224
+ embeddedCount += batch.length;
225
+ }
226
+
227
+ if (verbose) {
228
+ process.stderr.write('\r');
229
+ console.log(` ${ui.green('✓')} Embedded ${fmtNum(embeddedCount)} chunks (${fmtNum(totalApiTokens)} tokens)`);
230
+ console.log('');
231
+ }
232
+
233
+ // Step 4: Store in MongoDB
234
+ if (verbose) console.log(ui.bold('Step 3/3 — Storing in MongoDB'));
235
+
236
+ const { client: c, collection: coll } = await getMongoCollection(db, collection);
237
+ client = c;
238
+
239
+ const documents = allChunks.map((chunk, i) => ({
240
+ text: chunk.text,
241
+ [field]: embeddings[i],
242
+ metadata: chunk.metadata,
243
+ _model: model,
244
+ _embeddedAt: new Date(),
245
+ }));
246
+
247
+ const insertResult = await coll.insertMany(documents);
248
+
249
+ if (verbose) {
250
+ console.log(` ${ui.green('✓')} Inserted ${fmtNum(insertResult.insertedCount)} documents`);
251
+ }
252
+
253
+ // Optional: create index
254
+ if (opts.createIndex) {
255
+ if (verbose) console.log('');
256
+ try {
257
+ const dim = embeddings[0]?.length || dimensions || 1024;
258
+ const indexDef = {
259
+ name: index,
260
+ type: 'vectorSearch',
261
+ definition: {
262
+ fields: [{
263
+ type: 'vector',
264
+ path: field,
265
+ numDimensions: dim,
266
+ similarity: 'cosine',
267
+ }],
268
+ },
269
+ };
270
+ await coll.createSearchIndex(indexDef);
271
+ if (verbose) console.log(` ${ui.green('✓')} Created vector index "${index}" (${dim} dims, cosine)`);
272
+ } catch (err) {
273
+ if (err.message?.includes('already exists')) {
274
+ if (verbose) console.log(` ${ui.dim('ℹ Index "' + index + '" already exists — skipping')}`);
275
+ } else {
276
+ if (verbose) console.error(` ${ui.yellow('⚠')} Index creation failed: ${err.message}`);
277
+ }
278
+ }
279
+ }
280
+
281
+ // Summary
282
+ if (opts.json) {
283
+ console.log(JSON.stringify({
284
+ files: files.length,
285
+ fileErrors: fileErrors.length,
286
+ chunks: allChunks.length,
287
+ tokens: totalApiTokens,
288
+ inserted: insertResult.insertedCount,
289
+ model, db, collection, field, strategy, chunkSize,
290
+ index: opts.createIndex ? index : null,
291
+ }, null, 2));
292
+ } else if (verbose) {
293
+ console.log('');
294
+ console.log(ui.success('Pipeline complete'));
295
+ console.log(ui.label('Files', `${fmtNum(files.length)}${fileErrors.length ? ` (${fileErrors.length} failed)` : ''}`));
296
+ console.log(ui.label('Chunks', fmtNum(allChunks.length)));
297
+ console.log(ui.label('Tokens', fmtNum(totalApiTokens)));
298
+ console.log(ui.label('Stored', `${fmtNum(insertResult.insertedCount)} docs → ${db}.${collection}`));
299
+ console.log('');
300
+ console.log(ui.dim(' Next: vai query "your search" --db ' + db + ' --collection ' + collection));
301
+ }
302
+ } catch (err) {
303
+ console.error(ui.error(err.message));
304
+ process.exit(1);
305
+ } finally {
306
+ if (client) await client.close();
307
+ }
308
+ });
309
+ }
310
+
311
+ module.exports = { registerPipeline };
@@ -0,0 +1,266 @@
1
+ 'use strict';
2
+
3
+ const { getDefaultModel, DEFAULT_RERANK_MODEL } = require('../lib/catalog');
4
+ const { generateEmbeddings, apiRequest } = require('../lib/api');
5
+ const { getMongoCollection } = require('../lib/mongo');
6
+ const { loadProject } = require('../lib/project');
7
+ const ui = require('../lib/ui');
8
+
9
+ /**
10
+ * Register the query command on a Commander program.
11
+ * @param {import('commander').Command} program
12
+ */
13
+ function registerQuery(program) {
14
+ program
15
+ .command('query <text>')
16
+ .description('Search + rerank in one shot — the two-stage retrieval pattern')
17
+ .option('--db <database>', 'Database name')
18
+ .option('--collection <name>', 'Collection name')
19
+ .option('--index <name>', 'Vector search index name')
20
+ .option('--field <name>', 'Embedding field name')
21
+ .option('-m, --model <model>', 'Embedding model for query')
22
+ .option('-d, --dimensions <n>', 'Output dimensions', (v) => parseInt(v, 10))
23
+ .option('-l, --limit <n>', 'Number of vector search candidates', (v) => parseInt(v, 10), 20)
24
+ .option('-k, --top-k <n>', 'Final results to return (after rerank)', (v) => parseInt(v, 10), 5)
25
+ .option('--rerank', 'Enable reranking (recommended)')
26
+ .option('--no-rerank', 'Skip reranking — vector search only')
27
+ .option('--rerank-model <model>', 'Reranking model')
28
+ .option('--text-field <name>', 'Document text field for reranking and display', 'text')
29
+ .option('--filter <json>', 'Pre-filter JSON for $vectorSearch')
30
+ .option('--num-candidates <n>', 'ANN candidates (default: limit × 15)', (v) => parseInt(v, 10))
31
+ .option('--show-vectors', 'Include embedding vectors in output')
32
+ .option('--json', 'Machine-readable JSON output')
33
+ .option('-q, --quiet', 'Suppress non-essential output')
34
+ .action(async (text, opts) => {
35
+ let client;
36
+ try {
37
+ // Merge project config
38
+ const { config: proj } = loadProject();
39
+ const db = opts.db || proj.db;
40
+ const collection = opts.collection || proj.collection;
41
+ const index = opts.index || proj.index || 'vector_index';
42
+ const field = opts.field || proj.field || 'embedding';
43
+ const model = opts.model || proj.model || getDefaultModel();
44
+ const rerankModel = opts.rerankModel || DEFAULT_RERANK_MODEL;
45
+ const textField = opts.textField || 'text';
46
+ const dimensions = opts.dimensions || proj.dimensions;
47
+ const doRerank = opts.rerank !== false;
48
+
49
+ if (!db || !collection) {
50
+ console.error(ui.error('Database and collection required. Use --db and --collection, or create .vai.json with "vai init".'));
51
+ process.exit(1);
52
+ }
53
+
54
+ const useColor = !opts.json;
55
+ const useSpinner = useColor && !opts.quiet;
56
+
57
+ // Step 1: Embed query
58
+ let spin;
59
+ if (useSpinner) {
60
+ spin = ui.spinner('Embedding query...');
61
+ spin.start();
62
+ }
63
+
64
+ const embedOpts = { model, inputType: 'query' };
65
+ if (dimensions) embedOpts.dimensions = dimensions;
66
+ const embedResult = await generateEmbeddings([text], embedOpts);
67
+ const queryVector = embedResult.data[0].embedding;
68
+ const embedTokens = embedResult.usage?.total_tokens || 0;
69
+
70
+ if (spin) spin.stop();
71
+
72
+ // Step 2: Vector search
73
+ if (useSpinner) {
74
+ spin = ui.spinner(`Searching ${db}.${collection}...`);
75
+ spin.start();
76
+ }
77
+
78
+ const { client: c, coll } = await connectCollection(db, collection);
79
+ client = c;
80
+
81
+ const numCandidates = opts.numCandidates || Math.min(opts.limit * 15, 10000);
82
+ const vectorSearchStage = {
83
+ index,
84
+ path: field,
85
+ queryVector,
86
+ numCandidates,
87
+ limit: opts.limit,
88
+ };
89
+
90
+ if (opts.filter) {
91
+ try {
92
+ vectorSearchStage.filter = JSON.parse(opts.filter);
93
+ } catch {
94
+ if (spin) spin.stop();
95
+ console.error(ui.error('Invalid --filter JSON.'));
96
+ process.exit(1);
97
+ }
98
+ }
99
+
100
+ const pipeline = [
101
+ { $vectorSearch: vectorSearchStage },
102
+ { $addFields: { _vsScore: { $meta: 'vectorSearchScore' } } },
103
+ ];
104
+
105
+ const searchResults = await coll.aggregate(pipeline).toArray();
106
+ if (spin) spin.stop();
107
+
108
+ if (searchResults.length === 0) {
109
+ if (opts.json) {
110
+ console.log(JSON.stringify({ query: text, results: [], stages: { search: 0, rerank: 0 } }, null, 2));
111
+ } else {
112
+ console.log(ui.yellow('No results found.'));
113
+ }
114
+ return;
115
+ }
116
+
117
+ // Step 3: Rerank (optional)
118
+ let finalResults;
119
+ let rerankTokens = 0;
120
+
121
+ if (doRerank && searchResults.length > 1) {
122
+ if (useSpinner) {
123
+ spin = ui.spinner(`Reranking ${searchResults.length} results...`);
124
+ spin.start();
125
+ }
126
+
127
+ // Extract text for reranking
128
+ const documents = searchResults.map(doc => {
129
+ const txt = doc[textField];
130
+ if (!txt) return JSON.stringify(doc);
131
+ return typeof txt === 'string' ? txt : JSON.stringify(txt);
132
+ });
133
+
134
+ const rerankBody = {
135
+ query: text,
136
+ documents,
137
+ model: rerankModel,
138
+ top_k: opts.topK,
139
+ };
140
+
141
+ const rerankResult = await apiRequest('/rerank', rerankBody);
142
+ rerankTokens = rerankResult.usage?.total_tokens || 0;
143
+
144
+ if (spin) spin.stop();
145
+
146
+ // Map reranked indices back to original docs
147
+ finalResults = (rerankResult.data || []).map(item => {
148
+ const doc = searchResults[item.index];
149
+ return {
150
+ ...doc,
151
+ _vsScore: doc._vsScore,
152
+ _rerankScore: item.relevance_score,
153
+ _finalScore: item.relevance_score,
154
+ };
155
+ });
156
+ } else {
157
+ // No rerank — just take top-k from vector search
158
+ finalResults = searchResults.slice(0, opts.topK).map(doc => ({
159
+ ...doc,
160
+ _finalScore: doc._vsScore,
161
+ }));
162
+ }
163
+
164
+ // Build output
165
+ const output = finalResults.map((doc, i) => {
166
+ const clean = {};
167
+ // Include key fields
168
+ if (doc._id) clean._id = doc._id;
169
+ if (doc[textField]) {
170
+ clean[textField] = doc[textField];
171
+ }
172
+ // Include metadata fields (skip embedding and internal scores)
173
+ for (const key of Object.keys(doc)) {
174
+ if (key === field || key === '_vsScore' || key === '_rerankScore' || key === '_finalScore') continue;
175
+ if (key === '_id' || key === textField) continue;
176
+ if (!opts.showVectors) clean[key] = doc[key];
177
+ else clean[key] = doc[key];
178
+ }
179
+ // Scores
180
+ clean.score = doc._finalScore;
181
+ if (doc._vsScore !== undefined) clean.vectorScore = doc._vsScore;
182
+ if (doc._rerankScore !== undefined) clean.rerankScore = doc._rerankScore;
183
+ clean.rank = i + 1;
184
+ return clean;
185
+ });
186
+
187
+ if (opts.json) {
188
+ console.log(JSON.stringify({
189
+ query: text,
190
+ model,
191
+ rerankModel: doRerank ? rerankModel : null,
192
+ db,
193
+ collection,
194
+ stages: {
195
+ searchCandidates: searchResults.length,
196
+ finalResults: output.length,
197
+ reranked: doRerank && searchResults.length > 1,
198
+ },
199
+ tokens: { embed: embedTokens, rerank: rerankTokens },
200
+ results: output,
201
+ }, null, 2));
202
+ return;
203
+ }
204
+
205
+ // Pretty output
206
+ if (!opts.quiet) {
207
+ console.log('');
208
+ console.log(ui.label('Query', ui.cyan(`"${text}"`)));
209
+ console.log(ui.label('Search', `${searchResults.length} candidates from ${ui.dim(`${db}.${collection}`)}`));
210
+ if (doRerank && searchResults.length > 1) {
211
+ console.log(ui.label('Rerank', `Top ${output.length} via ${ui.dim(rerankModel)}`));
212
+ }
213
+ console.log(ui.label('Model', ui.dim(model)));
214
+ console.log('');
215
+ }
216
+
217
+ for (let i = 0; i < output.length; i++) {
218
+ const r = output[i];
219
+ const scoreStr = r.score != null ? ui.score(r.score) : 'N/A';
220
+ const vsStr = r.vectorScore != null ? ui.dim(`vs:${r.vectorScore.toFixed(3)}`) : '';
221
+ const rrStr = r.rerankScore != null ? ui.dim(`rr:${r.rerankScore.toFixed(3)}`) : '';
222
+ const scores = [vsStr, rrStr].filter(Boolean).join(' ');
223
+
224
+ console.log(`${ui.bold(`#${i + 1}`)} ${scoreStr} ${scores}`);
225
+
226
+ // Show text preview
227
+ const textVal = r[textField];
228
+ if (textVal) {
229
+ const preview = textVal.substring(0, 200);
230
+ const ellipsis = textVal.length > 200 ? '...' : '';
231
+ console.log(` ${preview}${ellipsis}`);
232
+ }
233
+
234
+ // Show source metadata if present
235
+ if (r.source) console.log(` ${ui.dim('source: ' + r.source)}`);
236
+ if (r.metadata?.source) console.log(` ${ui.dim('source: ' + r.metadata.source)}`);
237
+
238
+ console.log(` ${ui.dim('_id: ' + r._id)}`);
239
+ console.log('');
240
+ }
241
+
242
+ if (!opts.quiet) {
243
+ const totalTokens = embedTokens + rerankTokens;
244
+ console.log(ui.dim(` Tokens: ${totalTokens} (embed: ${embedTokens}${rerankTokens ? `, rerank: ${rerankTokens}` : ''})`));
245
+ }
246
+ } catch (err) {
247
+ console.error(ui.error(err.message));
248
+ process.exit(1);
249
+ } finally {
250
+ if (client) await client.close();
251
+ }
252
+ });
253
+ }
254
+
255
+ /**
256
+ * Connect to a MongoDB collection.
257
+ * @param {string} db
258
+ * @param {string} collName
259
+ * @returns {Promise<{client: MongoClient, coll: Collection}>}
260
+ */
261
+ async function connectCollection(db, collName) {
262
+ const { client, collection } = await getMongoCollection(db, collName);
263
+ return { client, coll: collection };
264
+ }
265
+
266
+ module.exports = { registerQuery };