auto-coder 0.1.231__py3-none-any.whl → 0.1.233__py3-none-any.whl

autocoder/index/filter/normal_filter.py

@@ -0,0 +1,168 @@
+from typing import List, Union,Dict,Any
+from autocoder.index.types import IndexItem
+from autocoder.common import SourceCode, AutoCoderArgs
+import byzerllm
+import time
+from autocoder.index.index import IndexManager
+from autocoder.index.types import (
+    IndexItem,
+    TargetFile,
+    VerifyFileRelevance,
+    FileList,
+    FileNumberList
+)
+from loguru import logger
+from autocoder.utils.queue_communicate import (
+    queue_communicate,
+    CommunicateEvent,
+    CommunicateEventType,
+)
+from concurrent.futures import ThreadPoolExecutor, as_completed
+import json
+
+def get_file_path(file_path):
+    if file_path.startswith("##"):
+        return file_path.strip()[2:]
+    return file_path
+
+
+class NormalFilter():
+    def __init__(self, index_manager: IndexManager,stats:Dict[str,Any],sources:List[SourceCode]):
+        self.index_manager = index_manager
+        self.args = index_manager.args
+        self.stats = stats
+        self.sources = sources
+
+    def filter(self, index_items: List[IndexItem], query: str) -> Dict[str, TargetFile]:
+        final_files: Dict[str, TargetFile] = {}
+        if not self.args.skip_filter_index:
+            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_INDEX_FILTER_START.value,
+                        data=json.dumps({})
+                    )
+                )
+            # Phase 3: Level 1 filtering - Query-based
+            logger.info(
+                "Phase 3: Performing Level 1 filtering (query-based)...")
+
+            phase_start = time.monotonic()
+            target_files = self.index_manager.get_target_files_by_query(self.args.query)
+
+            if target_files:
+                for file in target_files.file_list:
+                    file_path = file.file_path.strip()
+                    final_files[get_file_path(file_path)] = file
+                self.stats["level1_filtered"] = len(target_files.file_list)
+            phase_end = time.monotonic()
+            self.stats["timings"]["normal_filter"]["level1_filter"] = phase_end - phase_start
+
+            # Phase 4: Level 2 filtering - Related files
+            if target_files is not None and self.args.index_filter_level >= 2:
+                logger.info(
+                    "Phase 4: Performing Level 2 filtering (related files)...")
+                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_INDEX_FILTER_START.value,
+                            data=json.dumps({})
+                        )
+                    )
+                phase_start = time.monotonic()
+                related_files = self.index_manager.get_related_files(
+                    [file.file_path for file in target_files.file_list]
+                )
+                if related_files is not None:
+                    for file in related_files.file_list:
+                        file_path = file.file_path.strip()
+                        final_files[get_file_path(file_path)] = file
+                    self.stats["level2_filtered"] = len(related_files.file_list)
+                phase_end = time.monotonic()
+                self.stats["timings"]["normal_filter"]["level2_filter"] = phase_end - phase_start
+
+            if not final_files:
+                logger.warning("No related files found, using all files")
+                for source in self.sources:
+                    final_files[get_file_path(source.module_name)] = TargetFile(
+                        file_path=source.module_name,
+                        reason="No related files found, use all files",
+                    )
+
+            # Phase 5: Relevance verification
+            logger.info("Phase 5: Performing relevance verification...")
+            if self.args.index_filter_enable_relevance_verification:
+                phase_start = time.monotonic()
+                verified_files = {}
+                temp_files = list(final_files.values())
+                verification_results = []
+                
+                def print_verification_results(results):
+                    from rich.table import Table
+                    from rich.console import Console
+                    
+                    console = Console()
+                    table = Table(title="File Relevance Verification Results", show_header=True, header_style="bold magenta")
+                    table.add_column("File Path", style="cyan", no_wrap=True)
+                    table.add_column("Score", justify="right", style="green")
+                    table.add_column("Status", style="yellow")
+                    table.add_column("Reason/Error")
+                    
+                    for file_path, score, status, reason in results:
+                        table.add_row(
+                            file_path,
+                            str(score) if score is not None else "N/A",
+                            status,
+                            reason
+                        )
+                    
+                    console.print(table)
+
+                def verify_single_file(file: TargetFile):
+                    for source in self.sources:
+                        if source.module_name == file.file_path:
+                            file_content = source.source_code
+                            try:
+                                result = self.index_manager.verify_file_relevance.with_llm(llm).with_return_type(VerifyFileRelevance).run(
+                                    file_content=file_content,
+                                    query=self.args.query
+                                )
+                                if result.relevant_score >= self.args.verify_file_relevance_score:
+                                    verified_files[file.file_path] = TargetFile(
+                                        file_path=file.file_path,
+                                        reason=f"Score:{result.relevant_score}, {result.reason}"
+                                    )
+                                    return file.file_path, result.relevant_score, "PASS", result.reason
+                                else:
+                                    return file.file_path, result.relevant_score, "FAIL", result.reason
+                            except Exception as e:
+                                error_msg = str(e)
+                                verified_files[file.file_path] = TargetFile(
+                                    file_path=file.file_path,
+                                    reason=f"Verification failed: {error_msg}"
+                                )
+                                return file.file_path, None, "ERROR", error_msg
+                    return None
+
+                with ThreadPoolExecutor(max_workers=self.args.index_filter_workers) as executor:
+                    futures = [executor.submit(verify_single_file, file)
+                            for file in temp_files]
+                    for future in as_completed(futures):
+                        result = future.result()
+                        if result:
+                            verification_results.append(result)
+                            time.sleep(self.args.anti_quota_limit)
+
+                # Print verification results in a table
+                print_verification_results(verification_results)
+                
+                self.stats["verified_files"] = len(verified_files)
+                phase_end = time.monotonic()
+                self.stats["timings"]["normal_filter"]["relevance_verification"] = phase_end - phase_start
+
+                # Keep all files, not just verified ones
+                final_files = verified_files
+        
+        return final_files

