""" 流式检索包装器 - 为现有检索器添加状态回调支持 保持向后兼容，不影响原有业务逻辑 """ import time from typing import Dict, Any, Callable, Optional import traceback def stream_retrieve( retriever, query: str, mode: str = "0", status_callback: Optional[Callable] = None ) -> Dict[str, Any]: """ 带状态回调的检索包装器 在不修改原有检索器的情况下，通过包装添加状态回调功能 Args: retriever: 原始检索器实例 query: 查询字符串 mode: 调试模式 status_callback: 状态回调函数，接收(status_type, data)参数 Returns: 检索结果字典 """ # 如果没有提供回调，直接调用原始方法 if not status_callback: return retriever.retrieve(query, mode) try: # 开始检索 status_callback("starting", {"message": "开始分析查询"}) # 创建增强的初始状态 from retriver.langgraph.graph_state import create_initial_state initial_state = create_initial_state( original_query=query, max_iterations=retriever.max_iterations, debug_mode=mode ) # 配置工作流执行 config = { "recursion_limit": 50, "metadata": { "stream_mode": True, "callback_enabled": True }, "run_name": f"Stream-Retrieval-{query[:20]}..." } # 执行工作流，并在关键点插入回调 result_state = execute_workflow_with_callbacks( retriever.workflow, initial_state, config, status_callback ) # 构建最终结果 total_time = result_state.get("debug_info", {}).get("total_time", 0) # 构建结果（使用原有逻辑） result = build_result(retriever, result_state, total_time) # 保存完整结果到JSON文件（如果retriever有这个方法） if hasattr(retriever, '_save_full_result_to_file'): # 构建完整结果以便保存 full_result = { **result, "all_documents": result_state.get('all_documents', []), "all_passages": result_state.get('all_passages', []), "passage_sources": result_state.get('passage_sources', []), "retrieval_path": result_state.get('retrieval_path', 'unknown'), "initial_retrieval_details": result_state.get('initial_retrieval_details', {}), "current_sub_queries": result_state.get('current_sub_queries', []) } retriever._save_full_result_to_file(full_result) print("[INFO] 流式检索结果已保存到JSON文件") return result except Exception as e: print(f"[ERROR] 流式检索出错: {e}") traceback.print_exc() # 出错时返回基础结果 return { "query": query, "answer": f"检索过程中遇到错误: {str(e)}", "error": str(e), "total_passages": 0, "iterations": 0, "sub_queries": [], "supporting_facts": [], "supporting_events": [] } def execute_workflow_with_callbacks( workflow, initial_state: Dict[str, Any], config: Dict[str, Any], status_callback: Callable ) -> Dict[str, Any]: """ 执行工作流并在关键节点插入状态回调 使用LangGraph的stream功能实时获取状态更新 """ # 将回调函数注入到state中，供答案生成节点使用 initial_state['stream_callback'] = status_callback print(f"[DEBUG] 已将stream_callback注入到initial_state") # 创建一个状态监控器 class StateMonitor: def __init__(self, callback): self.callback = callback self.last_iteration = 0 self.last_doc_count = 0 self.last_sub_queries = [] self.seen_nodes = set() self.last_sufficiency_check = None # 记录上次的充分性检查结果 def check_state(self, state, node_name=None): """检查状态变化并触发回调""" # 记录节点访问 if node_name and node_name not in self.seen_nodes: self.seen_nodes.add(node_name) # 检查子查询 - 节点完成后发送完整数据 if "decomposed_sub_queries" in state and state["decomposed_sub_queries"]: sub_queries = state["decomposed_sub_queries"] if sub_queries != self.last_sub_queries: self.callback("sub_queries", sub_queries) self.last_sub_queries = sub_queries # 检查文档数量 - 节点完成后发送完整数据 if "all_passages" in state: doc_count = len(state.get("all_passages", [])) if doc_count != self.last_doc_count: # 提取来源文档名，优先使用evidence字段 sources = extract_sources_from_documents(state.get("all_documents", []), state.get("passage_sources", [])) # 判断检索类型和原因 current_iteration = state.get("current_iteration", 0) sub_queries = state.get("current_sub_queries", []) # 确定检索原因 if self.last_doc_count == 0: # 初始检索 retrieval_type = "初始检索" retrieval_reason = "原始查询" elif sub_queries and len(sub_queries) > 0: # 基于子查询的检索 retrieval_type = f"子查询检索" retrieval_reason = f"基于{len(sub_queries)}个新子查询" else: # 迭代检索 retrieval_type = f"第{current_iteration}轮迭代" retrieval_reason = "迭代深化" self.callback("documents", { "count": doc_count, "sources": sources, "new_docs": doc_count - self.last_doc_count, "retrieval_type": retrieval_type, "retrieval_reason": retrieval_reason, "is_incremental": self.last_doc_count > 0 # 是否为增量检索 }) self.last_doc_count = doc_count # 检查迭代轮次 if "current_iteration" in state: current = state["current_iteration"] max_iter = state.get("max_iterations", 3) if current != self.last_iteration: self.callback("iteration", { "current": current, "max": max_iter }) self.last_iteration = current # 检查充分性判断 - 在每次文档检索后都应该显示 if "sufficiency_check" in state and state["sufficiency_check"]: check = state["sufficiency_check"] # 创建一个可比较的字符串表示 check_str = f"{check.get('is_sufficient')}_{check.get('confidence', 0):.2f}_{check.get('iteration', 0)}" # 如果是新的充分性检查结果，或者文档数量有变化（表示刚完成检索） doc_count = len(state.get("all_passages", [])) if (check and check_str != self.last_sufficiency_check) or (doc_count > 0 and doc_count != self.last_doc_count): self.callback("sufficiency_check", { "is_sufficient": check.get("is_sufficient", False), "confidence": check.get("confidence", 0), "reason": check.get("reason", "") }) self.last_sufficiency_check = check_str # 如果充分，或者达到最大迭代次数，即将生成答案 current_iteration = state.get("current_iteration", 0) max_iterations = state.get("max_iterations", 2) if check.get("is_sufficient", False) or current_iteration >= max_iterations - 1: self.callback("generating", {"message": "正在生成最终答案..."}) return state monitor = StateMonitor(status_callback) final_state = None # 尝试使用stream方法获取实时更新 if hasattr(workflow, 'stream'): try: # 使用stream方法逐步获取状态 for chunk in workflow.stream(initial_state, config=config): # chunk通常是 {node_name: state} 的字典 if isinstance(chunk, dict): for node_name, node_state in chunk.items(): # 实时检查每个节点的状态 monitor.check_state(node_state, node_name) final_state = node_state # 根据节点名称在节点开始时发送状态（只发送一次） # 注意：这些状态只在节点开始时发送，不要和节点完成后的状态重复 if "query_complexity_check" in node_name: # complexity_check在节点完成后会发送，这里不发送 pass elif "query_decomposition" in node_name: # 分解查询节点 - 不发送中间状态，结果会很快出现 pass elif "initial_retrieval" in node_name: # 初始检索节点 - 不发送中间状态，结果会很快出现 pass elif "sufficiency_check" in node_name: # 充分性检查节点开始 - 不发送，避免重复 pass elif "sub_query_generation" in node_name: # 子查询生成节点开始 - 不显示，因为用户不需要看到这个细节 pass # status_callback("generating_subqueries", {"message": "正在生成子查询..."}) elif "parallel_retrieval" in node_name: # 并行检索节点开始 - 不显示，因为用户不需要看到这个细节 pass # status_callback("parallel_retrieving", {"message": "正在并行检索..."}) elif "final" in node_name.lower() or "answer" in node_name.lower(): # 最终答案节点 - 状态已在充分性检查后发送，这里不重复 pass except Exception as e: # 如果stream失败，回退到invoke final_state = workflow.invoke(initial_state, config=config) monitor.check_state(final_state) else: # 不支持stream，使用普通invoke final_state = workflow.invoke(initial_state, config=config) monitor.check_state(final_state) return final_state def extract_sources_from_documents(all_documents, passage_sources): """ 从documents中提取真实的文档名（优先使用evidence字段） Args: all_documents: 文档列表（包含metadata中的evidence字段） passage_sources: 段落来源列表（作为后备） Returns: 去重的文档名列表 """ sources = set() # 优先从all_documents中提取evidence字段 if all_documents: for doc in all_documents: if hasattr(doc, 'metadata'): # 优先使用evidence字段（段落节点的真实文档名） evidence = doc.metadata.get('evidence', '') if evidence: sources.add(evidence) continue # 其次使用source_text_id（事件节点的来源） source_text_id = doc.metadata.get('source_text_id', '') if source_text_id: # 如果source_text_id是列表格式的字符串，解析它 if source_text_id.startswith('[') and source_text_id.endswith(']'): try: import ast text_ids = ast.literal_eval(source_text_id) if isinstance(text_ids, list) and text_ids: sources.add(text_ids[0]) # 使用第一个来源 except: sources.add(source_text_id) else: sources.add(source_text_id) continue # 如果没有找到evidence或source_text_id，使用原始的passage_sources if not sources and passage_sources: for source in passage_sources: if isinstance(source, str): # 提取文件名 if "/" in source: filename = source.split("/")[-1] else: filename = source sources.add(filename) elif isinstance(source, dict): # 如果是字典，尝试获取文件名字段 filename = source.get("filename") or source.get("source") or str(source) sources.add(filename) return list(sources)[:5] # 最多返回5个 def extract_sources(passage_sources): """ 从passage_sources提取文档名列表（保留以兼容旧代码） Args: passage_sources: 段落来源列表 Returns: 去重的文档名列表 """ return extract_sources_from_documents([], passage_sources) def build_result(retriever, final_state: Dict[str, Any], total_time: float) -> Dict[str, Any]: """ 构建最终结果（复用原有逻辑） """ # 获取Token使用信息（如果有的话） token_info = {} if hasattr(retriever, '_get_token_usage_info'): token_info = retriever._get_token_usage_info() result = { "query": final_state.get('original_query', ''), "answer": final_state.get('final_answer', '') or "未能生成答案", # 查询复杂度信息 "query_complexity": final_state.get('query_complexity', {}), "is_complex_query": final_state.get('is_complex_query', False), # 检索统计 "iterations": final_state.get('current_iteration', 0), "total_passages": len(final_state.get('all_passages', [])), "sub_queries": final_state.get('sub_queries', []), "decomposed_sub_queries": final_state.get('decomposed_sub_queries', []), # 充分性信息 "sufficiency_check": final_state.get('sufficiency_check', {}), "is_sufficient": final_state.get('is_sufficient', False), # 支撑信息 "supporting_facts": extract_supporting_facts(final_state), "supporting_events": extract_supporting_events(final_state), # 调试信息 "debug_info": { "total_time": total_time, "retrieval_calls": final_state.get('debug_info', {}).get('retrieval_calls', 0), "llm_calls": final_state.get('debug_info', {}).get('llm_calls', 0), "token_usage": token_info } } return result def extract_supporting_facts(state): """提取支撑事实""" passages = state.get('all_passages', []) facts = [] for i, passage in enumerate(passages[:10]): # 最多10个 if isinstance(passage, str): facts.append([f"Fact_{i+1}", passage[:200]]) elif isinstance(passage, dict): content = passage.get('content', '') or passage.get('text', '') facts.append([f"Fact_{i+1}", content[:200]]) return facts def extract_supporting_events(state): """提取支撑事件""" # 从状态中提取事件信息 events = [] # 可以从iteration_history中提取关键事件 iteration_history = state.get('iteration_history', []) for item in iteration_history: if item.get('action') == 'sub_query_generation': events.append([ "子查询生成", f"第{item.get('iteration', 0)}轮: 生成{len(item.get('sub_queries', []))}个子查询" ]) elif item.get('action') == 'parallel_retrieval': events.append([ "并行检索", f"第{item.get('iteration', 0)}轮: 检索{item.get('documents_retrieved', 0)}篇文档" ]) return events[:5] # 最多5个事件