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

jarvis/tools/file_ops.py

@@ -0,0 +1,106 @@
+from typing import Dict, Any
+import os
+from pathlib import Path
+from ..utils import PrettyOutput, OutputType
+
+class FileOperationTool:
+    name = "file_operation"
+    description = "Execute file operations (read/write/append/exists)"
+    parameters = {
+        "type": "object",
+        "properties": {
+            "operation": {
+                "type": "string",
+                "enum": ["read", "write", "append", "exists"],
+                "description": "Type of file operation to perform"
+            },
+            "filepath": {
+                "type": "string",
+                "description": "Absolute or relative path to the file"
+            },
+            "content": {
+                "type": "string",
+                "description": "Content to write (required for write/append operations)",
+                "default": ""
+            },
+            "encoding": {
+                "type": "string",
+                "description": "File encoding (default: utf-8)",
+                "default": "utf-8"
+            }
+        },
+        "required": ["operation", "filepath"]
+    }
+
+    def execute(self, args: Dict) -> Dict[str, Any]:
+        """执行文件操作"""
+        try:
+            operation = args["operation"]
+            filepath = args["filepath"]
+            encoding = args.get("encoding", "utf-8")
+            
+            # 记录操作和完整路径
+            abs_path = os.path.abspath(filepath)
+            PrettyOutput.print(f"文件操作: {operation} - {abs_path}", OutputType.INFO)
+            
+            if operation == "exists":
+                exists = os.path.exists(filepath)
+                return {
+                    "success": True,
+                    "stdout": str(exists),
+                    "stderr": ""
+                }
+                
+            elif operation == "read":
+                if not os.path.exists(filepath):
+                    return {
+                        "success": False,
+                        "error": f"文件不存在: {filepath}"
+                    }
+                    
+                # 检查文件大小
+                if os.path.getsize(filepath) > 10 * 1024 * 1024:  # 10MB
+                    return {
+                        "success": False,
+                        "error": "文件过大 (>10MB)"
+                    }
+                    
+                with open(filepath, 'r', encoding=encoding) as f:
+                    content = f.read()
+                return {
+                    "success": True,
+                    "stdout": content,
+                    "stderr": ""
+                }
+                
+            elif operation in ["write", "append"]:
+                if not args.get("content"):
+                    return {
+                        "success": False,
+                        "error": "写入/追加操作需要提供content参数"
+                    }
+                
+                # 创建目录（如果不存在）
+                os.makedirs(os.path.dirname(os.path.abspath(filepath)), exist_ok=True)
+                    
+                mode = 'a' if operation == "append" else 'w'
+                with open(filepath, mode, encoding=encoding) as f:
+                    f.write(args["content"])
+                    
+                return {
+                    "success": True,
+                    "stdout": f"成功{operation}内容到 {filepath}",
+                    "stderr": ""
+                }
+                
+            else:
+                return {
+                    "success": False,
+                    "error": f"未知操作: {operation}"
+                }
+                
+        except Exception as e:
+            return {
+                "success": False,
+                "error": f"文件操作失败: {str(e)}"
+            } 

jarvis/tools/python_script.py

