GraphRAG：知识图谱 + 向量检索的检索增强生成实践

from neo4j import GraphDatabase

driver = GraphDatabase.driver("bolt://localhost:7687", auth=("neo4j", "password"))

def link_entities(question_entities):
    """把 NER/LLM 抽取的实体名链接到图内实体（别名优先）"""
    with driver.session() as session:
        result = session.run("""
            MATCH (a:EntityAlias {alias: $alias})
            RETURN a.entityId AS eid, a.entityName AS name, a.type AS type
            UNION
            MATCH (e:Entity {name: $alias})
            RETURN e.id AS eid, e.name AS name, e.type AS type
        """, alias=question_entities[0])  # 循环处理每个实体
        return [r.data() for r in result]

def hybrid_retrieve(question, top_k_text=5, top_k_graph=20):
    # 1. 向量召回：找语义相关的实体描述块
    vector_hits = vector_search(question, top_k_text)

    # 2. 从文本块中提取实体，链接到图谱
    entities = extract_entities(vector_hits + question)
    entity_ids = link_to_graph(entities)

    # 3. 图谱展开：沿关系扩展 1-2 跳
    subgraph = expand_subgraph(entity_ids, top_k_graph)

    # 4. 组合上下文（文本块 + 三元组 + 子图描述）
    return build_prompt(vector_hits, subgraph)

# Docker 启动时启用
docker run \
  --name neo4j-graphrag \
  -p 7474:7474 -p 7687:7687 \
  -e NEO4J_AUTH=neo4j/password \
  -e NEO4J_PLUGINS='["apoc"]' \
  neo4j:5.15-community

# 生成实体描述嵌入并写入
import openai
from neo4j import GraphDatabase

openai_client = openai.OpenAI()
driver = GraphDatabase.driver("bolt://localhost:7687", auth=("neo4j", "password"))

def embed(text):
    resp = openai_client.embeddings.create(
        model="text-embedding-3-small", input=text
    )
    return resp.data[0].embedding

def index_entity(eid, description):
    vec = embed(description)
    with driver.session() as session:
        session.run(
            "MATCH (e:Entity {id: $eid}) SET e.embedding = $vec",
            eid=eid, vec=vec
        )

from langchain_community.graphs import Neo4jGraph
from langchain_community.chains.graph_qa.cypher import GraphCypherQAChain
from langchain_openai import ChatOpenAI

graph = Neo4jGraph(
    url="bolt://localhost:7687",
    username="neo4j", password="password",
    database="neo4j"
)

llm = ChatOpenAI(model="gpt-4o-mini", temperature=0)

chain = GraphCypherQAChain.from_llm(
    llm=llm, graph=graph, verbose=True,
    allow_dangerous_requests=True,   # 允许 LLM 生成 Cypher
    validate_cypher=True             # 用 EXPLAIN 验证生成语句
)

answer = chain.invoke("华为和阿里云在 AI 领域有什么关系？")
print(answer)

class GraphRAG:
    def __init__(self, uri, user, pw, llm_client):
        self.driver = GraphDatabase.driver(uri, auth=(user, pw))
        self.llm = llm_client

    def answer(self, question):
        # Step 1: 实体抽取（LLM 解析）
        entities = self.extract_entities(question)

        # Step 2: 向量召回 + 图扩展 混合检索
        qvec = self.embed(question)
        graph_context = self.hybrid_search(qvec, entities)

        # Step 3: 组织上下文提示
        prompt = self.build_prompt(question, graph_context)

        # Step 4: 生成（要求给出事实来源）
        return self.llm.chat.completions.create(
            model="gpt-4o",
            messages=[{"role": "user", "content": prompt}]
        ).choices[0].message.content

    def hybrid_search(self, qvec, entities):
        facts = []
        # 向量召回
        with self.driver.session() as s:
            r = s.run(
                "CALL db.index.vector.queryNodes('entity_embedding', 5, $qvec)"
                " YIELD node, score RETURN node.name, score",
                qvec=qvec
            )
            facts += [x.data() for x in r]
            # 图扩展
            r = s.run(
                "MATCH (e:Entity)-[rel]-(n) WHERE e.name IN $names"
                " RETURN e.name, type(rel) AS rel, n.name LIMIT 30",
                names=entities
            )
            facts += [x.data() for x in r]
        return facts

    def build_prompt(self, question, facts):
        return f"""基于以下知识图谱事实回答。每个事实都是 {head}-{rel}-{tail} 结构。
事实：
{chr(10).join(str(f) for f in facts)}

问题：{question}
要求：如果事实不足，明确回答"知识不足"；引用事实编号说明依据。"""

# 构建评测集：2-hop 问答（A 与 C 通过 B 关联）
eval_set = [
    {
        "question": "任正非创立的公司在深圳发布了什么芯片？",
        "expected_path": [("任正非","FOUNDED_BY","华为"),("华为","DEVELOPED","昇腾AI芯片")],
        "answer_keywords": ["昇腾AI芯片"]
    },
    # ... 覆盖 1-hop、2-hop、全局三类
]

def evaluate(chain, eval_set):
    correct = 0
    for item in eval_set:
        answer = chain(item["question"])
        hit = any(kw in answer for kw in item["answer_keywords"])
        correct += hit
        print(f"{'✓' if hit else '✗'} {item['question']}")
    print(f"Acc: {correct}/{len(eval_set)}")