25 - Graph RAG 与知识图谱增强检索¶

# 知识图谱由三元组 (Subject, Predicate, Object) 组成
triples = [
    ("DeepSeek", "开发了", "DeepSeek-V3"),
    ("DeepSeek", "总部位于", "杭州"),
    ("DeepSeek-V3", "参数量", "671B"),
    ("DeepSeek-V3", "架构", "MoE"),
    ("MoE", "全称", "Mixture of Experts"),
    ("DeepSeek-V3", "激活参数", "37B"),
    ("DeepSeek-V3", "训练数据", "14.8T tokens"),
]

# 知识图谱可以表示为有向图
# 节点 = 实体 (Entities)
# 边 = 关系 (Relations)

import networkx as nx

def build_knowledge_graph(triples):
    """从三元组构建知识图谱"""
    G = nx.DiGraph()

    for subj, pred, obj in triples:
        G.add_node(subj, type='entity')
        G.add_node(obj, type='entity')
        G.add_edge(subj, obj, relation=pred)

    return G

# 构建图谱
kg = build_knowledge_graph(triples)

# 查询：DeepSeek 的所有关系
print("DeepSeek 的所有关系：")
for neighbor in kg.successors("DeepSeek"):
    relation = kg.edges["DeepSeek", neighbor]["relation"]
    print(f"  DeepSeek --[{relation}]--> {neighbor}")

# 最短路径查询
path = nx.shortest_path(kg, "DeepSeek", "MoE")
print(f"\nDeepSeek 到 MoE 的路径: {' → '.join(path)}")

# Neo4j：生产级图数据库，适合百万级以上节点的大规模知识图谱
# 与 NetworkX（内存图库，适合小规模实验）不同，Neo4j 使用 Cypher 查询语言，
# 原生支持索引、持久化、事务和分布式，适合生产环境部署
# 安装：pip install neo4j

from neo4j import GraphDatabase

class KnowledgeGraphDB:
    def __init__(self, uri="bolt://localhost:7687", user="neo4j", password="password"):
        self.driver = GraphDatabase.driver(uri, auth=(user, password))

    def close(self):
        self.driver.close()

    def add_triple(self, subject, predicate, object_):
        """添加三元组"""
        with self.driver.session() as session:
            session.run(
                """
                MERGE (s:Entity {name: $subject})
                MERGE (o:Entity {name: $object})
                MERGE (s)-[r:RELATION {type: $predicate}]->(o)
                """,
                subject=subject, predicate=predicate, object=object_
            )

    def query_neighbors(self, entity_name, max_depth=2):
        """查询实体的邻居"""
        with self.driver.session() as session:
            result = session.run(
                """
                MATCH (e:Entity {name: $name})-[r*1..%d]-(neighbor)
                RETURN neighbor.name AS name, 
                       type(r[-1]) AS relation,
                       length(r) AS depth
                """ % max_depth,
                name=entity_name
            )
            return [record.data() for record in result]

    def multi_hop_query(self, start_entity, end_entity, max_hops=3):
        """多跳查询"""
        with self.driver.session() as session:
            result = session.run(
                """
                MATCH path = (s:Entity {name: $start})-[r*1..%d]-(e:Entity {name: $end})
                RETURN [node in nodes(path) | node.name] AS entities,
                       [rel in relationships(path) | rel.type] AS relations
                """ % max_hops,
                start=start_entity, end=end_entity
            )
            return [record.data() for record in result]

# 使用
db = KnowledgeGraphDB()
for s, p, o in triples:
    db.add_triple(s, p, o)

# 查询 DeepSeek 的邻居
neighbors = db.query_neighbors("DeepSeek")

