auto-coder 0.1.374__py3-none-any.whl → 0.1.376__py3-none-any.whl

autocoder/rag/agentic_rag.py

@@ -0,0 +1,217 @@
+import json
+import os
+import time
+from typing import Any, Dict, Generator, List, Optional, Tuple, Union
+
+import pathspec
+from byzerllm import ByzerLLM
+import byzerllm
+from loguru import logger
+import traceback
+
+from autocoder.common import AutoCoderArgs, SourceCode
+from importlib.metadata import version
+from pydantic import BaseModel
+from autocoder.common import openai_content as OpenAIContentProcessor
+from autocoder.rag.long_context_rag import LongContextRAG
+import json, os
+from autocoder.agent.base_agentic.base_agent import BaseAgent
+from autocoder.agent.base_agentic.types import AgentRequest
+from autocoder.common import SourceCodeList
+from autocoder.rag.tools import register_search_tool, register_recall_tool
+from byzerllm.utils.types import SingleOutputMeta
+from autocoder.utils.llms import get_single_llm
+try:
+    from autocoder_pro.rag.llm_compute import LLMComputeEngine
+    pro_version = version("auto-coder-pro")
+    autocoder_version = version("auto-coder")
+    logger.warning(
+        f"auto-coder-pro({pro_version}) plugin is enabled in auto-coder.rag({autocoder_version})")
+except ImportError:
+    logger.warning(
+        "Please install auto-coder-pro to enhance llm compute ability")
+    LLMComputeEngine = None
+
+
+class RAGAgent(BaseAgent):    
+    def __init__(self, name: str, 
+        llm: Union[byzerllm.ByzerLLM, byzerllm.SimpleByzerLLM], 
+        files: SourceCodeList, 
+        args: AutoCoderArgs, 
+        rag: LongContextRAG,
+        conversation_history: Optional[List[Dict[str, Any]]] = None):
+
+        self.default_llm = self.llm
+        self.context_prune_llm = self.default_llm
+        if self.default_llm.get_sub_client("context_prune_model"):
+            self.context_prune_llm = self.default_llm.get_sub_client("context_prune_model")
+
+        self.llm = self.default_llm
+        if self.default_llm.get_sub_client("agentic_model"):
+            self.llm = self.default_llm.get_sub_client("agentic_model")
+        
+        self.rag = rag        
+        super().__init__(name, self.llm, files, args, conversation_history, default_tools_list=["read_file"])        
+        # 注册RAG工具
+        # register_search_tool()
+        register_recall_tool()
+
+class AgenticRAG:
+    def __init__(
+        self,
+        llm: ByzerLLM,
+        args: AutoCoderArgs,
+        path: str,
+        tokenizer_path: Optional[str] = None,
+    ) -> None:
+       self.llm = llm
+       self.args = args
+       self.path = path
+       self.tokenizer_path = tokenizer_path
+       self.rag = LongContextRAG(llm=self.llm, args=self.args, path=self.path, tokenizer_path=self.tokenizer_path) 
+
+
+    def build(self):
+        pass
+
+    def search(self, query: str) -> List[SourceCode]:
+        return []
+    
+
+    def stream_chat_oai(
+        self,
+        conversations,
+        model: Optional[str] = None,
+        role_mapping=None,
+        llm_config: Dict[str, Any] = {},
+        extra_request_params: Dict[str, Any] = {}
+    ):        
+        try:
+            return self._stream_chat_oai(
+                conversations,
+                model=model,
+                role_mapping=role_mapping,
+                llm_config=llm_config,
+                extra_request_params=extra_request_params
+            )
+        except Exception as e:
+            logger.error(f"Error in stream_chat_oai: {str(e)}")
+            traceback.print_exc()
+            return ["出现错误，请稍后再试。"], []
+
+    @byzerllm.prompt()
+    def conversation_to_query(self,messages: List[Dict[str, Any]]):
+        '''        
+        【历史对话】按时间顺序排列，从旧到新：
+        {% for message in messages %}
+        <message>
+        {% if message.role == "user" %}【用户】{% else %}【助手】{% endif %}    
+        <content>
+        {{ message.content }}
+        </content>
+        </message>
+        {% endfor %}
+        
+        【当前问题】用户的最新需求如下:
+        <current_query>
+        {{ query }}
+        </current_query>            
+        ''' 
+        temp_messages = messages[0:-1]
+        message = messages[-1]
+
+        return {
+            "messages": temp_messages,
+            "query":message["content"]
+        }
+
+
+    def system_prompt(self):
+        '''
+        你是一个基于知识库的智能助手，我的核心能力是通过检索增强生成（RAG）技术来回答用户问题。
+
+        你的工作流程如下：
+        1. 当用户提出问题时，我会首先理解问题的核心意图和关键信息需求
+        2. 你会从多个角度分析问题，确定最佳的检索策略和关键词，然后召回工具 recall 获取与问题最相关的详细内容，只有在特别有必要的情况下，你才回使用 read_file 来获得相关文件更详细的信息。
+        5. 如果获得的信息足够回答用户问题，你会直接生成回答。
+        6. 如果获得的信息不足以回答用户问题，你会继续使用 recall 工具，直到你确信已经获取了足够的信息来回答用户问题。
+        7. 有的问题可能需要拆解成多个问题，分别进行recall,然后最终得到的结果才是完整信息，最后才能进行回答。                
+
+        此外，你回答会遵循以下要求：
+
+        1. 严格基于召回的文档内容回答        
+        - 如果召回的文档提供的信息无法回答问题,请明确回复:"抱歉,文档中没有足够的信息来回答这个问题。" 
+        - 不要添加、推测或扩展文档未提及的信息
+
+        2. 格式如 ![image](/path/to/images/path.png) 的 Markdown 图片处理
+        - 根据Markdown 图片前后文本内容推测改图片与问题的相关性，有相关性则在回答中输出该Markdown图片路径
+        - 根据相关图片在文档中的位置，自然融入答复内容,保持上下文连贯
+        - 完整保留原始图片路径,不省略任何部分
+
+        3. 回答格式要求
+        - 使用markdown格式提升可读性        
+        {% if local_image_host %}
+        4. 图片路径处理
+        - 图片地址需返回绝对路径, 
+        - 对于Windows风格的路径，需要转换为Linux风格， 例如：C:\\Users\\user\\Desktop\\image.png 转换为 C:/Users/user/Desktop/image.png
+        - 为请求图片资源 需增加 http://{{ local_image_host }}/static/ 作为前缀
+        例如：/path/to/images/image.png， 返回 http://{{ local_image_host }}/static/path/to/images/image.png
+        {% endif %} 
+        '''    
+        return {
+            "local_image_host": self.args.local_image_host
+        }
+
+    
+    def _stream_chat_oai(
+        self,
+        conversations,
+        model: Optional[str] = None,
+        role_mapping=None,
+        llm_config: Dict[str, Any] = {},
+        extra_request_params: Dict[str, Any] = {}
+    ):
+        if not llm_config:
+            llm_config = {}
+        
+        if extra_request_params:
+            llm_config.update(extra_request_params)
+        
+        conversations = OpenAIContentProcessor.process_conversations(conversations)                
+
+        context = []
+
+        def _generate_sream():
+                           
+            recall_request = AgentRequest(user_input=self.conversation_to_query.prompt(conversations))              
+            rag_agent = RAGAgent(
+                name="RAGAgent",
+                llm=self.llm,
+                files=SourceCodeList(sources=[]),
+                args=self.args,
+                rag=self.rag,
+                conversation_history=[]
+            ) 
+
+            rag_agent.who_am_i(self.system_prompt.prompt())
+                                     
+            events =rag_agent.run_with_generator(recall_request)
+            for (t,content) in events:    
+                if t == "thinking":                          
+                    yield ("", SingleOutputMeta(
+                    generated_tokens_count=0,
+                    input_tokens_count=0,
+                    reasoning_content=content,                    
+                ))
+                else:
+                    yield (content, SingleOutputMeta(
+                        generated_tokens_count=0,
+                        input_tokens_count=0,
+                        reasoning_content="",                    
+                    ))
+
+        return _generate_sream(), context
+
+        
+
+    

