auto-coder 0.1.226__py3-none-any.whl → 0.1.227__py3-none-any.whl

autocoder/common/code_auto_generate_diff.py

@@ -6,6 +6,7 @@ from autocoder.utils.queue_communicate import queue_communicate, CommunicateEven
 from autocoder.common import sys_prompt
 from concurrent.futures import ThreadPoolExecutor
 import json
+from autocoder.common.utils_code_auto_generate import chat_with_continue
 
 
 class CodeAutoGenerateDiff:
@@ -336,35 +337,48 @@ class CodeAutoGenerateDiff:
 
         conversations_list = []
         results = []
+        input_tokens_count = 0
+        generated_tokens_count = 0
         if not self.args.human_as_model:
             with ThreadPoolExecutor(max_workers=len(self.llms) * self.generate_times_same_model) as executor:
                 futures = []
                 for llm in self.llms:
                     for _ in range(self.generate_times_same_model):
                         futures.append(executor.submit(
-                            llm.chat_oai, conversations=conversations, llm_config=llm_config))
-                results = [future.result()[0].output for future in futures]
+                            chat_with_continue, llm=llm, conversations=conversations, llm_config=llm_config))
+                temp_results = [future.result() for future in futures]
+                for result in temp_results:
+                    results.append(result.content)
+                    input_tokens_count += result.input_tokens_count
+                    generated_tokens_count += result.generated_tokens_count
+            
             for result in results:
                 conversations_list.append(
                     conversations + [{"role": "assistant", "content": result}])
         else:            
             for _ in range(self.args.human_model_num):
-                v = self.llms[0].chat_oai(
-                    conversations=conversations, llm_config=llm_config)
-                results.append(v[0].output)
+                single_result = chat_with_continue(llm=self.llms[0], conversations=conversations, llm_config=llm_config)                
+                results.append(single_result.content)
+                input_tokens_count += single_result.input_tokens_count
+                generated_tokens_count += single_result.generated_tokens_count
                 conversations_list.append(
-                    conversations + [{"role": "assistant", "content": v[0].output}])
-
+                    conversations + [{"role": "assistant", "content": single_result.content}])
+        
+        statistics = {
+            "input_tokens_count": input_tokens_count,
+            "generated_tokens_count": generated_tokens_count
+        }        
+    
         if self.args.request_id and not self.args.skip_events:
             _ = queue_communicate.send_event(
                 request_id=self.args.request_id,
                 event=CommunicateEvent(
                     event_type=CommunicateEventType.CODE_GENERATE_END.value,
-                    data=json.dumps({}, ensure_ascii=False),
+                    data=json.dumps(statistics, ensure_ascii=False),
                 ),
             )
 
-        return CodeGenerateResult(contents=results, conversations=conversations_list)
+        return CodeGenerateResult(contents=results, conversations=conversations_list, metadata=statistics)
 
     def multi_round_run(
         self, query: str, source_content: str, max_steps: int = 10

autocoder/common/code_auto_generate_editblock.py

@@ -10,6 +10,7 @@ from autocoder.utils.queue_communicate import (
 )
 import json
 from concurrent.futures import ThreadPoolExecutor
+from autocoder.common.utils_code_auto_generate import chat_with_continue
 
 
 class CodeAutoGenerateEditBlock:
@@ -418,34 +419,47 @@ class CodeAutoGenerateEditBlock:
 
         conversations_list = []
         results = []
+        input_tokens_count = 0
+        generated_tokens_count = 0
         if not self.args.human_as_model:
             with ThreadPoolExecutor(max_workers=len(self.llms) * self.generate_times_same_model) as executor:
                 futures = []
                 for llm in self.llms:
                     for _ in range(self.generate_times_same_model):
                         futures.append(executor.submit(
-                            llm.chat_oai, conversations=conversations, llm_config=llm_config))
-                results = [future.result()[0].output for future in futures]
+                            chat_with_continue,llm=llm, conversations=conversations, llm_config=llm_config))
+                temp_results = [future.result() for future in futures]
+                for result in temp_results:
+                    results.append(result.content)
+                    input_tokens_count += result.input_tokens_count
+                    generated_tokens_count += result.generated_tokens_count
+            
             for result in results:
                 conversations_list.append(
                     conversations + [{"role": "assistant", "content": result}])
         else:            
             for _ in range(self.args.human_model_num):
