auto-coder 0.1.280__py3-none-any.whl → 0.1.282__py3-none-any.whl

autocoder/rag/cache/local_byzer_storage_cache.py

@@ -0,0 +1,457 @@
+from autocoder.rag.cache.base_cache import (
+    BaseCacheManager,
+    DeleteEvent,
+    AddOrUpdateEvent,
+    FileInfo,
+    CacheItem
+)
+from typing import Generator, List, Dict, Any, Optional, Tuple
+from autocoder.common import SourceCode
+from loguru import logger
+import pathspec
+import os
+import uuid
+import json
+from autocoder.rag.utils import process_file_in_multi_process, process_file_local
+from byzerllm.apps.byzer_storage.local_simple_api import (
+    LocalByzerStorage,
+    DataType,
+    FieldOption,
+    SortOption,
+)
+from autocoder.common import AutoCoderArgs
+import threading
+from multiprocessing import Pool
+from concurrent.futures import ThreadPoolExecutor, as_completed
+from autocoder.rag.variable_holder import VariableHolder
+import platform
+import hashlib
+from typing import Union
+from byzerllm import SimpleByzerLLM, ByzerLLM
+
+if platform.system() != "Windows":
+    import fcntl
+else:
+    fcntl = None
+
+
+def generate_file_md5(file_path: str) -> str:
+    md5_hash = hashlib.md5()
+    with open(file_path, "rb") as f:
+        for chunk in iter(lambda: f.read(4096), b""):
+            md5_hash.update(chunk)
+    return md5_hash.hexdigest()
+
+
+def generate_content_md5(content: Union[str, bytes]) -> str:
+    if isinstance(content, str):
+        content = content.encode("utf-8")
+    md5_hash = hashlib.md5()
+    md5_hash.update(content)
+    return md5_hash.hexdigest()
+
+
+default_ignore_dirs = [
+    "__pycache__",
+    "node_modules",
+    "_images"
+]
+
+
+class LocalByzerStorageCache(BaseCacheManager):
+    def __init__(
+        self,
+        path,
+        ignore_spec,
+        required_exts,
+        extra_params: Optional[AutoCoderArgs] = None,
+        emb_llm: Union[ByzerLLM, SimpleByzerLLM] = None,
+    ):
+        self.path = path
+        self.ignore_spec = ignore_spec
+        self.required_exts = required_exts
+        self.rag_build_name = extra_params.rag_build_name
+        self.storage = LocalByzerStorage("byzerai_store",
+            "rag_test", self.rag_build_name, host="127.0.0.1", port=33333,emb_llm=emb_llm)
+        self.queue = []
+        self.chunk_size = 1000
+        self._init_schema()
+
+        if not extra_params:
+            raise ValueError("extra_params is required for ByzerStorageCache")
+
+        self.max_output_tokens = extra_params.hybrid_index_max_output_tokens
+
+        # 设置缓存文件路径
+        self.cache_dir = os.path.join(self.path, ".cache")
+        self.cache_file = os.path.join(
+            self.cache_dir, "byzer_storage_speedup.jsonl")
+        self.cache: Dict[str, CacheItem] = {}
+
+        self.lock = threading.Lock()
+        self.stop_event = threading.Event()
+        self.thread = threading.Thread(target=self.process_queue)
+        self.thread.daemon = True
+        self.thread.start()
+
+        # 创建缓存目录
+        if not os.path.exists(self.cache_dir):
+            os.makedirs(self.cache_dir)
+
+        # 加载缓存
+        self.cache = self._load_cache()
+
+    def _chunk_text(self, text, max_length=1000):
+        """Split text into chunks"""
+        chunks = []
+        current_chunk = []
+        current_length = 0
+
+        for line in text.split("\n"):
+            if current_length + len(line) > max_length and current_chunk:
+                chunks.append("\n".join(current_chunk))
+                current_chunk = []
+                current_length = 0
+            current_chunk.append(line)
+            current_length += len(line)
+
+        if current_chunk:
+            chunks.append("\n".join(current_chunk))
+
+        return chunks
+
+    def _init_schema(self):
+        """Initialize the Byzer Storage schema"""
+        _ = (
+            self.storage.schema_builder()
+            .add_field("_id", DataType.STRING)
+            .add_field("file_path", DataType.STRING)
+            .add_field("content", DataType.STRING, [FieldOption.ANALYZE])
+            .add_field("raw_content", DataType.STRING, [FieldOption.NO_INDEX])
+            .add_array_field("vector", DataType.FLOAT)
+            .add_field("mtime", DataType.DOUBLE, [FieldOption.SORT])
+            .execute()
+        )
+
+    def _load_cache(self) -> Dict[str, CacheItem]:
+        """Load cache from file"""
+        if os.path.exists(self.cache_file):
+            try:
+                with open(self.cache_file, "r", encoding="utf-8") as f:
+                    lines = f.readlines()
+                    cache = {}
+                    for line in lines:
+                        try:
+                            data = json.loads(line.strip())
+                            if isinstance(data, dict) and "file_path" in data:
+                                # 转换为 CacheItem 对象
+                                cache_item = CacheItem.model_validate(data)
+                                cache[data["file_path"]] = cache_item
+                        except json.JSONDecodeError:
+                            continue
+                    return cache
+            except Exception as e:
+                logger.error(f"Error loading cache file: {str(e)}")
+                return {}
+        return {}
+
+    def write_cache(self):
+        cache_file = self.cache_file
+
+        if not fcntl:
+            try:
+                with open(cache_file, "w", encoding="utf-8") as f:
+                    for cache_item in self.cache.values():
+                        # 确保序列化 Pydantic 模型
+                        json.dump(cache_item.model_dump(),
+                                  f, ensure_ascii=False)
+                        f.write("\n")
+            except IOError as e:
+                logger.error(f"Error writing cache file: {str(e)}")
+        else:
+            lock_file = cache_file + ".lock"
+            with open(lock_file, "w", encoding="utf-8") as lockf:
+                try:
+                    # 获取文件锁
+                    fcntl.flock(lockf, fcntl.LOCK_EX | fcntl.LOCK_NB)
+                    # 写入缓存文件
+                    with open(cache_file, "w", encoding="utf-8") as f:
+                        for cache_item in self.cache.values():
+                            # 确保序列化 Pydantic 模型
+                            json.dump(cache_item.model_dump(),
+                                      f, ensure_ascii=False)
+                            f.write("\n")
+
+                finally:
+                    # 释放文件锁
+                    fcntl.flock(lockf, fcntl.LOCK_UN)
+
+    def fileinfo_to_tuple(self, file_info: FileInfo) -> Tuple[str, str, float, str]:
+        return (file_info.file_path, file_info.relative_path, file_info.modify_time, file_info.file_md5)
+
+    def build_cache(self):
+        """Build the cache by reading files and storing in Byzer Storage"""
+        logger.info(f"Building cache for path: {self.path}")
+
+        files_to_process = []
+        for file_info in self.get_all_files():
+            if (
+                file_info.file_path not in self.cache
+                or self.cache[file_info.file_path].md5 != file_info.file_md5
+            ):
+                files_to_process.append(file_info)
+
+        if not files_to_process:
+            return
+
+        from autocoder.rag.token_counter import initialize_tokenizer
+
+        with Pool(
+            processes=os.cpu_count(),
+            initializer=initialize_tokenizer,
+            initargs=(VariableHolder.TOKENIZER_PATH,),
+        ) as pool:
+            target_files_to_process = []
+            for file_info in files_to_process:
+                target_files_to_process.append(
+                    self.fileinfo_to_tuple(file_info))
+            results = pool.map(process_file_in_multi_process,
+                               target_files_to_process)
+
+        items = []
+        for file_info, result in zip(files_to_process, results):
+            content: List[SourceCode] = result
+            self.cache[file_info.file_path] = CacheItem(
+                file_path=file_info.file_path,
+                relative_path=file_info.relative_path,
+                content=[c.model_dump() for c in content],
+                modify_time=file_info.modify_time,
+                md5=file_info.file_md5,
+            )
+
+            for doc in content:
+                logger.info(f"Processing file: {doc.module_name}")
+                doc.module_name
+                chunks = self._chunk_text(doc.source_code, self.chunk_size)
+                for chunk_idx, chunk in enumerate(chunks):
+                    chunk_item = {
+                        "_id": f"{doc.module_name}_{chunk_idx}",
+                        "file_path": file_info.file_path,
+                        "content": chunk,
+                        "raw_content": chunk,
+                        "vector": chunk,
+                        "mtime": file_info.modify_time,
+                    }
+                    items.append(chunk_item)
+
+        # Save to local cache
+        logger.info("Saving cache to local file")
+        self.write_cache()
+
+        if items:
+            logger.info("Clear cache from Byzer Storage")
+            self.storage.truncate_table()
+            logger.info("Save new cache to Byzer Storage")
+            max_workers = 5
+            chunk_size = max(1, len(items) // max_workers)
+            item_chunks = [items[i:i + chunk_size]
+                           for i in range(0, len(items), chunk_size)]
+
+            total_chunks = len(item_chunks)
+            completed_chunks = 0
+
+            logger.info(f"Progress: {0}/{total_chunks} chunks completed")
+
+            with ThreadPoolExecutor(max_workers=max_workers) as executor:
+                futures = []
+                for chunk in item_chunks:
+                    futures.append(
+                        executor.submit(
+                            lambda x: self.storage.write_builder().add_items(
+                                x, vector_fields=["vector"], search_fields=["content"]
+                            ).execute(),
+                            chunk
+                        )
+                    )
+                # Wait for all futures to complete
+                for future in as_completed(futures):
+                    try:
+                        future.result()
+                        completed_chunks += 1
+                        logger.info(
+                            f"Progress: {completed_chunks}/{total_chunks} chunks completed")
+                    except Exception as e:
+                        logger.error(f"Error in saving chunk: {str(e)}")
+
+            self.storage.commit()
+
+    def update_storage(self, file_info: FileInfo, is_delete: bool):
+        query = self.storage.query_builder()
+        query.and_filter().add_condition("file_path", file_info.file_path).build()
+        results = query.execute()
+        if results:
+            for result in results:
+                self.storage.delete_by_ids([result["_id"]])
+        items = []
+
+        if not is_delete:
+            content = [
+                SourceCode.model_validate(doc)
+                for doc in self.cache[file_info.file_path].content
+            ]
+            modify_time = self.cache[file_info.file_path].modify_time
+            for doc in content:
+                logger.info(f"Processing file: {doc.module_name}")
+                doc.module_name
+                chunks = self._chunk_text(doc.source_code, self.chunk_size)
+                for chunk_idx, chunk in enumerate(chunks):
+                    chunk_item = {
+                        "_id": f"{doc.module_name}_{chunk_idx}",
+                        "file_path": file_info.file_path,
+                        "content": chunk,
+                        "raw_content": chunk,
+                        "vector": chunk,
+                        "mtime": modify_time,
+                    }
+                    items.append(chunk_item)
+        if items:
+            self.storage.write_builder().add_items(
+                items, vector_fields=["vector"], search_fields=["content"]
+            ).execute()
+            self.storage.commit()
+
+    def process_queue(self):
+        while self.queue:
+            file_list = self.queue.pop(0)
+            if isinstance(file_list, DeleteEvent):
+                for item in file_list.file_paths:
+                    logger.info(f"{item} is detected to be removed")
+                    del self.cache[item]
+                    # 创建一个临时的 FileInfo 对象
+                    file_info = FileInfo(
+                        file_path=item, relative_path="", modify_time=0, file_md5="")
+                    self.update_storage(file_info, is_delete=True)
+
+            elif isinstance(file_list, AddOrUpdateEvent):
+                for file_info in file_list.file_infos:
+                    logger.info(
+                        f"{file_info.file_path} is detected to be updated")
+                    # 处理文件并创建 CacheItem
+                    content = process_file_local(
+                        self.fileinfo_to_tuple(file_info))
+                    self.cache[file_info.file_path] = CacheItem(
+                        file_path=file_info.file_path,
+                        relative_path=file_info.relative_path,
+                        content=[c.model_dump() for c in content],
+                        modify_time=file_info.modify_time,
+                        md5=file_info.file_md5,
+                    )
+                    self.update_storage(file_info, is_delete=False)
+            self.write_cache()
+
+    def trigger_update(self):
+        logger.info("检查文件是否有更新.....")
+        files_to_process = []
+        current_files = set()
+        for file_info in self.get_all_files():
+            current_files.add(file_info.file_path)
+            if (
+                file_info.file_path not in self.cache
+                or self.cache[file_info.file_path].md5 != file_info.file_md5
+            ):
+                files_to_process.append(file_info)
+
+        deleted_files = set(self.cache.keys()) - current_files
+        logger.info(f"files_to_process: {files_to_process}")
+        logger.info(f"deleted_files: {deleted_files}")
+        if deleted_files:
+            with self.lock:
+                self.queue.append(DeleteEvent(file_paths=deleted_files))
+        if files_to_process:
+            with self.lock:
+                self.queue.append(AddOrUpdateEvent(
+                    file_infos=files_to_process))
+
+    def get_cache(self, options: Dict[str, Any]) -> Dict[str, Dict]:
+        """Search cached documents using query"""
+
+        self.trigger_update()
+
+        if options is None or "query" not in options:
+            return {file_path: self.cache[file_path].model_dump() for file_path in self.cache}
+
+        query = options.get("query", "")
+        total_tokens = 0
+
+        # Build query with both vector search and text search
+        query_builder = self.storage.query_builder()
+        query_builder.set_limit(100000)
+
+        # Add vector search if enabled
+        if options.get("enable_vector_search", True):
+            query_builder.set_vector_query(query, fields=["vector"])
+
+        # Add text search
+        if options.get("enable_text_search", True):
+            query_builder.set_search_query(query, fields=["content"])
+
+        results = query_builder.execute()
+
+        # Group results by file_path and reconstruct documents while preserving order
+        # 这里还可以有排序优化，综合考虑一篇内容出现的次数以及排序位置
+        file_paths = []
+        seen = set()
+        for result in results:
+            file_path = result["file_path"]
+            if file_path not in seen:
+                seen.add(file_path)
+                file_paths.append(file_path)
+
+        # 从缓存中获取文件内容
+        result = {}
+        for file_path in file_paths:
+            if file_path in self.cache:
+                cached_data = self.cache[file_path]
+                for doc in cached_data.content:
+                    if total_tokens + doc["tokens"] > self.max_output_tokens:
+                        return result
+                    total_tokens += doc["tokens"]
+                result[file_path] = cached_data.model_dump()
+
+        return result
+
+    def get_all_files(self) -> List[FileInfo]:
+        all_files = []
+        for root, dirs, files in os.walk(self.path, followlinks=True):
+            dirs[:] = [d for d in dirs if not d.startswith(
+                ".") and d not in default_ignore_dirs]
+
+            if self.ignore_spec:
+                relative_root = os.path.relpath(root, self.path)
+                dirs[:] = [
+                    d
+                    for d in dirs
+                    if not self.ignore_spec.match_file(os.path.join(relative_root, d))
+                ]
+                files = [
+                    f
+                    for f in files
+                    if not self.ignore_spec.match_file(os.path.join(relative_root, f))
+                ]
+
+            for file in files:
+                if self.required_exts and not any(
+                    file.endswith(ext) for ext in self.required_exts
+                ):
+                    continue
+
+                file_path = os.path.join(root, file)
+                relative_path = os.path.relpath(file_path, self.path)
+                modify_time = os.path.getmtime(file_path)
+                file_md5 = generate_file_md5(file_path)
+                all_files.append(FileInfo(file_path=file_path,
+                                          relative_path=relative_path,
+                                          modify_time=modify_time,
+                                          file_md5=file_md5))
+
+        return all_files

autocoder/rag/document_retriever.py

@@ -1,52 +1,20 @@
-import json
-import os
-import platform
-import time
-import traceback
-
 import threading
-from multiprocessing import Pool
-from typing import Dict, Generator, List, Tuple, Any, Optional
+from typing import Dict, Generator, List, Tuple, Any, Optional,Union
 
-import ray
-from loguru import logger
-from pydantic import BaseModel
+from byzerllm import ByzerLLM, SimpleByzerLLM
 
+from loguru import logger
 from autocoder.common import SourceCode
 from uuid import uuid4
-from autocoder.rag.variable_holder import VariableHolder
 from abc import ABC, abstractmethod
-from autocoder.rag.cache.base_cache import BaseCacheManager
 from autocoder.rag.cache.simple_cache import AutoCoderRAGAsyncUpdateQueue
 from autocoder.rag.cache.file_monitor_cache import AutoCoderRAGDocListener
 from autocoder.rag.cache.byzer_storage_cache import ByzerStorageCache
-from autocoder.rag.utils import process_file_in_multi_process, process_file_local
+from autocoder.rag.cache.local_byzer_storage_cache import LocalByzerStorageCache
 from autocoder.common import AutoCoderArgs
 
 cache_lock = threading.Lock()
 
-
-def get_or_create_actor(path: str, ignore_spec, required_exts: list, cacher={}):
-    with cache_lock:
-        # 处理路径名
-        actor_name = "AutoCoderRAGAsyncUpdateQueue_" + path.replace(
-            os.sep, "_"
-        ).replace(" ", "")
-        try:
-            actor = ray.get_actor(actor_name)
-        except ValueError:
-            actor = None
-        if actor is None:
-            actor = (
-                ray.remote(AutoCoderRAGAsyncUpdateQueue)
-                .options(name=actor_name, num_cpus=0)
-                .remote(path, ignore_spec, required_exts)
-            )
-            ray.get(actor.load_first.remote())
-        cacher[actor_name] = actor
-        return actor
-
-
 class BaseDocumentRetriever(ABC):
     """Abstract base class for document retrieval."""
 
@@ -77,6 +45,7 @@ class LocalDocumentRetriever(BaseDocumentRetriever):
         disable_auto_window: bool = False,
         enable_hybrid_index: bool = False,
         extra_params: Optional[AutoCoderArgs] = None,
+        emb_llm: Union[ByzerLLM, SimpleByzerLLM] = None,
     ) -> None:
         self.path = path
         self.ignore_spec = ignore_spec
