「学习笔记」大模型GraphRAG（知识图谱增强检索）

输入文档 → 知识图谱构建层 → 图存储层 → 检索推理层 → LLM生成层

输入文档
    ↓
实体提取（人物、组织、地点、概念等）
    ↓
关系提取（人物A-工作于-公司B）
    ↓
构建知识图谱（节点+边+属性）
    ↓
图嵌入学习（Node2Vec、GraphSAGE）
    ↓
用户查询 → 图查询 → 相关子图检索
    ↓
子图解释 + 文本证据 → 提示工程
    ↓
LLM生成 → 最终答案 + 解释

# 基础库
pip install neo4j py2neo networkx
pip install sentence-transformers spacy
pip install openai langchain
# 图数据库（可选）
# 下载并安装Neo4j Desktop或使用云服务

import spacy
from typing import List, Dict, Tuple
import networkx as nx
from sentence_transformers import SentenceTransformer
import openai

class SimpleGraphRAG:
    def __init__(self, llm_api_key):
        self.nlp = spacy.load("en_core_web_sm")
        self.graph = nx.Graph()
        self.embedder = SentenceTransformer('all-MiniLM-L6-v2')
        self.llm_client = openai.OpenAI(api_key=llm_api_key)
        
    def extract_entities_relations(self, text: str) -> Tuple[List[str], List[Tuple]]:
        """使用spacy进行基础的实体关系抽取"""
        doc = self.nlp(text)
        
        entities = []
        for ent in doc.ents:
            entities.append({
                'text': ent.text,
                'label': ent.label_,
                'start': ent.start_char,
                'end': ent.end_char
            })
        
        # 简化的关系抽取：基于依存句法
        relations = []
        for token in doc:
            if token.dep_ in ['nsubj', 'dobj', 'attr']:
                relations.append((
                    token.head.text,
                    token.dep_,
                    token.text
                ))
        
        return entities, relations
    
    def build_graph_from_docs(self, documents: List[str]):
        """从文档构建知识图谱"""
        for doc_id, text in enumerate(documents):
            entities, relations = self.extract_entities_relations(text)
            
            # 添加实体节点
            for entity in entities:
                node_id = f"{entity['text']}_{entity['label']}"
                if not self.graph.has_node(node_id):
                    self.graph.add_node(node_id, 
                                      text=entity['text'],
                                      type=entity['label'],
                                      embeddings=self.embedder.encode(entity['text']))
            
            # 添加关系边
            for rel in relations:
                source = f"{rel[0]}_ENTITY"
                target = f"{rel[2]}_ENTITY"
                if self.graph.has_node(source) and self.graph.has_node(target):
                    self.graph.add_edge(source, target, 
                                      relation=rel[1],
                                      doc_id=doc_id)
    
    def graph_retrieval(self, query: str, top_k: int = 3):
        """基于图的检索"""
        # 1. 查询中的实体识别
        query_doc = self.nlp(query)
        query_entities = [ent.text for ent in query_doc.ents]
        
        relevant_nodes = []
        
        # 2. 实体为中心的检索
        for entity in query_entities:
            entity_node = f"{entity}_ENTITY"
            if entity_node in self.graph:
                # 获取一度邻居
                neighbors = list(self.graph.neighbors(entity_node))
                relevant_nodes.extend(neighbors[:top_k])
                
                # 获取关系路径
                for neighbor in neighbors[:2]:
                    if self.graph.has_edge(entity_node, neighbor):
                        edge_data = self.graph.get_edge_data(entity_node, neighbor)
                        relevant_nodes.append({
                            'source': entity_node,
                            'target': neighbor,
                            'relation': edge_data['relation']
                        })
        
        # 3. 向量检索回退
        if not relevant_nodes:
            all_texts = [self.graph.nodes[n]['text'] 
                        for n in self.graph.nodes]
            query_embedding = self.embedder.encode(query)
            # 简化的相似性计算（实际应使用更高效的向量索引）
            
        return relevant_nodes
    
    def generate_answer(self, query: str, context: List):
        """使用LLM生成答案"""
        context_str = "\n".join([str(c) for c in context])
        
        prompt = f"""基于以下知识图谱信息回答问题：
        
        相关知识：
        {context_str}
        
        问题：{query}
        
        答案："""
        
        response = self.llm_client.chat.completions.create(
            model="gpt-3.5-turbo",
            messages=[
                {"role": "system", "content": "你是一个基于知识图谱的智能助手。"},
                {"role": "user", "content": prompt}
            ]
        )
        
        return response.choices[0].message.content

