Skip to content

fix(token-estimation): refine logic #921

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

Open
wants to merge 3 commits into
base: main
Choose a base branch
from
Open

fix(token-estimation): refine logic #921

Changes from 1 commit
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
c2f7be5
refactor(chat): refine token estimation logic in estimateTokens function
steebchen Sep 22, 2025
3e265ff
fix(chat): round token estimates to nearest integer
steebchen Sep 23, 2025
2008509
refactor(chat): streamline token and cost logic
steebchen Sep 23, 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
73 changes: 34 additions & 39 deletions apps/gateway/src/chat/tools/estimate-tokens.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -19,47 +19,42 @@ export function estimateTokens(
let calculatedPromptTokens = promptTokens;
let calculatedCompletionTokens = completionTokens;

// Always estimate missing tokens for any provider
if (!promptTokens || !completionTokens) {
// Estimate prompt tokens using encodeChat for better accuracy
if (!promptTokens && messages && messages.length > 0) {
try {
// Convert messages to the format expected by gpt-tokenizer
const chatMessages: ChatMessage[] = messages.map((m) => ({
role: m.role,
content:
typeof m.content === "string"
? m.content
: JSON.stringify(m.content),
name: m.name,
}));
calculatedPromptTokens = encodeChat(
chatMessages,
DEFAULT_TOKENIZER_MODEL,
).length;
} catch (error) {
// Fallback to simple estimation if encoding fails
logger.error(
"Failed to encode chat messages in estimate tokens",
error instanceof Error ? error : new Error(String(error)),
);
calculatedPromptTokens =
messages.reduce((acc, m) => acc + (m.content?.length || 0), 0) / 4;
}
// Estimate prompt tokens only if not provided by the API
if (!promptTokens && messages && messages.length > 0) {
try {
// Convert messages to the format expected by gpt-tokenizer
const chatMessages: ChatMessage[] = messages.map((m) => ({
role: m.role,
content:
typeof m.content === "string" ? m.content : JSON.stringify(m.content),
name: m.name,
}));
calculatedPromptTokens = encodeChat(
chatMessages,
DEFAULT_TOKENIZER_MODEL,
).length;
} catch (error) {
// Fallback to simple estimation if encoding fails
logger.error(
"Failed to encode chat messages in estimate tokens",
error instanceof Error ? error : new Error(String(error)),
);
calculatedPromptTokens =
messages.reduce((acc, m) => acc + (m.content?.length || 0), 0) / 4;
}
}
Comment on lines 22 to 46
Copy link
Contributor

Choose a reason for hiding this comment

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

⚠️ Potential issue

Guard for “missing” must not treat 0 as missing; fix fallback undercount and return integers

  • Line 23: Using !promptTokens re-estimates when the API legitimately returns 0. Use a null/undefined check.
  • Lines 42-44: Fallback undercounts when content is non-string (objects/arrays). Mirror the encoding path and stringify per‑message before length. Also round to an integer token count.

Apply this diff:

- // Estimate prompt tokens only if not provided by the API
- if (!promptTokens && messages && messages.length > 0) {
+ // Estimate prompt tokens only if not provided by the API
+ if (promptTokens == null && messages?.length > 0) {
   try {
@@
-      calculatedPromptTokens =
-        messages.reduce((acc, m) => acc + (m.content?.length || 0), 0) / 4;
+      calculatedPromptTokens = Math.round(
+        messages.reduce((acc, m) => {
+          const text =
+            typeof m.content === "string" ? m.content : JSON.stringify(m.content);
+          return acc + text.length;
+        }, 0) / 4
+      );
   }
 }

Additionally: the parameter messages: any[] violates the project guideline to avoid any. Consider typing it (e.g., a ChatMessageLike with content: unknown) and narrowing in place.

📝 Committable suggestion

‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.

Suggested change
// Estimate prompt tokens only if not provided by the API
if (!promptTokens && messages && messages.length > 0) {
try {
// Convert messages to the format expected by gpt-tokenizer
const chatMessages: ChatMessage[] = messages.map((m) => ({
role: m.role,
content:
typeof m.content === "string" ? m.content : JSON.stringify(m.content),
name: m.name,
}));
calculatedPromptTokens = encodeChat(
chatMessages,
DEFAULT_TOKENIZER_MODEL,
).length;
} catch (error) {
// Fallback to simple estimation if encoding fails
logger.error(
"Failed to encode chat messages in estimate tokens",
error instanceof Error ? error : new Error(String(error)),
);
calculatedPromptTokens =
messages.reduce((acc, m) => acc + (m.content?.length || 0), 0) / 4;
}
}
// Estimate prompt tokens only if not provided by the API
if (promptTokens == null && messages?.length > 0) {
try {
// Convert messages to the format expected by gpt-tokenizer
const chatMessages: ChatMessage[] = messages.map((m) => ({
role: m.role,
content:
typeof m.content === "string" ? m.content : JSON.stringify(m.content),
name: m.name,
}));
calculatedPromptTokens = encodeChat(
chatMessages,
DEFAULT_TOKENIZER_MODEL,
).length;
} catch (error) {
// Fallback to simple estimation if encoding fails
logger.error(
"Failed to encode chat messages in estimate tokens",
error instanceof Error ? error : new Error(String(error)),
);
calculatedPromptTokens = Math.round(
messages.reduce((acc, m) => {
const text =
typeof m.content === "string" ? m.content : JSON.stringify(m.content);
return acc + text.length;
}, 0) / 4
);
}
}


// Estimate completion tokens using encode for better accuracy
if (!completionTokens && content) {
try {
calculatedCompletionTokens = encode(JSON.stringify(content)).length;
} catch (error) {
// Fallback to simple estimation if encoding fails
logger.error(
"Failed to encode completion text",
error instanceof Error ? error : new Error(String(error)),
);
calculatedCompletionTokens = content.length / 4;
}
// Estimate completion tokens only if not provided by the API
if (!completionTokens && content) {
try {
calculatedCompletionTokens = encode(JSON.stringify(content)).length;
} catch (error) {
// Fallback to simple estimation if encoding fails
logger.error(
"Failed to encode completion text",
error instanceof Error ? error : new Error(String(error)),
);
calculatedCompletionTokens = content.length / 4;
}
}
Comment on lines 48 to 60
Copy link
Contributor

Choose a reason for hiding this comment

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

⚠️ Potential issue

Fix completion-path guard and encoding; avoid JSON.stringify and round fallback

  • Line 48: !completionTokens incorrectly treats 0 as missing and skips empty-string content. Use a null/undefined check and allow empty strings.
  • Line 50: encode(JSON.stringify(content)) overcounts by adding quotes/escapes; pass the string directly.
  • Line 57: Round the heuristic to an integer.

Apply this diff:

- // Estimate completion tokens only if not provided by the API
- if (!completionTokens && content) {
+ // Estimate completion tokens only if not provided by the API
+ if (completionTokens == null && content != null) {
   try {
-    calculatedCompletionTokens = encode(JSON.stringify(content)).length;
+    calculatedCompletionTokens = encode(content).length;
   } catch (error) {
@@
-    calculatedCompletionTokens = content.length / 4;
+    calculatedCompletionTokens = Math.round(content.length / 4);
   }
 }
📝 Committable suggestion

‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.

Suggested change
// Estimate completion tokens only if not provided by the API
if (!completionTokens && content) {
try {
calculatedCompletionTokens = encode(JSON.stringify(content)).length;
} catch (error) {
// Fallback to simple estimation if encoding fails
logger.error(
"Failed to encode completion text",
error instanceof Error ? error : new Error(String(error)),
);
calculatedCompletionTokens = content.length / 4;
}
}
// Estimate completion tokens only if not provided by the API
if (completionTokens == null && content != null) {
try {
calculatedCompletionTokens = encode(content).length;
} catch (error) {
// Fallback to simple estimation if encoding fails
logger.error(
"Failed to encode completion text",
error instanceof Error ? error : new Error(String(error)),
);
calculatedCompletionTokens = Math.round(content.length / 4);
}
}
🤖 Prompt for AI Agents 
In apps/gateway/src/chat/tools/estimate-tokens.ts around lines 47 to 59, the
guard `!completionTokens` treats 0 as missing and the encoder call uses
JSON.stringify which inflates token count and the fallback heuristic isn't
rounded; change the guard to check for null/undefined (e.g., completionTokens ==
null) so zero is allowed, only run the block when content is not null/undefined
(allow empty strings), call encode on the string content directly (e.g.,
encode(String(content))) instead of JSON.stringify(content), and round the
fallback heuristic to an integer (e.g., Math.round(content.length / 4)).


Expand Down
Loading