@guardaian/sdk/interceptors/gemini.js

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

/**
 * Google AI (Gemini) API Interceptor - Bulletproof Production Version
 * 
 * CRITICAL GUARANTEES:
 * 1. Customer's Gemini 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. Perfect for free testing (Gemini has generous free tier)
 */

function patchGemini(GoogleGenerativeAI, guardaianClient) {
  // Early validation with error protection
  if (!GoogleGenerativeAI || typeof GoogleGenerativeAI.GoogleGenerativeAI !== 'function') {
    safeLog(guardaianClient, '⚠️  Invalid Google Generative AI module - skipping patch (non-blocking)');
    return;
  }

  // Prevent double-patching
  if (GoogleGenerativeAI.__guardaianPatched) {
    safeLog(guardaianClient, '⚠️  Google Generative AI already patched - skipping');
    return;
  }

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

    // Patch the GoogleGenerativeAI class
    const OriginalGoogleGenerativeAI = GoogleGenerativeAI.GoogleGenerativeAI;
    
    GoogleGenerativeAI.GoogleGenerativeAI = class extends OriginalGoogleGenerativeAI {
      getGenerativeModel(modelParams) {
        let model;
        
        try {
          // CRITICAL: Always call original method first
          model = super.getGenerativeModel(modelParams);
        } catch (originalError) {
          // Re-throw original errors - don't interfere
          throw originalError;
        }
        
        // CRITICAL: Patch model methods in isolated context
        setImmediate(() => {
          try {
            patchModelMethods(model, modelParams, guardaianClient);
          } catch (patchError) {
            // Patching errors are isolated from customer code
            safeLog(guardaianClient, `❌ Gemini model patching error (isolated): ${patchError.message}`);
          }
        });
        
        // CRITICAL: Always return original model unchanged
        return model;
      }
    };

    // Copy static properties and methods from original class
    Object.setPrototypeOf(GoogleGenerativeAI.GoogleGenerativeAI, OriginalGoogleGenerativeAI);
    Object.assign(GoogleGenerativeAI.GoogleGenerativeAI, OriginalGoogleGenerativeAI);

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

/**
 * Patch model methods with bulletproof error handling
 */
function patchModelMethods(model, modelParams, guardaianClient) {
  try {
    // Patch generateContent method
    if (model.generateContent && typeof model.generateContent === 'function') {
      const originalGenerateContent = model.generateContent.bind(model);
      
      model.generateContent = async function(request) {
        const startTime = Date.now();
        
        // CRITICAL: Customer's API call always happens first
        let response, originalError;
        
        try {
          response = await originalGenerateContent(request);
        } catch (error) {
          originalError = error;
        }

        // CRITICAL: Track usage in completely isolated context
        setImmediate(() => {
          try {
            trackGeminiUsage(modelParams, request, response, originalError, startTime, guardaianClient, 'generateContent');
          } catch (trackingError) {
            // Tracking errors are completely isolated
            safeLog(guardaianClient, `❌ Gemini tracking error (isolated): ${trackingError.message}`);
          }
        });

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

    // Patch generateContentStream method if it exists
    if (model.generateContentStream && typeof model.generateContentStream === 'function') {
      const originalGenerateContentStream = model.generateContentStream.bind(model);
      
      model.generateContentStream = async function(request) {
        const startTime = Date.now();
        
        // CRITICAL: Customer's streaming call always happens first
        let stream, originalError;
        
        try {
          stream = await originalGenerateContentStream(request);
        } catch (error) {
          originalError = error;
        }

        // Track streaming usage (more complex, but still isolated)
        if (stream && !originalError) {
          setImmediate(() => {
            try {
              trackGeminiStreamUsage(modelParams, request, stream, startTime, guardaianClient);
            } catch (trackingError) {
              safeLog(guardaianClient, `❌ Gemini stream tracking error (isolated): ${trackingError.message}`);
            }
          });
        }

        // Return original result
        if (originalError) throw originalError;
        return stream;
      };
    }

  } catch (methodPatchError) {
    safeLog(guardaianClient, `❌ Failed to patch Gemini model methods: ${methodPatchError.message}`);
  }
}

/**
 * Track Gemini usage with complete error isolation
 */
async function trackGeminiUsage(modelParams, request, response, error, startTime, guardaianClient, operation) {
  try {
    const endTime = Date.now();
    const duration = endTime - startTime;
    const model = modelParams.model || 'gemini-pro';

    if (response) {
      // Successful request tracking
      const responseText = extractGeminiResponseText(response);
      const requestText = extractGeminiRequestText(request);
      
      // Estimate tokens (Gemini doesn't always provide exact counts)
      const inputTokens = estimateTokens(requestText);
      const outputTokens = estimateTokens(responseText);
      const totalTokens = inputTokens + outputTokens;
      
      // Gemini cost calculation (often free for testing)
      const cost = calculateGeminiCost(model, { inputTokens, outputTokens });

      // Fire-and-forget tracking call
      guardaianClient.track({
        service: 'google',
        model: model,
        operation: operation,
        inputTokens: inputTokens,
        outputTokens: outputTokens,
        totalTokens: totalTokens,
        cost: cost,
        duration: duration,
        requestData: sanitizeGeminiRequest(request, modelParams),
        responseData: sanitizeGeminiResponse(response),
        metadata: {
          success: true,
          finishReason: extractFinishReason(response),
          candidatesCount: response.response?.candidates?.length || 1
        }
      }).catch(trackingError => {
        // Even the track() call errors are isolated
        safeLog(guardaianClient, `❌ Gemini track call failed (isolated): ${trackingError.message}`);
      });

      safeLog(guardaianClient, `✅ Tracked Gemini call: ${model} (${totalTokens} tokens, $${cost.toFixed(6)})`);

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

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

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

/**
 * Track streaming usage (simplified version)
 */
async function trackGeminiStreamUsage(modelParams, request, stream, startTime, guardaianClient) {
  try {
    // For streaming, we'll track when the stream starts
    // Final tracking would happen when stream completes, but that's complex
    // For now, just track the initiation
    
    const model = modelParams.model || 'gemini-pro';
    const requestText = extractGeminiRequestText(request);
    const inputTokens = estimateTokens(requestText);
    
    guardaianClient.track({
      service: 'google',
      model: model,
      operation: 'generateContentStream',
      inputTokens: inputTokens,
      outputTokens: 0, // Will be updated when stream completes
      totalTokens: inputTokens,
      cost: 0, // Will be calculated when complete
      duration: Date.now() - startTime,
      requestData: sanitizeGeminiRequest(request, modelParams),
      responseData: { streaming: true },
      metadata: {
        success: true,
        streaming: true
      }
    }).catch(() => {}); // Silent fail for streaming
    
  } catch (streamTrackingError) {
    safeLog(guardaianClient, `❌ Gemini stream tracking error (isolated): ${streamTrackingError.message}`);
  }
}

/**
 * Extract text from Gemini response with error protection
 */
function extractGeminiResponseText(response) {
  try {
    if (response?.response?.text && typeof response.response.text === 'function') {
      return response.response.text();
    }
    
    if (response?.response?.candidates?.[0]?.content?.parts?.[0]?.text) {
      return response.response.candidates[0].content.parts[0].text;
    }
    
    return '';
  } catch (extractError) {
    return '';
  }
}

/**
 * Extract text from Gemini request with error protection
 */
function extractGeminiRequestText(request) {
  try {
    if (typeof request === 'string') {
      return request;
    }
    
    if (request?.contents?.[0]?.parts?.[0]?.text) {
      return request.contents[0].parts[0].text;
    }
    
    if (request?.prompt) {
      return request.prompt;
    }
    
    return JSON.stringify(request).substring(0, 500);
  } catch (extractError) {
    return '';
  }
}

/**
 * Extract finish reason with error protection
 */
function extractFinishReason(response) {
  try {
    return response?.response?.candidates?.[0]?.finishReason || 'unknown';
  } catch (extractError) {
    return 'unknown';
  }
}

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

/**
 * Calculate Gemini costs with error protection
 */
function calculateGeminiCost(model, usage) {
  try {
    // Gemini pricing (many models are free up to rate limits)
    const pricing = {
      'gemini-pro': { input: 0, output: 0 }, // Free tier
      'gemini-1.5-pro': { input: 0.0035 / 1000, output: 0.0105 / 1000 },
      'gemini-1.5-flash': { input: 0.00035 / 1000, output: 0.00105 / 1000 },
      'gemini-1.0-pro': { input: 0, output: 0 } // Free tier
    };

    const modelPricing = pricing[model] || pricing['gemini-pro']; // Default to free
    
    const inputCost = (usage.inputTokens || 0) * modelPricing.input;
    const outputCost = (usage.outputTokens || 0) * modelPricing.output;
    
    return Math.round((inputCost + outputCost) * 10000) / 10000;
  
  } catch (costError) {
    return 0; // Gemini is often free anyway
  }
}

/**
 * Sanitize Gemini request data
 */
function sanitizeGeminiRequest(request, modelParams) {
  try {
    const sanitized = {
      model: modelParams.model || 'gemini-pro',
      temperature: modelParams.generationConfig?.temperature,
      maxOutputTokens: modelParams.generationConfig?.maxOutputTokens,
      topP: modelParams.generationConfig?.topP,
      topK: modelParams.generationConfig?.topK
    };

    if (typeof request === 'string') {
      sanitized.prompt = request.substring(0, 200) + (request.length > 200 ? '...' : '');
    } else if (request?.contents) {
      sanitized.contents = request.contents.slice(0, 5).map(content => ({
        role: content.role,
        parts: content.parts?.slice(0, 3).map(part => ({
          text: part.text ? part.text.substring(0, 200) + (part.text.length > 200 ? '...' : '') : undefined,
          type: part.inlineData ? 'inline_data' : part.text ? 'text' : 'unknown'
        }))
      }));
    }

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

/**
 * Sanitize Gemini response data
 */
function sanitizeGeminiResponse(response) {
  try {
    if (!response) return {};

    const sanitized = {
      candidates_count: response.response?.candidates?.length || 0,
      finish_reason: extractFinishReason(response),
      has_text: !!extractGeminiResponseText(response)
    };

    if (response.response?.candidates?.[0]) {
      const candidate = response.response.candidates[0];
      sanitized.safety_ratings = candidate.safetyRatings?.length || 0;
      sanitized.citation_metadata = !!candidate.citationMetadata;
    }

    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 = {
  patchGemini
};