jarvis-ai-assistant 0.1.3__py3-none-any.whl → 0.1.4__py3-none-any.whl

jarvis/__init__.py

@@ -1,3 +1,3 @@
 """Jarvis AI Assistant"""
 
-__version__ = "0.1.3" 
+__version__ = "0.1.4" 

jarvis/agent.py

@@ -3,21 +3,32 @@ import subprocess
 from typing import Dict, Any, List, Optional
 from .tools import ToolRegistry
 from .utils import Spinner, PrettyOutput, OutputType, get_multiline_input
-from .models import BaseModel, OllamaModel
+from .models import BaseModel
 import re
 import os
 from datetime import datetime
 
 class Agent:
-    def __init__(self, model: BaseModel, tool_registry: ToolRegistry):
+    def __init__(self, model: BaseModel, tool_registry: Optional[ToolRegistry] = None, name: str = "Jarvis"):
+        """Initialize Agent with a model, optional tool registry and name"""
         self.model = model
-        self.tool_registry = tool_registry
+        self.tool_registry = tool_registry or ToolRegistry(model)
+        self.name = name
         # 编译正则表达式
         self.tool_call_pattern = re.compile(r'<tool_call>\s*({[^}]+})\s*</tool_call>')
         self.messages = [
             {
                 "role": "system",
-                "content": """You are a rigorous AI assistant, all data must be obtained through tools, and no fabrication or speculation is allowed. """ + "\n" + self.tool_registry.tool_help_text()
+                "content": f"""You are {name}, a rigorous AI assistant that executes tasks step by step.
+
+Key Principles:
+1. Execute ONE step at a time
+2. Wait for each step's result before planning the next
+3. Use tools to obtain all data, no fabrication
+4. Create sub-agents for independent subtasks
+5. Think carefully before each action
+
+""" + self.tool_registry.tool_help_text()
             }
         ]
         self.spinner = Spinner()
@@ -28,44 +39,38 @@ class Agent:
         try:
             return self.model.chat(
                 messages=messages,
-                tools=self.tool_registry.get_all_tools() if use_tools else None
+                tools=self.tool_registry.get_all_tools() if use_tools else []
             )
         except Exception as e:
-            raise Exception(f"模型调用失败: {str(e)}")
+            raise Exception(f"{self.name}: 模型调用失败: {str(e)}")
         finally:
             self.spinner.stop()
 
-
-    def run(self, user_input: str) :
-        """处理用户输入并返回响应"""
-        # 检查是否是结束命令
+    def run(self, user_input: str) -> str:
+        """处理用户输入并返回响应，返回任务总结报告"""
         self.clear_history()
         self.messages.append({
             "role": "user",
             "content": user_input
         })
+        
         while True:
             try:
-                # 获取初始响应
                 response = self._call_model(self.messages)
                 current_response = response
 
-                # 将工具执行结果添加到对话
                 self.messages.append({
                     "role": "assistant",
                     "content": response["message"].get("content", ""),
                     "tool_calls": current_response["message"]["tool_calls"]
                 })
                 
-                # 处理可能的多轮工具调用
                 if len(current_response["message"]["tool_calls"]) > 0:
-                    # 添加当前助手响应到输出（如果有内容）
                     if current_response["message"].get("content"):
-                        PrettyOutput.print(current_response["message"]["content"], OutputType.SYSTEM)
+                        PrettyOutput.print(f"{self.name}: {current_response['message']['content']}", OutputType.SYSTEM)
                         
-                    # 使用 ToolRegistry 的 handle_tool_calls 方法处理工具调用
                     tool_result = self.tool_registry.handle_tool_calls(current_response["message"]["tool_calls"])
-                    PrettyOutput.print(tool_result, OutputType.RESULT)
+                    PrettyOutput.print(f"{self.name} Tool Result: {tool_result}", OutputType.RESULT)
 
                     self.messages.append({
                         "role": "tool",
@@ -73,18 +78,46 @@ class Agent:
                     })
                     continue
                 
-                # 添加最终响应到对话历史和输出
                 final_content = current_response["message"].get("content", "")
-                
                 if final_content:
-                    PrettyOutput.print(final_content, OutputType.SYSTEM)
-                    
+                    PrettyOutput.print(f"{self.name}: {final_content}", OutputType.SYSTEM)
                 
