Skip to content

fix(cost-estimation): fix zero output tokens #913

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
steebchen merged 3 commits into from
Sep 22, 2025
Merged

fix(cost-estimation): fix zero output tokens #913

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
5f1c0cc
feat(costs): include tool results in token and cost calculations
steebchen Sep 22, 2025
56152b5
Merge branch 'main' into terragon/fix-cost-estimation-zero-output
steebchen Sep 22, 2025
4b7b3c1
test(gateway): add type annotations for toolResults in cost calculati…
steebchen Sep 22, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 2 additions & 0 deletions apps/gateway/src/chat/chat.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2647,6 +2647,7 @@ chat.openapi(completions, async (c) => {
{
prompt: messages.map((m) => m.content).join("\n"),
completion: fullContent,
toolResults: streamingToolCalls || undefined,
},
);

Expand Down Expand Up @@ -3033,6 +3034,7 @@ chat.openapi(completions, async (c) => {
{
prompt: messages.map((m) => m.content).join("\n"),
completion: content,
toolResults: toolResults,
},
Comment on lines 3035 to 3038

@coderabbitai coderabbitai Bot Sep 22, 2025

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue

Non‑streaming path parity — good; consider redacting secrets before logging

You’re also attaching toolResults to logs elsewhere; tool args may contain secrets/PII. Redact keys like token, authorization, apiKey, password before persisting.

Example helper (place in a shared util and use before insertLog/calculateCosts if desired):

const SENSITIVE_KEYS = /^(authorization|api[-_]?key|token|access[_-]?token|secret|password)$/i;
function redactDeep(value: unknown): unknown {
  if (Array.isArray(value)) return value.map(redactDeep);
  if (value && typeof value === "object") {
    return Object.fromEntries(
      Object.entries(value as Record<string, unknown>).map(([k, v]) => [
        k,
        SENSITIVE_KEYS.test(k) ? "***" : redactDeep(v),
      ]),
    );
  }
  return value;
}
// usage: toolResults?.map(tr => ({ ...tr, function: tr.function ? { ...tr.function, arguments: typeof tr.function.arguments === 'string' ? tr.function.arguments : redactDeep(tr.function.arguments) } : tr.function }))
🤖 Prompt for AI Agents 
In apps/gateway/src/chat/chat.ts around lines 3041 to 3044, toolResults are
being attached to logs and may contain secrets/PII; before persisting or calling
insertLog/calculateCosts, deep-redact sensitive keys (authorization, apiKey,
token, access_token, secret, password, etc.) from toolResults and any nested
objects or arrays. Add or reuse a shared util (redactDeep) that recursively
traverses arrays/objects, replaces matching keys with a placeholder (e.g.,
"***"), and apply it to toolResults and to tool.function.arguments (handling
string vs structured args) so logs never store raw secrets. Ensure the redaction
runs in the non-streaming and streaming paths just prior to any log or storage
call.

);

Expand Down
79 changes: 79 additions & 0 deletions apps/gateway/src/lib/costs.spec.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -171,4 +171,83 @@ 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",
type: "function",
function: {
name: "get_weather",
arguments: '{"location": "San Francisco"}',
},
},
{
id: "call_2",
type: "function",
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", 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"}',
},
},
],
});

// Should still work with partial tool result data
expect(result.promptTokens).toBeGreaterThan(0);
expect(result.completionTokens).toBeGreaterThan(0);
expect(result.estimatedCost).toBe(true);
});
});
52 changes: 41 additions & 11 deletions apps/gateway/src/lib/costs.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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 {
Expand All @@ -27,6 +32,7 @@ export function calculateCosts(
messages?: ChatMessage[];
prompt?: string;
completion?: string;
toolResults?: ToolCall[];
},
) {
// Find the model info - try both base model name and provider model name
Expand Down Expand Up @@ -89,20 +95,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,
Expand All @@ -117,6 +142,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,
Expand Down