class GraphRAG:
    def __init__(self, llm, graph_db, vector_store):
        self.llm = llm
        self.graph_db = graph_db
        self.vector_store = vector_store

    def local_search(self, query: str, top_k: int = 5) -> str:
        """局部搜索：精确实体匹配 + 子图扩展

        适合：具体事实型问题
        例如："DeepSeek-V3 的架构是什么？"
        """
        # 1. 从查询中抽取实体
        entities = self._extract_entities(query)

        # 2. 在图谱中查找匹配实体
        subgraph_context = []
        for entity in entities:
            neighbors = self.graph_db.query_neighbors(entity, max_depth=2)
            subgraph_context.extend(neighbors)

        # 3. 向量检索补充
        vector_results = self.vector_store.search(query, top_k=top_k)

        # 4. 组合上下文生成回答
        context = self._format_context(subgraph_context, vector_results)
        return self.llm.generate(query, context)

    def global_search(self, query: str) -> str:
        """全局搜索：社区摘要 + Map-Reduce

        适合：全局概览型问题
        例如："所有大模型的技术趋势是什么？"
        """
        # 1. 检索相关社区摘要
        community_summaries = self._search_communities(query)

        # 2. Map: 每个社区独立回答
        intermediate_answers = []
        for summary in community_summaries:
            answer = self.llm.generate(
                query, 
                f"基于以下社区信息回答：\n{summary}"
            )
            intermediate_answers.append(answer)

        # 3. Reduce: 合并所有中间回答
        final_context = "\n\n".join(intermediate_answers)
        return self.llm.generate(
            query,
            f"基于以下多个来源的信息，综合回答问题：\n{final_context}"
        )

    def _extract_entities(self, text: str) -> list:
        """使用 LLM 从文本中抽取实体"""
        prompt = f"""从以下文本中提取关键实体名称，每行一个：

文本：{text}

实体列表："""
        response = self.llm.generate(prompt)
        return [e.strip() for e in response.strip().split('\n') if e.strip()]

import json

def extract_triples_from_text(llm, text: str) -> list:
    """使用 LLM 从文本中抽取三元组"""

    prompt = f"""请从以下文本中抽取知识三元组 (subject, predicate, object)。

要求：
1. 实体名称要规范统一
2. 关系要简洁明确
3. 以 JSON 数组格式返回

文本：
{text}

输出格式示例：
[
    {{"subject": "实体1", "predicate": "关系", "object": "实体2"}},
    ...
]

三元组："""

    response = llm.generate(prompt)

    try:
        triples = json.loads(response)
        return [(t["subject"], t["predicate"], t["object"]) for t in triples]
    except json.JSONDecodeError:
        return []

# 示例
sample_text = """
DeepSeek-V3 是深度求索公司于2024年12月发布的大语言模型。
该模型采用 MoE（混合专家）架构，总参数量671B，激活参数37B。
训练使用了14.8T tokens的数据，在多个基准测试上接近GPT-4水平。
"""

# triples = extract_triples_from_text(llm, sample_text)
# 预期输出：
# [
#   ("DeepSeek-V3", "发布者", "深度求索"),
#   ("DeepSeek-V3", "发布时间", "2024年12月"),
#   ("DeepSeek-V3", "架构", "MoE"),
#   ("DeepSeek-V3", "总参数量", "671B"),
#   ("DeepSeek-V3", "激活参数", "37B"),
#   ("DeepSeek-V3", "训练数据量", "14.8T tokens"),
#   ("DeepSeek-V3", "性能水平", "接近GPT-4"),
# ]

class EntityResolver:
    """实体消歧与对齐"""

    def __init__(self, llm):
        self.llm = llm
        self.entity_aliases = {}  # 实体别名映射

    def resolve(self, entity_name: str) -> str:
        """将实体名称规范化"""
        # 检查已知别名
        for canonical, aliases in self.entity_aliases.items():
            if entity_name.lower() in [a.lower() for a in aliases]:
                return canonical

        # 使用 LLM 进行消歧
        prompt = f"""以下实体名称是否指向同一个实体？

已知实体：{list(self.entity_aliases.keys())}
待消歧实体：{entity_name}

如果是已知实体的别名，返回规范名称。否则返回原始名称。"""

        resolved = self.llm.generate(prompt).strip()
        return resolved

    def add_alias(self, canonical: str, alias: str):
        """添加实体别名"""
        if canonical not in self.entity_aliases:
            self.entity_aliases[canonical] = [canonical]
        self.entity_aliases[canonical].append(alias)

# 使用
resolver = EntityResolver(llm)
resolver.add_alias("DeepSeek", "深度求索")
resolver.add_alias("DeepSeek", "深度求索公司")
resolver.add_alias("DeepSeek-V3", "DeepSeek V3")

print(resolver.resolve("深度求索"))  # → "DeepSeek"

