pediatric-ai-scribe-v3/src/utils/embeddings.js

// ============================================================
// EMBEDDINGS UTILITY — Generate & search with Vertex AI embeddings
// Supports: Vertex AI (direct), LiteLLM proxy, OpenAI fallback
// ============================================================

var axios = require('axios');
var { gatewayUrl } = require('./errors');
var { getLiteLLMHeaders } = require('./litellm');

// Vertex AI embedding models (via LiteLLM or direct)
// gemini-embedding-001: 768 dims, multilingual + code, best quality
// text-embedding-005: 768 dims, English + code optimized
// text-multilingual-embedding-002: 768 dims, multilingual focus
var DEFAULT_MODEL = 'vertex_ai/text-embedding-005';
var DEFAULT_DIMS = 768;

/**
 * Generate embedding for text using configured provider
 * @param {string} text - Text to embed (max ~2000 tokens)
 * @param {object} opts - Options: { model, dimensions }
 * @returns {Promise<number[]>} - Embedding vector
 */
async function generateEmbedding(text, opts) {
  opts = opts || {};
  var dbModel, dbDims;
  try {
    var db = require('../db/database');
    dbModel = await db.getSetting('embeddings.model') || '';
    dbDims  = await db.getSetting('embeddings.dimensions') || '';
  } catch(e) { /* DB not available during startup */ }
  var model = opts.model || dbModel || process.env.EMBEDDING_MODEL || DEFAULT_MODEL;
  var dimensions = opts.dimensions || (dbDims ? parseInt(dbDims) : 0) || parseInt(process.env.EMBEDDING_DIMENSIONS) || DEFAULT_DIMS;

  // Truncate text to ~2000 tokens (~8000 chars) to avoid API errors
  // NOTE: Large PDFs (e.g., 100MB) will be truncated to first ~8000 chars for embedding.
  // The full PDF content is still extracted and stored in the database body field.
  // This is expected behavior - embeddings are semantic representations, not full-text storage.
  var truncated = text.substring(0, 8000);
  if (!truncated.trim()) {
    throw new Error('Empty text provided for embedding');
  }

  // Try LiteLLM first if configured
  if (process.env.LITELLM_API_BASE) {
    return await generateEmbeddingLiteLLM(truncated, model, dimensions);
  }

  // Try Vertex AI direct if configured
  if (process.env.GOOGLE_APPLICATION_CREDENTIALS || process.env.VERTEX_PROJECT) {
    return await generateEmbeddingVertexDirect(truncated, model, dimensions);
  }

  // Fallback to OpenAI if configured
  if (process.env.OPENAI_API_KEY) {
    return await generateEmbeddingOpenAI(truncated, model, dimensions);
  }

  throw new Error('No embedding provider configured. Set LITELLM_API_BASE, VERTEX_PROJECT, or OPENAI_API_KEY');
}

/**
 * Generate embedding via LiteLLM proxy
 */
async function generateEmbeddingLiteLLM(text, model, dimensions) {
  try {
    var payload = {
      model: model,
      input: text
    };

    // Only include dimensions if model supports it (some models have fixed dims)
    if (dimensions && model.includes('text-embedding-005')) {
      payload.dimensions = dimensions;
    }

    var response = await axios.post(gatewayUrl('/embeddings'), payload, {
      headers: getLiteLLMHeaders('application/json'),
      timeout: 30000
    });

    if (!response.data || !response.data.data || !response.data.data[0]) {
      throw new Error('Invalid response from LiteLLM embeddings API');
    }

    return response.data.data[0].embedding;
  } catch (err) {
    console.error('[Embeddings] LiteLLM error:', err.response?.data || err.message);
    throw new Error('LiteLLM embedding failed: ' + (err.response?.data?.error || err.message));
  }
}

/**
 * Generate embedding via Vertex AI direct (using @google-cloud/vertexai)
 */
async function generateEmbeddingVertexDirect(text, model, dimensions) {
  try {
    var { VertexAI } = require('@google-cloud/vertexai');

    var project = process.env.VERTEX_PROJECT || process.env.GOOGLE_CLOUD_PROJECT;
    var location = process.env.VERTEX_LOCATION || 'us-central1';

    if (!project) {
      throw new Error('VERTEX_PROJECT or GOOGLE_CLOUD_PROJECT not set');
    }

    var vertexAI = new VertexAI({ project: project, location: location });

    // Extract model name (strip vertex_ai/ prefix if present)
    var modelName = model.replace(/^vertex_ai\//, '');

    // For text-embedding-005, we can specify output dimensions
    var request = {
      instances: [{ content: text }]
    };

    if (dimensions && modelName.includes('text-embedding-005')) {
      request.parameters = { outputDimensionality: dimensions };
    }

    // Use predictText API for embeddings
    var predictionClient = vertexAI.preview.getPredictionServiceClient();
    var endpoint = `projects/${project}/locations/${location}/publishers/google/models/${modelName}`;

    var [response] = await predictionClient.predict({
      endpoint: endpoint,
      instances: [{ content: text }],
      parameters: request.parameters || {}
    });

    if (!response || !response.predictions || !response.predictions[0]) {
      throw new Error('Invalid response from Vertex AI');
    }

    var prediction = response.predictions[0];
    return prediction.embeddings?.values || prediction.values || prediction;

  } catch (err) {
    console.error('[Embeddings] Vertex AI direct error:', err.message);
    throw new Error('Vertex AI embedding failed: ' + err.message);
  }
}

/**
 * Generate embedding via OpenAI (fallback)
 */
async function generateEmbeddingOpenAI(text, model, dimensions) {
  try {
    var openai = require('openai');
    var client = new openai.OpenAI({ apiKey: process.env.OPENAI_API_KEY });

    // Use OpenAI's text-embedding-3-small model (1536 dims by default)
    var embModel = 'text-embedding-3-small';
    var response = await client.embeddings.create({
      model: embModel,
      input: text,
      dimensions: dimensions || 768 // OpenAI supports custom dimensions
    });

    return response.data[0].embedding;
  } catch (err) {
    console.error('[Embeddings] OpenAI error:', err.message);
    throw new Error('OpenAI embedding failed: ' + err.message);
  }
}

/**
 * Search for similar content using cosine similarity
 * @param {string} queryText - Search query
 * @param {object} opts - Options: { limit, threshold, contentType }
 * @returns {Promise<Array>} - Matching content with similarity scores
 */
async function searchSimilar(queryText, opts) {
  opts = opts || {};
  var limit = opts.limit || 10;
  var threshold = opts.threshold || 0.5; // Cosine similarity threshold (0-1)

  var db = require('../db/database');

  // Generate embedding for query
  var queryEmbedding = await generateEmbedding(queryText);

  // Build WHERE clause for filtering
  var whereClause = 'WHERE c.published = true AND c.embedding IS NOT NULL';
  var params = [JSON.stringify(queryEmbedding), threshold, limit];
  var paramIdx = 4;

  if (opts.contentType) {
    whereClause += ' AND c.content_type = $' + paramIdx;
    params.push(opts.contentType);
    paramIdx++;
  }

  if (opts.categoryId) {
    whereClause += ' AND c.category_id = $' + paramIdx;
    params.push(opts.categoryId);
  }

  // Query with cosine similarity using pgvector
  // 1 - (a <=> b) converts distance to similarity (higher = more similar)
  var sql = `
    SELECT
      c.id, c.title, c.slug, c.subject, c.content_type, c.created_at,
      cat.name as category_name, cat.slug as category_slug,
      1 - (c.embedding <=> $1::vector) as similarity,
      (SELECT COUNT(*) FROM learning_questions q WHERE q.content_id = c.id) as question_count
    FROM learning_content c
    LEFT JOIN learning_categories cat ON c.category_id = cat.id
    ${whereClause}
    AND 1 - (c.embedding <=> $1::vector) >= $2
    ORDER BY c.embedding <=> $1::vector
    LIMIT $3
  `;

  var results = await db.all(sql, params);
  return results;
}

/**
 * Generate embedding for learning content (combines title + subject + body)
 * @param {object} content - { title, subject, body }
 * @returns {Promise<number[]>} - Embedding vector
 */
async function generateContentEmbedding(content) {
  // Combine title, subject, and body (weighted toward title)
  var text = [
    content.title || '',
    content.title || '', // Title twice for emphasis
    content.subject || '',
    stripHtml(content.body || '').substring(0, 6000)
  ].filter(Boolean).join('\n\n');

  return await generateEmbedding(text);
}

/**
 * Strip HTML tags from string
 */
function stripHtml(html) {
  return html.replace(/<[^>]*>/g, ' ').replace(/\s+/g, ' ').trim();
}

/**
 * Check if embeddings are available (provider configured)
 */
function isEmbeddingsAvailable() {
  return !!(
    process.env.LITELLM_API_BASE ||
    process.env.VERTEX_PROJECT ||
    process.env.GOOGLE_CLOUD_PROJECT ||
    process.env.OPENAI_API_KEY
  );
}

module.exports = {
  generateEmbedding,
  generateContentEmbedding,
  searchSimilar,
  isEmbeddingsAvailable,
  DEFAULT_MODEL,
  DEFAULT_DIMS
};