graph_retriever.py

"""
GraphRAG检索器
实现基于知识图谱的检索策略，包括本地查询和全局查询
"""

from typing import List, Dict, Set, Tuple
try:
    from langchain_core.prompts import PromptTemplate
except ImportError:
    from langchain.prompts import PromptTemplate

from langchain_community.chat_models import ChatOllama
from langchain_core.output_parsers import StrOutputParser, JsonOutputParser

from knowledge_graph import KnowledgeGraph
from config import LOCAL_LLM


class GraphRetriever:
    """基于知识图谱的检索器"""
    
    def __init__(self, knowledge_graph: KnowledgeGraph):
        self.kg = knowledge_graph
        self.llm = ChatOllama(model=LOCAL_LLM, temperature=0.3)
        
        # 实体识别提示
        self.entity_recognition_prompt = PromptTemplate(
            template="""从以下问题中识别关键实体和概念:

问题: {question}

已知实体示例: {sample_entities}

请识别问题中提到的实体，返回JSON格式:
{{
    "entities": ["实体1", "实体2", ...]
}}

只返回JSON，不要其他内容。
""",
            input_variables=["question", "sample_entities"]
        )
        
        # 全局查询生成提示
        self.global_query_prompt = PromptTemplate(
            template="""你是一个知识图谱分析专家。基于以下社区摘要，回答用户问题。

用户问题: {question}

相关社区摘要:
{community_summaries}

请基于这些摘要提供一个综合性的答案。如果摘要中没有相关信息，请说明。

答案:
""",
            input_variables=["question", "community_summaries"]
        )
        
        # 本地查询生成提示
        self.local_query_prompt = PromptTemplate(
            template="""基于以下实体及其关系信息，回答用户问题。

用户问题: {question}

相关实体信息:
{entity_info}

实体间的关系:
{relations}

请基于这些信息提供答案。

答案:
""",
            input_variables=["question", "entity_info", "relations"]
        )
        
        self.entity_recognition_chain = self.entity_recognition_prompt | self.llm | JsonOutputParser()
        self.global_query_chain = self.global_query_prompt | self.llm | StrOutputParser()
        self.local_query_chain = self.local_query_prompt | self.llm | StrOutputParser()
    
    def recognize_entities(self, question: str) -> List[str]:
        """
        从问题中识别实体
        
        Args:
            question: 用户问题
            
        Returns:
            识别到的实体列表
        """
        # 获取一些示例实体
        sample_entities = list(self.kg.entities.keys())[:10]
        sample_text = ", ".join(sample_entities)
        
        try:
            result = self.entity_recognition_chain.invoke({
                "question": question,
                "sample_entities": sample_text
            })
            entities = result.get("entities", [])
            
            # 匹配到图谱中的实体
            matched_entities = []
            for entity in entities:
                # 精确匹配
                if entity in self.kg.entities:
                    matched_entities.append(entity)
                else:
                    # 模糊匹配
                    for kg_entity in self.kg.entities.keys():
                        if entity.lower() in kg_entity.lower() or kg_entity.lower() in entity.lower():
                            matched_entities.append(kg_entity)
                            break
            
            print(f"🔍 识别到实体: {matched_entities}")
            return matched_entities
            
        except Exception as e:
            print(f"❌ 实体识别失败: {e}")
            return []
    
    def local_query(self, question: str, max_hops: int = 2, top_k: int = 10) -> str:
        """
        本地查询 - 基于问题中的实体及其邻域进行检索
        
        适用场景: 针对特定实体的详细问题
        例如: "AlphaCodium的作者是谁？"
        
        Args:
            question: 用户问题
            max_hops: 最大跳数
            top_k: 返回的最大实体数
            
        Returns:
            答案文本
        """
        print(f"\n🔎 执行本地查询...")
        
        # 1. 识别问题中的实体
        mentioned_entities = self.recognize_entities(question)
        
        if not mentioned_entities:
            return "未能在知识图谱中找到相关实体。"
        
        # 2. 获取实体的邻域
        relevant_entities = set()
        for entity in mentioned_entities:
            neighbors = self.kg.get_node_neighbors(entity, depth=max_hops)
            relevant_entities.update(neighbors)
        
        relevant_entities = list(relevant_entities)[:top_k]
        
        # 3. 收集实体信息
        entity_info_list = []
        for entity in relevant_entities:
            info = self.kg.get_entity_info(entity)
            if info:
                entity_info_list.append(
                    f"- {info['name']} ({info.get('type', 'UNKNOWN')}): {info.get('description', '无描述')}"
                )
        
        # 4. 收集关系信息
        relation_list = []
        for u, v, data in self.kg.graph.edges(data=True):
            if u in relevant_entities and v in relevant_entities:
                relation_list.append(
                    f"- {u} --[{data.get('relation_type', 'RELATED')}]--> {v}: {data.get('description', '')}"
                )
        
        entity_info_text = "\n".join(entity_info_list) if entity_info_list else "无相关实体信息"
        relations_text = "\n".join(relation_list[:20]) if relation_list else "无相关关系"
        
        # 5. 生成答案
        try:
            answer = self.local_query_chain.invoke({
                "question": question,
                "entity_info": entity_info_text,
                "relations": relations_text
            })
            print(f"✅ 本地查询完成")
            return answer.strip()
        except Exception as e:
            print(f"❌ 本地查询失败: {e}")
            return "查询失败，请重试。"
    
    def global_query(self, question: str, top_k_communities: int = 5) -> str:
        """
        全局查询 - 基于社区摘要进行高层次查询
        
        适用场景: 需要整体理解的概括性问题
        例如: "这些文档主要讨论什么主题？"
        
        Args:
            question: 用户问题
            top_k_communities: 使用的社区数量
            
        Returns:
            答案文本
        """
        print(f"\n🌍 执行全局查询...")
        
        if not self.kg.community_summaries:
            return "知识图谱尚未生成社区摘要，请先运行索引流程。"
        
        # 获取社区摘要
        community_summaries = []
        for cid, summary in list(self.kg.community_summaries.items())[:top_k_communities]:
            community_summaries.append(f"社区 {cid}:\n{summary}\n")
        
        summaries_text = "\n".join(community_summaries)
        
        # 生成答案
        try:
            answer = self.global_query_chain.invoke({
                "question": question,
                "community_summaries": summaries_text
            })
            print(f"✅ 全局查询完成")
            return answer.strip()
        except Exception as e:
            print(f"❌ 全局查询失败: {e}")
            return "查询失败，请重试。"
    
    def hybrid_query(self, question: str) -> Dict[str, str]:
        """
        混合查询 - 同时执行本地和全局查询，返回两种结果
        
        Args:
            question: 用户问题
            
        Returns:
            包含本地和全局查询结果的字典
        """
        print(f"\n🔀 执行混合查询...")
        
        local_answer = self.local_query(question)
        global_answer = self.global_query(question)
        
        return {
            "local": local_answer,
            "global": global_answer,
            "question": question
        }
    
    def smart_query(self, question: str) -> str:
        """
        智能查询 - 根据问题类型自动选择查询策略
        
        Args:
            question: 用户问题
            
        Returns:
            答案文本
        """
        # 判断问题类型
        question_lower = question.lower()
        
        # 包含具体实体名称的问题 -> 本地查询
        mentioned_entities = self.recognize_entities(question)
        if mentioned_entities:
            print("📍 检测到具体实体，使用本地查询")
            return self.local_query(question)
        
        # 概括性问题 -> 全局查询
        global_keywords = ["主要", "总体", "概述", "整体", "主题", "讨论", "内容", "what", "overview", "main", "topics"]
        if any(keyword in question_lower for keyword in global_keywords):
            print("🌐 检测到概括性问题，使用全局查询")
            return self.global_query(question)
        
        # 默认使用本地查询
        print("📍 使用本地查询作为默认策略")
        return self.local_query(question)


def initialize_graph_retriever(knowledge_graph: KnowledgeGraph):
    """初始化GraphRAG检索器"""
    return GraphRetriever(knowledge_graph)