From c51f59e50b6642494ed5c5addf6897e9a9827ae5 Mon Sep 17 00:00:00 2001
From: Daniel Trugman <dtrugman@gmail.com>
Date: Fri, 7 Mar 2025 14:57:11 +0000
Subject: [PATCH] Requesty: Correctly calculate request costs

---
 src/api/providers/__tests__/requesty.test.ts | 26 +++++++++++-----
 src/api/providers/openai.ts                  |  6 ++--
 src/api/providers/requesty.ts                | 32 ++++++++++++++++----
 3 files changed, 47 insertions(+), 17 deletions(-)

diff --git a/src/api/providers/__tests__/requesty.test.ts b/src/api/providers/__tests__/requesty.test.ts
index 7867b15ebc..47921a1c53 100644
--- a/src/api/providers/__tests__/requesty.test.ts
+++ b/src/api/providers/__tests__/requesty.test.ts
@@ -22,8 +22,10 @@ describe("RequestyHandler", () => {
 			contextWindow: 4000,
 			supportsPromptCache: false,
 			supportsImages: true,
-			inputPrice: 0,
-			outputPrice: 0,
+			inputPrice: 1,
+			outputPrice: 10,
+			cacheReadsPrice: 0.1,
+			cacheWritesPrice: 1.5,
 		},
 		openAiStreamingEnabled: true,
 		includeMaxTokens: true, // Add this to match the implementation
@@ -83,8 +85,12 @@ describe("RequestyHandler", () => {
 						yield {
 							choices: [{ delta: { content: " world" } }],
 							usage: {
-								prompt_tokens: 10,
-								completion_tokens: 5,
+								prompt_tokens: 30,
+								completion_tokens: 10,
+								prompt_tokens_details: {
+									cached_tokens: 15,
+									caching_tokens: 5,
+								},
 							},
 						}
 					},
@@ -105,10 +111,11 @@ describe("RequestyHandler", () => {
 					{ type: "text", text: " world" },
 					{
 						type: "usage",
-						inputTokens: 10,
-						outputTokens: 5,
-						cacheWriteTokens: undefined,
-						cacheReadTokens: undefined,
+						inputTokens: 30,
+						outputTokens: 10,
+						cacheWriteTokens: 5,
+						cacheReadTokens: 15,
+						totalCost: 0.000119, // (10 * 1 / 1,000,000) + (5 * 1.5 / 1,000,000) + (15 * 0.1 / 1,000,000) + (10 * 10 / 1,000,000)
 					},
 				])
 
@@ -182,6 +189,9 @@ describe("RequestyHandler", () => {
 						type: "usage",
 						inputTokens: 10,
 						outputTokens: 5,
+						cacheWriteTokens: 0,
+						cacheReadTokens: 0,
+						totalCost: 0.00006, // (10 * 1 / 1,000,000) + (5 * 10 / 1,000,000)
 					},
 				])
 
diff --git a/src/api/providers/openai.ts b/src/api/providers/openai.ts
index 9262f3b75a..5d85a86a5a 100644
--- a/src/api/providers/openai.ts
+++ b/src/api/providers/openai.ts
@@ -111,7 +111,7 @@ export class OpenAiHandler extends BaseProvider implements SingleCompletionHandl
 					}
 				}
 				if (chunk.usage) {
-					yield this.processUsageMetrics(chunk.usage)
+					yield this.processUsageMetrics(chunk.usage, modelInfo)
 				}
 			}
 		} else {
@@ -134,11 +134,11 @@ export class OpenAiHandler extends BaseProvider implements SingleCompletionHandl
 				type: "text",
 				text: response.choices[0]?.message.content || "",
 			}
-			yield this.processUsageMetrics(response.usage)
+			yield this.processUsageMetrics(response.usage, modelInfo)
 		}
 	}
 
-	protected processUsageMetrics(usage: any): ApiStreamUsageChunk {
+	protected processUsageMetrics(usage: any, modelInfo?: ModelInfo): ApiStreamUsageChunk {
 		return {
 			type: "usage",
 			inputTokens: usage?.prompt_tokens || 0,
diff --git a/src/api/providers/requesty.ts b/src/api/providers/requesty.ts
index 0060bfc5d8..434d6f4316 100644
--- a/src/api/providers/requesty.ts
+++ b/src/api/providers/requesty.ts
@@ -1,9 +1,20 @@
 import axios from "axios"
 
 import { ModelInfo, requestyModelInfoSaneDefaults, requestyDefaultModelId } from "../../shared/api"
-import { parseApiPrice } from "../../utils/cost"
+import { calculateApiCostOpenAI, parseApiPrice } from "../../utils/cost"
 import { ApiStreamUsageChunk } from "../transform/stream"
 import { OpenAiHandler, OpenAiHandlerOptions } from "./openai"
+import OpenAI from "openai"
+
+// Requesty usage includes an extra field for Anthropic use cases.
+// Safely cast the prompt token details section to the appropriate structure.
+interface RequestyUsage extends OpenAI.CompletionUsage {
+	prompt_tokens_details?: {
+		caching_tokens?: number
+		cached_tokens?: number
+	}
+	total_cost?: number
+}
 
 export class RequestyHandler extends OpenAiHandler {
 	constructor(options: OpenAiHandlerOptions) {
@@ -27,13 +38,22 @@ export class RequestyHandler extends OpenAiHandler {
 		}
 	}
 
-	protected override processUsageMetrics(usage: any): ApiStreamUsageChunk {
+	protected override processUsageMetrics(usage: any, modelInfo?: ModelInfo): ApiStreamUsageChunk {
+		const requestyUsage = usage as RequestyUsage
+		const inputTokens = requestyUsage?.prompt_tokens || 0
+		const outputTokens = requestyUsage?.completion_tokens || 0
+		const cacheWriteTokens = requestyUsage?.prompt_tokens_details?.caching_tokens || 0
+		const cacheReadTokens = requestyUsage?.prompt_tokens_details?.cached_tokens || 0
+		const totalCost = modelInfo
+			? calculateApiCostOpenAI(modelInfo, inputTokens, outputTokens, cacheWriteTokens, cacheReadTokens)
+			: 0
 		return {
 			type: "usage",
-			inputTokens: usage?.prompt_tokens || 0,
-			outputTokens: usage?.completion_tokens || 0,
-			cacheWriteTokens: usage?.cache_creation_input_tokens,
-			cacheReadTokens: usage?.cache_read_input_tokens,
+			inputTokens: inputTokens,
+			outputTokens: outputTokens,
+			cacheWriteTokens: cacheWriteTokens,
+			cacheReadTokens: cacheReadTokens,
+			totalCost: totalCost,
 		}
 	}
 }