auto-coder 0.1.251__py3-none-any.whl → 0.1.253__py3-none-any.whl

autocoder/index/index.py

@@ -22,7 +22,8 @@ from autocoder.index.types import (
     TargetFile,
     FileList,
 )
-
+from autocoder.common.global_cancel import global_cancel
+from autocoder.utils.llms import get_llm_names
 class IndexManager:
     def __init__(
         self, llm: byzerllm.ByzerLLM, sources: List[SourceCode], args: AutoCoderArgs
@@ -195,7 +196,10 @@ class IndexManager:
             return True
         return False
 
-    def build_index_for_single_source(self, source: SourceCode):        
+    def build_index_for_single_source(self, source: SourceCode):   
+        if global_cancel.requested:
+            return None
+
         file_path = source.module_name
         if not os.path.exists(file_path):
             return None
@@ -205,9 +209,7 @@ class IndexManager:
 
         md5 = hashlib.md5(source.source_code.encode("utf-8")).hexdigest()
 
-        model_name = getattr(self.index_llm, 'default_model_name', None)
-        if not model_name:
-            model_name = "unknown(without default model name)"
+        model_name = ",".join(get_llm_names(self.index_llm))
 
         try:
             start_time = time.monotonic()
@@ -314,6 +316,9 @@ class IndexManager:
                 ):
                     wait_to_build_files.append(source)
 
+            # Remove duplicates based on module_name
+            wait_to_build_files = list({source.module_name: source for source in wait_to_build_files}.values())
+
             counter = 0
             num_files = len(wait_to_build_files)
             total_files = len(self.sources)
@@ -329,6 +334,8 @@ class IndexManager:
                 for source in wait_to_build_files
             ]
             for future in as_completed(futures):
+                if global_cancel.requested:
+                    break
                 result = future.result()
                 if result is not None:
                     counter += 1
@@ -345,7 +352,7 @@ class IndexManager:
                         with open(self.index_file, "w") as file:
                             json.dump(index_data, file, ensure_ascii=False, indent=2)
                         updated_sources = []
-                        
+        
         # 如果 updated_sources 或 keys_to_remove 有值，则保存索引文件
         if updated_sources or keys_to_remove:
             with open(self.index_file, "w") as file:

autocoder/models.py

@@ -2,7 +2,6 @@ import os
 import json
 from typing import List, Dict
 from urllib.parse import urlparse
-from autocoder.common.auto_coder_lang import get_message_with_format
 
 MODELS_JSON = os.path.expanduser("~/.auto-coder/keys/models.json")
 
@@ -15,7 +14,10 @@ default_models_list = [
         "model_type": "saas/openai",
         "base_url": "https://api.deepseek.com/v1",
         "api_key_path": "api.deepseek.com",
-        "is_reasoning": True
+        "is_reasoning": True,
+        "input_price": 0.0,  # 单位:M/百万 input tokens
+        "output_price": 0.0,  # 单位:M/百万 output tokens
+        "average_speed": 0.0  # 单位:秒/请求
     },    
     {
         "name": "deepseek_chat",
@@ -24,7 +26,10 @@ default_models_list = [
         "model_type": "saas/openai",
         "base_url": "https://api.deepseek.com/v1",
         "api_key_path": "api.deepseek.com",
-        "is_reasoning": False
+        "is_reasoning": False,
+        "input_price": 0.0,
+        "output_price": 0.0,
+        "average_speed": 0.0
     },
     {
         "name":"o1",
@@ -33,7 +38,10 @@ default_models_list = [
         "model_type": "saas/openai",
         "base_url": "https://api.openai.com/v1",
         "api_key_path": "",
-        "is_reasoning": True
+        "is_reasoning": True,
+        "input_price": 0.0,
+        "output_price": 0.0,
+        "average_speed": 0.0
     }
 ]
 
@@ -106,6 +114,7 @@ def get_model_by_name(name: str) -> Dict:
     """
     根据模型名称查找模型
     """
+    from autocoder.common.auto_coder_lang import get_message_with_format
     models = load_models()
     v = [m for m in models if m["name"] == name.strip()]
     
@@ -114,6 +123,78 @@ def get_model_by_name(name: str) -> Dict:
     return v[0]
 
 
+def update_model_input_price(name: str, price: float) -> bool:
+    """更新模型输入价格
+    
+    Args:
+        name: 模型名称
+        price: 输入价格(M/百万input tokens)
+        
+    Returns:
+        bool: 是否更新成功
+    """
+    if price < 0:
+        raise ValueError("Price cannot be negative")
+        
+    models = load_models()
+    updated = False
+    for model in models:
+        if model["name"] == name:
+            model["input_price"] = float(price)
+            updated = True
+            break
+    if updated:
+        save_models(models)
+    return updated
+
+def update_model_output_price(name: str, price: float) -> bool:
+    """更新模型输出价格
+    
+    Args:
+        name: 模型名称
+        price: 输出价格(M/百万output tokens)
+        
+    Returns:
+        bool: 是否更新成功
+    """
+    if price < 0:
+        raise ValueError("Price cannot be negative")
+        
+    models = load_models()
+    updated = False
+    for model in models:
+        if model["name"] == name:
+            model["output_price"] = float(price)
+            updated = True
+            break
+    if updated:
+        save_models(models)
+    return updated
+
+def update_model_speed(name: str, speed: float) -> bool:
+    """更新模型平均速度
+    
+    Args:
+        name: 模型名称
+        speed: 速度(秒/请求)
+        
+    Returns:
+        bool: 是否更新成功
+    """
+    if speed <= 0:
+        raise ValueError("Speed must be positive")
+        
+    models = load_models()
+    updated = False
+    for model in models:
+        if model["name"] == name:
+            model["average_speed"] = float(speed)
+            updated = True
+            break
+    if updated:
+        save_models(models)
+    return updated
+
 def check_model_exists(name: str) -> bool:
     """
     检查模型是否存在
@@ -124,14 +205,14 @@ def check_model_exists(name: str) -> bool:
 def update_model_with_api_key(name: str, api_key: str) -> Dict:
     """
     根据模型名称查找并更新模型的 api_key_path。
-    如果找到模型，会根据其 base_url 处理 api_key_path。
+    如果找到模型,会根据其 base_url 处理 api_key_path。
     
     Args:
         name: 模型名称
         api_key: API密钥
         
     Returns:
-        Dict: 更新后的模型信息，如果未找到则返回None
+        Dict: 更新后的模型信息,如果未找到则返回None
     """
     models = load_models()
     

autocoder/rag/doc_filter.py

@@ -91,9 +91,7 @@ class DocFilter:
                 def _run(conversations, docs):
                     submit_time_1 = time.time()
                     try:
-                        llm = ByzerLLM()
-                        llm.skip_nontext_check = True
-                        llm.setup_default_model_name(self.recall_llm.default_model_name)
+                        llm = self.recall_llm
 
                         v = (
                             _check_relevance_with_conversation.with_llm(
@@ -150,7 +148,7 @@ class DocFilter:
                 )
                 if (
                     relevance
-                    and relevance.is_relevant
+                    # and relevance.is_relevant
                     and relevance.relevant_score >= self.relevant_score
                 ):
                     relevant_docs.append(

autocoder/rag/long_context_rag.py

@@ -52,11 +52,13 @@ class LongContextRAG:
     ) -> None:
         self.llm = llm
         self.args = args
-
-        self.index_model = byzerllm.ByzerLLM()
-        self.index_model.setup_default_model_name(
-            args.index_model or self.llm.default_model_name
-        )
+        if args.product_mode == "pro":
+            self.index_model = byzerllm.ByzerLLM()
+            self.index_model.setup_default_model_name(
+                args.index_model or self.llm.default_model_name
+            )
+        else:
+            self.index_model = self.llm
 
         self.path = path
         self.relevant_score = self.args.rag_doc_filter_relevance or 5
@@ -437,7 +439,7 @@ class LongContextRAG:
 
             logger.info(f"Query: {query} only_contexts: {only_contexts}")
             start_time = time.time()
-            relevant_docs: List[FilterDoc] = self._filter_docs(conversations)
+            relevant_docs: List[FilterDoc] = self._filter_docs(conversations)            
             filter_time = time.time() - start_time
 
             # Filter relevant_docs to only include those with is_relevant=True

autocoder/rag/token_limiter.py

@@ -224,9 +224,7 @@ class TokenLimiter:
                 for idx, line in enumerate(source_code_lines):
                     source_code_with_line_number += f"{idx+1} {line}\n"
                 
-                llm = ByzerLLM()
-                llm.skip_nontext_check = True
-                llm.setup_default_model_name(self.chunk_llm.default_model_name)
+                llm = self.chunk_llm
 
                 extracted_info = (
                     self.extract_relevance_range_from_docs_with_conversation.options(

autocoder/utils/auto_coder_utils/chat_stream_out.py

@@ -1,4 +1,5 @@
 from rich.console import Console
+from autocoder.common.printer import Printer
 from rich.live import Live
 from rich.panel import Panel
 from rich.markdown import Markdown
@@ -11,6 +12,7 @@ from autocoder.utils.request_queue import request_queue
 import time
 from byzerllm.utils.types import SingleOutputMeta
 from autocoder.common import AutoCoderArgs
+from autocoder.common.global_cancel import global_cancel
 
 MAX_HISTORY_LINES = 40  # 最大保留历史行数
 
@@ -167,12 +169,18 @@ def stream_out(
     keep_reasoning_content = True
     if args:
         keep_reasoning_content = args.keep_reasoning_content            
+    
+    keep_only_reasoning_content = False
+    if args:
+        keep_only_reasoning_content = args.keep_only_reasoning_content
 
     lines_buffer = []  # 存储历史行
     current_line = ""  # 当前行
     assistant_response = ""
     last_meta = None
-    panel_title = title if title is not None else f"Response[ {model_name} ]"    
+    panel_title = title if title is not None else f"Response[ {model_name} ]"  
+    first_token_time = 0.0
+    first_token_time_start = time.time()
     try:        
         with Live(
             Panel("", title=panel_title, border_style="green"),
@@ -180,28 +188,41 @@ def stream_out(
             console=console
         ) as live:
             for res in stream_generator:
+                if global_cancel.requested:
+                    printer = Printer(console)
+                    printer.print_in_terminal("generation_cancelled")                    
+                    break
                 last_meta = res[1]                
                 content = res[0]
-                reasoning_content = last_meta.reasoning_content
+
+                reasoning_content = ""
+                if last_meta:
+                    reasoning_content = last_meta.reasoning_content
 
                 if reasoning_content == "" and content == "":
                     continue
                 
-                if keep_reasoning_content:
-                    # 处理思考内容
-                    if reasoning_content:
-                        if assistant_response == "":  # 首次遇到思考内容时添加开标签
-                            assistant_response = "<thinking>"
-                        assistant_response += reasoning_content
-                        
-                    # 处理正式内容
-                    if content:
-                        # 如果之前有思考内容,需要先关闭thinking标签
-                        if "</thinking>" not in assistant_response and "<thinking>" in assistant_response:
-                            assistant_response += "</thinking>"
+                if first_token_time == 0.0:
+                    first_token_time = time.time() - first_token_time_start
+
+                if keep_only_reasoning_content:
+                    assistant_response += reasoning_content
+                else:    
+                    if keep_reasoning_content:
+                        # 处理思考内容
+                        if reasoning_content:
+                            if assistant_response == "":  # 首次遇到思考内容时添加开标签
+                                assistant_response = "<thinking>"
+                            assistant_response += reasoning_content
+                            
+                        # 处理正式内容
+                        if content:
+                            # 如果之前有思考内容,需要先关闭thinking标签
+                            if "</thinking>" not in assistant_response and "<thinking>" in assistant_response:
+                                assistant_response += "</thinking>"
+                            assistant_response += content
+                    else:
                         assistant_response += content
-                else:
-                    assistant_response += content
 
                 display_delta = reasoning_content if reasoning_content else content
                 
@@ -280,5 +301,7 @@ def stream_out(
                     status=RequestOption.COMPLETED
                 ),
             )
-            
+
+    if last_meta:
+        last_meta.first_token_time = first_token_time
     return assistant_response, last_meta

autocoder/utils/llms.py

@@ -1,7 +1,21 @@
 import byzerllm
-from autocoder.auto_coder import models_module
+from typing import Union,Optional
+
+def get_llm_names(llm: Union[byzerllm.ByzerLLM, byzerllm.SimpleByzerLLM,str],target_model_type:Optional[str]=None):
+   if target_model_type is None:
+      return [llm.default_model_name for llm in [llm] if llm.default_model_name]
+   llms = llm.get_sub_client(target_model_type) 
+   if llms is None:
+      return [llm.default_model_name for llm in [llm] if llm.default_model_name]
+   elif isinstance(llms, list):
+       return [llm.default_model_name for llm in llms if llm.default_model_name]
+   elif isinstance(llms,str) and llms:
+      return llms.split(",")      
+   else:
+      return [llm.default_model_name for llm in [llms] if llm.default_model_name]
 
 def get_single_llm(model_names: str, product_mode: str):
+    from autocoder import models as models_module
     if product_mode == "pro":
         if "," in model_names:
             # Multiple code models specified

autocoder/utils/thread_utils.py

@@ -0,0 +1,201 @@
+from concurrent.futures import ThreadPoolExecutor, TimeoutError, CancelledError
+from threading import Event
+from inspect import signature
+from functools import wraps
+from typing import Any, Optional
+import threading
+import logging
+import time
+from autocoder.common.global_cancel import global_cancel
+
+class CancellationRequested(Exception):
+    """Raised when a task is requested to be cancelled."""
+    pass
+
+
+def run_in_thread(timeout: Optional[float] = None):
+    """Decorator that runs a function in a thread with signal handling.
+    
+    Args:
+        timeout (float, optional): Maximum time to wait for thread completion in seconds.
+            If None, will wait indefinitely.
+            
+    The decorated function will run in a separate thread and can be interrupted by
+    signals like Ctrl+C (KeyboardInterrupt). When interrupted, it will log the event
+    and clean up gracefully.
+    """
+    def decorator(func):
+        @wraps(func)
+        def wrapper(*args, **kwargs):
+            with ThreadPoolExecutor(max_workers=1) as executor:
+                future = executor.submit(func, *args, **kwargs)
+                start_time = time.time()
+                
+                while True:
+                    try:
+                        # 使用较短的超时时间进行轮询，确保能够响应中断信号
+                        poll_timeout = 0.1
+                        if timeout is not None:
+                            remaining = timeout - (time.time() - start_time)
+                            if remaining <= 0:
+                                future.cancel()
+                                raise TimeoutError(f"Timeout after {timeout}s in {func.__name__}")
+                            poll_timeout = min(poll_timeout, remaining)
+                            
+                        try:
+                            return future.result(timeout=poll_timeout)
+                        except TimeoutError:
+                            continue  # 继续轮询
+                            
+                    except KeyboardInterrupt:
+                        logging.warning("KeyboardInterrupt received, attempting to cancel task...")
+                        future.cancel()
+                        raise
+                    except Exception as e:
+                        logging.error(f"Error occurred in thread: {str(e)}")
+                        raise
+        return wrapper
+    return decorator
+
+def run_in_thread_with_cancel(timeout: Optional[float] = None):
+    """Decorator that runs a function in a thread with explicit cancellation support.
+    
+    Args:
+        timeout (float, optional): Maximum time to wait for thread completion in seconds.
+            If None, will wait indefinitely.
+            
+    The decorated function MUST accept 'cancel_event' as its first parameter.
+    This cancel_event is a threading.Event object that can be used to check if
+    cancellation has been requested.
+    
+    The decorated function can be called with an external cancel_event passed as a keyword argument.
+    If not provided, a new Event will be created.
+    
+    Example:
+        @run_in_thread_with_cancel(timeout=10)
+        def long_task(cancel_event, arg1, arg2):
+            while not cancel_event.is_set():
+                # do work
+                if cancel_event.is_set():
+                    raise CancellationRequested()
+                    
+        # 使用外部传入的cancel_event
+        external_cancel = Event()
+        try:
+            result = long_task(arg1, arg2, cancel_event=external_cancel)
+        except CancelledError:
+            print("Task was cancelled")
+            
+        # 在其他地方取消任务
+        external_cancel.set()
+    """
+    def decorator(func):
+        # 检查函数签名
+        sig = signature(func)
+        params = list(sig.parameters.keys())
+        if not params or params[0] != 'cancel_event':
+            raise ValueError(
+                f"Function {func.__name__} must have 'cancel_event' as its first parameter. "
+                f"Current parameters: {params}"
+            )
+        
+        @wraps(func)
+        def wrapper(*args, **kwargs):
+            # 从kwargs中提取或创建cancel_event
+            cancel_event = kwargs.pop('cancel_event', None) or Event()
+            
+            def cancellable_task():
+                try:
+                    return func(cancel_event, *args, **kwargs)
+                except CancellationRequested:
+                    logging.info(f"Task {func.__name__} was cancelled")
+                    raise
+                except Exception as e:
+                    logging.error(f"Error in {func.__name__}: {str(e)}")
+                    raise
+            
+            with ThreadPoolExecutor(max_workers=1) as executor:
+                future = executor.submit(cancellable_task)
+                start_time = time.time()
+                
+                while True:
+                    try:
+                        # 使用较短的超时时间进行轮询，确保能够响应中断信号
+                        poll_timeout = 0.1
+                        if timeout is not None:
+                            remaining = timeout - (time.time() - start_time)
+                            if remaining <= 0:
+                                cancel_event.set()
+                                future.cancel()
+                                raise TimeoutError(f"Timeout after {timeout}s in {func.__name__}")
+                            poll_timeout = min(poll_timeout, remaining)
+                            
+                        try:
+                            return future.result(timeout=poll_timeout)
+                        except TimeoutError:
+                            continue  # 继续轮询
+                            
+                    except KeyboardInterrupt:
+                        logging.warning(f"KeyboardInterrupt received, cancelling {func.__name__}...")
+                        cancel_event.set()
+                        future.cancel()
+                        raise CancelledError("Task cancelled by user")
+                    except CancellationRequested:
+                        logging.info(f"Task {func.__name__} was cancelled")
+                        raise CancelledError("Task cancelled by request")
+                    except Exception as e:
+                        logging.error(f"Error occurred in thread: {str(e)}")
+                        raise
+                        
+        return wrapper
+    return decorator
+
+
+def run_in_raw_thread():
+    """A decorator that runs a function in a separate thread and handles exceptions.
+    
+    Args:
+        func: The function to run in a thread
+        
+    Returns:
+        A wrapper function that executes the decorated function in a thread
+        
+    The decorator will:
+    1. Run the function in a separate thread
+    2. Handle KeyboardInterrupt properly
+    3. Propagate exceptions from the thread
+    4. Support function arguments
+    5. Preserve function metadata
+    """
+    def decorator(func):
+
+        @wraps(func)
+        def wrapper(*args, **kwargs):
+            # Store thread results
+            result = []
+            exception = []            
+            def worker():            
+                try:
+                    ret = func(*args, **kwargs)
+                    result.append(ret)
+                    global_cancel.reset()                
+                except Exception as e:
+                    global_cancel.reset()
+                    raise
+            
+            # Create and start thread with a meaningful name
+            thread = threading.Thread(target=worker, name=f"{func.__name__}_thread")
+            thread.daemon = True  # Make thread daemon so it doesn't prevent program exit
+            
+            try:
+                thread.start()
+                while thread.is_alive():
+                    thread.join(0.1)
+
+                return result[0] if result else None            
+            except KeyboardInterrupt:  
+                global_cancel.set()                              
+                raise KeyboardInterrupt("Task was cancelled by user")
+                
+        return wrapper
+    return decorator    

autocoder/version.py

@@ -1 +1 @@
-__version__ = "0.1.251"
+__version__ = "0.1.253"