-                # 如果没有工具调用且响应很短，可能需要继续对话
-                user_input = get_multiline_input("您可以继续输入，或输入空行结束当前任务")
-                if  not user_input:
-                    PrettyOutput.print("===============任务结束===============", OutputType.INFO)
-                    break
+                user_input = get_multiline_input(f"{self.name}: 您可以继续输入，或输入空行结束当前任务")
+                if not user_input:
+                    PrettyOutput.print(f"{self.name}: 正在生成任务总结...", OutputType.INFO)
+                    
+                    # 请求模型生成任务总结
+                    summary_prompt = {
+                        "role": "user",
+                        "content": """Please provide a concise task summary focusing on:
+
+1. Key Information:
+   - Important data points
+   - Critical findings
+   - Significant tool outputs
+
+2. Task Results:
+   - Final outcomes
+   - Task completion status
+   - Any important conclusions
+
+Keep the summary focused and brief, highlighting only the most essential information."""
+                    }
+                    
+                    try:
+                        summary_response = self._call_model(self.messages + [summary_prompt], use_tools=False)
+                        summary = summary_response["message"].get("content", "")
+                        
+                        PrettyOutput.print(f"==============={self.name} 任务总结===============", OutputType.INFO)
+                        PrettyOutput.print(summary, OutputType.SYSTEM)
+                        PrettyOutput.print("=" * (len(self.name) + 16), OutputType.INFO)
+                        
+                        return summary
+                        
+                    except Exception as e:
+                        error_msg = f"{self.name}: 生成任务总结时出错: {str(e)}"
+                        PrettyOutput.print(error_msg, OutputType.ERROR)
+                        return "Task completed but summary generation failed."
                 
                 self.messages.append({
                     "role": "user",
@@ -92,8 +125,9 @@ class Agent:
                 })
                 
             except Exception as e:
-                error_msg = f"处理响应时出错: {str(e)}"
+                error_msg = f"{self.name}: 处理响应时出错: {str(e)}"
                 PrettyOutput.print(error_msg, OutputType.ERROR)
+                return f"Task failed with error: {str(e)}"
 
     def clear_history(self):
         """清除对话历史，只保留系统提示"""

jarvis/main.py

@@ -13,7 +13,8 @@ sys.path.insert(0, str(Path(__file__).parent.parent))
 from jarvis.agent import Agent
 from jarvis.tools import ToolRegistry
 from jarvis.models import DDGSModel, OllamaModel
-from jarvis.utils import PrettyOutput, OutputType, get_multiline_input
+from jarvis.utils import PrettyOutput, OutputType, get_multiline_input, load_env_from_file
+from jarvis.zte_llm import create_zte_llm
 
 # 定义支持的平台和模型
 SUPPORTED_PLATFORMS = {
@@ -24,6 +25,10 @@ SUPPORTED_PLATFORMS = {
     "ddgs": {
         "models": ["gpt-4o-mini", "claude-3-haiku", "llama-3.1-70b", "mixtral-8x7b"],
         "default": "gpt-4o-mini"
+    },
+    "zte": {
+        "models": ["NebulaBiz", "nebulacoder", "NTele-72B"],
+        "default": "NebulaBiz"
     }
 }
 
@@ -97,6 +102,8 @@ def main():
     )
     
     args = parser.parse_args()
+
+    load_env_from_file()
     
     # 验证并设置默认模型
     if args.model:
@@ -119,9 +126,12 @@ def main():
                 api_base=args.api_base
             )
             platform_name = f"Ollama ({args.model})"
-        else:  # ddgs
+        elif args.platform == "ddgs":  # ddgs
             model = DDGSModel(model_name=args.model)
             platform_name = f"DuckDuckGo Search ({args.model})"
+        elif args.platform == "zte":  # zte
+            model = create_zte_llm(model_name=args.model)
+            platform_name = f"ZTE ({args.model})"
 
         tool_registry = ToolRegistry()
         agent = Agent(model, tool_registry)

jarvis/tools/__init__.py

@@ -5,7 +5,6 @@ from .search import SearchTool
 from .shell import ShellTool
 from .user_interaction import UserInteractionTool
 from .user_confirmation import UserConfirmationTool
