@coworker-agency/rag/src/cav/embeddings.js

Retrieval Augmented Generation (RAG) library for document indexing, vector storage, and AI-powered question answering

/**
 * Utility for generating embeddings from text using OpenAI's embedding models
 */
import OpenAI from 'openai';

/**
 * Generate embeddings for document chunks
 * @param {Array<{context: string, content: string}>} chunks - Array of chunks with context and content
 * @param {string} openaiApiKey - OpenAI API key
 * @param {Object} options - Additional options
 * @returns {Promise<Array<{context: string, content: string, vector: number[]}>>} Chunks with vectors
 */
export async function generateEmbeddings(chunks, openaiApiKey, options = {}) {
  const {
    embeddingsModel = 'text-embedding-3-small',
    batchSize = 20
  } = options;

  const openai = new OpenAI({ apiKey: openaiApiKey });
  const chunksWithVectors = [];
  
  // Process in batches for memory efficiency
  const batchCount = Math.ceil(chunks.length / batchSize);
  
  for (let batchIndex = 0; batchIndex < batchCount; batchIndex++) {
    const batchStart = batchIndex * batchSize;
    const batchEnd = Math.min((batchIndex + 1) * batchSize, chunks.length);
    const currentBatch = chunks.slice(batchStart, batchEnd);
    
    console.log(`Generating embeddings for batch ${batchIndex + 1}/${batchCount} with ${currentBatch.length} chunks`);
    
    // Prepare texts to embed (combine context and content for better embeddings)
    const textsToEmbed = currentBatch.map(chunk => 
      `Context: ${chunk.context}\nContent: ${chunk.content}`
    );
    
    // Get embeddings for the batch
    const embedResponse = await openai.embeddings.create({
      model: embeddingsModel,
      input: textsToEmbed,
      encoding_format: "float"
    });
    
    // Add vectors to chunks
    for (let i = 0; i < currentBatch.length; i++) {
      chunksWithVectors.push({
        ...currentBatch[i],
        vector: embedResponse.data[i].embedding
      });
    }
    
    // Add a small delay between batches to avoid rate limiting
    if (batchIndex < batchCount - 1) {
      await new Promise(resolve => setTimeout(resolve, 200));
    }
  }
  
  return chunksWithVectors;
}