main_graphrag.py

"""
GraphRAG集成示例
展示如何在自适应RAG系统中使用知识图谱功能
"""

import os
from pprint import pprint

from config import (
    setup_environment, 
    ENABLE_GRAPHRAG,
    GRAPHRAG_INDEX_PATH,
    GRAPHRAG_BATCH_SIZE
)
from document_processor import initialize_document_processor
from graph_indexer import initialize_graph_indexer
from graph_retriever import initialize_graph_retriever


class AdaptiveRAGWithGraph:
    """集成GraphRAG的自适应RAG系统"""
    
    def __init__(self, enable_graphrag=True, rebuild_index=False):
        print("🚀 初始化集成GraphRAG的自适应RAG系统...")
        print("="*60)
        
        # 设置环境
        try:
            setup_environment()
            print("✅ 环境配置完成")
        except ValueError as e:
            print(f"❌ {e}")
            raise
        
        # 初始化文档处理器
        print("\n📚 初始化文档处理器...")
        self.doc_processor, self.vectorstore, self.retriever, self.doc_splits = \
            initialize_document_processor()
        
        # GraphRAG组件
        self.enable_graphrag = enable_graphrag and ENABLE_GRAPHRAG
        self.graph_indexer = None
        self.graph_retriever = None
        self.knowledge_graph = None
        
        if self.enable_graphrag:
            self._setup_graphrag(rebuild_index)
        
        print("\n" + "="*60)
        print("✅ 系统初始化完成!")
        print("="*60)
    
    def _setup_graphrag(self, rebuild_index=False):
        """设置GraphRAG组件"""
        print("\n🔷 设置GraphRAG组件...")
        
        # 初始化索引器
        self.graph_indexer = initialize_graph_indexer()
        
        # 检查是否已有索引
        index_exists = os.path.exists(GRAPHRAG_INDEX_PATH)
        
        if index_exists and not rebuild_index:
            print(f"📂 发现现有索引: {GRAPHRAG_INDEX_PATH}")
            print("   加载现有索引...")
            self.knowledge_graph = self.graph_indexer.load_index(GRAPHRAG_INDEX_PATH)
        else:
            if rebuild_index:
                print("🔄 重新构建索引...")
            else:
                print("📝 首次构建索引...")
            
            if self.doc_splits is None:
                try:
                    docs_from_vs = self.doc_processor.get_all_documents_from_vectorstore()
                    if docs_from_vs:
                        self.doc_splits = docs_from_vs
                    else:
                        docs = self.doc_processor.load_documents()
                        self.doc_splits = self.doc_processor.split_documents(docs)
                except Exception as e:
                    print(f"   ❌ 准备GraphRAG文档块失败: {e}")
                    raise

            # 构建索引
            self.knowledge_graph = self.graph_indexer.index_documents(
                documents=self.doc_splits,
                batch_size=GRAPHRAG_BATCH_SIZE,
                save_path=GRAPHRAG_INDEX_PATH
            )
        
        # 初始化检索器
        self.graph_retriever = initialize_graph_retriever(self.knowledge_graph)
        print("✅ GraphRAG组件设置完成")
    
    def query_vector_only(self, question: str) -> str:
        """仅使用向量检索"""
        print(f"\n{'='*60}")
        print(f"🔍 向量检索模式")
        print(f"问题: {question}")
        print(f"{'='*60}")
        
        docs = self.retriever.get_relevant_documents(question)
        
        print(f"\n📄 检索到 {len(docs)} 个文档片段:")
        for i, doc in enumerate(docs[:3], 1):
            print(f"\n片段 {i}:")
            print(f"{doc.page_content[:200]}...")
        
        return self.doc_processor.format_docs(docs)
    
    def query_graph_local(self, question: str) -> str:
        """使用图谱本地查询"""
        if not self.enable_graphrag:
            return "GraphRAG未启用"
        
        print(f"\n{'='*60}")
        print(f"🔎 图谱本地查询模式")
        print(f"问题: {question}")
        print(f"{'='*60}")
        
        answer = self.graph_retriever.local_query(question)
        
        print(f"\n💡 答案:")
        print(answer)
        
        return answer
    
    def query_graph_global(self, question: str) -> str:
        """使用图谱全局查询"""
        if not self.enable_graphrag:
            return "GraphRAG未启用"
        
        print(f"\n{'='*60}")
        print(f"🌍 图谱全局查询模式")
        print(f"问题: {question}")
        print(f"{'='*60}")
        
        answer = self.graph_retriever.global_query(question)
        
        print(f"\n💡 答案:")
        print(answer)
        
        return answer
    
    def query_hybrid(self, question: str) -> dict:
        """混合查询：向量 + 图谱"""
        if not self.enable_graphrag:
            return {"error": "GraphRAG未启用"}
        
        print(f"\n{'='*60}")
        print(f"🔀 混合查询模式")
        print(f"问题: {question}")
        print(f"{'='*60}")
        
        # 向量检索
        vector_docs = self.retriever.get_relevant_documents(question)
        vector_context = self.doc_processor.format_docs(vector_docs[:3])
        
        # 图谱查询
        graph_results = self.graph_retriever.hybrid_query_with_metrics(question)
        
        result = {
            "question": question,
            "vector_retrieval": {
                "doc_count": len(vector_docs),
                "context": vector_context[:500] + "..." if len(vector_context) > 500 else vector_context
            },
            "graph_local": graph_results["local"],
            "graph_global": graph_results["global"],
            "graph_local_hallucination": graph_results.get("local_hallucination"),
            "graph_global_hallucination": graph_results.get("global_hallucination"),
            "graph_local_metrics": graph_results.get("local_metrics"),
            "graph_global_metrics": graph_results.get("global_metrics")
        }
        
        print("\n📊 结果汇总:")
        print(f"  • 向量检索: {len(vector_docs)} 个文档")
        print(f"  • 图谱本地查询完成")
        print(f"  • 图谱全局查询完成")
        
        return result
    
    def query_smart(self, question: str) -> str:
        """智能查询：自动选择最佳策略"""
        if not self.enable_graphrag:
            return self.query_vector_only(question)
        
        print(f"\n{'='*60}")
        print(f"🧠 智能查询模式")
        print(f"问题: {question}")
        print(f"{'='*60}")
        
        answer = self.graph_retriever.smart_query(question)
        
        print(f"\n💡 答案:")
        print(answer)
        
        return answer
    
    def get_graph_statistics(self):
        """获取知识图谱统计信息"""
        if not self.enable_graphrag or not self.knowledge_graph:
            print("GraphRAG未启用或图谱未构建")
            return
        
        stats = self.knowledge_graph.get_statistics()
        
        print("\n" + "="*60)
        print("📊 知识图谱统计信息")
        print("="*60)
        print(f"节点数: {stats['num_nodes']}")
        print(f"边数: {stats['num_edges']}")
        print(f"社区数: {stats['num_communities']}")
        print(f"图密度: {stats['density']:.4f}")
        print("\n实体类型分布:")
        for etype, count in stats['entity_types'].items():
            print(f"  • {etype}: {count}")
        print("="*60)
        
        return stats
    
    def interactive_mode(self):
        """交互模式"""
        print("\n" + "="*60)
        print("🤖 欢迎使用GraphRAG增强的自适应RAG系统!")
        print("="*60)
        print("\n查询模式:")
        print("  1️⃣  vector   - 仅向量检索")
        print("  2️⃣  local    - 图谱本地查询")
        print("  3️⃣  global   - 图谱全局查询")
        print("  4️⃣  hybrid   - 混合查询")
        print("  5️⃣  smart    - 智能查询（推荐）")
        print("  6️⃣  stats    - 显示图谱统计")
        print("  7️⃣  quit     - 退出")
        print("-"*60)
        
        while True:
            try:
                mode = input("\n选择模式 (1-7): ").strip()
                
                if mode in ['7', 'quit', 'exit', '退出', 'q']:
                    print("👋 感谢使用，再见!")
                    break
                
                if mode in ['6', 'stats']:
                    self.get_graph_statistics()
                    continue
                
                question = input("❓ 请输入问题: ").strip()
                
                if not question:
                    print("⚠️  请输入有效问题")
                    continue
                
                if mode in ['1', 'vector']:
                    self.query_vector_only(question)
                elif mode in ['2', 'local']:
                    self.query_graph_local(question)
                elif mode in ['3', 'global']:
                    self.query_graph_global(question)
                elif mode in ['4', 'hybrid']:
                    result = self.query_hybrid(question)
                    pprint(result)
                else:  # 默认智能模式
                    self.query_smart(question)
                
            except KeyboardInterrupt:
                print("\n👋 感谢使用，再见!")
                break
            except Exception as e:
                print(f"❌ 发生错误: {e}")
                print("请重试或输入 'quit' 退出")


def main():
    """主函数"""
    try:
        # 初始化系统（首次运行设置rebuild_index=True）
        rag_system = AdaptiveRAGWithGraph(
            enable_graphrag=True,
            rebuild_index=False  # 设为True重新构建索引
        )
        
        # 显示图谱统计
        rag_system.get_graph_statistics()
        
        # 测试查询
        print("\n" + "="*60)
        print("🧪 测试查询示例")
        print("="*60)
        
        # 示例1: 本地查询
        rag_system.query_graph_local("LLM Agent的主要组成部分是什么？")
        
        # 示例2: 全局查询  
        rag_system.query_graph_global("这些文档主要讨论了什么主题？")
        
        # 启动交互模式
        rag_system.interactive_mode()
        
    except Exception as e:
        print(f"❌ 系统初始化失败: {e}")
        import traceback
        traceback.print_exc()


if __name__ == "__main__":
    main()