GraphRAG 检索增强生成

NEO4J_URI = "bolt://localhost:7687"
NEO4J_USERNAME = "neo4j"
NEO4J_PASSWORD = "neo4j#2026"
NEO4J_DATABASE = "neo4j"

QWEN_API_KEY = os.getenv("QWEN_API_KEY")               # 调用前须 export
QWEN_BASE_URL = "https://dashscope.aliyuncs.com/compatible-mode/v1"
QWEN_EMBEDDING_MODEL = "text-embedding-v3"
EMBEDDING_DIMS = 1024                                  # 必须与向量索引维度保持一致

def create_driver() -> neo4j.Driver:
    return neo4j.GraphDatabase.driver(NEO4J_URI, auth=(NEO4J_USERNAME, NEO4J_PASSWORD))

def create_qwen_embedder() -> OpenAIEmbeddings:
    return OpenAIEmbeddings(model=QWEN_EMBEDDING_MODEL, api_key=QWEN_API_KEY, base_url=QWEN_BASE_URL)

def setup_index() -> None:
    kg.query(f"DROP INDEX {INDEX_NAME} IF EXISTS")
    kg.query(f"""
        CREATE VECTOR INDEX {INDEX_NAME} IF NOT EXISTS
        FOR (m:Movie) ON (m.taglineEmbedding)
        
OPTIONS {{ indexConfig: {{
            `vector.dimensions`: {EMBEDDING_DIMS},
            `vector.similarity_function`: 'cosine'
        }}}}

    """)

movies = kg.query("""
    MATCH (movie:Movie) WHERE movie.tagline IS NOT NULL
    RETURN movie.title AS title, movie.tagline AS tagline
""")
vectors = qwen_embeddings.embed_documents([m["tagline"] for m in movies])
for title, vec in zip(titles, vectors):
    kg.query("""
        MATCH (m:Movie {title: $title})
        CALL db.create.setNodeVectorProperty(m, 'taglineEmbedding', $vector)
    """, params={"title": title, "vector": vec})

def vector_search(question: str, top_k: int = 5, embeddings = None):
    if embeddings is None:
        embeddings = QwenEmbeddings(model=EMBEDDING_MODEL)
    q_vec = embeddings.embed_query(question)
    return kg.query("""
        CALL db.index.vector.queryNodes($index_name, $top_k, $question_embedding)
        YIELD node AS movie, score
        RETURN movie.title AS title, movie.tagline AS tagline, score
    """, params={"index_name": INDEX_NAME, "top_k": top_k, "question_embedding": q_vec})

SimpleKGPipeline(
    llm=llm,
    driver=driver,
    text_splitter=FixedSizeSplitter(chunk_size=500, chunk_overlap=100),
    embedder=embedder,
    entities=NODE_LABELS,                # 27 个医学/学术/基础标签
    relations=REL_TYPES,                 # 6 种关系类型
    prompt_template=PROMPT_TEMPLATE,     # 指导 LLM 输出标准 JSON
    from_file=True,
)

BASIC_NODE_LABELS = ["Object", "Entity", "Group", "Researcher", "Organization", "Place"]
ACADEMIC_NODE_LABELS = ["ArticleOrPaper", "PublicationOrJournal"]
MEDICAL_NODE_LABELS = [
    "Anatomy", "BiologicalProcess", "Cell", "CellularComponent",
    "CellType", "Condition", "Disease", "Drug",
    "EffectOrPhenotype", "Exposure", "GeneOrProtein", "Molecule",
    "MolecularFunction", "Pathway",
]
NODE_LABELS = BASIC_NODE_LABELS + ACADEMIC_NODE_LABELS + MEDICAL_NODE_LABELS

REL_TYPES = [
    "ACTIVATES", "AFFECTS", "ASSESSES", "ASSOCIATED_WITH",
    "AUTHORED", "BIOMARKER_FOR",
]

results = await asyncio.gather(
    *(process_pdf(pipeline, path) for path in pdf_paths),
    return_exceptions=True,             # 单个文件失败不中断其他文件的处理
)

def main():
    driver = create_driver()
    try:
        embedder = create_qwen_embedder()
        ensure_vector_index(driver, embedder)

        # 演示 1: 纯向量召回
        retriever = build_retriever(driver, embedder)
        vector_result = retriever.search(query_text=vector_question, top_k=5)
        print_vector_search_result(vector_result, vector_question, 5)

        # 演示 2: 向量召回叠加 Cypher 拉取关联实体关系
        vc_retriever = build_vector_cypher_retriever(driver, embedder)
        vc_res = vc_retriever.get_search_results(query_text=cypher_question, top_k=3)
        print_vector_cypher_result(vc_res, cypher_question, 3)

        # 演示 3: GraphRAG（基于前述两种 retriever，由 DeepSeek 生成回答）
        llm = create_llm()
        prompt_template = build_rag_prompt_template()
        v_rag  = build_graph_rag(llm, retriever, prompt_template)
        vc_rag = build_graph_rag(llm, vc_retriever, prompt_template)
        print_graph_rag_answers(rag_question,
            v_rag.search(rag_question, retriever_config={"top_k": 5}).answer,
            vc_rag.search(rag_question, retriever_config={"top_k": 5}).answer)
    finally:
        driver.close()

VECTOR_CYPHER_RETRIEVAL_QUERY = """
WITH node AS chunk
MATCH (chunk)<-[:FROM_CHUNK]-()-[relList:!FROM_CHUNK]-{1,2}()
UNWIND relList AS rel
WITH collect(DISTINCT chunk) AS chunks,
     collect(DISTINCT rel) AS rels
RETURN '=== text ===\\n' + apoc.text.join([c in chunks | c.text], '\\n---\\n')
     + '\\n\\n=== kg_rels ===\\n'
     + apoc.text.join([r in rels |
         startNode(r).name + ' - ' + type(r) + '(' + coalesce(r.details, '') + ') -> '
         + endNode(r).name], '\\n---\\n') AS info
"""

DEEPSEEK_MODEL = "deepseek-v4-pro"
DEEPSEEK_BASE_URL = "https://api.deepseek.com"
DEEPSEEK_API_KEY = os.getenv("DEEPSEEK_API_KEY")

def create_llm() -> LLM:
    return LLM(
        model_name=DEEPSEEK_MODEL,
        model_params={"temperature": 0},
        api_key=DEEPSEEK_API_KEY,
        base_url=DEEPSEEK_BASE_URL,
    )

RAG_TEMPLATE = """Answer the Question using the following Context. Only respond with
information mentioned in the Context. Do not inject any speculative information not mentioned.

# Question:
{query_text}

# Context:
{context}

# Answer:
"""

def build_graph_rag(llm, retriever, prompt_template) -> GraphRAG:
    return GraphRAG(llm=llm, retriever=retriever, prompt_template=prompt_template)