def retrieve_subgraph(graph_db, query_entities: list, max_depth: int = 2) -> str:
    """基于查询实体的子图检索"""
    all_context = []

    for entity in query_entities:
        # 获取实体周围 N 跳的子图
        neighbors = graph_db.query_neighbors(entity, max_depth=max_depth)

        for record in neighbors:
            context_line = f"- {entity} --[{record['relation']}]--> {record['name']}"
            all_context.append(context_line)

    return "\n".join(all_context)

# 示例
context = retrieve_subgraph(db, ["DeepSeek-V3"])
print(context)
# - DeepSeek-V3 --[发布者]--> 深度求索
# - DeepSeek-V3 --[架构]--> MoE
# - DeepSeek-V3 --[参数量]--> 671B
# - MoE --[全称]--> Mixture of Experts

class HybridGraphRAG:
    """混合检索：向量检索 + 图谱检索"""

    def __init__(self, vector_store, graph_db, llm):
        self.vector_store = vector_store
        self.graph_db = graph_db
        self.llm = llm

    def search(self, query: str, top_k: int = 5) -> dict:
        """混合检索"""

        # 1. 向量检索（语义相似度）
        vector_results = self.vector_store.search(query, top_k=top_k)

        # 2. 图谱检索（结构化关系）
        entities = self._extract_entities(query)
        graph_context = []
        for entity in entities:
            neighbors = self.graph_db.query_neighbors(entity, max_depth=2)
            graph_context.extend(neighbors)

        # 3. 合并上下文
        combined_context = self._merge_contexts(vector_results, graph_context)

        # 4. 生成回答
        answer = self.llm.generate(
            query=query,
            context=combined_context
        )

        return {
            "answer": answer,
            "vector_results": len(vector_results),
            "graph_triples": len(graph_context),
            "entities_found": entities,
        }

    def _merge_contexts(self, vector_results, graph_context):
        """合并向量检索和图谱检索的上下文"""
        context_parts = ["=== 相关文档片段 ==="]
        for i, result in enumerate(vector_results):
            context_parts.append(f"[文档{i+1}] {result['text']}")

        context_parts.append("\n=== 知识图谱关系 ===")
        for record in graph_context:
            context_parts.append(f"- {record}")

        return "\n".join(context_parts)

from llama_index.core import KnowledgeGraphIndex, SimpleDirectoryReader
from llama_index.core.graph_stores import SimpleGraphStore
from llama_index.llms.openai import OpenAI

# 1. 加载文档
documents = SimpleDirectoryReader("./data").load_data()

# 2. 构建知识图谱索引
graph_store = SimpleGraphStore()
llm = OpenAI(model="gpt-5-mini")

kg_index = KnowledgeGraphIndex.from_documents(
    documents,
    graph_store=graph_store,
    llm=llm,
    max_triplets_per_chunk=10,  # 每个文本块最多抽取10个三元组
    show_progress=True,
)

# 3. 查询
query_engine = kg_index.as_query_engine(
    include_text=True,      # 包含原文
    response_mode="tree_summarize",
    embedding_mode="hybrid",  # 混合模式：向量 + 图谱
    sim_score_threshold=0.5,
)

response = query_engine.query("DeepSeek-V3 的技术特点是什么？")
print(response)

class AdaptiveRAG:
    """自适应 RAG：根据查询类型选择策略"""

    def __init__(self, vector_rag, graph_rag, llm):
        self.vector_rag = vector_rag
        self.graph_rag = graph_rag
        self.llm = llm

    def query(self, question: str) -> str:
        # 1. 分类查询类型
        query_type = self._classify_query(question)

        if query_type == "simple_fact":
            return self.vector_rag.query(question)
        elif query_type == "multi_hop":
            return self.graph_rag.local_search(question)
        elif query_type == "global_overview":
            return self.graph_rag.global_search(question)
        else:
            # 混合模式
            vector_answer = self.vector_rag.query(question)
            graph_answer = self.graph_rag.local_search(question)
            return self._merge_answers(question, vector_answer, graph_answer)

    def _classify_query(self, question: str) -> str:
        prompt = f"""将以下问题分类为：
- simple_fact: 简单事实查询
- multi_hop: 需要多步推理
- global_overview: 全局概览

问题：{question}
分类："""
        return self.llm.generate(prompt).strip().lower()