auto-coder 0.1.297__tar.gz → 0.1.299__tar.gz

{auto_coder-0.1.297 → auto_coder-0.1.299}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: auto-coder
-Version: 0.1.297
+Version: 0.1.299
 Summary: AutoCoder: AutoCoder
 Author: allwefantasy
 Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
@@ -26,7 +26,8 @@ Requires-Dist: tabulate
 Requires-Dist: jupyter_client
 Requires-Dist: prompt-toolkit
 Requires-Dist: tokenizers
-Requires-Dist: byzerllm[saas]>=0.1.174
+Requires-Dist: aiofiles
+Requires-Dist: byzerllm[saas]>=0.1.177
 Requires-Dist: patch
 Requires-Dist: diff_match_patch
 Requires-Dist: GitPython

{auto_coder-0.1.297 → auto_coder-0.1.299}/src/auto_coder.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: auto-coder
-Version: 0.1.297
+Version: 0.1.299
 Summary: AutoCoder: AutoCoder
 Author: allwefantasy
 Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
@@ -26,7 +26,8 @@ Requires-Dist: tabulate
 Requires-Dist: jupyter_client
 Requires-Dist: prompt-toolkit
 Requires-Dist: tokenizers
-Requires-Dist: byzerllm[saas]>=0.1.174
+Requires-Dist: aiofiles
+Requires-Dist: byzerllm[saas]>=0.1.177
 Requires-Dist: patch
 Requires-Dist: diff_match_patch
 Requires-Dist: GitPython

{auto_coder-0.1.297 → auto_coder-0.1.299}/src/auto_coder.egg-info/requires.txt

@@ -16,7 +16,8 @@ tabulate
 jupyter_client
 prompt-toolkit
 tokenizers
-byzerllm[saas]>=0.1.174
+aiofiles
+byzerllm[saas]>=0.1.177
 patch
 diff_match_patch
 GitPython

{auto_coder-0.1.297 → auto_coder-0.1.299}/src/autocoder/auto_coder_rag.py

@@ -187,6 +187,13 @@ def main(input_args: Optional[List[str]] = None):
         help="The storage type of the RAG, duckdb or byzer-storage",
     )
 
+    build_index_parser.add_argument(
+        "--rag_index_build_workers",
+        type=int,
+        default=5,
+        help="The number of workers to build the RAG index",
+    )
+
     build_index_parser.add_argument(
         "--quick", action="store_true", help="Skip system initialization"
     )

{auto_coder-0.1.297 → auto_coder-0.1.299}/src/autocoder/common/__init__.py

@@ -300,6 +300,7 @@ class AutoCoderArgs(pydantic.BaseModel):
     rag_duckdb_vector_dim: Optional[int] = 1024  # DuckDB 向量化存储的维度
     rag_duckdb_query_similarity: Optional[float] = 0.1  # DuckDB 向量化检索 相似度 阈值
     rag_duckdb_query_top_k: Optional[int] = 10000  # DuckDB 向量化检索 返回 TopK个结果(且大于相似度)
+    rag_index_build_workers: Optional[int] = 10
     # rag 本地图床地址
     local_image_host: Optional[str] = ""
     rag_recall_max_queries: Optional[int] = 5

{auto_coder-0.1.297 → auto_coder-0.1.299}/src/autocoder/rag/cache/byzer_storage_cache.py

@@ -132,6 +132,7 @@ class ByzerStorageCache(BaseCacheManager):
         self.path = path
         self.ignore_spec = ignore_spec
         self.required_exts = required_exts
+        self.extra_params = extra_params
         self.rag_build_name = extra_params.rag_build_name
         self.storage = ByzerStorage("byzerai_store", "rag", self.rag_build_name)
         self.queue = []
@@ -329,7 +330,7 @@ class ByzerStorageCache(BaseCacheManager):
             start_time = time.time()
             
             # Use more workers to process the smaller batches efficiently
-            max_workers = min(10, total_batches)  # Cap at 10 workers or total batch count
+            max_workers = min(self.extra_params.rag_index_build_workers, total_batches)  # Cap at 10 workers or total batch count
             logger.info(f"[BUILD CACHE] Using {max_workers} parallel workers for processing")
             
             with ThreadPoolExecutor(max_workers=max_workers) as executor:

{auto_coder-0.1.297 → auto_coder-0.1.299}/src/autocoder/rag/cache/local_byzer_storage_cache.py

@@ -338,7 +338,7 @@ class LocalByzerStorageCache(BaseCacheManager):
             start_time = time.time()
 
             # Use more workers to process the smaller batches efficiently
-            max_workers = min(10, total_batches)  # Cap at 10 workers or total batch count
+            max_workers = min(self.extra_params.rag_index_build_workers, total_batches)  # Cap at 10 workers or total batch count
             logger.info(f"[BUILD CACHE] Using {max_workers} parallel workers for processing")
 
             with ThreadPoolExecutor(max_workers=max_workers) as executor:

