From bb12b488871051ab7d03a0134b12b6ab9d959e5a Mon Sep 17 00:00:00 2001
From: RockChinQ <1010553892@qq.com>
Date: Fri, 22 Dec 2023 12:38:27 +0800
Subject: [PATCH] =?UTF-8?q?feat:=20usage.query=E5=AE=8C=E6=88=90?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 main.py               |  2 +-
 pkg/openai/session.py | 22 ++++++++++++++++++++++
 pkg/qqbot/manager.py  |  2 +-
 3 files changed, 24 insertions(+), 2 deletions(-)

diff --git a/main.py b/main.py
index 48971941..5ceb67f6 100644
--- a/main.py
+++ b/main.py
@@ -222,7 +222,7 @@ async def start_process(first_time_init=False):
                     "platform": sys.platform,
                 },
                 runtime_info={
-                    "admin_qq": cfg['admin_qq'],
+                    "admin_id": "{}".format(cfg['admin_qq']),
                     "msg_source": cfg['msg_source_adapter'],
                 }
             )
diff --git a/pkg/openai/session.py b/pkg/openai/session.py
index 197ccfd0..1f9c76de 100644
--- a/pkg/openai/session.py
+++ b/pkg/openai/session.py
@@ -261,6 +261,8 @@ class Session:
 
         pending_res_text = ""
 
+        start_time = time.time()
+
         # TODO 对不起，我知道这样非常非常屎山，但我之后会重构的
         for resp in context.get_openai_manager().request_completion(prompts):
 
@@ -349,6 +351,26 @@ class Session:
             self.just_switched_to_exist_session = False
             self.set_ongoing()
 
+        # 上报使用量数据
+        session_type = session_name_spt[0]
+        session_id = session_name_spt[1]
+
+        ability_provider = "QChatGPT.Text"
+        usage = total_tokens
+        model_name = context.get_config_manager().data['completion_api_params']['model']
+        response_seconds = int(time.time() - start_time)
+        retry_times = -1 # 暂不记录
+
+        context.get_center_v2_api().usage.post_query_record(
+            session_type=session_type,
+            session_id=session_id,
+            query_ability_provider=ability_provider,
+            usage=usage,
+            model_name=model_name,
+            response_seconds=response_seconds,
+            retry_times=retry_times
+        )
+
         return res_ans if res_ans[0] != '\n' else res_ans[1:], finish_reason, funcs
 
     # 删除上一回合并返回上一回合的问题
diff --git a/pkg/qqbot/manager.py b/pkg/qqbot/manager.py
index b8ba8a9c..8cd663ff 100644
--- a/pkg/qqbot/manager.py
+++ b/pkg/qqbot/manager.py
@@ -128,7 +128,7 @@ class QQBotManager:
         
         # 保存 account_id 到审计模块
         from ..utils.center import apigroup
-        apigroup.APIGroup._runtime_info['account_id'] = self.bot_account_id
+        apigroup.APIGroup._runtime_info['account_id'] = "{}".format(self.bot_account_id)
 
         context.set_qqbot_manager(self)