voyageai-cli 1.15.0 → 1.18.0

package/src/commands/pipeline.js

@@ -0,0 +1,311 @@
+'use strict';
+
+const fs = require('fs');
+const path = require('path');
+const { chunk, estimateTokens, STRATEGIES } = require('../lib/chunker');
+const { readFile, scanDirectory, isSupported, getReaderType } = require('../lib/readers');
+const { loadProject } = require('../lib/project');
+const { getDefaultModel } = require('../lib/catalog');
+const { generateEmbeddings } = require('../lib/api');
+const { getMongoCollection } = require('../lib/mongo');
+const ui = require('../lib/ui');
+
+/**
+ * Format number with commas.
+ */
+function fmtNum(n) {
+  return n.toLocaleString('en-US');
+}
+
+/**
+ * Resolve input path(s) to file list.
+ */
+function resolveFiles(input, opts) {
+  const resolved = path.resolve(input);
+  if (!fs.existsSync(resolved)) {
+    throw new Error(`Not found: ${input}`);
+  }
+
+  const stat = fs.statSync(resolved);
+  if (stat.isFile()) return [resolved];
+
+  if (stat.isDirectory()) {
+    const scanOpts = {};
+    if (opts.extensions) scanOpts.extensions = opts.extensions.split(',').map(e => e.trim());
+    if (opts.ignore) scanOpts.ignore = opts.ignore.split(',').map(d => d.trim());
+    return scanDirectory(resolved, scanOpts);
+  }
+
+  return [];
+}
+
+/**
+ * Register the pipeline command on a Commander program.
+ * @param {import('commander').Command} program
+ */
+function registerPipeline(program) {
+  program
+    .command('pipeline <input>')
+    .description('End-to-end: chunk → embed → store in MongoDB Atlas')
+    .option('--db <database>', 'Database name')
+    .option('--collection <name>', 'Collection name')
+    .option('--field <name>', 'Embedding field name')
+    .option('--index <name>', 'Vector search index name')
+    .option('-m, --model <model>', 'Embedding model')
+    .option('-d, --dimensions <n>', 'Output dimensions', (v) => parseInt(v, 10))
+    .option('-s, --strategy <strategy>', 'Chunking strategy')
+    .option('-c, --chunk-size <n>', 'Target chunk size in characters', (v) => parseInt(v, 10))
+    .option('--overlap <n>', 'Overlap between chunks', (v) => parseInt(v, 10))
+    .option('--batch-size <n>', 'Texts per embedding API call', (v) => parseInt(v, 10), 25)
+    .option('--text-field <name>', 'Text field for JSON/JSONL input', 'text')
+    .option('--extensions <exts>', 'File extensions to include')
+    .option('--ignore <dirs>', 'Directory names to skip', 'node_modules,.git,__pycache__')
+    .option('--create-index', 'Auto-create vector search index if it doesn\'t exist')
+    .option('--dry-run', 'Show what would happen without executing')
+    .option('--json', 'Machine-readable JSON output')
+    .option('-q, --quiet', 'Suppress non-essential output')
+    .action(async (input, opts) => {
+      let client;
+      try {
+        // Merge project config
+        const { config: proj } = loadProject();
+        const projChunk = proj.chunk || {};
+
+        const db = opts.db || proj.db;
+        const collection = opts.collection || proj.collection;
+        const field = opts.field || proj.field || 'embedding';
+        const index = opts.index || proj.index || 'vector_index';
+        const model = opts.model || proj.model || getDefaultModel();
+        const dimensions = opts.dimensions || proj.dimensions;
+        const strategy = opts.strategy || projChunk.strategy || 'recursive';
+        const chunkSize = opts.chunkSize || projChunk.size || 512;
+        const overlap = opts.overlap != null ? opts.overlap : (projChunk.overlap != null ? projChunk.overlap : 50);
+        const batchSize = opts.batchSize || 25;
+        const textField = opts.textField || 'text';
+
+        if (!db || !collection) {
+          console.error(ui.error('Database and collection required. Use --db/--collection or "vai init".'));
+          process.exit(1);
+        }
+
+        if (!STRATEGIES.includes(strategy)) {
+          console.error(ui.error(`Unknown strategy: "${strategy}". Available: ${STRATEGIES.join(', ')}`));
+          process.exit(1);
+        }
+
+        // Step 1: Resolve files
+        const files = resolveFiles(input, opts);
+        if (files.length === 0) {
+          console.error(ui.error('No supported files found.'));
+          process.exit(1);
+        }
+
+        const basePath = fs.statSync(path.resolve(input)).isDirectory()
+          ? path.resolve(input)
+          : process.cwd();
+
+        const verbose = !opts.json && !opts.quiet;
+
+        if (verbose) {
+          console.log('');
+          console.log(ui.bold('🚀 Pipeline: chunk → embed → store'));
+          console.log(ui.dim(`  Files: ${files.length} | Strategy: ${strategy} | Model: ${model}`));
+          console.log(ui.dim(`  Target: ${db}.${collection} (field: ${field})`));
+          console.log('');
+        }
+
+        // Step 2: Chunk all files
+        if (verbose) console.log(ui.bold('Step 1/3 — Chunking'));
+
+        const allChunks = [];
+        let totalInputChars = 0;
+        const fileErrors = [];
+
+        for (const filePath of files) {
+          const relPath = path.relative(basePath, filePath);
+          try {
+            const content = await readFile(filePath, { textField });
+            const texts = typeof content === 'string'
+              ? [{ text: content, metadata: {} }]
+              : content;
+
+            for (const item of texts) {
+              const useStrategy = (strategy === 'recursive' && filePath.endsWith('.md'))
+                ? 'markdown' : strategy;
+
+              const chunks = chunk(item.text, {
+                strategy: useStrategy,
+                size: chunkSize,
+                overlap,
+              });
+
+              totalInputChars += item.text.length;
+
+              for (let ci = 0; ci < chunks.length; ci++) {
+                allChunks.push({
+                  text: chunks[ci],
+                  metadata: {
+                    ...item.metadata,
+                    source: relPath,
+                    chunk_index: ci,
+                    total_chunks: chunks.length,
+                  },
+                });
+              }
+            }
+
+            if (verbose) console.log(`  ${ui.green('✓')} ${relPath} → ${allChunks.length} chunks total`);
+          } catch (err) {
+            fileErrors.push({ file: relPath, error: err.message });
+            if (verbose) console.error(`  ${ui.red('✗')} ${relPath}: ${err.message}`);
+          }
+        }
+
+        if (allChunks.length === 0) {
+          console.error(ui.error('No chunks produced. Check your files and chunk settings.'));
+          process.exit(1);
+        }
+
+        const totalTokens = allChunks.reduce((sum, c) => sum + estimateTokens(c.text), 0);
+
+        if (verbose) {
+          console.log(ui.dim(`  ${fmtNum(allChunks.length)} chunks, ~${fmtNum(totalTokens)} tokens`));
+          console.log('');
+        }
+
+        // Dry run — stop here
+        if (opts.dryRun) {
+          if (opts.json) {
+            console.log(JSON.stringify({
+              dryRun: true,
+              files: files.length,
+              chunks: allChunks.length,
+              estimatedTokens: totalTokens,
+              strategy, chunkSize, overlap, model, db, collection, field,
+            }, null, 2));
+          } else {
+            console.log(ui.success(`Dry run complete: ${fmtNum(allChunks.length)} chunks from ${files.length} files.`));
+            const cost = (totalTokens / 1e6) * 0.12;
+            console.log(ui.dim(`  Estimated embedding cost: ~$${cost.toFixed(4)} with ${model}`));
+          }
+          return;
+        }
+
+        // Step 3: Embed in batches
+        if (verbose) console.log(ui.bold('Step 2/3 — Embedding'));
+
+        const batches = [];
+        for (let i = 0; i < allChunks.length; i += batchSize) {
+          batches.push(allChunks.slice(i, i + batchSize));
+        }
+
+        let embeddedCount = 0;
+        let totalApiTokens = 0;
+        const embeddings = new Array(allChunks.length);
+
+        for (let bi = 0; bi < batches.length; bi++) {
+          const batch = batches[bi];
+          const texts = batch.map(c => c.text);
+
+          if (verbose) {
+            const pct = Math.round(((bi + 1) / batches.length) * 100);
+            process.stderr.write(`\r  Batch ${bi + 1}/${batches.length} (${pct}%)...`);
+          }
+
+          const embedOpts = { model, inputType: 'document' };
+          if (dimensions) embedOpts.dimensions = dimensions;
+
+          const result = await generateEmbeddings(texts, embedOpts);
+          totalApiTokens += result.usage?.total_tokens || 0;
+
+          for (let j = 0; j < result.data.length; j++) {
+            embeddings[embeddedCount + j] = result.data[j].embedding;
+          }
+          embeddedCount += batch.length;
+        }
+
+        if (verbose) {
+          process.stderr.write('\r');
+          console.log(`  ${ui.green('✓')} Embedded ${fmtNum(embeddedCount)} chunks (${fmtNum(totalApiTokens)} tokens)`);
+          console.log('');
+        }
+
+        // Step 4: Store in MongoDB
+        if (verbose) console.log(ui.bold('Step 3/3 — Storing in MongoDB'));
+
+        const { client: c, collection: coll } = await getMongoCollection(db, collection);
+        client = c;
+
+        const documents = allChunks.map((chunk, i) => ({
+          text: chunk.text,
+          [field]: embeddings[i],
+          metadata: chunk.metadata,
+          _model: model,
+          _embeddedAt: new Date(),
+        }));
+
+        const insertResult = await coll.insertMany(documents);
+
+        if (verbose) {
+          console.log(`  ${ui.green('✓')} Inserted ${fmtNum(insertResult.insertedCount)} documents`);
+        }
+
+        // Optional: create index
+        if (opts.createIndex) {
+          if (verbose) console.log('');
+          try {
+            const dim = embeddings[0]?.length || dimensions || 1024;
+            const indexDef = {
+              name: index,
+              type: 'vectorSearch',
+              definition: {
+                fields: [{
+                  type: 'vector',
+                  path: field,
+                  numDimensions: dim,
+                  similarity: 'cosine',
+                }],
+              },
+            };
+            await coll.createSearchIndex(indexDef);
+            if (verbose) console.log(`  ${ui.green('✓')} Created vector index "${index}" (${dim} dims, cosine)`);
+          } catch (err) {
+            if (err.message?.includes('already exists')) {
+              if (verbose) console.log(`  ${ui.dim('ℹ Index "' + index + '" already exists — skipping')}`);
+            } else {
+              if (verbose) console.error(`  ${ui.yellow('⚠')} Index creation failed: ${err.message}`);
+            }
+          }
+        }
+
+        // Summary
+        if (opts.json) {
+          console.log(JSON.stringify({
+            files: files.length,
+            fileErrors: fileErrors.length,
+            chunks: allChunks.length,
+            tokens: totalApiTokens,
+            inserted: insertResult.insertedCount,
+            model, db, collection, field, strategy, chunkSize,
+            index: opts.createIndex ? index : null,
+          }, null, 2));
+        } else if (verbose) {
+          console.log('');
+          console.log(ui.success('Pipeline complete'));
+          console.log(ui.label('Files', `${fmtNum(files.length)}${fileErrors.length ? ` (${fileErrors.length} failed)` : ''}`));
+          console.log(ui.label('Chunks', fmtNum(allChunks.length)));
+          console.log(ui.label('Tokens', fmtNum(totalApiTokens)));
+          console.log(ui.label('Stored', `${fmtNum(insertResult.insertedCount)} docs → ${db}.${collection}`));
+          console.log('');
+          console.log(ui.dim('  Next: vai query "your search" --db ' + db + ' --collection ' + collection));
+        }
+      } catch (err) {
+        console.error(ui.error(err.message));
+        process.exit(1);
+      } finally {
+        if (client) await client.close();
+      }
+    });
+}
+
+module.exports = { registerPipeline };

