add async

document_processor.py

@@ -307,6 +307,119 @@ class DocumentProcessor:
         # 返回doc_splits用于GraphRAG索引
         return vectorstore, retriever, doc_splits
     
+    async def async_expand_query(self, query: str) -> List[str]:
+        """异步扩展查询"""
+        if not self.query_expansion_model:
+            return [query]
+            
+        try:
+            # 使用LLM生成扩展查询
+            prompt = QUERY_EXPANSION_PROMPT.format(query=query)
+            expanded_queries_text = await self.query_expansion_model.ainvoke(prompt)
+            
+            # 解析扩展查询
+            expanded_queries = [query]  # 包含原始查询
+            for line in expanded_queries_text.strip().split('\n'):
+                line = line.strip()
+                if line and not line.startswith('#') and not line.startswith('//'):
+                    # 移除可能的编号前缀
+                    if line[0].isdigit() and '.' in line[:5]:
+                        line = line.split('.', 1)[1].strip()
+                    expanded_queries.append(line)
+            
+            # 限制扩展查询数量
+            return expanded_queries[:MAX_EXPANDED_QUERIES + 1]
+        except Exception as e:
+            print(f"⚠️ 异步查询扩展失败: {e}")
+            return [query]
+
+    async def async_hybrid_retrieve(self, query: str, top_k: int = 5) -> List:
+        """异步混合检索"""
+        if not ENABLE_HYBRID_SEARCH or not self.ensemble_retriever:
+            return await self.retriever.ainvoke(query)
+            
+        try:
+            results = await self.ensemble_retriever.ainvoke(query)
+            return results[:top_k]
+        except Exception as e:
+            print(f"⚠️ 异步混合检索失败: {e}")
+            print("回退到向量检索")
+            return await self.retriever.ainvoke(query)
+
+    async def async_enhanced_retrieve(self, query: str, top_k: int = 5, rerank_candidates: int = 20, 
+                         image_paths: List[str] = None, use_query_expansion: bool = None):
+        """异步增强检索"""
+        import asyncio
+        
+        # 确定是否使用查询扩展
+        if use_query_expansion is None:
+            use_query_expansion = ENABLE_QUERY_EXPANSION
+            
+        # 如果启用查询扩展，生成扩展查询
+        if use_query_expansion:
+            expanded_queries = await self.async_expand_query(query)
+            print(f"查询扩展: {len(expanded_queries)} 个查询")
+        else:
+            expanded_queries = [query]
+            
+        # 多模态检索（暂时保持同步，使用线程池）
+        if image_paths and ENABLE_MULTIMODAL:
+            loop = asyncio.get_running_loop()
+            return await loop.run_in_executor(None, self.multimodal_retrieve, query, image_paths, top_k)
+            
+        # 混合检索或向量检索
+        all_candidate_docs = []
+        
+        async def retrieve_single(q):
+            if ENABLE_HYBRID_SEARCH:
+                docs = await self.async_hybrid_retrieve(q, rerank_candidates)
+            else:
+                docs = await self.retriever.ainvoke(q)
+                if len(docs) > rerank_candidates:
+                    docs = docs[:rerank_candidates]
+            return docs
+
+        # 并发执行所有查询的检索
+        results = await asyncio.gather(*[retrieve_single(q) for q in expanded_queries])
+        
+        for docs in results:
+            all_candidate_docs.extend(docs)
+            
+        # 去重（基于文档内容）
+        unique_docs = []
+        seen_content = set()
+        for doc in all_candidate_docs:
+            content = doc.page_content
+            if content not in seen_content:
+                seen_content.add(content)
+                unique_docs.append(doc)
+                
+        print(f"检索获得 {len(unique_docs)} 个候选文档")
+        
+        # 重排（如果重排器可用）
+        # 注意：重排通常是计算密集型，建议放入线程池
+        if self.reranker and len(unique_docs) > top_k:
+            try:
+                loop = asyncio.get_running_loop()
+                # rerank 方法内部可能也比较耗时
+                reranked_results = await loop.run_in_executor(
+                    None, 
+                    self.reranker.rerank, 
+                    query, unique_docs, top_k
+                )
+                final_docs = [doc for doc, score in reranked_results]
+                scores = [score for doc, score in reranked_results]
+                
+                print(f"重排后返回 {len(final_docs)} 个文档")
+                print(f"重排分数范围: {min(scores):.4f} - {max(scores):.4f}")
+                
+                return final_docs
+            except Exception as e:
+                print(f"⚠️ 重排失败: {e}，使用原始检索结果")
+                return unique_docs[:top_k]
+        else:
+            return unique_docs[:top_k]
+    
     def expand_query(self, query: str) -> List[str]:
         """扩展查询，生成相关查询"""
         if not self.query_expansion_model:

main.py

@@ -159,9 +159,9 @@ class AdaptiveRAGSystem:
             debug=False
         )
     
-    def query(self, question: str, verbose: bool = True):
+    async def query(self, question: str, verbose: bool = True):
         """
-        处理查询
+        处理查询 (异步版本)
         
         Args:
             question (str): 用户问题
@@ -170,6 +170,7 @@ class AdaptiveRAGSystem:
         Returns:
             dict: 包含最终答案和评估指标的字典
         """
+        import asyncio
         print(f"\n🔍 处理问题: {question}")
         print("=" * 50)
         