@@ -91,23 +60,26 @@ class LocalDocumentRetriever(BaseDocumentRetriever):
         # 合并后的最大文件大小
         self.small_file_merge_limit = self.single_file_token_limit / 2
 
-        self.on_ray = on_ray
-        if self.on_ray:
-            self.cacher = get_or_create_actor(path, ignore_spec, required_exts)
-        else:
-            if self.enable_hybrid_index:
+        self.on_ray = on_ray        
+        if self.enable_hybrid_index:
+            if self.on_ray:
                 self.cacher = ByzerStorageCache(
                     path, ignore_spec, required_exts, extra_params
                 )
+            else:                                
+                self.cacher = LocalByzerStorageCache(
+                    path, ignore_spec, required_exts, extra_params,
+                    emb_llm = emb_llm
+                )
+        else:
+            if self.monitor_mode:
+                self.cacher = AutoCoderRAGDocListener(
+                    path, ignore_spec, required_exts
+                )
             else:
-                if self.monitor_mode:
-                    self.cacher = AutoCoderRAGDocListener(
-                        path, ignore_spec, required_exts
-                    )
-                else:
-                    self.cacher = AutoCoderRAGAsyncUpdateQueue(
-                        path, ignore_spec, required_exts
-                    )
+                self.cacher = AutoCoderRAGAsyncUpdateQueue(
+                    path, ignore_spec, required_exts
+                )
 
         logger.info(f"DocumentRetriever initialized with:")
         logger.info(f"  Path: {self.path}")
