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

jarvis/models/ai8.py

@@ -6,13 +6,12 @@ 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 = ""
@@ -22,23 +21,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'):
@@ -49,11 +48,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)
@@ -61,22 +60,19 @@ 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:
@@ -89,36 +85,25 @@ 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,
@@ -129,13 +114,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:
@@ -143,24 +128,16 @@ 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)
@@ -172,17 +149,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
         """
@@ -191,7 +168,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',
@@ -201,13 +178,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:
@@ -216,19 +193,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():
@@ -245,30 +222,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',
@@ -278,12 +255,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:
@@ -298,14 +275,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]: 可用模型名称列表
         """
@@ -318,37 +295,30 @@ 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")