jarvis-ai-assistant 0.1.224__py3-none-any.whl → 0.2.0__py3-none-any.whl

jarvis/__init__.py

@@ -1,4 +1,4 @@
 # -*- coding: utf-8 -*-
 """Jarvis AI Assistant"""
 
-__version__ = "0.1.224"
+__version__ = "0.2.0"

jarvis/jarvis_agent/__init__.py

@@ -102,6 +102,7 @@ class Agent:
         name: str = "Jarvis",
         description: str = "",
         llm_type: str = "normal",
+        model_group: Optional[str] = None,
         summary_prompt: Optional[str] = None,
         auto_complete: bool = False,
         output_handler: List[OutputHandlerProtocol] = [],
@@ -137,11 +138,11 @@ class Agent:
         self.description = description
         # 初始化平台和模型
         if llm_type == "thinking":
-            platform_name = get_thinking_platform_name()
-            model_name = get_thinking_model_name()
+            platform_name = get_thinking_platform_name(model_group)
+            model_name = get_thinking_model_name(model_group)
         else:  # 默认为 normal
-            platform_name = get_normal_platform_name()
-            model_name = get_normal_model_name()
+            platform_name = get_normal_platform_name(model_group)
+            model_name = get_normal_model_name(model_group)
 
         self.model = PlatformRegistry().create_platform(platform_name)
         if self.model is None:
@@ -153,6 +154,8 @@ class Agent:
         if model_name:
             self.model.set_model_name(model_name)
 
+        self.model.set_model_group(model_group)
+
         self.user_data: Dict[str, Any] = {}
 
         self.model.set_suppress_output(False)
@@ -197,7 +200,7 @@ class Agent:
             summary_prompt if summary_prompt else DEFAULT_SUMMARY_PROMPT
         )
 
-        self.max_token_count = get_max_token_count()
+        self.max_token_count = get_max_token_count(model_group)
         self.auto_complete = auto_complete
         welcome_message = f"{name} 初始化完成 - 使用 {self.model.name()} 模型"
 
@@ -213,6 +216,7 @@ class Agent:
 """
         )
         self.first = True
+        self.run_input_handlers_next_turn = False
 
     def set_user_data(self, key: str, value: Any):
         """Sets user data in the session."""
@@ -236,6 +240,10 @@ class Agent:
         """Sets the addon prompt in the session."""
         self.session.set_addon_prompt(addon_prompt)
 
+    def set_run_input_handlers_next_turn(self, value: bool):
+        """Sets the flag to run input handlers on the next turn."""
+        self.run_input_handlers_next_turn = value
+
     def set_after_tool_call_cb(self, cb: Callable[[Any], None]):  # type: ignore
         """设置工具调用后回调函数。
 
@@ -264,7 +272,9 @@ class Agent:
                 return handler
         return None
 
-    def _call_model(self, message: str, need_complete: bool = False) -> str:
+    def _call_model(
+        self, message: str, need_complete: bool = False, run_input_handlers: bool = True
+    ) -> str:
         """调用AI模型并实现重试逻辑
 
         参数:
@@ -280,10 +290,11 @@ class Agent:
             3. 会自动添加附加提示
             4. 会检查并处理上下文长度限制
         """
-        for handler in self.input_handler:
-            message, need_return = handler(message, self)
-            if need_return:
-                return message
+        if run_input_handlers:
+            for handler in self.input_handler:
+                message, need_return = handler(message, self)
+                if need_return:
+                    return message
 
         if self.session.addon_prompt:
             message += f"\n\n{self.session.addon_prompt}"
@@ -478,12 +489,20 @@ class Agent:
         try:
             set_agent(self.name, self)
 
+            run_input_handlers = True
             while True:
+                if self.run_input_handlers_next_turn:
+                    run_input_handlers = True
+                    self.run_input_handlers_next_turn = False
+
                 if self.first:
                     self._first_run()
                 try:
-                    current_response = self._call_model(self.session.prompt, True)
+                    current_response = self._call_model(
+                        self.session.prompt, True, run_input_handlers
+                    )
                     self.session.prompt = ""
+                    run_input_handlers = False
 
                     if get_interrupt():
                         set_interrupt(False)
@@ -491,6 +510,7 @@ class Agent:
                             f"模型交互期间被中断，请输入用户干预信息："
                         )
                         if user_input:
