Add cost calculation support for LiteLLM provider (RooCodeInc#2403)

* Add cost calculation support for LiteLLM provider

* Calculate input/output cost up front once

Author: mrexodia

.changeset/purple-queens-teach.md

@@ -0,0 +1,5 @@
+---
+"claude-dev": minor
+---
+
+Add cost calculation support for LiteLLM provider

src/api/providers/litellm.ts

@@ -17,6 +17,40 @@ export class LiteLlmHandler implements ApiHandler {
 		})
 	}
 
+	async calculateCost(prompt_tokens: number, completion_tokens: number): Promise<number | undefined> {
+		// Reference: https://github.com/BerriAI/litellm/blob/122ee634f434014267af104814022af1d9a0882f/litellm/proxy/spend_tracking/spend_management_endpoints.py#L1473
+		const modelId = this.options.liteLlmModelId || liteLlmDefaultModelId
+		try {
+			const response = await fetch(`${this.client.baseURL}/spend/calculate`, {
+				method: "POST",
+				headers: {
+					"Content-Type": "application/json",
+					Authorization: `Bearer ${this.options.liteLlmApiKey}`,
+				},
+				body: JSON.stringify({
+					completion_response: {
+						model: modelId,
+						usage: {
+							prompt_tokens,
+							completion_tokens,
+						},
+					},
+				}),
+			})
+
+			if (response.ok) {
+				const data: { cost: number } = await response.json()
+				return data.cost
+			} else {
+				console.error("Error calculating spend:", response.statusText)
+				return undefined
+			}
+		} catch (error) {
+			console.error("Error calculating spend:", error)
+			return undefined
+		}
+	}
+
 	async *createMessage(systemPrompt: string, messages: Anthropic.Messages.MessageParam[]): ApiStream {
 		const formattedMessages = convertToOpenAiMessages(messages)
 		const systemMessage: OpenAI.Chat.ChatCompletionSystemMessageParam = {
@@ -39,6 +73,9 @@ export class LiteLlmHandler implements ApiHandler {
 			stream_options: { include_usage: true },
 		})
 
+		const inputCost = (await this.calculateCost(1e6, 0)) || 0
+		const outputCost = (await this.calculateCost(0, 1e6)) || 0
+
 		for await (const chunk of stream) {
 			const delta = chunk.choices[0]?.delta
 			if (delta?.content) {
@@ -49,10 +86,13 @@ export class LiteLlmHandler implements ApiHandler {
 			}
 
 			if (chunk.usage) {
+				const totalCost =
+					(inputCost * chunk.usage.prompt_tokens) / 1e6 + (outputCost * chunk.usage.completion_tokens) / 1e6
 				yield {
 					type: "usage",
 					inputTokens: chunk.usage.prompt_tokens || 0,
 					outputTokens: chunk.usage.completion_tokens || 0,
+					totalCost,
 				}
 			}
 		}