Save input/output tokens in AiLLM

TrafficGuard · Dec 25, 2024 · 21a8f80 · 21a8f80
1 parent 5b8bd1d
commit 21a8f80
Showing 1 changed file with 2 additions and 4 deletions.
diff --git a/src/llm/services/ai-llm.ts b/src/llm/services/ai-llm.ts
@@ -92,10 +92,6 @@ export abstract class AiLLM<Provider extends ProviderV1> extends BaseLLM {
 				const finishTime = Date.now();
 				const llmCall: LlmCall = await llmCallSave;
 
-				// TODO calculate costs from response tokens
-				result.usage.totalTokens;
-				result.usage.promptTokens;
-				result.usage.completionTokens;
 				const inputCost = this.calculateInputCost(prompt);
 				const outputCost = this.calculateOutputCost(responseText);
 				const cost = inputCost + outputCost;
@@ -104,6 +100,8 @@ export abstract class AiLLM<Provider extends ProviderV1> extends BaseLLM {
 				llmCall.timeToFirstToken = null; // Not available in this implementation
 				llmCall.totalTime = finishTime - requestTime;
 				llmCall.cost = cost;
+				llmCall.inputTokens = result.usage.promptTokens;
+				llmCall.outputTokens = result.usage.completionTokens;
 				addCost(cost);
 
 				span.setAttributes({