feat(token-usage): enhance token estimation for images - #1202
Conversation
WalkthroughToken estimation logic is extended to account for image data. An Changes
Sequence DiagramsequenceDiagram
participant Client
participant Chat as chat.ts
participant Extract as extract-token-usage.ts
Client->>Chat: Request with images
Chat->>Chat: Accumulate imageByteSize
Chat->>Extract: extractTokenUsage(..., imageByteSize)
Extract->>Extract: Check if fullContent OR imageByteSize present
alt Image data present
Extract->>Extract: Calculate imageTokens<br/>(258 + ceil(imageByteSize/750))
Extract->>Extract: Combine with textTokens<br/>completionTokens = textTokens + imageTokens
else Text only
Extract->>Extract: Use estimateTokens<br/>for completionTokens
end
Extract-->>Chat: Return usage with<br/>image-aware tokens
Chat-->>Client: Response with final<br/>token estimates
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~20 minutes
Possibly related PRs
Pre-merge checks and finishing touches✅ Passed checks (3 passed)
✨ Finishing touches
🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 4
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (2)
apps/gateway/src/chat/chat.ts(8 hunks)apps/gateway/src/chat/tools/extract-token-usage.ts(2 hunks)
🧰 Additional context used
🧬 Code graph analysis (2)
apps/gateway/src/chat/tools/extract-token-usage.ts (1)
apps/gateway/src/chat/tools/estimate-tokens.ts (1)
estimateTokens(12-70)
apps/gateway/src/chat/chat.ts (2)
apps/gateway/src/chat/tools/extract-token-usage.ts (1)
extractTokenUsage(8-92)packages/logger/src/index.ts (2)
error(153-160)logger(181-181)
⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (10)
- GitHub Check: test / run
- GitHub Check: build / run
- GitHub Check: generate / run
- GitHub Check: lint / run
- GitHub Check: e2e-shards (4)
- GitHub Check: e2e-shards (1)
- GitHub Check: e2e-shards (5)
- GitHub Check: e2e-shards (2)
- GitHub Check: e2e-shards (3)
- GitHub Check: autofix
| 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 |
There was a problem hiding this comment.
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.
| 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; | ||
| } |
There was a problem hiding this comment.
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
imageCountalongsideimageByteSize, 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.
| // 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, | ||
| ); | ||
| } | ||
| } | ||
| } |
There was a problem hiding this comment.
🧩 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
inlineDatablobs 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 -100Length 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 -150Length 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 -30Length 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 -200Length 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 -80Length 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 -A3Length 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.tsLength 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 -80Length 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 -20Length of output: 0
🏁 Script executed:
#!/bin/bash
# Check how imageByteSize is used in token calculations
sed -n '2500,2530p' apps/gateway/src/chat/chat.tsLength 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 -30Length 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 -A5Length 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.tsLength 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 -20Length 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.
| data: any, | ||
| provider: Provider, | ||
| fullContent?: string, | ||
| imageByteSize?: number, |
There was a problem hiding this comment.
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 usesimageTokens = 258 + Math.ceil(imageByteSize / 750), which applies the 258‑token base only once, even ifimageByteSizerepresents multiple images.
If Google responses can contain multiple images, this will under‑estimate completion tokens. Either:
- also pass an
imageCountand do258 * 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.
This reverts commit 93ca2b0.
1 participant
Summary by CodeRabbit
New Features
Improvements
✏️ Tip: You can customize this high-level summary in your review settings.