simple-ai-provider/src/providers/gemini.ts

/**
 * Google Gemini Provider Implementation
 *
 * Integrates with Google's Gemini models through the official @google/genai SDK
 * (the successor to the deprecated @google/generative-ai package).
 *
 * Key Features:
 * - Support for Gemini 1.5 / 2.0 / 2.5 model families
 * - Streaming responses with real-time chunk delivery
 * - Multimodal capabilities and configurable safety filtering
 * - System instruction handling separate from conversation flow
 *
 * @author Jan-Marlon Leibl
 * @see https://ai.google.dev/docs
 */

import { GoogleGenAI } from '@google/genai';
import type {
  Content,
  GenerateContentConfig,
  GenerateContentResponse,
  SafetySetting
} from '@google/genai';
import type {
  AIProviderConfig,
  CompletionParams,
  CompletionResponse,
  CompletionChunk,
  ProviderInfo,
  AIMessage
} from '../types/index.js';
import { BaseAIProvider } from './base.js';
import { AIProviderError, AIErrorType } from '../types/index.js';
import {
  DEFAULT_MAX_TOKENS,
  DEFAULT_MODELS,
  DEFAULT_TEMPERATURE,
  VALIDATION_PROMPT
} from '../constants.js';

// ============================================================================
// TYPES AND INTERFACES
// ============================================================================

/**
 * Configuration interface for Gemini provider with Google-specific options.
 */
export interface GeminiConfig extends AIProviderConfig {
  /**
   * Default Gemini model to use for requests.
   *
   * Recommended models:
   * - 'gemini-2.5-flash': Latest fast model, great default
   * - 'gemini-2.5-pro': Flagship reasoning model
   * - 'gemini-2.0-flash': Previous generation fast model
   * - 'gemini-1.5-pro': Long-context flagship from the 1.5 generation
   *
   * @default 'gemini-2.5-flash'
   */
  defaultModel?: string;

  /** Safety settings for content filtering and harm prevention. */
  safetySettings?: SafetySetting[];

  /** Default generation configuration applied to every request. */
  generationConfig?: {
    temperature?: number;
    topP?: number;
    topK?: number;
    maxOutputTokens?: number;
    stopSequences?: string[];
  };
}

/**
 * Processed messages structure for Gemini's conversation format.
 */
interface ProcessedMessages {
  /** System instruction string (if any) */
  systemInstruction?: string;
  /** Conversation content in Gemini's format */
  contents: Content[];
}

// ============================================================================
// GEMINI PROVIDER IMPLEMENTATION
// ============================================================================

/**
 * Google Gemini provider implementation backed by @google/genai.
 *
 * The new SDK is stateless w.r.t. model selection — each `generateContent` call
 * specifies the model directly, so this class only holds a single client.
 */
export class GeminiProvider extends BaseAIProvider {
  private client: GoogleGenAI | null = null;
  private readonly defaultModel: string;
  private readonly safetySettings?: SafetySetting[];
  private readonly defaultGenerationConfig?: GeminiConfig['generationConfig'];

  constructor(config: GeminiConfig) {
    super(config);

    this.defaultModel = config.defaultModel || DEFAULT_MODELS.gemini;
    this.safetySettings = config.safetySettings;
    this.defaultGenerationConfig = config.generationConfig;

    this.validateModelName(this.defaultModel);
  }

  // ========================================================================
  // PROTECTED TEMPLATE METHOD IMPLEMENTATIONS
  // ========================================================================

  protected async doInitialize(): Promise<void> {
    try {
      this.client = new GoogleGenAI({ apiKey: this.config.apiKey });

      await this.validateConnection();
    } catch (error) {
      this.client = null;

      throw new AIProviderError(
        `Failed to initialize Gemini provider: ${(error as Error).message}`,
        AIErrorType.AUTHENTICATION,
        undefined,
        error as Error
      );
    }
  }

  protected async doComplete(params: CompletionParams): Promise<CompletionResponse<string>> {
    if (!this.client) {
      throw new AIProviderError('Gemini client not initialized', AIErrorType.INVALID_REQUEST);
    }

    const { systemInstruction, contents } = this.convertMessages(params.messages);
    const model = params.model || this.defaultModel;

    const response = await this.client.models.generateContent({
      model,
      contents,
      config: this.buildConfig(params, systemInstruction)
    });

    return this.formatCompletionResponse(response, model);
  }

  protected async *doStream<T = any>(params: CompletionParams<T>): AsyncIterable<CompletionChunk> {
    if (!this.client) {
      throw new AIProviderError('Gemini client not initialized', AIErrorType.INVALID_REQUEST);
    }

    const { systemInstruction, contents } = this.convertMessages(params.messages);
    const model = params.model || this.defaultModel;

    const stream = await this.client.models.generateContentStream({
      model,
      contents,
      config: this.buildConfig(params, systemInstruction)
    });

    yield* this.processStreamChunks(stream);
  }

  // ========================================================================
  // PUBLIC INTERFACE METHODS
  // ========================================================================

  public getInfo(): ProviderInfo {
    return {
      name: 'Gemini',
      version: '2.0.0',
      models: [
        'gemini-2.5-pro',
        'gemini-2.5-flash',
        'gemini-2.0-flash',
        'gemini-1.5-pro',
        'gemini-1.5-flash'
      ],
      maxContextLength: 1048576,
      supportsStreaming: true,
      capabilities: {
        vision: true,
        functionCalling: true,
        jsonMode: true,
        systemMessages: true,
        reasoning: true,
        codeGeneration: true,
        multimodal: true,
        safetyFiltering: true
      }
    };
  }

  // ========================================================================
  // PRIVATE UTILITY METHODS
  // ========================================================================

  protected override getModelNamePatterns(): RegExp[] {
    return [
      /^gemini-2\.5-(?:pro|flash)(?:-lite)?(?:-\d{4}-\d{2}-\d{2})?$/,
      /^gemini-2\.0-(?:pro|flash)(?:-lite)?(?:-\d{4}-\d{2}-\d{2})?$/,
      /^gemini-1\.5-pro(?:-latest|-vision)?$/,
      /^gemini-1\.5-flash(?:-latest|-8b)?$/,
      /^gemini-1\.0-pro(?:-latest|-vision)?$/,
      /^gemini-pro(?:-vision)?$/,
      /^models\/gemini-.+$/
    ];
  }

  protected override async sendValidationProbe(): Promise<void> {
    if (!this.client) {
      throw new Error('Client not initialized');
    }
    await this.client.models.generateContent({
      model: this.defaultModel,
      contents: [{ role: 'user', parts: [{ text: VALIDATION_PROMPT }] }],
      config: { maxOutputTokens: 1 }
    });
  }

  protected override providerErrorMessages(): Partial<Record<number, string>> {
    return {
      401: 'Authentication failed with Gemini API. Please check your API key and permissions.',
      403: 'Authentication failed with Gemini API. Please check your API key and permissions.',
      404: 'Gemini API endpoint not found. Please check the model name and API version.',
      429: 'Rate limit exceeded. Please reduce request frequency.',
      500: 'Gemini service temporarily unavailable. Please try again in a few moments.',
      502: 'Gemini service temporarily unavailable. Please try again in a few moments.',
      503: 'Gemini service temporarily unavailable. Please try again in a few moments.'
    };
  }

  protected override mapProviderError(error: any): AIProviderError | null {
    if (!error) return null;
    const message: string = error.message || '';
    const status = error.status || error.statusCode;

    if (message.includes('API key')) {
      return new AIProviderError(
        'Invalid Google API key. Please verify your API key from https://aistudio.google.com/app/apikey',
        AIErrorType.AUTHENTICATION, status, error
      );
    }

    if (message.includes('quota') || message.includes('limit exceeded')) {
      return new AIProviderError(
        'API quota exceeded. Please check your Google Cloud quotas and billing.',
        AIErrorType.RATE_LIMIT, status, error
      );
    }

    if (message.includes('model not found') || message.includes('invalid model')) {
      return new AIProviderError(
        'Model not found. The specified Gemini model may not exist or be available to your API key.',
        AIErrorType.MODEL_NOT_FOUND, status, error
      );
    }

    if (message.includes('safety') || message.includes('blocked')) {
      return new AIProviderError(
        'Content blocked by safety filters. Please modify your input or adjust safety settings.',
        AIErrorType.INVALID_REQUEST, status, error
      );
    }

    if (message.includes('length') || message.includes('token')) {
      return new AIProviderError(
        'Input too long or exceeds token limits. Please reduce input size or max tokens.',
        AIErrorType.INVALID_REQUEST, status, error
      );
    }

    return null;
  }

