voyageai-cli 1.22.0 → 1.23.0

package/src/mcp/tools/embedding.js

@@ -0,0 +1,67 @@
+'use strict';
+
+const { generateEmbeddings } = require('../../lib/api');
+const { cosineSimilarity } = require('../../lib/math');
+
+/**
+ * Register embedding tools: vai_embed, vai_similarity
+ * @param {import('@modelcontextprotocol/sdk/server/mcp.js').McpServer} server
+ * @param {object} schemas
+ */
+function registerEmbeddingTools(server, schemas) {
+  // vai_embed — embed text and return the vector
+  server.tool(
+    'vai_embed',
+    'Embed text using a Voyage AI model and return the vector representation. Use when you need the raw embedding vector for custom similarity logic, storing in another system, or debugging.',
+    schemas.embedSchema,
+    async (input) => {
+      const embedOpts = { model: input.model, inputType: input.inputType };
+      if (input.dimensions) embedOpts.dimensions = input.dimensions;
+
+      const result = await generateEmbeddings([input.text], embedOpts);
+      const vector = result.data[0].embedding;
+
+      const structured = {
+        text: input.text.slice(0, 100) + (input.text.length > 100 ? '...' : ''),
+        model: input.model,
+        vector,
+        dimensions: vector.length,
+        inputType: input.inputType,
+      };
+
+      return {
+        structuredContent: structured,
+        content: [{ type: 'text', text: `Embedded text (${vector.length} dimensions, model: ${input.model}, type: ${input.inputType}). Vector: [${vector.slice(0, 5).map(v => v.toFixed(4)).join(', ')}, ... ${vector.length - 5} more]` }],
+      };
+    }
+  );
+
+  // vai_similarity — compare two texts
+  server.tool(
+    'vai_similarity',
+    'Compare two texts semantically by embedding both and computing cosine similarity. Returns a score from -1 (opposite) to 1 (identical). Use for duplicate detection, relevance checking, or topic comparison.',
+    schemas.similaritySchema,
+    async (input) => {
+      const result = await generateEmbeddings([input.text1, input.text2], {
+        model: input.model,
+        inputType: 'document',
+      });
+
+      const vec1 = result.data[0].embedding;
+      const vec2 = result.data[1].embedding;
+      const similarity = cosineSimilarity(vec1, vec2);
+
+      return {
+        structuredContent: {
+          text1: input.text1.slice(0, 100) + (input.text1.length > 100 ? '...' : ''),
+          text2: input.text2.slice(0, 100) + (input.text2.length > 100 ? '...' : ''),
+          similarity,
+          model: input.model,
+        },
+        content: [{ type: 'text', text: `Similarity: ${similarity.toFixed(4)} (model: ${input.model})\nText 1: "${input.text1.slice(0, 80)}..."\nText 2: "${input.text2.slice(0, 80)}..."` }],
+      };
+    }
+  );
+}
+
+module.exports = { registerEmbeddingTools };

package/src/mcp/tools/ingest.js

@@ -0,0 +1,89 @@
+'use strict';
+
+const { chunk } = require('../../lib/chunker');
+const { generateEmbeddings } = require('../../lib/api');
+const { getMongoCollection } = require('../../lib/mongo');
+const { loadProject } = require('../../lib/project');
+const { getDefaultModel } = require('../../lib/catalog');
+
+/**
+ * Register the vai_ingest tool (write operation).
+ * @param {import('@modelcontextprotocol/sdk/server/mcp.js').McpServer} server
+ * @param {object} schemas
+ */
+function registerIngestTool(server, schemas) {
+  server.tool(
+    'vai_ingest',
+    'Add a document to a collection: chunks the text, embeds each chunk with Voyage AI, and stores them in MongoDB Atlas. Use when the user provides new content to add to the knowledge base.',
+    schemas.ingestSchema,
+    async (input) => {
+      const { config: proj } = loadProject();
+      const db = input.db || proj.db;
+      const collName = input.collection || proj.collection;
+      if (!db) throw new Error('No database specified. Pass db parameter or configure via vai init.');
+      if (!collName) throw new Error('No collection specified. Pass collection parameter or configure via vai init.');
+
+      const model = input.model || proj.model || getDefaultModel();
+      const start = Date.now();
+
+      // Step 1: Chunk the text
+      const chunks = chunk(input.text, {
+        strategy: input.chunkStrategy,
+        size: input.chunkSize,
+      });
+
+      if (chunks.length === 0) {
+        return {
+          structuredContent: { source: input.source || 'unknown', chunksCreated: 0, collection: collName },
+          content: [{ type: 'text', text: 'No chunks produced — text may be too short or empty.' }],
+        };
+      }
+
+      // Step 2: Embed all chunks
+      const embedResult = await generateEmbeddings(chunks, {
+        model,
+        inputType: 'document',
+      });
+
+      // Step 3: Store in MongoDB
+      const { client, collection: coll } = await getMongoCollection(db, collName);
+      try {
+        const docs = chunks.map((text, i) => ({
+          text,
+          embedding: embedResult.data[i].embedding,
+          source: input.source || 'mcp-ingest',
+          metadata: {
+            ...(input.metadata || {}),
+            ingestedAt: new Date().toISOString(),
+            chunkIndex: i,
+            totalChunks: chunks.length,
+            model,
+            chunkStrategy: input.chunkStrategy,
+          },
+        }));
+
+        await coll.insertMany(docs);
+        const timeMs = Date.now() - start;
+
+        const structured = {
+          source: input.source || 'mcp-ingest',
+          chunksCreated: chunks.length,
+          collection: collName,
+          database: db,
+          model,
+          timeMs,
+          metadata: input.metadata || {},
+        };
+
+        return {
+          structuredContent: structured,
+          content: [{ type: 'text', text: `Ingested "${input.source || 'document'}" into ${db}.${collName}: ${chunks.length} chunks embedded with ${model} (${timeMs}ms)` }],
+        };
+      } finally {
+        await client.close();
+      }
+    }
+  );
+}
+
+module.exports = { registerIngestTool };

