From 5f1c0cc93db876d62c6751a302841de79689f425 Mon Sep 17 00:00:00 2001 From: Luca Steeb Date: Mon, 22 Sep 2025 20:26:13 +0000 Subject: [PATCH 1/2] feat(costs): include tool results in token and cost calculations - Add support for toolResults in calculateCosts to count tokens from tool calls - Update chat.ts to pass toolResults in completions - Add tests for toolResults handling and edge cases in costs.spec.ts - Handle null and zero token counts gracefully to ensure accurate cost calculation This enhancement allows more accurate cost estimation by including tokens from tool function calls in the calculation. Co-authored-by: terragon-labs[bot] --- apps/gateway/src/chat/chat.ts | 2 + apps/gateway/src/lib/costs.spec.ts | 76 ++++++++++++++++++++++++++++++ apps/gateway/src/lib/costs.ts | 45 ++++++++++++++---- 3 files changed, 113 insertions(+), 10 deletions(-) diff --git a/apps/gateway/src/chat/chat.ts b/apps/gateway/src/chat/chat.ts index 7b1fbe0fc2..a810d30d8d 100644 --- a/apps/gateway/src/chat/chat.ts +++ b/apps/gateway/src/chat/chat.ts @@ -2653,6 +2653,7 @@ chat.openapi(completions, async (c) => { { prompt: messages.map((m) => m.content).join("\n"), completion: fullContent, + toolResults: streamingToolCalls || undefined, }, ); @@ -3039,6 +3040,7 @@ chat.openapi(completions, async (c) => { { prompt: messages.map((m) => m.content).join("\n"), completion: content, + toolResults: toolResults, }, ); diff --git a/apps/gateway/src/lib/costs.spec.ts b/apps/gateway/src/lib/costs.spec.ts index c4b365e8b0..8037407983 100644 --- a/apps/gateway/src/lib/costs.spec.ts +++ b/apps/gateway/src/lib/costs.spec.ts @@ -171,4 +171,80 @@ describe("calculateCosts", () => { expect(result.discount).toBeUndefined(); // Should not include discount field when discount is 1 }); + + it("should calculate input costs even when output tokens are zero", () => { + const result = calculateCosts("gpt-4", "openai", 100, 0, null); + + expect(result.inputCost).toBeCloseTo(0.001); // 100 * 0.00001 + expect(result.outputCost).toBeCloseTo(0); // 0 * 0.00003 + expect(result.totalCost).toBeCloseTo(0.001); // 0.001 + 0 + expect(result.promptTokens).toBe(100); + expect(result.completionTokens).toBe(0); + expect(result.estimatedCost).toBe(false); + }); + + it("should calculate input costs when completion tokens are null but prompt tokens exist", () => { + const result = calculateCosts("gpt-4", "openai", 100, null, null); + + expect(result.inputCost).toBeCloseTo(0.001); // 100 * 0.00001 + expect(result.outputCost).toBeCloseTo(0); // 0 * 0.00003 (completion tokens set to 0) + expect(result.totalCost).toBeCloseTo(0.001); // 0.001 + 0 + expect(result.promptTokens).toBe(100); + expect(result.completionTokens).toBe(0); // Should default to 0 + expect(result.estimatedCost).toBe(false); + }); + + it("should include tool results in completion token estimation", () => { + const result = calculateCosts("gpt-4", "openai", null, null, null, { + prompt: "What's the weather like?", + completion: "", // Empty completion + toolResults: [ + { + id: "call_1", + function: { + name: "get_weather", + arguments: '{"location": "San Francisco"}', + }, + }, + { + id: "call_2", + function: { + name: "get_temperature", + arguments: '{"location": "New York", "units": "celsius"}', + }, + }, + ], + }); + + // Should calculate tokens for tool calls even with empty completion + expect(result.promptTokens).toBeGreaterThan(0); + expect(result.completionTokens).toBeGreaterThan(0); // Should include tool call tokens + expect(result.inputCost).toBeGreaterThan(0); + expect(result.outputCost).toBeGreaterThan(0); // Should have cost from tool calls + expect(result.totalCost).toBeGreaterThan(0); + expect(result.estimatedCost).toBe(true); + }); + + it("should handle tool results with missing function data gracefully", () => { + const result = calculateCosts("gpt-4", "openai", null, null, null, { + prompt: "What's the weather like?", + completion: "Here's the weather:", + toolResults: [ + { id: "call_1" }, // Missing function data + { id: "call_2", function: {} }, // Missing name and arguments + { + id: "call_3", + function: { + name: "get_weather", + arguments: '{"location": "Paris"}', + }, + }, + ], + }); + + // Should still work with partial tool result data + expect(result.promptTokens).toBeGreaterThan(0); + expect(result.completionTokens).toBeGreaterThan(0); + expect(result.estimatedCost).toBe(true); + }); }); diff --git a/apps/gateway/src/lib/costs.ts b/apps/gateway/src/lib/costs.ts index 8fbb6612f3..aa59f1f9f7 100644 --- a/apps/gateway/src/lib/costs.ts +++ b/apps/gateway/src/lib/costs.ts @@ -27,6 +27,7 @@ export function calculateCosts( messages?: ChatMessage[]; prompt?: string; completion?: string; + toolResults?: any[]; }, ) { // Find the model info - try both base model name and provider model name @@ -89,20 +90,39 @@ export function calculateCosts( } // Calculate completion tokens - if (!completionTokens && fullOutput && fullOutput.completion) { - try { - calculatedCompletionTokens = encode( - JSON.stringify(fullOutput.completion), - ).length; - } catch (error) { - // If encoding fails, leave as null - logger.error(`Failed to encode completion text: ${error}`); + if (!completionTokens && fullOutput) { + let completionText = ""; + + // Include main completion content + if (fullOutput.completion) { + completionText += fullOutput.completion; + } + + // Include tool results if available + if (fullOutput.toolResults && Array.isArray(fullOutput.toolResults)) { + for (const toolResult of fullOutput.toolResults) { + if (toolResult.function?.name) { + completionText += toolResult.function.name; + } + if (toolResult.function?.arguments) { + completionText += JSON.stringify(toolResult.function.arguments); + } + } + } + + if (completionText) { + try { + calculatedCompletionTokens = encode(completionText).length; + } catch (error) { + // If encoding fails, leave as null + logger.error(`Failed to encode completion text: ${error}`); + } } } } - // If we still don't have token counts, return null costs - if (!calculatedPromptTokens || !calculatedCompletionTokens) { + // If we don't have prompt tokens, we can't calculate any costs + if (!calculatedPromptTokens) { return { inputCost: null, outputCost: null, @@ -117,6 +137,11 @@ export function calculateCosts( }; } + // Set completion tokens to 0 if not available (but still calculate input costs) + if (!calculatedCompletionTokens) { + calculatedCompletionTokens = 0; + } + // Find the provider-specific pricing const providerInfo = modelInfo.providers.find( (p) => p.providerId === provider, From 4b7b3c1214ccc5b54a5dffc96d2bac9f11dfdfad Mon Sep 17 00:00:00 2001 From: Luca Steeb Date: Mon, 22 Sep 2025 21:37:39 +0000 Subject: [PATCH 2/2] test(gateway): add type annotations for toolResults in cost calculation tests - Added 'type: function' to toolResults entries in costs.spec.ts to improve test clarity and type safety. - Updated costs.ts to use ToolCall type for toolResults parameter, enhancing type correctness in cost calculations. Co-authored-by: terragon-labs[bot] --- apps/gateway/src/lib/costs.spec.ts | 7 +++++-- apps/gateway/src/lib/costs.ts | 9 +++++++-- 2 files changed, 12 insertions(+), 4 deletions(-) diff --git a/apps/gateway/src/lib/costs.spec.ts b/apps/gateway/src/lib/costs.spec.ts index 8037407983..4730ea4111 100644 --- a/apps/gateway/src/lib/costs.spec.ts +++ b/apps/gateway/src/lib/costs.spec.ts @@ -201,6 +201,7 @@ describe("calculateCosts", () => { toolResults: [ { id: "call_1", + type: "function", function: { name: "get_weather", arguments: '{"location": "San Francisco"}', @@ -208,6 +209,7 @@ describe("calculateCosts", () => { }, { id: "call_2", + type: "function", function: { name: "get_temperature", arguments: '{"location": "New York", "units": "celsius"}', @@ -230,10 +232,11 @@ describe("calculateCosts", () => { prompt: "What's the weather like?", completion: "Here's the weather:", toolResults: [ - { id: "call_1" }, // Missing function data - { id: "call_2", function: {} }, // Missing name and arguments + { id: "call_1", type: "function" } as any, // Missing function data + { id: "call_2", type: "function", function: {} as any }, // Missing name and arguments { id: "call_3", + type: "function", function: { name: "get_weather", arguments: '{"location": "Paris"}', diff --git a/apps/gateway/src/lib/costs.ts b/apps/gateway/src/lib/costs.ts index aa59f1f9f7..864db788fb 100644 --- a/apps/gateway/src/lib/costs.ts +++ b/apps/gateway/src/lib/costs.ts @@ -1,7 +1,12 @@ import { encode, encodeChat } from "gpt-tokenizer"; import { logger } from "@llmgateway/logger"; -import { type Model, type ModelDefinition, models } from "@llmgateway/models"; +import { + type Model, + type ModelDefinition, + models, + type ToolCall, +} from "@llmgateway/models"; // Define ChatMessage type to match what gpt-tokenizer expects interface ChatMessage { @@ -27,7 +32,7 @@ export function calculateCosts( messages?: ChatMessage[]; prompt?: string; completion?: string; - toolResults?: any[]; + toolResults?: ToolCall[]; }, ) { // Find the model info - try both base model name and provider model name