wip

data-pipeline/utils/llm-adapters/gemini.ts

@@ -0,0 +1,92 @@
+import type { LlmAdapter } from "./types.js";
+
+interface GeminiResponse {
+  candidates: Array<{ content: { parts: Array<{ text: string }> } }>;
+  usageMetadata: {
+    promptTokenCount: number;
+    candidatesTokenCount: number;
+    totalTokenCount: number;
+  };
+}
+
+export class GeminiAdapter implements LlmAdapter {
+  private apiKey: string;
+  private model: string;
+
+  constructor(apiKey: string, model: string) {
+    this.apiKey = apiKey;
+    this.model = model;
+  }
+
+  async call(
+    words: string[],
+    systemPrompt: string,
+  ): Promise<{
+    content: string;
+    promptTokens: number;
+    completionTokens: number;
+    totalTokens: number;
+    promptTimeMs: number;
+    completionTimeMs: number;
+    promptTokensPerSecond: number;
+    completionTokensPerSecond: number;
+  }> {
+    const url = `https://generativelanguage.googleapis.com/v1beta/models/${this.model}:generateContent?key=${this.apiKey}`;
+
+    const payload = {
+      contents: [
+        {
+          role: "user",
+          parts: [
+            { text: systemPrompt + "\n\nWords: " + JSON.stringify(words) },
+          ],
+        },
+      ],
+      generationConfig: {
+        temperature: 0.1,
+        topP: 0.9,
+        maxOutputTokens: Math.ceil(words.length * 250 * 1.2),
+      },
+    };
+
+    const startTime = Date.now();
+
+    const response = await fetch(url, {
+      method: "POST",
+      headers: { "Content-Type": "application/json" },
+      body: JSON.stringify(payload),
+    });
+
+    const totalTimeMs = Date.now() - startTime;
+
+    if (!response.ok) {
+      throw new Error(`Gemini API responded with status: ${response.status}`);
+    }
+
+    const json = (await response.json()) as GeminiResponse;
+
+    const content = json.candidates[0]?.content?.parts[0]?.text;
+    if (!content) {
+      throw new Error("Gemini response content is empty");
+    }
+
+    const promptTokens = json.usageMetadata.promptTokenCount;
+    const completionTokens = json.usageMetadata.candidatesTokenCount;
+    const totalTokens = json.usageMetadata.totalTokenCount;
+
+    // Gemini doesn't provide timing breakdown, so we estimate
+    const promptTimeMs = totalTimeMs * 0.3; // rough estimate
+    const completionTimeMs = totalTimeMs * 0.7; // rough estimate
+
+    return {
+      content,
+      promptTokens,
+      completionTokens,
+      totalTokens,
+      promptTimeMs,
+      completionTimeMs,
+      promptTokensPerSecond: promptTokens / (promptTimeMs / 1000),
+      completionTokensPerSecond: completionTokens / (completionTimeMs / 1000),
+    };
+  }
+}