{auto_coder-0.1.297 → auto_coder-0.1.299}/src/autocoder/rag/cache/local_duckdb_storage_cache.py

@@ -133,15 +133,31 @@ class LocalDuckdbStorage:
         return reduced
 
     def _embedding(self, context: str, norm: bool = True, dim: int | None = None) -> List[float]:
-        embedding = self.llm.emb_query(context)[0].output
-
-        if dim:
-            embedding = self._apply_pca(embedding, target_dim=dim)  # 降维后形状 (1024,)
-
-        if norm:
-            embedding = embedding / np.linalg.norm(embedding)
-
-        return embedding.tolist()
+        max_retries = 3
+        retry_count = 0
+        
+        while retry_count < max_retries:
+            try:
+                embedding = self.llm.emb_query(context)[0].output
+                
+                if dim:
+                    embedding = self._apply_pca(embedding, target_dim=dim)  # 降维后形状 (1024,)
+                
+                if norm:
+                    embedding = embedding / np.linalg.norm(embedding)
+                
+                return embedding.tolist()
+            except Exception as e:
+                retry_count += 1
+                if retry_count >= max_retries:
+                    logger.error(f"Failed to get embedding after {max_retries} attempts: {str(e)}")
+                    raise
+                
+                # Sleep between 1-5 seconds before retrying
+                sleep_time = 1 + (retry_count * 1.5)
+                logger.warning(f"Embedding API call failed (attempt {retry_count}/{max_retries}). "
+                              f"Error: {str(e)}. Retrying in {sleep_time:.1f} seconds...")
+                time.sleep(sleep_time)
 
     def _initialize(self) -> None:
         if self.embed_dim is None:
@@ -452,9 +468,9 @@ class LocalDuckDBStorageCache(BaseCacheManager):
         self.write_cache()
 
         if items:
-            logger.info("[BUILD CACHE] Clearing existing cache from Byzer DuckDB Storage")
+            logger.info("[BUILD CACHE] Clearing existing cache from DuckDB  Storage")
             self.storage.truncate_table()