-from .rag import RAGTool
 from .webpage import WebpageTool
 
 __all__ = [

jarvis/tools/base.py

@@ -1,6 +1,7 @@
 from typing import Dict, Any, List, Optional, Callable
 import json
 from ..utils import PrettyOutput, OutputType
+from ..models import BaseModel
 
 class Tool:
     def __init__(self, name: str, description: str, parameters: Dict, func: Callable):
@@ -25,8 +26,9 @@ class Tool:
         return self.func(arguments)
 
 class ToolRegistry:
-    def __init__(self):
+    def __init__(self, model: Optional[BaseModel] = None):
         self.tools: Dict[str, Tool] = {}
+        self.model = model
         self._register_default_tools()
 
     def _register_default_tools(self):
@@ -37,8 +39,8 @@ class ToolRegistry:
         from .user_confirmation import UserConfirmationTool
         from .python_script import PythonScriptTool
         from .file_ops import FileOperationTool
-        from .rag import RAGTool
         from .webpage import WebpageTool
+        from .sub_agent import SubAgentTool
 
         tools = [
             SearchTool(),
@@ -47,7 +49,6 @@ class ToolRegistry:
             UserConfirmationTool(),
             PythonScriptTool(),
             FileOperationTool(),
-            RAGTool(),
             WebpageTool(),
         ]
 
@@ -59,6 +60,14 @@ class ToolRegistry:
                 func=tool.execute
             )
 
+        sub_agent_tool = SubAgentTool(self.model)
+        self.register_tool(
+            name=sub_agent_tool.name,
+            description=sub_agent_tool.description,
+            parameters=sub_agent_tool.parameters,
+            func=sub_agent_tool.execute
+        )
+
     def register_tool(self, name: str, description: str, parameters: Dict, func: Callable):
         """注册新工具"""
         self.tools[name] = Tool(name, description, parameters, func)
@@ -125,36 +134,30 @@ class ToolRegistry:
 5. ask_user: Get input from user with options support
 6. ask_user_confirmation: Get yes/no confirmation from user
 7. file_operation: Read/write files in workspace directory
-8. rag_query: Query documents using RAG
-
-Output Guidelines:
-1. Focus on Essential Information
-   - Filter out irrelevant or redundant information
-   - Only output the most relevant search results
-   - Summarize long content when possible
-
-2. Context Management
-   - Avoid repeating information that's already known
-   - Don't echo back the entire content of files
-   - Use snippets instead of full documents
-
-3. Error Handling
-   - Only show relevant parts of error messages
-   - Include troubleshooting steps when appropriate
-   - Omit stack traces unless specifically requested
-
-4. Search Results
-   - Focus on the most relevant 2-3 results
-   - Extract key information instead of full articles
-   - Skip duplicate or similar content
-
-5. RAG Queries
-   - Return only the most relevant passages
-   - Combine similar information
-   - Skip redundant context
-
-Tool Call Format Requirements:
-1. MUST use exact format:
+8. create_sub_agent: Create a sub-agent for independent tasks (RECOMMENDED for subtasks)
+
+Core Rules:
+1. ONE Step at a Time
+   - Execute only one action per response
+   - Explain what you're going to do before doing it
+   - Wait for the result before planning next step
+   - No multiple actions or parallel execution
+
+2. Sub-Agent Usage
+   - Use create_sub_agent for independent subtasks
+   - Let sub-agents handle complex task sequences
+   - Examples of good sub-agent tasks:
+     * Code analysis and documentation
+     * File system operations
+     * Data processing
+     * Research tasks
+
+3. Output Guidelines
+   - Focus on essential information
+   - Clear step-by-step explanations
+   - Summarize results concisely
+
+Tool Call Format:
 <tool_call>
 {
     "name": "tool_name",
@@ -164,48 +167,26 @@ Tool Call Format Requirements:
 }
 </tool_call>
 
-2. Format Rules:
-   - NO comments or explanations inside tool_call blocks
-   - NO additional text or formatting within JSON
-   - NO markdown or other markup inside tool_call
-   - ONLY valid JSON allowed in arguments
-   - ALL arguments must match tool parameters exactly
-
-3. Invalid Examples (DO NOT USE):
-   ❌ <tool_call>
-   {
-       "name": "search",  # Search tool
-       "arguments": {
-           "query": "Python GIL"  // Search query
-       }
-   }
-   </tool_call>
-
-   ❌ <tool_call>
-   ```json
-   {
-       "name": "search",
-       "arguments": {
-           "query": "Python GIL"
-       }
-   }
-   ```
-   </tool_call>
-
-4. Valid Example:
-   ✅ <tool_call>
-   {
-       "name": "search",
-       "arguments": {
-           "query": "Python GIL",
-           "max_results": 2
-       }
-   }
-   </tool_call>
+Format Rules:
+1. ONE tool call per response
+2. Valid JSON only in arguments
+3. No comments or extra formatting
+4. Match parameters exactly
+
+Example (Correct - Single Step):
+Assistant: I'll start by searching for the documentation.
+<tool_call>
+{
+    "name": "search",
+    "arguments": {
+        "query": "Python GIL documentation",
+        "max_results": 2
+    }
+}
+</tool_call>
 
 Remember:
-1. Quality over quantity
-2. Relevance over completeness
-3. Conciseness over verbosity
-4. Context efficiency is crucial
-5. STRICT adherence to tool call format"""
+1. ONE step at a time
+2. Explain before acting
+3. Wait for results
+4. Use sub-agents for complex tasks"""

jarvis/tools/sub_agent.py

@@ -0,0 +1,79 @@
+from typing import Dict, Any
+from ..agent import Agent
+from ..models import BaseModel
+from ..utils import PrettyOutput, OutputType
+from .base import ToolRegistry
+
+class SubAgentTool:
+    name = "create_sub_agent"
+    description = """Create a sub-agent to handle independent tasks.
+    
+Use this tool when:
+1. A subtask can be executed independently
+2. The task requires separate context management
+3. To optimize token usage of the main agent
+4. For parallel task processing
+
+The sub-agent will:
+1. Inherit all tools from the parent agent
+2. Maintain its own conversation history
+3. Return a comprehensive task summary
+"""
+    parameters = {
+        "type": "object",
+        "properties": {
+            "name": {
+                "type": "string",
+                "description": "Name for the sub-agent (e.g., 'FileAnalyzer', 'CodeReviewer')"
+            },
+            "task": {
+                "type": "string",
+                "description": "Task description with complete context"
+            },
+            "context": {
+                "type": "string",
+                "description": "Additional context or background information",
+                "default": ""
+            }
+        },
+        "required": ["name", "task"]
+    }
+
+    def __init__(self, model: BaseModel):
+        """Initialize with the same model as parent agent"""
+        self.model = model
+
+    def execute(self, args: Dict) -> Dict[str, Any]:
+        """Create and run a sub-agent for the specified task"""
+        try:
+            name = args["name"]
+            task = args["task"]
+            context = args.get("context", "")
+
+            PrettyOutput.print(f"Creating sub-agent '{name}'...", OutputType.INFO)
+            
+            # Create a new tool registry for the sub-agent
+            tool_registry = ToolRegistry()
+            
+            # Create the sub-agent with the specified name
+            sub_agent = Agent(self.model, tool_registry, name=name)
+            
+            # Prepare the task with context if provided
+            full_task = f"{context}\n\nTask: {task}" if context else task
+            
+            PrettyOutput.print(f"Sub-agent '{name}' executing task...", OutputType.INFO)
+            
+            # Execute the task and get the summary
+            summary = sub_agent.run(full_task)
+            
+            return {
+                "success": True,
+                "stdout": f"Sub-agent '{name}' completed the task.\n\nSummary:\n{summary}",
+                "stderr": ""
+            }
+
+        except Exception as e:
+            return {
+                "success": False,
+                "error": f"Sub-agent execution failed: {str(e)}"
+            } 

jarvis/utils.py

@@ -1,11 +1,13 @@
+from pathlib import Path
 import sys
 import time
 import threading
-from typing import Optional
+from typing import Dict, Optional
 from enum import Enum
 from datetime import datetime
 import colorama
 from colorama import Fore, Style
+import os
 
 # 初始化colorama
 colorama.init()
@@ -102,4 +104,25 @@ def get_multiline_input(tip: str) -> str:
             PrettyOutput.print("\n输入已取消", OutputType.ERROR)
             return ""
     
-    return "\n".join(lines).strip()
+    return "\n".join(lines).strip()
+
+
+def load_env_from_file():
+    """Load environment variables from ~/.jarvis_env file"""
+    env_file = Path.home() / ".jarvis_env"
+    
+    if env_file.exists():
+        try:
+            with open(env_file, "r", encoding="utf-8") as f:
+                for line in f:
+                    line = line.strip()
+                    if line and not line.startswith("#"):
+                        try:
+                            key, value = line.split("=", 1)
+                            os.environ[key.strip()] = value.strip().strip("'").strip('"')
+                        except ValueError:
+                            continue
+        except Exception as e:
+            print(f"Warning: Failed to read ~/.jarvis_env: {e}")
+    
+    

jarvis/zte_llm.py

@@ -0,0 +1,132 @@
+import requests
+import json
+import os
+from pathlib import Path
+from typing import Dict, Any, List, Optional
+from .models import BaseModel
+
+class ZteLLM(BaseModel):
+    """ZTE Nebula LLM implementation"""
+    
+    def __init__(self, 
+                 app_id: str,
+                 app_key: str,
+                 emp_no: str,
+                 auth_value: str,
+                 model: str = "nebulacoder",
+                 ):
+        """Initialize ZTE LLM with required credentials"""
+        self.app_id = str(app_id)
+        self.app_key = str(app_key)
+        self.emp_no = str(emp_no)
+        self.auth_value = str(auth_value)
+        self.model = model
+        self.base_url = "https://studio.zte.com.cn/zte-studio-ai-platform/openapi/v1"
+        
+    def _make_request(self, endpoint: str, data: Dict[str, Any]) -> Dict[str, Any]:
+        """Make request to ZTE API"""
+        headers = {
+            'Content-Type': 'application/json',
+            'Authorization': f'Bearer {self.app_id}-{self.app_key}',
+            'X-Emp-No': self.emp_no,
+            'X-Auth-Value': self.auth_value
+        }
+        
+        response = requests.post(
+            f"{self.base_url}/{endpoint}",
+            headers=headers,
+            json=data
+        )
+        
+        response.raise_for_status()
+        result = response.json()
+        
+        if result["code"]["code"] != "0000":
+            raise Exception(f"API Error: {result['code']['msg']}")
+            
+        return result["bo"]
+
+    def chat(self, messages: List[Dict[str, Any]], tools: Optional[List[Dict]] = None) -> Dict[str, Any]:
+        """Chat with ZTE LLM"""
+        # Convert messages to prompt
+        prompt = self._convert_messages_to_prompt(messages)
+        
+        # Prepare data for API call
+        data = {
+            "chatUuid": "",
+            "chatName": "",
+            "stream": False,
+            "keep": False,
+            "text": prompt,
+            "model": self.model
+        }
+        
+        # If tools are provided, add them to the prompt
+        if tools:
+            tools_desc = "Available tools:\n\n" + json.dumps(tools, indent=2, ensure_ascii=False)
+            data["text"] = tools_desc + "\n\n" + data["text"]
+        
+        try:
+            result = self._make_request("chat", data)
+            
+            # Parse the response to extract potential tool calls
+            response_text = result["result"]
+            tool_calls = BaseModel.extract_tool_calls(response_text)
+            
+            return {
+                "message": {
+                    "content": response_text,
+                    "tool_calls": tool_calls
+                }
+            }
+            
+        except Exception as e:
+            raise Exception(f"ZTE LLM chat failed: {str(e)}")
+
+    def _convert_messages_to_prompt(self, messages: List[Dict[str, Any]]) -> str:
+        """Convert message list to a single prompt string"""
+        prompt_parts = []
+        
+        for message in messages:
+            role = message["role"]
+            content = message.get("content", "")
+            
+            if role == "system":
+                prompt_parts.append(f"System: {content}")
+            elif role == "user":
+                prompt_parts.append(f"User: {content}")
+            elif role == "assistant":
+                prompt_parts.append(f"Assistant: {content}")
+            elif role == "tool":
+                prompt_parts.append(f"Tool Result: {content}")
+                
+        return "\n\n".join(prompt_parts)
+
+
+def create_zte_llm(model_name: str = "NebulaBiz") -> ZteLLM:
+    """Create ZTE LLM instance with provided parameters"""
+    # Load environment variables from file
+    
+    # Get credentials from parameters, env file, or system environment variables
+    app_id = os.getenv('ZTE_APP_ID') 
+    app_key = os.getenv('ZTE_APP_KEY') 
+    emp_no = os.getenv('ZTE_EMP_NO') 
+    auth_value = os.getenv('ZTE_AUTH_VALUE') 
+    
+    # Validate required credentials
+    if not all([app_id, app_key, emp_no, auth_value]):
+        raise ValueError(
+            "Missing required credentials. Please provide through either:\n"
+            "1. Function parameters\n"
+            "2. ~/.jarvis_env file\n"
+            "3. System environment variables\n\n"
+            "Required variables: ZTE_APP_ID, ZTE_APP_KEY, ZTE_EMP_NO, ZTE_AUTH_VALUE"
+        )
+    
+    return ZteLLM(
+        app_id=app_id,
+        app_key=app_key,
+        emp_no=emp_no,
+        auth_value=auth_value,
+        model=model_name
+    ) 

{jarvis_ai_assistant-0.1.3.dist-info → jarvis_ai_assistant-0.1.4.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: jarvis-ai-assistant
-Version: 0.1.3
+Version: 0.1.4
 Summary: Jarvis: An AI assistant that uses tools to interact with the system
 Home-page: https://github.com/skyfireitdiy/Jarvis
 Author: skyfire
@@ -21,8 +21,7 @@ Requires-Dist: beautifulsoup4>=4.9.3
 Requires-Dist: duckduckgo-search>=3.0.0
 Requires-Dist: pyyaml>=5.1
 Requires-Dist: ollama>=0.1.6
-Requires-Dist: sentence-transformers>=2.5.1
-Requires-Dist: chromadb>=0.4.24
+Requires-Dist: colorama>=0.4.6
 Provides-Extra: dev
 Requires-Dist: pytest; extra == "dev"
 Requires-Dist: black; extra == "dev"

{jarvis_ai_assistant-0.1.3.dist-info → jarvis_ai_assistant-0.1.4.dist-info}/RECORD

@@ -1,21 +1,22 @@
 jarvis/.jarvis,sha256=S4ZMmLqlVLHAPmnwHFQ3vYt0gnDFp-KQRH4okBh8Hpw,209
-jarvis/__init__.py,sha256=rLUKzHgrbk82WAMvg4u4GlFDRal1Yw3XpMg_741mQXw,49
-jarvis/agent.py,sha256=hEatrIothG6C_RuOsygZz4ez1aF1EycnUPwY2Krwiwk,4092
-jarvis/main.py,sha256=wAeeGoRq8fFiTWUCQ31MYVeU9o6SZ48ORUow0jlMJJE,5259
+jarvis/__init__.py,sha256=cIUjE04zhBHLfqRFdRHAeTF1YjKzNHQ6iFpTwPWsrA0,49
+jarvis/agent.py,sha256=2KUmbWDqcyCEC5QgN1YXCliJvO6HGxfQBzLz1IPgG9Y,5538
+jarvis/main.py,sha256=LFymNA5xFzB48IlSzWYSkrhNrUq0fXAukVfu0cV-aQQ,5635
 jarvis/models.py,sha256=ZQAyc39e_UsNmHkME6ATp6053safQ7Gog-oWwETsrMM,4415
-jarvis/utils.py,sha256=-fOPAgiKVPC6DevHM7A7IJvvIzvvFlGM0pHBKyBXDcA,3044
+jarvis/utils.py,sha256=-5tugkGABZdi2QXVStVtCQfag9CRTfCsGUvReitmHUM,3816
+jarvis/zte_llm.py,sha256=ygUeGd8WGm0qrnkt5TrCLnaf5ZFZ0tdfwrA7nrByt4g,4504
 jarvis/__pycache__/__init__.cpython-313.pyc,sha256=u64PLu4KyOPrirhIqNXZSkj9B6yB1L_ohFYhPX8p334,208
 jarvis/__pycache__/agent.cpython-313.pyc,sha256=I1JYUKC-zedURwuxQQjd9saj_yQKUhpdbb1KNWlY6Ss,4940
 jarvis/__pycache__/models.cpython-313.pyc,sha256=fot2VOV0lgODg4b_WuD-kkq-g7uZGLuZlYXiljlW13U,6350
 jarvis/__pycache__/tools.cpython-313.pyc,sha256=lAD4LrnnWzNZQmHXGfZ_2l7oskOpr2_2OC-gdFhxQY8,33933
 jarvis/__pycache__/utils.cpython-313.pyc,sha256=o8Fubq8Dcrvv3ATzQ9BnZoZjQmU3_FNPrfTt0ANdKdI,5804
-jarvis/tools/__init__.py,sha256=FNF5X32UzKvlT08-hPz7-7vcIyqqMngQYj36LXyNbV0,553
-jarvis/tools/base.py,sha256=Y1XInc2skUVZYZzLRPlCpH3JyY3pk1XNZgys6MSCudQ,6586
+jarvis/tools/__init__.py,sha256=FPXTV4QQy5KrUFyhoMW1H9lBlIhTHAAyT0P-QxTBnuY,528
+jarvis/tools/base.py,sha256=iiA2M6IeAG6edVSTvXDBevlQAq6yMrkUINujNW6lg2s,6274
 jarvis/tools/file_ops.py,sha256=05Vc4u-NGMjLD15WI52eL_nt_RyYt-Z_cXJnnBepf-c,3781
 jarvis/tools/python_script.py,sha256=_eK8LqNs-Mz50zdcgwbjdd8-qAeOl6kJ_qRDvWawGMw,5006
-jarvis/tools/rag.py,sha256=eBrn1fdqy-nd2nR0b4oH1EQpthTYVen-sYDhC5Ypl38,5647
 jarvis/tools/search.py,sha256=dyJmeP_s2tWv5KQejOl-TuVF12B6SXViiXEyTemD1Wo,1412
 jarvis/tools/shell.py,sha256=VhEqEdoUGIRW6CT6F0YNMAlzJxpHcsKS9hSS0VC3ezw,2644
+jarvis/tools/sub_agent.py,sha256=I_OmSq1EIcQO0u8zCDqj4hMtsSx4NHiBNTcu9UOxbvQ,2673
 jarvis/tools/user_confirmation.py,sha256=p0JverifHJfnALeIhtKtaVBBVEGkgSpUzT-PSycG4NY,1850
 jarvis/tools/user_interaction.py,sha256=xOaoYsCKTa9-K3mg5tsu1zegn1HKcO40k2aJ3VygZqM,3073
 jarvis/tools/webpage.py,sha256=UTEomu5j7jbOw8c5az2jsjv5E7LeokWKj1QahvZO7xc,3077
@@ -29,8 +30,8 @@ jarvis/tools/__pycache__/shell.cpython-313.pyc,sha256=Xtqt-LzMRfsCXeCE7QxFY8TR8X
 jarvis/tools/__pycache__/user_confirmation.cpython-313.pyc,sha256=wK3Ev10lHSUSRvoYmi7A0GzxYkzU-C4Wfhs5qW_HBqs,2271
 jarvis/tools/__pycache__/user_interaction.cpython-313.pyc,sha256=RuVZ-pmiPBDywY3efgXSfohMAciC1avMGPmBK5qlnew,3305
 jarvis/tools/__pycache__/webpage.cpython-313.pyc,sha256=VcpkaV8IyOOtebedXjn8ybjr8AglU-3-Cs80yzE4FYo,3628
-jarvis_ai_assistant-0.1.3.dist-info/METADATA,sha256=eVKO4jecgKkHOIRUXZGvfvi5KiR9CGaZuQLhDR9W3NY,3675
-jarvis_ai_assistant-0.1.3.dist-info/WHEEL,sha256=A3WOREP4zgxI0fKrHUG8DC8013e3dK3n7a6HDbcEIwE,91
-jarvis_ai_assistant-0.1.3.dist-info/entry_points.txt,sha256=iKu7OMfew9dtfGhW71gIMTg4wvafuPqKb4wyQOnMAGU,44
-jarvis_ai_assistant-0.1.3.dist-info/top_level.txt,sha256=1BOxyWfzOP_ZXj8rVTDnNCJ92bBGB0rwq8N1PCpoMIs,7
-jarvis_ai_assistant-0.1.3.dist-info/RECORD,,
+jarvis_ai_assistant-0.1.4.dist-info/METADATA,sha256=hiXlGsoga_MEvAZtKDEz0LmrWMUlV-ye8OBbHYQqwUc,3630
+jarvis_ai_assistant-0.1.4.dist-info/WHEEL,sha256=A3WOREP4zgxI0fKrHUG8DC8013e3dK3n7a6HDbcEIwE,91
+jarvis_ai_assistant-0.1.4.dist-info/entry_points.txt,sha256=iKu7OMfew9dtfGhW71gIMTg4wvafuPqKb4wyQOnMAGU,44
+jarvis_ai_assistant-0.1.4.dist-info/top_level.txt,sha256=1BOxyWfzOP_ZXj8rVTDnNCJ92bBGB0rwq8N1PCpoMIs,7
+jarvis_ai_assistant-0.1.4.dist-info/RECORD,,

jarvis/tools/rag.py

@@ -1,154 +0,0 @@
-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)}"
-            }