From 57b36e16b8fbafa0ba31bb92cc6cb3afcfdb138b Mon Sep 17 00:00:00 2001 From: Saoud Rizwan <7799382+saoudrizwan@users.noreply.github.com> Date: Sat, 1 Mar 2025 03:11:45 -0800 Subject: [PATCH] Fix pricing calculation --- src/api/providers/deepseek.ts | 42 ++++++++++++++++++------ src/api/providers/openai-native.ts | 51 ++++++++++++++++++------------ 2 files changed, 63 insertions(+), 30 deletions(-) diff --git a/src/api/providers/deepseek.ts b/src/api/providers/deepseek.ts index 9049e646db..35373b7cce 100644 --- a/src/api/providers/deepseek.ts +++ b/src/api/providers/deepseek.ts @@ -3,6 +3,7 @@ import OpenAI from "openai" import { withRetry } from "../retry" import { ApiHandler } from "../" import { ApiHandlerOptions, DeepSeekModelId, ModelInfo, deepSeekDefaultModelId, deepSeekModels } from "../../shared/api" +import { calculateApiCostOpenAI } from "../../utils/cost" import { convertToOpenAiMessages } from "../transform/openai-format" import { ApiStream } from "../transform/stream" import { convertToR1Format } from "../transform/r1-format" @@ -19,6 +20,37 @@ export class DeepSeekHandler implements ApiHandler { }) } + private async *yieldUsage(info: ModelInfo, usage: OpenAI.Completions.CompletionUsage | undefined): ApiStream { + // Deepseek reports total input AND cache reads/writes, + // see context caching: https://api-docs.deepseek.com/guides/kv_cache) + // where the input tokens is the sum of the cache hits/misses, just like OpenAI. + // This affects: + // 1) context management truncation algorithm, and + // 2) cost calculation + + // Deepseek usage includes extra fields. + // Safely cast the prompt token details section to the appropriate structure. + interface DeepSeekUsage extends OpenAI.CompletionUsage { + prompt_cache_hit_tokens?: number + prompt_cache_miss_tokens?: number + } + const deepUsage = usage as DeepSeekUsage + + const inputTokens = deepUsage?.prompt_tokens || 0 + const outputTokens = deepUsage?.completion_tokens || 0 + const cacheReadTokens = deepUsage?.prompt_cache_hit_tokens || 0 + const cacheWriteTokens = deepUsage?.prompt_cache_miss_tokens || 0 + const totalCost = calculateApiCostOpenAI(info, inputTokens, outputTokens, cacheWriteTokens, cacheReadTokens) + yield { + type: "usage", + inputTokens: inputTokens, + outputTokens: outputTokens, + cacheWriteTokens: cacheWriteTokens, + cacheReadTokens: cacheReadTokens, + totalCost: totalCost, + } + } + @withRetry() async *createMessage(systemPrompt: string, messages: Anthropic.Messages.MessageParam[]): ApiStream { const model = this.getModel() @@ -61,15 +93,7 @@ export class DeepSeekHandler implements ApiHandler { } if (chunk.usage) { - yield { - type: "usage", - inputTokens: chunk.usage.prompt_tokens || 0, // (deepseek reports total input AND cache reads/writes, see context caching: https://api-docs.deepseek.com/guides/kv_cache) where the input tokens is the sum of the cache hits/misses, while anthropic reports them as separate tokens. This is important to know for 1) context management truncation algorithm, and 2) cost calculation (NOTE: we report both input and cache stats but for now set input price to 0 since all the cost calculation will be done using cache hits/misses) - outputTokens: chunk.usage.completion_tokens || 0, - // @ts-ignore-next-line - cacheReadTokens: chunk.usage.prompt_cache_hit_tokens || 0, - // @ts-ignore-next-line - cacheWriteTokens: chunk.usage.prompt_cache_miss_tokens || 0, - } + yield* this.yieldUsage(model.info, chunk.usage) } } } diff --git a/src/api/providers/openai-native.ts b/src/api/providers/openai-native.ts index c90aa6f862..8fadc60485 100644 --- a/src/api/providers/openai-native.ts +++ b/src/api/providers/openai-native.ts @@ -10,6 +10,7 @@ import { openAiNativeModels, } from "../../shared/api" import { convertToOpenAiMessages } from "../transform/openai-format" +import { calculateApiCostOpenAI } from "../../utils/cost" import { ApiStream } from "../transform/stream" import { ChatCompletionReasoningEffort } from "openai/resources/chat/completions.mjs" @@ -24,31 +25,47 @@ export class OpenAiNativeHandler implements ApiHandler { }) } + private async *yieldUsage(info: ModelInfo, usage: OpenAI.Completions.CompletionUsage | undefined): ApiStream { + const inputTokens = usage?.prompt_tokens || 0 + const outputTokens = usage?.completion_tokens || 0 + const cacheReadTokens = usage?.prompt_tokens_details?.cached_tokens || 0 + const cacheWriteTokens = 0 + const totalCost = calculateApiCostOpenAI(info, inputTokens, outputTokens, cacheWriteTokens, cacheReadTokens) + yield { + type: "usage", + inputTokens: inputTokens, + outputTokens: outputTokens, + cacheWriteTokens: cacheWriteTokens, + cacheReadTokens: cacheReadTokens, + totalCost: totalCost, + } + } + @withRetry() async *createMessage(systemPrompt: string, messages: Anthropic.Messages.MessageParam[]): ApiStream { - switch (this.getModel().id) { + const model = this.getModel() + + switch (model.id) { case "o1": case "o1-preview": case "o1-mini": { // o1 doesnt support streaming, non-1 temp, or system prompt const response = await this.client.chat.completions.create({ - model: this.getModel().id, + model: model.id, messages: [{ role: "user", content: systemPrompt }, ...convertToOpenAiMessages(messages)], }) yield { type: "text", text: response.choices[0]?.message.content || "", } - yield { - type: "usage", - inputTokens: response.usage?.prompt_tokens || 0, - outputTokens: response.usage?.completion_tokens || 0, - } + + yield* this.yieldUsage(model.info, response.usage) + break } case "o3-mini": { const stream = await this.client.chat.completions.create({ - model: this.getModel().id, + model: model.id, messages: [{ role: "developer", content: systemPrompt }, ...convertToOpenAiMessages(messages)], stream: true, stream_options: { include_usage: true }, @@ -63,18 +80,15 @@ export class OpenAiNativeHandler implements ApiHandler { } } if (chunk.usage) { - yield { - type: "usage", - inputTokens: chunk.usage.prompt_tokens || 0, - outputTokens: chunk.usage.completion_tokens || 0, - } + // Only last chunk contains usage + yield* this.yieldUsage(model.info, chunk.usage) } } break } default: { const stream = await this.client.chat.completions.create({ - model: this.getModel().id, + model: model.id, // max_completion_tokens: this.getModel().info.maxTokens, temperature: 0, messages: [{ role: "system", content: systemPrompt }, ...convertToOpenAiMessages(messages)], @@ -90,14 +104,9 @@ export class OpenAiNativeHandler implements ApiHandler { text: delta.content, } } - - // contains a null value except for the last chunk which contains the token usage statistics for the entire request if (chunk.usage) { - yield { - type: "usage", - inputTokens: chunk.usage.prompt_tokens || 0, - outputTokens: chunk.usage.completion_tokens || 0, - } + // Only last chunk contains usage + yield* this.yieldUsage(model.info, chunk.usage) } } }