← back to Goodquestion

src/lib/ai-provider.ts

250 lines

import Anthropic from '@anthropic-ai/sdk'
import { GoogleGenerativeAI } from '@google/generative-ai'
import { mcpClient } from './mcp-client'

/**
 * Unified AI provider interface with automatic fallback
 * Tracks token usage via MCP and switches to Gemini when needed
 */

export interface AIResponse {
  content: string
  provider: 'anthropic' | 'gemini'
  tokensUsed?: number
}

export interface AIProviderStats {
  provider: 'anthropic' | 'gemini'
  requestCount: number
  totalTokens: number
  lastError?: string
  lastErrorTime?: Date
}

class AIProviderManager {
  private anthropic: Anthropic
  private gemini: GoogleGenerativeAI
  private stats: Map<string, AIProviderStats> = new Map()

  // Token tracking per minute
  private tokenUsageWindow: { timestamp: number; tokens: number; provider: string }[] = []

  // Thresholds
  private readonly ANTHROPIC_TOKEN_LIMIT_PER_MINUTE = 4000
  private readonly RATE_LIMIT_COOLDOWN_MS = 60000 // 1 minute

  constructor() {
    this.anthropic = new Anthropic({
      apiKey: process.env.ANTHROPIC_API_KEY || '',
    })

    this.gemini = new GoogleGenerativeAI(
      process.env.GOOGLE_API_KEY || ''
    )

    // Initialize stats
    this.stats.set('anthropic', {
      provider: 'anthropic',
      requestCount: 0,
      totalTokens: 0,
    })
    this.stats.set('gemini', {
      provider: 'gemini',
      requestCount: 0,
      totalTokens: 0,
    })
  }

  /**
   * Track token usage in a rolling window
   */
  private trackTokenUsage(provider: string, tokens: number) {
    const now = Date.now()
    this.tokenUsageWindow.push({ timestamp: now, tokens, provider })

    // Remove entries older than 1 minute
    this.tokenUsageWindow = this.tokenUsageWindow.filter(
      entry => now - entry.timestamp < 60000
    )

    // Update stats
    const stats = this.stats.get(provider)
    if (stats) {
      stats.requestCount++
      stats.totalTokens += tokens
      this.stats.set(provider, stats)
    }

    // Log to MCP for monitoring (optional - MCP doesn't have built-in monitoring)
    console.log(`[AI Provider] ${provider} used ${tokens} tokens. Total in last minute: ${this.getTokensInLastMinute(provider)}`)
  }

  /**
   * Get total tokens used in the last minute for a provider
   */
  private getTokensInLastMinute(provider: string): number {
    const now = Date.now()
    return this.tokenUsageWindow
      .filter(entry => entry.provider === provider && now - entry.timestamp < 60000)
      .reduce((sum, entry) => sum + entry.tokens, 0)
  }

  /**
   * Check if Anthropic is currently rate limited
   */
  private isAnthropicRateLimited(): boolean {
    const stats = this.stats.get('anthropic')
    if (!stats?.lastError) return false

    // Check if the last error was a rate limit and within cooldown period
    if (stats.lastError.includes('rate_limit') || stats.lastError.includes('credit balance')) {
      const timeSinceError = stats.lastErrorTime ? Date.now() - stats.lastErrorTime.getTime() : Infinity
      return timeSinceError < this.RATE_LIMIT_COOLDOWN_MS
    }

    return false
  }

  /**
   * Check if we're approaching Anthropic's token limit
   */
  private isApproachingTokenLimit(): boolean {
    const tokensUsed = this.getTokensInLastMinute('anthropic')
    // Use 80% of limit as threshold
    return tokensUsed >= this.ANTHROPIC_TOKEN_LIMIT_PER_MINUTE * 0.8
  }

  /**
   * Determine which provider to use
   */
  private selectProvider(): 'anthropic' | 'gemini' {
    // Check if Anthropic is rate limited
    if (this.isAnthropicRateLimited()) {
      console.log('[AI Provider] Anthropic is rate limited, using Gemini')
      return 'gemini'
    }

    // Check if approaching token limit
    if (this.isApproachingTokenLimit()) {
      console.log('[AI Provider] Approaching Anthropic token limit, using Gemini')
      return 'gemini'
    }

    // Default to Anthropic (preferred for quality)
    return 'anthropic'
  }

  /**
   * Generate a completion using the best available provider
   */
  async generateCompletion(
    prompt: string,
    options: {
      maxTokens?: number
      temperature?: number
      forceProvider?: 'anthropic' | 'gemini'
    } = {}
  ): Promise<AIResponse> {
    const provider = options.forceProvider || this.selectProvider()
    const maxTokens = options.maxTokens || 5000
    const temperature = options.temperature || 1.0

    try {
      if (provider === 'anthropic') {
        const message = await this.anthropic.messages.create({
          model: 'claude-opus-4-8',
          max_tokens: maxTokens,
          temperature,
          messages: [{ role: 'user', content: prompt }],
        })

        const content = message.content[0].type === 'text' ? message.content[0].text : ''
        const tokensUsed = message.usage?.output_tokens || 0

        this.trackTokenUsage('anthropic', tokensUsed)

        return {
          content,
          provider: 'anthropic',
          tokensUsed,
        }
      } else {
        // Use Gemini as fallback
        const model = this.gemini.getGenerativeModel({
          model: 'gemini-1.5-pro',
          generationConfig: {
            maxOutputTokens: maxTokens,
            temperature,
          },
        })

        const result = await model.generateContent(prompt)
        const response = result.response
        const content = response.text()

        // Gemini doesn't provide exact token usage in the same way
        // Estimate: ~4 chars per token
        const estimatedTokens = Math.ceil(content.length / 4)

        this.trackTokenUsage('gemini', estimatedTokens)

        return {
          content,
          provider: 'gemini',
          tokensUsed: estimatedTokens,
        }
      }
    } catch (error: any) {
      // Record error for rate limiting detection
      const stats = this.stats.get(provider)
      if (stats) {
        stats.lastError = error?.message || String(error)
        stats.lastErrorTime = new Date()
        this.stats.set(provider, stats)
      }

      // If Anthropic fails with rate limit or credit error, try Gemini
      if (provider === 'anthropic' && !options.forceProvider) {
        const errorMsg = error?.message || String(error)
        if (errorMsg.includes('rate_limit') ||
            errorMsg.includes('credit balance') ||
            error?.status === 429 ||
            error?.status === 400) {
          console.log(`[AI Provider] Anthropic failed (${errorMsg}), falling back to Gemini`)
          return this.generateCompletion(prompt, { ...options, forceProvider: 'gemini' })
        }
      }

      throw error
    }
  }

  /**
   * Get current provider statistics
   */
  getStats(): Record<string, AIProviderStats> {
    const result: Record<string, AIProviderStats> = {}
    this.stats.forEach((value, key) => {
      result[key] = { ...value }
    })
    return result
  }

  /**
   * Reset statistics
   */
  resetStats() {
    this.stats.forEach((value, key) => {
      this.stats.set(key, {
        provider: value.provider,
        requestCount: 0,
        totalTokens: 0,
      })
    })
    this.tokenUsageWindow = []
  }
}

// Export singleton instance
export const aiProvider = new AIProviderManager()