AIEC-new/AIEC-RAG/retriver/langgraph/iterative_retriever.py

"""
基于LangGraph的迭代检索器
实现智能迭代检索工作流
"""

import time
from typing import Dict, Any, Optional
from langgraph.graph import StateGraph, END
from exceptions import TaskCancelledException, WorkflowStoppedException

from retriver.langgraph.graph_state import QueryState, create_initial_state, get_state_summary
from retriver.langgraph.graph_nodes import GraphNodes
from retriver.langgraph.routing_functions import (
    should_continue_retrieval, 
    route_by_complexity, 
    route_by_debug_mode
)
from retriver.langgraph.langchain_hipporag_retriever import create_langchain_hipporag_retriever
from retriver.langgraph.langchain_components import create_oneapi_llm


class IterativeRetriever:
    """
    基于LangGraph的迭代检索器
    
    实现迭代检索流程：
    1. 初始检索 -> 2. 充分性检查 -> 3. 子查询生成(如需要) -> 4. 并行检索 -> 5. 重复2-4直到充分或达到最大迭代次数 -> 6. 生成最终答案
    """
    
    def __init__(
        self,
        keyword: str,
        top_k: int = 2,
        max_iterations: int = 2,
        max_parallel_retrievals: int = 2,
        oneapi_key: Optional[str] = None,
        oneapi_base_url: Optional[str] = None,
        model_name: Optional[str] = None,
        embed_model_name: Optional[str] = None,
        complexity_model_name: Optional[str] = None,
        sufficiency_model_name: Optional[str] = None,
        skip_llm_generation: bool = False
    ):
        """
        初始化迭代检索器
        
        Args:
            keyword: ES索引关键词
            top_k: 每次检索返回的文档数量
            max_iterations: 最大迭代次数
            max_parallel_retrievals: 最大并行检索数
            oneapi_key: OneAPI密钥
            oneapi_base_url: OneAPI基础URL
            model_name: 主LLM模型名称(用于生成答案)
            embed_model_name: 嵌入模型名称
            complexity_model_name: 复杂度判断模型(如不指定则使用model_name)
            sufficiency_model_name: 充分性检查模型(如不指定则使用model_name)
            skip_llm_generation: 是否跳过LLM生成答案（仅返回检索结果）
        """
        self.keyword = keyword
        self.top_k = top_k
        self.max_iterations = max_iterations
        self.max_parallel_retrievals = max_parallel_retrievals
        self.skip_llm_generation = skip_llm_generation
        
        # 使用默认值（如果没有指定特定模型）
        complexity_model_name = complexity_model_name or model_name
        sufficiency_model_name = sufficiency_model_name or model_name
        
        # 创建组件
        print("[CONFIG] 初始化检索器组件...")
        self.retriever = create_langchain_hipporag_retriever(
            keyword=keyword,
            top_k=top_k,
            oneapi_key=oneapi_key,
            oneapi_base_url=oneapi_base_url,
            oneapi_model_gen=model_name,
            oneapi_model_embed=embed_model_name
        )
        
        # 创建主LLM（用于生成答案）
        self.llm = create_oneapi_llm(
            oneapi_key=oneapi_key,
            oneapi_base_url=oneapi_base_url,
            model_name=model_name
        )
        
        # 创建复杂度判断LLM（如果模型不同）
        if complexity_model_name != model_name:
            print(f"   [INFO] 使用独立的复杂度判断模型: {complexity_model_name}")
            self.complexity_llm = create_oneapi_llm(
                oneapi_key=oneapi_key,
                oneapi_base_url=oneapi_base_url,
                model_name=complexity_model_name
            )
        else:
            self.complexity_llm = self.llm
            
        # 创建充分性检查LLM（如果模型不同）
        if sufficiency_model_name != model_name:
            print(f"   [INFO] 使用独立的充分性检查模型: {sufficiency_model_name}")
            self.sufficiency_llm = create_oneapi_llm(
                oneapi_key=oneapi_key,
                oneapi_base_url=oneapi_base_url,
                model_name=sufficiency_model_name
            )
        else:
            self.sufficiency_llm = self.llm
        
        # 创建节点处理器
        self.nodes = GraphNodes(
            retriever=self.retriever,
            llm=self.llm,
            complexity_llm=self.complexity_llm,
            sufficiency_llm=self.sufficiency_llm,
            keyword=keyword,
            max_parallel_retrievals=max_parallel_retrievals,
            skip_llm_generation=skip_llm_generation
        )
        
        # 构建工作流图
        self.workflow = self._build_workflow()
        
        print("[OK] 迭代检索器初始化完成")
    
    def _build_workflow(self) -> StateGraph:
        """构建LangGraph工作流"""
        print("[?] 构建工作流图...")
        
        # 创建状态图
        workflow = StateGraph(QueryState)
        
        # 添加节点
        # 新增：查询复杂度判断节点
        workflow.add_node("query_complexity_check", self.nodes.query_complexity_check_node)
        
        # 新增：调试模式节点
        workflow.add_node("debug_mode_node", self.nodes.debug_mode_node)
        
        # 简单查询路径
        workflow.add_node("simple_vector_retrieval", self.nodes.simple_vector_retrieval_node)
        workflow.add_node("simple_answer_generation", self.nodes.simple_answer_generation_node)
        
        # 复杂查询路径（现有hipporag2逻辑）
        workflow.add_node("query_decomposition", self.nodes.query_decomposition_node)
        workflow.add_node("initial_retrieval", self.nodes.initial_retrieval_node)
        workflow.add_node("event_triples_extraction", self.nodes.event_triples_extraction_node)  # 新增事件三元组提取节点
        workflow.add_node("sufficiency_check", self.nodes.sufficiency_check_node)
        workflow.add_node("sub_query_generation", self.nodes.sub_query_generation_node)
        workflow.add_node("parallel_retrieval", self.nodes.parallel_retrieval_node)
        workflow.add_node("next_iteration", self.nodes.next_iteration_node)
        workflow.add_node("final_answer", self.nodes.final_answer_generation_node)
        
        # 设置入口点：从查询复杂度判断开始
        workflow.set_entry_point("query_complexity_check")
        
        # 复杂度检查后进入调试模式节点
        workflow.add_edge("query_complexity_check", "debug_mode_node")
        
        # 条件边：根据调试模式和复杂度判断结果决定路径
        workflow.add_conditional_edges(
            "debug_mode_node",
            route_by_debug_mode,
            {
                "simple_vector_retrieval": "simple_vector_retrieval",
                "initial_retrieval": "query_decomposition"  # 复杂路径先进入查询分解节点
            }
        )
        
        # 简单查询路径的边
        workflow.add_edge("simple_vector_retrieval", "simple_answer_generation")
        workflow.add_edge("simple_answer_generation", END)
        
        # 复杂查询路径的边（包含查询分解逻辑）
        workflow.add_edge("query_decomposition", "initial_retrieval")  # 查询分解后进入并行初始检索
        workflow.add_edge("initial_retrieval", "event_triples_extraction")  # 初始检索后提取事件三元组
        workflow.add_edge("event_triples_extraction", "sufficiency_check")  # 三元组提取后进入充分性检查

        # 条件边：根据充分性检查结果决定下一步
        workflow.add_conditional_edges(
            "sufficiency_check",
            should_continue_retrieval,
            {
                "final_answer": "final_answer",
                "parallel_retrieval": "sub_query_generation"
            }
        )

        workflow.add_edge("sub_query_generation", "parallel_retrieval")
        workflow.add_edge("parallel_retrieval", "next_iteration")  # 并行检索后直接进入下一迭代
        workflow.add_edge("next_iteration", "event_triples_extraction")  # 迭代后再提取事件三元组
        # 注意：这将形成循环：event_triples_extraction → sufficiency_check → (可能的) parallel_retrieval → next_iteration → event_triples_extraction
        
        # 结束节点
        workflow.add_edge("final_answer", END)
        
        return workflow.compile()
    
    def retrieve(self, query: str, mode: str = "0", task_id: Optional[str] = None) -> Dict[str, Any]:
        """
        执行迭代检索
        
        Args:
            query: 用户查询
            mode: 调试模式，"0"=自动判断，"simple"=强制简单路径，"complex"=强制复杂路径
            task_id: 任务ID（可选）

        Returns:
            包含最终答案和详细信息的字典
        """
        print(f"[STARTING] 开始迭代检索: {query}")
        start_time = time.time()
        
        # 创建初始状态
        initial_state = create_initial_state(
            original_query=query,
            max_iterations=self.max_iterations,
            debug_mode=mode,
            task_id=task_id  # 传递任务ID
        )
        initial_state["debug_info"]["start_time"] = start_time
        
        try:
            # 执行工作流
            final_state = self.workflow.invoke(initial_state)
            
            # 记录结束时间
            end_time = time.time()
            final_state["debug_info"]["end_time"] = end_time
            final_state["debug_info"]["total_time"] = end_time - start_time
            
            print(f"[SUCCESS] 迭代检索完成，耗时 {end_time - start_time:.2f}秒")
            
            # 返回结果
            return {
                "query": query,
                "answer": final_state["final_answer"],
                "query_complexity": final_state["query_complexity"],
                "is_complex_query": final_state["is_complex_query"],
                "iterations": final_state["current_iteration"],
                "total_passages": len(final_state["all_passages"]),
                "sub_queries": final_state["sub_queries"],
                "decomposed_sub_queries": final_state.get("decomposed_sub_queries", []),
                "initial_retrieval_details": final_state.get("initial_retrieval_details", {}),
                "sufficiency_check": final_state["sufficiency_check"],
                "all_passages": final_state["all_passages"],
                "debug_info": final_state["debug_info"],
                "state_summary": get_state_summary(final_state),
                "iteration_history": final_state["iteration_history"]
            }
            
        except (TaskCancelledException, WorkflowStoppedException) as e:
            print(f"[CANCELLED] 迭代检索被取消: {e}")
            return {
                "query": query,
                "answer": "任务已被取消",
                "cancelled": True,
                "iterations": 0,
                "total_passages": 0,
                "sub_queries": [],
                "debug_info": {"cancelled": True, "message": str(e), "total_time": time.time() - start_time}
            }

        except Exception as e:
            print(f"[ERROR] 迭代检索失败: {e}")
            return {
                "query": query,
                "answer": f"抱歉，检索过程中遇到错误: {str(e)}",
                "error": str(e),
                "iterations": 0,
                "total_passages": 0,
                "sub_queries": [],
                "debug_info": {"error": str(e), "total_time": time.time() - start_time}
            }
    
    def retrieve_simple(self, query: str) -> str:
        """
        简单检索接口，只返回答案
        
        Args:
            query: 用户查询
            
        Returns:
            最终答案字符串
        """
        result = self.retrieve(query)
        return result["answer"]
    
    def get_retrieval_stats(self) -> Dict[str, Any]:
        """
        获取检索器统计信息
        
        Returns:
            统计信息字典
        """
        return {
            "keyword": self.keyword,
            "top_k": self.top_k,
            "max_iterations": self.max_iterations,
            "max_parallel_retrievals": self.max_parallel_retrievals,
            "retriever_type": "IterativeRetriever with LangGraph",
            "model_info": {
                "llm_model": getattr(self.llm.oneapi_generator, 'model_name', 'unknown'),
                "embed_model": getattr(self.retriever.embedding_model, 'model_name', 'unknown')
            }
        }


def create_iterative_retriever(
    keyword: str,
    top_k: int = 2,
    max_iterations: int = 2,
    max_parallel_retrievals: int = 2,
    **kwargs
) -> IterativeRetriever:
    """
    创建迭代检索器的便捷函数
    
    Args:
        keyword: ES索引关键词
        top_k: 每次检索返回的文档数量
        max_iterations: 最大迭代次数
        max_parallel_retrievals: 最大并行检索数
        **kwargs: 其他参数
        
    Returns:
        迭代检索器实例
    """
    return IterativeRetriever(
        keyword=keyword,
        top_k=top_k,
        max_iterations=max_iterations,
        max_parallel_retrievals=max_parallel_retrievals,
        **kwargs
    )