Skip to content

feat(token-usage): enhance token estimation for images #1202

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 1 commit into from
Nov 20, 2025
Merged

feat(token-usage): enhance token estimation for images #1202

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
83 changes: 65 additions & 18 deletions apps/gateway/src/chat/chat.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2286,11 +2286,12 @@ chat.openapi(completions, async (c) => {
let reasoningTokens = null;
let cachedTokens = null;
let streamingToolCalls = null;
let imageByteSize = 0; // Track total image data size for token estimation
let doneSent = false; // Track if [DONE] has been sent
let buffer = ""; // Buffer for accumulating partial data across chunks (string for SSE)
let binaryBuffer = new Uint8Array(0); // Buffer for binary event streams (AWS Bedrock)
let rawUpstreamData = ""; // Raw data received from upstream provider
const MAX_BUFFER_SIZE = 10 * 1024 * 1024; // 10MB limit
// const MAX_BUFFER_SIZE = 10 * 1024 * 1024; // 10MB limit
Comment on lines +2289 to +2294

@coderabbitai coderabbitai Bot Nov 20, 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 | 🟠 Major

Commented‑out MAX_BUFFER_SIZE removes protection against unbounded streaming buffer growth

You’ve left buffer unbounded by commenting out the 10MB guard:

let buffer = "";
// const MAX_BUFFER_SIZE = 10 * 1024 * 1024; // 10MB limit
...
// if (buffer.length > MAX_BUFFER_SIZE) { ... }

This means a misbehaving or malicious upstream stream that never yields a complete SSE event can grow buffer without limit, risking high memory usage or OOM under load.

Consider re‑enabling the guard (even at a higher limit or gated behind debugMode) so production still has a hard cap, for example:

-			let buffer = ""; // Buffer for accumulating partial data across chunks (string for SSE)
+			let buffer = ""; // Buffer for accumulating partial data across chunks (string for SSE)
+			const MAX_BUFFER_SIZE = 10 * 1024 * 1024; // 10MB limit
...
-					// // Check buffer size to prevent memory exhaustion
-					// if (buffer.length > MAX_BUFFER_SIZE) {
-					// 	logger.warn(
-					// 		"Buffer size exceeded 10MB, clearing buffer to prevent memory exhaustion",
-					// 	);
-					// 	buffer = "";
-					// 	continue;
-					// }
+					// Check buffer size to prevent memory exhaustion
+					if (buffer.length > MAX_BUFFER_SIZE) {
+						logger.warn(
+							"Buffer size exceeded 10MB, clearing buffer to prevent memory exhaustion",
+						);
+						buffer = "";
+						continue;
+					}

Also applies to: 2335-2342

🤖 Prompt for AI Agents 
In apps/gateway/src/chat/chat.ts around lines 2289-2294, the MAX_BUFFER_SIZE
guard was commented out leaving `buffer` (and similarly `binaryBuffer`)
unbounded; restore a sensible hard cap (e.g. const MAX_BUFFER_SIZE = 10 * 1024 *
1024) and enforce it whenever you append to the buffers: if the new size would
exceed the cap, stop/close the upstream stream or discard the event and emit an
error/cleanup to prevent unbounded growth. Optionally gate the cap via a runtime
flag (debugMode) or make it configurable, but ensure a production default limit
is always enforced and apply the same check to binaryBuffer/rawUpstreamData
regions noted at lines ~2335-2342.

const isAwsBedrock = usedProvider === "aws-bedrock";

try {
Expand Down Expand Up @@ -2331,14 +2332,14 @@ chat.openapi(completions, async (c) => {
rawUpstreamData += chunk;
}

// Check buffer size to prevent memory exhaustion
if (buffer.length > MAX_BUFFER_SIZE) {
logger.warn(
"Buffer size exceeded 10MB, clearing buffer to prevent memory exhaustion",
);
buffer = "";
continue;
}
// // Check buffer size to prevent memory exhaustion
// if (buffer.length > MAX_BUFFER_SIZE) {
// logger.warn(
// "Buffer size exceeded 10MB, clearing buffer to prevent memory exhaustion",
// );
// buffer = "";
// continue;
// }

// Process SSE events from buffer
let processedLength = 0;
Expand Down Expand Up @@ -2504,7 +2505,15 @@ chat.openapi(completions, async (c) => {
}

if (finalCompletionTokens === null) {
finalCompletionTokens = estimateTokensFromContent(fullContent);
let textTokens = estimateTokensFromContent(fullContent);
// For images, estimate ~258 tokens per image + 1 token per 750 bytes
// This is based on Google's image token calculation
let imageTokens = 0;
if (imageByteSize > 0) {
// Base tokens per image (258) + additional tokens based on size
imageTokens = 258 + Math.ceil(imageByteSize / 750);
}
finalCompletionTokens = textTokens + imageTokens;
}
Comment on lines 2507 to 2517

@coderabbitai coderabbitai Bot Nov 20, 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 | 🟡 Minor

Image token estimation applies 258‑token base only once, regardless of number of images

Across the streaming paths you now compute completion tokens as:

// [DONE] path
let textTokens = estimateTokensFromContent(fullContent);
// ...
if (imageByteSize > 0) {
	imageTokens = 258 + Math.ceil(imageByteSize / 750);
}
finalCompletionTokens = textTokens + imageTokens;

and later:

// finishReason-based estimation
let textTokens = estimateTokensFromContent(fullContent);
if (imageByteSize > 0) {
	imageTokens = 258 + Math.ceil(imageByteSize / 750);
}
completionTokens = textTokens + imageTokens;

and in finally:

if (!completionTokens && (fullContent || imageByteSize > 0)) {
	// encode(JSON.stringify(fullContent)) or estimateTokensFromContent
	// ...
	if (imageByteSize > 0) {
		imageTokens = 258 + Math.ceil(imageByteSize / 750);
	}
	calculatedCompletionTokens = textTokens + imageTokens;
}

As in extract-token-usage.ts, the comments say “~258 tokens per image + 1 token per 750 bytes”, but this code applies the 258‑token base only once, using imageByteSize as an aggregate over all images.

If a response can contain multiple images, this will under‑estimate completion tokens (and thus costs/limits). Consider:

  • tracking imageCount alongside imageByteSize, and using:
const imageTokens =
	imageCount > 0
		? 258 * imageCount + Math.ceil(imageByteSize / 750)
		: 0;
  • or, if the formula is intentionally “once per response”, update comments to remove “per image” to avoid confusion.

Also, for consistency, you may want to factor out a small helper like estimateImageTokens(imageByteSize, imageCount?) and reuse it in both streaming and extract-token-usage.ts.

Also applies to: 2899-2905, 3015-3038

🤖 Prompt for AI Agents 
In apps/gateway/src/chat/chat.ts around lines 2507 to 2517 (and also at
2899-2905, 3015-3038), the image token estimation applies the 258-token base
only once using aggregate imageByteSize, which undercounts when multiple images
are present; update the logic to track imageCount alongside imageByteSize and
compute imageTokens as (imageCount > 0 ? 258 * imageCount +
Math.ceil(imageByteSize / 750) : 0), replace the repeated inline calculations
with a single helper function (e.g., estimateImageTokens(imageByteSize,
imageCount)) and use it in all three places (and in extract-token-usage.ts) to
ensure correct per-image base tokens and avoid duplication.


if (finalTotalTokens === null) {
Expand Down Expand Up @@ -2648,6 +2657,7 @@ chat.openapi(completions, async (c) => {
data,
usedProvider,
fullContent,
imageByteSize,
);

// If we have usage data from Google, add it to the streaming chunk
Expand Down Expand Up @@ -2742,6 +2752,22 @@ chat.openapi(completions, async (c) => {
}
}

// Track image data size for Google providers (for token estimation)
if (
usedProvider === "google-ai-studio" ||
usedProvider === "google-vertex"
) {
const parts = data.candidates?.[0]?.content?.parts || [];
for (const part of parts) {
if (part.inlineData?.data) {
// Base64 string length * 0.75 ≈ actual byte size
imageByteSize += Math.ceil(
part.inlineData.data.length * 0.75,
);
}
}
}
Comment on lines +2755 to +2769

@coderabbitai coderabbitai Bot Nov 20, 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 | 🟡 Minor

🧩 Analysis chain

Potential double‑counting of image bytes across Google streaming chunks

imageByteSize accumulation for Google streaming:

if (
	usedProvider === "google-ai-studio" ||
	usedProvider === "google-vertex"
) {
	const parts = data.candidates?.[0]?.content?.parts || [];
	for (const part of parts) {
		if (part.inlineData?.data) {
			// Base64 string length * 0.75 ≈ actual byte size
			imageByteSize += Math.ceil(part.inlineData.data.length * 0.75);
		}
	}
}

If Google’s streaming responses repeat the same inlineData in multiple SSE chunks (e.g., resending the full content state each delta), this logic will add the same image bytes multiple times, inflating imageByteSize and thus image token estimates.

You might want to:

  • either track which inlineData blobs you’ve already seen (e.g., by a hash or reference) and only count new ones, or
  • constrain accumulation to the first chunk that contains inline image data (e.g., only when imageByteSize === 0).

Example minimal change:

-	for (const part of parts) {
-		if (part.inlineData?.data) {
+	for (const part of parts) {
+		if (part.inlineData?.data && imageByteSize === 0) {

I need to verify the context of this code to properly assess the double-counting concern. Let me explore the codebase structure.


---

🏁 Script executed:

```shell
#!/bin/bash
# Locate the function containing lines 2755-2769 in chat.ts
# and understand the streaming loop structure
cat apps/gateway/src/chat/chat.ts | sed -n '2700,2800p'

Length of output: 3553

🏁 Script executed:

#!/bin/bash
# Search for where imageByteSize is declared and how it's used
rg -n "imageByteSize" apps/gateway/src/chat/chat.ts -B5 -A5 | head -100

Length of output: 3954

🏁 Script executed:

#!/bin/bash
# Look for Google streaming chunk processing to understand the loop
rg -n "google-ai-studio|google-vertex" apps/gateway/src/chat/chat.ts -B10 -A10 | head -150

Length of output: 5604

🏁 Script executed:

#!/bin/bash
# Check for SSE or streaming event handling
rg -n "for.*chunk|while.*chunk|\.on\(" apps/gateway/src/chat/chat.ts | head -30

Length of output: 327

🏁 Script executed:

#!/bin/bash
# Find function signature or context around line 2755
sed -n '2600,2770p' apps/gateway/src/chat/chat.ts | tail -200

Length of output: 5452

I need to verify Google's streaming behavior and check for any existing deduplication logic.

🏁 Script executed:

#!/bin/bash
# Check the extractTokenUsage function to see how imageByteSize is used
rg -n "function extractTokenUsage|const extractTokenUsage" apps/gateway/src/chat/chat.ts -A50 | head -80

Length of output: 0

🏁 Script executed:

#!/bin/bash
# Look for any image deduplication or tracking logic
rg -n "inlineData|image.*hash|image.*track" apps/gateway/src/chat/chat.ts -B3 -A3

Length of output: 379

🏁 Script executed:

#!/bin/bash
# Check how streaming chunks are processed - look for the loop structure
sed -n '2350,2420p' apps/gateway/src/chat/chat.ts

Length of output: 2512

🏁 Script executed:

#!/bin/bash
# Find where the streaming loop starts to understand chunk processing
rg -n "while.*buffer|for.*chunk" apps/gateway/src/chat/chat.ts -B5 -A15 | head -80

Length of output: 3106

🌐 Web query:

Google Generative AI API streaming response inlineData repeated chunks

💡 Result:

Short answer: what you're seeing is expected — inlineData (binary/base64) can be split across multiple streaming chunks and may appear duplicated by transport buffering; you must treat the stream as a sequence of partial messages, reassemble/concatenate the inlineData pieces, and dedupe/ignore repeats before decoding. See below.

Why it happens (brief)

  • The Generative API returns streaming responses (SSE / gRPC stream) that may split large fields (including inlineData/base64 blobs) across multiple partial chunks. The client/transport (proxies, SSE buffering, SDK wrapper) can also produce repeated or overlapping fragments. Handle the stream as a sequence of partial JSON/SSE frames, not as final objects. [1][2]

How to handle it (recommended)

  • Parse SSE / stream frames robustly: buffer incomplete frame text, skip empty lines and the [DONE] marker, and only JSON-parse complete frames. [3][1]
  • Reassemble inlineData: when a part contains inlineData.data that is a fragment of a base64 blob, append fragments in order to a buffer (do not decode until fully reassembled). Use any provided identifiers/offsets if present; otherwise append in arrival order. [2]
  • Dedupe/ignore repeats: if you receive identical inlineData fragments more than once (same base64 substring or same part id/offset), ignore duplicates before final decode. If you get overlapping fragments, assemble using offsets or a sliding-window compare to avoid double-counting.
  • Protect memory: stream to disk or stream-decode incrementally for very large media rather than buffering entire file in memory.
  • Test against network/proxy buffering: cross-region/edge proxies (or Cloud Run / load balancers) can change chunking/latency — run tests close to service region and avoid extra buffering layers. (Common streaming advice). [3][4]

Sources

  • Google Generative AI streaming examples and SDK generate_content_stream docs. [1]
  • Protobuf/Blob (inlineData) definitions and guidance for inline binary parts. [2]
  • SSE / streaming parsing best-practices (buffer incomplete frames, skip invalid JSON / [DONE]). [3]

If you want, tell me which client/SDK (fetch, @google/genai, Python SDK, gRPC) you're using and I’ll give a short example showing robust parsing + inlineData reassembly for that environment.

The web search confirms the double-counting concern is real. Google's streaming API can repeat inlineData across chunks. Let me verify the token estimation logic and the adequacy of the proposed fix.


---

🏁 Script executed:

```shell
#!/bin/bash
# Search for extractTokenUsage implementation across all files
fd -e ts -e js | xargs grep -l "extractTokenUsage" | head -5

Length of output: 80

🏁 Script executed:

#!/bin/bash
# Look for extractTokenUsage definition more broadly
grep -rn "extractTokenUsage.*=" apps/gateway/src/ --include="*.ts" --include="*.js" | head -20

Length of output: 0

🏁 Script executed:

#!/bin/bash
# Check how imageByteSize is used in token calculations
sed -n '2500,2530p' apps/gateway/src/chat/chat.ts

Length of output: 1037

🏁 Script executed:

#!/bin/bash
# Look for any part.id or image identification in Google response handling
rg -n "part\." apps/gateway/src/chat/chat.ts | grep -v "part of" | head -30

Length of output: 94

🏁 Script executed:

#!/bin/bash
# Check if there are multiple inlineData entries possible in a single chunk
rg -n "candidates\[0\]\.content\.parts" apps/gateway/src/chat/chat.ts -B3 -A5

Length of output: 0

Let me examine the extractTokenUsage function and confirm whether the proposed fix is adequate.


---

🏁 Script executed:

```shell
#!/bin/bash
# Read the extractTokenUsage implementation
cat apps/gateway/src/chat/tools/extract-token-usage.ts

Length of output: 3055

🏁 Script executed:

#!/bin/bash
# Look at surrounding context of the image byte size code to understand chunk processing better
sed -n '2740,2800p' apps/gateway/src/chat/chat.ts

Length of output: 2130

🏁 Script executed:

#!/bin/bash
# Check if data.candidates is queried for multiple indices or always [0]
rg -n "candidates\?" apps/gateway/src/chat/chat.ts | head -20

Length of output: 133

Concern validated: double-counting of image bytes is possible, but proposed fix is overly simplistic

Google's streaming API can repeat or split inlineData across multiple SSE chunks. Your code at lines 2755-2769 accumulates imageByteSize without deduplication across all chunks, inflating token estimates via extractTokenUsage.

The suggested guard && imageByteSize === 0 prevents counting after the first chunk—functionally adequate for single-image responses but architecturally fragile. A more robust approach would track seen inlineData by hash or content reference to handle:

  • Multiple distinct images across chunks
  • Inlinedata splits or repeats (as Google documents)

Improved fix:

-	for (const part of parts) {
-		if (part.inlineData?.data) {
+	for (const part of parts) {
+		if (part.inlineData?.data) {
+			// Deduplicate by checking if this base64 blob was already counted
+			const dataHash = crypto.createHash('sha256').update(part.inlineData.data).digest('hex');
+			if (!seenImageHashes.has(dataHash)) {
+				seenImageHashes.add(dataHash);

Alternatively, if images only appear in the first chunk per Google's typical behavior, the simpler imageByteSize === 0 guard is acceptable but should include a comment explaining the assumption.

Committable suggestion skipped: line range outside the PR's diff.


// Extract reasoning content for logging using helper function
// For providers with custom extraction logic (google-ai-studio, anthropic),
// use raw data. For others, use transformed OpenAI format.
Expand Down Expand Up @@ -2834,7 +2860,12 @@ chat.openapi(completions, async (c) => {
}

// Extract token usage using helper function
const usage = extractTokenUsage(data, usedProvider, fullContent);
const usage = extractTokenUsage(
data,
usedProvider,
fullContent,
imageByteSize,
);
if (usage.promptTokens !== null) {
promptTokens = usage.promptTokens;
}
Expand Down Expand Up @@ -2865,7 +2896,13 @@ chat.openapi(completions, async (c) => {
}

if (!completionTokens) {
completionTokens = estimateTokensFromContent(fullContent);
let textTokens = estimateTokensFromContent(fullContent);
// For images, estimate ~258 tokens per image + 1 token per 750 bytes
let imageTokens = 0;
if (imageByteSize > 0) {
imageTokens = 258 + Math.ceil(imageByteSize / 750);
}
completionTokens = textTokens + imageTokens;
}

totalTokens = (promptTokens || 0) + (completionTokens || 0);
Expand Down Expand Up @@ -2975,19 +3012,29 @@ chat.openapi(completions, async (c) => {
}
}

if (!completionTokens && fullContent) {
if (!completionTokens && (fullContent || imageByteSize > 0)) {
try {
calculatedCompletionTokens = encode(
JSON.stringify(fullContent),
).length;
let textTokens = fullContent
? encode(JSON.stringify(fullContent)).length
: 0;
// For images, estimate ~258 tokens per image + 1 token per 750 bytes
let imageTokens = 0;
if (imageByteSize > 0) {
imageTokens = 258 + Math.ceil(imageByteSize / 750);
}
calculatedCompletionTokens = textTokens + imageTokens;
} catch (error) {
// Fallback to simple estimation if encoding fails
logger.error(
"Failed to encode completion text in streaming",
error instanceof Error ? error : new Error(String(error)),
);
calculatedCompletionTokens =
estimateTokensFromContent(fullContent);
let textTokens = estimateTokensFromContent(fullContent);
let imageTokens = 0;
if (imageByteSize > 0) {
imageTokens = 258 + Math.ceil(imageByteSize / 750);
}
calculatedCompletionTokens = textTokens + imageTokens;
}
}

Expand Down
18 changes: 14 additions & 4 deletions apps/gateway/src/chat/tools/extract-token-usage.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,7 @@ export function extractTokenUsage(
data: any,
provider: Provider,
fullContent?: string,
imageByteSize?: number,

@coderabbitai coderabbitai Bot Nov 20, 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 | 🟡 Minor

Image token estimation formula ignores image count vs comment (“per image”)

The new imageByteSize parameter and fallback estimation path look good and are safely gated on completionTokens === null. One thing to watch:

  • The comment says ~258 tokens per image + 1 token per 750 bytes, but the implementation uses imageTokens = 258 + Math.ceil(imageByteSize / 750), which applies the 258‑token base only once, even if imageByteSize represents multiple images.

If Google responses can contain multiple images, this will under‑estimate completion tokens. Either:

  • also pass an imageCount and do 258 * imageCount + ceil(totalBytes / 750), or
  • adjust the comment to clarify that 258 is applied once per response, not per image.

Also applies to: 29-48

🤖 Prompt for AI Agents 
In apps/gateway/src/chat/tools/extract-token-usage.ts around lines 12 and 29-48,
the image token estimation applies the ~258 token base only once even when
multiple images may be present, but the comment states it’s "per image"; update
the implementation to accept an imageCount (or derive it if available) and
compute imageTokens as (258 * imageCount) + Math.ceil(imageByteSize / 750), or
if you intend 258 to be per response, change the comment to state that 258 is
applied once per response; implement the chosen change consistently in the
fallback path and update the comment accordingly.

) {
let promptTokens = null;
let completionTokens = null;
Expand All @@ -25,16 +26,25 @@ export function extractTokenUsage(
// Don't use Google's totalTokenCount as it doesn't include reasoning tokens
reasoningTokens = data.usageMetadata.thoughtsTokenCount ?? null;

// If candidatesTokenCount is missing and we have content, estimate it
if (completionTokens === null && fullContent) {
// If candidatesTokenCount is missing and we have content or images, estimate it
if (
completionTokens === null &&
(fullContent || (imageByteSize && imageByteSize > 0))
) {
const estimation = estimateTokens(
provider,
[],
fullContent,
fullContent || "",
null,
null,
);
completionTokens = estimation.calculatedCompletionTokens;
let textTokens = estimation.calculatedCompletionTokens || 0;
// For images, estimate ~258 tokens per image + 1 token per 750 bytes
let imageTokens = 0;
if (imageByteSize && imageByteSize > 0) {
imageTokens = 258 + Math.ceil(imageByteSize / 750);
}
completionTokens = textTokens + imageTokens;
}
// Calculate total including reasoning tokens (after potential estimation)
totalTokens =
Expand Down
Loading