@@ -0,0 +1,150 @@
+import os
+import time
+import subprocess
+import sys
+import pkgutil
+from pathlib import Path
+from typing import Dict, Any, Optional
+from ..utils import PrettyOutput, OutputType
+
+class PythonScript:
+    """Python脚本管理类"""
+    SCRIPTS_DIR = "/tmp/ai_scripts"
+
+    @classmethod
+    def init_scripts_dir(cls):
+        """初始化脚本目录"""
+        Path(cls.SCRIPTS_DIR).mkdir(parents=True, exist_ok=True)
+
+    @classmethod
+    def generate_script_path(cls, name: Optional[str] = None) -> str:
+        """生成脚本文件路径"""
+        if name:
+            safe_name = "".join(c for c in name if c.isalnum() or c in "._- ")
+            filename = f"{int(time.time())}_{safe_name}.py"
+        else:
+            filename = f"{int(time.time())}_script.py"
+        return str(Path(cls.SCRIPTS_DIR) / filename)
+
+class PythonScriptTool:
+    name = "execute_python"
+    description = """Execute Python code and return the results.
+    Notes:
+    1. Use print() to output results
+    2. Automatic dependency management
+    3. Code saved to temporary file
+    """
+    parameters = {
+        "type": "object",
+        "properties": {
+            "code": {
+                "type": "string",
+                "description": "Python code to execute"
+            },
+            "dependencies": {
+                "type": "array",
+                "items": {"type": "string"},
+                "description": "Python package dependencies",
+                "default": []
+            },
+            "name": {
+                "type": "string",
+                "description": "Script name",
+                "default": ""
+            }
+        },
+        "required": ["code"]
+    }
+
+    @staticmethod
+    def _is_builtin_package(package_name: str) -> bool:
+        package_name = package_name.split("==")[0].strip()
+        if hasattr(sys.modules, package_name) or package_name in sys.stdlib_module_names:
+            return True
+        try:
+            return pkgutil.find_spec(package_name) is not None
+        except Exception:
+            return False
+
+    def execute(self, args: Dict) -> Dict[str, Any]:
+        """执行Python代码"""
+        try:
+            code = args["code"]
+            dependencies = args.get("dependencies", [])
+            script_name = args.get("name", "")
+            
+            # 初始化脚本目录
+            PythonScript.init_scripts_dir()
+            
+            # 生成脚本路径
+            script_path = PythonScript.generate_script_path(script_name)
+            
+            # 安装依赖
+            missing_deps = []
+            for dep in dependencies:
+                if not self._is_builtin_package(dep):
+                    missing_deps.append(dep)
+            
+            if missing_deps:
+                PrettyOutput.print("正在安装依赖...", OutputType.INFO)
+                try:
+                    subprocess.run(
+                        [sys.executable, "-m", "pip", "install"] + missing_deps,
+                        check=True,
+                        capture_output=True,
+                        text=True
+                    )
+                    PrettyOutput.print("依赖安装完成", OutputType.INFO)
+                except subprocess.CalledProcessError as e:
+                    return {
+                        "success": False,
+                        "error": f"依赖安装失败: {e.stderr}"
+                    }
+            
+            # 写入脚本文件
+            with open(script_path, 'w', encoding='utf-8') as f:
+                f.write(code)
+            
+            # 执行脚本
+            PrettyOutput.print(f"执行脚本: {script_path}", OutputType.INFO)
+            result = subprocess.run(
+                [sys.executable, script_path],
+                capture_output=True,
+                text=True
+            )
+            
+            # 构建输出
+            output = []
+            
+            # 添加脚本信息
+            output.append(f"脚本路径: {script_path}")
+            if dependencies:
+                output.append(f"依赖项: {', '.join(dependencies)}")
+            output.append("")
+            
+            # 添加执行结果
+            if result.stdout:
+                output.append("输出:")
+                output.append(result.stdout)
+            
+            # 添加错误信息（如果有）
+            if result.stderr:
+                output.append("错误:")
+                output.append(result.stderr)
+            
+            # 添加返回码
+            output.append(f"返回码: {result.returncode}")
+            
+            return {
+                "success": result.returncode == 0,
+                "stdout": "\n".join(output),
+                "stderr": result.stderr,
+                "return_code": result.returncode,
+                "script_path": script_path
+            }
+            
+        except Exception as e:
+            return {
+                "success": False,
+                "error": f"执行Python代码失败: {str(e)}"
+            }

jarvis/tools/rag.py

