auto-coder 0.1.374__py3-none-any.whl → 0.1.375__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/tools/__init__.py

@@ -0,0 +1,10 @@
+# 导出 SearchTool 相关类和函数
+from .search_tool import SearchTool, SearchToolResolver, register_search_tool
+
+# 导出 RecallTool 相关类和函数
+from .recall_tool import RecallTool, RecallToolResolver, register_recall_tool
+
+__all__ = [
+    'SearchTool', 'SearchToolResolver', 'register_search_tool',
+    'RecallTool', 'RecallToolResolver', 'register_recall_tool'
+]

autocoder/rag/tools/recall_tool.py

@@ -0,0 +1,162 @@
+"""
+RecallTool 模块
+
+该模块实现了 RecallTool 和 RecallToolResolver 类，用于在 BaseAgent 框架中
+提供基于 LongContextRAG 的文档内容召回功能。
+"""
+
+import os
+import traceback
+from typing import Dict, Any, List, Optional
+
+from loguru import logger
+
+from autocoder.agent.base_agentic.types import BaseTool, ToolResult
+from autocoder.agent.base_agentic.tool_registry import ToolRegistry
+from autocoder.agent.base_agentic.tools.base_tool_resolver import BaseToolResolver
+from autocoder.agent.base_agentic.types import ToolDescription, ToolExample
+from autocoder.common import AutoCoderArgs
+from autocoder.rag.long_context_rag import LongContextRAG, RecallStat, ChunkStat, AnswerStat, RAGStat
+from autocoder.rag.relevant_utils import FilterDoc, DocRelevance, DocFilterResult
+from autocoder.common import SourceCode
+from autocoder.rag.relevant_utils import TaskTiming
+
+
+class RecallTool(BaseTool):
+    """召回工具，用于获取与查询相关的文档内容"""
+    query: str  # 用户查询
+    file_paths: Optional[List[str]] = None  # 指定要处理的文件路径列表，如果为空则自动搜索
+
+
+class RecallToolResolver(BaseToolResolver):
+    """召回工具解析器，实现召回逻辑"""
+    def __init__(self, agent, tool, args):
+        super().__init__(agent, tool, args)
+        self.tool: RecallTool = tool
+        
+    def resolve(self) -> ToolResult:
+        """实现召回工具的解析逻辑"""
+        try:
+            # 获取参数
+            query = self.tool.query                        
+            file_paths = self.tool.file_paths
+            rag:LongContextRAG = self.agent.rag                                    
+            # 构建对话历史
+            conversations = [
+                {"role": "user", "content": query}
+            ]
+            
+            # 创建 RAGStat 对象
+            
+            rag_stat = RAGStat(
+                recall_stat=RecallStat(total_input_tokens=0, total_generated_tokens=0),
+                chunk_stat=ChunkStat(total_input_tokens=0, total_generated_tokens=0),
+                answer_stat=AnswerStat(total_input_tokens=0, total_generated_tokens=0)
+            )
+            
+            # 如果提供了文件路径，则直接使用；否则，执行搜索
+            if file_paths:                
+                
+                # 创建 FilterDoc 对象
+                relevant_docs = []
+                for file_path in file_paths:
+                    try:
+                        with open(file_path, 'r', encoding='utf-8') as f:
+                            content = f.read()
+                        
+                        source_code = SourceCode(
+                            module_name=file_path,
+                            source_code=content
+                        )
+                        
+                        doc = FilterDoc(
+                            source_code=source_code,
+                            relevance=DocRelevance(is_relevant=True, relevant_score=5),  # 默认相关性
+                            task_timing=TaskTiming()
+                        )
+                        relevant_docs.append(doc)
+                    except Exception as e:
+                        logger.error(f"读取文件 {file_path} 失败: {str(e)}")
+            else:
+                # 调用文档检索处理
+                generator = rag._process_document_retrieval(conversations, query, rag_stat)
+                
+                # 获取检索结果
+                relevant_docs = None
+                for item in generator:
+                    if isinstance(item, dict) and "result" in item:
+                        relevant_docs = item["result"]
+                
+                if not relevant_docs:
+                    return ToolResult(
+                        success=False,
+                        message="未找到相关文档",
+                        content=[]
+                    )
+            
+            # 调用文档分块处理
+            relevant_docs = [doc.source_code for doc in relevant_docs]
+            doc_chunking_generator = rag._process_document_chunking(relevant_docs, conversations, rag_stat, 0)
+            
+            # 获取分块结果
+            final_relevant_docs = None
+            for item in doc_chunking_generator:
+                if isinstance(item, dict) and "result" in item:
+                    final_relevant_docs = item["result"]
+            
+            if not final_relevant_docs:
+                return ToolResult(
+                    success=False,
+                    message="文档分块处理失败",
+                    content=[]
+                )
+            
+            # 格式化结果
+            doc_contents = []
+            for doc in final_relevant_docs:
+                doc_contents.append({
+                    "path": doc.module_name,
+                    "content": doc.source_code                    
+                })
+            
+            return ToolResult(
+                success=True,
+                message=f"成功召回 {len(doc_contents)} 个相关文档片段",
+                content=doc_contents
+            )
+            
+        except Exception as e:
+            import traceback
+            return ToolResult(
+                success=False,
+                message=f"召回工具执行失败: {str(e)}",
+                content=traceback.format_exc()
+            )
+
+
+def register_recall_tool():
+    """注册召回工具"""
+    # 准备工具描述
+    description = ToolDescription(
+        description="召回与查询相关的文档内容",
+        parameters="query: 搜索查询\nfile_paths: 指定要处理的文件路径列表（可选）",
+        usage="用于根据查询获取相关文档的内容片段"
+    )
+    
+    # 准备工具示例
+    example = ToolExample(
+        title="召回工具使用示例",
+        body="""<recall>
+<query>如何实现文件监控功能</query>
+</recall>"""
+    )
+    
+    # 注册工具
+    ToolRegistry.register_tool(
+        tool_tag="recall",  # XML标签名
+        tool_cls=RecallTool,  # 工具类
+        resolver_cls=RecallToolResolver,  # 解析器类
+        description=description,  # 工具描述
+        example=example,  # 工具示例
+        use_guideline="此工具用于根据用户查询召回相关文档内容，返回经过分块和重排序的文档片段。适用于需要深入了解特定功能实现细节的场景。"  # 使用指南
+    )