autocoder/rag/cache/local_duckdb_storage_cache.py

@@ -10,7 +10,6 @@ from concurrent.futures import ThreadPoolExecutor, as_completed
 from typing import List, Dict, Any, Optional, Tuple, Union
 import numpy as np
 from loguru import logger
-from typing import Union
 from byzerllm import SimpleByzerLLM, ByzerLLM
 from autocoder.utils.llms import get_llm_names
 
@@ -31,10 +30,12 @@ from autocoder.rag.cache.base_cache import (
     FileInfo,
     CacheItem,
 )
-from autocoder.rag.utils import process_file_in_multi_process, process_file_local
+from autocoder.rag.utils import (
+    process_file_in_multi_process,
+    process_file_local,
+)
 from autocoder.rag.variable_holder import VariableHolder
-from byzerllm import SimpleByzerLLM, ByzerLLM
-from .failed_files_utils import save_failed_files, load_failed_files
+from .failed_files_utils import save_failed_files
 
 if platform.system() != "Windows":
     import fcntl
@@ -66,7 +67,8 @@ class DuckDBLocalContext:
     def __enter__(self) -> "duckdb.DuckDBPyConnection":
         if not os.path.exists(os.path.dirname(self.database_path)):
             raise ValueError(
-                f"Directory {os.path.dirname(self.database_path)} does not exist."
+                f"Directory {os.path.dirname(self.database_path)} "
+                f"does not exist."
             )
 
         self._conn = duckdb.connect(self.database_path)