@@ -0,0 +1,154 @@
+from typing import Dict, Any, List
+import chromadb
+from chromadb.utils import embedding_functions
+import hashlib
+import re
+from ..utils import PrettyOutput, OutputType
+
+class RAGTool:
+    name = "rag_query"
+    description = """Execute RAG queries on documents.
+    Features:
+    1. Auto-creates document embeddings
+    2. Returns relevant passages
+    3. Maintains embedding database
+    """
+    parameters = {
+        "type": "object",
+        "properties": {
+            "files": {
+                "type": "array",
+                "items": {"type": "string"},
+                "description": "Files to process"
+            },
+            "query": {
+                "type": "string",
+                "description": "Query text"
+            },
+            "num_passages": {
+                "type": "integer",
+                "description": "Number of passages",
+                "default": 3
+            },
+            "chunk_size": {
+                "type": "integer",
+                "description": "Chunk size",
+                "default": 500
+            }
+        },
+        "required": ["files", "query"]
+    }
+
+    def _get_document_hash(self, content: str) -> str:
+        return hashlib.md5(content.encode()).hexdigest()
+
+    def _chunk_text(self, text: str, chunk_size: int = 500) -> List[str]:
+        """将文本分割成适当大小的块"""
+        # 按句子分割
+        sentences = re.split(r'(?<=[.!?])\s+', text)
+        chunks = []
+        current_chunk = []
+        current_length = 0
+        
+        for sentence in sentences:
+            sentence_length = len(sentence)
+            if current_length + sentence_length > chunk_size and current_chunk:
+                chunks.append(" ".join(current_chunk))
+                current_chunk = []
+                current_length = 0
+            current_chunk.append(sentence)
+            current_length += sentence_length
+            
+        if current_chunk:
+            chunks.append(" ".join(current_chunk))
+            
+        return chunks
+
+    def execute(self, args: Dict) -> Dict[str, Any]:
+        """执行RAG查询"""
+        try:
+            files = args["files"]
+            query = args["query"]
+            num_passages = args.get("num_passages", 3)
+            chunk_size = args.get("chunk_size", 500)
+            
+            # 初始化ChromaDB
+            chroma_client = chromadb.PersistentClient(path="./data/chromadb")
+            
+            # 使用sentence-transformers作为嵌入模型
+            embedding_function = embedding_functions.SentenceTransformerEmbeddingFunction(
+                model_name="all-MiniLM-L6-v2"
+            )
+            
+            # 获取或创建集合
+            collection = chroma_client.get_or_create_collection(
+                name="document_store",
+                embedding_function=embedding_function
+            )
+            
+            # 处理每个文件
+            for file_path in files:
+                try:
+                    with open(file_path, 'r', encoding='utf-8') as f:
+                        content = f.read()
+                    
+                    # 计算文档哈希值
+                    doc_hash = self._get_document_hash(content)
+                    
+                    # 检查文档是否已经存在且未更改
+                    existing_ids = collection.get(
+                        where={"doc_hash": doc_hash}
+                    )
+                    
+                    if not existing_ids["ids"]:
+                        # 分块处理文档
+                        chunks = self._chunk_text(content, chunk_size)
+                        
+                        # 为每个块生成唯一ID
+                        chunk_ids = [f"{doc_hash}_{i}" for i in range(len(chunks))]
+                        
+                        # 添加到数据库
+                        collection.add(
+                            documents=chunks,
+                            ids=chunk_ids,
+                            metadatas=[{
+                                "file_path": file_path,
+                                "doc_hash": doc_hash,
+                                "chunk_index": i
+                            } for i in range(len(chunks))]
+                        )
+                        
+                        PrettyOutput.print(f"已添加文档: {file_path}", OutputType.INFO)
+                    else:
+                        PrettyOutput.print(f"文档已存在且未更改: {file_path}", OutputType.INFO)
+                        
+                except Exception as e:
+                    PrettyOutput.print(f"处理文件 {file_path} 时出错: {str(e)}", OutputType.ERROR)
+            
+            # 执行查询
+            results = collection.query(
+                query_texts=[query],
+                n_results=num_passages
+            )
+            
+            # 格式化输出
+            output = [f"查询: {query}\n"]
+            output.append(f"找到 {len(results['documents'][0])} 个相关段落:\n")
+            
+            for i, (doc, metadata) in enumerate(zip(results['documents'][0], results['metadatas'][0]), 1):
+                output.append(f"\n段落 {i}:")
+                output.append(f"来源: {metadata['file_path']}")
+                output.append(f"相关内容:\n{doc}\n")
+                output.append("-" * 50)
+            
+            return {
+                "success": True,
+                "stdout": "\n".join(output),
+                "stderr": ""
+            }
+
+        except Exception as e:
+            return {
+                "success": False,
+                "error": f"RAG查询失败: {str(e)}"
+            } 

