Update chat_service: support Coze OAuth

agent_service.py

@@ -7,6 +7,7 @@ import time
 from typing import Dict, List, Tuple, Any, Optional
 import logging
 from datetime import datetime, timedelta
+import traceback
 
 import apscheduler.triggers.cron
 from apscheduler.schedulers.background import BackgroundScheduler
@@ -51,9 +52,12 @@ class AgentService:
         )
         # DeepSeek on Volces
         self.model_name = chat_service.VOLCENGINE_MODEL_DEEPSEEK_V3
+        coze_config = configs.get()['chat_api']['coze']
+        coze_oauth_app = CozeChat.get_oauth_app(
+            coze_config['oauth_client_id'], coze_config['private_key_path'], str(coze_config['public_key_id']))
         self.coze_client = CozeChat(
-            token=chat_service.COZE_API_TOKEN,
-            base_url=chat_service.COZE_CN_BASE_URL
+            base_url=chat_service.COZE_CN_BASE_URL,
+            auth_app=coze_oauth_app
         )
         self.chat_service_type = chat_service_type
 
@@ -86,6 +90,7 @@ class AgentService:
                     self.process_single_message(message)
                 except Exception as e:
                     logging.error("Error processing message: {}".format(e))
+                    traceback.print_exc()
                 # 无论处理是否有异常，都ACK消息
                 self.receive_queue.ack(message)
             time.sleep(1)
@@ -252,7 +257,8 @@ if __name__ == "__main__":
         chat_service_type=ChatServiceType.COZE_CHAT
     )
     # 只有企微场景需要主动发起
-    # service.setup_initiative_conversations({'second': '5,35'})
+    if not config['debug_flags'].get('disable_active_conversation', False):
+        service.setup_initiative_conversations({'second': '5,35'})
 
     process_thread = threading.Thread(target=service.process_messages)
     process_thread.start()

chat_service.py

@@ -4,10 +4,12 @@
 #
 
 import os
-from typing import List, Dict
+import threading
+from typing import List, Dict, Optional
 from enum import Enum, auto
 import logging
-from cozepy import Coze, TokenAuth, Message, ChatStatus, MessageContentType, ChatEventType, MessageType
+from cozepy import Coze, TokenAuth, Message, ChatStatus, MessageType, JWTOAuthApp, JWTAuth
+import time
 
 COZE_API_TOKEN = os.getenv("COZE_API_TOKEN")
 COZE_CN_BASE_URL = 'https://api.coze.cn'
@@ -21,8 +23,16 @@ class ChatServiceType(Enum):
     COZE_CHAT = auto()
 
 class CozeChat:
-    def __init__(self, token, base_url: str):
-        self.coze = Coze(auth=TokenAuth(token), base_url=base_url)
+    def __init__(self, base_url: str, auth_token: Optional[str] = None, auth_app: Optional[JWTOAuthApp] = None):
+        if not auth_token and not auth_app:
+            raise ValueError("Either auth_token or auth_app must be provided.")
+        if auth_token:
+            self.coze = Coze(auth=TokenAuth(auth_token), base_url=base_url)
+        else:
+            self.auth_app = auth_app
+            oauth_token = auth_app.get_access_token(ttl=12*3600)
+            self.coze = Coze(auth=JWTAuth(oauth_app=auth_app), base_url=base_url)
+            self.setup_token_refresh()
 
     def create(self, bot_id: str, user_id: str, messages: List, custom_variables: Dict):
         response = self.coze.chat.create_and_poll(
@@ -34,6 +44,30 @@ class CozeChat:
                 return message.content
         return None
 
+    def setup_token_refresh(self):
+        thread = threading.Thread(target=self.refresh_token_loop)
+        thread.start()
+
+    def refresh_token_loop(self):
+        while True:
+            time.sleep(11*3600)
+            if self.auth_app:
+                self.auth_app.get_access_token(ttl=12*3600)
+
+    @staticmethod
+    def get_oauth_app(client_id, private_key_path, public_key_id, base_url=None) -> JWTOAuthApp:
+        if not base_url:
+            base_url = COZE_CN_BASE_URL
+        with open(private_key_path, "r") as f:
+            private_key = f.read()
+        jwt_oauth_app = JWTOAuthApp(
+            client_id=str(client_id),
+            private_key=private_key,
+            public_key_id=public_key_id,
+            base_url=base_url,
+        )
+        return jwt_oauth_app
+
 if __name__ == '__main__':
     # Init the Coze client through the access_token.
     coze = Coze(auth=TokenAuth(token=COZE_API_TOKEN), base_url=COZE_CN_BASE_URL)

configs/dev.yaml

@@ -25,10 +25,17 @@ storage:
   staff:
     table: qywx_employee
 
+chat_api:
+  coze:
+    oauth_client_id: 1194838245616
+    public_key_id: xafitzyxY0OBCFJFzmhBxauo8LKe2pe2YjlTNYfEsns
+    private_key_path: oauth/coze_privkey.pem
+
 debug_flags:
-  disable_llm_api_call: True
+  disable_llm_api_call: False
   use_local_user_storage: False
   console_input: False
+  disable_active_conversation: True
 
 use_aliyun_mq: True
 mq: