jarvis-ai-assistant 0.1.45__py3-none-any.whl → 0.1.46__py3-none-any.whl

jarvis/models/ai8.py

@@ -6,12 +6,13 @@ import requests
 import json
 import base64
 
+
 class AI8Model(BasePlatform):
     """AI8 model implementation"""
-    
+
     platform_name = "ai8"
     BASE_URL = "https://ai8.rcouyi.com"
-    
+
     def __init__(self):
         """Initialize model"""
         self.system_message = ""
@@ -21,23 +22,23 @@ class AI8Model(BasePlatform):
 
         # 获取可用模型列表
         available_models = self.get_available_models()
-        
+
         if available_models:
             PrettyOutput.section("支持的模型", OutputType.SUCCESS)
             for model in self.models.values():
                 # 格式化显示模型信息
                 model_str = f"{model['value']:<30}"
-                
+
                 # 添加标签
                 model_str += f"{model['label']}"
-                
+
                 # 添加标签和积分信息
                 attrs = []
                 if model['attr'].get('tag'):
                     attrs.append(model['attr']['tag'])
                 if model['attr'].get('integral'):
                     attrs.append(model['attr']['integral'])
-                    
+
                 # 添加特性标记
                 features = []
                 if model['attr'].get('multimodal'):
@@ -48,11 +49,11 @@ class AI8Model(BasePlatform):
                     features.append("图像支持")
                 if features:
                     model_str += f" [{'|'.join(features)}]"
-                    
+
                 # 添加备注
                 if model['attr'].get('note'):
                     model_str += f" - {model['attr']['note']}"
-                    
+
                 PrettyOutput.print(model_str, OutputType.INFO)
         else:
             PrettyOutput.print("获取模型列表失败", OutputType.WARNING)
@@ -60,19 +61,22 @@ class AI8Model(BasePlatform):
         self.token = os.getenv("AI8_API_KEY")
         if not self.token:
             raise Exception("AI8_API_KEY is not set")
-        
+
         PrettyOutput.print("使用AI8_MODEL环境变量配置模型", OutputType.SUCCESS)
-        
+
         self.model_name = os.getenv("AI8_MODEL") or "deepseek-chat"
         if self.model_name not in self.models:
-            PrettyOutput.print(f"警告: 当前选择的模型 {self.model_name} 不在可用列表中", OutputType.WARNING)
-        
+            PrettyOutput.print(
+                f"警告: 当前选择的模型 {
+                    self.model_name} 不在可用列表中",
+                OutputType.WARNING)
+
         PrettyOutput.print(f"当前使用模型: {self.model_name}", OutputType.SYSTEM)
 
     def set_model_name(self, model_name: str):
         """设置模型名称"""
         self.model_name = model_name
-            
+
     def create_conversation(self) -> bool:
         """Create a new conversation"""
         try:
@@ -85,25 +89,36 @@ class AI8Model(BasePlatform):
                 'Origin': self.BASE_URL,
                 'Referer': f'{self.BASE_URL}/chat?_userMenuKey=chat'
             }
-            
+
             # 1. 创建会话
             response = requests.post(
                 f"{self.BASE_URL}/api/chat/session",
                 headers=headers
             )
-            
+
             if response.status_code != 200:
-                PrettyOutput.print(f"创建会话失败: {response.status_code}", OutputType.ERROR)
+                PrettyOutput.print(
+                    f"创建会话失败: {
+                        response.status_code}",
+                    OutputType.ERROR)
                 return False
-            
+
             data = response.json()
             if data['code'] != 0:
-                PrettyOutput.print(f"创建会话失败: {data.get('msg', '未知错误')}", OutputType.ERROR)
+                PrettyOutput.print(
+                    f"创建会话失败: {
+                        data.get(
+                            'msg',
+                            '未知错误')}",
+                    OutputType.ERROR)
                 return False
-            
+
             self.conversation = data['data']
