@guardaian/sdk/interceptors/cohere.js

Zero-friction AI governance and monitoring SDK for Node.js applications

/**
 * Cohere API Interceptor - Bulletproof Production Version
 * 
 * CRITICAL GUARANTEES:
 * 1. Customer's Cohere calls ALWAYS work, even if GuardAIan is down
 * 2. Zero performance impact on customer's AI operations
 * 3. All tracking errors are isolated and non-blocking
 * 4. Supports all Cohere endpoints (chat, generate, embed)
 */

function patchCohere(CohereAI, guardaianClient) {
  // Early validation with error protection
  if (!CohereAI || typeof CohereAI !== 'function') {
    safeLog(guardaianClient, '⚠️  Invalid Cohere module - skipping patch (non-blocking)');
    return;
  }

  // Prevent double-patching
  if (CohereAI.__guardaianPatched) {
    safeLog(guardaianClient, '⚠️  Cohere already patched - skipping');
    return;
  }

  try {
    safeLog(guardaianClient, '🔧 Patching Cohere SDK with bulletproof protection...');

    // Patch various Cohere endpoints
    patchCohereChat(CohereAI, guardaianClient);
    patchCohereGenerate(CohereAI, guardaianClient);
    patchCohereEmbed(CohereAI, guardaianClient);

    // Mark as patched
    CohereAI.__guardaianPatched = true;
    
    safeLog(guardaianClient, '✅ Cohere SDK patched successfully with failsafe protection');
  } catch (patchError) {
    // Even patching errors should not break customer code
    safeLog(guardaianClient, `⚠️  Cohere patching failed (non-blocking): ${patchError.message}`);
  }
}

/**
 * Patch Cohere Chat endpoint with bulletproof error handling
 */
function patchCohereChat(CohereAI, guardaianClient) {
  try {
    const originalChat = CohereAI.prototype.chat;
    
    if (!originalChat) {
      safeLog(guardaianClient, '⚠️  Cohere chat method not found - skipping patch');
      return;
    }

    CohereAI.prototype.chat = async function(params, options) {
      const startTime = Date.now();
      
      // CRITICAL: Always call original Cohere API first
      let response, originalError;
      
      try {
        response = await originalChat.call(this, params, options);
      } catch (error) {
        originalError = error;
      }

      // CRITICAL: Track usage in completely isolated context
      setImmediate(() => {
        try {
          trackCohereUsage(params, response, originalError, startTime, guardaianClient, 'chat');
        } catch (trackingError) {
          safeLog(guardaianClient, `❌ Cohere chat tracking error (isolated): ${trackingError.message}`);
        }
      });

      // CRITICAL: Always return original result or throw original error
      if (originalError) throw originalError;
      return response;
    };

  } catch (patchError) {
    safeLog(guardaianClient, `❌ Failed to patch Cohere chat: ${patchError.message}`);
  }
}

/**
 * Patch Cohere Generate endpoint with bulletproof error handling
 */
function patchCohereGenerate(CohereAI, guardaianClient) {
  try {
    const originalGenerate = CohereAI.prototype.generate;
    
    if (!originalGenerate) {
      safeLog(guardaianClient, '⚠️  Cohere generate method not found - skipping patch');
      return;
    }

    CohereAI.prototype.generate = async function(params, options) {
      const startTime = Date.now();
      
      // CRITICAL: Always call original Cohere API first
      let response, originalError;
      
      try {
        response = await originalGenerate.call(this, params, options);
      } catch (error) {
        originalError = error;
      }

      // CRITICAL: Track usage in completely isolated context
      setImmediate(() => {
        try {
          trackCohereUsage(params, response, originalError, startTime, guardaianClient, 'generate');
        } catch (trackingError) {
          safeLog(guardaianClient, `❌ Cohere generate tracking error (isolated): ${trackingError.message}`);
        }
      });

      // CRITICAL: Always return original result or throw original error
      if (originalError) throw originalError;
      return response;
    };

  } catch (patchError) {
    safeLog(guardaianClient, `❌ Failed to patch Cohere generate: ${patchError.message}`);
  }
}

/**
 * Patch Cohere Embed endpoint with bulletproof error handling
 */