package/src/mcp/tools/management.js

@@ -0,0 +1,132 @@
+'use strict';
+
+const { MODEL_CATALOG } = require('../../lib/catalog');
+const { loadProject } = require('../../lib/project');
+const { requireMongoUri } = require('../../lib/mongo');
+
+/**
+ * Introspect MongoDB collections — list collections with vector index info.
+ * @param {string} dbName
+ * @returns {Promise<Array<{ name: string, documentCount: number, hasVectorIndex: boolean, embeddingField?: string, dimensions?: number }>>}
+ */
+async function introspectCollections(dbName) {
+  const { MongoClient } = require('mongodb');
+  const uri = requireMongoUri();
+  const client = new MongoClient(uri);
+  await client.connect();
+
+  try {
+    const db = client.db(dbName);
+    const collections = await db.listCollections().toArray();
+    const results = [];
+
+    for (const collInfo of collections) {
+      if (collInfo.name.startsWith('system.')) continue;
+      const coll = db.collection(collInfo.name);
+      const documentCount = await coll.estimatedDocumentCount();
+
+      let hasVectorIndex = false;
+      let embeddingField;
+      let dimensions;
+
+      try {
+        const indexes = await coll.listSearchIndexes().toArray();
+        for (const idx of indexes) {
+          // Atlas Search index definitions vary; look for vector type
+          const fields = idx.latestDefinition?.fields || [];
+          for (const f of fields) {
+            if (f.type === 'vector') {
+              hasVectorIndex = true;
+              embeddingField = f.path;
+              dimensions = f.numDimensions;
+              break;
+            }
+          }
+          if (hasVectorIndex) break;
+        }
+      } catch {
+        // listSearchIndexes may not be available on non-Atlas deployments
+      }
+
+      results.push({
+        name: collInfo.name,
+        documentCount,
+        hasVectorIndex,
+        ...(embeddingField && { embeddingField }),
+        ...(dimensions && { dimensions }),
+      });
+    }
+
+    return results;
+  } finally {
+    await client.close();
+  }
+}
+
+/**
+ * Register management tools: vai_collections, vai_models
+ * @param {import('@modelcontextprotocol/sdk/server/mcp.js').McpServer} server
+ * @param {object} schemas
+ */
+function registerManagementTools(server, schemas) {
+  // vai_collections — list collections with vector index info
+  server.tool(
+    'vai_collections',
+    'List available MongoDB collections with document counts and vector index information. Use at the start of a task to discover which knowledge bases exist, or when the user mentions a topic and you need to find the right collection.',
+    schemas.collectionsSchema,
+    async (input) => {
+      const { config: proj } = loadProject();
+      const dbName = input.db || proj.db;
+      if (!dbName) throw new Error('No database specified. Pass db parameter or configure via vai init.');
+
+      const collections = await introspectCollections(dbName);
+
+      return {
+        structuredContent: { database: dbName, collections },
+        content: [{
+          type: 'text',
+          text: `Database: ${dbName}\n\n${collections.map(c =>
+            `• ${c.name} — ${c.documentCount} docs${c.hasVectorIndex ? ` ✓ vector index (${c.embeddingField}, ${c.dimensions}d)` : ''}`
+          ).join('\n')}`,
+        }],
+      };
+    }
+  );
+
+  // vai_models — list Voyage AI models
+  server.tool(
+    'vai_models',
+    'List available Voyage AI models with capabilities, benchmarks, and pricing. Use when selecting a model for embedding or reranking, or when the user asks about model tradeoffs.',
+    schemas.modelsSchema,
+    async (input) => {
+      let models = MODEL_CATALOG.filter(m => !m.legacy && !m.unreleased);
+
+      if (input.category !== 'all') {
+        models = models.filter(m => m.type === input.category);
+      }
+
+      const mapped = models.map(m => ({
+        id: m.name,
+        name: m.name,
+        type: m.type,
+        dimensions: m.dimensions,
+        maxTokens: m.maxTokens,
+        pricePerMToken: m.pricePerMToken,
+        ...(m.architecture && { architecture: m.architecture }),
+        ...(m.sharedSpace && { sharedSpace: m.sharedSpace }),
+      }));
+
+      return {
+        structuredContent: { category: input.category, models: mapped },
+        content: [{
+          type: 'text',
+          text: `Available ${input.category === 'all' ? '' : input.category + ' '}models:\n\n${mapped.map(m =>
+            `• ${m.name} (${m.type}) — ${m.dimensions}d, $${m.pricePerMToken}/M tokens`
+          ).join('\n')}`,
+        }],
+      };
+    }
+  );
+}
+
+module.exports = { registerManagementTools, introspectCollections };

