auto-coder 0.1.347__py3-none-any.whl → 0.1.349__py3-none-any.whl

autocoder/common/model_speed_tester.py

@@ -0,0 +1,392 @@
+import time
+import byzerllm
+from typing import Dict, Any, List, Optional
+from rich.console import Console
+from rich.table import Table
+from rich.panel import Panel
+from autocoder.common.printer import Printer
+from autocoder import models as models_module
+from autocoder.utils.llms import get_single_llm
+import byzerllm
+import pkg_resources
+from concurrent.futures import ThreadPoolExecutor
+from typing import Dict, List, Tuple
+from pydantic import BaseModel
+
+class ModelSpeedTestResult(BaseModel):
+    model_name: str
+    tokens_per_second: float
+    first_token_time: float
+    input_tokens_count: float
+    generated_tokens_count: float
+    input_tokens_cost: float
+    generated_tokens_cost: float
+    status: str
+    error: Optional[str] = None
+
+class SpeedTestResults(BaseModel):
+    results: List[ModelSpeedTestResult]
+
+byzerllm_content = ""
+try:
+    byzerllm_conten_path = pkg_resources.resource_filename(
+        "autocoder", "data/byzerllm.md"
+    )
+    with open(byzerllm_conten_path, "r",encoding="utf-8") as f:
+        byzerllm_content = f.read()
+except FileNotFoundError:
+    pass
+
+@byzerllm.prompt()
+def long_context_prompt() -> str:
+    '''
+    下面是我们提供的一份文档：
+    <document>
+    {{ content }}
+    </document>
+    
+    请根据上述文档，实现用户的需求：
+
+    <query>
+    我想开发一个翻译程序，使用prompt 函数实现。
+    </query>
+    '''
+    return {
+        "content": byzerllm_content
+    }
+
+@byzerllm.prompt()
+def short_context_prompt() -> str:
+    '''
+    Hello, can you help me test the response speed?
+    '''
+    return {}
+
+def test_model_speed(model_name: str, 
+                    product_mode: str, 
+                    test_rounds: int = 3,
+                    enable_long_context: bool = False
+                    ) -> Dict[str, Any]:
+    from autocoder.models import get_model_by_name
+    """
+    测试单个模型的速度
+    
+    Args:
+        model_name: 模型名称
+        product_mode: 产品模式 (lite/pro)
+        test_rounds: 测试轮数
+        
+    Returns:
+        Dict包含测试结果:
+            - avg_time: 平均响应时间
+            - min_time: 最小响应时间
+            - max_time: 最大响应时间
+            - first_token_time: 首token时间
+            - success: 是否测试成功
+            - error: 错误信息(如果有)
+    """
+    try:
+        llm = get_single_llm(model_name, product_mode)
+        model_info = get_model_by_name(model_name)
+
+        times = []
+        first_token_times = []
+        tokens_per_seconds = []
+        input_tokens_counts = []
+        generated_tokens_counts = []
+        
+        input_tokens_costs = []
+        generated_tokens_costs = []
+
+        input_tokens_cost_per_m = model_info.get("input_price", 0.0) / 1000000
+        output_tokens_cost_per_m = model_info.get("output_price", 0.0) / 1000000
+
+        test_query = short_context_prompt.prompt()
+        if enable_long_context:
+            test_query = long_context_prompt.prompt()                
+        
+        content = ""
+        for _ in range(test_rounds):
+            start_time = time.time()
+            first_token_received = False
+            first_token_time = None
+            last_meta = None
+            input_tokens_count = 0
+            generated_tokens_count = 0 
+            input_tokens_cost = 0
+            generated_tokens_cost = 0           
+            for chunk,meta in llm.stream_chat_oai(conversations=[{
+                "role": "user",
+                "content": test_query
+            }],delta_mode=True):
+                content += chunk
+                last_meta = meta                
+                current_time = time.time()
+                if not first_token_received:
+                    first_token_time = current_time - start_time
+                    first_token_received = True
+                    first_token_times.append(first_token_time)
+            
+            end_time = time.time()            
+            generated_tokens_count = 0
+            if last_meta:                
+                generated_tokens_count = last_meta.generated_tokens_count
+                input_tokens_count = last_meta.input_tokens_count
+                input_tokens_cost = input_tokens_count * input_tokens_cost_per_m
+                generated_tokens_cost = generated_tokens_count * output_tokens_cost_per_m
+
+                input_tokens_costs.append(input_tokens_cost)
+                generated_tokens_costs.append(generated_tokens_cost)
+                generated_tokens_counts.append(generated_tokens_count)
+                input_tokens_counts.append(input_tokens_count)
+            
+            tokens_per_seconds.append(generated_tokens_count / (end_time - start_time))    
+            times.append(end_time - start_time)
+            
+            
+        avg_time = sum(times) / len(times)
+        return {
+            "tokens_per_second": sum(tokens_per_seconds) / len(tokens_per_seconds),
+            "avg_time": avg_time,
+            "min_time": min(times),
+            "max_time": max(times),
+            "first_token_time": sum(first_token_times) / len(first_token_times),
+            "input_tokens_count": sum(input_tokens_counts) / len(input_tokens_counts),
+            "generated_tokens_count": sum(generated_tokens_counts) / len(generated_tokens_counts),
+            "success": True,
+            "error": None,
+            "input_tokens_cost": sum(input_tokens_costs) / len(input_tokens_costs),
+            "generated_tokens_cost": sum(generated_tokens_costs) / len(generated_tokens_costs)
+        }
+    except Exception as e:
+        return {
+            "tokens_per_second": 0,
+            "avg_time": 0,
+            "min_time": 0,
+            "max_time": 0,
+            "first_token_time": 0,
+            "input_tokens_count": 0,
+            "generated_tokens_count": 0,
+            "success": False,
+            "error": str(e),
+            "input_tokens_cost": 0.0,
+            "generated_tokens_cost": 0.0
+        }
+
+def test_model_speed_wrapper(args: Tuple[str, str, int, bool]) -> Tuple[str, Dict[str, Any]]:
+    """
+    包装测试函数以适应线程池调用
+    
+    Args:
+        args: (model_name, product_mode, test_rounds)的元组
+        
+    Returns:
+        (model_name, test_results)的元组
+    """
+    model_name, product_mode, test_rounds,enable_long_context = args
+    results = test_model_speed(model_name, product_mode, test_rounds,enable_long_context)
+    return (model_name, results)
+
+
+def run_speed_test(product_mode: str, test_rounds: int = 3, max_workers: Optional[int] = None, enable_long_context: bool = False) -> SpeedTestResults:
+    """
+    运行所有已激活模型的速度测试
+    
+    Args:
+        product_mode: 产品模式 (lite/pro)
+        test_rounds: 每个模型测试的轮数
+        max_workers: 最大线程数,默认为None(ThreadPoolExecutor会自动设置)
+        enable_long_context: 是否启用长文本上下文测试
+    
+    Returns:
+        SpeedTestResults: 包含所有模型测试结果的pydantic模型
+    """
+    # 获取所有模型
+    models_data = models_module.load_models()
+    active_models = [m for m in models_data if "api_key" in m] if product_mode == "lite" else models_data
+    
+    if not active_models:
+        return SpeedTestResults(results=[])
+    
+    # 准备测试参数
+    test_args = [(model["name"], product_mode, test_rounds, enable_long_context) for model in active_models]
+    
+    # 存储结果用于排序
+    results_list = []
+    
+    # 使用线程池并发测试
+    with ThreadPoolExecutor(max_workers=max_workers) as executor:
+        # 提交所有测试任务并获取future对象
+        future_to_model = {executor.submit(test_model_speed_wrapper, args): args[0] 
+                          for args in test_args}
+        
+        # 收集结果
+        for future in future_to_model:
+            model_name = future_to_model[future]
+            
+            try:
+                _, results = future.result()
+                
+                if results["success"]:
+                    status = "✓"
+                    results_list.append((
+                        results['tokens_per_second'],
+                        ModelSpeedTestResult(
+                            model_name=model_name,
+                            tokens_per_second=results['tokens_per_second'],
+                            first_token_time=results['first_token_time'],
+                            input_tokens_count=results['input_tokens_count'],
+                            generated_tokens_count=results['generated_tokens_count'],
+                            status=status,
+                            input_tokens_cost=results['input_tokens_cost'],
+                            generated_tokens_cost=results['generated_tokens_cost'],                            
+                        )
+                    ))
+                    try:
+                        # 更新模型的平均速度
+                        models_module.update_model_speed(model_name, results['tokens_per_second'])
+                    except Exception:
+                        pass
+                else:                    
+                    results_list.append((
+                        0,
+                        ModelSpeedTestResult(
+                            model_name=model_name,
+                            tokens_per_second=0,
+                            first_token_time=0,
+                            input_tokens_count=0,
+                            generated_tokens_count=0,
+                            status=f"✗ {results['error']}",
+                            error=results['error'],
+                            input_tokens_cost=0.0,
+                            generated_tokens_cost=0.0
+                        )
+                    ))
+            except Exception as e:
+                results_list.append((
+                    0,
+                    ModelSpeedTestResult(
+                        model_name=model_name,
+                        tokens_per_second=0,
+                        first_token_time=0,
+                        input_tokens_count=0,
+                        generated_tokens_count=0,
+                        status=f"✗ {str(e)}",
+                        error=str(e),
+                        input_tokens_cost=0.0,
+                        generated_tokens_cost=0.0
+                    )
+                ))
+
+    # 按速度排序
+    results_list.sort(key=lambda x: x[0], reverse=True)
+    
+    return SpeedTestResults(results=[result[1] for result in results_list])
+
+def render_speed_test_in_terminal(product_mode: str, test_rounds: int = 3, max_workers: Optional[int] = None,enable_long_context: bool = False) -> None:
+    """
+    运行所有已激活模型的速度测试
+    
+    Args:
+        product_mode: 产品模式 (lite/pro)
+        test_rounds: 每个模型测试的轮数
+        max_workers: 最大线程数,默认为None(ThreadPoolExecutor会自动设置)
+    """
+    printer = Printer()
+    console = Console()
+    
+    # 获取所有模型
+    models_data = models_module.load_models()
+    active_models = [m for m in models_data if "api_key" in m] if product_mode == "lite" else models_data
+    
+    if not active_models:
+        printer.print_in_terminal("models_no_active", style="yellow")
+        return
+        
+    # 创建结果表格
+    table = Table(
+        title=printer.get_message_from_key("models_speed_test_results"),
+        show_header=True,
+        header_style="bold magenta",
+        show_lines=True
+    )
+    
+    table.add_column("Model", style="cyan", width=30)
+    table.add_column("Tokens/s", style="green", width=15)
+    table.add_column("First Token(s)", style="magenta", width=15)
+    table.add_column("Input Tokens", style="magenta", width=15)
+    table.add_column("Generated Tokens", style="magenta", width=15)            
+    table.add_column("Input Tokens Cost", style="yellow", width=15)
+    table.add_column("Generated Tokens Cost", style="yellow", width=15)
+    table.add_column("Status", style="red", width=20)
+    
+    # 准备测试参数
+    test_args = [(model["name"], product_mode, test_rounds, enable_long_context) for model in active_models]
+    
+    # 存储结果用于排序
+    results_list = []
+    
+    # 使用线程池并发测试
+    with ThreadPoolExecutor(max_workers=max_workers) as executor:
+        printer.print_in_terminal("models_testing_start", style="yellow")
+        
+        # 提交所有测试任务并获取future对象
+        future_to_model = {executor.submit(test_model_speed_wrapper, args): args[0] 
+                          for args in test_args}
+        
+        # 收集结果
+        completed = 0
+        total = len(future_to_model)
+        for future in future_to_model:
+            completed += 1
+            printer.print_in_terminal("models_testing_progress", style="yellow", completed=completed, total=total)
+            model_name = future_to_model[future]
+            printer.print_in_terminal("models_testing", style="yellow", name=model_name)
+            
+            try:
+                _, results = future.result()
+                
+                if results["success"]:
+                    status = "✓"
+                    results['status'] = status
+                    results_list.append((
+                        results['tokens_per_second'],
+                        model_name,
+                        results
+                    ))                  
+                    try:  
+                        # 更新模型的平均速度
+                        models_module.update_model_speed(model_name, results['tokens_per_second'])
+                    except Exception as e:
+                        pass
+                else:
+                    status = f"✗ ({results['error']})"
+                    results_list.append((
+                        0,
+                        model_name,
+                        {"tokens_per_second":0,"avg_time": 0, "input_tokens_count":0, "generated_tokens_count":0, "min_time": 0, "max_time": 0, "first_token_time": 0, "input_tokens_cost": 0.0, "generated_tokens_cost": 0.0, "status": status}
+                    ))
+            except Exception as e:
+                results_list.append((
+                    0,
+                    model_name,
+                        {"tokens_per_second":0,"avg_time": 0, "input_tokens_count":0, "generated_tokens_count":0, "min_time": 0, "max_time": 0, "first_token_time": 0, "input_tokens_cost": 0.0, "generated_tokens_cost": 0.0, "status": f"✗ ({str(e)})"}
+                ))
+
+    # 按速度排序
+    results_list.sort(key=lambda x: x[0], reverse=True)
+    
+    # 添加排序后的结果到表格
+    for tokens_per_second, model_name, results in results_list:        
+        table.add_row(
+            model_name,  
+            f"{tokens_per_second:.2f}",
+            f"{results['first_token_time']:.2f}",
+            f"{results['input_tokens_count']}",
+            f"{results['generated_tokens_count']}",
+            f"{results['input_tokens_cost']:.4f}",
+            f"{results['generated_tokens_cost']:.4f}",            
+            results['status']
+        )        
+    
+    console.print(Panel(table, border_style="blue"))

autocoder/common/printer.py

@@ -15,19 +15,18 @@ class Printer:
 
     def get_message_from_key(self, msg_key: str):
         try:
-            return get_message(msg_key)
+            v = get_message(msg_key)
+            if not v:
+                return get_chat_message(msg_key)
         except Exception as e:
             return get_chat_message(msg_key)
 
-    def get_message_from_key_with_format(self, msg_key: str, **kwargs):
-        try:
-            return format_str_jinja2(self.get_message_from_key(msg_key), **kwargs)
-        except Exception as e:
-            return format_str_jinja2(self.get_chat_message_from_key(msg_key), **kwargs)
-
+    def get_message_from_key_with_format(self, msg_key: str, **kwargs):        
+        return format_str_jinja2(self.get_message_from_key(msg_key), **kwargs)
+        
     def print_in_terminal(self, msg_key: str, style: str = None,**kwargs):     
         try:
-            if style:
+            if style:                
                 self.console.print(format_str_jinja2(self.get_message_from_key(msg_key),**kwargs), style=style)
             else:
                 self.console.print(format_str_jinja2(self.get_message_from_key(msg_key),**kwargs))

autocoder/common/run_cmd.py

@@ -0,0 +1,247 @@
+import os
+import platform
+import subprocess
+import sys
+from io import BytesIO
+
+import pexpect
+import psutil
+
+
+def run_cmd(command, verbose=False, error_print=None, cwd=None):
+    """
+    执行一条命令，根据不同系统和环境选择最合适的执行方式（交互式或非交互式）。
+
+    适用场景：
+    - 需要跨平台运行命令，自动判断是否使用pexpect（支持交互式）还是subprocess。
+    - 希望获得命令的执行状态及输出内容。
+    - 在CLI工具、自动化脚本、REPL中执行shell命令。
+
+    参数：
+    - command (str): 需要执行的命令字符串。
+    - verbose (bool): 是否打印详细调试信息，默认为False。
+    - error_print (callable|None): 自定义错误打印函数，默认为None，使用print。
+    - cwd (str|None): 指定命令的工作目录，默认为None。
+
+    返回：
+    - tuple: (exit_code, output)，其中exit_code为整型退出码，output为命令输出内容。
+
+    异常：
+    - 捕获OSError异常，返回错误信息。
+    """
+    try:
+        if sys.stdin.isatty() and hasattr(pexpect, "spawn") and platform.system() != "Windows":
+            return run_cmd_pexpect(command, verbose, cwd)
+
+        return run_cmd_subprocess(command, verbose, cwd)
+    except OSError as e:
+        error_message = f"Error occurred while running command '{command}': {str(e)}"
+        if error_print is None:
+            print(error_message)
+        else:
+            error_print(error_message)
+        return 1, error_message
+
+
+def get_windows_parent_process_name():
+    """
+    获取当前进程的父进程名（仅在Windows系统下有意义）。
+
+    适用场景：
+    - 判断命令是否由PowerShell或cmd.exe启动，以便调整命令格式。
+    - 在Windows平台上进行父进程分析。
+
+    参数：
+    - 无
+
+    返回：
+    - str|None: 父进程名（小写字符串，如"powershell.exe"或"cmd.exe"），如果无法获取则为None。
+
+    异常：
+    - 捕获所有异常，返回None。
+    """
+    try:
+        current_process = psutil.Process()
+        while True:
+            parent = current_process.parent()
+            if parent is None:
+                break
+            parent_name = parent.name().lower()
+            if parent_name in ["powershell.exe", "cmd.exe"]:
+                return parent_name
+            current_process = parent
+        return None
+    except Exception:
+        return None
+
+
+def run_cmd_subprocess(command, verbose=False, cwd=None, encoding=sys.stdout.encoding):
+    if verbose:
+        print("Using run_cmd_subprocess:", command)
+
+    try:
+        shell = os.environ.get("SHELL", "/bin/sh")
+        parent_process = None
+
+        # Determine the appropriate shell
+        if platform.system() == "Windows":
+            parent_process = get_windows_parent_process_name()
+            if parent_process == "powershell.exe":
+                command = f"powershell -Command {command}"
+
+        if verbose:
+            print("Running command:", command)
+            print("SHELL:", shell)
+            if platform.system() == "Windows":
+                print("Parent process:", parent_process)
+
+        process = subprocess.Popen(
+            command,
+            stdout=subprocess.PIPE,
+            stderr=subprocess.STDOUT,
+            text=True,
+            shell=True,
+            encoding=encoding,
+            errors="replace",
+            bufsize=0,  # Set bufsize to 0 for unbuffered output
+            universal_newlines=True,
+            cwd=cwd,
+        )
+
+        output = []
+        while True:
+            chunk = process.stdout.read(1)
+            if not chunk:
+                break
+            print(chunk, end="", flush=True)  # Print the chunk in real-time
+            output.append(chunk)  # Store the chunk for later use
+
+        process.wait()
+        return process.returncode, "".join(output)
+    except Exception as e:
+        return 1, str(e)
+
+def run_cmd_subprocess_generator(command, verbose=False, cwd=None, encoding=sys.stdout.encoding):
+    """
+    使用subprocess运行命令，将命令输出逐步以生成器方式yield出来。
+
+    适用场景：
+    - 运行无需交互的命令。
+    - 希望实时逐步处理命令输出（如日志打印、进度监控）。
+    - 在Linux、macOS、Windows等多平台环境下安全运行命令。
+
+    参数：
+    - command (str): 需要执行的命令字符串。
+    - verbose (bool): 是否打印详细调试信息，默认为False。
+    - cwd (str|None): 指定命令的工作目录，默认为None。
+    - encoding (str): 输出解码使用的字符编码，默认为当前stdout编码。
+
+    返回：
+    - 生成器: 逐块yield命令的输出字符串。
+
+    异常：
+    - 捕获所有异常，yield错误信息字符串。
+    """
+    if verbose:
+        print("Using run_cmd_subprocess:", command)
+
+    try:
+        shell = os.environ.get("SHELL", "/bin/sh")
+        parent_process = None
+
+        # Windows下调整命令
+        if platform.system() == "Windows":
+            parent_process = get_windows_parent_process_name()
+            if parent_process == "powershell.exe":
+                command = f"powershell -Command {command}"
+
+        if verbose:
+            print("Running command:", command)
+            print("SHELL:", shell)
+            if platform.system() == "Windows":
+                print("Parent process:", parent_process)
+
+        process = subprocess.Popen(
+            command,
+            stdout=subprocess.PIPE,
+            stderr=subprocess.STDOUT,
+            text=True,
+            shell=True,
+            encoding=encoding,
+            errors="replace",
+            bufsize=0,
+            universal_newlines=True,
+            cwd=cwd,
+        )
+
+        while True:
+            chunk = process.stdout.read(1)
+            if not chunk:
+                break            
+            # 确保始终yield字符串，避免因字节或其他类型导致异常
+            if not isinstance(chunk, str):
+                chunk = str(chunk)
+            yield chunk
+
+        process.wait()
+    except Exception as e:
+        # 出错时yield异常信息，也可以raise
+        yield f"[run_cmd_subprocess error]: {str(e)}"
+
+
+def run_cmd_pexpect(command, verbose=False, cwd=None):
+    """
+    使用pexpect以交互方式运行命令，捕获完整输出。
+
+    适用场景：
+    - 执行需要用户交互的命令（如登录、密码输入等）。
+    - 在Linux、macOS等Unix系统下模拟终端操作。
+    - 希望完整捕获交互式命令的输出。
+
+    参数：
+    - command (str): 需要执行的命令字符串。
+    - verbose (bool): 是否打印详细调试信息，默认为False。
+    - cwd (str|None): 指定命令的工作目录，默认为None。
+
+    返回：
+    - tuple: (exit_code, output)，exit_code为退出状态码，output为命令完整输出内容。
+
+    异常：
+    - 捕获pexpect相关异常，返回错误信息。
+    """
+    if verbose:
+        print("Using run_cmd_pexpect:", command)
+
+    output = BytesIO()
+
+    def output_callback(b):
+        output.write(b)
+        return b
+
+    try:
+        # Use the SHELL environment variable, falling back to /bin/sh if not set
+        shell = os.environ.get("SHELL", "/bin/sh")
+        if verbose:
+            print("With shell:", shell)
+
+        if os.path.exists(shell):
+            # Use the shell from SHELL environment variable
+            if verbose:
+                print("Running pexpect.spawn with shell:", shell)
+            child = pexpect.spawn(shell, args=["-i", "-c", command], encoding="utf-8", cwd=cwd)
+        else:
+            # Fall back to spawning the command directly
+            if verbose:
+                print("Running pexpect.spawn without shell.")
+            child = pexpect.spawn(command, encoding="utf-8", cwd=cwd)
+
+        # Transfer control to the user, capturing output
+        child.interact(output_filter=output_callback)
+
+        # Wait for the command to finish and get the exit status
+        child.close()
+        return child.exitstatus, output.getvalue().decode("utf-8", errors="replace")
+
+    except (pexpect.ExceptionPexpect, TypeError, ValueError) as e:
+        error_msg = f"Error running command {command}: {e}"
+        return 1, error_msg