bugfixing #1

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

@@ -28,8 +28,6 @@ export class GeminiAdapter implements LlmAdapter {
     totalTokens: number;
     promptTimeMs: number;
     completionTimeMs: number;
-    promptTokensPerSecond: number;
-    completionTokensPerSecond: number;
   }> {
     const url = `https://generativelanguage.googleapis.com/v1beta/models/${this.model}:generateContent?key=${this.apiKey}`;
 
@@ -72,21 +70,14 @@ export class GeminiAdapter implements LlmAdapter {
 
     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),
+      totalTokens: json.usageMetadata.totalTokenCount,
+      promptTimeMs: totalTimeMs * 0.3,
+      completionTimeMs: totalTimeMs * 0.7,
     };
   }
 }

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

@@ -7,12 +7,7 @@ interface OpenAiResponse {
     completion_tokens: number;
     total_tokens: number;
   };
-  timings: {
-    prompt_ms: number;
-    predicted_ms: number;
-    prompt_per_second: number;
-    predicted_per_second: number;
-  };
+  timings?: { prompt_ms: number; predicted_ms: number };
 }
 
 export class OpenAiCompatibleAdapter implements LlmAdapter {
@@ -36,8 +31,6 @@ export class OpenAiCompatibleAdapter implements LlmAdapter {
     totalTokens: number;
     promptTimeMs: number;
     completionTimeMs: number;
-    promptTokensPerSecond: number;
-    completionTokensPerSecond: number;
   }> {
     const payload: Record<string, unknown> = {
       messages: [
@@ -61,12 +54,16 @@ export class OpenAiCompatibleAdapter implements LlmAdapter {
       headers["Authorization"] = `Bearer ${this.apiKey}`;
     }
 
+    const startTime = Date.now();
+
     const response = await fetch(this.url, {
       method: "POST",
       headers,
       body: JSON.stringify(payload),
     });
 
+    const totalTimeMs = Date.now() - startTime;
+
     if (!response.ok) {
       throw new Error(`LLM server responded with status: ${response.status}`);
     }
@@ -78,15 +75,19 @@ export class OpenAiCompatibleAdapter implements LlmAdapter {
       throw new Error("LLM response content is empty");
     }
 
+    const promptTokens = json.usage.prompt_tokens;
+    const completionTokens = json.usage.completion_tokens;
+
+    const promptTimeMs = json.timings?.prompt_ms ?? totalTimeMs * 0.3;
+    const completionTimeMs = json.timings?.predicted_ms ?? totalTimeMs * 0.7;
+
     return {
       content,
-      promptTokens: json.usage.prompt_tokens,
-      completionTokens: json.usage.completion_tokens,
+      promptTokens,
+      completionTokens,
       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,
+      promptTimeMs,
+      completionTimeMs,
     };
   }
 }

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

@@ -9,7 +9,5 @@ export interface LlmAdapter {
     totalTokens: number;
     promptTimeMs: number;
     completionTimeMs: number;
-    promptTokensPerSecond: number;
-    completionTokensPerSecond: number;
   }>;
 }