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

jarvis/__init__.py

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

jarvis/agent.py

@@ -2,70 +2,137 @@ import json
 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 .utils import PrettyOutput, OutputType, get_multiline_input
+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: ToolRegistry, name: str = "Jarvis"):
+        """Initialize Agent with a model, optional tool registry and name"""
         self.model = model
-        self.tool_registry = tool_registry
-        # 编译正则表达式
-        self.tool_call_pattern = re.compile(r'<tool_call>\s*({[^}]+})\s*</tool_call>')
+        self.tool_registry = tool_registry or ToolRegistry(model)
+        self.name = name
+
+        # 构建工具说明
+        tools_prompt = "Available Tools:\n"
+        for tool in self.tool_registry.get_all_tools():
+            tools_prompt += f"- Tool: {tool['function']['name']}\n"
+            tools_prompt += f"  Description: {tool['function']['description']}\n"
+            tools_prompt += f"  Arguments: {tool['function']['parameters']}\n"
+
         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}, an AI assistant that follows the ReAct (Reasoning + Acting) framework to solve tasks step by step.
+
+FRAMEWORK:
+1. Thought: Analyze the current situation and plan the next step
+2. Action: Execute ONE specific tool call
+3. Observation: Review the result
+4. Next: Plan the next step or conclude
+
+FORMAT:
+Thought: I need to [reasoning about the current situation]...
+Action: I will use [tool] to [purpose]...
+<START_TOOL_CALL>
+name: tool_name
+arguments:
+    param1: value1
+<END_TOOL_CALL>
+
+After receiving result:
+Observation: The tool returned [analyze result]...
+Next: Based on this, I will [next step]...
+
+CORE RULES:
+1. ONE Action Per Response
+   - Only ONE tool call per response
+   - Additional tool calls will be ignored
+   - Complete current step before next
+
+2. Clear Reasoning
+   - Explain your thought process
+   - Justify tool selection
+   - Analyze results thoroughly
+
+Examples:
+✓ Good Response:
+Thought: I need to check the content of utils.py first to understand its structure.
+Action: I will read the file content.
+<START_TOOL_CALL>
+name: file_operation
+arguments:
+    operation: read
+    filepath: src/utils.py
+<END_TOOL_CALL>
+
+✗ Bad Response:
+Thought: Let's analyze the code.
+Action: I'll read and check everything.
+[Multiple or vague tool calls...]
+
+Remember:
+- Always start with "Thought:"
+- Use exactly ONE tool per response
+- Wait for results before next step
+- Clearly explain your reasoning
+
+{tools_prompt}"""
             }
         ]
-        self.spinner = Spinner()
 
     def _call_model(self, messages: List[Dict], use_tools: bool = True) -> Dict:
         """调用模型获取响应"""
-        self.spinner.start()
         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)}")
-        finally:
-            self.spinner.stop()
-
+            raise Exception(f"{self.name}: 模型调用失败: {str(e)}")
 
-    def run(self, user_input: str) :
-        """处理用户输入并返回响应"""
-        # 检查是否是结束命令
+    def run(self, user_input: str) -> str:
+        """处理用户输入并返回响应，返回任务总结报告"""
         self.clear_history()
+        
+        # 显示任务开始
+        PrettyOutput.section(f"开始新任务: {self.name}", OutputType.PLANNING)
+        
         self.messages.append({
             "role": "user",
             "content": user_input
         })
+        
         while True:
             try:
-                # 获取初始响应
+                # 显示思考状态
+                PrettyOutput.print("分析任务...", OutputType.PROGRESS)
                 response = self._call_model(self.messages)
                 current_response = response
 
-                # 将工具执行结果添加到对话
+                # 流式输出已经在model中处理，这里添加换行
+                PrettyOutput.print_stream_end()
+
                 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)
                         
-                    # 使用 ToolRegistry 的 handle_tool_calls 方法处理工具调用
-                    tool_result = self.tool_registry.handle_tool_calls(current_response["message"]["tool_calls"])
-                    PrettyOutput.print(tool_result, OutputType.RESULT)
+                    try:
+                        # 显示工具调用
+                        PrettyOutput.print("执行工具调用...", OutputType.PROGRESS)
+                        tool_result = self.tool_registry.handle_tool_calls(current_response["message"]["tool_calls"])
+                        PrettyOutput.print(tool_result, OutputType.RESULT)
+                    except Exception as e:
+                        PrettyOutput.print(str(e), OutputType.ERROR)
+                        tool_result = f"Tool call failed: {str(e)}"
 
                     self.messages.append({
                         "role": "tool",
@@ -73,18 +140,47 @@ class Agent:
                     })
                     continue
                 
-                # 添加最终响应到对话历史和输出
-                final_content = current_response["message"].get("content", "")
-                
-                if final_content:
-                    PrettyOutput.print(final_content, OutputType.SYSTEM)
+                # 获取用户输入
+                user_input = get_multiline_input(f"{self.name}: 您可以继续输入，或输入空行结束当前任务")
+                if not user_input:
+                    PrettyOutput.print("生成任务总结...", OutputType.PROGRESS)
+                    
+                    # 生成任务总结
+                    summary_prompt = {
+                        "role": "user",
+                        "content": """The task has been completed. Based on the previous analysis and execution results, provide a task summary including:
+
+1. Key Information:
+   - Essential findings from analysis
+   - Important results from tool executions
+   - Critical data discovered
+
+2. Task Results:
+   - Final outcome
+   - Actual achievements
+   - Concrete results
+
+Focus only on facts and actual results. Be direct and concise."""
+                    }
                     
+                    while True:
+                        try:
+                            summary_response = self._call_model(self.messages + [summary_prompt], use_tools=False)
+                            summary = summary_response["message"].get("content", "")
+                            
+                            # 显示任务总结
+                            PrettyOutput.section("任务总结", OutputType.SUCCESS)
+                            PrettyOutput.print(summary, OutputType.SYSTEM)
+                            PrettyOutput.section("任务完成", OutputType.SUCCESS)
+                            
+                            return summary
+                            
+                        except Exception as e:
+                            PrettyOutput.print(str(e), OutputType.ERROR)
                 
-                # 如果没有工具调用且响应很短，可能需要继续对话
-                user_input = get_multiline_input("您可以继续输入，或输入空行结束当前任务")
-                if  not user_input:
-                    PrettyOutput.print("===============任务结束===============", OutputType.INFO)
-                    break
+                if user_input == "__interrupt__":
+                    PrettyOutput.print("任务已取消", OutputType.WARNING)
+                    return "Task cancelled by user"
                 
                 self.messages.append({
                     "role": "user",
@@ -92,8 +188,7 @@ class Agent:
                 })
                 
             except Exception as e:
-                error_msg = f"处理响应时出错: {str(e)}"
-                PrettyOutput.print(error_msg, OutputType.ERROR)
+                PrettyOutput.print(str(e), OutputType.ERROR)
 
     def clear_history(self):
         """清除对话历史，只保留系统提示"""

jarvis/main.py

@@ -13,46 +13,60 @@ 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 = {
     "ollama": {
-        "models": ["llama3.2", "qwen2.5:14b"],
+        "models": ["qwen2.5:14b", "qwq"],
         "default": "qwen2.5:14b"
     },
     "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"
     }
 }
 
-def load_tasks() -> list:
+def load_tasks() -> dict:
     """Load tasks from .jarvis file if it exists."""
     if not os.path.exists(".jarvis"):
-        return []
+        return {}
     
     try:
         with open(".jarvis", "r", encoding="utf-8") as f:
             tasks = yaml.safe_load(f)
             
-        if not isinstance(tasks, list):
-            PrettyOutput.print("Warning: .jarvis file should contain a list of tasks", OutputType.ERROR)
-            return []
+        if not isinstance(tasks, dict):
+            PrettyOutput.print("Warning: .jarvis file should contain a dictionary of task_name: task_description", OutputType.ERROR)
+            return {}
             
-        return [str(task) for task in tasks if task]  # Convert all tasks to strings and filter out empty ones
+        # Validate format and convert all values to strings
+        validated_tasks = {}
+        for name, desc in tasks.items():
+            if desc:  # Ensure description is not empty
+                validated_tasks[str(name)] = str(desc)
+                
+        return validated_tasks
     except Exception as e:
         PrettyOutput.print(f"Error loading .jarvis file: {str(e)}", OutputType.ERROR)
-        return []
+        return {}
 
-def select_task(tasks: list) -> str:
-    """Let user select a task from the list or skip."""
+def select_task(tasks: dict) -> str:
+    """Let user select a task from the list or skip. Returns task description if selected."""
     if not tasks:
         return ""
     
-    PrettyOutput.print("\nFound predefined tasks:", OutputType.INFO)
-    for i, task in enumerate(tasks, 1):
-        PrettyOutput.print(f"[{i}] {task}", OutputType.INFO)
+    # Convert tasks to list for ordered display
+    task_names = list(tasks.keys())
+    
+    PrettyOutput.print("\nAvailable tasks:", OutputType.INFO)
+    for i, name in enumerate(task_names, 1):
+        PrettyOutput.print(f"[{i}] {name}", OutputType.INFO)
     PrettyOutput.print("[0] Skip predefined tasks", OutputType.INFO)
     
     while True:
@@ -64,8 +78,9 @@ def select_task(tasks: list) -> str:
             choice = int(choice)
             if choice == 0:
                 return ""
-            elif 1 <= choice <= len(tasks):
-                return tasks[choice - 1]
+            elif 1 <= choice <= len(task_names):
+                selected_name = task_names[choice - 1]
+                return tasks[selected_name]  # Return the task description
             else:
                 PrettyOutput.print("Invalid choice. Please try again.", OutputType.ERROR)
         except ValueError:
@@ -73,13 +88,16 @@ def select_task(tasks: list) -> str:
 
 def main():
     """Main entry point for Jarvis."""
+
+    load_env_from_file()
+
     parser = argparse.ArgumentParser(description="Jarvis AI Assistant")
     
     # 添加平台选择参数
     parser.add_argument(
         "--platform",
         choices=list(SUPPORTED_PLATFORMS.keys()),
-        default="ollama",
+        default=os.getenv("JARVIS_PLATFORM") or "ddgs",
         help="选择运行平台 (默认: ollama)"
     )
     
@@ -92,11 +110,13 @@ def main():
     # 添加API基础URL参数
     parser.add_argument(
         "--api-base",
-        default="http://localhost:11434",
+        default=os.getenv("JARVIS_OLLAMA_API_BASE") or "http://localhost:11434",
         help="Ollama API基础URL (仅用于Ollama平台, 默认: http://localhost:11434)"
     )
     
     args = parser.parse_args()
+
+    args.model = args.model or os.getenv("JARVIS_MODEL")
     
     # 验证并设置默认模型
     if args.model:
@@ -119,11 +139,14 @@ 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()
+        tool_registry = ToolRegistry(model)
         agent = Agent(model, tool_registry)
 
         # 欢迎信息
@@ -142,12 +165,9 @@ def main():
         while True:
             try:
                 user_input = get_multiline_input("请输入您的任务(输入空行退出):")
-                if not user_input:
+                if not user_input or user_input == "__interrupt__":
                     break
                 agent.run(user_input)
-            except KeyboardInterrupt:
-                print("\n正在退出...")
-                break
             except Exception as e:
                 PrettyOutput.print(f"错误: {str(e)}", OutputType.ERROR)
 

jarvis/models.py

@@ -1,10 +1,10 @@
-import json
 import re
 import time
 from typing import Dict, List, Optional
 from duckduckgo_search import DDGS
 import ollama
 from abc import ABC, abstractmethod
+import yaml
 
 from .utils import OutputType, PrettyOutput
 
@@ -18,34 +18,46 @@ class BaseModel(ABC):
 
     @staticmethod
     def extract_tool_calls(content: str) -> List[Dict]:
-        """从内容中提取工具调用"""
-        tool_calls = []
-        # 修改正则表达式以更好地处理多行内容
-        pattern = re.compile(
-            r'<tool_call>\s*({(?:[^{}]|(?:{[^{}]*})|(?:{(?:[^{}]|{[^{}]*})*}))*})\s*</tool_call>',
-            re.DOTALL  # 添加DOTALL标志以匹配跨行内容
-        )
+        """从内容中提取工具调用，只返回第一个有效的工具调用"""
+        # 分割内容为行
+        lines = content.split('\n')
+        tool_call_lines = []
+        in_tool_call = False
         
-        matches = pattern.finditer(content)
-        for match in matches:
-            try:
-                tool_call_str = match.group(1).strip()
-                tool_call = json.loads(tool_call_str)
-                if isinstance(tool_call, dict) and "name" in tool_call and "arguments" in tool_call:
-                    tool_calls.append({
-                        "function": {
-                            "name": tool_call["name"],
-                            "arguments": tool_call["arguments"]
-                        }
-                    })
-                else:
-                    PrettyOutput.print(f"无效的工具调用格式: {tool_call_str}", OutputType.ERROR)
-            except json.JSONDecodeError as e:
-                PrettyOutput.print(f"JSON解析错误: {str(e)}", OutputType.ERROR)
-                PrettyOutput.print(f"解析内容: {tool_call_str}", OutputType.ERROR)
+        # 逐行处理
+        for line in lines:
+            if not line:
                 continue
+                
+            if line == '<START_TOOL_CALL>':
+                tool_call_lines = []
+                in_tool_call = True
+                continue
+            elif line == '<END_TOOL_CALL>':
+                if in_tool_call and tool_call_lines:
+                    try:
+                        # 解析工具调用内容
+                        tool_call_text = '\n'.join(tool_call_lines)
+                        tool_call_data = yaml.safe_load(tool_call_text)
+                        
+                        # 验证必要的字段
+                        if "name" in tool_call_data and "arguments" in tool_call_data:
+                            # 只返回第一个有效的工具调用
+                            return [{
+                                "function": {
+                                    "name": tool_call_data["name"],
+                                    "arguments": tool_call_data["arguments"]
+                                }
+                            }]
+                    except yaml.YAMLError:
+                        pass  # 跳过无效的YAML
+                    except Exception:
+                        pass  # 跳过其他错误
+                in_tool_call = False
+            elif in_tool_call:
+                tool_call_lines.append(line)
         
-        return tool_calls
+        return []  # 如果没有找到有效的工具调用，返回空列表
 
 
 class DDGSModel(BaseModel):
@@ -59,11 +71,7 @@ class DDGSModel(BaseModel):
         self.model_name = model_name
 
     def __make_prompt(self, messages: List[Dict], tools: Optional[List[Dict]] = None) -> str:
-        prompt = "You are an AI Agent skilled in utilizing tools and planning tasks. Based on the task input by the user and the list of available tools, you output the tool invocation methods in a specified format. The user will provide feedback on the results of the tool execution, allowing you to continue analyzing and ultimately complete the user's designated task. Below is the list of tools and their usage methods. Let's use them step by step to accomplish the user's task.\n"
-        for tool in tools:
-            prompt += f"- Tool: {tool['function']['name']}\n"
-            prompt += f"  Description: {tool['function']['description']}\n"
-            prompt += f"  Arguments: {tool['function']['parameters']}\n"
+        prompt = ""
         for message in messages:
             prompt += f"[{message['role']}]: {message['content']}\n"
         return prompt
@@ -72,6 +80,7 @@ class DDGSModel(BaseModel):
         ddgs = DDGS()
         prompt = self.__make_prompt(messages, tools)
         content = ddgs.chat(prompt)
+        PrettyOutput.print_stream(content, OutputType.SYSTEM)
         tool_calls = BaseModel.extract_tool_calls(content)
         return {
             "message": {
@@ -92,16 +101,25 @@ class OllamaModel(BaseModel):
     def chat(self, messages: List[Dict], tools: Optional[List[Dict]] = None) -> Dict:
         """调用Ollama API获取响应"""
         try:
-            response = self.client.chat(
+            # 使用流式调用
+            stream = self.client.chat(
                 model=self.model_name,
                 messages=messages,
-                tools=tools
+                stream=True
             )
 
-            content = response.message.content
-            tool_calls = response.message.tool_calls or BaseModel.extract_tool_calls(content)
+            # 收集完整响应
+            content_parts = []
+            for chunk in stream:
+                if chunk.message.content:
+                    content_parts.append(chunk.message.content)
+                    # 实时打印内容
+                    PrettyOutput.print_stream(chunk.message.content, OutputType.SYSTEM)
+
+            # 合并完整内容
+            content = "".join(content_parts)
+            tool_calls = BaseModel.extract_tool_calls(content)
             
-            # 转换响应格式
             return {
                 "message": {
                     "content": content,

jarvis/tools/__init__.py

@@ -1,22 +1,16 @@
 from .base import Tool, ToolRegistry
-from .python_script import PythonScript
 from .file_ops import FileOperationTool
 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
+from .user_input import UserInputTool
 
 __all__ = [
     'Tool',
     'ToolRegistry',
-    'PythonScript',
     'FileOperationTool',
     'SearchTool',
     'ShellTool',
-    'UserInteractionTool',
-    'UserConfirmationTool',
-    'RAGTool',
     'WebpageTool',
-] 
+    'UserInputTool',
+]