KGStudio/src/agent/llm/GeminiProvider.ts

import { LLMProvider } from './LLMProvider';
import type { StreamChunk } from './StreamingTypes';
import type { Message } from '../core/AgentState';
import { ConfigManager } from '../../core/config/ConfigManager';

/**
 * Google Gemini API provider implementation
 */
export class GeminiProvider extends LLMProvider {
  readonly name = 'Gemini';
  
  private apiKey: string;
  private model: string;
  private baseURL: string = 'https://generativelanguage.googleapis.com';
  private apiEndpoint: string;
  
  constructor() {
    super();
    
    const configManager = ConfigManager.instance();
    this.apiKey = configManager.get('general.gemini.api_key') as string;
    this.model = configManager.get('general.gemini.model') as string;
    this.apiEndpoint = `${this.baseURL}/v1beta/models/${this.model}`;
  }

  /**
   * Convert internal messages to Gemini's format
   */
  private convertMessages(messages: Message[], systemPrompt?: string): {
    systemInstruction?: { parts: Array<{ text: string }> };
    contents: Array<{ role: 'user' | 'model'; parts: Array<{ text: string }> }>;
  } {
    const geminiMessages: Array<{ role: 'user' | 'model'; parts: Array<{ text: string }> }> = [];
    
    for (const msg of messages) {
      let role: 'user' | 'model';
      
      if (msg.role === 'assistant') {
        role = 'model';
      } else {
        // Treat system and user messages as 'user' role
        role = 'user';
      }
      
      geminiMessages.push({
        role,
        parts: [{ text: msg.content }]
      });
    }
    
    const result: {
      systemInstruction?: { parts: Array<{ text: string }> };
      contents: Array<{ role: 'user' | 'model'; parts: Array<{ text: string }> }>;
    } = {
      contents: geminiMessages
    };
    
    if (systemPrompt) {
      result.systemInstruction = {
        parts: [{ text: systemPrompt }]
      };
    }
    
    return result;
  }

  /**
   * Parse Gemini's streaming response chunks
   */
  private parseGeminiStreamChunk(chunk: string): { content?: string; isDone?: boolean } {
    try {
      const json = JSON.parse(chunk.trim());
      
      // Gemini streaming format
      if (json.candidates && json.candidates.length > 0) {
        const candidate = json.candidates[0];
        
        // Check if generation is finished
        if (candidate.finishReason && candidate.finishReason !== 'STOP') {
          return { isDone: true };
        }
        
        // Extract text content
        const content = candidate.content?.parts?.[0]?.text;
        if (content) {
          return { content, isDone: false };
        }
      }
      
      // Check for explicit done signal
      if (json.done === true) {
        return { isDone: true };
      }
      
      return {};
    } catch {
      return {}; // Skip invalid JSON
    }
  }
  
  async *generateStream(
    messages: Message[], 
    systemPrompt?: string,
    tools?: Record<string, unknown>[]
  ): AsyncIterableIterator<StreamChunk> {
    const { systemInstruction, contents } = this.convertMessages(messages, systemPrompt);
    
    const requestBody: {
      contents: Array<{ role: 'user' | 'model'; parts: Array<{ text: string }> }>;
      generationConfig: { temperature: number; maxOutputTokens: number };
      systemInstruction?: { parts: Array<{ text: string }> };
      tools?: Record<string, unknown>[];
    } = {
      contents,
      generationConfig: {
        temperature: 0.7,
        maxOutputTokens: 8192,
      }
    };
    
    if (systemInstruction) {
      requestBody.systemInstruction = systemInstruction;
    }
    
    if (tools && tools.length > 0) {
      requestBody.tools = tools;
    }

    const response = await fetch(`${this.apiEndpoint}:streamGenerateContent?key=${this.apiKey}`, {
      method: 'POST',
      headers: {
        'Content-Type': 'application/json',
      },
      body: JSON.stringify(requestBody),
    });
    
    if (!response.ok) {
      throw new Error(`Gemini API error: ${response.status} ${response.statusText}`);
    }
    
    const reader = response.body?.getReader();
    if (!reader) {
      throw new Error('Failed to get response reader');
    }
    
    const decoder = new TextDecoder();
    let buffer = '';
    
    try {
      while (true) {
        const { done, value } = await reader.read();
        if (done) break;
        
        buffer += decoder.decode(value, { stream: true });
        
        // Gemini sends JSON objects separated by newlines
        const lines = buffer.split('\n');
        buffer = lines.pop() || '';
        
        for (const line of lines) {
          const trimmedLine = line.trim();
          if (!trimmedLine) continue;
          
          const parseResult = this.parseGeminiStreamChunk(trimmedLine);
          
          if (parseResult.isDone) {
            yield { type: 'done', content: '' };
            return;
          }
          
          if (parseResult.content) {
            yield {
              type: 'text',
              content: parseResult.content
            };
          }
        }
      }
    } finally {
      reader.releaseLock();
    }
  }
}