jarvis/tools/search.py

@@ -0,0 +1,48 @@
+from typing import Dict, Any
+from duckduckgo_search import DDGS
+from ..utils import PrettyOutput, OutputType
+
+class SearchTool:
+    name = "search"
+    description = "Search for information using DuckDuckGo search engine"
+    parameters = {
+        "type": "object",
+        "properties": {
+            "query": {
+                "type": "string",
+                "description": "Search query string"
+            },
+            "max_results": {
+                "type": "integer",
+                "description": "Maximum number of results to return",
+                "default": 5
+            }
+        },
+        "required": ["query"]
+    }
+
+    def execute(self, args: Dict) -> Dict[str, Any]:
+        """使用DuckDuckGo进行搜索"""
+        try:
+            # 打印搜索查询
+            PrettyOutput.print(f"搜索查询: {args['query']}", OutputType.INFO)
+            
+            # 获取搜索结果
+            with DDGS() as ddgs:
+                results = ddgs.text(
+                    keywords=args["query"],
+                    max_results=args.get("max_results", 5)
+                )
+            
+            
+            return {
+                "success": True,
+                "stdout": results,
+                "stderr": ""
+            }
+
+        except Exception as e:
+            return {
+                "success": False,
+                "error": f"搜索失败: {str(e)}"
+            } 

jarvis/tools/shell.py

@@ -0,0 +1,67 @@
+from typing import Dict, Any
+import subprocess
+from ..utils import PrettyOutput, OutputType
+
+class ShellTool:
+    name = "execute_shell"
+    description = "Execute shell commands and return the results"
+    parameters = {
+        "type": "object",
+        "properties": {
+            "command": {
+                "type": "string",
+                "description": "Shell command to execute"
+            },
+            "timeout": {
+                "type": "integer",
+                "description": "Command execution timeout in seconds",
+                "default": 30
+            }
+        },
+        "required": ["command"]
+    }
+
+    def execute(self, args: Dict) -> Dict[str, Any]:
+        """执行shell命令"""
+        try:
+            # 获取参数
+            command = args["command"]
+            timeout = args.get("timeout", 30)
+            
+            # 执行命令
+            result = subprocess.run(
+                command,
+                shell=True,
+                capture_output=True,
+                text=True,
+                timeout=timeout
+            )
+            
+            # 构建输出
+            output = []
+            
+            # 添加命令信息
+            PrettyOutput.print(f"执行命令: {command}", OutputType.INFO)
+            output.append(f"命令: {command}")
+            output.append("")
+            
+            # 添加输出
+            if result.stdout:
+                output.append(result.stdout)
+            
+            return {
+                "success": True,
+                "stdout": "\n".join(output),
+                "stderr": result.stderr,
+                "return_code": result.returncode
+            }
+        except subprocess.TimeoutExpired:
+            return {
+                "success": False,
+                "error": f"命令执行超时 (>{timeout}秒)"
+            }
+        except Exception as e:
+            return {
+                "success": False,
+                "error": str(e)
+            } 

jarvis/tools/user_confirmation.py