@@ -123,10 +95,7 @@ class LocalDocumentRetriever(BaseDocumentRetriever):
             )
 
     def get_cache(self, options: Optional[Dict[str, Any]] = None):
-        if self.on_ray:
-            return ray.get(self.cacher.get_cache.remote())
-        else:
-            return self.cacher.get_cache(options=options)
+        return self.cacher.get_cache(options=options)
 
     def retrieve_documents(
         self, options: Optional[Dict[str, Any]] = None

autocoder/rag/long_context_rag.py

@@ -86,6 +86,7 @@ class LongContextRAG:
         self.recall_llm = self.llm
         self.chunk_llm = self.llm
         self.qa_llm = self.llm
+        self.emb_llm = None
 
         if self.llm.get_sub_client("qa_model"):
             self.qa_llm = self.llm.get_sub_client("qa_model")
@@ -96,6 +97,9 @@ class LongContextRAG:
         if self.llm.get_sub_client("chunk_model"):
             self.chunk_llm = self.llm.get_sub_client("chunk_model")
 
+        if self.llm.get_sub_client("emb_model"):
+            self.emb_llm = self.llm.get_sub_client("emb_model")
+
         self.args = args
 
         self.path = path
@@ -169,6 +173,11 @@ class LongContextRAG:
 
         self.token_limit = self.args.rag_context_window_limit or 120000
         retriever_class = self._get_document_retriever_class()
+        
+        if self.args.enable_hybrid_index and not self.on_ray:
+            if self.emb_llm is None:
+                raise ValueError("emb_llm is required for local byzer storage cache")
+
         self.document_retriever = retriever_class(
             self.path,
             self.ignore_spec,
@@ -179,7 +188,8 @@ class LongContextRAG:
             single_file_token_limit=self.full_text_limit - 100,
             disable_auto_window=self.args.disable_auto_window,
             enable_hybrid_index=self.args.enable_hybrid_index,
-            extra_params=self.args
+            extra_params=self.args,
+            emb_llm=self.emb_llm
         )
 
         self.doc_filter = DocFilter(
@@ -778,6 +788,12 @@ class LongContextRAG:
                                             tokens=request_tokens
                                         )
                                         ))
+            
+            yield ("", SingleOutputMeta(input_tokens_count=rag_stat.recall_stat.total_input_tokens + rag_stat.chunk_stat.total_input_tokens,
+                                        generated_tokens_count=rag_stat.recall_stat.total_generated_tokens +
+                                        rag_stat.chunk_stat.total_generated_tokens,
+                                        reasoning_content="qa_model_thinking"
+                                        ))
 
             if LLMComputeEngine is not None and not self.args.disable_inference_enhance:
                 llm_compute_engine = LLMComputeEngine(
@@ -844,6 +860,7 @@ class LongContextRAG:
                         chunk[1].generated_tokens_count = rag_stat.recall_stat.total_generated_tokens + \
                             rag_stat.chunk_stat.total_generated_tokens + \
                             rag_stat.answer_stat.total_generated_tokens
+                        
                     yield chunk
 
                 self._print_rag_stats(rag_stat)

autocoder/version.py

@@ -1 +1 @@
-__version__ = "0.1.280"
+__version__ = "0.1.282"