function patchCohereEmbed(CohereAI, guardaianClient) {
  try {
    const originalEmbed = CohereAI.prototype.embed;
    
    if (!originalEmbed) {
      safeLog(guardaianClient, '⚠️  Cohere embed method not found - skipping patch');
      return;
    }

    CohereAI.prototype.embed = async function(params, options) {
      const startTime = Date.now();
      
      // CRITICAL: Always call original Cohere API first
      let response, originalError;
      
      try {
        response = await originalEmbed.call(this, params, options);
      } catch (error) {
        originalError = error;
      }

      // CRITICAL: Track usage in completely isolated context
      setImmediate(() => {
        try {
          trackCohereUsage(params, response, originalError, startTime, guardaianClient, 'embed');
        } catch (trackingError) {
          safeLog(guardaianClient, `❌ Cohere embed tracking error (isolated): ${trackingError.message}`);
        }
      });

      // CRITICAL: Always return original result or throw original error
      if (originalError) throw originalError;
      return response;
    };

  } catch (patchError) {
    safeLog(guardaianClient, `❌ Failed to patch Cohere embed: ${patchError.message}`);
  }
}

/**
 * Track Cohere usage with complete error isolation
 */
async function trackCohereUsage(params, response, error, startTime, guardaianClient, operation) {
  try {
    const endTime = Date.now();
    const duration = endTime - startTime;
    const model = params.model || 'command';

    if (response) {
      // Successful request tracking
      const usage = extractCohereUsage(params, response, operation);
      const cost = calculateCohereCost(model, usage, operation);

      // Fire-and-forget tracking call
      guardaianClient.track({
        service: 'cohere',
        model: model,
        operation: operation,
        inputTokens: usage.inputTokens || 0,
        outputTokens: usage.outputTokens || 0,
        totalTokens: (usage.inputTokens || 0) + (usage.outputTokens || 0),
        cost: cost,
        duration: duration,
        requestData: sanitizeCohereRequest(params, operation),
        responseData: sanitizeCohereResponse(response, operation),
        metadata: {
          success: true,
          operation: operation,
          finishReason: response.finish_reason,
          responseId: response.response_id
        }
      }).catch(trackingError => {
        safeLog(guardaianClient, `❌ Cohere track call failed (isolated): ${trackingError.message}`);
      });

      safeLog(guardaianClient, `✅ Tracked Cohere ${operation}: ${model} (${(usage.inputTokens || 0) + (usage.outputTokens || 0)} tokens, $${cost.toFixed(6)})`);

    } else if (error) {
      // Failed request tracking
      guardaianClient.track({
        service: 'cohere',
        model: model,
        operation: operation,
        inputTokens: 0,
        outputTokens: 0,
        totalTokens: 0,
        cost: 0,
        duration: duration,
        requestData: sanitizeCohereRequest(params, operation),
        responseData: {
          error: true,
          status: error.status || 'unknown',
          message: error.message?.substring(0, 200) || 'Unknown error'
        },
        metadata: {
          success: false,
          errorType: error.constructor.name,
          operation: operation
        }
      }).catch(trackingError => {
        safeLog(guardaianClient, `❌ Cohere error tracking failed (isolated): ${trackingError.message}`);
      });

      safeLog(guardaianClient, `📊 Tracked Cohere ${operation} error: ${model} - ${error.message}`);
    }

  } catch (trackingError) {
    // Ultimate safety net
    safeLog(guardaianClient, `❌ Cohere usage tracking completely failed (isolated): ${trackingError.message}`);
  }
}

/**
 * Extract usage data from Cohere response with error protection
 */
function extractCohereUsage(params, response, operation) {
  try {
    switch (operation) {
      case 'chat':
        const chatUsage = response.meta?.tokens || {};
        return {
          inputTokens: chatUsage.input_tokens || 0,
          outputTokens: chatUsage.output_tokens || 0
        };
        
      case 'generate':
        const generateUsage = response.meta?.tokens || {};
        return {
          inputTokens: generateUsage.input_tokens || 0,
          outputTokens: generateUsage.output_tokens || 0
        };
        
      case 'embed':
        // Embeddings only have input tokens
        const texts = Array.isArray(params.texts) ? params.texts : [params.texts];
        const inputTokens = texts.reduce((total, text) => total + estimateTokens(text), 0);
        return {
          inputTokens: inputTokens,
          outputTokens: 0
        };
        
      default:
        return { inputTokens: 0, outputTokens: 0 };
    }
  } catch (extractError) {
    return { inputTokens: 0, outputTokens: 0 };
  }
}

/**
 * Estimate tokens from text with error protection
 */
function estimateTokens(text) {
  try {
    if (!text || typeof text !== 'string') return 0;
    return Math.ceil(text.length / 4); // Rough estimate: 1 token ≈ 4 characters
  } catch (estimateError) {
    return 0;
  }
}

