Skip to content

feat(completions): ensure prompt tokens are always calculated #588

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 6 commits into from
Aug 8, 2025
Merged

feat(completions): ensure prompt tokens are always calculated #588

Show file tree
Hide file tree
Changes from 2 commits
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
f730412
Ensure prompt tokens are never zero, add fallback token calculation
cursoragent Aug 7, 2025
3558f62
fix: ensure prompt tokens are never zero and manually calculated when…
cursoragent Aug 7, 2025
2d58105
chore: gen
steebchen Aug 7, 2025
d1a1fe7
fix(api): replace mock token with real token in e2e tests
steebchen Aug 7, 2025
44670c5
fix(models): remove test skip flags for Zai providers
steebchen Aug 7, 2025
14becd8
Merge branch 'main' into cursor/ensure-prompt-tokens-are-always-calcu…
steebchen Aug 8, 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
77 changes: 77 additions & 0 deletions apps/gateway/src/api.e2e.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -991,6 +991,83 @@ describe("e2e", () => {
expect(logs[0].usedModel).toBe("custom");
});

test("Prompt tokens are never zero even when provider returns 0", async () => {
const res = await app.request("/v1/chat/completions", {
method: "POST",
headers: {
"Content-Type": "application/json",
Authorization: `Bearer mock-token`,
},
body: JSON.stringify({
model: "openai/gpt-4o-mini",
messages: [
{
role: "user",
content: "ZERO_TOKENS test message",
},
],
}),
});

expect(res.status).toBe(200);

const json = await res.json();

// Verify we have usage information
expect(json).toHaveProperty("usage");
expect(json.usage).toHaveProperty("prompt_tokens");
expect(json.usage).toHaveProperty("completion_tokens");
expect(json.usage).toHaveProperty("total_tokens");

// Verify types are numbers
expect(typeof json.usage.prompt_tokens).toBe("number");
expect(typeof json.usage.completion_tokens).toBe("number");
expect(typeof json.usage.total_tokens).toBe("number");

// Most importantly: prompt_tokens should never be 0, even if provider returns 0
expect(json.usage.prompt_tokens).toBeGreaterThan(0);

// Completion tokens can be non-zero as set by mock
expect(json.usage.completion_tokens).toBeGreaterThan(0);

// Total tokens should be at least as large as prompt tokens
expect(json.usage.total_tokens).toBeGreaterThanOrEqual(json.usage.prompt_tokens);
});

test("Prompt tokens are calculated for streaming when provider returns 0", async () => {
const res = await app.request("/v1/chat/completions", {
method: "POST",
headers: {
"Content-Type": "application/json",
Authorization: `Bearer mock-token`,
},
body: JSON.stringify({
model: "openai/gpt-4o-mini",
messages: [
{
role: "user",
content: "ZERO_TOKENS streaming test message",
},
],
stream: true,
}),
});

expect(res.status).toBe(200);

const result = await readAll(res.body);

// Find a usage chunk
const usageChunk = result.chunks.find((chunk: any) => chunk.usage);
expect(usageChunk).toBeDefined();

if (usageChunk) {
// Verify prompt tokens are calculated and greater than 0
expect(usageChunk.usage.prompt_tokens).toBeGreaterThan(0);
expect(typeof usageChunk.usage.prompt_tokens).toBe("number");
}
});

