mirror of
https://github.com/RooVetGit/Roo-Code.git
synced 2026-08-28 05:27:24 +00:00
fix: improve Vertex AI context tracking and token counting
- Enhanced GeminiHandler.countTokens() to handle zero token counts and provide better fallback - Added countTokens override in AnthropicVertexHandler to use reliable tiktoken for Claude models - Added robust token counting with validation in VertexHandler for Gemini models - Improved error handling and logging for token counting failures - Ensures consistent context tracking for both Vertex AI Claude and Gemini models Fixes #5830
This commit is contained in:
parent
4b1c9021d6
commit
56224459f5
3 changed files with 44 additions and 3 deletions
|
|
@ -217,4 +217,13 @@ export class AnthropicVertexHandler extends BaseProvider implements SingleComple
|
|||
throw error
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Override token counting to use tiktoken fallback for Vertex AI Claude models.
|
||||
* The Vertex AI token counting API can be unreliable, so we use the base provider's
|
||||
* tiktoken implementation which is more consistent for context tracking.
|
||||
*/
|
||||
override async countTokens(content: Anthropic.Messages.ContentBlockParam[]): Promise<number> {
|
||||
return super.countTokens(content)
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -176,14 +176,14 @@ export class GeminiHandler extends BaseProvider implements SingleCompletionHandl
|
|||
contents: convertAnthropicContentToGemini(content),
|
||||
})
|
||||
|
||||
if (response.totalTokens === undefined) {
|
||||
console.warn("Gemini token counting returned undefined, using fallback")
|
||||
if (response.totalTokens === undefined || response.totalTokens === 0) {
|
||||
console.warn("Gemini token counting returned undefined or zero, using fallback")
|
||||
return super.countTokens(content)
|
||||
}
|
||||
|
||||
return response.totalTokens
|
||||
} catch (error) {
|
||||
console.warn("Gemini token counting failed, using fallback", error)
|
||||
console.warn("Gemini token counting failed, using fallback:", error)
|
||||
return super.countTokens(content)
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -1,4 +1,5 @@
|
|||
import { type ModelInfo, type VertexModelId, vertexDefaultModelId, vertexModels } from "@roo-code/types"
|
||||
import type { Anthropic } from "@anthropic-ai/sdk"
|
||||
|
||||
import type { ApiHandlerOptions } from "../../shared/api"
|
||||
|
||||
|
|
@ -24,4 +25,35 @@ export class VertexHandler extends GeminiHandler implements SingleCompletionHand
|
|||
// suffix.
|
||||
return { id: id.endsWith(":thinking") ? id.replace(":thinking", "") : id, info, ...params }
|
||||
}
|
||||
|
||||
/**
|
||||
* Override token counting to add additional error handling for Vertex AI.
|
||||
* Falls back to tiktoken if the Gemini API token counting fails or returns unreliable results.
|
||||
*/
|
||||
override async countTokens(content: Array<Anthropic.Messages.ContentBlockParam>): Promise<number> {
|
||||
try {
|
||||
// Try the parent GeminiHandler's token counting first
|
||||
const tokenCount = await super.countTokens(content)
|
||||
|
||||
// Additional validation: if token count seems unreasonably low for non-empty content,
|
||||
// fall back to tiktoken
|
||||
if (content.length > 0 && tokenCount === 0) {
|
||||
console.warn("Vertex AI token counting returned 0 for non-empty content, using tiktoken fallback")
|
||||
return this.fallbackTokenCount(content)
|
||||
}
|
||||
|
||||
return tokenCount
|
||||
} catch (error) {
|
||||
console.warn("Vertex AI token counting failed completely, using tiktoken fallback:", error)
|
||||
return this.fallbackTokenCount(content)
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Fallback token counting using the base provider's tiktoken implementation
|
||||
*/
|
||||
private async fallbackTokenCount(content: Array<Anthropic.Messages.ContentBlockParam>): Promise<number> {
|
||||
// Call the base provider's countTokens method (which uses tiktoken)
|
||||
return super.countTokens(content)
|
||||
}
|
||||
}
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue