auto-coder 0.1.263__py3-none-any.whl → 0.1.264__py3-none-any.whl

autocoder/common/code_auto_merge_editblock.py

@@ -164,15 +164,35 @@ class CodeAutoMergeEditBlock:
     def choose_best_choice(self, generate_result: CodeGenerateResult) -> CodeGenerateResult:
         if len(generate_result.contents) == 1:
             return generate_result
-
+        
+        merge_results = []
+        for content,conversations in zip(generate_result.contents,generate_result.conversations):
+            merge_result = self._merge_code_without_effect(content)
+            merge_results.append(merge_result)
+
+        # If all merge results are None, return first one
+        if all(len(result.failed_blocks) != 0 for result in merge_results):
+            self.printer.print_in_terminal("all_merge_results_failed")
+            return CodeGenerateResult(contents=[generate_result.contents[0]], conversations=[generate_result.conversations[0]])
+        
+        # If only one merge result is not None, return that one
+        not_none_indices = [i for i, result in enumerate(merge_results) if len(result.failed_blocks) == 0]
+        if len(not_none_indices) == 1:
+            idx = not_none_indices[0]
+            self.printer.print_in_terminal("only_one_merge_result_success")
+            return CodeGenerateResult(contents=[generate_result.contents[idx]], conversations=[generate_result.conversations[idx]])        
+
+        # 最后，如果有多个，那么根据质量排序再返回
         ranker = CodeModificationRanker(self.llm, self.args)
-        ranked_result = ranker.rank_modifications(generate_result)
-        # Filter out contents with failed blocks
+        ranked_result = ranker.rank_modifications(generate_result,merge_results)        
+         
+        ## 得到的结果，再做一次合并，第一个通过的返回 , 返回做合并有点重复低效，未来修改。
         for content,conversations in zip(ranked_result.contents,ranked_result.conversations):
             merge_result = self._merge_code_without_effect(content)
             if not merge_result.failed_blocks:
                 return CodeGenerateResult(contents=[content], conversations=[conversations])
-        # If all have failed blocks, return the first one
+
+        # 最后保底，但实际不会出现
         return CodeGenerateResult(contents=[ranked_result.contents[0]], conversations=[ranked_result.conversations[0]])
 
     @byzerllm.prompt()

autocoder/common/code_auto_merge_strict_diff.py

@@ -138,15 +138,35 @@ class CodeAutoMergeStrictDiff:
     def choose_best_choice(self, generate_result: CodeGenerateResult) -> CodeGenerateResult:
         if len(generate_result.contents) == 1:
             return generate_result
+        
+        merge_results = []
+        for content,conversations in zip(generate_result.contents,generate_result.conversations):
+            merge_result = self._merge_code_without_effect(content)
+            merge_results.append(merge_result)
+
+        # If all merge results are None, return first one
+        if all(len(result.failed_blocks) != 0 for result in merge_results):
+            self.printer.print_in_terminal("all_merge_results_failed")
+            return CodeGenerateResult(contents=[generate_result.contents[0]], conversations=[generate_result.conversations[0]])
+        
+        # If only one merge result is not None, return that one
+        not_none_indices = [i for i, result in enumerate(merge_results) if len(result.failed_blocks) == 0]
+        if len(not_none_indices) == 1:
+            idx = not_none_indices[0]
+            self.printer.print_in_terminal("only_one_merge_result_success")
+            return CodeGenerateResult(contents=[generate_result.contents[idx]], conversations=[generate_result.conversations[idx]])        
 
+        # 最后，如果有多个，那么根据质量排序再返回
         ranker = CodeModificationRanker(self.llm, self.args)
-        ranked_result = ranker.rank_modifications(generate_result)
-        # Filter out contents with failed blocks
+        ranked_result = ranker.rank_modifications(generate_result,merge_results)        
+         
+        ## 得到的结果，再做一次合并，第一个通过的返回 , 返回做合并有点重复低效，未来修改。
         for content,conversations in zip(ranked_result.contents,ranked_result.conversations):
             merge_result = self._merge_code_without_effect(content)
             if not merge_result.failed_blocks:
                 return CodeGenerateResult(contents=[content], conversations=[conversations])
-        # If all have failed blocks, return the first one
+
+        # 最后保底，但实际不会出现
         return CodeGenerateResult(contents=[ranked_result.contents[0]], conversations=[ranked_result.conversations[0]])
     
 