-                v = self.llms[0].chat_oai(
-                    conversations=conversations, llm_config=llm_config)
-                results.append(v[0].output)
-                conversations_list.append(conversations + [{"role": "assistant", "content": v[0].output}])
-
+                single_result = chat_with_continue(llm=self.llms[0], conversations=conversations, llm_config=llm_config)                
+                results.append(single_result.content)
+                input_tokens_count += single_result.input_tokens_count
+                generated_tokens_count += single_result.generated_tokens_count
+                conversations_list.append(conversations + [{"role": "assistant", "content": single_result.content}])
+        
+        statistics = {
+            "input_tokens_count": input_tokens_count,
+            "generated_tokens_count": generated_tokens_count
+        }        
+    
         if self.args.request_id and not self.args.skip_events:
             _ = queue_communicate.send_event(
                 request_id=self.args.request_id,
                 event=CommunicateEvent(
                     event_type=CommunicateEventType.CODE_GENERATE_END.value,
-                    data=json.dumps({}, ensure_ascii=False),
+                    data=json.dumps(statistics, ensure_ascii=False),
                 ),
             )
 
-        return CodeGenerateResult(contents=results, conversations=conversations_list)
+        return CodeGenerateResult(contents=results, conversations=conversations_list, metadata=statistics)
 
     def multi_round_run(
         self, query: str, source_content: str, max_steps: int = 10

autocoder/common/code_auto_generate_strict_diff.py

@@ -6,6 +6,7 @@ from autocoder.utils.queue_communicate import queue_communicate, CommunicateEven
 from autocoder.common import sys_prompt
 from concurrent.futures import ThreadPoolExecutor
 import json
+from autocoder.common.utils_code_auto_generate import chat_with_continue
 
 class CodeAutoGenerateStrictDiff:
     def __init__(
@@ -306,34 +307,47 @@ class CodeAutoGenerateStrictDiff:
         
         conversations_list = []
         results = []
+        input_tokens_count = 0
+        generated_tokens_count = 0
         if not self.args.human_as_model:
             with ThreadPoolExecutor(max_workers=len(self.llms) * self.generate_times_same_model) as executor:
                 futures = []
                 for llm in self.llms:
                     for _ in range(self.generate_times_same_model):
                         futures.append(executor.submit(
-                            llm.chat_oai, conversations=conversations, llm_config=llm_config))
-                results = [future.result()[0].output for future in futures]
+                            chat_with_continue, llm=llm, conversations=conversations, llm_config=llm_config))
+                temp_results = [future.result() for future in futures]
+                for result in temp_results:
+                    results.append(result.content)
+                    input_tokens_count += result.input_tokens_count
+                    generated_tokens_count += result.generated_tokens_count
+            
             for result in results:
                 conversations_list.append(
                     conversations + [{"role": "assistant", "content": result}])
         else:            
             for _ in range(self.args.human_model_num):
-                v = self.llms[0].chat_oai(
-                    conversations=conversations, llm_config=llm_config)
-                results.append(v[0].output)
-                conversations_list.append(conversations + [{"role": "assistant", "content": v[0].output}])
-
+                single_result = chat_with_continue(llm=self.llms[0], conversations=conversations, llm_config=llm_config)                
+                results.append(single_result.content)
+                input_tokens_count += single_result.input_tokens_count
+                generated_tokens_count += single_result.generated_tokens_count
+                conversations_list.append(conversations + [{"role": "assistant", "content": single_result.content}])
+        
+        statistics = {
+            "input_tokens_count": input_tokens_count,
+            "generated_tokens_count": generated_tokens_count
+        }        
+    
         if self.args.request_id and not self.args.skip_events:
             _ = queue_communicate.send_event(
                 request_id=self.args.request_id,
                 event=CommunicateEvent(
                     event_type=CommunicateEventType.CODE_GENERATE_END.value,
-                    data=json.dumps({}, ensure_ascii=False),
+                    data=json.dumps(statistics, ensure_ascii=False),
                 ),
             )
 
-        return CodeGenerateResult(contents=results, conversations=conversations_list)
+        return CodeGenerateResult(contents=results, conversations=conversations_list, metadata=statistics)
 
     def multi_round_run(
         self, query: str, source_content: str, max_steps: int = 10

autocoder/common/command_completer.py

@@ -29,6 +29,12 @@ COMMANDS = {
         "/refresh": "",
         "/get": "",
     },
+    "/models": {
+        "/add": "",
+        "/add_model": "",
+        "/remove": "",
+        "/list": ""        
+    }
 }
 
 

autocoder/common/types.py

@@ -12,6 +12,7 @@ class StepNum(pydantic.BaseModel):
 class CodeGenerateResult(pydantic.BaseModel):
     contents:List[str]
     conversations:List[List[Dict[str, Any]]]
+    metadata:Dict[str, Any] = {}
 
 class MergeCodeWithoutEffect(pydantic.BaseModel):
     success_blocks: List[Tuple[str, str]]

autocoder/common/utils_code_auto_generate.py

@@ -0,0 +1,38 @@
+from byzerllm import ByzerLLM
+from typing import List,Any,Union
+from pydantic import BaseModel
+from loguru import logger
+class ChatWithContinueResult(BaseModel):
+    content: str
+    input_tokens_count: int
+    generated_tokens_count: int
+    
+
+def chat_with_continue(llm: ByzerLLM, conversations: List[dict], llm_config: dict) -> ChatWithContinueResult:
+    final_result = ChatWithContinueResult(content="", input_tokens_count=0, generated_tokens_count=0)
+    v = llm.chat_oai(
+        conversations=conversations, llm_config=llm_config)
+                
+    single_result = v[0].output
+    metadata = v[0].metadata
+
+    final_result.input_tokens_count += metadata.get("input_tokens_count", 0)
+    final_result.generated_tokens_count += metadata.get("generated_tokens_count", 0)
+
+    temp_conversations = conversations + \
+        [{"role": "assistant", "content": single_result}]
+    
+    count = 1
+    while (metadata.get("finish_reason", "stop") == "length" and count < 6):        
+        v = llm.chat_oai(
+            conversations=temp_conversations, llm_config={**llm_config, "gen.response_prefix": True})
+        metadata = v[0].metadata
+        single_result += v[0].output
+        final_result.input_tokens_count += metadata.get("input_tokens_count", 0)
+        final_result.generated_tokens_count += metadata.get("generated_tokens_count", 0)
+        count += 1
+    
+    # if count >= 2:
+    #   logger.info(f"The code generation is exceed the max length, continue to generate the code {count -1 } times")
+    final_result.content = single_result    
+    return final_result

autocoder/dispacher/actions/action.py

@@ -116,7 +116,7 @@ class ActionTSProject(BaseAction):
                 generate_result = generate.single_round_run(
                     query=args.query, source_content=content
                 )
-            logger.info(f"Code generation completed in {time.time() - start_time:.2f} seconds")    
+            logger.info(f"Code generation completed in {time.time() - start_time:.2f} seconds, input_tokens_count: {generate_result.metadata.get('input_tokens_count', 0)}, generated_tokens_count: {generate_result.metadata.get('generated_tokens_count', 0)}")    
             merge_result = None
             if args.execute and args.auto_merge:
                 logger.info("Auto merge the code...")
@@ -200,7 +200,7 @@ class ActionPyScriptProject(BaseAction):
                     query=args.query, source_content=content
                 )
 
-            logger.info(f"Code generation completed in {time.time() - start_time:.2f} seconds")    
+            logger.info(f"Code generation completed in {time.time() - start_time:.2f} seconds, input_tokens_count: {generate_result.metadata.get('input_tokens_count', 0)}, generated_tokens_count: {generate_result.metadata.get('generated_tokens_count', 0)}")    
             merge_result = None
             if args.execute and args.auto_merge:
                 logger.info("Auto merge the code...")
