jarvis-ai-assistant 0.1.134__py3-none-any.whl → 0.1.138__py3-none-any.whl

jarvis/jarvis_platform/yuanbao.py

@@ -14,7 +14,7 @@ class YuanbaoPlatform(BasePlatform):
     def get_model_list(self) -> List[Tuple[str, str]]:
         """获取支持的模型列表"""
         return [("deep_seek", "DeepSeek-R1"), ("deep_seek_v3", "DeepSeek-v3"), ("hunyuan_gpt_175B_0404", "Tencent Hunyuan"), ("hunyuan_t1", "Tencent Hunyuan-T1")]
-    
+
     def __init__(self):
         """
         初始化Hunyuan模型
@@ -25,13 +25,13 @@ class YuanbaoPlatform(BasePlatform):
         self.cookies = os.getenv("YUANBAO_COOKIES")  # 认证cookies
         self.agent_id = os.getenv("YUANBAO_AGENT_ID")  # 代理ID
         self.web = os.getenv("YUANBAO_WEB", "false") == "true"  # 是否启用网页功能
-        
+
         if not self.cookies:
             message = (
                 "需要设置 YUANBAO_COOKIES 和 YUANBAO_AGENT_ID 才能使用 Jarvis 的元宝功能。请按照以下步骤操作：\n"
                 "1. 获取元宝 API 参数:\n"
                 "   • 访问元宝平台: https://yuanbao.tencent.com\n"
-                "   • 登录您的账户\n" 
+                "   • 登录您的账户\n"
                 "   • 打开浏览器开发者工具 (F12 或右键 -> 检查)\n"
                 "   • 切换到网络标签\n"
                 "   • 发送任意消息\n"
@@ -47,7 +47,7 @@ class YuanbaoPlatform(BasePlatform):
             )
             PrettyOutput.print(message, OutputType.INFO)
             PrettyOutput.print("YUANBAO_COOKIES 未设置", OutputType.WARNING)
-        
+
         self.system_message = ""  # 系统消息，用于初始化对话
         self.first_chat = True  # 标识是否为第一次对话
         self.model_name = "deep_seek_v3"  # 默认模型名称，使用下划线保持一致
@@ -59,18 +59,17 @@ class YuanbaoPlatform(BasePlatform):
     def set_model_name(self, model_name: str):
         # 模型映射表，可以根据需要扩展
         model_mapping = [m[0] for m in self.get_model_list()]
-        
+
         if model_name in model_mapping:
             self.model_name = model_name
         else:
             PrettyOutput.print(f"错误：不支持的模型: {model_name}", OutputType.ERROR)
-        self.reset()
 
     def _get_base_headers(self):
         """Get base headers for API requests"""
         return {
             'Host': 'yuanbao.tencent.com',
-            'X-Language': 'zh-CN', 
+            'X-Language': 'zh-CN',
             'X-Requested-With': 'XMLHttpRequest',
             'chat_version': 'v1',
             'X-Instance-ID': '5',
@@ -94,17 +93,17 @@ class YuanbaoPlatform(BasePlatform):
     def _create_conversation(self) -> bool:
         """Create a new conversation session"""
         url = "https://yuanbao.tencent.com/api/user/agent/conversation/create"
-        
+
         headers = self._get_base_headers()
-        
+
         payload = json.dumps({
             "agentId": self.agent_id
         })
-        
+
         try:
             response = while_success(lambda: requests.post(url, headers=headers, data=payload), sleep_time=5)
             response_json = response.json()
-            
+
             if "id" in response_json:
                 self.conversation_id = response_json["id"]
                 return True
@@ -120,11 +119,11 @@ class YuanbaoPlatform(BasePlatform):
         if not self.conversation_id:
             if not self._create_conversation():
                 raise Exception("Failed to create conversation session")
-        
+
         url = f"https://yuanbao.tencent.com/api/chat/{self.conversation_id}"
-        
+
         headers = self._get_base_headers()
-        
+
         # 准备消息内容
         payload = {
             "model": "gpt_175B_0404",
@@ -149,44 +148,44 @@ class YuanbaoPlatform(BasePlatform):
 
         if self.web:
             payload["supportFunctions"] = ["supportInternetSearch"]
-        
-        
+
+
         # 添加系统消息（如果是第一次对话）
         if self.first_chat and self.system_message:
             payload["prompt"] = f"{self.system_message}\n\n{message}"
             payload["displayPrompt"] = payload["prompt"]
             self.first_chat = False
-        
-        try:            
+
+        try:
             # 发送消息请求，获取流式响应
             response = while_success(
                 lambda: requests.post(url, headers=headers, json=payload, stream=True),
                 sleep_time=5
             )
-            
+
             # 检查响应状态
             if response.status_code != 200:
                 error_msg = f"发送消息失败，状态码: {response.status_code}"
                 if hasattr(response, 'text'):
                     error_msg += f", 响应: {response.text}"
                 raise Exception(error_msg)
-            
+
             full_response = ""
             is_text_block = False
-            
+
             # 处理SSE流响应
             for line in response.iter_lines():
                 if not line:
                     continue
-                
+
                 line_str = line.decode('utf-8')
-                
+
                 # SSE格式的行通常以"data: "开头
                 if line_str.startswith("data: "):
                     try:
                         data_str = line_str[6:]  # 移除"data: "前缀
                         data = json.loads(data_str)
-                        
+
                         # 处理文本类型的消息
                         if data.get("type") == "text":
                             is_text_block = True
@@ -195,60 +194,57 @@ class YuanbaoPlatform(BasePlatform):
                                 if not self.suppress_output:
                                     PrettyOutput.print_stream(msg)
                                 full_response += msg
-                        
+
                         # 处理思考中的消息（可选展示）
                         elif data.get("type") == "think" and not self.suppress_output:
                             think_content = data.get("content", "")
                             # 可以选择性地显示思考过程，但不加入最终响应
                             PrettyOutput.print_stream(f"{think_content}", is_thinking=True)
                             pass
-                            
+
                     except json.JSONDecodeError:
                         pass
-                
+
                 # 检测结束标志
                 elif line_str == "data: [DONE]":
                     break
-            
+
             if not self.suppress_output:
                 PrettyOutput.print_stream_end()
-                
+
             return full_response
-            
+
         except Exception as e:
             raise Exception(f"对话失败: {str(e)}")
 
-    def reset(self):
-        """Reset chat"""
-        self.conversation_id = ""
-        self.first_chat = True
 
     def delete_chat(self) -> bool:
         """Delete current session"""
         if not self.conversation_id:
             return True  # 如果没有会话ID，视为删除成功
-            
+
         # Hunyuan使用专门的clear API来清除会话
         url = "https://yuanbao.tencent.com/api/user/agent/conversation/v1/clear"
-        
+
         # 为这个请求获取基础头部
         headers = self._get_base_headers()
-        
+
         # 更新X-AgentID头部，需要包含会话ID
         headers.update({
             'X-AgentID': f"{self.agent_id}/{self.conversation_id}"
         })
-        
+
         # 创建请求体，包含要删除的会话ID
         payload = {
             "conversationIds": [self.conversation_id]
         }
-        
+
         try:
             response = while_success(lambda: requests.post(url, headers=headers, json=payload), sleep_time=5)
-            
+
             if response.status_code == 200:
-                self.reset()
+                self.conversation_id = ""
+                self.first_chat = True
                 return True
             else:
                 PrettyOutput.print(f"删除会话失败: HTTP {response.status_code}", OutputType.WARNING)
@@ -261,4 +257,4 @@ class YuanbaoPlatform(BasePlatform):
 
     def name(self) -> str:
         """Model name"""
-        return "yuanbao"
+        return self.model_name