diff --git a/src/api/providers/anthropic.ts b/src/api/providers/anthropic.ts index 412f5de621..52dec1ae55 100644 --- a/src/api/providers/anthropic.ts +++ b/src/api/providers/anthropic.ts @@ -17,6 +17,7 @@ import { getModelParams } from "../transform/model-params" import { BaseProvider } from "./base-provider" import type { SingleCompletionHandler, ApiHandlerCreateMessageMetadata } from "../index" +import { calculateApiCostAnthropic } from "../../shared/cost" export class AnthropicHandler extends BaseProvider implements SingleCompletionHandler { private options: ApiHandlerOptions @@ -132,20 +133,35 @@ export class AnthropicHandler extends BaseProvider implements SingleCompletionHa } } + let inputTokens = 0 + let outputTokens = 0 + let cacheWriteTokens = 0 + let cacheReadTokens = 0 + for await (const chunk of stream) { switch (chunk.type) { case "message_start": { // Tells us cache reads/writes/input/output. - const usage = chunk.message.usage + const { + input_tokens = 0, + output_tokens = 0, + cache_creation_input_tokens, + cache_read_input_tokens, + } = chunk.message.usage yield { type: "usage", - inputTokens: usage.input_tokens || 0, - outputTokens: usage.output_tokens || 0, - cacheWriteTokens: usage.cache_creation_input_tokens || undefined, - cacheReadTokens: usage.cache_read_input_tokens || undefined, + inputTokens: input_tokens, + outputTokens: output_tokens, + cacheWriteTokens: cache_creation_input_tokens || undefined, + cacheReadTokens: cache_read_input_tokens || undefined, } + inputTokens += input_tokens + outputTokens += output_tokens + cacheWriteTokens += cache_creation_input_tokens || 0 + cacheReadTokens += cache_read_input_tokens || 0 + break } case "message_delta": @@ -198,6 +214,21 @@ export class AnthropicHandler extends BaseProvider implements SingleCompletionHa break } } + + if (inputTokens > 0 || outputTokens > 0 || cacheWriteTokens > 0 || cacheReadTokens > 0) { + yield { + type: "usage", + inputTokens: 0, + outputTokens: 0, + totalCost: calculateApiCostAnthropic( + this.getModel().info, + inputTokens, + outputTokens, + cacheWriteTokens, + cacheReadTokens, + ), + } + } } getModel() { diff --git a/src/core/task/Task.ts b/src/core/task/Task.ts index 6996b60647..6f6f2d684a 100644 --- a/src/core/task/Task.ts +++ b/src/core/task/Task.ts @@ -1439,19 +1439,22 @@ export class Task extends EventEmitter { } finally { this.isStreaming = false } - if ( - inputTokens > 0 || - outputTokens > 0 || - cacheWriteTokens > 0 || - cacheReadTokens > 0 || - typeof totalCost !== "undefined" - ) { + + if (inputTokens > 0 || outputTokens > 0 || cacheWriteTokens > 0 || cacheReadTokens > 0) { TelemetryService.instance.captureLlmCompletion(this.taskId, { inputTokens, outputTokens, cacheWriteTokens, cacheReadTokens, - cost: totalCost, + cost: + totalCost ?? + calculateApiCostAnthropic( + this.api.getModel().info, + inputTokens, + outputTokens, + cacheWriteTokens, + cacheReadTokens, + ), }) } diff --git a/src/services/checkpoints/__tests__/ShadowCheckpointService.spec.ts b/src/services/checkpoints/__tests__/ShadowCheckpointService.spec.ts index ddfca7fc6d..a403c4e851 100644 --- a/src/services/checkpoints/__tests__/ShadowCheckpointService.spec.ts +++ b/src/services/checkpoints/__tests__/ShadowCheckpointService.spec.ts @@ -12,7 +12,7 @@ import * as fileSearch from "../../../services/search/file-search" import { RepoPerTaskCheckpointService } from "../RepoPerTaskCheckpointService" -vitest.setConfig({ testTimeout: 10_000 }) +vitest.setConfig({ testTimeout: 20_000 }) const tmpDir = path.join(os.tmpdir(), "CheckpointService") diff --git a/src/shared/cost.ts b/src/shared/cost.ts index 3257cab16c..a628756b0d 100644 --- a/src/shared/cost.ts +++ b/src/shared/cost.ts @@ -15,7 +15,8 @@ function calculateApiCostInternal( return totalCost } -// For Anthropic compliant usage, the input tokens count does NOT include the cached tokens +// For Anthropic compliant usage, the input tokens count does NOT include the +// cached tokens. export function calculateApiCostAnthropic( modelInfo: ModelInfo, inputTokens: number, @@ -23,18 +24,16 @@ export function calculateApiCostAnthropic( cacheCreationInputTokens?: number, cacheReadInputTokens?: number, ): number { - const cacheCreationInputTokensNum = cacheCreationInputTokens || 0 - const cacheReadInputTokensNum = cacheReadInputTokens || 0 return calculateApiCostInternal( modelInfo, inputTokens, outputTokens, - cacheCreationInputTokensNum, - cacheReadInputTokensNum, + cacheCreationInputTokens || 0, + cacheReadInputTokens || 0, ) } -// For OpenAI compliant usage, the input tokens count INCLUDES the cached tokens +// For OpenAI compliant usage, the input tokens count INCLUDES the cached tokens. export function calculateApiCostOpenAI( modelInfo: ModelInfo, inputTokens: number, @@ -45,6 +44,7 @@ export function calculateApiCostOpenAI( const cacheCreationInputTokensNum = cacheCreationInputTokens || 0 const cacheReadInputTokensNum = cacheReadInputTokens || 0 const nonCachedInputTokens = Math.max(0, inputTokens - cacheCreationInputTokensNum - cacheReadInputTokensNum) + return calculateApiCostInternal( modelInfo, nonCachedInputTokens,