  private convertMessages(messages: AIMessage[]): ProcessedMessages {
    let systemInstruction: string | undefined;
    const contents: Content[] = [];

    for (const message of messages) {
      if (message.role === 'system') {
        systemInstruction = systemInstruction
          ? `${systemInstruction}\n\n${message.content}`
          : message.content;
      } else {
        contents.push({
          role: message.role === 'assistant' ? 'model' : 'user',
          parts: [{ text: message.content }]
        });
      }
    }

    return { systemInstruction, contents };
  }

  private buildConfig(
    params: CompletionParams,
    systemInstruction?: string
  ): GenerateContentConfig {
    const defaults = this.defaultGenerationConfig;

    const config: GenerateContentConfig = {
      temperature: params.temperature ?? defaults?.temperature ?? DEFAULT_TEMPERATURE,
      maxOutputTokens: params.maxTokens ?? defaults?.maxOutputTokens ?? DEFAULT_MAX_TOKENS
    };

    const topP = params.topP ?? defaults?.topP;
    if (topP !== undefined) config.topP = topP;

    if (defaults?.topK !== undefined) config.topK = defaults.topK;

    const stopSequences = params.stopSequences ?? defaults?.stopSequences;
    if (stopSequences) config.stopSequences = stopSequences;

    if (systemInstruction) config.systemInstruction = systemInstruction;
    if (this.safetySettings) config.safetySettings = this.safetySettings;

    return config;
  }

  private async *processStreamChunks(
    stream: AsyncGenerator<GenerateContentResponse>
  ): AsyncIterable<CompletionChunk> {
    const requestId = `gemini-${Date.now()}-${Math.random().toString(36).slice(2, 11)}`;
    let lastUsage: GenerateContentResponse['usageMetadata'] | undefined;

    try {
      for await (const chunk of stream) {
        if (chunk.usageMetadata) {
          lastUsage = chunk.usageMetadata;
        }

        const text = chunk.text;
        if (text) {
          yield {
            content: text,
            isComplete: false,
            id: requestId
          };
        }
      }

      yield {
        content: '',
        isComplete: true,
        id: requestId,
        usage: {
          promptTokens: lastUsage?.promptTokenCount || 0,
          completionTokens: lastUsage?.candidatesTokenCount || 0,
          totalTokens: lastUsage?.totalTokenCount || 0
        }
      };
    } catch (error) {
      throw new AIProviderError(
        `Streaming interrupted: ${(error as Error).message}`,
        AIErrorType.NETWORK,
        undefined,
        error as Error
      );
    }
  }

  private formatCompletionResponse(
    response: GenerateContentResponse,
    model: string
  ): CompletionResponse<string> {
    const text = response.text;

    if (!text) {
      throw new AIProviderError(
        'No text content found in Gemini response',
        AIErrorType.UNKNOWN
      );
    }

    const candidate = response.candidates?.[0];

    return {
      content: text,
      model,
      usage: {
        promptTokens: response.usageMetadata?.promptTokenCount || 0,
        completionTokens: response.usageMetadata?.candidatesTokenCount || 0,
        totalTokens: response.usageMetadata?.totalTokenCount || 0
      },
      id: `gemini-${Date.now()}`,
      metadata: {
        finishReason: candidate?.finishReason,
        safetyRatings: candidate?.safetyRatings,
        citationMetadata: candidate?.citationMetadata
      }
    };
  }

}