knowledge_graph.py

"""
知识图谱模块
实现GraphRAG的核心功能：图谱构建、社区检测、层次化摘要
"""

import networkx as nx
from typing import List, Dict, Set, Tuple, Optional
from collections import defaultdict
import json

try:
    from community import community_louvain  # python-louvain
    LOUVAIN_AVAILABLE = True
except ImportError:
    LOUVAIN_AVAILABLE = False
    print("⚠️ python-louvain未安装，社区检测功能受限")

try:
    from langchain_core.prompts import PromptTemplate
except ImportError:
    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
from config import LOCAL_LLM


class KnowledgeGraph:
    """知识图谱类 - 使用NetworkX构建和管理图谱"""
    
    def __init__(self):
        self.graph = nx.Graph()  # 无向图
        self.entities = {}  # 实体详细信息
        self.communities = {}  # 社区划分结果
        self.community_summaries = {}  # 社区摘要
        
    def add_entity(self, name: str, entity_type: str, description: str = "", **kwargs):
        """添加实体节点"""
        self.graph.add_node(
            name,
            type=entity_type,
            description=description,
            **kwargs
        )
        self.entities[name] = {
            "name": name,
            "type": entity_type,
            "description": description,
            **kwargs
        }
    
    def add_relation(self, source: str, target: str, relation_type: str, 
                    description: str = "", weight: float = 1.0):
        """添加关系边"""
        self.graph.add_edge(
            source,
            target,
            relation_type=relation_type,
            description=description,
            weight=weight
        )
    
    def build_from_extractions(self, extraction_results: List[Dict]):
        """
        从实体提取结果构建图谱
        
        Args:
            extraction_results: 实体和关系提取结果列表
        """
        print("🔨 开始构建知识图谱...")
        
        total_entities = 0
        total_relations = 0
        
        for result in extraction_results:
            # 添加实体
            entities = result.get("entities", [])
            for entity in entities:
                self.add_entity(
                    name=entity["name"],
                    entity_type=entity.get("type", "UNKNOWN"),
                    description=entity.get("description", "")
                )
                total_entities += 1
            
            # 添加关系
            relations = result.get("relations", [])
            for relation in relations:
                source = relation.get("source")
                target = relation.get("target")
                
                # 确保节点存在
                if source in self.graph and target in self.graph:
                    self.add_relation(
                        source=source,
                        target=target,
                        relation_type=relation.get("relation_type", "RELATED_TO"),
                        description=relation.get("description", "")
                    )
                    total_relations += 1
        
        print(f"✅ 图谱构建完成: {total_entities} 个实体, {total_relations} 个关系")
        print(f"   实际节点数: {self.graph.number_of_nodes()}")
        print(f"   实际边数: {self.graph.number_of_edges()}")
    
    def detect_communities(self, algorithm: str = "louvain") -> Dict[str, int]:
        """
        社区检测 - GraphRAG的核心组件
        
        Args:
            algorithm: 社区检测算法 ('louvain', 'greedy', 'label_propagation')
            
        Returns:
            节点到社区ID的映射
        """
        print(f"🔍 开始社区检测 (算法: {algorithm})...")
        
        if self.graph.number_of_nodes() == 0:
            print("⚠️ 图谱为空，跳过社区检测")
            return {}
        
        try:
            if algorithm == "louvain" and LOUVAIN_AVAILABLE:
                communities = community_louvain.best_partition(self.graph)
            elif algorithm == "greedy":
                communities_generator = nx.community.greedy_modularity_communities(self.graph)
                communities = {}
                for idx, community_set in enumerate(communities_generator):
                    for node in community_set:
                        communities[node] = idx
            elif algorithm == "label_propagation":
                communities_generator = nx.community.label_propagation_communities(self.graph)
                communities = {}
                for idx, community_set in enumerate(communities_generator):
                    for node in community_set:
                        communities[node] = idx
            else:
                print(f"⚠️ 未知算法 {algorithm}，使用贪婪算法")
                communities_generator = nx.community.greedy_modularity_communities(self.graph)
                communities = {}
                for idx, community_set in enumerate(communities_generator):
                    for node in community_set:
                        communities[node] = idx
            
            self.communities = communities
            num_communities = len(set(communities.values()))
            print(f"✅ 检测到 {num_communities} 个社区")
            
            return communities
            
        except Exception as e:
            print(f"❌ 社区检测失败: {e}")
            return {}
    
    def get_community_members(self, community_id: int) -> List[str]:
        """获取指定社区的所有成员"""
        return [node for node, cid in self.communities.items() if cid == community_id]
    
    def get_community_subgraph(self, community_id: int) -> nx.Graph:
        """获取指定社区的子图"""
        members = self.get_community_members(community_id)
        return self.graph.subgraph(members)
    
    def get_node_neighbors(self, node: str, depth: int = 1) -> Set[str]:
        """获取节点的邻居（支持多跳）"""
        if node not in self.graph:
            return set()
        
        neighbors = {node}
        current_layer = {node}
        
        for _ in range(depth):
            next_layer = set()
            for n in current_layer:
                next_layer.update(self.graph.neighbors(n))
            neighbors.update(next_layer)
            current_layer = next_layer
        
        return neighbors
    
    def get_entity_info(self, entity_name: str) -> Optional[Dict]:
        """获取实体详细信息"""
        return self.entities.get(entity_name)
    
    def search_entities_by_type(self, entity_type: str) -> List[str]:
        """按类型搜索实体"""
        return [
            name for name, data in self.entities.items()
            if data.get("type") == entity_type
        ]
    
    def get_statistics(self) -> Dict:
        """获取图谱统计信息"""
        stats = {
            "num_nodes": self.graph.number_of_nodes(),
            "num_edges": self.graph.number_of_edges(),
            "num_communities": len(set(self.communities.values())) if self.communities else 0,
            "density": nx.density(self.graph) if self.graph.number_of_nodes() > 0 else 0,
            "entity_types": {}
        }
        
        # 统计实体类型分布
        for entity in self.entities.values():
            etype = entity.get("type", "UNKNOWN")
            stats["entity_types"][etype] = stats["entity_types"].get(etype, 0) + 1
        
        return stats
    
    def save_to_file(self, filepath: str):
        """保存图谱到文件"""
        data = {
            "entities": self.entities,
            "edges": [
                {
                    "source": u,
                    "target": v,
                    "data": data
                }
                for u, v, data in self.graph.edges(data=True)
            ],
            "communities": self.communities,
            "community_summaries": self.community_summaries
        }
        
        with open(filepath, 'w', encoding='utf-8') as f:
            json.dump(data, f, ensure_ascii=False, indent=2)
        
        print(f"✅ 图谱已保存到: {filepath}")
    
    def load_from_file(self, filepath: str):
        """从文件加载图谱"""
        with open(filepath, 'r', encoding='utf-8') as f:
            data = json.load(f)
        
        self.entities = data.get("entities", {})
        self.communities = data.get("communities", {})
        self.community_summaries = data.get("community_summaries", {})
        
        # 重建图
        self.graph.clear()
        for name, entity in self.entities.items():
            self.add_entity(**entity)
        
        for edge in data.get("edges", []):
            self.graph.add_edge(
                edge["source"],
                edge["target"],
                **edge["data"]
            )
        
        print(f"✅ 图谱已从文件加载: {filepath}")