autocoder/rag/tools/search_tool.py

@@ -0,0 +1,125 @@
+"""
+SearchTool 模块
+
+该模块实现了 SearchTool 和 SearchToolResolver 类，用于在 BaseAgent 框架中
+提供基于 LongContextRAG 的文档搜索功能。
+"""
+
+import os
+from typing import Dict, Any, List, Optional
+
+from loguru import logger
+
+from autocoder.agent.base_agentic.types import BaseTool, ToolResult
+from autocoder.agent.base_agentic.tool_registry import ToolRegistry
+from autocoder.agent.base_agentic.tools.base_tool_resolver import BaseToolResolver
+from autocoder.agent.base_agentic.types import ToolDescription, ToolExample
+from autocoder.common import AutoCoderArgs
+from autocoder.rag.long_context_rag import LongContextRAG, RecallStat, ChunkStat, AnswerStat, RAGStat
+from autocoder.rag.relevant_utils import FilterDoc, DocRelevance, DocFilterResult
+
+
+class SearchTool(BaseTool):
+    """搜索工具，用于获取与查询相关的文件列表"""
+    query: str  # 用户查询
+    max_files: Optional[int] = 10  # 最大返回文件数量
+
+
+class SearchToolResolver(BaseToolResolver):
+    """搜索工具解析器，实现搜索逻辑"""
+    def __init__(self, agent, tool, args):
+        super().__init__(agent, tool, args)
+        self.tool: SearchTool = tool
+        
+    def resolve(self) -> ToolResult:
+        """实现搜索工具的解析逻辑"""
+        try:
+            # 获取参数
+            query = self.tool.query            
+            max_files = self.tool.max_files            
+            rag = self.agent.rag            
+            # 构建对话历史
+            conversations = [
+                {"role": "user", "content": query}
+            ]
+            
+            # 创建 RAGStat 对象            
+            rag_stat = RAGStat(
+                recall_stat=RecallStat(total_input_tokens=0, total_generated_tokens=0),
+                chunk_stat=ChunkStat(total_input_tokens=0, total_generated_tokens=0),
+                answer_stat=AnswerStat(total_input_tokens=0, total_generated_tokens=0)
+            )
+            
+            # 调用文档检索处理
+            generator = rag._process_document_retrieval(conversations, query, rag_stat)
+            
+            # 获取最终结果
+            result = None
+            for item in generator:
+                if isinstance(item, dict) and "result" in item:
+                    result = item["result"]
+            
+            if not result:
+                return ToolResult(
+                    success=False,
+                    message="未找到相关文档",
+                    content=[]
+                )
+            
+            # 格式化结果
+            file_list = []
+            for doc in result:
+                file_list.append({
+                    "path": doc.source_code.module_name,
+                    "relevance": doc.relevance.relevant_score if doc.relevance else 0,
+                    "is_relevant": doc.relevance.is_relevant if doc.relevance else False
+                })
+            
+            # 按相关性排序
+            file_list.sort(key=lambda x: x["relevance"], reverse=True)
+            
+            # 限制返回数量
+            file_list = file_list[:max_files]
+            
+            return ToolResult(
+                success=True,
+                message=f"成功检索到 {len(file_list)} 个相关文件",
+                content=file_list
+            )
+            
+        except Exception as e:
+            import traceback
+            return ToolResult(
+                success=False,
+                message=f"搜索工具执行失败: {str(e)}",
+                content=traceback.format_exc()
+            )
+
+
+def register_search_tool():
+    """注册搜索工具"""
+    # 准备工具描述
+    description = ToolDescription(
+        description="搜索与查询相关的文件",
+        parameters="query: 搜索查询\nmax_files: 最大返回文件数量（可选，默认为10）",
+        usage="用于根据查询找到相关的代码文件"
+    )
+    
+    # 准备工具示例
+    example = ToolExample(
+        title="搜索工具使用示例",
+        body="""<search>
+<query>如何实现文件监控功能</query>
+<max_files>5</max_files>
+</search>"""
+    )
+    
+    # 注册工具
+    ToolRegistry.register_tool(
+        tool_tag="search",  # XML标签名
+        tool_cls=SearchTool,  # 工具类
+        resolver_cls=SearchToolResolver,  # 解析器类
+        description=description,  # 工具描述
+        example=example,  # 工具示例
+        use_guideline="此工具用于根据用户查询搜索相关代码文件，返回文件路径及其相关性分数。适用于需要快速找到与特定功能或概念相关的代码文件的场景。"  # 使用指南
+    )

autocoder/version.py

@@ -1,2 +1,2 @@
 
-__version__ = "0.1.374"
+__version__ = "0.1.375"