@@ -302,7 +302,7 @@ class ActionPyProject(BaseAction):
                 generate_result = generate.single_round_run(
                     query=args.query, source_content=content
                 )
-            logger.info(f"Code generation completed in {time.time() - start_time:.2f} seconds")   
+            logger.info(f"Code generation completed in {time.time() - start_time:.2f} seconds, input_tokens_count: {generate_result.metadata.get('input_tokens_count', 0)}, generated_tokens_count: {generate_result.metadata.get('generated_tokens_count', 0)}")    
             merge_result = None
             if args.execute and args.auto_merge:
                 logger.info("Auto merge the code...")
@@ -396,7 +396,7 @@ class ActionSuffixProject(BaseAction):
                     query=args.query, source_content=content
                 )
               
-        logger.info(f"Code generation completed in {time.time() - start_time:.2f} seconds")
+        logger.info(f"Code generation completed in {time.time() - start_time:.2f} seconds, input_tokens_count: {generate_result.metadata.get('input_tokens_count', 0)}, generated_tokens_count: {generate_result.metadata.get('generated_tokens_count', 0)}")
         merge_result = None
         if args.execute and args.auto_merge:
             logger.info("Auto merge the code...")

autocoder/dispacher/actions/plugins/action_regex_project.py

@@ -13,7 +13,7 @@ from autocoder.index.index import build_index_and_filter_files
 from autocoder.regexproject import RegexProject
 from autocoder.utils.conversation_store import store_code_model_conversation
 from loguru import logger
-
+import time
 
 class ActionRegexProject:
     def __init__(
@@ -51,8 +51,10 @@ class ActionRegexProject:
                 )
                 content = content[: self.args.model_max_input_length]
 
+        start_time = time.time()
         if args.execute:
             logger.info("Auto generate the code...")
