All files / src/gemini gemini-llm.provider.ts

26.08% Statements 12/46
0% Branches 0/17
30% Functions 3/10
23.25% Lines 10/43

Press n or j to go to the next uncovered block, b, p or k for the previous block.

1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 1576x             6x         6x           6x                                         6x     6x     6x             6x       6x 6x                                                                                                                                                                                                      
import { Injectable, OnModuleInit } from '@nestjs/common';
import {
  LlmProvider,
  LlmProviderRequest,
  LlmModel,
  LlmResponse,
} from '../llm-provider/llm-provider.interface';
import {
  Config,
  AuthType,
  GeminiChat,
} from '../../packages/gemini-core/src/index';
import { randomUUID } from 'crypto';
import {
  GenerateContentResponse,
  GenerateContentConfig,
  Content,
} from '@google/genai';
import { TiktokenTokenizer } from '../../packages/tokenpatch/strategies/tiktoken-tokenizer';
 
function getResponseText(response: GenerateContentResponse): string | null {
  Iif (response.candidates && response.candidates.length > 0) {
    const candidate = response.candidates[0];
 
    Iif (
      candidate.content &&
      candidate.content.parts &&
      candidate.content.parts.length > 0
    ) {
      return candidate.content.parts
        .filter((part: any) => part.text)
        .map((part: any) => part.text)
        .join('');
    }
  }
  return null;
}
 
@Injectable()
export class GeminiLlmProvider implements LlmProvider, OnModuleInit {
  private config: Config;
  private initializationPromise: Promise<void>;
  private readonly tokenizer = new TiktokenTokenizer();
 
  constructor() {
    this.config = new Config({
      sessionId: `repoburg-backend-session-${randomUUID()}`,
      model: 'gemini-2.0-flash',
    });
  }
 
  onModuleInit() {
    this.initializationPromise = this.initialize();
  }
 
  private async initialize(): Promise<void> {
    try {
      await this.config.refreshAuth(AuthType.LOGIN_WITH_GOOGLE);
    } catch (error) {
      // Initialization failed - Gemini might not work
    }
  }
 
  async getModels(): Promise<LlmModel[]> {
    return [
      {
        id: 'gemini-2.0-flash',
        name: 'Gemini 2.0 Flash (Local)',
        provider: 'gemini',
      },
      {
        id: 'gemini-2.0-pro-exp-02-05',
        name: 'Gemini 2.0 Pro Exp (Local)',
        provider: 'gemini',
      },
      {
        id: 'gemini-1.5-pro',
        name: 'Gemini 1.5 Pro (Local)',
        provider: 'gemini',
      },
      {
        id: 'gemini-1.5-flash',
        name: 'Gemini 1.5 Flash (Local)',
        provider: 'gemini',
      },
    ];
  }
 
  async generateContent(request: LlmProviderRequest): Promise<LlmResponse> {
    try {
      await this.initializationPromise;
    } catch (error) {
      throw new Error(
        'Gemini LLM Provider failed to initialize. Please check Google authentication.',
      );
    }
 
    const { prompt, systemInstruction, history, modelId, generationConfig } =
      request;
 
    // Remove prefix if exists for the core generator
    const model = (modelId || this.config.getModel()).replace('gemini/', '');
 
    // Estimate input tokens
    let inputTokenText = (systemInstruction || '') + '\n' + prompt;
    Iif (history) {
      inputTokenText += history
        .map((h) => h.parts.map((p) => p.text).join(''))
        .join('\n');
    }
    const inputTokenCount = this.tokenizer.tokenize(inputTokenText).length;
 
    try {
      const finalGenerationConfig: GenerateContentConfig = {
        ...generationConfig,
      };
      Iif (systemInstruction) {
        finalGenerationConfig.systemInstruction = systemInstruction;
      }
 
      const chat = new GeminiChat(
        this.config,
        finalGenerationConfig,
        history as Content[],
      );
 
      const result = await chat.sendMessageStream(
        model,
        { message: prompt },
        `prompt-${randomUUID()}`,
      );
 
      let fullResponse = '';
      for await (const chunk of result) {
        Iif (chunk.type === 'chunk' && chunk.value) {
          const text = getResponseText(chunk.value);
          Iif (text) {
            fullResponse += text;
          }
        }
      }
 
      const outputTokenCount = this.tokenizer.tokenize(fullResponse).length;
 
      return {
        text: fullResponse,
        usage: {
          inputTokens: inputTokenCount,
          outputTokens: outputTokenCount,
        },
      };
    } catch (error) {
      throw new Error(`Gemini Error: ${error.message}`);
    }
  }
}