autocoder/index/filter/quick_filter.py

@@ -0,0 +1,98 @@
+from typing import List, Union,Dict,Any
+from autocoder.index.types import IndexItem
+from autocoder.common import AutoCoderArgs,SourceCode
+import byzerllm
+import time
+from autocoder.index.index import IndexManager
+from autocoder.index.types import (
+    IndexItem,
+    TargetFile,    
+    FileNumberList
+)
+from autocoder.rag.token_counter import count_tokens
+from loguru import logger
+
+
+def get_file_path(file_path):
+    if file_path.startswith("##"):
+        return file_path.strip()[2:]
+    return file_path
+
+
+class QuickFilter():
+    def __init__(self, index_manager: IndexManager,stats:Dict[str,Any],sources:List[SourceCode]):
+        self.index_manager = index_manager
+        self.args = index_manager.args
+        self.stats = stats
+        self.sources = sources
+
+    @byzerllm.prompt()
+    def quick_filter_files(self,file_meta_list:List[IndexItem],query:str) -> str:
+        '''
+        当用户提一个需求的时候，我们需要找到相关的文件，然后阅读这些文件，并且修改其中部分文件。
+        现在，给定下面的索引文件：
+
+        <index>
+        {{ content }}
+        </index>
+
+        索引文件包含文件序号(##[]括起来的部分)，文件路径，文件符号信息等。
+        下面是用户的查询需求：
+        
+        <query>
+        {{ query }}
+        </query>
+
+        请根据用户的需求，找到相关的文件，并给出文件序号列表。请返回如下json格式：
+
+        ```json
+        {
+            "file_list": [
+                file_index1,
+                file_index2,
+                ...
+            ]
+        }
+        ```
+
+        特别注意,如果用户的query里 @文件 或者 @@符号，那么被@的文件或者@@的符号必须要返回，并且查看他们依赖的文件是否相关。        
+        '''
+        file_meta_str = "\n".join([f"##[{index}]{item.module_name}\n{item.symbols}" for index,item in enumerate(file_meta_list)])
+        context = {
+            "content": file_meta_str,
+            "query": query
+        }
+        return context    
+
+    def filter(self, index_items: List[IndexItem], query: str) -> Dict[str, TargetFile]:
+        final_files: Dict[str, TargetFile] = {}
+        if not self.args.skip_filter_index and self.args.index_filter_model:
+            start_time = time.monotonic()
+            index_items = self.index_manager.read_index()
+
+            prompt_str = self.quick_filter_files.prompt(index_items,query)
+
+            print(prompt_str)
+            
+            tokens_len = count_tokens(prompt_str)
+            
+            if tokens_len > 55*1024:
+                logger.warning(f"Quick filter prompt is too long, tokens_len: {tokens_len}/{55*1024} fallback to normal filter")
+                return final_files
+            
+            try:
+                file_number_list = self.quick_filter_files.with_llm(
+                    self.index_manager.index_filter_llm).with_return_type(FileNumberList).run(index_items, self.args.query)
+            except Exception as e:
+                logger.error(f"Quick filter failed, error: {str(e)} fallback to normal filter")
+                return final_files
+            
+            if file_number_list:
+                for file_number in file_number_list.file_list:
+                    final_files[get_file_path(index_items[file_number].module_name)] = TargetFile(
+                        file_path=index_items[file_number].module_name,
+                        reason="Quick Filter"
+                    )
+            end_time = time.monotonic()            
+            self.stats["timings"]["quick_filter"] = end_time - start_time            
+        return final_files