auto-coder 0.1.270__py3-none-any.whl → 0.1.272__py3-none-any.whl

autocoder/rag/doc_filter.py

@@ -5,8 +5,9 @@ from concurrent.futures import ThreadPoolExecutor, as_completed
 
 from autocoder.rag.relevant_utils import (
     parse_relevance,
-    FilterDoc,    
+    FilterDoc,
     TaskTiming,
+    DocFilterResult
 )
 
 from autocoder.common import SourceCode, AutoCoderArgs
@@ -48,7 +49,6 @@ def _check_relevance_with_conversation(
     其中， <relevant> 是你认为文档中和问题的相关度，0-10之间的数字，数字越大表示相关度越高。
     """
 
-
 class DocFilter:
     def __init__(
         self,
@@ -62,40 +62,57 @@ class DocFilter:
             self.recall_llm = self.llm.get_sub_client("recall_model")
         else:
             self.recall_llm = self.llm
-        
+
         self.args = args
         self.relevant_score = self.args.rag_doc_filter_relevance
         self.on_ray = on_ray
-        self.path = path        
+        self.path = path
 
     def filter_docs(
         self, conversations: List[Dict[str, str]], documents: List[SourceCode]
-    ) -> List[FilterDoc]:
-        return self.filter_docs_with_threads(conversations, documents)            
+    ) -> DocFilterResult:
+        return self.filter_docs_with_threads(conversations, documents)
 
     def filter_docs_with_threads(
         self, conversations: List[Dict[str, str]], documents: List[SourceCode]
-    ) -> List[FilterDoc]:
-        
+    ) -> DocFilterResult:
+
+        start_time = time.time()
+        logger.info(f"=== DocFilter Starting ===")
+        logger.info(
+            f"Configuration: relevance_threshold={self.relevant_score}, thread_workers={self.args.index_filter_workers or 5}")
+
         rag_manager = RagConfigManager(path=self.path)
         rag_config = rag_manager.load_config()
-        documents = list(documents)   
-        logger.info(f"Filtering {len(documents)} documents....")
+
+        documents = list(documents)
+        logger.info(f"Filtering {len(documents)} documents...")
+
+        submitted_tasks = 0
+        completed_tasks = 0
+        relevant_count = 0
+        model_name = self.recall_llm.default_model_name or "unknown"
+
         with ThreadPoolExecutor(
             max_workers=self.args.index_filter_workers or 5
         ) as executor:
             future_to_doc = {}
+
+            # 提交所有任务
             for doc in documents:
                 submit_time = time.time()
+                submitted_tasks += 1
 
                 def _run(conversations, docs):
                     submit_time_1 = time.time()
+                    meta = None
                     try:
                         llm = self.recall_llm
+                        meta_holder = byzerllm.MetaHolder()
 
                         v = (
                             _check_relevance_with_conversation.with_llm(
-                                llm)
+                                llm).with_meta(meta_holder)
                             .options({"llm_config": {"max_length": 10}})
                             .run(
                                 conversations=conversations,
@@ -103,14 +120,16 @@ class DocFilter:
                                 filter_config=rag_config.filter_config,
                             )
                         )
+
+                        meta = meta_holder.get_meta_model()
                     except Exception as e:
                         logger.error(
                             f"Error in _check_relevance_with_conversation: {str(e)}"
                         )
-                        return (None, submit_time_1, time.time())
+                        return (None, submit_time_1, time.time(), meta)
 
                     end_time_2 = time.time()
-                    return (v, submit_time_1, end_time_2)
+                    return (v, submit_time_1, end_time_2, meta)
 
                 m = executor.submit(
                     _run,
@@ -119,57 +138,144 @@ class DocFilter:
                 )
                 future_to_doc[m] = (doc, submit_time)
 
-        relevant_docs = []
-        for future in as_completed(list(future_to_doc.keys())):
-            try:
-                doc, submit_time = future_to_doc[future]
-                end_time = time.time()
-                v, submit_time_1, end_time_2 = future.result()
-                task_timing = TaskTiming(
-                    submit_time=submit_time,
-                    end_time=end_time,
-                    duration=end_time - submit_time,
-                    real_start_time=submit_time_1,
-                    real_end_time=end_time_2,
-                    real_duration=end_time_2 - submit_time_1,
-                )                
-
-                relevance = parse_relevance(v)
-                logger.info(
-                    f"Document filtering progress:\n"
-                    f"  - File: {doc.module_name}\n"
-                    f"  - Relevance: {'Relevant' if relevance and relevance.is_relevant else 'Not Relevant'}\n"
-                    f"  - Score: {relevance.relevant_score if relevance else 'N/A'}\n"
-                    f"  - Score Threshold: {self.relevant_score}\n"
-                    f"  - Raw Response: {v}\n"
-                    f"  - Timing:\n"
-                    f"    * Total Duration: {task_timing.duration:.2f}s\n"
-                    f"    * Real Duration: {task_timing.real_duration:.2f}s\n"
-                    f"    * Queue Time: {(task_timing.real_start_time - task_timing.submit_time):.2f}s"
-                )
-                if (
-                    relevance
-                    # and relevance.is_relevant
-                    and relevance.relevant_score >= self.relevant_score
-                ):
-                    relevant_docs.append(
-                        FilterDoc(
+            logger.info(
+                f"Submitted {submitted_tasks} document filtering tasks to thread pool")
+
+            # 处理完成的任务
+            doc_filter_result = DocFilterResult(
+                docs=[],
+                raw_docs=[],
+                input_tokens_counts=[],
+                generated_tokens_counts=[],
+                durations=[],
+                model_name=model_name
+            )            
+            relevant_docs = doc_filter_result.docs
+            for future in as_completed(list(future_to_doc.keys())):
+                try:
+                    doc, submit_time = future_to_doc[future]
+                    end_time = time.time()
+                    completed_tasks += 1
+                    progress_percent = (completed_tasks / len(documents)) * 100
+
+                    v, submit_time_1, end_time_2, meta = future.result()
+                    task_timing = TaskTiming(
+                        submit_time=submit_time,
+                        end_time=end_time,
+                        duration=end_time - submit_time,
+                        real_start_time=submit_time_1,
+                        real_end_time=end_time_2,
+                        real_duration=end_time_2 - submit_time_1,
+                    )
+
+                    relevance = parse_relevance(v)
+                    is_relevant = relevance and relevance.relevant_score >= self.relevant_score
+
+                    if is_relevant:
+                        relevant_count += 1
+                        status_text = f"RELEVANT (Score: {relevance.relevant_score:.1f})"
+                    else:
+                        score_text = f"{relevance.relevant_score:.1f}" if relevance else "N/A"
+                        status_text = f"NOT RELEVANT (Score: {score_text})"
+
+                    queue_time = task_timing.real_start_time - task_timing.submit_time
+
+                    input_tokens_count = meta.input_tokens_count if meta else 0
+                    generated_tokens_count = meta.generated_tokens_count if meta else 0
+
+                    logger.info(
+                        f"Document filtering [{progress_percent:.1f}%] - {completed_tasks}/{len(documents)}:"
+                        f"\n  - File: {doc.module_name}"
+                        f"\n  - Status: {status_text}"
+                        f"\n  - Model: {model_name}"
+                        f"\n  - Threshold: {self.relevant_score}"
+                        f"\n  - Input tokens: {input_tokens_count}"
+                        f"\n  - Generated tokens: {generated_tokens_count}"
+                        f"\n  - Timing: Duration={task_timing.duration:.2f}s, Processing={task_timing.real_duration:.2f}s, Queue={queue_time:.2f}s"
+                        f"\n  - Response: {v}"
+                    )
+                    
+                    if "rag" not in doc.metadata:
+                        doc.metadata["rag"] = {}
+                    doc.metadata["rag"]["recall"] = {
+                        "input_tokens_count": input_tokens_count,
+                        "generated_tokens_count": generated_tokens_count,
+                        "recall_model": model_name,
+                        "duration": task_timing.real_duration                                                                                                            
+                    }
+                    
+                    doc_filter_result.input_tokens_counts.append(input_tokens_count)
+                    doc_filter_result.generated_tokens_counts.append(generated_tokens_count)
+                    doc_filter_result.durations.append(task_timing.real_duration)
+                    
+                    new_filter_doc = FilterDoc(
                             source_code=doc,
                             relevance=relevance,
                             task_timing=task_timing,
                         )
-                    )
-            except Exception as exc:
-                try:
-                    doc, submit_time = future_to_doc[future]
-                    logger.error(
-                        f"Filtering document generated an exception (doc: {doc.module_name}): {exc}")
-                except Exception as e:
-                    logger.error(
-                        f"Filtering document generated an exception: {exc}")
+                    
+                    doc_filter_result.raw_docs.append(new_filter_doc)
+
+                    if is_relevant:
+                        relevant_docs.append(
+                            new_filter_doc
+                        )
+                except Exception as exc:
+                    try:
+                        doc, submit_time = future_to_doc[future]
+                        completed_tasks += 1
+                        progress_percent = (
+                            completed_tasks / len(documents)) * 100
+                        logger.error(
+                            f"Document filtering [{progress_percent:.1f}%] - {completed_tasks}/{len(documents)}:"
+                            f"\n  - File: {doc.module_name}"
+                            f"\n  - Error: {exc}"
+                            f"\n  - Duration: {time.time() - submit_time:.2f}s"
+                        )
+                        doc_filter_result.raw_docs.append(
+                            FilterDoc(
+                                source_code=doc,
+                                relevance=None,
+                                task_timing=TaskTiming(),                                
+                            )
+                        )
+                    except Exception as e:
+                        logger.error(
+                            f"Document filtering error in task tracking: {exc}"
+                        )
 
         # Sort relevant_docs by relevance score in descending order
         relevant_docs.sort(
             key=lambda x: x.relevance.relevant_score, reverse=True)
-        return relevant_docs
-    
+
+        total_time = time.time() - start_time
+
+        avg_processing_time = sum(
+            doc.task_timing.real_duration for doc in relevant_docs) / len(relevant_docs) if relevant_docs else 0
+        avg_queue_time = sum(doc.task_timing.real_start_time -
+                             doc.task_timing.submit_time for doc in relevant_docs) / len(relevant_docs) if relevant_docs else 0
+        
+        total_input_tokens = sum(doc_filter_result.input_tokens_counts)
+        total_generated_tokens = sum(doc_filter_result.generated_tokens_counts)
+
+        logger.info(
+            f"=== DocFilter Complete ==="
+            f"\n  * Total time: {total_time:.2f}s"
+            f"\n  * Documents processed: {completed_tasks}/{len(documents)}"
+            f"\n  * Relevant documents: {relevant_count} (threshold: {self.relevant_score})"
+            f"\n  * Average processing time: {avg_processing_time:.2f}s"
+            f"\n  * Average queue time: {avg_queue_time:.2f}s"
+            f"\n  * Total input tokens: {total_input_tokens}"
+            f"\n  * Total generated tokens: {total_generated_tokens}"
+        )
+
+        if relevant_docs:
+            logger.info(
+                f"Top 5 relevant documents:"
+                + "".join([f"\n  * {doc.source_code.module_name} (Score: {doc.relevance.relevant_score:.1f})"
+                          for doc in relevant_docs[:5]])
+            )
+        else:
+            logger.warning("No relevant documents found!")
+
+        return doc_filter_result

autocoder/rag/llm_wrapper.py

@@ -44,13 +44,15 @@ class LLWrapper:
         res,contexts = self.rag.stream_chat_oai(conversations,llm_config=llm_config)        
         for t in res:                        
             yield (t,SingleOutputMeta(0,0))
+    
 
     async def async_stream_chat_oai(self,conversations, 
                         model:Optional[str]=None, 
                         role_mapping=None,
                         delta_mode=False,
                         llm_config:Dict[str,Any]={}): 
-        res,contexts = await asyncfy_with_semaphore(lambda: self.rag.stream_chat_oai(conversations,llm_config=llm_config))()                
+        res,contexts = await asyncfy_with_semaphore(lambda: self.rag.stream_chat_oai(conversations,llm_config=llm_config))()
+        # res,contexts = await self.llm.async_stream_chat_oai(conversations,llm_config=llm_config)
         for t in res:                                    
             yield (t,SingleOutputMeta(0,0))
                          

autocoder/rag/long_context_rag.py

@@ -31,6 +31,8 @@ from tokenizers import Tokenizer
 from autocoder.rag.variable_holder import VariableHolder
 from importlib.metadata import version
 from autocoder.rag.stream_event import event_writer
+from autocoder.rag.relevant_utils import DocFilterResult
+from pydantic import BaseModel
 
 try:        
     from autocoder_pro.rag.llm_compute import LLMComputeEngine
@@ -42,6 +44,24 @@ 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__(
         self,
@@ -305,7 +325,7 @@ class LongContextRAG:
                 url = ",".join(contexts)
                 return [SourceCode(module_name=f"RAG:{url}", source_code="".join(v))]
 
-    def _filter_docs(self, conversations: List[Dict[str, str]]) -> List[FilterDoc]:
+    def _filter_docs(self, conversations: List[Dict[str, str]]) -> DocFilterResult:
         query = conversations[-1]["content"]
         documents = self._retrieve_documents(options={"query":query})
         return self.doc_filter.filter_docs(
@@ -439,7 +459,32 @@ class LongContextRAG:
 
             logger.info(f"Query: {query} only_contexts: {only_contexts}")
             start_time = time.time()
-            relevant_docs: List[FilterDoc] = self._filter_docs(conversations)            
+
+            rag_stat = RAGStat(
+                recall_stat=RecallStat(
+                    total_input_tokens=0,
+                    total_generated_tokens=0,
+                    model_name=self.llm.default_model_name,
+                ),
+                chunk_stat=ChunkStat(
+                    total_input_tokens=0,
+                    total_generated_tokens=0,
+                    model_name=self.llm.default_model_name,
+                ),
+                answer_stat=AnswerStat(
+                    total_input_tokens=0,
+                    total_generated_tokens=0,
+                    model_name=self.llm.default_model_name,
+                ),
+            )
+
+            doc_filter_result = self._filter_docs(conversations)            
+
+            rag_stat.recall_stat.total_input_tokens += sum(doc_filter_result.input_tokens_counts)
+            rag_stat.recall_stat.total_generated_tokens += sum(doc_filter_result.generated_tokens_counts)
+            rag_stat.recall_stat.model_name = doc_filter_result.model_name
+
+            relevant_docs: List[FilterDoc] = doc_filter_result.docs
             filter_time = time.time() - start_time
 
             # Filter relevant_docs to only include those with is_relevant=True
@@ -469,17 +514,15 @@ class LongContextRAG:
             # 将 FilterDoc 转化为 SourceCode 方便后续的逻辑继续做处理
             relevant_docs = [doc.source_code for doc in relevant_docs]
 
-            console = Console()
+            logger.info(f"=== RAG Search Results ===")
+            logger.info(f"Query: {query}")
+            logger.info(f"Found relevant docs: {len(relevant_docs)}")
 
-            # Create a table for the query information
-            query_table = Table(title="Query Information", show_header=False)
-            query_table.add_row("Query", query)
-            query_table.add_row("Relevant docs", str(len(relevant_docs)))
-
-            # Add relevant docs information
+            # 记录相关文档信息
             relevant_docs_info = []
-            for doc in relevant_docs:
-                info = f"- {doc.module_name.replace(self.path,'',1)}"
+            for i, doc in enumerate(relevant_docs):
+                doc_path = doc.module_name.replace(self.path, '', 1)
+                info = f"{i+1}. {doc_path}"
                 if "original_docs" in doc.metadata:
                     original_docs = ", ".join(
                         [
@@ -490,8 +533,11 @@ class LongContextRAG:
                     info += f" (Original docs: {original_docs})"
                 relevant_docs_info.append(info)
 
-            relevant_docs_info = "\n".join(relevant_docs_info)
-            query_table.add_row("Relevant docs list", relevant_docs_info)
+            if relevant_docs_info:
+                logger.info(
+                    f"Relevant documents list:"
+                    + "".join([f"\n  * {info}" for info in relevant_docs_info])
+                )
 
             first_round_full_docs = []
             second_round_extracted_docs = []
@@ -507,11 +553,18 @@ class LongContextRAG:
                     llm=self.llm,
                     disable_segment_reorder=self.args.disable_segment_reorder,
                 )
-                final_relevant_docs = token_limiter.limit_tokens(
+                
+                token_limiter_result = token_limiter.limit_tokens(
                     relevant_docs=relevant_docs,
                     conversations=conversations,
                     index_filter_workers=self.args.index_filter_workers or 5,
                 )
+
+                rag_stat.chunk_stat.total_input_tokens += sum(token_limiter_result.input_tokens_counts)
+                rag_stat.chunk_stat.total_generated_tokens += sum(token_limiter_result.generated_tokens_counts)
+                rag_stat.chunk_stat.model_name = token_limiter_result.model_name
+                
+                final_relevant_docs = token_limiter_result.docs
                 first_round_full_docs = token_limiter.first_round_full_docs
                 second_round_extracted_docs = token_limiter.second_round_extracted_docs
                 sencond_round_time = token_limiter.sencond_round_time
@@ -522,57 +575,64 @@ class LongContextRAG:
 
             logger.info(f"Finally send to model: {len(relevant_docs)}")
 
-            query_table.add_row("Only contexts", str(only_contexts))
-            query_table.add_row("Filter time", f"{filter_time:.2f} seconds")
-            query_table.add_row("Final relevant docs", str(len(relevant_docs)))
-            query_table.add_row(
-                "first_round_full_docs", str(len(first_round_full_docs))
-            )
-            query_table.add_row(
-                "second_round_extracted_docs", str(len(second_round_extracted_docs))
-            )
-            query_table.add_row(
-                "Second round time", f"{sencond_round_time:.2f} seconds"
+            # 记录分段处理的统计信息
+            logger.info(
+                f"=== Token Management ===\n"
+                f"  * Only contexts: {only_contexts}\n"
+                f"  * Filter time: {filter_time:.2f} seconds\n" 
+                f"  * Final relevant docs: {len(relevant_docs)}\n"
+                f"  * First round full docs: {len(first_round_full_docs)}\n"
+                f"  * Second round extracted docs: {len(second_round_extracted_docs)}\n"
+                f"  * Second round time: {sencond_round_time:.2f} seconds"
             )
 
-            # Add relevant docs information
+            # 记录最终选择的文档详情
             final_relevant_docs_info = []
-            for doc in relevant_docs:
-                info = f"- {doc.module_name.replace(self.path,'',1)}"
+            for i, doc in enumerate(relevant_docs):
+                doc_path = doc.module_name.replace(self.path, '', 1)
+                info = f"{i+1}. {doc_path}"
+                
+                metadata_info = []
                 if "original_docs" in doc.metadata:
                     original_docs = ", ".join(
                         [
-                            doc.replace(self.path, "", 1)
-                            for doc in doc.metadata["original_docs"]
+                            od.replace(self.path, "", 1)
+                            for od in doc.metadata["original_docs"]
                         ]
                     )
-                    info += f" (Original docs: {original_docs})"
+                    metadata_info.append(f"Original docs: {original_docs}")
+                    
                 if "chunk_ranges" in doc.metadata:
                     chunk_ranges = json.dumps(
                         doc.metadata["chunk_ranges"], ensure_ascii=False
                     )
-                    info += f" (Chunk ranges: {chunk_ranges})"
+                    metadata_info.append(f"Chunk ranges: {chunk_ranges}")
+                    
+                if "processing_time" in doc.metadata:
+                    metadata_info.append(f"Processing time: {doc.metadata['processing_time']:.2f}s")
+                    
+                if metadata_info:
+                    info += f" ({'; '.join(metadata_info)})"
+                    
                 final_relevant_docs_info.append(info)
 
-            final_relevant_docs_info = "\n".join(final_relevant_docs_info)
-            query_table.add_row("Final Relevant docs list", final_relevant_docs_info)
-
-            # Create a panel to contain the table
-            panel = Panel(
-                query_table,
-                title="RAG Search Results",
-                expand=False,
+            if final_relevant_docs_info:
+                logger.info(
+                    f"Final documents to be sent to model:"
+                    + "".join([f"\n  * {info}" for info in final_relevant_docs_info])
             )
 
-            # Log the panel using rich
-            console.print(panel)
-
+            # 记录令牌统计
             request_tokens = sum([doc.tokens for doc in relevant_docs])
             target_model = model or self.llm.default_model_name
             logger.info(
-                f"Start to send to model {target_model} with {request_tokens} tokens"
+                f"=== LLM Request ===\n"
+                f"  * Target model: {target_model}\n"
+                f"  * Total tokens: {request_tokens}"
             )
 
+            logger.info(f"Start to send to model {target_model} with {request_tokens} tokens")
+
             if LLMComputeEngine is not None and not self.args.disable_inference_enhance:
                 llm_compute_engine = LLMComputeEngine(
                     llm=target_llm,
@@ -585,17 +645,22 @@ class LongContextRAG:
                 new_conversations = llm_compute_engine.process_conversation(
                     conversations, query, [doc.source_code for doc in relevant_docs]
                 )
-
-                return (
-                    llm_compute_engine.stream_chat_oai(
+                chunks = llm_compute_engine.stream_chat_oai(
                         conversations=new_conversations,
                         model=model,
                         role_mapping=role_mapping,
                         llm_config=llm_config,
                         delta_mode=True,
-                    ),
-                    context,
-                )
+                    )
+                
+                def generate_chunks():
+                    for chunk in chunks:
+                        yield chunk[0]
+                        if chunk[1] is not None:
+                            rag_stat.answer_stat.total_input_tokens += chunk[1].input_tokens_count
+                            rag_stat.answer_stat.total_generated_tokens += chunk[1].generated_tokens_count                      
+                    self._print_rag_stats(rag_stat)        
+                return generate_chunks(), context
 
             new_conversations = conversations[:-1] + [
                 {
@@ -614,5 +679,85 @@ class LongContextRAG:
                 llm_config=llm_config,
                 delta_mode=True,
             )
+            
+            def generate_chunks():
+                for chunk in chunks:
+                    yield chunk[0]
+                    if chunk[1] is not None:
+                        rag_stat.answer_stat.total_input_tokens += chunk[1].input_tokens_count
+                        rag_stat.answer_stat.total_generated_tokens += chunk[1].generated_tokens_count                      
+                self._print_rag_stats(rag_stat)        
+            return generate_chunks(), context
+            
+            
 
-            return (chunk[0] for chunk in chunks), context
+    def _print_rag_stats(self, rag_stat: RAGStat) -> None:
+        """打印RAG执行的详细统计信息"""
+        total_input_tokens = (
+            rag_stat.recall_stat.total_input_tokens +
+            rag_stat.chunk_stat.total_input_tokens +
+            rag_stat.answer_stat.total_input_tokens
+        )
+        total_generated_tokens = (
+            rag_stat.recall_stat.total_generated_tokens +
+            rag_stat.chunk_stat.total_generated_tokens + 
+            rag_stat.answer_stat.total_generated_tokens
+        )
+        total_tokens = total_input_tokens + total_generated_tokens
+        
+        # 避免除以零错误
+        if total_tokens == 0:
+            recall_percent = chunk_percent = answer_percent = 0
+        else:
+            recall_percent = (rag_stat.recall_stat.total_input_tokens + rag_stat.recall_stat.total_generated_tokens) / total_tokens * 100
+            chunk_percent = (rag_stat.chunk_stat.total_input_tokens + 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"
+            f"总令牌使用: {total_tokens} 令牌\n"
+            f"  * 输入令牌总数: {total_input_tokens}\n"
+            f"  * 生成令牌总数: {total_generated_tokens}\n"
+            f"\n"
+            f"阶段统计:\n"
+            f"  1. 文档检索阶段:\n"
+            f"     - 模型: {rag_stat.recall_stat.model_name}\n"
+            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"\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"\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"\n"
+            f"令牌分布百分比:\n"
+            f"  - 文档检索: {recall_percent:.1f}%\n"
+            f"  - 文档分块: {chunk_percent:.1f}%\n"
+            f"  - 答案生成: {answer_percent:.1f}%\n"
+        )
+        
+        # 记录原始统计数据，以便调试
+        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} 人民币")
+
+    def _estimate_token_cost(self, input_tokens: int, output_tokens: int) -> 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
+