Fix usage tracking

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<BasetenModelId>
 	): 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<BasetenModelId>
 				}
 			}
 
-			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,
 		}
 	}
 }