All files / src/providers gemini.ts

92% Statements 46/50
59.25% Branches 16/27
100% Functions 14/14
95.83% Lines 46/48

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 157    1x             1x         4x   4x 3x 3x 3x   3x       3x 3x 3x 3x   3x 3x 2x 2x   1x 1x       3x     4x                     2x 2x 2x 2x 2x 2x 2x   2x                                                   1x 1x                                     2x 2x 2x       2x 2x   2x 2x 2x             3x 3x             9x       3x       5x       5x       4x       1x    
import type { Batcher, GeminiClientLike, LLMCallRecord, UnknownFunction } from '../types.js';
 
const GEMINI_COST_PER_1K: Record<string, { input: number; output: number }> = {
  'gemini-2.5-pro': { input: 0.00125, output: 0.01 },
  'gemini-2.5-flash': { input: 0.000075, output: 0.0003 },
  'gemini-1.5-pro': { input: 0.00125, output: 0.005 },
  'gemini-1.5-flash': { input: 0.000075, output: 0.0003 },
};
 
export function wrapGemini<TClient extends GeminiClientLike>(
  client: TClient,
  batcher: Batcher,
  projectId: string | null,
): TClient {
  const originalGetGenerativeModel = client.getGenerativeModel.bind(client);
 
  client.getGenerativeModel = (params: unknown, ...rest: unknown[]): unknown => {
    const model = originalGetGenerativeModel(params, ...rest);
    const modelRecord = asMutableRecord(model);
    const originalGenerate = getFunction(modelRecord?.generateContent);
 
    Iif (!modelRecord || !originalGenerate) {
      return model;
    }
 
    const modelName = getString(asRecord(params)?.model) ?? 'unknown';
    modelRecord.generateContent = async (...args: unknown[]): Promise<unknown> => {
      const start = Date.now();
      const timestamp = new Date(start).toISOString();
 
      try {
        const response = await originalGenerate.apply(model, args);
        pushGeminiSuccess(response, modelName, batcher, projectId, timestamp, Date.now() - start);
        return response;
      } catch (error) {
        pushGeminiError(error, modelName, batcher, projectId, timestamp, Date.now() - start);
        throw error;
      }
    };
 
    return model;
  };
 
  return client;
}
 
function pushGeminiSuccess(
  response: unknown,
  model: string,
  batcher: Batcher,
  projectId: string | null,
  timestamp: string,
  latencyMs: number,
): void {
  try {
    const responseRecord = asRecord(response);
    const innerResponse = asRecord(responseRecord?.response);
    const usageMetadata = asRecord(innerResponse?.usageMetadata);
    const promptTokens = getNumber(usageMetadata?.promptTokenCount);
    const completionTokens = getNumber(usageMetadata?.candidatesTokenCount);
    const preview = getGeminiPreview(innerResponse);
 
    pushSafely(batcher, {
      project_id: projectId,
      timestamp,
      provider: 'gemini',
      model,
      prompt_tokens: promptTokens,
      completion_tokens: completionTokens,
      latency_ms: latencyMs,
      cost_usd: computeCost(model, promptTokens, completionTokens),
      response_preview: preview,
      error: null,
      metadata: {},
    });
  } catch {
    // Logging must never affect the provider call.
  }
}
 
function pushGeminiError(
  error: unknown,
  model: string,
  batcher: Batcher,
  projectId: string | null,
  timestamp: string,
  latencyMs: number,
): void {
  try {
    pushSafely(batcher, {
      project_id: projectId,
      timestamp,
      provider: 'gemini',
      model,
      prompt_tokens: 0,
      completion_tokens: 0,
      latency_ms: latencyMs,
      cost_usd: 0,
      response_preview: '',
      error: stringifyError(error),
      metadata: {},
    });
  } catch {
    // Preserve the original provider error exactly.
  }
}
 
function computeCost(model: string, promptTokens: number, completionTokens: number): number {
  const pricing = GEMINI_COST_PER_1K[model];
  Iif (!pricing) return 0;
  return (promptTokens / 1000) * pricing.input + (completionTokens / 1000) * pricing.output;
}
 
function getGeminiPreview(innerResponse: Record<string, unknown> | undefined): string {
  const text = getFunction(innerResponse?.text);
  Iif (!text) return '';
 
  try {
    const value = text.apply(innerResponse);
    return (getString(value) ?? '').slice(0, 200);
  } catch {
    return '';
  }
}
 
function pushSafely(batcher: Batcher, record: LLMCallRecord): void {
  try {
    batcher.push(record);
  } catch {
    // Explicitly isolate provider behavior from logging failures.
  }
}
 
function asRecord(value: unknown): Record<string, unknown> | undefined {
  return typeof value === 'object' && value !== null ? (value as Record<string, unknown>) : undefined;
}
 
function asMutableRecord(value: unknown): Record<string, unknown> | undefined {
  return typeof value === 'object' && value !== null ? (value as Record<string, unknown>) : undefined;
}
 
function getFunction(value: unknown): UnknownFunction | undefined {
  return typeof value === 'function' ? (value as UnknownFunction) : undefined;
}
 
function getString(value: unknown): string | undefined {
  return typeof value === 'string' ? value : undefined;
}
 
function getNumber(value: unknown): number {
  return typeof value === 'number' && Number.isFinite(value) ? value : 0;
}
 
function stringifyError(error: unknown): string {
  return error instanceof Error ? error.message : String(error);
}