-            PrettyOutput.print(f"创建会话成功: {data['data']['id']}", OutputType.SUCCESS)
-            
+            PrettyOutput.print(
+                f"创建会话成功: {
+                    data['data']['id']}",
+                OutputType.SUCCESS)
+
             # 2. 更新会话设置
             session_data = {
                 **self.conversation,
@@ -114,13 +129,13 @@ class AI8Model(BasePlatform):
                 "localPlugins": None,
                 "useAppId": 0
             }
-            
+
             response = requests.put(
                 f"{self.BASE_URL}/api/chat/session/{self.conversation['id']}",
                 headers=headers,
                 json=session_data
             )
-            
+
             if response.status_code == 200:
                 data = response.json()
                 if data['code'] == 0:
@@ -128,16 +143,24 @@ class AI8Model(BasePlatform):
                     PrettyOutput.print("会话设置更新成功", OutputType.SUCCESS)
                     return True
                 else:
-                    PrettyOutput.print(f"更新会话设置失败: {data.get('msg', '未知错误')}", OutputType.ERROR)
+                    PrettyOutput.print(
+                        f"更新会话设置失败: {
+                            data.get(
+                                'msg',
+                                '未知错误')}",
+                        OutputType.ERROR)
                     return False
             else:
-                PrettyOutput.print(f"更新会话设置失败: {response.status_code}", OutputType.ERROR)
+                PrettyOutput.print(
+                    f"更新会话设置失败: {
+                        response.status_code}",
+                    OutputType.ERROR)
                 return False
-            
+
         except Exception as e:
             PrettyOutput.print(f"创建会话异常: {str(e)}", OutputType.ERROR)
             return False
-        
+
     def upload_files(self, file_list: List[str]) -> List[Dict]:
         for file_path in file_list:
             name = os.path.basename(file_path)
@@ -149,17 +172,17 @@ class AI8Model(BasePlatform):
                 "data": f"data:image/png;base64,{base64_data}"
             })
             PrettyOutput.print(f"文件 {name} 已准备好发送", OutputType.SUCCESS)
-    
+
     def set_system_message(self, message: str):
         """Set system message"""
         self.system_message = message
-        
+
     def chat(self, message: str) -> str:
         """Execute chat with the model
-        
+
         Args:
             message: User input message
-            
+
         Returns:
             str: Model response
         """
@@ -168,7 +191,7 @@ class AI8Model(BasePlatform):
             if not self.conversation:
                 if not self.create_conversation():
                     raise Exception("Failed to create conversation")
-                
+
             headers = {
                 'Authorization': self.token,
                 'Content-Type': 'application/json',
@@ -178,13 +201,13 @@ class AI8Model(BasePlatform):
                 'Origin': self.BASE_URL,
                 'Referer': f'{self.BASE_URL}/chat?_userMenuKey=chat'
             }
-            
+
             payload = {
                 "text": message,
                 "sessionId": self.conversation['id'],
                 "files": []
             }
-            
+
             # 如果有文件需要发送
             if self.files:
                 for file_data in self.files:
@@ -193,19 +216,19 @@ class AI8Model(BasePlatform):
                         "data": file_data["data"]
                     })
                 self.files = []  # 清空已使用的文件
-            
+
             response = requests.post(
                 f"{self.BASE_URL}/api/chat/completions",
                 headers=headers,
                 json=payload,
                 stream=True
             )
-            
+
             if response.status_code != 200:
                 error_msg = f"聊天请求失败: {response.status_code}"
                 PrettyOutput.print(error_msg, OutputType.ERROR)
                 raise Exception(error_msg)
-            
+
             # 处理流式响应
             full_response = ""
             for line in response.iter_lines():
@@ -222,30 +245,30 @@ class AI8Model(BasePlatform):
 
                         except json.JSONDecodeError:
                             continue
-            
+
             PrettyOutput.print_stream_end()
 
             return full_response
-            
+
         except Exception as e:
             PrettyOutput.print(f"聊天异常: {str(e)}", OutputType.ERROR)
             raise e