@@ -0,0 +1,58 @@
+from typing import Dict, Any
+from ..utils import PrettyOutput, OutputType
+
+class UserConfirmationTool:
+    name = "ask_user_confirmation"
+    description = "Request confirmation from user, returns yes/no"
+    parameters = {
+        "type": "object",
+        "properties": {
+            "question": {
+                "type": "string",
+                "description": "Question to ask for confirmation"
+            },
+            "details": {
+                "type": "string",
+                "description": "Additional details or context",
+                "default": ""
+            }
+        },
+        "required": ["question"]
+    }
+
+    def execute(self, args: Dict) -> Dict[str, Any]:
+        """获取用户确认"""
+        try:
+            question = args["question"]
+            details = args.get("details", "")
+            
+            # 打印详细信息
+            if details:
+                PrettyOutput.print(details, OutputType.INFO)
+                PrettyOutput.print("", OutputType.INFO)  # 空行
+            
+            # 打印问题
+            PrettyOutput.print(f"{question} (y/n)", OutputType.INFO)
+            
+            while True:
+                response = input(">>> ").strip().lower()
+                if response in ['y', 'yes']:
+                    result = True
+                    break
+                elif response in ['n', 'no']:
+                    result = False
+                    break
+                else:
+                    PrettyOutput.print("请输入 y/yes 或 n/no", OutputType.ERROR)
+            
+            return {
+                "success": True,
+                "stdout": str(result),
+                "stderr": ""
+            }
+
+        except Exception as e:
+            return {
+                "success": False,
+                "error": f"获取用户确认失败: {str(e)}"
+            } 

jarvis/tools/user_interaction.py

@@ -0,0 +1,86 @@
+from typing import Dict, Any, List
+from ..utils import PrettyOutput, OutputType, get_multiline_input
+
+class UserInteractionTool:
+    name = "ask_user"
+    description = "Ask user for information, supports option selection and multiline input"
+    parameters = {
+        "type": "object",
+        "properties": {
+            "question": {
+                "type": "string",
+                "description": "Question to ask the user"
+            },
+            "options": {
+                "type": "array",
+                "items": {"type": "string"},
+                "description": "List of options for user to choose from",
+                "default": []
+            },
+            "multiline": {
+                "type": "boolean",
+                "description": "Allow multiline input",
+                "default": False
+            },
+            "description": {
+                "type": "string",
+                "description": "Additional description or context",
+                "default": ""
+            }
+        },
+        "required": ["question"]
+    }
+
+    def execute(self, args: Dict) -> Dict[str, Any]:
+        """获取用户输入"""
+        try:
+            question = args["question"]
+            options = args.get("options", [])
+            multiline = args.get("multiline", False)
+            description = args.get("description", "")
+            
+            # 打印问题描述
+            if description:
+                PrettyOutput.print(description, OutputType.INFO)
+            
+            # 如果有选项，显示选项列表
+            if options:
+                PrettyOutput.print("\n可选项:", OutputType.INFO)
+                for i, option in enumerate(options, 1):
+                    PrettyOutput.print(f"[{i}] {option}", OutputType.INFO)
+                
+                while True:
+                    try:
+                        choice = input("\n请选择 (输入数字): ").strip()
+                        if not choice:
+                            continue
+                        
+                        choice = int(choice)
+                        if 1 <= choice <= len(options):
+                            response = options[choice - 1]
+                            break
+                        else:
+                            PrettyOutput.print("无效选择，请重试", OutputType.ERROR)
+                    except ValueError:
+                        PrettyOutput.print("请输入有效数字", OutputType.ERROR)
+            
+            # 多行输入
+            elif multiline:
+                response = get_multiline_input(question)
+            
+            # 单行输入
+            else:
+                PrettyOutput.print(question, OutputType.INFO)
+                response = input(">>> ").strip()
+            
+            return {
+                "success": True,
+                "stdout": response,
+                "stderr": ""
+            }
+
+        except Exception as e:
+            return {
+                "success": False,
+                "error": f"获取用户输入失败: {str(e)}"
+            }