From 3d9587f128a20b786c464a9f77ace143f4f426d8 Mon Sep 17 00:00:00 2001
From: neotf <10400594+neotf@users.noreply.github.com>
Date: Thu, 29 May 2025 22:24:29 +0800
Subject: [PATCH] feat: enhance cache_create_tokens calculation for OpenRouter

---
 dto/openai_response.go |  2 ++
 service/quota.go       | 27 +++++++++++++++++++++++++++
 2 files changed, 29 insertions(+)

diff --git a/dto/openai_response.go b/dto/openai_response.go
index 790d4df8..fb4aeb4c 100644
--- a/dto/openai_response.go
+++ b/dto/openai_response.go
@@ -178,6 +178,8 @@ type Usage struct {
 	InputTokens            int                `json:"input_tokens"`
 	OutputTokens           int                `json:"output_tokens"`
 	InputTokensDetails     *InputTokenDetails `json:"input_tokens_details"`
+	// OpenRouter Params
+	Cost float64 `json:"cost,omitempty"`
 }
 
 type InputTokenDetails struct {
diff --git a/service/quota.go b/service/quota.go
index 0d11b4a0..43297b4a 100644
--- a/service/quota.go
+++ b/service/quota.go
@@ -3,6 +3,7 @@ package service
 import (
 	"errors"
 	"fmt"
+	"math"
 	"one-api/common"
 	constant2 "one-api/constant"
 	"one-api/dto"
@@ -214,6 +215,11 @@ func PostClaudeConsumeQuota(ctx *gin.Context, relayInfo *relaycommon.RelayInfo,
 	cacheCreationRatio := priceData.CacheCreationRatio
 	cacheCreationTokens := usage.PromptTokensDetails.CachedCreationTokens
 
+	if relayInfo.ChannelType == common.ChannelTypeOpenRouter && priceData.CacheCreationRatio != 1 {
+		cacheCreationTokens = CalcOpenRouterCacheCreateTokens(*usage, priceData)
+		promptTokens = promptTokens - cacheCreationTokens - cacheTokens
+	}
+
 	calculateQuota := 0.0
 	if !priceData.UsePrice {
 		calculateQuota = float64(promptTokens)
@@ -261,6 +267,27 @@ func PostClaudeConsumeQuota(ctx *gin.Context, relayInfo *relaycommon.RelayInfo,
 		tokenName, quota, logContent, relayInfo.TokenId, userQuota, int(useTimeSeconds), relayInfo.IsStream, relayInfo.Group, other)
 }
 
+func CalcOpenRouterCacheCreateTokens(usage dto.Usage, priceData helper.PriceData) int {
+	if priceData.CacheCreationRatio == 1 {
+		return 0
+	}
+	quotaPrice := priceData.ModelRatio / common.QuotaPerUnit
+	promptCacheCreatePrice := quotaPrice * priceData.CacheCreationRatio
+	promptCacheReadPrice := quotaPrice * priceData.CacheRatio
+	completionPrice := quotaPrice * priceData.CompletionRatio
+
+	cost := usage.Cost
+	totalPromptTokens := float64(usage.PromptTokens)
+	completionTokens := float64(usage.CompletionTokens)
+	promptCacheReadTokens := float64(usage.PromptTokensDetails.CachedTokens)
+
+	return int(math.Round((cost -
+		totalPromptTokens*quotaPrice +
+		promptCacheReadTokens*(quotaPrice-promptCacheReadPrice) -
+		completionTokens*completionPrice) /
+		(promptCacheCreatePrice - quotaPrice)))
+}
+
 func PostAudioConsumeQuota(ctx *gin.Context, relayInfo *relaycommon.RelayInfo,
 	usage *dto.Usage, preConsumedQuota int, userQuota int, priceData helper.PriceData, extraContent string) {