class CommunitySummarizer:
    """社区摘要生成器 - GraphRAG的关键组件"""
    
    def __init__(self):
        self.llm = ChatOllama(model=LOCAL_LLM, temperature=0.3)
        
        self.summary_prompt = PromptTemplate(
            template="""你是一个知识图谱分析专家。请为以下社区生成一个综合摘要。

社区成员（实体）:
{entities}

实体间的关系:
{relations}

请生成一个简洁的摘要，描述：
1. 这个社区的主题是什么
2. 主要包含哪些核心概念
3. 实体之间的关键关系

摘要（2-3句话）:
""",
            input_variables=["entities", "relations"]
        )
        
        self.summary_chain = self.summary_prompt | self.llm | StrOutputParser()
    
    def summarize_community(self, kg: KnowledgeGraph, community_id: int) -> str:
        """
        为指定社区生成摘要
        
        Args:
            kg: 知识图谱对象
            community_id: 社区ID
            
        Returns:
            社区摘要文本
        """
        members = kg.get_community_members(community_id)
        subgraph = kg.get_community_subgraph(community_id)
        
        # 准备实体信息
        entity_info = []
        for member in members[:20]:  # 限制数量
            info = kg.get_entity_info(member)
            if info:
                entity_info.append(
                    f"- {info['name']} ({info.get('type', 'UNKNOWN')}): {info.get('description', '无描述')}"
                )
        
        # 准备关系信息
        relation_info = []
        for u, v, data in subgraph.edges(data=True):
            relation_info.append(
                f"- {u} --[{data.get('relation_type', 'RELATED')}]--> {v}"
            )
        
        entities_text = "\n".join(entity_info) if entity_info else "无实体"
        relations_text = "\n".join(relation_info[:15]) if relation_info else "无关系"
        
        try:
            summary = self.summary_chain.invoke({
                "entities": entities_text,
                "relations": relations_text
            })
            return summary.strip()
        except Exception as e:
            print(f"❌ 社区 {community_id} 摘要生成失败: {e}")
            return f"社区{community_id}: 包含{len(members)}个实体"
    
    def summarize_all_communities(self, kg: KnowledgeGraph) -> Dict[int, str]:
        """为所有社区生成摘要"""
        if not kg.communities:
            print("⚠️ 未检测到社区，请先运行社区检测")
            return {}
        
        community_ids = set(kg.communities.values())
        print(f"📝 开始为 {len(community_ids)} 个社区生成摘要...")
        
        summaries = {}
        for cid in community_ids:
            print(f"   处理社区 {cid}...")
            summary = self.summarize_community(kg, cid)
            summaries[cid] = summary
            kg.community_summaries[cid] = summary
        
        print("✅ 所有社区摘要生成完成")
        return summaries


def initialize_knowledge_graph():
    """初始化知识图谱"""
    return KnowledgeGraph()


def initialize_community_summarizer():
    """初始化社区摘要生成器"""
    return CommunitySummarizer()