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

autocoder/rag/long_context_rag.py

@@ -29,6 +29,9 @@ from autocoder.rag.searchable import SearchableResults
 from autocoder.rag.conversation_to_queries import extract_search_queries
 from autocoder.common import openai_content as OpenAIContentProcessor
 from autocoder.common.save_formatted_log import save_formatted_log
+from autocoder.rag.types import (
+    RecallStat,ChunkStat,AnswerStat,OtherStat,RAGStat
+)
 import json, os
 try:
     from autocoder_pro.rag.llm_compute import LLMComputeEngine
@@ -42,29 +45,6 @@ except ImportError:
     LLMComputeEngine = None
 
 
-class RecallStat(BaseModel):
-    total_input_tokens: int
-    total_generated_tokens: int
-    model_name: str = "unknown"
-
-
-class ChunkStat(BaseModel):
-    total_input_tokens: int
-    total_generated_tokens: int
-    model_name: str = "unknown"
-
-
-class AnswerStat(BaseModel):
-    total_input_tokens: int
-    total_generated_tokens: int
-    model_name: str = "unknown"
-
-
-class RAGStat(BaseModel):
-    recall_stat: RecallStat
-    chunk_stat: ChunkStat
-    answer_stat: AnswerStat
-
 
 class LongContextRAG:
     def __init__(
@@ -690,7 +670,7 @@ class LongContextRAG:
                             yield gen_item
                         
                         # 打印最终的统计信息
-                        self._print_rag_stats(rag_stat)
+                        self._print_rag_stats(rag_stat, conversations)
                         return
 
     def _process_document_retrieval(self, conversations, 
@@ -716,7 +696,7 @@ class LongContextRAG:
         
         # 提取查询并检索候选文档
         queries = extract_search_queries(
-            conversations=conversations, args=self.args, llm=self.llm, max_queries=self.args.rag_recall_max_queries)
+            conversations=conversations, args=self.args, llm=self.llm, max_queries=self.args.rag_recall_max_queries,rag_stat=rag_stat)
         documents = self._retrieve_documents(
             options={"queries": [query] + [query.query for query in queries]})
 
@@ -913,7 +893,7 @@ class LongContextRAG:
                         rag_stat.answer_stat.total_generated_tokens
                 yield chunk
 
-    def _print_rag_stats(self, rag_stat: RAGStat) -> None:
+    def _print_rag_stats(self, rag_stat: RAGStat, conversations: Optional[List[Dict[str, str]]] = None) -> None:
         """打印RAG执行的详细统计信息"""
         total_input_tokens = (
             rag_stat.recall_stat.total_input_tokens +
@@ -937,12 +917,46 @@ class LongContextRAG:
                              rag_stat.chunk_stat.total_generated_tokens) / total_tokens * 100
             answer_percent = (rag_stat.answer_stat.total_input_tokens +
                               rag_stat.answer_stat.total_generated_tokens) / total_tokens * 100
-
-        logger.info(
-            f"=== RAG 执行统计信息 ===\n"
+        
+        # 计算其他阶段的令牌占比
+        other_percents = []
+        if total_tokens > 0 and rag_stat.other_stats:
+            for other_stat in rag_stat.other_stats:
+                other_percent = (other_stat.total_input_tokens +
+                                other_stat.total_generated_tokens) / total_tokens * 100
+                other_percents.append(other_percent)
+        
+        # 计算成本分布百分比
+        if rag_stat.cost == 0:
+            recall_cost_percent = chunk_cost_percent = answer_cost_percent = 0
+        else:
+            recall_cost_percent = rag_stat.recall_stat.cost / rag_stat.cost * 100
+            chunk_cost_percent = rag_stat.chunk_stat.cost / rag_stat.cost * 100
+            answer_cost_percent = rag_stat.answer_stat.cost / rag_stat.cost * 100
+            
+        # 计算其他阶段的成本占比
+        other_costs_percent = []
+        if rag_stat.cost > 0 and rag_stat.other_stats:
+            for other_stat in rag_stat.other_stats:
+                other_costs_percent.append(other_stat.cost / rag_stat.cost * 100)
+
+        ## 这里会计算每个阶段的成本
+        estimated_cost = self._estimate_token_cost(rag_stat)
+        # 构建统计信息字符串
+        query_content = ""
+        if conversations and len(conversations) > 0:
+            query_content = conversations[-1].get("content", "")
+            if len(query_content) > 100:
+                query_content = query_content[:100] + "..."
+            query_content = f"查询内容: {query_content}\n"
+        
+        stats_str = (
+            f"=== (RAG 执行统计信息) ===\n"
+            f"{query_content}"
             f"总令牌使用: {total_tokens} 令牌\n"
             f"  * 输入令牌总数: {total_input_tokens}\n"
             f"  * 生成令牌总数: {total_generated_tokens}\n"
+            f"  * 总成本: {rag_stat.cost:.6f}\n"
             f"\n"
             f"阶段统计:\n"
             f"  1. 文档检索阶段:\n"
@@ -950,40 +964,146 @@ class LongContextRAG:
             f"     - 输入令牌: {rag_stat.recall_stat.total_input_tokens}\n"
             f"     - 生成令牌: {rag_stat.recall_stat.total_generated_tokens}\n"
             f"     - 阶段总计: {rag_stat.recall_stat.total_input_tokens + rag_stat.recall_stat.total_generated_tokens}\n"
+            f"     - 阶段成本: {rag_stat.recall_stat.cost:.6f}\n"
             f"\n"
             f"  2. 文档分块阶段:\n"
             f"     - 模型: {rag_stat.chunk_stat.model_name}\n"
             f"     - 输入令牌: {rag_stat.chunk_stat.total_input_tokens}\n"
             f"     - 生成令牌: {rag_stat.chunk_stat.total_generated_tokens}\n"
             f"     - 阶段总计: {rag_stat.chunk_stat.total_input_tokens + rag_stat.chunk_stat.total_generated_tokens}\n"
+            f"     - 阶段成本: {rag_stat.chunk_stat.cost:.6f}\n"
             f"\n"
             f"  3. 答案生成阶段:\n"
             f"     - 模型: {rag_stat.answer_stat.model_name}\n"
             f"     - 输入令牌: {rag_stat.answer_stat.total_input_tokens}\n"
             f"     - 生成令牌: {rag_stat.answer_stat.total_generated_tokens}\n"
             f"     - 阶段总计: {rag_stat.answer_stat.total_input_tokens + rag_stat.answer_stat.total_generated_tokens}\n"
+            f"     - 阶段成本: {rag_stat.answer_stat.cost:.6f}\n"
             f"\n"
+        )
+        
+        # 如果存在 other_stats，添加其统计信息
+        if rag_stat.other_stats:
+            for i, other_stat in enumerate(rag_stat.other_stats):
+                stats_str += (
+                    f"  {i+4}. 其他阶段 {i+1}:\n"
+                    f"     - 模型: {other_stat.model_name}\n"
+                    f"     - 输入令牌: {other_stat.total_input_tokens}\n"
+                    f"     - 生成令牌: {other_stat.total_generated_tokens}\n"
+                    f"     - 阶段总计: {other_stat.total_input_tokens + other_stat.total_generated_tokens}\n"
+                    f"     - 阶段成本: {other_stat.cost:.6f}\n"
+                    f"\n"
+                )
+        
+        # 添加令牌分布百分比
+        stats_str += (
             f"令牌分布百分比:\n"
             f"  - 文档检索: {recall_percent:.1f}%\n"
             f"  - 文档分块: {chunk_percent:.1f}%\n"
             f"  - 答案生成: {answer_percent:.1f}%\n"
         )
+        
+        # 如果存在 other_stats，添加其令牌占比
+        if rag_stat.other_stats:
+            for i, other_percent in enumerate(other_percents):
+                if other_percent > 0:
+                    stats_str += f"  - 其他阶段 {i+1}: {other_percent:.1f}%\n"
+        
+        # 添加成本分布百分比
+        stats_str += (
+            f"\n"
+            f"成本分布百分比:\n"
+            f"  - 文档检索: {recall_cost_percent:.1f}%\n"
+            f"  - 文档分块: {chunk_cost_percent:.1f}%\n"
+            f"  - 答案生成: {answer_cost_percent:.1f}%\n"
+        )
+        
+        # 如果存在 other_stats，添加其成本占比
+        if rag_stat.other_stats:
+            for i, other_cost_percent in enumerate(other_costs_percent):
+                if other_cost_percent > 0:
+                    stats_str += f"  - 其他阶段 {i+1}: {other_cost_percent:.1f}%\n"
+        
+        # 输出统计信息
+        logger.info(stats_str)
 
         # 记录原始统计数据，以便调试
         logger.debug(f"RAG Stat 原始数据: {rag_stat}")
 
-        # 返回成本估算
-        estimated_cost = self._estimate_token_cost(
-            total_input_tokens, total_generated_tokens)
+        
         if estimated_cost > 0:
-            logger.info(f"估计成本: 约 ${estimated_cost:.4f} 人民币")
+            logger.info(f"估计成本: 约 {estimated_cost:.4f} ")
 
-    def _estimate_token_cost(self, input_tokens: int, output_tokens: int) -> float:
+    def _estimate_token_cost(self, rag_stat: RAGStat) -> float:
         """估算当前请求的令牌成本（人民币）"""
-        # 实际应用中，可以根据不同模型设置不同价格
-        input_cost_per_1m = 2.0/1000000   # 每百万输入令牌的成本
-        output_cost_per_1m = 8.0/100000   # 每百万输出令牌的成本
-
-        cost = (input_tokens * input_cost_per_1m / 1000000) + \
-            (output_tokens * output_cost_per_1m/1000000)
-        return cost
+        from autocoder.models import get_model_by_name
+        
+        total_cost = 0.0
+        
+        # 计算召回阶段成本
+        if rag_stat.recall_stat.model_name != "unknown":
+            try:
+                recall_model = get_model_by_name(rag_stat.recall_stat.model_name)
+                input_cost = recall_model.get("input_price", 0.0) / 1000000
+                output_cost = recall_model.get("output_price", 0.0) / 1000000
+                recall_cost = (rag_stat.recall_stat.total_input_tokens * input_cost) + \
+                             (rag_stat.recall_stat.total_generated_tokens * output_cost)
+                total_cost += recall_cost
+            except Exception as e:
+                logger.warning(f"计算召回阶段成本时出错: {str(e)}")                
+                recall_cost = 0.0
+                total_cost += recall_cost
+            rag_stat.recall_stat.cost = recall_cost    
+        
+        # 计算分块阶段成本
+        if rag_stat.chunk_stat.model_name != "unknown":
+            try:
+                chunk_model = get_model_by_name(rag_stat.chunk_stat.model_name)
+                input_cost = chunk_model.get("input_price", 0.0) / 1000000
+                output_cost = chunk_model.get("output_price", 0.0) / 1000000
+                chunk_cost = (rag_stat.chunk_stat.total_input_tokens * input_cost) + \
+                            (rag_stat.chunk_stat.total_generated_tokens * output_cost)
+                total_cost += chunk_cost
+            except Exception as e:
+                logger.warning(f"计算分块阶段成本时出错: {str(e)}")
+                # 使用默认值
+                chunk_cost = 0.0
+                total_cost += chunk_cost
+            rag_stat.chunk_stat.cost = chunk_cost    
+        
+        # 计算答案生成阶段成本
+        if rag_stat.answer_stat.model_name != "unknown":
+            try:
+                answer_model = get_model_by_name(rag_stat.answer_stat.model_name)
+                input_cost = answer_model.get("input_price", 0.0) / 1000000
+                output_cost = answer_model.get("output_price", 0.0) / 1000000
+                answer_cost = (rag_stat.answer_stat.total_input_tokens * input_cost) + \
+                             (rag_stat.answer_stat.total_generated_tokens * output_cost)
+                total_cost += answer_cost
+            except Exception as e:
+                logger.warning(f"计算答案生成阶段成本时出错: {str(e)}")
+                # 使用默认值
+                answer_cost = 0.0
+                total_cost += answer_cost
+            rag_stat.answer_stat.cost = answer_cost
+            
+        # 计算其他阶段成本（如果存在）
+        for i, other_stat in enumerate(rag_stat.other_stats):
+            if other_stat.model_name != "unknown":
+                try:
+                    other_model = get_model_by_name(other_stat.model_name)
+                    input_cost = other_model.get("input_price", 0.0) / 1000000
+                    output_cost = other_model.get("output_price", 0.0) / 1000000
+                    other_cost = (other_stat.total_input_tokens * input_cost) + \
+                                (other_stat.total_generated_tokens * output_cost)
+                    total_cost += other_cost
+                except Exception as e:
+                    logger.warning(f"计算其他阶段 {i+1} 成本时出错: {str(e)}")
+                    # 使用默认值
+                    other_cost = 0.0
+                    total_cost += other_cost
+                rag_stat.other_stats[i].cost = other_cost
+        
+        # 将总成本保存到 rag_stat
+        rag_stat.cost = total_cost
+        return total_cost

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,163 @@
+"""
+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
+from autocoder.rag.types import 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,126 @@
+"""
+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
+from autocoder.rag.types import 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/rag/types.py

@@ -3,10 +3,46 @@ import os
 import json
 import time
 import pydantic
+from pydantic import BaseModel
 from typing import Dict, Any, Optional, List
 import psutil
 import glob
 
+class RecallStat(BaseModel):
+    total_input_tokens: int
+    total_generated_tokens: int
+    model_name: str = "unknown"
+    cost:float = 0.0
+
+
+class ChunkStat(BaseModel):
+    total_input_tokens: int
+    total_generated_tokens: int
+    model_name: str = "unknown"
+    cost:float = 0.0
+
+
+class AnswerStat(BaseModel):
+    total_input_tokens: int
+    total_generated_tokens: int
+    model_name: str = "unknown"
+    cost:float = 0.0
+
+
+class OtherStat(BaseModel):
+    total_input_tokens: int = 0
+    total_generated_tokens: int = 0
+    model_name: str = "unknown"
+    cost:float = 0.0
+
+
+class RAGStat(BaseModel):
+    recall_stat: RecallStat
+    chunk_stat: ChunkStat
+    answer_stat: AnswerStat
+    other_stats: List[OtherStat] = []
+    cost:float = 0.0
+
 class RAGServiceInfo(pydantic.BaseModel):
     host: str
     port: int