jarvis-ai-assistant 0.1.32__py3-none-any.whl

jarvis/models/kimi.py

@@ -0,0 +1,389 @@
+from typing import Dict, List
+import requests
+import json
+import os
+import mimetypes
+import time
+from jarvis.models.base import BaseModel
+from jarvis.utils import PrettyOutput, OutputType
+from jarvis.utils import while_success
+
+class KimiModel(BaseModel):
+    """Kimi模型实现"""
+
+    model_name = "kimi"
+    
+    def __init__(self):
+        """
+        初始化Kimi模型
+        """
+        self.api_key = os.getenv("KIMI_API_KEY")
+        if not self.api_key:
+            PrettyOutput.info("\n需要设置 KIMI_API_KEY 才能使用 Jarvis。请按以下步骤操作：")
+            PrettyOutput.info("\n1. 获取 Kimi API Key:")
+            PrettyOutput.info("   • 访问 Kimi AI 平台: https://kimi.moonshot.cn")
+            PrettyOutput.info("   • 登录您的账号")
+            PrettyOutput.info("   • 打开浏览器开发者工具 (F12 或右键 -> 检查)")
+            PrettyOutput.info("   • 切换到 Network 标签页")
+            PrettyOutput.info("   • 发送任意消息")
+            PrettyOutput.info("   • 在请求中找到 Authorization 头部")
+            PrettyOutput.info("   • 复制 token 值（去掉 'Bearer ' 前缀）")
+            PrettyOutput.info("\n2. 设置环境变量:")
+            PrettyOutput.info("   方法 1: 创建或编辑 ~/.jarvis_env 文件:")
+            PrettyOutput.info("   echo 'KIMI_API_KEY=your_key_here' > ~/.jarvis_env")
+            PrettyOutput.info("\n   方法 2: 直接设置环境变量:")
+            PrettyOutput.info("   export KIMI_API_KEY=your_key_here")
+            PrettyOutput.info("\n设置完成后重新运行 Jarvis。")
+            raise Exception("KIMI_API_KEY is not set")
+        self.auth_header = f"Bearer {self.api_key}"
+        self.chat_id = ""
+        self.uploaded_files = []  # 存储已上传文件的信息
+        self.first_chat = True  # 添加标记，用于判断是否是第一次对话
+        self.system_message = ""
+
+    def set_system_message(self, message: str):
+        """设置系统消息"""
+        self.system_message = message
+
+    def _create_chat(self) -> bool:
+        """创建新的对话会话"""
+        url = "https://kimi.moonshot.cn/api/chat"
+        payload = json.dumps({
+            "name": "未命名会话",
+            "is_example": False,
+            "kimiplus_id": "kimi"
+        })
+        headers = {
+            'Authorization': self.auth_header,
+            'Content-Type': 'application/json'
+        }
+        try:
+            response = while_success(lambda: requests.request("POST", url, headers=headers, data=payload), sleep_time=5)
+            self.chat_id = response.json()["id"]
+            return True
+        except Exception as e:
+            PrettyOutput.print(f"Error: Failed to create chat: {e}", OutputType.ERROR)
+            return False
+
+    def _get_presigned_url(self, filename: str, action: str) -> Dict:
+        """获取预签名上传URL"""
+        url = "https://kimi.moonshot.cn/api/pre-sign-url"
+        
+        
+        
+        payload = json.dumps({
+            "action": action,
+            "name": os.path.basename(filename)
+        }, ensure_ascii=False)
+        
+        headers = {
+            'Authorization': self.auth_header,
+            'Content-Type': 'application/json'
+        }
+        
+        response = while_success(lambda: requests.post(url, headers=headers, data=payload), sleep_time=5)
+        return response.json()
+
+    def _upload_file(self, file_path: str, presigned_url: str) -> bool:
+        """上传文件到预签名URL"""
+        try:
+            with open(file_path, 'rb') as f:
+                content = f.read()
+                response = while_success(lambda: requests.put(presigned_url, data=content), sleep_time=5)
+                return response.status_code == 200
+        except Exception as e:
+            PrettyOutput.print(f"Error: Failed to upload file: {e}", OutputType.ERROR)
+            return False
+
+    def _get_file_info(self, file_data: Dict, name: str, file_type: str) -> Dict:
+        """获取文件信息"""
+        url = "https://kimi.moonshot.cn/api/file"
+        payload = json.dumps({
+            "type": file_type,
+            "name": name,
+            "object_name": file_data["object_name"],
+            "chat_id": self.chat_id,
+            "file_id": file_data.get("file_id", "")
+        }, ensure_ascii=False)
+        
+        headers = {
+            'Authorization': self.auth_header,
+            'Content-Type': 'application/json'
+        }
+        
+        response = while_success(lambda: requests.post(url, headers=headers, data=payload), sleep_time=5)
+        return response.json()
+
+    def _wait_for_parse(self, file_id: str) -> bool:
+        """等待文件解析完成"""
+        url = "https://kimi.moonshot.cn/api/file/parse_process"
+        headers = {
+            'Authorization': self.auth_header,
+            'Content-Type': 'application/json'
+        }
+        
+        max_retries = 30
+        retry_count = 0
+        
+        while retry_count < max_retries:
+            payload = json.dumps({"ids": [file_id]}, ensure_ascii=False)
+            response = while_success(lambda: requests.post(url, headers=headers, data=payload, stream=True), sleep_time=5)
+            
+            for line in response.iter_lines():
+                if not line:
+                    continue
+                    
+                line = line.decode('utf-8')
+                if not line.startswith("data: "):
+                    continue
+                    
+                try:
+                    data = json.loads(line[6:])
+                    if data.get("event") == "resp":
+                        status = data.get("file_info", {}).get("status")
+                        if status == "parsed":
+                            return True
+                        elif status == "failed":
+                            return False
+                except json.JSONDecodeError:
+                    continue
+            
+            retry_count += 1
+            time.sleep(1)
+        
+        return False
+    def upload_files(self, file_list: List[str]) -> List[Dict]:
+        """上传文件列表并返回文件信息"""
+        if not file_list:
+            return []
+
+        PrettyOutput.print("Progress: 开始处理文件上传...", OutputType.PROGRESS)
+        
+        if not self.chat_id:
+            PrettyOutput.print("创建新的对话会话...", OutputType.PROGRESS)
+            if not self._create_chat():
+                raise Exception("Failed to create chat session")
+
+        uploaded_files = []
+        for index, file_path in enumerate(file_list, 1):
+            try:
+                PrettyOutput.print(f"处理文件 [{index}/{len(file_list)}]: {file_path}", OutputType.PROGRESS)
+
+                mime_type, _ = mimetypes.guess_type(file_path)
+                action = "image" if mime_type and mime_type.startswith('image/') else "file"
+                
+                # 获取预签名URL
+                PrettyOutput.print("获取上传URL...", OutputType.PROGRESS)
+                presigned_data = self._get_presigned_url(file_path, action)
+                
+                # 上传文件
+                PrettyOutput.print("上传文件内容...", OutputType.PROGRESS)
+                if self._upload_file(file_path, presigned_data["url"]):
+                    # 获取文件信息
+                    PrettyOutput.print("获取文件信息...", OutputType.PROGRESS)
+                    file_info = self._get_file_info(presigned_data, os.path.basename(file_path), action)
+                    # 等待文件解析
+                    PrettyOutput.print("等待文件解析完成...", OutputType.PROGRESS)
+
+                    # 只有文件需要解析
+                    if action == "file":
+                        if self._wait_for_parse(file_info["id"]):
+                            uploaded_files.append(file_info)
+                            PrettyOutput.print(f"Success: 文件处理成功: {file_path}", OutputType.SUCCESS)
+                        else:
+                            PrettyOutput.print(f"✗ 文件解析失败: {file_path}", OutputType.ERROR)
+                    else:
+                        uploaded_files.append(file_info)
+                        PrettyOutput.print(f"Success: 文件处理成功: {file_path}", OutputType.SUCCESS)
+                else:
+                    PrettyOutput.print(f"Error: 文件上传失败: {file_path}", OutputType.ERROR)
+                    
+            except Exception as e:
+                PrettyOutput.print(f"✗ 处理文件出错 {file_path}: {str(e)}", OutputType.ERROR)
+                continue
+        
+        if uploaded_files:
+            PrettyOutput.print(f"成功处理 {len(uploaded_files)}/{len(file_list)} 个文件", OutputType.SUCCESS)
+        else:
+            PrettyOutput.print("没有文件成功处理", OutputType.ERROR)
+        
+        self.uploaded_files = uploaded_files
+        return uploaded_files
+
+    def chat(self, message: str) -> str:
+        """发送消息并获取响应"""
+        if not self.chat_id:
+            PrettyOutput.print("创建新的对话会话...", OutputType.PROGRESS)
+            if not self._create_chat():
+                raise Exception("Failed to create chat session")
+
+        url = f"https://kimi.moonshot.cn/api/chat/{self.chat_id}/completion/stream"
+        
+        # 只在第一次对话时带上文件引用
+        refs = []
+        refs_file = []
+        if self.first_chat:
+            if self.uploaded_files:
+                PrettyOutput.print(f"首次对话，引用 {len(self.uploaded_files)} 个文件...", OutputType.PROGRESS)
+                refs = [f["id"] for f in self.uploaded_files]
+                refs_file = self.uploaded_files
+            message = self.system_message + "\n" + message
+            self.first_chat = False
+        
+        PrettyOutput.print("发送请求...", OutputType.PROGRESS)
+        payload = {
+            "messages": [{"role": "user", "content": message}],
+            "use_search": True,
+            "extend": {"sidebar": True},
+            "kimiplus_id": "kimi",
+            "use_research": False,
+            "use_math": False,
+            "refs": refs,
+            "refs_file": refs_file
+        }
+
+        headers = {
+            'Authorization': self.auth_header,
+            'Content-Type': 'application/json'
+        }
+
+        try:
+            response = while_success(lambda: requests.post(url, headers=headers, json=payload, stream=True), sleep_time=5)
+            full_response = ""
+            
+            # 收集搜索和引用结果
+            search_results = []
+            ref_sources = []
+            
+            PrettyOutput.print("接收响应...", OutputType.PROGRESS)
+            for line in response.iter_lines():
+                if not line:
+                    continue
+                    
+                line = line.decode('utf-8')
+                if not line.startswith("data: "):
+                    continue
+                    
+                try:
+                    data = json.loads(line[6:])
+                    event = data.get("event")
+                    
+                    if event == "cmpl":
+                        # 处理补全文本
+                        text = data.get("text", "")
+                        if text:
+                            PrettyOutput.print_stream(text)
+                            full_response += text
+                            
+                    elif event == "search_plus":
+                        # 收集搜索结果
+                        msg = data.get("msg", {})
+                        if msg.get("type") == "get_res":
+                            search_results.append({
+                                "date": msg.get("date", ""),
+                                "site_name": msg.get("site_name", ""),
+                                "snippet": msg.get("snippet", ""),
+                                "title": msg.get("title", ""),
+                                "type": msg.get("type", ""),
+                                "url": msg.get("url", "")
+                            })
+                                
+                    elif event == "ref_docs":
+                        # 收集引用来源
+                        ref_cards = data.get("ref_cards", [])
+                        for card in ref_cards:
+                            ref_sources.append({
+                                "idx_s": card.get("idx_s", ""),
+                                "idx_z": card.get("idx_z", ""),
+                                "ref_id": card.get("ref_id", ""),
+                                "url": card.get("url", ""),
+                                "title": card.get("title", ""),
+                                "abstract": card.get("abstract", ""),
+                                "source": card.get("source_label", ""),
+                                "rag_segments": card.get("rag_segments", []),
+                                "origin": card.get("origin", {})
+                            })
+                                    
+                except json.JSONDecodeError:
+                    continue
+                    
+            PrettyOutput.print_stream_end()
+            
+
+            # 显示搜索结果摘要
+            if search_results:
+                PrettyOutput.print("\n搜索结果:", OutputType.PROGRESS)
+                for result in search_results:
+                    PrettyOutput.print(f"- {result['title']}", OutputType.PROGRESS)
+                    if result['date']:
+                        PrettyOutput.print(f"  日期: {result['date']}", OutputType.PROGRESS)
+                    PrettyOutput.print(f"  来源: {result['site_name']}", OutputType.PROGRESS)
+                    if result['snippet']:
+                        PrettyOutput.print(f"  摘要: {result['snippet']}", OutputType.PROGRESS)
+                    PrettyOutput.print(f"  链接: {result['url']}", OutputType.PROGRESS)
+                    PrettyOutput.print("", OutputType.PROGRESS)
+                        
+            # 显示引用来源
+            if ref_sources:
+                PrettyOutput.print("\n引用来源:", OutputType.PROGRESS)
+                for source in ref_sources:
+                    PrettyOutput.print(f"- [{source['ref_id']}] {source['title']} ({source['source']})", OutputType.PROGRESS)
+                    PrettyOutput.print(f"  链接: {source['url']}", OutputType.PROGRESS)
+                    if source['abstract']:
+                        PrettyOutput.print(f"  摘要: {source['abstract']}", OutputType.PROGRESS)
+                    
+                    # 显示相关段落
+                    if source['rag_segments']:
+                        PrettyOutput.print("  相关段落:", OutputType.PROGRESS)
+                        for segment in source['rag_segments']:
+                            text = segment.get('text', '').replace('\n', ' ').strip()
+                            if text:
+                                PrettyOutput.print(f"    - {text}", OutputType.PROGRESS)
+                    
+                    # 显示原文引用
+                    origin = source['origin']
+                    if origin:
+                        text = origin.get('text', '')
+                        if text:
+                            PrettyOutput.print(f"  原文: {text}", OutputType.PROGRESS)
+                    
+                    PrettyOutput.print("", OutputType.PROGRESS)
+            
+            return full_response
+
+        except Exception as e:
+            raise Exception(f"Chat failed: {str(e)}")
+
+    def delete_chat(self) -> bool:
+        """删除当前会话"""
+        if not self.chat_id:
+            return True  # 如果没有会话ID，视为删除成功
+            
+        url = f"https://kimi.moonshot.cn/api/chat/{self.chat_id}"
+        headers = {
+            'Authorization': self.auth_header,
+            'Content-Type': 'application/json'
+        }
+        
+        try:
+            response = while_success(lambda: requests.delete(url, headers=headers), sleep_time=5)
+            if response.status_code == 200:
+                PrettyOutput.print("会话已删除", OutputType.SUCCESS)
+                self.reset()
+                return True
+            else:
+                PrettyOutput.print(f"删除会话失败: HTTP {response.status_code}", OutputType.ERROR)
+                return False
+        except Exception as e:
+            PrettyOutput.print(f"删除会话时发生错误: {str(e)}", OutputType.ERROR)
+            return False
+
+    def reset(self):
+        """重置对话"""
+        self.chat_id = ""
+        self.uploaded_files = []
+        self.first_chat = True  # 重置first_chat标记
+
+    def name(self) -> str:
+        """模型名称"""
+        return "kimi"

