Fix pricing calculation

This commit is contained in:
Saoud Rizwan 2025-03-01 03:11:45 -08:00
parent c891d24ef6
commit 57b36e16b8
2 changed files with 63 additions and 30 deletions

View file

@ -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)
}
}
}

View file

@ -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)
}
}
}