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

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

# 实体-关系-实体
(iPhone 15, 采用, 钛合金中框)
(钛合金中框, 提升, 散热效率)
(iPhone 15, 对比, iPhone 14)
(iPhone 14, 采用, 铝合金中框)

# 实体-属性-值
(钛合金中框, 散热性能, 优秀)
(铝合金中框, 散热性能, 一般)

第1跳：从【糖尿病】节点出发
 → 沿着【常见并发症】边找到【肾功能不全】节点

第2跳：从【布洛芬】节点出发
 → 沿着【禁忌症】边找到【肾功能损伤】节点

第3跳：发现两条路径在"肾脏相关风险"概念上产生交集
 → 推理得出结论：糖尿病患者服用布洛芬需要谨慎

# 基础库
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

第1跳：从【某某科技公司】出发
 → 沿着【法人代表】边找到【张三】
 → 沿着【股东】边找到【某某投资公司】

第2跳：从【张三】出发
 → 沿着【担任法人】边找到【某某贸易公司】（有逾期记录）
 从【某某投资公司】出发
 → 沿着【控股】边找到【某某地产公司】（被列为失信被执行人）

第3跳：通过关系传播计算风险评分
 → 虽然某某科技公司本身没有不良记录
 → 但通过关联路径发现其法人和股东涉及高风险实体
 → 系统自动提升风险评级

检索路径：
1. 找到【阿司匹林】节点
2. 沿着【适应症】边找到【心肌梗死】、【心绞痛】、【脑卒中】等疾病节点
3. 从这些疾病节点出发，沿着【症状】、【预防措施】等边扩展
4. 返回完整的适应症列表和相关医学知识

检索路径：
1. 找到【相机A】节点
2. 沿着【兼容】边找到【镜头型号B】、【镜头型号C】
3. 沿着【常用配件】边找到【存储卡】、【三脚架】、【相机包】
4. 沿着【适用场景】边找到【夜景拍摄】、【人像摄影】
5. 从【夜景拍摄】出发，沿着【推荐设备】边找到【闪光灯】

检索路径：
1. 找到【请假流程】节点
2. 沿着【适用人群】边找到【全职员工】、【实习生】
3. 沿着【审批流程】边找到【部门经理审批】→【人事部备案】
4. 沿着【相关规定】边找到【年假政策】、【病假政策】
5. 从【年假政策】出发，找到具体的休假天数和申请条件

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

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