@@ -181,24 +182,19 @@ class AdaptiveRAGSystem:
         config = {"recursion_limit": 50}  # 增加到 50，默认是 25
         
         print("\n🤖 思考过程:")
-        for output in self.app.stream(inputs, config=config):
+        async for output in self.app.astream(inputs, config=config):
             for key, value in output.items():
                 if verbose:
                     # 简单的节点执行提示，模拟流式感
                     print(f"  ↳ 执行节点: {key}...", end="\r")
-                    time.sleep(0.1) # 视觉暂停
+                    # 异步暂停
+                    await asyncio.sleep(0.1) 
                     print(f"  ✅ 完成节点: {key}      ")
                     
-                    # pprint(f"节点 '{key}':")
-                    # 可选：在每个节点打印完整状态
-                    # pprint(value, indent=2, width=80, depth=None)
                 final_generation = value.get("generation", final_generation)
                 # 保存检索评估指标
                 if "retrieval_metrics" in value:
                     retrieval_metrics = value["retrieval_metrics"]
-            if verbose:
-                # pprint("\n---\n")
-                pass
         
         print("\n" + "=" * 50)
         print("🎯 最终答案:")
@@ -207,11 +203,11 @@ class AdaptiveRAGSystem:
         # 模拟流式输出效果 (打字机效果)
         if final_generation:
             import sys
-            import time
             for char in final_generation:
                 sys.stdout.write(char)
                 sys.stdout.flush()
-                time.sleep(0.01) # 控制打字速度
+                # 异步暂停
+                await asyncio.sleep(0.01) # 控制打字速度
             print() # 换行
         else:
             print("未生成答案")
@@ -226,6 +222,7 @@ class AdaptiveRAGSystem:
     
     def interactive_mode(self):
         """交互模式，允许用户持续提问"""
+        import asyncio
         print("\n🤖 欢迎使用自适应RAG系统!")
         print("💡 输入问题开始对话，输入 'quit' 或 'exit' 退出")
         print("-" * 50)
@@ -242,7 +239,8 @@ class AdaptiveRAGSystem:
                     print("⚠️  请输入一个有效的问题")
                     continue
                 
-                result = self.query(question)
+                # 使用 asyncio.run 执行异步查询
+                result = asyncio.run(self.query(question))
                 
                 # 显示检索评估摘要
                 if result.get("retrieval_metrics"):
@@ -259,11 +257,14 @@ class AdaptiveRAGSystem:
                 break
             except Exception as e:
                 print(f"❌ 发生错误: {e}")
+                import traceback
+                traceback.print_exc()
                 print("请重试或输入 'quit' 退出")
 
 
 def main():
     """主函数"""
+    import asyncio
     try:
         # 初始化系统
         rag_system: AdaptiveRAGSystem = AdaptiveRAGSystem()
@@ -272,7 +273,9 @@ def main():
         # test_question = "AlphaCodium论文讲的是什么？"
         test_question = "LangGraph的作者目前在哪家公司工作？"
         # test_question = "解释embedding嵌入的原理，最好列举实现过程的具体步骤"
-        result = rag_system.query(test_question)
+        
+        # 使用 asyncio.run 执行异步查询
+        result = asyncio.run(rag_system.query(test_question))
         
         # 显示测试查询的检索评估摘要
         if result.get("retrieval_metrics"):
@@ -289,6 +292,8 @@ def main():
         
     except Exception as e:
         print(f"❌ 系统初始化失败: {e}")
+        import traceback
+        traceback.print_exc()
         print("请检查配置和依赖是否正确安装")
 
 

workflow_nodes.py

@@ -118,9 +118,9 @@ class WorkflowNodes:
             "retry_count": 0
         }
 
-    def retrieve(self, state):
+    async def retrieve(self, state):
         """
-        检索文档
+        检索文档 (异步版本)
         
         Args:
             state (dict): 当前图状态
@@ -138,8 +138,8 @@ class WorkflowNodes:
             # 检查是否有图像路径（多模态检索）
             image_paths = state.get("image_paths", None)
             
-            # 使用增强检索
-            documents = self.doc_processor.enhanced_retrieve(
+            # 使用异步增强检索
+            documents = await self.doc_processor.async_enhanced_retrieve(
                 question, 
                 top_k=5, 
                 rerank_candidates=20,
@@ -157,15 +157,15 @@ class WorkflowNodes:
                 
         except Exception as e:
             print(f"⚠️ 增强检索失败: {e}，回退到基本检索")
-            # 回退到基本检索
+            # 回退到基本检索 (同步回退，如果需要也可以改为异步)
             try:
                 if self.retriever is not None:
-                    documents = self.retriever.invoke(question)
+                    documents = await self.retriever.ainvoke(question)
                 elif hasattr(self.doc_processor, 'vector_retriever') and self.doc_processor.vector_retriever is not None:
-                    documents = self.doc_processor.vector_retriever.invoke(question)
+                    documents = await self.doc_processor.vector_retriever.ainvoke(question)
                     print("   使用 vector_retriever 作为备选")
                 elif hasattr(self.doc_processor, 'retriever') and self.doc_processor.retriever is not None:
-                    documents = self.doc_processor.retriever.invoke(question)
+                    documents = await self.doc_processor.retriever.ainvoke(question)
                     print("   使用 doc_processor.retriever 作为备选")
                 else:
                     print("❌ 检索器未正确初始化，返回空文档列表")