package/src/mcp/tools/retrieval.js

@@ -0,0 +1,209 @@
+'use strict';
+
+const { generateEmbeddings, apiRequest } = require('../../lib/api');
+const { getMongoCollection } = require('../../lib/mongo');
+const { getDefaultModel, DEFAULT_RERANK_MODEL } = require('../../lib/catalog');
+const { loadProject } = require('../../lib/project');
+
+/**
+ * Resolve db/collection from tool input, falling back to project config.
+ * @param {object} input
+ * @returns {{ db: string, collection: string }}
+ */
+function resolveDbCollection(input) {
+  const { config: proj } = loadProject();
+  const db = input.db || proj.db;
+  const collection = input.collection || proj.collection;
+  if (!db) throw new Error('No database specified. Pass db parameter or configure via vai init.');
+  if (!collection) throw new Error('No collection specified. Pass collection parameter or configure via vai init.');
+  return { db, collection };
+}
+
+/**
+ * Register retrieval tools: vai_query, vai_search, vai_rerank
+ * @param {import('@modelcontextprotocol/sdk/server/mcp.js').McpServer} server
+ * @param {object} schemas
+ */
+function registerRetrievalTools(server, schemas) {
+  // vai_query — full RAG query: embed → vector search → rerank
+  server.tool(
+    'vai_query',
+    'Full RAG query: embeds the question with Voyage AI, runs vector search against MongoDB Atlas, and reranks results. Use this when you need to answer a question using the knowledge base.',
+    schemas.querySchema,
+    async (input) => {
+      const { db, collection: collName } = resolveDbCollection(input);
+      const { config: proj } = loadProject();
+      const model = input.model || proj.model || getDefaultModel();
+      const index = proj.index || 'vector_index';
+      const field = proj.field || 'embedding';
+      const dimensions = proj.dimensions;
+      const limit = input.limit;
+      const candidateLimit = Math.min(limit * 4, 20);
+      const start = Date.now();
+
+      // Step 1: Embed query
+      const embedOpts = { model, inputType: 'query' };
+      if (dimensions) embedOpts.dimensions = dimensions;
+      const embedResult = await generateEmbeddings([input.query], embedOpts);
+      const queryVector = embedResult.data[0].embedding;
+
+      // Step 2: Vector search
+      const { client, collection: coll } = await getMongoCollection(db, collName);
+      try {
+        const vectorSearchStage = {
+          index,
+          path: field,
+          queryVector,
+          numCandidates: Math.min(candidateLimit * 15, 10000),
+          limit: candidateLimit,
+        };
+        if (input.filter) vectorSearchStage.filter = input.filter;
+
+        const searchResults = await coll.aggregate([
+          { $vectorSearch: vectorSearchStage },
+          { $addFields: { _vsScore: { $meta: 'vectorSearchScore' } } },
+        ]).toArray();
+
+        if (searchResults.length === 0) {
+          return {
+            structuredContent: { query: input.query, results: [], metadata: { collection: collName, model, reranked: false, retrievalTimeMs: Date.now() - start, resultCount: 0 } },
+            content: [{ type: 'text', text: `No results found for "${input.query}" in ${db}.${collName}` }],
+          };
+        }
+
+        // Step 3: Rerank (optional)
+        let finalResults;
+        let reranked = false;
+
+        if (input.rerank && searchResults.length > 1) {
+          const documents = searchResults.map(doc => doc.text || JSON.stringify(doc));
+          const rerankResult = await apiRequest('/rerank', {
+            query: input.query,
+            documents,
+            model: DEFAULT_RERANK_MODEL,
+            top_k: limit,
+          });
+          reranked = true;
+          finalResults = (rerankResult.data || []).map(item => {
+            const doc = searchResults[item.index];
+            return {
+              source: doc.metadata?.source || doc.source || 'unknown',
+              content: doc.text || '',
+              score: doc._vsScore,
+              rerankedScore: item.relevance_score,
+              metadata: doc.metadata || {},
+            };
+          });
+        } else {
+          finalResults = searchResults.slice(0, limit).map(doc => ({
+            source: doc.metadata?.source || doc.source || 'unknown',
+            content: doc.text || '',
+            score: doc._vsScore,
+            metadata: doc.metadata || {},
+          }));
+        }
+
+        const retrievalTimeMs = Date.now() - start;
+        const structured = {
+          query: input.query,
+          results: finalResults,
+          metadata: { collection: collName, model, reranked, retrievalTimeMs, resultCount: finalResults.length },
+        };
+
+        const textLines = finalResults.map((r, i) =>
+          `[${i + 1}] ${r.source} (score: ${(r.rerankedScore || r.score || 0).toFixed(3)})\n${r.content.slice(0, 500)}`
+        );
+
+        return {
+          structuredContent: structured,
+          content: [{ type: 'text', text: `Found ${finalResults.length} results for "${input.query}" (${retrievalTimeMs}ms):\n\n${textLines.join('\n\n')}` }],
+        };
+      } finally {
+        await client.close();
+      }
+    }
+  );
+
+  // vai_search — raw vector similarity search (no reranking)
+  server.tool(
+    'vai_search',
+    'Raw vector similarity search without reranking. Faster than vai_query but results are ordered by vector distance only. Use for exploratory searches or when you plan to rerank separately.',
+    schemas.searchSchema,
+    async (input) => {
+      const { db, collection: collName } = resolveDbCollection(input);
+      const { config: proj } = loadProject();
+      const model = input.model || proj.model || getDefaultModel();
+      const index = proj.index || 'vector_index';
+      const field = proj.field || 'embedding';
+      const dimensions = proj.dimensions;
+      const start = Date.now();
+
+      const embedOpts = { model, inputType: 'query' };
+      if (dimensions) embedOpts.dimensions = dimensions;
+      const embedResult = await generateEmbeddings([input.query], embedOpts);
+      const queryVector = embedResult.data[0].embedding;
+
+      const { client, collection: coll } = await getMongoCollection(db, collName);
+      try {
+        const vectorSearchStage = {
+          index,
+          path: field,
+          queryVector,
+          numCandidates: Math.min(input.limit * 15, 10000),
+          limit: input.limit,
+        };
+        if (input.filter) vectorSearchStage.filter = input.filter;
+
+        const results = await coll.aggregate([
+          { $vectorSearch: vectorSearchStage },
+          { $addFields: { _vsScore: { $meta: 'vectorSearchScore' } } },
+        ]).toArray();
+
+        const mapped = results.map(doc => ({
+          source: doc.metadata?.source || doc.source || 'unknown',
+          content: doc.text || '',
+          score: doc._vsScore,
+          metadata: doc.metadata || {},
+        }));
+
+        const retrievalTimeMs = Date.now() - start;
+
+        return {
+          structuredContent: { query: input.query, results: mapped, metadata: { collection: collName, model, retrievalTimeMs, resultCount: mapped.length } },
+          content: [{ type: 'text', text: `Found ${mapped.length} results for "${input.query}" (${retrievalTimeMs}ms):\n\n${mapped.map((r, i) => `[${i + 1}] ${r.source} (${(r.score || 0).toFixed(3)})\n${r.content.slice(0, 500)}`).join('\n\n')}` }],
+        };
+      } finally {
+        await client.close();
+      }
+    }
+  );
+
+  // vai_rerank — standalone reranking
+  server.tool(
+    'vai_rerank',
+    'Rerank documents against a query using Voyage AI reranker. Takes a query and candidate documents, returns them reordered by relevance. Use when you have documents from another source and want to order them by relevance.',
+    schemas.rerankSchema,
+    async (input) => {
+      const start = Date.now();
+      const result = await apiRequest('/rerank', {
+        query: input.query,
+        documents: input.documents,
+        model: input.model,
+        top_k: input.documents.length,
+      });
+
+      const ranked = (result.data || []).map(item => ({
+        index: item.index,
+        relevanceScore: item.relevance_score,
+        document: input.documents[item.index].slice(0, 200) + (input.documents[item.index].length > 200 ? '...' : ''),
+      }));
+
+      return {
+        structuredContent: { query: input.query, results: ranked, model: input.model, timeMs: Date.now() - start },
+        content: [{ type: 'text', text: `Reranked ${input.documents.length} documents:\n\n${ranked.map((r, i) => `[${i + 1}] Score: ${r.relevanceScore.toFixed(3)} — ${r.document}`).join('\n')}` }],
+      };
+    }
+  );
+}
+
+module.exports = { registerRetrievalTools };