mirror of
https://github.com/simstudioai/sim.git
synced 2026-01-28 16:27:55 -05:00
fix(gemini): token count (#3039)
* fix(gemini): token count * fix to include tool call tokens
This commit is contained in:
committed by
GitHub
parent
b4a389a71f
commit
c8ffda1616
@@ -936,8 +936,12 @@ export class AgentBlockHandler implements BlockHandler {
|
||||
systemPrompt: validMessages ? undefined : inputs.systemPrompt,
|
||||
context: validMessages ? undefined : stringifyJSON(messages),
|
||||
tools: formattedTools,
|
||||
temperature: inputs.temperature != null ? Number(inputs.temperature) : undefined,
|
||||
maxTokens: inputs.maxTokens != null ? Number(inputs.maxTokens) : undefined,
|
||||
temperature:
|
||||
inputs.temperature != null && inputs.temperature !== ''
|
||||
? Number(inputs.temperature)
|
||||
: undefined,
|
||||
maxTokens:
|
||||
inputs.maxTokens != null && inputs.maxTokens !== '' ? Number(inputs.maxTokens) : undefined,
|
||||
apiKey: inputs.apiKey,
|
||||
azureEndpoint: inputs.azureEndpoint,
|
||||
azureApiVersion: inputs.azureApiVersion,
|
||||
|
||||
@@ -14,8 +14,8 @@ export interface AgentInputs {
|
||||
slidingWindowSize?: string // For message-based sliding window
|
||||
slidingWindowTokens?: string // For token-based sliding window
|
||||
// LLM parameters
|
||||
temperature?: number
|
||||
maxTokens?: number
|
||||
temperature?: string
|
||||
maxTokens?: string
|
||||
apiKey?: string
|
||||
azureEndpoint?: string
|
||||
azureApiVersion?: string
|
||||
|
||||
@@ -123,17 +123,21 @@ export function extractFunctionCallPart(candidate: Candidate | undefined): Part
|
||||
}
|
||||
|
||||
/**
|
||||
* Converts usage metadata from SDK response to our format
|
||||
* Converts usage metadata from SDK response to our format.
|
||||
* Per Gemini docs, total = promptTokenCount + candidatesTokenCount + toolUsePromptTokenCount + thoughtsTokenCount
|
||||
* We include toolUsePromptTokenCount in input and thoughtsTokenCount in output for correct billing.
|
||||
*/
|
||||
export function convertUsageMetadata(
|
||||
usageMetadata: GenerateContentResponseUsageMetadata | undefined
|
||||
): GeminiUsage {
|
||||
const promptTokenCount = usageMetadata?.promptTokenCount ?? 0
|
||||
const candidatesTokenCount = usageMetadata?.candidatesTokenCount ?? 0
|
||||
const thoughtsTokenCount = usageMetadata?.thoughtsTokenCount ?? 0
|
||||
const toolUsePromptTokenCount = usageMetadata?.toolUsePromptTokenCount ?? 0
|
||||
const promptTokenCount = (usageMetadata?.promptTokenCount ?? 0) + toolUsePromptTokenCount
|
||||
const candidatesTokenCount = (usageMetadata?.candidatesTokenCount ?? 0) + thoughtsTokenCount
|
||||
return {
|
||||
promptTokenCount,
|
||||
candidatesTokenCount,
|
||||
totalTokenCount: usageMetadata?.totalTokenCount ?? promptTokenCount + candidatesTokenCount,
|
||||
totalTokenCount: usageMetadata?.totalTokenCount ?? 0,
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
Reference in New Issue
Block a user