autocoder/common/code_modification_ranker.py

@@ -9,6 +9,8 @@ import traceback
 from autocoder.common.utils_code_auto_generate import chat_with_continue
 from byzerllm.utils.str2model import to_model
 from autocoder.utils.llms import get_llm_names, get_model_info
+from autocoder.common.types import CodeGenerateResult, MergeCodeWithoutEffect
+import os
 
 class RankResult(BaseModel):
     rank_result: List[int]
@@ -51,12 +53,67 @@ class CodeModificationRanker:
         }
         ```
 
-        注意：   
+        注意：                   
         1. id 为 edit_block 的 id,按质量从高到低排序，并且 id 必须是数字        
         2. 只输出前面要求的 Json 格式就好，不要输出其他内容，Json 需要使用 ```json ```包裹                
         '''
 
-    def rank_modifications(self, generate_result: CodeGenerateResult) -> CodeGenerateResult:
+    @byzerllm.prompt()
+    def _rank_modifications_with_merge_result(self, s: CodeGenerateResult,merge_results: List[MergeCodeWithoutEffect]) -> str:
+        '''
+        对一组代码修改进行质量评估并排序。
+
+        下面是修改需求：
+
+        <edit_requirement>
+        {{ s.conversations[0][-2]["content"] }}
+        </edit_requirement>
+
+        下面是相应的代码修改，如果Before 为空，那么表示是新增文件，如果After 为空，那么表示是删除文件，如果Before 和 After 都不为空，那么表示是修改文件：        
+        {% for change in changes %}
+        <edit_file id="{{ loop.index0 }}">
+        {{change}}
+        </edit_file>
+        {% endfor %}
+        
+        请输出如下格式的评估结果,只包含 JSON 数据:
+
+        ```json
+        {
+            "rank_result": [id1, id2, id3] 
+        }
+        ```
+
+        注意：     
+        1. 像python的缩进，前端诸如 reacjs,vue 的标签闭合匹配，这些很重要，需要在排序中作为重点考虑对象之一。   
+        1. id 为 edit_file 的 id,按质量从高到低排序，并且 id 必须是数字        
+        2. 只输出前面要求的 Json 格式就好，不要输出其他内容，Json 需要使用 ```json ```包裹                
+        '''
+        changes = []
+        for merge_result in merge_results:
+            s = ""
+            for block in merge_result.success_blocks:
+                file_path,content = block
+                s += f"##File: {file_path}\n\n"
+                if not os.path.exists(file_path):                    
+                    s += f"##Before: \n\n"
+                    s += f"##File: {file_path}\n\n"
+                    s += f"##After: \n\n"
+                    s += content
+                else:
+                    with open(file_path, "r",encoding="utf-8") as f:
+                        original_content = f.read()
+                    s += f"##Before: \n\n"
+                    s += original_content
+                    s += f"##File: {file_path}\n\n"
+                    s += f"##After: \n\n"
+                    s += content 
+            changes.append(s)                       
+        return {
+            "changes": changes
+        }    
+
+    def rank_modifications(self, generate_result: CodeGenerateResult, merge_result: List[MergeCodeWithoutEffect]) -> CodeGenerateResult:
         import time
         from collections import defaultdict
 
@@ -69,8 +126,13 @@ class CodeModificationRanker:
         
         rank_times = self.args.rank_times_same_model
         total_tasks = len(self.llms) * rank_times
+        if self.args.rank_strategy == "block":
+            query = self._rank_modifications.prompt(generate_result)
+        elif self.args.rank_strategy == "file":
+            query = self._rank_modifications_with_merge_result.prompt(generate_result, merge_result)            
+        else:
+            raise Exception(f"Invalid rank strategy: {self.args.rank_strategy}")
 
-        query = self._rank_modifications.prompt(generate_result)
         input_tokens_count = 0
         generated_tokens_count = 0
         try:

autocoder/common/conf_validator.py

@@ -132,6 +132,12 @@ class ConfigValidator:
             "type": str,
             "default": "v3_chat",
             "description": "提交信息生成模型名称"
+        },
+        "rank_strategy": {
+            "type": str,
+            "allowed": ["block", "file"],
+            "default": "block",
+            "description": "排序策略(block/file)"
         }
     }
 

autocoder/common/context_pruner.py

@@ -0,0 +1,305 @@
+from typing import List, Dict, Any, Union
+from pathlib import Path
+import json
+from loguru import logger
+from autocoder.rag.token_counter import count_tokens
+from autocoder.common import AutoCoderArgs,SourceCode
+from byzerllm.utils.client.code_utils import extract_code
+from autocoder.index.types import VerifyFileRelevance
+import byzerllm
+from concurrent.futures import ThreadPoolExecutor, as_completed
+
+class PruneContext:
+    def __init__(self, max_tokens: int, args: AutoCoderArgs, llm: Union[byzerllm.ByzerLLM, byzerllm.SimpleByzerLLM]):
+        self.max_tokens = max_tokens
+        self.args = args
+        self.llm = llm
+
+    def _delete_overflow_files(self, file_paths: List[str]) -> List[SourceCode]:
+        """直接删除超出 token 限制的文件"""
+        total_tokens = 0
+        selected_files = []
+        token_count = 0
+        for file_path in file_paths:
+            try:
+                with open(file_path, "r", encoding="utf-8") as f:
+                    content = f.read()
+                    token_count = count_tokens(content)
+                    if total_tokens + token_count <= self.max_tokens:
+                        total_tokens += token_count
+                        print(f"{file_path} {token_count} {content}")
+                        selected_files.append(SourceCode(module_name=file_path,source_code=content,tokens=token_count))
+                    else:
+                        break
+            except Exception as e:
+                logger.error(f"Failed to read file {file_path}: {e}")
+                selected_files.append(SourceCode(module_name=file_path,source_code=content,tokens=token_count))
+
+        return selected_files
+
+    def _extract_code_snippets(self, file_paths: List[str], conversations: List[Dict[str, str]]) -> List[SourceCode]:
+        """抽取关键代码片段策略"""
+        token_count = 0
+        selected_files = []
+        full_file_tokens = int(self.max_tokens * 0.8)
+
+        @byzerllm.prompt()
+        def extract_code_snippets(conversations: List[Dict[str, str]], content: str) -> str:
+            """
+            根据提供的代码文件和对话历史提取相关代码片段。            
+
+            处理示例：
+            <examples>
+            1.  代码文件：
+            <code_file>
+                1 def add(a, b):
+                2     return a + b
+                3 def sub(a, b):
+                4     return a - b
+            </code_file>
+            <conversation_history>
+                <user>: 如何实现加法？                
+            </conversation_history>
+
+            输出：
+            ```json
+            [
+                {"start_line": 1, "end_line": 2}                
+            ]
+            ```
+
+            2.  代码文件：
+                1 class User:
+                2     def __init__(self, name):
+                3         self.name = name
+                4     def greet(self):
+                5         return f"Hello, {self.name}"
+            </code_file>
+            <conversation_history>
+                <user>: 如何创建一个User对象？                
+            </conversation_history>
+
+            输出：
+            ```json
+            [
+                {"start_line": 1, "end_line": 3}
+            ]
+            ```
+
+            3.  代码文件：
+            <code_file>
+                1 def foo():
+                2     pass
+            </code_file>
+            <conversation_history>
+                <user>: 如何实现减法？                
+            </conversation_history>
+
+            输出：
+            ```json
+            []
+            ```
+            </examples>
+
+            输入:
+            1. 代码文件内容:
+            <code_file>
+            {{ content }}
+            </code_file>
+
+            2. 对话历史:
+            <conversation_history>
+            {% for msg in conversations %}
+            <{{ msg.role }}>: {{ msg.content }}
+            {% endfor %}
+            </conversation_history>
+
+            任务:
+            1. 分析最后一个用户问题及其上下文。
+            2. 在代码文件中找出与问题相关的一个或多个重要代码段。
+            3. 对每个相关代码段，确定其起始行号(start_line)和结束行号(end_line)。
+            4. 代码段数量不超过4个。
+
+            输出要求:
+            1. 返回一个JSON数组，每个元素包含"start_line"和"end_line"。
+            2. start_line和end_line必须是整数，表示代码文件中的行号。
+            3. 行号从1开始计数。
+            4. 如果没有相关代码段，返回空数组[]。
+
+            输出格式:
+            严格的JSON数组，不包含其他文字或解释。
+
+            ```json
+            [
+                {"start_line": 第一个代码段的起始行号, "end_line": 第一个代码段的结束行号},
+                {"start_line": 第二个代码段的起始行号, "end_line": 第二个代码段的结束行号}
+            ]
+            ```            
+            """
+
+        for file_path in file_paths:
+            try:
+                with open(file_path, "r", encoding="utf-8") as f:
+                    content = f.read()
+
+                    # 完整文件优先
+                    tokens = count_tokens(content)
+                    if token_count + tokens <= full_file_tokens:
+                        selected_files.append(SourceCode(module_name=file_path,source_code=content,tokens=tokens))
+                        token_count += tokens
+                        continue
+
+                    # 抽取关键片段
+                    extracted = extract_code_snippets.with_llm(self.llm).run(
+                        conversations=conversations, 
+                        content=content
+                    )                
+
+                    if extracted:
+                        json_str = extract_code(extracted)[0][1]
+                        snippets = json.loads(json_str)
+                        new_content = self._build_snippet_content(file_path, content, snippets)
+
+                        snippet_tokens = count_tokens(new_content)
+                        if token_count + snippet_tokens <= self.max_tokens:
+                            selected_files.append(SourceCode(module_name=file_path,source_code=new_content,tokens=snippet_tokens))
+                            token_count += snippet_tokens
+                        else:
+                            break
+            except Exception as e:
+                logger.error(f"Failed to process {file_path}: {e}")
+                continue
+
+        return selected_files
+
+    def _build_snippet_content(self, file_path: str, full_content: str, snippets: List[dict]) -> str:
+        """构建包含代码片段的文件内容"""
+        lines = full_content.split("\n")
+        header = f"Snippets:\n"
+
+        content = []
+        for snippet in snippets:
+            start = max(0, snippet["start_line"] - 1)
+            end = min(len(lines), snippet["end_line"])
+            content.append(f"# Lines {start+1}-{end} ({snippet.get('reason','')})")
+            content.extend(lines[start:end])
+
+        return header + "\n".join(content)
+
+    def handle_overflow(
+        self,
+        file_paths: List[str],
+        conversations: List[Dict[str, str]],
+        strategy: str = "score"        
+    ) -> List[SourceCode]:
+        """
+        处理超出 token 限制的文件
+        :param file_paths: 要处理的文件路径列表
+        :param conversations: 对话上下文（用于提取策略）
+        :param strategy: 处理策略 (delete/extract/score)        
+        """
+        total_tokens,sources = self._count_tokens(file_paths)
+        if total_tokens <= self.max_tokens:
+            return sources
+
+        if strategy == "score":            
+            return self._score_and_filter_files(file_paths, conversations)
+        if strategy == "delete":
+            return self._delete_overflow_files(file_paths)
+        elif strategy == "extract":
+            return self._extract_code_snippets(file_paths, conversations)
+        else:
+            raise ValueError(f"无效策略: {strategy}. 可选值: delete/extract/score")
+    
+    def _count_tokens(self, file_paths: List[str]) -> int:
+        """计算文件总token数"""
+        total_tokens = 0
+        sources = []
+        for file_path in file_paths:
+            try:
+                with open(file_path, "r", encoding="utf-8") as f:
+                    content = f.read()
+                    sources.append(SourceCode(module_name=file_path,source_code=content,tokens=count_tokens(content)))
+                    total_tokens += count_tokens(content)
+            except Exception as e:
+                logger.error(f"Failed to read file {file_path}: {e}")
+                total_tokens += 0
+        return total_tokens,sources
+
+    def _score_and_filter_files(self, file_paths: List[str], conversations: List[Dict[str, str]]) -> List[SourceCode]:
+        """根据文件相关性评分过滤文件，直到token数大于max_tokens 停止追加"""
+        selected_files = []
+        total_tokens = 0
+        scored_files = []        
+
+        @byzerllm.prompt()
+        def verify_file_relevance(file_content: str, conversations: List[Dict[str, str]]) -> str:
+            """
+            请验证下面的文件内容是否与用户对话相关:
+
+            文件内容:
+            {{ file_content }}
+
+            历史对话:
+            <conversation_history>
+            {% for msg in conversations %}
+            <{{ msg.role }}>: {{ msg.content }}
+            {% endfor %}
+            </conversation_history>
+
+            相关是指，需要依赖这个文件提供上下文，或者需要修改这个文件才能解决用户的问题。
+            请给出相应的可能性分数：0-10，并结合用户问题，理由控制在50字以内。格式如下:
+
+            ```json
+            {
+                "relevant_score": 0-10,
+                "reason": "这是相关的原因（不超过10个中文字符）..."
+            }
+            ```
+            """
+
+        def _score_file(file_path: str) -> dict:
+            try:
+                with open(file_path, "r", encoding="utf-8") as f:
+                    content = f.read()
+                    tokens = count_tokens(content)
+                    result = verify_file_relevance.with_llm(self.llm).with_return_type(VerifyFileRelevance).run(
+                        file_content=content,
+                        conversations=conversations
+                    )
+                    return {
+                        "file_path": file_path,
+                        "score": result.relevant_score,
+                        "tokens": tokens,
+                        "content": content
+                    }
+            except Exception as e:
+                logger.error(f"Failed to score file {file_path}: {e}")
+                return None
+
+        # 使用线程池并行打分
+        with ThreadPoolExecutor() as executor:
+            futures = [executor.submit(_score_file, file_path) for file_path in file_paths]
+            for future in as_completed(futures):
+                result = future.result()
+                print(f"score file {result['file_path']} {result['score']}")
+                if result:
+                    scored_files.append(result)
+
+        # 第二步：按分数从高到低排序
+        scored_files.sort(key=lambda x: x["score"], reverse=True)
+
+        # 第三步：从高分开始过滤，直到token数大于max_tokens 停止追加
+        for file_info in scored_files:
+            if total_tokens + file_info["tokens"] <= self.max_tokens:
+                selected_files.append(SourceCode(
+                    module_name=file_info["file_path"],
+                    source_code=file_info["content"],
+                    tokens=file_info["tokens"]
+                ))
+                total_tokens += file_info["tokens"]
+            else:
+                break
+
+        return selected_files
+

autocoder/index/entry.py

@@ -24,6 +24,7 @@ from autocoder.index.filter.normal_filter import NormalFilter
 from autocoder.index.index import IndexManager
 from loguru import logger
 from autocoder.common import SourceCodeList
+from autocoder.common.context_pruner import PruneContext
 
 def build_index_and_filter_files(
     llm, args: AutoCoderArgs, sources: List[SourceCode]
@@ -113,8 +114,13 @@ def build_index_and_filter_files(
                 raise KeyboardInterrupt(printer.get_message_from_key_with_format("quick_filter_failed",error=quick_filter_result.error_message))
 
             # Merge quick filter results into final_files
-            final_files.update(quick_filter_result.files)    
-
+            if args.context_prune:
+                context_pruner = PruneContext(max_tokens=args.conversation_prune_safe_zone_tokens, args=args, llm=llm)
+                pruned_files = context_pruner.handle_overflow(quick_filter_result.files, [{"role":"user","content":args.query}], args.context_prune_strategy)
+                for source_file in pruned_files:
+                    final_files[source_file.module_name] = quick_filter_result.files[source_file.module_name]
+            else:
+                final_files.update(quick_filter_result.files)                                     
         
         if not args.skip_filter_index and not args.index_filter_model:
             model_name = getattr(index_manager.llm, 'default_model_name', None)

autocoder/index/filter/normal_filter.py

@@ -1,4 +1,6 @@
-from typing import List, Union,Dict,Any
+from typing import List, Union,Dict,Any,Optional
+
+from pydantic import BaseModel
 from autocoder.index.types import IndexItem
 from autocoder.common import SourceCode, AutoCoderArgs
 import byzerllm
@@ -25,6 +27,11 @@ def get_file_path(file_path):
         return file_path.strip()[2:]
     return file_path
 
+class NormalFilterResult(BaseModel):
+    files: Dict[str, TargetFile]
+    has_error: bool
+    error_message: Optional[str] = None
+    file_positions: Optional[Dict[str, int]]
 
 class NormalFilter():
     def __init__(self, index_manager: IndexManager,stats:Dict[str,Any],sources:List[SourceCode]):
@@ -167,4 +174,8 @@ class NormalFilter():
                 # Keep all files, not just verified ones
                 final_files = verified_files
         
-        return final_files
+        return NormalFilterResult(
+            files=final_files,
+            has_error=False,
+            error_message=None
+        )