@@ -97,10 +99,12 @@ class LocalDuckdbStorage:
         self.persist_dir = persist_dir
         self.cache_dir = os.path.join(self.persist_dir, ".cache")
         self.args = args
-        logger.info(f"正在启动 DuckDBVectorStore.")
+        logger.info("正在启动 DuckDBVectorStore.")
 
         if self.database_name != ":memory:":
-            self.database_path = os.path.join(self.cache_dir, self.database_name)
+            self.database_path = os.path.join(
+                self.cache_dir, self.database_name
+            )
 
         if self.database_name == ":memory:":
             self._conn = duckdb.connect(self.database_name)
@@ -114,7 +118,8 @@ class LocalDuckdbStorage:
             self._conn = None
         logger.info(
             f"DuckDBVectorStore 初始化完成, 存储目录: {self.cache_dir}, "
-            f"数据库名称: {self.database_name}, 数据表名称: {self.table_name}"
+            f"数据库名称: {self.database_name}, "
+            f"数据表名称: {self.table_name}"
         )
 
     @classmethod
@@ -167,15 +172,17 @@ class LocalDuckdbStorage:
                 retry_count += 1
                 if retry_count >= max_retries:
                     logger.error(
-                        f"Failed to get embedding after {max_retries} attempts: {str(e)}"
+                        f"Failed to get embedding after {max_retries} "
+                        f"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..."
+                    f"Embedding API call failed (attempt {retry_count}/"
+                    f"{max_retries}). Error: {str(e)}. Retrying in "
+                    f"{sleep_time:.1f} seconds..."
                 )
                 time.sleep(sleep_time)
 
@@ -244,7 +251,9 @@ class LocalDuckdbStorage:
         
         if not context_chunk["raw_content"]:
             context_chunk["raw_content"] = "empty"        
-        context_chunk["raw_content"] = context_chunk["raw_content"][: self.args.rag_emb_text_size]
+        context_chunk["raw_content"] = context_chunk["raw_content"][
+            : self.args.rag_emb_text_size
+        ]
             
         return (
             context_chunk["_id"],
@@ -343,11 +352,14 @@ class LocalDuckDBStorageCache(BaseCacheManager):
         )
         self.queue = []
         self.chunk_size = 1000
-        self.max_output_tokens = extra_params.hybrid_index_max_output_tokens
+        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, "duckdb_storage_speedup.jsonl")
+        self.cache_file = os.path.join(self.cache_dir,
+                                       "duckdb_storage_speedup.jsonl")
         self.cache: Dict[str, CacheItem] = {}
         # 创建缓存目录
         if not os.path.exists(self.cache_dir):
@@ -356,7 +368,9 @@ class LocalDuckDBStorageCache(BaseCacheManager):
         # failed files support
         from .failed_files_utils import load_failed_files
 
-        self.failed_files_path = os.path.join(self.cache_dir, "failed_files.json")
+        self.failed_files_path = os.path.join(
+            self.cache_dir, "failed_files.json"
+        )
         self.failed_files = load_failed_files(self.failed_files_path)
 
         self.lock = threading.Lock()
@@ -406,7 +420,8 @@ class LocalDuckDBStorageCache(BaseCacheManager):
                             continue
                     return cache
             except Exception as e:
-                logger.error(f"Error loading cache file: {str(e)}")
+                logger.warning(f"Error loading cache file: {str(e)}")
+                logger.exception(e)
                 return {}
         return {}
 
@@ -421,7 +436,8 @@ class LocalDuckDBStorageCache(BaseCacheManager):
                         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)}")
+                logger.warning(f"Error writing cache file: {str(e)}")
+                logger.exception(e)
         else:
             lock_file = cache_file + ".lock"
             with open(lock_file, "w", encoding="utf-8") as lockf:
@@ -510,11 +526,12 @@ class LocalDuckDBStorageCache(BaseCacheManager):
         self.write_cache()
 
         if items:
-            logger.info("[BUILD CACHE] Clearing existing cache from DuckDB  Storage")
+            logger.info("[BUILD CACHE] Clearing DuckDB Storage cache")
             self.storage.truncate_table()
+            logger.info(f"[BUILD CACHE] Preparing to write to DuckDB Storage.")
             logger.info(
-                f"[BUILD CACHE] Preparing to write to DuckDB  Storage, "
-                f"total chunks: {len(items)}, total files: {len(files_to_process)}"
+                f"[BUILD CACHE] Total chunks: {len(items)}, "
+                f"Total files: {len(files_to_process)}"
             )
 
             # Use a fixed optimal batch size instead of dividing by worker count
@@ -526,9 +543,10 @@ class LocalDuckDBStorageCache(BaseCacheManager):
             total_batches = len(item_batches)
             completed_batches = 0
 
+            logger.info(f"[BUILD CACHE] Writing to DuckDB Storage.")
             logger.info(
-                f"[BUILD CACHE] Starting to write to DuckDB Storage using {batch_size} items per batch, "
-                f"total batches: {total_batches}"
+                f"[BUILD CACHE] Batch size: {batch_size}, "
+                f"Total batches: {total_batches}"
             )
             start_time = time.time()
 