test("Success when requesting multi-provider model without prefix", async () => {
const multiProviderModel = models.find((m) => m.providers.length > 1);
if (!multiProviderModel) {
Expand Down
66 changes: 46 additions & 20 deletions apps/gateway/src/chat/chat.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -648,9 +648,9 @@ function transformToOpenAIFormat(
},
],
usage: {
prompt_tokens: promptTokens,
completion_tokens: completionTokens,
total_tokens: totalTokens,
prompt_tokens: Math.max(1, promptTokens || 1),
completion_tokens: completionTokens || 0,
total_tokens: Math.max(1, totalTokens || Math.max(1, promptTokens || 1)),
...(reasoningTokens !== null && {
reasoning_tokens: reasoningTokens,
}),
Expand Down Expand Up @@ -689,9 +689,9 @@ function transformToOpenAIFormat(
},
],
usage: {
prompt_tokens: promptTokens,
completion_tokens: completionTokens,
total_tokens: totalTokens,
prompt_tokens: Math.max(1, promptTokens || 1),
completion_tokens: completionTokens || 0,
total_tokens: Math.max(1, totalTokens || Math.max(1, promptTokens || 1)),
...(reasoningTokens !== null && {
reasoning_tokens: reasoningTokens,
}),
Expand Down Expand Up @@ -727,9 +727,9 @@ function transformToOpenAIFormat(
},
],
usage: {
prompt_tokens: promptTokens,
completion_tokens: completionTokens,
total_tokens: totalTokens,
prompt_tokens: Math.max(1, promptTokens || 1),
completion_tokens: completionTokens || 0,
total_tokens: Math.max(1, totalTokens || Math.max(1, promptTokens || 1)),
Comment on lines 943 to -

@coderabbitai coderabbitai Bot Aug 7, 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.

🛠️ Refactor suggestion

Refactor repeated token normalization logic

The same token normalization pattern is repeated across multiple providers. Consider extracting this into a helper function to follow DRY principles.

Create a helper function:

function normalizeTokenUsage(
  promptTokens: number | null,
  completionTokens: number | null,
  totalTokens: number | null
) {
  const normalizedPrompt = Math.max(1, promptTokens || 0);
  const normalizedCompletion = completionTokens || 0;
  const normalizedTotal = Math.max(1, totalTokens || normalizedPrompt);
  
  return {
    prompt_tokens: normalizedPrompt,
    completion_tokens: normalizedCompletion,
    total_tokens: normalizedTotal,
  };
}

Then use it in all provider cases:

 usage: {
-  prompt_tokens: Math.max(1, promptTokens || 1),
-  completion_tokens: completionTokens || 0,
-  total_tokens: Math.max(1, totalTokens || Math.max(1, promptTokens || 1)),
+  ...normalizeTokenUsage(promptTokens, completionTokens, totalTokens),
   // ... other fields
 },
🤖 Prompt for AI Agents 
In apps/gateway/src/chat/chat.ts around lines 730 to 732, the token
normalization logic is repeated for multiple providers. Extract this repeated
logic into a single helper function named normalizeTokenUsage that takes
promptTokens, completionTokens, and totalTokens as parameters and returns an
object with normalized prompt_tokens, completion_tokens, and total_tokens
values. Replace all instances of the repeated normalization code with calls to
this helper function to adhere to DRY principles and improve maintainability.

...(reasoningTokens !== null && {
reasoning_tokens: reasoningTokens,
}),
Expand All @@ -746,10 +746,25 @@ function transformToOpenAIFormat(
/**
* Transforms streaming chunk to OpenAI format for non-OpenAI providers
*/
// Helper function to calculate prompt tokens when missing or 0
function calculatePromptTokensFromMessages(messages: any[]): number {
try {
const chatMessages: ChatMessage[] = messages.map((m: any) => ({
role: m.role,
content: typeof m.content === "string" ? m.content : JSON.stringify(m.content),
name: m.name,
}));
return encodeChat(chatMessages, DEFAULT_TOKENIZER_MODEL).length;
} catch (error) {
return Math.max(1, Math.round(messages.reduce((acc: number, m: any) => acc + (m.content?.length || 0), 0) / 4));
}
}
Comment thread
coderabbitai[bot] marked this conversation as resolved.
Resolved

function transformStreamingChunkToOpenAIFormat(
usedProvider: Provider,
usedModel: string,
data: any,
messages: any[],
): any {
let transformedData = data;

Expand Down Expand Up @@ -942,11 +957,16 @@ function transformStreamingChunkToOpenAIFormat(
],
usage: data.usageMetadata
? {
prompt_tokens: data.usageMetadata.promptTokenCount || 0,
prompt_tokens: data.usageMetadata.promptTokenCount > 0
? data.usageMetadata.promptTokenCount
: calculatePromptTokensFromMessages(messages),
completion_tokens: data.usageMetadata.candidatesTokenCount || 0,
// Calculate total including reasoning tokens for Google models
total_tokens:
(data.usageMetadata.promptTokenCount || 0) +
(data.usageMetadata.promptTokenCount > 0
? data.usageMetadata.promptTokenCount
: calculatePromptTokensFromMessages(messages)
) +
(data.usageMetadata.candidatesTokenCount || 0) +
(data.usageMetadata.thoughtsTokenCount || 0),
...(data.usageMetadata.thoughtsTokenCount && {
Expand Down Expand Up @@ -982,11 +1002,16 @@ function transformStreamingChunkToOpenAIFormat(
],
usage: data.usageMetadata
? {
prompt_tokens: data.usageMetadata.promptTokenCount || 0,
prompt_tokens: data.usageMetadata.promptTokenCount > 0
? data.usageMetadata.promptTokenCount
: calculatePromptTokensFromMessages(messages),
completion_tokens: data.usageMetadata.candidatesTokenCount || 0,
// Calculate total including reasoning tokens for Google models
total_tokens:
(data.usageMetadata.promptTokenCount || 0) +
(data.usageMetadata.promptTokenCount > 0
? data.usageMetadata.promptTokenCount
: calculatePromptTokensFromMessages(messages)
) +
(data.usageMetadata.candidatesTokenCount || 0) +
(data.usageMetadata.thoughtsTokenCount || 0),
...(data.usageMetadata.thoughtsTokenCount && {
Expand Down Expand Up @@ -2590,7 +2615,7 @@ chat.openapi(completions, async (c) => {
let finalTotalTokens = totalTokens;

// Estimate missing tokens if needed using helper function
if (finalPromptTokens === null) {
if (finalPromptTokens === null || finalPromptTokens === 0) {
const estimation = estimateTokens(
usedProvider,
messages,
Expand Down Expand Up @@ -2629,9 +2654,9 @@ chat.openapi(completions, async (c) => {
},
],
usage: {
prompt_tokens: finalPromptTokens || 0,
prompt_tokens: Math.max(1, finalPromptTokens || 1),
completion_tokens: finalCompletionTokens || 0,
total_tokens: finalTotalTokens || 0,
total_tokens: Math.max(1, finalTotalTokens || Math.max(1, finalPromptTokens || 1)),
},
};

Expand Down Expand Up @@ -2679,6 +2704,7 @@ chat.openapi(completions, async (c) => {
usedProvider,
usedModel,
data,
messages,
);

// For Anthropic, if we have partial usage data, complete it
Expand Down Expand Up @@ -2973,14 +2999,14 @@ chat.openapi(completions, async (c) => {
},
],
usage: {
prompt_tokens: Math.round(
promptTokens || calculatedPromptTokens || 0,
),
prompt_tokens: Math.max(1, Math.round(
(promptTokens && promptTokens > 0) ? promptTokens : calculatedPromptTokens || 1,
)),
completion_tokens: Math.round(
completionTokens || calculatedCompletionTokens || 0,
),
total_tokens: Math.round(
totalTokens || calculatedTotalTokens || 0,
totalTokens || calculatedTotalTokens || Math.max(1, (promptTokens && promptTokens > 0) ? promptTokens : calculatedPromptTokens || 1),
),
...(cachedTokens !== null && {
prompt_tokens_details: {
Expand Down
79 changes: 79 additions & 0 deletions apps/gateway/src/lib/prompt-tokens.spec.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,79 @@
import { describe, it, expect } from "vitest";

describe("Prompt token calculation", () => {
it("should calculate prompt tokens when provider returns 0", () => {
// Mock message
const messages = [
{
role: "user",
content: "This is a test message to calculate tokens",
name: undefined,
},
];

// Simulate calculation logic (similar to what we implemented)
const calculatePromptTokens = (
promptTokenCount: number,
messages: any[],
): number => {
if (promptTokenCount > 0) {
return promptTokenCount;
}

// Calculate prompt tokens if missing or 0
try {
// Simple estimation fallback (as in our implementation)
const totalChars = messages.reduce(
(acc, m) => acc + (m.content?.length || 0),
0,
);
return Math.max(1, Math.round(totalChars / 4));
} catch (error) {
Comment thread
coderabbitai[bot] marked this conversation as resolved.
Resolved
return 1; // Minimum fallback
}
};

// Test that 0 prompt tokens are calculated
const result = calculatePromptTokens(0, messages);
expect(result).toBeGreaterThan(0);
expect(typeof result).toBe("number");

// Test that existing prompt tokens are preserved
const result2 = calculatePromptTokens(50, messages);
expect(result2).toBe(50);
});

it("should always return at least 1 token", () => {
const calculateMinTokens = (promptTokens: number | null): number => {
return Math.max(1, promptTokens || 1);
};

expect(calculateMinTokens(0)).toBe(1);
expect(calculateMinTokens(null)).toBe(1);
expect(calculateMinTokens(undefined as any)).toBe(1);
expect(calculateMinTokens(10)).toBe(10);
});

it("should handle empty messages gracefully", () => {
const calculatePromptTokens = (
promptTokenCount: number,
messages: any[],
): number => {
if (promptTokenCount > 0) {
return promptTokenCount;
}

const totalChars = messages.reduce(
(acc, m) => acc + (m.content?.length || 0),
0,
);
return Math.max(1, Math.round(totalChars / 4));
};

const result = calculatePromptTokens(0, []);
expect(result).toBe(1); // Should return minimum of 1

const result2 = calculatePromptTokens(0, [{ role: "user", content: "" }]);
expect(result2).toBe(1); // Should return minimum of 1 for empty content
});
});
10 changes: 10 additions & 0 deletions apps/gateway/src/test-utils/mock-openai-server.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -53,6 +53,11 @@ mockOpenAIServer.post("/v1/chat/completions", async (c) => {
return c.json(sampleErrorResponse);
}

// Check if this request should trigger zero tokens response
const shouldReturnZeroTokens = body.messages.some(
(msg: any) => msg.role === "user" && msg.content.includes("ZERO_TOKENS"),
);

// Get the user's message to include in the response
const userMessage =
body.messages.find((msg: any) => msg.role === "user")?.content || "";
Expand All @@ -69,6 +74,11 @@ mockOpenAIServer.post("/v1/chat/completions", async (c) => {
},
},
],
usage: shouldReturnZeroTokens ? {
prompt_tokens: 0,
completion_tokens: 20,
total_tokens: 20,
} : sampleChatCompletionResponse.usage,
};

return c.json(response);
Expand Down