package/src/commands/query.js

@@ -0,0 +1,266 @@
+'use strict';
+
+const { getDefaultModel, DEFAULT_RERANK_MODEL } = require('../lib/catalog');
+const { generateEmbeddings, apiRequest } = require('../lib/api');
+const { getMongoCollection } = require('../lib/mongo');
+const { loadProject } = require('../lib/project');
+const ui = require('../lib/ui');
+
+/**
+ * Register the query command on a Commander program.
+ * @param {import('commander').Command} program
+ */
+function registerQuery(program) {
+  program
+    .command('query <text>')
+    .description('Search + rerank in one shot — the two-stage retrieval pattern')
+    .option('--db <database>', 'Database name')
+    .option('--collection <name>', 'Collection name')
+    .option('--index <name>', 'Vector search index name')
+    .option('--field <name>', 'Embedding field name')
+    .option('-m, --model <model>', 'Embedding model for query')
+    .option('-d, --dimensions <n>', 'Output dimensions', (v) => parseInt(v, 10))
+    .option('-l, --limit <n>', 'Number of vector search candidates', (v) => parseInt(v, 10), 20)
+    .option('-k, --top-k <n>', 'Final results to return (after rerank)', (v) => parseInt(v, 10), 5)
+    .option('--rerank', 'Enable reranking (recommended)')
+    .option('--no-rerank', 'Skip reranking — vector search only')
+    .option('--rerank-model <model>', 'Reranking model')
+    .option('--text-field <name>', 'Document text field for reranking and display', 'text')
+    .option('--filter <json>', 'Pre-filter JSON for $vectorSearch')
+    .option('--num-candidates <n>', 'ANN candidates (default: limit × 15)', (v) => parseInt(v, 10))
+    .option('--show-vectors', 'Include embedding vectors in output')
+    .option('--json', 'Machine-readable JSON output')
+    .option('-q, --quiet', 'Suppress non-essential output')
+    .action(async (text, opts) => {
+      let client;
+      try {
+        // Merge project config
+        const { config: proj } = loadProject();
+        const db = opts.db || proj.db;
+        const collection = opts.collection || proj.collection;
+        const index = opts.index || proj.index || 'vector_index';
+        const field = opts.field || proj.field || 'embedding';
+        const model = opts.model || proj.model || getDefaultModel();
+        const rerankModel = opts.rerankModel || DEFAULT_RERANK_MODEL;
+        const textField = opts.textField || 'text';
+        const dimensions = opts.dimensions || proj.dimensions;
+        const doRerank = opts.rerank !== false;
+
+        if (!db || !collection) {
+          console.error(ui.error('Database and collection required. Use --db and --collection, or create .vai.json with "vai init".'));
+          process.exit(1);
+        }
+
+        const useColor = !opts.json;
+        const useSpinner = useColor && !opts.quiet;
+
+        // Step 1: Embed query
+        let spin;
+        if (useSpinner) {
+          spin = ui.spinner('Embedding query...');
+          spin.start();
+        }
+
+        const embedOpts = { model, inputType: 'query' };
+        if (dimensions) embedOpts.dimensions = dimensions;
+        const embedResult = await generateEmbeddings([text], embedOpts);
+        const queryVector = embedResult.data[0].embedding;
+        const embedTokens = embedResult.usage?.total_tokens || 0;
+
+        if (spin) spin.stop();
+
+        // Step 2: Vector search
+        if (useSpinner) {
+          spin = ui.spinner(`Searching ${db}.${collection}...`);
+          spin.start();
+        }
+
+        const { client: c, coll } = await connectCollection(db, collection);
+        client = c;
+
+        const numCandidates = opts.numCandidates || Math.min(opts.limit * 15, 10000);
+        const vectorSearchStage = {
+          index,
+          path: field,
+          queryVector,
+          numCandidates,
+          limit: opts.limit,
+        };
+
+        if (opts.filter) {
+          try {
+            vectorSearchStage.filter = JSON.parse(opts.filter);
+          } catch {
+            if (spin) spin.stop();
+            console.error(ui.error('Invalid --filter JSON.'));
+            process.exit(1);
+          }
+        }
+
+        const pipeline = [
+          { $vectorSearch: vectorSearchStage },
+          { $addFields: { _vsScore: { $meta: 'vectorSearchScore' } } },
+        ];
+
+        const searchResults = await coll.aggregate(pipeline).toArray();
+        if (spin) spin.stop();
+
+        if (searchResults.length === 0) {
+          if (opts.json) {
+            console.log(JSON.stringify({ query: text, results: [], stages: { search: 0, rerank: 0 } }, null, 2));
+          } else {
+            console.log(ui.yellow('No results found.'));
+          }
+          return;
+        }
+
+        // Step 3: Rerank (optional)
+        let finalResults;
+        let rerankTokens = 0;
+
+        if (doRerank && searchResults.length > 1) {
+          if (useSpinner) {
+            spin = ui.spinner(`Reranking ${searchResults.length} results...`);
+            spin.start();
+          }
+
+          // Extract text for reranking
+          const documents = searchResults.map(doc => {
+            const txt = doc[textField];
+            if (!txt) return JSON.stringify(doc);
+            return typeof txt === 'string' ? txt : JSON.stringify(txt);
+          });
+
+          const rerankBody = {
+            query: text,
+            documents,
+            model: rerankModel,
+            top_k: opts.topK,
+          };
+
+          const rerankResult = await apiRequest('/rerank', rerankBody);
+          rerankTokens = rerankResult.usage?.total_tokens || 0;
+
+          if (spin) spin.stop();
+
+          // Map reranked indices back to original docs
+          finalResults = (rerankResult.data || []).map(item => {
+            const doc = searchResults[item.index];
+            return {
+              ...doc,
+              _vsScore: doc._vsScore,
+              _rerankScore: item.relevance_score,
+              _finalScore: item.relevance_score,
+            };
+          });
+        } else {
+          // No rerank — just take top-k from vector search
+          finalResults = searchResults.slice(0, opts.topK).map(doc => ({
+            ...doc,
+            _finalScore: doc._vsScore,
+          }));
+        }
+
+        // Build output
+        const output = finalResults.map((doc, i) => {
+          const clean = {};
+          // Include key fields
+          if (doc._id) clean._id = doc._id;
+          if (doc[textField]) {
+            clean[textField] = doc[textField];
+          }
+          // Include metadata fields (skip embedding and internal scores)
+          for (const key of Object.keys(doc)) {
+            if (key === field || key === '_vsScore' || key === '_rerankScore' || key === '_finalScore') continue;
+            if (key === '_id' || key === textField) continue;
+            if (!opts.showVectors) clean[key] = doc[key];
+            else clean[key] = doc[key];
+          }
+          // Scores
+          clean.score = doc._finalScore;
+          if (doc._vsScore !== undefined) clean.vectorScore = doc._vsScore;
+          if (doc._rerankScore !== undefined) clean.rerankScore = doc._rerankScore;
+          clean.rank = i + 1;
+          return clean;
+        });
+
+        if (opts.json) {
+          console.log(JSON.stringify({
+            query: text,
+            model,
+            rerankModel: doRerank ? rerankModel : null,
+            db,
+            collection,
+            stages: {
+              searchCandidates: searchResults.length,
+              finalResults: output.length,
+              reranked: doRerank && searchResults.length > 1,
+            },
+            tokens: { embed: embedTokens, rerank: rerankTokens },
+            results: output,
+          }, null, 2));
+          return;
+        }
+
+        // Pretty output
+        if (!opts.quiet) {
+          console.log('');
+          console.log(ui.label('Query', ui.cyan(`"${text}"`)));
+          console.log(ui.label('Search', `${searchResults.length} candidates from ${ui.dim(`${db}.${collection}`)}`));
+          if (doRerank && searchResults.length > 1) {
+            console.log(ui.label('Rerank', `Top ${output.length} via ${ui.dim(rerankModel)}`));
+          }
+          console.log(ui.label('Model', ui.dim(model)));
+          console.log('');
+        }
+
+        for (let i = 0; i < output.length; i++) {
+          const r = output[i];
+          const scoreStr = r.score != null ? ui.score(r.score) : 'N/A';
+          const vsStr = r.vectorScore != null ? ui.dim(`vs:${r.vectorScore.toFixed(3)}`) : '';
+          const rrStr = r.rerankScore != null ? ui.dim(`rr:${r.rerankScore.toFixed(3)}`) : '';
+          const scores = [vsStr, rrStr].filter(Boolean).join(' ');
+
+          console.log(`${ui.bold(`#${i + 1}`)} ${scoreStr} ${scores}`);
+
+          // Show text preview
+          const textVal = r[textField];
+          if (textVal) {
+            const preview = textVal.substring(0, 200);
+            const ellipsis = textVal.length > 200 ? '...' : '';
+            console.log(`  ${preview}${ellipsis}`);
+          }
+
+          // Show source metadata if present
+          if (r.source) console.log(`  ${ui.dim('source: ' + r.source)}`);
+          if (r.metadata?.source) console.log(`  ${ui.dim('source: ' + r.metadata.source)}`);
+
+          console.log(`  ${ui.dim('_id: ' + r._id)}`);
+          console.log('');
+        }
+
+        if (!opts.quiet) {
+          const totalTokens = embedTokens + rerankTokens;
+          console.log(ui.dim(`  Tokens: ${totalTokens} (embed: ${embedTokens}${rerankTokens ? `, rerank: ${rerankTokens}` : ''})`));
+        }
+      } catch (err) {
+        console.error(ui.error(err.message));
+        process.exit(1);
+      } finally {
+        if (client) await client.close();
+      }
+    });
+}
+
+/**
+ * Connect to a MongoDB collection.
+ * @param {string} db
+ * @param {string} collName
+ * @returns {Promise<{client: MongoClient, coll: Collection}>}
+ */
+async function connectCollection(db, collName) {
+  const { client, collection } = await getMongoCollection(db, collName);
+  return { client, coll: collection };
+}
+
+module.exports = { registerQuery };