wip

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