/**
 * Calculate cost for Cohere API calls with error protection
 */
function calculateCohereCost(model, usage, operation) {
  try {
    // Cohere pricing (as of 2024)
    const coherePricing = {
      // Chat models
      'command': { input: 0.0015, output: 0.002 },
      'command-light': { input: 0.0003, output: 0.0006 },
      'command-r': { input: 0.0005, output: 0.0015 },
      'command-r-plus': { input: 0.003, output: 0.015 },
      
      // Generate models (legacy)
      'command-xlarge': { input: 0.0015, output: 0.002 },
      'command-medium': { input: 0.0006, output: 0.0012 },
      'command-light-xlarge': { input: 0.0003, output: 0.0006 },
      
      // Embedding models (input only)
      'embed-english-v3.0': { input: 0.0001, output: 0 },
      'embed-multilingual-v3.0': { input: 0.0001, output: 0 },
      'embed-english-light-v3.0': { input: 0.0001, output: 0 },
      'embed-multilingual-light-v3.0': { input: 0.0001, output: 0 },
      'embed-english-v2.0': { input: 0.0001, output: 0 },
      'embed-english-light-v2.0': { input: 0.0001, output: 0 },
      'embed-multilingual-v2.0': { input: 0.0001, output: 0 }
    };

    // Default to command pricing if model not found
    const defaultPricing = { input: 0.0015, output: 0.002 };
    const modelPricing = coherePricing[model] || defaultPricing;
    
    const inputCost = (usage.inputTokens || 0) * modelPricing.input / 1000;
    const outputCost = (usage.outputTokens || 0) * modelPricing.output / 1000;
    
    return Math.round((inputCost + outputCost) * 10000) / 10000;
  
  } catch (costError) {
    safeLog(null, `❌ Cohere cost calculation error: ${costError.message}`);
    return 0;
  }
}

/**
 * Sanitize Cohere request data with privacy protection
 */
function sanitizeCohereRequest(params, operation) {
  try {
    const sanitized = {
      model: params.model || 'command',
      temperature: params.temperature,
      max_tokens: params.max_tokens,
      p: params.p,
      k: params.k,
      operation: operation
    };

    switch (operation) {
      case 'chat':
        sanitized.message = params.message ? params.message.substring(0, 200) + '...' : undefined;
        sanitized.chat_history_length = params.chat_history?.length || 0;
        sanitized.preamble = params.preamble ? 'present' : undefined;
        break;
        
      case 'generate':
        sanitized.prompt = params.prompt ? params.prompt.substring(0, 200) + '...' : undefined;
        sanitized.num_generations = params.num_generations || 1;
        sanitized.stop_sequences = params.stop_sequences?.length || 0;
        break;
        
      case 'embed':
        const texts = Array.isArray(params.texts) ? params.texts : [params.texts];
        sanitized.texts_count = texts.length;
        sanitized.total_text_length = texts.reduce((total, text) => total + (text?.length || 0), 0);
        sanitized.input_type = params.input_type;
        sanitized.embedding_types = params.embedding_types;
        break;
    }

    return sanitized;
  } catch (sanitizeError) {
    return { error: 'sanitization_failed' };
  }
}

/**
 * Sanitize Cohere response data with privacy protection
 */
function sanitizeCohereResponse(response, operation) {
  try {
    const sanitized = {
      response_id: response.response_id,
      operation: operation
    };

    switch (operation) {
      case 'chat':
        sanitized.text_length = response.text?.length || 0;
        sanitized.finish_reason = response.finish_reason;
        sanitized.generation_id = response.generation_id;
        break;
        
      case 'generate':
        const generation = response.generations?.[0] || {};
        sanitized.generation_id = generation.id;
        sanitized.text_length = generation.text?.length || 0;
        sanitized.finish_reason = generation.finish_reason;
        sanitized.likelihood = generation.likelihood;
        sanitized.generations_count = response.generations?.length || 0;
        break;
        
      case 'embed':
        sanitized.embeddings_count = response.embeddings?.length || 0;
        break;
    }

    return sanitized;
  } catch (sanitizeError) {
    return { error: 'sanitization_failed' };
  }
}

/**
 * Safe logging that never crashes
 */
function safeLog(guardaianClient, message) {
  try {
    if (guardaianClient?.options?.debug) {
      console.log(`🛡️  GuardAIan: ${message}`);
    }
  } catch (logError) {
    // Even logging can fail - do nothing
  }
}

module.exports = {
  patchCohere
};