jarvis/models/openai.py

@@ -0,0 +1,96 @@
+from typing import Dict, List
+import os
+from openai import OpenAI
+from jarvis.models.base import BaseModel
+from jarvis.utils import PrettyOutput, OutputType
+
+class OpenAIModel(BaseModel):
+    """DeepSeek模型实现"""
+
+    model_name = "openai"
+    
+    def __init__(self):
+        """
+        初始化DeepSeek模型
+        """
+        self.api_key = os.getenv("OPENAI_API_KEY")
+        if not self.api_key:
+            PrettyOutput.print("\n需要设置以下环境变量才能使用 OpenAI 模型：", OutputType.INFO)
+            PrettyOutput.print("  • OPENAI_API_KEY: API 密钥", OutputType.INFO)
+            PrettyOutput.print("  • OPENAI_API_BASE: (可选) API 基础地址，默认使用 https://api.deepseek.com", OutputType.INFO)
+            PrettyOutput.print("\n可以通过以下方式设置：", OutputType.INFO)
+            PrettyOutput.print("1. 创建或编辑 ~/.jarvis_env 文件:", OutputType.INFO)
+            PrettyOutput.print("   OPENAI_API_KEY=your_api_key", OutputType.INFO)
+            PrettyOutput.print("   OPENAI_API_BASE=your_api_base", OutputType.INFO)
+            PrettyOutput.print("   OPENAI_MODEL_NAME=your_model_name", OutputType.INFO)
+            PrettyOutput.print("\n2. 或者直接设置环境变量:", OutputType.INFO)
+            PrettyOutput.print("   export OPENAI_API_KEY=your_api_key", OutputType.INFO)
+            PrettyOutput.print("   export OPENAI_API_BASE=your_api_base", OutputType.INFO)
+            PrettyOutput.print("   export OPENAI_MODEL_NAME=your_model_name", OutputType.INFO)
+            raise Exception("OPENAI_API_KEY is not set")
+            
+        self.base_url = os.getenv("OPENAI_API_BASE", "https://api.deepseek.com")
+        self.model_name = os.getenv("OPENAI_MODEL_NAME", "deepseek-chat")
+            
+        self.client = OpenAI(
+            api_key=self.api_key,
+            base_url=self.base_url
+        )
+        self.messages: List[Dict[str, str]] = []
+        self.system_message = ""
+
+    def set_system_message(self, message: str):
+        """设置系统消息"""
+        self.system_message = message
+        self.messages.append({"role": "system", "content": self.system_message})
+
+    def chat(self, message: str) -> str:
+        """执行对话"""
+        try:
+            PrettyOutput.print("发送请求...", OutputType.PROGRESS)
+            
+            # 添加用户消息到历史记录
+            self.messages.append({"role": "user", "content": message})
+            
+            response = self.client.chat.completions.create(
+                model=self.model_name,  # 使用配置的模型名称
+                messages=self.messages,
+                stream=True
+            )
+            
+            PrettyOutput.print("接收响应...", OutputType.PROGRESS)
+            full_response = ""
+            
+            for chunk in response:
+                if chunk.choices[0].delta.content:
+                    text = chunk.choices[0].delta.content
+                    PrettyOutput.print_stream(text)
+                    full_response += text
+                    
+            PrettyOutput.print_stream_end()
+            
+            # 添加助手回复到历史记录
+            self.messages.append({"role": "assistant", "content": full_response})
+            
+            return full_response
+            
+        except Exception as e:
+            PrettyOutput.print(f"对话失败: {str(e)}", OutputType.ERROR)
+            raise Exception(f"Chat failed: {str(e)}")
+
+    def name(self) -> str:
+        """返回模型名称"""
+        return self.model_name
+
+    def reset(self):
+        """重置模型状态"""
+        # 清空对话历史，只保留system message
+        if self.system_message:
+            self.messages = [{"role": "system", "content": self.system_message}]
+        else:
+            self.messages = []
+
+    def delete_chat(self)->bool:
+        """删除对话"""
+        self.reset()
+        return True