wip

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

@@ -0,0 +1,30 @@
+import { LLM_CONFIG } from "../../config/llm.js";
+import { OpenAiCompatibleAdapter } from "./openai-compatible.js";
+import { GeminiAdapter } from "./gemini.js";
+import type { LlmAdapter } from "./types.js";
+
+export function createAdapter(): LlmAdapter {
+  switch (LLM_CONFIG.provider) {
+    case "local":
+      return new OpenAiCompatibleAdapter(LLM_CONFIG.url);
+    case "openrouter":
+      return new OpenAiCompatibleAdapter(
+        LLM_CONFIG.url,
+        process.env["OPENROUTER_API_KEY"],
+        LLM_CONFIG.model,
+      );
+    case "deepseek":
+      return new OpenAiCompatibleAdapter(
+        LLM_CONFIG.url,
+        process.env["DEEPSEEK_API_KEY"],
+        LLM_CONFIG.model,
+      );
+    case "gemini": {
+      const apiKey = process.env["GEMINI_API_KEY"];
+      if (!apiKey) throw new Error("GEMINI_API_KEY env var not set");
+      if (!LLM_CONFIG.model)
+        throw new Error("LLM_CONFIG.model required for gemini");
+      return new GeminiAdapter(apiKey, LLM_CONFIG.model);
+    }
+  }
+}

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),
+    };
+  }
+}

data-pipeline/utils/llm-adapters/openai-compatible.ts

@@ -0,0 +1,92 @@
+import type { LlmAdapter } from "./types.js";
+
+interface OpenAiResponse {
+  choices: Array<{ message: { content: string } }>;
+  usage: {
+    prompt_tokens: number;
+    completion_tokens: number;
+    total_tokens: number;
+  };
+  timings: {
+    prompt_ms: number;
+    predicted_ms: number;
+    prompt_per_second: number;
+    predicted_per_second: number;
+  };
+}
+
+export class OpenAiCompatibleAdapter implements LlmAdapter {
+  private url: string;
+  private apiKey: string | undefined;
+  private model: string | undefined;
+
+  constructor(url: string, apiKey?: string, model?: string) {
+    this.url = url;
+    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 payload: Record<string, unknown> = {
+      messages: [
+        { role: "system", content: systemPrompt },
+        { role: "user", content: JSON.stringify(words) },
+      ],
+      temperature: 0.1,
+      top_p: 0.9,
+      max_tokens: Math.ceil(words.length * 250 * 1.2),
+    };
+
+    if (this.model) {
+      payload["model"] = this.model;
+    }
+
+    const headers: Record<string, string> = {
+      "Content-Type": "application/json",
+    };
+
+    if (this.apiKey) {
+      headers["Authorization"] = `Bearer ${this.apiKey}`;
+    }
+
+    const response = await fetch(this.url, {
+      method: "POST",
+      headers,
+      body: JSON.stringify(payload),
+    });
+
+    if (!response.ok) {
+      throw new Error(`LLM server responded with status: ${response.status}`);
+    }
+
+    const json = (await response.json()) as OpenAiResponse;
+
+    const content = json.choices[0]?.message?.content;
+    if (!content) {
+      throw new Error("LLM response content is empty");
+    }
+
+    return {
+      content,
+      promptTokens: json.usage.prompt_tokens,
+      completionTokens: json.usage.completion_tokens,
+      totalTokens: json.usage.total_tokens,
+      promptTimeMs: json.timings.prompt_ms,
+      completionTimeMs: json.timings.predicted_ms,
+      promptTokensPerSecond: json.timings.prompt_per_second,
+      completionTokensPerSecond: json.timings.predicted_per_second,
+    };
+  }
+}

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

@@ -0,0 +1,15 @@
+export interface LlmAdapter {
+  call(
+    words: string[],
+    systemPrompt: string,
+  ): Promise<{
+    content: string;
+    promptTokens: number;
+    completionTokens: number;
+    totalTokens: number;
+    promptTimeMs: number;
+    completionTimeMs: number;
+    promptTokensPerSecond: number;
+    completionTokensPerSecond: number;
+  }>;
+}