@@ -569,18 +587,26 @@ class LocalDuckDBStorageCache(BaseCacheManager):
                             or (completed_batches == total_batches)
                             or (completed_batches % max(1, total_batches // 10) == 0)
                         ):
+                            progress_percent = (
+                                completed_batches / total_batches * 100
+                                if total_batches > 0
+                                else 0
+                            )
                             logger.info(
-                                f"[BUILD CACHE] Progress: {completed_batches}/{total_batches} batches completed "
-                                f"({(completed_batches / total_batches * 100):.1f}%) "
-                                f"Estimated time remaining: {remaining:.1f}s"
+                                f"[BUILD CACHE] Progress: {completed_batches}/"
+                                f"{total_batches} ({progress_percent:.1f}%). "
+                                f"ETA: {remaining:.1f}s"
                             )
                     except Exception as e:
                         logger.error(f"[BUILD CACHE] Error saving batch: {str(e)}")
                         # Add more detailed error information
+                        batch_len_info = (
+                            len(batch) if "batch" in locals() else "unknown"
+                        )
                         logger.error(
-                            f"[BUILD CACHE] Error details: batch size: "
-                            f"{len(batch) if 'batch' in locals() else 'unknown'}"
+                            f"[BUILD CACHE] Error details: batch size: {batch_len_info}"
                         )
+                        logger.exception(e)
 
             total_time = time.time() - start_time
             logger.info(
@@ -622,6 +648,7 @@ class LocalDuckDBStorageCache(BaseCacheManager):
                     time.sleep(self.extra_params.anti_quota_limit)
                 except Exception as err:
                     logger.error(f"Error in saving chunk: {str(err)}")
+                    logger.exception(err)
 
     def process_queue(self):
         while self.queue:
@@ -671,7 +698,8 @@ class LocalDuckDBStorageCache(BaseCacheManager):
                             self.failed_files.add(file_info.file_path)
                             save_failed_files(self.failed_files_path, self.failed_files)
                     except Exception as e:
-                        logger.error(f"Error in process_queue: {e}")
+                        logger.error(f"Error in process_queue: {str(e)}")
+                        logger.exception(e)
                         self.failed_files.add(file_info.file_path)
                         save_failed_files(self.failed_files_path, self.failed_files)
 
@@ -816,17 +844,18 @@ class LocalDuckDBStorageCache(BaseCacheManager):
                 for doc in cached_data.content:
                     if total_tokens + doc["tokens"] > self.max_output_tokens:
                         logger.info(
-                            f"当前检索已超出用户设置 Hybrid Index Max Tokens:{self.max_output_tokens}，"
-                            f"累计tokens: {total_tokens}, "
-                            f"经过向量搜索共检索出 {len(result.keys())} 个文档, 共 {len(self.cache.keys())} 个文档"
+                            f"当前检索已超出用户设置 Hybrid Index Max Tokens:"
+                            f"{self.max_output_tokens}，累计tokens: {total_tokens}, "
+                            f"经过向量搜索共检索出 {len(result.keys())} 个文档, "
+                            f"共 {len(self.cache.keys())} 个文档"
                         )
                         return result
                     total_tokens += doc["tokens"]
                 result[file_path] = cached_data.model_dump()
         logger.info(
             f"用户Hybrid Index Max Tokens设置为:{self.max_output_tokens}，"
-            f"累计tokens: {total_tokens}, "
-            f"经过向量搜索共检索出 {len(result.keys())} 个文档, 共 {len(self.cache.keys())} 个文档"
+            f"累计tokens: {total_tokens}, 经过向量搜索共检索出 "
+            f"{len(result.keys())} 个文档, 共 {len(self.cache.keys())} 个文档"
         )
         return result
 
@@ -904,6 +933,7 @@ class LocalDuckDBStorageCache(BaseCacheManager):
                     query_results.append((query, query_result))
                 except Exception as e:
                     logger.error(f"处理查询 '{query}' 时出错: {str(e)}")
+                    logger.exception(e)
 
         logger.info(f"所有查询共返回 {sum(len(r) for _, r in query_results)} 条记录")
 

autocoder/rag/conversation_to_queries.py

@@ -2,7 +2,9 @@ from typing import List, Dict, Any, Optional, Union
 import logging
 import byzerllm
 from pydantic import BaseModel
+from autocoder.rag.types import RAGStat
 from autocoder.common import AutoCoderArgs
+from byzerllm import MetaHolder
 
 logger = logging.getLogger(__name__)
 
@@ -86,7 +88,7 @@ class ConversationToQueries:
         ```
         """
     
-    def extract_queries(self, conversations: List[Dict[str, Any]], max_queries: int = 3) -> List[SearchQuery]:
+    def extract_queries(self, conversations: List[Dict[str, Any]], max_queries: int = 3,rag_stat:Optional[RAGStat] = None) -> List[SearchQuery]:
         """
         从对话历史中提取搜索查询。
         
@@ -99,10 +101,39 @@ class ConversationToQueries:
         """
         try:
             # 使用 prompt 函数生成搜索查询
-            queries = self.generate_search_queries.with_llm(self.llm).with_return_type(SearchQuery).run(
+            model_name = self.llm.default_model_name
+            meta_holder = MetaHolder()
+            queries = self.generate_search_queries.with_llm(self.llm).with_return_type(SearchQuery).with_meta(
+                        meta_holder).run(
                 conversations=conversations,
                 max_queries=max_queries
-            )            
+            ) 
+
+            # 如果有元数据且有 rag_stat，则记录模型使用情况
+            if meta_holder.get_meta() and rag_stat:
+                meta_dict = meta_holder.get_meta()                    
+                input_tokens_count = meta_dict.get("input_tokens_count", 0) 
+                generated_tokens_count = meta_dict.get("generated_tokens_count", 0)
+                
+                # 检查模型是否已存在于 other_stats 中
+                found = False
+                for other_stat in rag_stat.other_stats:
+                    if other_stat.model_name == model_name:
+                        # 模型已存在，累加统计数据
+                        other_stat.total_input_tokens += input_tokens_count
+                        other_stat.total_generated_tokens += generated_tokens_count
+                        found = True
+                        break
+                
+                # 如果模型不存在，添加新的 OtherStat
+                if not found and (input_tokens_count > 0 or generated_tokens_count > 0):
+                    from autocoder.rag.types import OtherStat
+                    new_stat = OtherStat(
+                        total_input_tokens=input_tokens_count,
+                        total_generated_tokens=generated_tokens_count,
+                        model_name=model_name
+                    )
+                    rag_stat.other_stats.append(new_stat)
             
             # 按重要性排序
             queries.sort(key=lambda x: x.importance, reverse=True)
@@ -116,7 +147,8 @@ def extract_search_queries(
     conversations: List[Dict[str, Any]], 
     args:AutoCoderArgs,
     llm: Union[byzerllm.ByzerLLM, byzerllm.SimpleByzerLLM],
-    max_queries: int = 3,    
+    max_queries: int = 3,  
+    rag_stat:Optional[RAGStat] = None  
 ) -> List[SearchQuery]:
     """
     从对话历史中提取搜索查询的便捷函数。
@@ -133,7 +165,7 @@ def extract_search_queries(
         return []
     try:    
         extractor = ConversationToQueries(llm)
-        return extractor.extract_queries(conversations, max_queries) 
+        return extractor.extract_queries(conversations, max_queries,rag_stat) 
     except Exception as e:
         logger.error(f"Error extracting search queries from conversation: {str(e)}")
         return []