From 8d4af1b8958ec9d1f6be9c023599843f0e6b0d76 Mon Sep 17 00:00:00 2001
From: dreamhunter2333 <dreamhunter2333@gmail.com>
Date: Fri, 27 Jun 2025 02:08:16 +0800
Subject: [PATCH] =?UTF-8?q?=E9=87=8D=E6=9E=84=20TokenUtils:=20=E4=BB=8E?=
 =?UTF-8?q?=E7=B1=BB=E6=94=B9=E4=B8=BA=E5=AF=B9=E8=B1=A1=EF=BC=8C=E6=B7=BB?=
 =?UTF-8?q?=E5=8A=A0=E6=96=B0=E7=9A=84=20pricing=20=E5=92=8C=20usage=20?=
 =?UTF-8?q?=E5=A4=84=E7=90=86=E6=96=B9=E6=B3=95?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

- 将 TokenUtils 从类重构为对象
- 添加 getPricing 方法，支持 channel-specific 和 global pricing
- 添加 processUsage 方法，集中处理使用量计算和记录
- 在 providers/index.ts 中简化 usage 处理逻辑
- 添加 API key 掩码功能以保护敏感信息
---
 src/admin/token_utils.ts | 46 ++++++++++++++++++++++++++++++++++++----
 src/providers/index.ts   | 37 ++++----------------------------
 2 files changed, 46 insertions(+), 37 deletions(-)

diff --git a/src/admin/token_utils.ts b/src/admin/token_utils.ts
index c4ba3c9..d96aaf4 100644
--- a/src/admin/token_utils.ts
+++ b/src/admin/token_utils.ts
@@ -1,8 +1,11 @@
 import { Context } from "hono";
 
-// Token 工具类
-export class TokenUtils {
-    static async updateUsage(c: Context<HonoCustomType>, key: string, usageAmount: number): Promise<boolean> {
+import { CONSTANTS } from "../constants";
+import utils from "../utils";
+
+// Token 工具对象
+export const TokenUtils = {
+    async updateUsage(c: Context<HonoCustomType>, key: string, usageAmount: number): Promise<boolean> {
         try {
             const result = await c.env.DB.prepare(
                 `UPDATE api_token SET usage = usage + ?, updated_at = datetime('now') WHERE key = ?`
@@ -13,5 +16,40 @@ export class TokenUtils {
             console.error('Error updating token usage:', error);
             return false;
         }
+    },
+    async getPricing(c: Context<HonoCustomType>, model: string, channelConfig: ChannelConfig): Promise<ModelPricing | null> {
+        // Check channel-specific pricing first
+        if (channelConfig?.model_pricing?.[model]) {
+            return channelConfig.model_pricing[model];
+        }
+
+        // Fallback to global pricing
+        const globalPricingMap = await utils.getJsonSetting(c, CONSTANTS.MODEL_PRICING_KEY);
+        return globalPricingMap?.[model] || null;
+    },
+
+    async processUsage(c: Context<HonoCustomType>, apiKey: string, model: string, targetChannelKey: string, targetChannelConfig: ChannelConfig, usage: Usage): Promise<void> {
+        console.log("Usage data:", usage);
+
+        // Get pricing and calculate cost
+        const pricing = await this.getPricing(c, model, targetChannelConfig);
+        const hasTokens = usage.prompt_tokens && usage.completion_tokens;
+
+        if (pricing && hasTokens) {
+            const inputCost = usage.prompt_tokens! * pricing.input;
+            const outputCost = usage.completion_tokens! * pricing.output;
+            const totalCost = inputCost + outputCost;
+            // Update token usage and log data
+            await this.updateUsage(c, apiKey, totalCost);
+            // 长度不够直接显示 *
+            var maskedApiKey = apiKey.length < 3 ? '*'.repeat(apiKey.length) : (
+                apiKey.slice(0, apiKey.length / 3)
+                + '*'.repeat(apiKey.length / 3)
+                + apiKey.slice((2 * apiKey.length) / 3)
+            );
+            console.log(`Model: ${model}, Channel: ${targetChannelKey}, apiKey: ${maskedApiKey}, Cost: ${totalCost}`);
+        } else {
+            console.warn(`No pricing found for model: ${model} in channel: ${targetChannelKey}`);
+        }
     }
-}
+};
diff --git a/src/providers/index.ts b/src/providers/index.ts
index bc98aa7..6cce98b 100644
--- a/src/providers/index.ts
+++ b/src/providers/index.ts
@@ -123,39 +123,10 @@ class ProxyEndpoint extends OpenAPIRoute {
 
         return proxyFetch(c, targetChannelConfig,
             async (usage: Usage) => {
-                console.log("Usage data:", usage);
-
-                // Calculate cost based on usage and model pricing
-                // Model was already extracted above
-
-                // Get global pricing as base configuration
-                let pricing = null;
-                const globalPricing = await utils.getSetting(c, CONSTANTS.MODEL_PRICING_KEY);
-                if (globalPricing) {
-                    const globalPricingMap = JSON.parse(globalPricing);
-                    pricing = globalPricingMap[model];
-                }
-
-                // Override with channel-specific pricing if available
-                if (targetChannelConfig.model_pricing?.[model]) {
-                    pricing = targetChannelConfig.model_pricing[model];
-                }
-
-                if (pricing && usage.prompt_tokens && usage.completion_tokens) {
-                    const inputCost = usage.prompt_tokens * pricing.input;
-                    const outputCost = usage.completion_tokens * pricing.output;
-                    const totalCost = inputCost + outputCost;
-
-                    // Update token usage using TokenUtils
-                    const updated = await TokenUtils.updateUsage(c, apiKey, totalCost);
-                    if (!updated) {
-                        console.error("Failed to update token usage");
-                    }
-
-                    console.log(`Model: ${model}, Channel: ${targetChannelKey}, Channel pricing: ${!!targetChannelConfig.model_pricing?.[model]}, Cost: ${totalCost}, Updated: ${updated}`);
-                }
-                else {
-                    console.warn(`No pricing found for model: ${model} in channel: ${targetChannelKey}`);
+                try {
+                    await TokenUtils.processUsage(c, apiKey, model, targetChannelKey, targetChannelConfig, usage);
+                } catch (error) {
+                    console.error('Error processing usage:', error);
                 }
             }
         );