+                            run_input_handlers = True
                             # 如果有工具调用且用户确认继续，则将干预信息和工具执行结果拼接为prompt
                             if any(
                                 handler.can_handle(current_response)
@@ -502,6 +522,7 @@ class Agent:
                                     self.session.prompt = (
                                         f"{user_input}\n\n{current_response}"
                                     )
+                                    run_input_handlers = False
                                     continue
                             self.session.prompt += f"{user_input}"
                             continue
@@ -529,6 +550,7 @@ class Agent:
 
                     if user_input:
                         self.session.prompt = user_input
+                        run_input_handlers = True
                         continue
 
                     if not user_input:

jarvis/jarvis_agent/jarvis.py

@@ -1,8 +1,11 @@
 # -*- coding: utf-8 -*-
 import argparse
 import os
+import shutil
+import subprocess
 import sys
-from typing import Dict
+from pathlib import Path
+from typing import Dict, Optional
 
 import yaml  # type: ignore
 from prompt_toolkit import prompt  # type: ignore
@@ -106,8 +109,8 @@ def _select_task(tasks: Dict[str, str]) -> str:
             PrettyOutput.print(f"选择任务失败: {str(val_err)}", OutputType.ERROR)
 
 
-def main() -> None:
-
+def _parse_args() -> argparse.Namespace:
+    """Parse command line arguments."""
     parser = argparse.ArgumentParser(description="Jarvis AI assistant")
     parser.add_argument(
         "--llm_type",
@@ -122,6 +125,11 @@ def main() -> None:
         type=str,
         help="Directly input task content from command line",
     )
+    parser.add_argument(
+        "--model_group",
+        type=str,
+        help="Model group to use, overriding config",
+    )
     parser.add_argument("-f", "--config", type=str, help="Path to custom config file")
     parser.add_argument(
         "--restore-session",
@@ -129,44 +137,91 @@ def main() -> None:
         help="Restore session from .jarvis/saved_session.json",
         default=False,
     )
-    args = parser.parse_args()
-    init_env(
-        "欢迎使用 Jarvis AI 助手，您的智能助理已准备就绪！", config_file=args.config
+    parser.add_argument(
+        "-e", "--edit", action="store_true", help="Edit the configuration file"
     )
+    return parser.parse_args()
 
-    try:
-        agent = Agent(
-            system_prompt=origin_agent_system_prompt,
-            llm_type=args.llm_type,
-            input_handler=[shell_input_handler, builtin_input_handler],
-            output_handler=[ToolRegistry()],
-            need_summary=False,
-        )
 
-        # 尝试恢复会话
-        if args.restore_session:
-            if agent.restore_session():
-                PrettyOutput.print("会话已成功恢复。", OutputType.SUCCESS)
-            else:
-                PrettyOutput.print("无法恢复会话。", OutputType.WARNING)
+def _handle_edit_mode(args: argparse.Namespace) -> None:
+    """If edit flag is set, open config file in editor and exit."""
+    if not args.edit:
+        return
+
+    config_file_path = (
+        Path(args.config)
+        if args.config
+        else Path(os.path.expanduser("~/.jarvis/config.yaml"))
+    )
+    editors = ["nvim", "vim", "vi"]
+    editor = next((e for e in editors if shutil.which(e)), None)
 
-        # 优先处理命令行直接传入的任务
-        if args.task:
-            agent.run(args.task)
+    if editor:
+        try:
+            subprocess.run([editor, str(config_file_path)], check=True)
             sys.exit(0)
+        except (subprocess.CalledProcessError, FileNotFoundError) as e:
+            PrettyOutput.print(f"Failed to open editor: {e}", OutputType.ERROR)
+            sys.exit(1)
+    else:
+        PrettyOutput.print(
+            "No suitable editor found (nvim, vim, vi).", OutputType.ERROR
+        )
+        sys.exit(1)
 
-        if agent.first:
-            tasks = _load_tasks()
-            if tasks and (selected_task := _select_task(tasks)):
-                PrettyOutput.print(f"开始执行任务: \n{selected_task}", OutputType.INFO)
-                agent.run(selected_task)
-                sys.exit(0)
 
-        user_input = get_multiline_input("请输入你的任务（输入空行退出）:")
-        if user_input:
-            agent.run(user_input)
+def _initialize_agent(args: argparse.Namespace) -> Agent:
+    """Initialize the agent and restore session if requested."""
+    agent = Agent(
+        system_prompt=origin_agent_system_prompt,
+        llm_type=args.llm_type,
+        model_group=args.model_group,
+        input_handler=[shell_input_handler, builtin_input_handler],
+        output_handler=[ToolRegistry()],  # type: ignore
+        need_summary=False,
+    )
+
+    # 尝试恢复会话
+    if args.restore_session:
+        if agent.restore_session():
+            PrettyOutput.print("会话已成功恢复。", OutputType.SUCCESS)
+        else:
+            PrettyOutput.print("无法恢复会话。", OutputType.WARNING)
+    return agent
+
+
+def _get_and_run_task(agent: Agent, task_content: Optional[str] = None) -> None:
+    """Get task from various sources and run it."""
+    # 优先处理命令行直接传入的任务
+    if task_content:
+        agent.run(task_content)
         sys.exit(0)
 
+    if agent.first:
+        tasks = _load_tasks()
+        if tasks and (selected_task := _select_task(tasks)):
+            PrettyOutput.print(f"开始执行任务: \n{selected_task}", OutputType.INFO)
+            agent.run(selected_task)
+            sys.exit(0)
+
+    user_input = get_multiline_input("请输入你的任务（输入空行退出）:")
+    if user_input:
+        agent.run(user_input)
+    sys.exit(0)
+
+
+def main() -> None:
+    """Main function for Jarvis AI assistant."""
+    args = _parse_args()
+    _handle_edit_mode(args)
+
+    init_env(
+        "欢迎使用 Jarvis AI 助手，您的智能助理已准备就绪！", config_file=args.config
+    )
+
+    try:
+        agent = _initialize_agent(args)
+        _get_and_run_task(agent, args.task)
     except Exception as err:  # pylint: disable=broad-except
         PrettyOutput.print(f"初始化错误: {str(err)}", OutputType.ERROR)
         sys.exit(1)

jarvis/jarvis_agent/main.py

@@ -20,7 +20,9 @@ def load_config(config_path: str) -> dict:
         dict: 配置字典
     """
     if not os.path.exists(config_path):
-        PrettyOutput.print(f"配置文件 {config_path} 不存在，使用默认配置", OutputType.WARNING)
+        PrettyOutput.print(
+            f"配置文件 {config_path} 不存在，使用默认配置", OutputType.WARNING
+        )
         return {}
 
     with open(config_path, "r", encoding="utf-8", errors="ignore") as f:
@@ -50,10 +52,17 @@ def main():
         choices=["normal", "thinking"],
         help="LLM type to use, overriding config",
     )
+    parser.add_argument(
+        "--model_group",
+        type=str,
+        help="Model group to use, overriding config",
+    )
     args = parser.parse_args()
 
     # Initialize environment
-    init_env("欢迎使用 Jarvis AI 助手，您的智能助理已准备就绪！", config_file=args.config)
+    init_env(
+        "欢迎使用 Jarvis AI 助手，您的智能助理已准备就绪！", config_file=args.config
+    )
 
     # Load configuration
     config = load_config(args.agent_definition) if args.agent_definition else {}
@@ -61,6 +70,8 @@ def main():
     # Override config with command-line arguments if provided
     if args.llm_type:
         config["llm_type"] = args.llm_type
+    if args.model_group:
+        config["model_group"] = args.model_group
 
     # Create and run agent
     try:

jarvis/jarvis_code_agent/code_agent.py

@@ -18,7 +18,10 @@ from jarvis.jarvis_code_agent.lint import get_lint_tools
 from jarvis.jarvis_git_utils.git_commiter import GitCommitTool
 from jarvis.jarvis_platform.registry import PlatformRegistry
 from jarvis.jarvis_tools.registry import ToolRegistry
-from jarvis.jarvis_utils.config import is_confirm_before_apply_patch, is_enable_static_analysis
+from jarvis.jarvis_utils.config import (
+    is_confirm_before_apply_patch,
+    is_enable_static_analysis,
+)
 from jarvis.jarvis_utils.git_utils import (
     confirm_add_new_files,
     find_git_root_and_cd,
@@ -44,6 +47,7 @@ class CodeAgent:
     def __init__(
         self,
         llm_type: str = "normal",
+        model_group: Optional[str] = None,
         need_summary: bool = True,
     ):
         self.root_dir = os.getcwd()
@@ -54,7 +58,6 @@ class CodeAgent:
                 "search_web",
                 "ask_user",
                 "read_code",
-                "methodology",
                 "rewrite_file",
             ]
         )
@@ -116,8 +119,9 @@ class CodeAgent:
             system_prompt=code_system_prompt,
             name="CodeAgent",
             auto_complete=False,
-            output_handler=[tool_registry, EditFileHandler()],
+            output_handler=[tool_registry, EditFileHandler()],  # type: ignore
             llm_type=llm_type,
+            model_group=model_group,
             input_handler=[shell_input_handler, builtin_input_handler],
             need_summary=need_summary,
             use_methodology=False,  # 禁用方法论
@@ -409,6 +413,11 @@ def main() -> None:
         choices=["normal", "thinking"],
         help="LLM type to use",
     )
+    parser.add_argument(
+        "--model_group",
+        type=str,
+        help="Model group to use, overriding config",
+    )
     parser.add_argument(
         "-r", "--requirement", type=str, help="Requirement to process", default=None
     )
@@ -425,7 +434,11 @@ def main() -> None:
     PrettyOutput.print(f"当前目录: {git_dir}", OutputType.INFO)
 
     try:
-        agent = CodeAgent(llm_type=args.llm_type, need_summary=False)
+        agent = CodeAgent(
+            llm_type=args.llm_type,
+            model_group=args.model_group,
+            need_summary=False,
+        )
 
         # 尝试恢复会话
         if args.restore_session:

jarvis/jarvis_code_analysis/checklists/loader.py

@@ -5,12 +5,26 @@ Utility module for loading language-specific code review checklists.
 from typing import Dict, Optional
 
 # Import checklist modules
-from jarvis.jarvis_code_analysis.checklists import (c_cpp, csharp, data_format,
-                                                    devops, docs, go,
-                                                    infrastructure, java,
-                                                    javascript, kotlin, php,
-                                                    python, ruby, rust, shell,
-                                                    sql, swift, web)
+from jarvis.jarvis_code_analysis.checklists import (
+    c_cpp,
+    csharp,
+    data_format,
+    devops,
+    docs,
+    go,
+    infrastructure,
+    java,
+    javascript,
+    kotlin,
+    php,
+    python,
+    ruby,
+    rust,
+    shell,
+    sql,
+    swift,
+    web,
+)
 
 # Map of language identifiers to their checklist content
 CHECKLIST_MAP = {

jarvis/jarvis_code_analysis/code_review.py

@@ -3,12 +3,13 @@ import os
 import re
 import subprocess
 import tempfile
-from typing import Any, Dict, List
+from typing import Any, Dict, List, Optional
 
 from jarvis.jarvis_agent import Agent
 from jarvis.jarvis_code_analysis.checklists.loader import get_language_checklist
 from jarvis.jarvis_platform.registry import PlatformRegistry
 from jarvis.jarvis_tools.read_code import ReadCodeTool
+from jarvis.jarvis_utils.globals import get_agent, current_agent_name
 from jarvis.jarvis_utils.output import OutputType, PrettyOutput
 from jarvis.jarvis_utils.tag import ct, ot
 from jarvis.jarvis_utils.utils import init_env, is_context_overflow
@@ -261,7 +262,9 @@ class CodeReviewTool:
         checklist = get_language_checklist(language)
         return checklist if checklist else ""
 
-    def execute(self, args: Dict[str, Any]) -> Dict[str, Any]:
+    def execute(
+        self, args: Dict[str, Any], agent: Optional["Agent"] = None
+    ) -> Dict[str, Any]:
         try:
             review_type = args.get("review_type", "current").strip()
             root_dir = args.get("root_dir", ".")
@@ -570,9 +573,17 @@ class CodeReviewTool:
 
                 tool_registry = ToolRegistry()
                 tool_registry.dont_use_tools(["code_review"])
+                
+                # Use the provided agent's model_group or get it from globals
+                calling_agent = agent or get_agent(current_agent_name)
+                model_group = None
+                if calling_agent and hasattr(calling_agent, "model") and calling_agent.model:
+                    model_group = calling_agent.model.model_group
+
                 agent = Agent(
                     system_prompt=system_prompt,
                     name="Code Review Agent",
+                    model_group=model_group,
                     summary_prompt=f"""<code_review_report>
 <overview>
 # 整体评估
@@ -675,7 +686,7 @@ class CodeReviewTool:
 
                 try:
                     # Check if content is too large
-                    is_large_content = is_context_overflow(diff_output)
+                    is_large_content = is_context_overflow(diff_output, model_group)
 
                     # Upload the file to the agent's model
                     if is_large_content:

jarvis/jarvis_data/config_schema.json

@@ -141,6 +141,47 @@
       "description": "思考操作模型名称",
       "default": "deep_seek"
     },
+    "JARVIS_MODEL_GROUP": {
+      "type": "string",
+      "description": "选择一个预定义的模型组"
+    },
+    "JARVIS_MODEL_GROUPS": {
+      "type": "array",
+      "description": "预定义的模型配置组",
+      "items": {
+        "type": "object",
+        "additionalProperties": {
+          "type": "object",
+          "properties": {
+            "JARVIS_PLATFORM": {
+              "type": "string"
+            },
+            "JARVIS_MODEL": {
+              "type": "string"
+            },
+            "JARVIS_THINKING_PLATFORM": {
+              "type": "string"
+            },
+            "JARVIS_THINKING_MODEL": {
+              "type": "string"
+            },
+            "JARVIS_MAX_TOKEN_COUNT": {
+              "type": "number"
+            },
+            "JARVIS_MAX_INPUT_TOKEN_COUNT": {
+              "type": "number"
+            },
+            "JARVIS_MAX_BIG_CONTENT_SIZE": {
+              "type": "number"
+            }
+          },
+          "required": [
+            "JARVIS_PLATFORM",
+            "JARVIS_MODEL"
+          ]
+        }
+      }
+    },
     "JARVIS_EXECUTE_TOOL_CONFIRM": {
       "type": "boolean",
       "description": "执行工具前是否需要确认",

jarvis/jarvis_git_utils/git_commiter.py

@@ -16,6 +16,7 @@ from jarvis.jarvis_utils.git_utils import (
     find_git_root_and_cd,
     has_uncommitted_changes,
 )
+from jarvis.jarvis_utils.globals import get_agent, current_agent_name
 from jarvis.jarvis_utils.output import OutputType, PrettyOutput
 from jarvis.jarvis_utils.tag import ct, ot
 from jarvis.jarvis_utils.utils import init_env, is_context_overflow
@@ -162,11 +163,17 @@ commit信息
                 """
 
                 # 获取模型并尝试上传文件
-                platform = PlatformRegistry().get_normal_platform()
+                agent = get_agent(current_agent_name)
+                if agent:
+                    platform = agent.model
+                    model_group = agent.model.model_group
+                else:
+                    platform = PlatformRegistry().get_normal_platform()
+                    model_group = None
                 upload_success = False
 
                 # Check if content is too large
-                is_large_content = is_context_overflow(diff)
+                is_large_content = is_context_overflow(diff, model_group)
 
                 if is_large_content:
                     if not platform.support_upload_files():

jarvis/jarvis_mcp/sse_mcp_client.py

@@ -2,7 +2,7 @@
 import json
 import threading
 import time
-from typing import Any, Callable, Dict, List
+from typing import Any, Callable, Dict, List, Optional
 from urllib.parse import parse_qs, urlencode, urljoin
 
 import requests
@@ -46,10 +46,10 @@ class SSEMcpClient(McpClient):
         self.session.headers.update(extra_headers)
 
         # SSE相关属性
-        self.sse_response = None
-        self.sse_thread = None
-        self.messages_endpoint = None
-        self.session_id = None  # 从SSE连接获取的会话ID
+        self.sse_response: Optional[requests.Response] = None
+        self.sse_thread: Optional[threading.Thread] = None
+        self.messages_endpoint: Optional[str] = None
+        self.session_id: Optional[str] = None
         self.pending_requests = {}  # 存储等待响应的请求 {id: Event}
         self.request_results = {}  # 存储请求结果 {id: result}
         self.notification_handlers = {}
@@ -123,13 +123,15 @@ class SSEMcpClient(McpClient):
             self.sse_response = self.session.get(
                 sse_url, stream=True, headers=sse_headers, timeout=30
             )
-            self.sse_response.raise_for_status()
+            if self.sse_response:
+                self.sse_response.raise_for_status()
 
             # 启动事件处理线程
             self.sse_thread = threading.Thread(
                 target=self._process_sse_events, daemon=True
             )
-            self.sse_thread.start()
+            if self.sse_thread:
+                self.sse_thread.start()
 
         except Exception as e:
             PrettyOutput.print(f"SSE连接失败: {str(e)}", OutputType.ERROR)

jarvis/jarvis_mcp/stdio_mcp_client.py

@@ -2,7 +2,7 @@
 import json
 import os
 import subprocess
-from typing import Any, Dict, List
+from typing import Any, Dict, List, Optional
 
 from jarvis.jarvis_mcp import McpClient
 from jarvis.jarvis_utils.output import OutputType, PrettyOutput
@@ -17,7 +17,7 @@ class StdioMcpClient(McpClient):
 
     def __init__(self, config: Dict[str, Any]):
         self.config = config
-        self.process = None
+        self.process: Optional[subprocess.Popen] = None
         self.protocol_version = "2025-03-26"  # MCP协议版本
         self._start_process()
         self._initialize()