-            logger.info(f"[BUILD CACHE] Preparing to write to Byzer DuckDB Storage, "
+            logger.info(f"[BUILD CACHE] Preparing to write to DuckDB  Storage, "
                         f"total chunks: {len(items)}, total files: {len(files_to_process)}")
 
             # Use a fixed optimal batch size instead of dividing by worker count
@@ -464,12 +480,12 @@ class LocalDuckDBStorageCache(BaseCacheManager):
             total_batches = len(item_batches)
             completed_batches = 0
 
-            logger.info(f"[BUILD CACHE] Starting to write to Byzer Storage using {batch_size} items per batch, "
+            logger.info(f"[BUILD CACHE] Starting to write to DuckDB Storage using {batch_size} items per batch, "
                         f"total batches: {total_batches}")
             start_time = time.time()
 
-            # Use more workers to process the smaller batches efficiently
-            max_workers = min(10, total_batches)  # Cap at 10 workers or total batch count
+            # Use more workers to process the smaller batches efficiently            
+            max_workers = min(self.extra_params.rag_index_build_workers, total_batches)  # Cap at 10 workers or total batch count
             logger.info(f"[BUILD CACHE] Using {max_workers} parallel workers for processing")
 
             def batch_add_doc(_batch):
@@ -636,40 +652,68 @@ class LocalDuckDBStorageCache(BaseCacheManager):
 
         return all_files
 
-    def get_cache(self, options: Optional[Dict[str, Any]] = None) -> Dict[str, Dict]:
-        """Search cached documents using query"""
-        self.trigger_update()  # 检查更新
-
-        if options is None or "queries" not in options:
-            return {file_path: self.cache[file_path].model_dump() for file_path in self.cache}
-
-        queries = options.get("queries", "")
-        query = queries[0]
+    def _get_single_cache(self, query: str, options: Dict[str, Any]) -> List[Dict[str, Any]]:
+        """
+        使用单个查询检索缓存文档
+        
+        参数:
+            query: 查询字符串
+            options: 包含查询选项的字典
+            
+        返回:
+            包含文档信息的字典列表，每个字典包含_id、file_path、mtime和score字段
+        """
         logger.info(f"正在使用向量搜索检索数据, 你的问题: {query}")
-        total_tokens = 0
         results = []
 
         # Add vector search if enabled
         if options.get("enable_vector_search", True):
             # 返回值包含  [(_id, file_path, mtime, score,),]
-            # results = self.storage.vector_search(query, similarity_value=0.7, similarity_top_k=200)
             search_results = self.storage.vector_search(
                 query,
                 similarity_value=self.extra_params.rag_duckdb_query_similarity,
                 similarity_top_k=self.extra_params.rag_duckdb_query_top_k,
                 query_dim=self.extra_params.rag_duckdb_vector_dim
             )
-            results.extend(search_results)
-
+            
+            # Convert tuples to dictionaries for the merger
+            for _id, file_path, mtime, score in search_results:
+                results.append({
+                    "_id": _id,
+                    "file_path": file_path,
+                    "mtime": mtime,
+                    "score": score
+                })
+        
+        logger.info(f"查询 '{query}' 返回 {len(results)} 条记录")
+        return results
+        
+    def _process_search_results(self, results: List[Dict[str, Any]]) -> Dict[str, Dict]:
+        """
+        处理搜索结果，提取文件路径并构建结果字典
+        
+        参数:
+            results: 搜索结果列表，每项包含文档信息的字典
+            
+        返回:
+            匹配文档的字典，键为文件路径，值为文件内容
+            
+        说明:
+            该方法会根据查询结果从缓存中提取文件内容，并记录累计token数，
+            当累计token数超过max_output_tokens时，将停止处理并返回已处理的结果。
+        """
+        # 记录被处理的总tokens数
+        total_tokens = 0
+        
         # Group results by file_path and reconstruct documents while preserving order
         # 这里还可以有排序优化，综合考虑一篇内容出现的次数以及排序位置
         file_paths = []
         seen = set()
         for result in results:
-            _id, _file_path, _mtime, _score = result
-            if _file_path not in seen:
-                seen.add(_file_path)
-                file_paths.append(_file_path)
+            file_path = result["file_path"]
+            if file_path not in seen:
+                seen.add(file_path)
+                file_paths.append(file_path)
 
         # 从缓存中获取文件内容
         result = {}
@@ -690,3 +734,69 @@ class LocalDuckDBStorageCache(BaseCacheManager):
             f"累计tokens: {total_tokens}, "
             f"经过向量搜索共检索出 {len(result.keys())} 个文档, 共 {len(self.cache.keys())} 个文档")
         return result
+
+    def get_cache(self, options: Optional[Dict[str, Any]] = None) -> Dict[str, Dict]:
+        """
+        获取缓存中的文档信息
+        
+        参数:
+            options: 包含查询参数的字典，可以包含以下键：
+                - queries: 查询列表，可以是单个查询或多个查询
+                - enable_vector_search: 是否启用向量搜索，默认为True
+                - merge_strategy: 多查询时的合并策略，默认为WEIGHTED_RANK
+                - max_results: 最大结果数，默认为None表示不限制
+                
+        返回:
+            匹配文档的字典，键为文件路径，值为文件内容
+        """
+        self.trigger_update()  # 检查更新
+
+        if options is None or "queries" not in options:
+            return {file_path: self.cache[file_path].model_dump() for file_path in self.cache}
+
+        queries = options.get("queries", [])
+        
+        # 如果没有查询或只有一个查询，使用原来的方法
+        if not queries:
+            return {file_path: self.cache[file_path].model_dump() for file_path in self.cache}
+        elif len(queries) == 1:
+            results = self._get_single_cache(queries[0], options)
+            return self._process_search_results(results)
+        
+        # 导入合并策略
+        from autocoder.rag.cache.cache_result_merge import CacheResultMerger, MergeStrategy
+        
+        # 获取合并策略
+        merge_strategy_name = options.get("merge_strategy", MergeStrategy.WEIGHTED_RANK.value)
+        try:
+            merge_strategy = MergeStrategy(merge_strategy_name)
+        except ValueError:
+            logger.warning(f"未知的合并策略: {merge_strategy_name}, 使用默认策略 WEIGHTED_RANK")
+            merge_strategy = MergeStrategy.WEIGHTED_RANK
+        
+        # 限制最大结果数
+        max_results = options.get("max_results", None)
+        merger = CacheResultMerger(max_results=max_results)
+        
+        # 并发处理多个查询
+        logger.info(f"处理多查询请求，查询数量: {len(queries)}, 合并策略: {merge_strategy}")
+        query_results = []
+        with ThreadPoolExecutor(max_workers=min(len(queries), 10)) as executor:
+            future_to_query = {executor.submit(self._get_single_cache, query, options): query for query in queries}
+            for future in as_completed(future_to_query):
+                query = future_to_query[future]
+                try:
+                    query_result = future.result()
+                    logger.info(f"查询 '{query}' 返回 {len(query_result)} 条记录")
+                    query_results.append((query, query_result))
+                except Exception as e:
+                    logger.error(f"处理查询 '{query}' 时出错: {str(e)}")
+        
+        logger.info(f"所有查询共返回 {sum(len(r) for _, r in query_results)} 条记录")
+        
+        # 使用策略合并结果
+        merged_results = merger.merge(query_results, strategy=merge_strategy)
+        logger.info(f"合并后的结果共 {len(merged_results)} 条记录")
+        
+        # 处理合并后的结果
+        return self._process_search_results(merged_results)

auto_coder-0.1.299/src/autocoder/version.py

@@ -0,0 +1 @@
+__version__ = "0.1.299"

auto_coder-0.1.297/src/autocoder/version.py

@@ -1 +0,0 @@
-__version__ = "0.1.297"