+            
             if args.auto_merge == "diff":
                 generate = CodeAutoGenerateDiff(
                     llm=self.llm, args=self.args, action=self
@@ -75,6 +77,8 @@ class ActionRegexProject:
                 generate_result = generate.single_round_run(
                     query=args.query, source_content=content
                 )
+
+            logger.info(f"Code generation completed in {time.time() - start_time:.2f} seconds, input_tokens_count: {generate_result.metadata.get('input_tokens_count', 0)}, generated_tokens_count: {generate_result.metadata.get('generated_tokens_count', 0)}")    
             merge_result = None
             if args.execute and args.auto_merge:
                 logger.info("Auto merge the code...")
@@ -90,7 +94,7 @@ class ActionRegexProject:
                 else:
                     code_merge = CodeAutoMerge(llm=self.llm, args=self.args)
                     merge_result = code_merge.merge_code(generate_result=generate_result)
-
+            
             if merge_result is not None:
                 content = merge_result.contents[0]
                 store_code_model_conversation(

autocoder/models.py

@@ -0,0 +1,158 @@
+import os
+import json
+from typing import List, Dict
+from urllib.parse import urlparse
+
+MODELS_JSON = os.path.expanduser("~/.auto-coder/keys/models.json")
+
+# Default built-in models
+default_models_list = [
+    {
+        "name": "deepseek_r1_chat",
+        "description": "DeepSeek Reasoner is for design/review",
+        "model_name": "deepseek-reasoner",
+        "model_type": "saas/openai",
+        "base_url": "https://api.deepseek.com/v1",
+        "api_key_path": "api.deepseek.com"
+    },    
+    {
+        "name": "deepsee_chat",
+        "description": "DeepSeek Chat is for coding",
+        "model_name": "deepseek-chat",
+        "model_type": "saas/openai",
+        "base_url": "https://api.deepseek.com/v1",
+        "api_key_path": "api.deepseek.com"
+    },
+    {
+        "name":"o1",
+        "description": "o1 is for design/review",
+        "model_name": "o1-2024-12-17",
+        "model_type": "saas/openai",
+        "base_url": "https://api.openai.com/v1",
+        "api_key_path": ""
+    }
+]
+
+def load_models() -> List[Dict]:
+    """
+    Load models from ~/.auto-coder/keys/models.json and merge with default_models_list.
+    Models are merged and deduplicated based on their name field.
+    If file doesn't exist or is invalid, use default_models_list.
+    """
+    os.makedirs(os.path.dirname(MODELS_JSON), exist_ok=True)
+    
+    # Start with default models
+    models_dict = {model["name"]: model for model in default_models_list}
+    
+    # If JSON file exists, read and merge with defaults
+    if os.path.exists(MODELS_JSON):
+        try:
+            with open(MODELS_JSON, 'r', encoding='utf-8') as f:
+                custom_models = json.load(f)
+                # Custom models will override defaults with same name
+                for model in custom_models:
+                    models_dict[model["name"]] = model
+        except json.JSONDecodeError:
+            # If JSON is invalid, just use defaults
+            save_models(default_models_list)
+    else:
+        # If file doesn't exist, create it with defaults
+        save_models(default_models_list)
+    
+    # Convert merged dictionary back to list
+    target_models = list(models_dict.values())
+    api_key_dir = os.path.expanduser("~/.auto-coder/keys")
+    for model in target_models:    
+        if model.get("api_key_path",""):           
+            api_key_file = os.path.join(api_key_dir, model["api_key_path"])
+            if os.path.exists(api_key_file):
+                with open(api_key_file, "r") as f:
+                    model["api_key"] = f.read()
+    return target_models
+
+def save_models(models: List[Dict]) -> None:
+    """
+    Save models to ~/.auto-coder/keys/models.json
+    """
+    os.makedirs(os.path.dirname(MODELS_JSON), exist_ok=True)
+    with open(MODELS_JSON, 'w', encoding='utf-8') as f:
+        json.dump(models, f, indent=2, ensure_ascii=False)
+
+
+def process_api_key_path(base_url: str) -> str:
+    """
+    从 base_url 中提取 host 部分并处理特殊字符
+    例如: https://api.example.com:8080/v1 -> api.example.com_8080
+    """
+    if not base_url:
+        return ""
+    
+    parsed = urlparse(base_url)
+    host = parsed.netloc
+    
+    # 将冒号替换为下划线
+    host = host.replace(":", "_")
+    
+    return host
+
+def get_model_by_name(name: str) -> Dict:
+    """
+    根据模型名称查找模型
+    """
+    models = load_models()
+    v = [m for m in models if m["name"] == name.strip()]
+    
+    if len(v) == 0:
+        raise Exception(f"Model {name} not found")
+    return v[0]
+
+def update_model_with_api_key(name: str, api_key: str) -> Dict:
+    """
+    根据模型名称查找并更新模型的 api_key_path。
+    如果找到模型，会根据其 base_url 处理 api_key_path。
+    
+    Args:
+        name: 模型名称
+        api_key: API密钥
+        
+    Returns:
+        Dict: 更新后的模型信息，如果未找到则返回None
+    """
+    models = load_models()
+    
+    # 在现有模型中查找
+    found_model = None
+    for model in models:
+        if model["name"] == name:
+            found_model = model
+            break
+                    
+    if not found_model:
+        return None
+        
+    # 从 base_url 中提取并处理 host
+    api_key_path = process_api_key_path(found_model["base_url"])
+    if api_key_path:
+        found_model["api_key_path"] = api_key_path
+        
+        # 保存 API 密钥
+        api_key_dir = os.path.expanduser("~/.auto-coder/keys")
+        os.makedirs(api_key_dir, exist_ok=True)
+        api_key_file = os.path.join(api_key_dir, api_key_path)
+        with open(api_key_file, "w") as f:
+            f.write(api_key)
+        
+        # 如果是新模型，添加到模型列表中
+        if all(model["name"] != name for model in models):
+            models.append(found_model)
+        else:
+            # 更新现有模型
+            for i, model in enumerate(models):
+                if model["name"] == name:
+                    models[i] = found_model
+                    break
+        
+        save_models(models)
+    
+    return found_model
+

autocoder/version.py

@@ -1 +1 @@
-__version__ = "0.1.226"
+__version__ = "0.1.227"