# 使用示例
if __name__ == "__main__":
    # 初始化
    graph_rag = SimpleGraphRAG("your-openai-api-key")
    
    # 文档示例
    documents = [
        "苹果公司由史蒂夫·乔布斯创立于1976年。",
        "苹果公司发布了iPhone和MacBook等产品。",
        "史蒂夫·乔布斯是苹果公司的前CEO。"
    ]
    
    # 构建图谱
    graph_rag.build_graph_from_docs(documents)
    
    # 查询
    query = "谁创立了苹果公司？"
    context = graph_rag.graph_retrieval(query)
    answer = graph_rag.generate_answer(query, context)
    
    print(f"问题：{query}")
    print(f"答案：{answer}")
    print(f"检索到的上下文：{context}")

第一层（知识图谱构建层）：
输入: "苹果公司由史蒂夫·乔布斯创立。乔布斯还创立了NeXT和皮克斯。"
输出: 知识图谱节点: [Person:史蒂夫·乔布斯, Company:苹果, Company:NeXT, Company:皮克斯]
      关系: [founded(乔布斯, 苹果), founded(乔布斯, NeXT), founded(乔布斯, 皮克斯)]

第二层（图存储与索引层）：
存储: 将图谱存入Neo4j，构建向量索引和属性索引
索引: 创建"founded"关系索引，公司名称全文索引

第三层（检索推理与生成层）：
查询分析: 识别实体["苹果公司", "创始人"]，关系"founded"，查询类型"多跳推理"
图检索: 从"苹果公司"找到创始人"乔布斯"，从乔布斯找到所有founded关系
推理: 排除苹果公司本身，获取其他公司
LLM生成: "史蒂夫·乔布斯除了创立苹果公司，还创立了NeXT公司和皮克斯动画工作室。"

# 多策略融合检索
class AdvancedGraphRetriever:
    def __init__(self):
        self.strategies = {
            'entity_expansion': self.entity_centric_retrieval,
            'path_finding': self.path_based_retrieval,
            'community_detection': self.community_aware_retrieval,
            'embedding_similarity': self.graph_embedding_search
        }
    
    def hybrid_retrieval(self, query, graph):
        results = []
        
        # 并行执行多种检索策略
        for strategy_name, strategy_func in self.strategies.items():
            strategy_results = strategy_func(query, graph)
            results.extend(strategy_results)
        
        # 结果融合与重排序
        fused_results = self.fuse_and_rerank(results)
        
        return fused_results

GraphRAG实施路线图:
  阶段1: 验证阶段
    - 选择核心文档子集
    - 构建最小可行图谱
    - 测试基本查询
  
  阶段2: 扩展阶段
    - 完善实体关系模型
    - 优化图构建流水线
    - 建立混合检索系统
  
  阶段3: 生产阶段
    - 实现增量更新
    - 添加缓存机制
    - 监控和评估

# 分布式GraphRAG架构草图
class DistributedGraphRAG:
    def __init__(self):
        self.subgraphs = {}  # 分片存储子图
        self.master_graph = None  # 元数据图
        
    def query_distributed(self, query):
        # 1. 元数据路由
        relevant_shards = self.route_query(query)
        
        # 2. 并行子图查询
        shard_results = parallel_query_shards(query, relevant_shards)
        
        # 3. 结果聚合
        final_result = aggregate_results(shard_results)
        
        return final_result

图数据库:
  - Neo4j (最流行)
  - NebulaGraph (高性能)
  - TigerGraph (企业级)
  - JanusGraph (开源)

图嵌入算法:
  - Node2Vec
  - GraphSAGE
  - GAT (图注意力网络)
  - TransE (知识图谱嵌入)

开发框架:
  - LangChain Graph integrations
  - LlamaIndex GraphIndex
  - 自定义pipeline

许多实际场景采用混合方案：

文档 → 并行处理 → 
    ├→ 文本切分 → 向量化 → 向量存储
    └→ 实体关系提取 → 知识图谱 → 图存储
    
查询 → 并行检索 →
    ├→ 向量相似性检索 → 文本结果
    └→ 图查询 → 关系结果
    
结果融合 → 重排序 → LLM生成 → 最终答案