开发者中心 » 编程语言 » Python » 代码指南 » Python、OpenAI 与 GraphRAG

Python、OpenAI 与 GraphRAG

本指南展示了如何使用 Neo4j 的官方 GraphRAG 包并集成 OpenAI 来设置和运行 GraphRAG（图检索增强生成）。该示例使用了一个包含 Book 节点及其相关实体的 Goodreads 数据集的公共演示数据库。

什么是 GraphRAG？

GraphRAG 将知识图谱与人工智能相结合，通过以下方式提供智能问答：

  1. 嵌入（Embedding）您的问题：使用 OpenAI 的模型
  2. 查找相关信息：利用向量相似度在 Neo4j 图数据库中进行检索
  3. 生成综合答案：使用 OpenAI 的大语言模型

创建项目

创建一个名为 app.py 的文件。
接下来，创建一个虚拟环境，启动它，并下载 Neo4j 相关包。

python -m venv venv
source venv/bin/activate  # On Windows: venv\Scripts\activate
pip install neo4j "neo4j-graphrag[openai]"

配置凭据和嵌入详情

import os
from neo4j import GraphDatabase
from neo4j_graphrag.embeddings import OpenAIEmbeddings
from neo4j_graphrag.generation import GraphRAG
from neo4j_graphrag.llm import OpenAILLM
from neo4j_graphrag.retrievers import VectorRetriever

NEO4J_URI = "neo4j+s://demo.neo4jlabs.com"
NEO4J_USERNAME = "goodreads"
NEO4J_PASSWORD = "goodreads"
NEO4J_DATABASE = "goodreads"
VECTOR_INDEX_NAME = "book-descriptions"
LLM_MODEL = "gpt-4o"
EMBEDDING_MODEL = "text-embedding-3-small"
OPENAI_API_KEY = "your-openai-api-key-here"  # Replace with your key

创建流水线（Pipeline）

GraphRAG 系统结合了 Neo4j 和 OpenAI 来提供智能问答。流水线设置会创建所有必要的 AI 组件。在本示例中，我们定义了一个函数，用于初始化包含嵌入模型、大语言模型和向量检索器的完整 GraphRAG 流水线。

def create_graphrag_pipeline():
    """Initialize the GraphRAG components"""
    # Set OpenAI API key
    os.environ["OPENAI_API_KEY"] = OPENAI_API_KEY

    # Connect to Neo4j
    driver = GraphDatabase.driver(NEO4J_URI, auth=(NEO4J_USERNAME, NEO4J_PASSWORD))

    # Create embeddings and LLM
    embedder = OpenAIEmbeddings(model=EMBEDDING_MODEL)
    llm = OpenAILLM(model_name=LLM_MODEL, model_params={"temperature": 0.0})

    # Create retriever
    retriever = VectorRetriever(
        driver=driver,
        index_name=VECTOR_INDEX_NAME,
        embedder=embedder,
        neo4j_database=NEO4J_DATABASE
    )

    # Create GraphRAG pipeline
    rag = GraphRAG(retriever=retriever, llm=llm)

    return rag, driver

该函数创建了四个关键组件：

  • embedder（嵌入器）：将问题转换为向量，以便进行相似度搜索
  • llm（大语言模型）：使用 LLM 模型生成智能答案
  • retriever（检索器）：在 Neo4j 数据库中查找相关信息
  • rag：将所有组件组合成一个完整的问答系统

运行查询

GraphRAG 搜索 API 会自动处理向量相似度搜索、上下文检索和答案生成。在本示例中，我们定义了一个主函数，它接受自然语言问题并返回 AI 生成的答案。

def ask_question(question, top_k=5):
    """Ask a question and get an AI answer from the book database"""
    rag, driver = create_graphrag_pipeline()

    try:
        print(f"🔍 Question: {question}")

        # Search and get answer
        response = rag.search(
            query_text=question,
            retriever_config={"top_k": top_k}
        )

        print(f"🤖 Answer: {response.answer}")
        return response.answer

    except Exception as e:
        print(f"❌ Error: {e}")
        return None
    finally:
        driver.close()

该方法接受两个参数：

  • question：您关于书籍的自然语言问题
  • top_k：用于生成答案的相关书籍数量（默认值：5）

该函数根据数据库中找到的最相关的书籍，返回 AI 生成的答案。

示例查询

以下示例展示了您可以向 Goodreads 书籍数据库提出的不同类型的问题。

if __name__ == "__main__":
    # Make sure to set your OpenAI API key above!

    # Ask questions about books
    questions = [
        "Which books are about motherhood and friendship?",
        "What are the best science fiction novels?",
        "Tell me about books that deal with mental health"
    ]

    for question in questions:
        print("=" * 60)
        ask_question(question, top_k=3)
        print()

运行脚本

您可以通过命令行运行脚本：

python app.py

运行前请确保已在代码中设置了您的 OpenAI API 密钥。

响应示例

第一个问题应该返回类似以下的响应：

🔍 Question: Which books are about motherhood and friendship?
🤖 Answer: The books about motherhood and friendship are "Friendship Crisis" and "Good Harbor." "Friendship Crisis" explores the keys to forming emotionally supportive connections, which often include themes of motherhood and friendship. "Good Harbor" examines the balance of marriage, career, motherhood, and friendship, highlighting the redemptive power of friendship.Code language: PHP (php)

资源

分享文章