v0.5.10 在根目录生成统计信息

src/plugins/models/utils_model.py

@@ -8,6 +8,8 @@ from nonebot import get_driver
 from loguru import logger
 from ..chat.config import global_config
 from ..chat.utils_image import compress_base64_image_by_scale
+from datetime import datetime
+from ...common.database import Database
 
 driver = get_driver()
 config = driver.config
@@ -24,6 +26,75 @@ class LLM_request:
             raise ValueError(f"配置错误：找不到对应的配置项 - {str(e)}") from e
         self.model_name = model["name"]
         self.params = kwargs
+        
+        self.pri_in = model.get("pri_in", 0)
+        self.pri_out = model.get("pri_out", 0)
+        
+        # 获取数据库实例
+        self.db = Database.get_instance()
+        self._init_database()
+
+    def _init_database(self):
+        """初始化数据库集合"""
+        try:
+            # 创建llm_usage集合的索引
+            self.db.db.llm_usage.create_index([("timestamp", 1)])
+            self.db.db.llm_usage.create_index([("model_name", 1)])
+            self.db.db.llm_usage.create_index([("user_id", 1)])
+            self.db.db.llm_usage.create_index([("request_type", 1)])
+        except Exception as e:
+            logger.error(f"创建数据库索引失败: {e}")
+
+    def _record_usage(self, prompt_tokens: int, completion_tokens: int, total_tokens: int, 
+                     user_id: str = "system", request_type: str = "chat", 
+                     endpoint: str = "/chat/completions"):
+        """记录模型使用情况到数据库
+        Args:
+            prompt_tokens: 输入token数
+            completion_tokens: 输出token数
+            total_tokens: 总token数
+            user_id: 用户ID，默认为system
+            request_type: 请求类型(chat/embedding/image等)
+            endpoint: API端点
+        """
+        try:
+            usage_data = {
+                "model_name": self.model_name,
+                "user_id": user_id,
+                "request_type": request_type,
+                "endpoint": endpoint,
+                "prompt_tokens": prompt_tokens,
+                "completion_tokens": completion_tokens,
+                "total_tokens": total_tokens,
+                "cost": self._calculate_cost(prompt_tokens, completion_tokens),
+                "status": "success",
+                "timestamp": datetime.now()
+            }
+            self.db.db.llm_usage.insert_one(usage_data)
+            logger.info(
+                f"Token使用情况 - 模型: {self.model_name}, "
+                f"用户: {user_id}, 类型: {request_type}, "
+                f"提示词: {prompt_tokens}, 完成: {completion_tokens}, "
+                f"总计: {total_tokens}"
+            )
+        except Exception as e:
+            logger.error(f"记录token使用情况失败: {e}")
+
+    def _calculate_cost(self, prompt_tokens: int, completion_tokens: int) -> float:
+        """计算API调用成本
+        使用模型的pri_in和pri_out价格计算输入和输出的成本
+        
+        Args:
+            prompt_tokens: 输入token数量
+            completion_tokens: 输出token数量
+            
+        Returns:
+            float: 总成本（元）
+        """
+        # 使用模型的pri_in和pri_out计算成本
+        input_cost = (prompt_tokens / 1000000) * self.pri_in
+        output_cost = (completion_tokens / 1000000) * self.pri_out
+        return round(input_cost + output_cost, 6)
 
     async def _execute_request(
             self,
@@ -33,6 +104,8 @@ class LLM_request:
             payload: dict = None,
             retry_policy: dict = None,
             response_handler: callable = None,
+            user_id: str = "system",
+            request_type: str = "chat"
     ):
         """统一请求执行入口
         Args:
@@ -40,10 +113,10 @@ class LLM_request:
             prompt: prompt文本
             image_base64: 图片的base64编码
             payload: 请求体数据
-            is_async: 是否异步
             retry_policy: 自定义重试策略
-                (示例: {"max_retries":3, "base_wait":15, "retry_codes":[429,500]})
             response_handler: 自定义响应处理器
+            user_id: 用户ID
+            request_type: 请求类型
         """
         # 合并重试策略
         default_retry = {
@@ -105,7 +178,7 @@ class LLM_request:
                         result = await response.json()
 
                         # 使用自定义处理器或默认处理
-                        return response_handler(result) if response_handler else self._default_response_handler(result)
+                        return response_handler(result) if response_handler else self._default_response_handler(result, user_id, request_type, endpoint)
 
             except Exception as e:
                 if retry < policy["max_retries"] - 1:
@@ -145,7 +218,8 @@ class LLM_request:
                 **self.params
             }
 
-    def _default_response_handler(self, result: dict) -> Tuple:
+    def _default_response_handler(self, result: dict, user_id: str = "system", 
+                                request_type: str = "chat", endpoint: str = "/chat/completions") -> Tuple:
         """默认响应解析"""
         if "choices" in result and result["choices"]:
             message = result["choices"][0]["message"]
@@ -157,6 +231,21 @@ class LLM_request:
                 if not reasoning_content:
                     reasoning_content = reasoning
 
+            # 记录token使用情况
+            usage = result.get("usage", {})
+            if usage:
+                prompt_tokens = usage.get("prompt_tokens", 0)
+                completion_tokens = usage.get("completion_tokens", 0)
+                total_tokens = usage.get("total_tokens", 0)
+                self._record_usage(
+                    prompt_tokens=prompt_tokens,
+                    completion_tokens=completion_tokens,
+                    total_tokens=total_tokens,
+                    user_id=user_id,
+                    request_type=request_type,
+                    endpoint=endpoint
+                )
+
             return content, reasoning_content
 
         return "没有返回结果", ""
@@ -244,3 +333,4 @@ class LLM_request:
             response_handler=embedding_handler
         )
         return embedding
+