diff --git a/src/api/providers/baseten.ts b/src/api/providers/baseten.ts index d33f059ab6..22389e6276 100644 --- a/src/api/providers/baseten.ts +++ b/src/api/providers/baseten.ts @@ -4,7 +4,7 @@ import OpenAI from "openai" import type { ApiHandlerOptions } from "../../shared/api" import type { ApiHandlerCreateMessageMetadata } from "../index" -import { ApiStream } from "../transform/stream" +import { ApiStream, ApiStreamUsageChunk } from "../transform/stream" import { calculateApiCostOpenAI } from "../../shared/cost" import { BaseOpenAiCompatibleProvider } from "./base-openai-compatible-provider" @@ -29,6 +29,8 @@ export class BasetenHandler extends BaseOpenAiCompatibleProvider ): ApiStream { const stream = await this.createStream(systemPrompt, messages, metadata) + let lastUsage: OpenAI.CompletionUsage | undefined + for await (const chunk of stream) { const delta = chunk.choices[0]?.delta @@ -39,37 +41,41 @@ export class BasetenHandler extends BaseOpenAiCompatibleProvider } } - if (chunk.usage) { - // Check if usage has cached token details (some models support this) - const usage = chunk.usage as any - const inputTokens = usage?.prompt_tokens || 0 - const outputTokens = usage?.completion_tokens || 0 - - // Check for cached tokens in various possible locations - const cacheReadTokens = - usage?.prompt_tokens_details?.cached_tokens || usage?.prompt_cache_hit_tokens || 0 - - // Baseten currently doesn't track cache writes - const cacheWriteTokens = 0 - - // Calculate cost using OpenAI-compatible cost calculation - const { totalCost } = calculateApiCostOpenAI( - this.getModel().info, - inputTokens, - outputTokens, - cacheWriteTokens, - cacheReadTokens, - ) - - yield { - type: "usage", - inputTokens, - outputTokens, - cacheWriteTokens, - cacheReadTokens, - totalCost, - } + if (delta && "reasoning_content" in delta && delta.reasoning_content) { + yield { type: "reasoning", text: (delta.reasoning_content as string | undefined) || "" } } + + if (chunk.usage) { + lastUsage = chunk.usage + } + } + + if (lastUsage) { + yield this.processUsageMetrics(lastUsage) + } + } + + protected processUsageMetrics(usage: any): ApiStreamUsageChunk { + const inputTokens = usage?.prompt_tokens || 0 + const outputTokens = usage?.completion_tokens || 0 + const cacheWriteTokens = usage?.prompt_tokens_details?.cache_write_tokens || 0 + const cacheReadTokens = usage?.prompt_tokens_details?.cached_tokens || 0 + + const { totalCost } = calculateApiCostOpenAI( + this.getModel().info, + inputTokens, + outputTokens, + cacheWriteTokens, + cacheReadTokens, + ) + + return { + type: "usage", + inputTokens, + outputTokens, + cacheWriteTokens: cacheWriteTokens || undefined, + cacheReadTokens: cacheReadTokens || undefined, + totalCost, } } }