-            
+
     def name(self) -> str:
         """Return model name"""
         return self.model_name
-        
+
     def reset(self):
         """Reset model state"""
         self.conversation = None
         self.files = []  # 清空文件列表
-            
+
     def delete_chat(self) -> bool:
         """Delete current chat session"""
         try:
             if not self.conversation:
                 return True
-            
+
             headers = {
                 'Authorization': self.token,
                 'Content-Type': 'application/json',
@@ -255,12 +278,12 @@ class AI8Model(BasePlatform):
                 'Origin': self.BASE_URL,
                 'Referer': f'{self.BASE_URL}/chat?_userMenuKey=chat'
             }
-            
+
             response = requests.delete(
                 f"{self.BASE_URL}/api/chat/session/{self.conversation['id']}",
                 headers=headers
             )
-            
+
             if response.status_code == 200:
                 data = response.json()
                 if data['code'] == 0:
@@ -275,14 +298,14 @@ class AI8Model(BasePlatform):
                 error_msg = f"删除会话请求失败: {response.status_code}"
                 PrettyOutput.print(error_msg, OutputType.ERROR)
                 return False
-            
+
         except Exception as e:
             PrettyOutput.print(f"删除会话异常: {str(e)}", OutputType.ERROR)
             return False
-        
+
     def get_available_models(self) -> List[str]:
         """获取可用的模型列表
-        
+
         Returns:
             List[str]: 可用模型名称列表
         """
@@ -295,30 +318,37 @@ class AI8Model(BasePlatform):
                 'Origin': self.BASE_URL,
                 'Referer': f'{self.BASE_URL}/chat?_userMenuKey=chat'
             }
-            
+
             response = requests.get(
                 f"{self.BASE_URL}/api/chat/template",
                 headers=headers
             )
-            
+
             if response.status_code != 200:
-                PrettyOutput.print(f"获取模型列表失败: {response.status_code}", OutputType.ERROR)
+                PrettyOutput.print(
+                    f"获取模型列表失败: {
+                        response.status_code}",
+                    OutputType.ERROR)
                 return []
-            
+
             data = response.json()
             if data['code'] != 0:
-                PrettyOutput.print(f"获取模型列表失败: {data.get('msg', '未知错误')}", OutputType.ERROR)
+                PrettyOutput.print(
+                    f"获取模型列表失败: {
+                        data.get(
+                            'msg',
+                            '未知错误')}",
+                    OutputType.ERROR)
                 return []
-            
+
             # 保存模型信息
             self.models = {
-                model['value']: model 
+                model['value']: model
                 for model in data['data']['models']
             }
-            
+
             return list(self.models.keys())
-            
+
         except Exception as e:
             PrettyOutput.print(f"获取模型列表异常: {str(e)}", OutputType.ERROR)
             return []
-        

jarvis/models/base.py

@@ -4,7 +4,7 @@ from typing import Dict, List
 
 class BasePlatform(ABC):
     """大语言模型基类"""
-    
+
     def __init__(self):
         """初始化模型"""
         pass
@@ -12,7 +12,7 @@ class BasePlatform(ABC):
     def set_model_name(self, model_name: str):
         """设置模型名称"""
         raise NotImplementedError("set_model_name is not implemented")
-        
+
     @abstractmethod
     def chat(self, message: str) -> str:
         """执行对话"""
@@ -25,17 +25,17 @@ class BasePlatform(ABC):
     def reset(self):
         """重置模型"""
         raise NotImplementedError("reset is not implemented")
-        
+
     @abstractmethod
     def name(self) -> str:
         """模型名称"""
         raise NotImplementedError("name is not implemented")
-    
+
     @abstractmethod
-    def delete_chat(self)->bool:
+    def delete_chat(self) -> bool:
         """删除对话"""
         raise NotImplementedError("delete_chat is not implemented")
-    
+
     def set_system_message(self, message: str):
         """设置系统消息"""
         raise NotImplementedError("set_system_message is not implemented")