基于知识图谱增强的RAG系统阅读笔记(四)自然语言问题生成Cypher查询

第四章 从自然语言问题生成Cypher查询

[!NOTE]

Cypher 是 Neo4j 图数据库专用的声明式查询语言，专为高效地查询和操作图数据（节点、关系、属性）而设计。它的语法直观、可读性强，灵感来源于 SQL，但更贴近图结构的可视化表达。

4.1 基础知识

查询语言生成的目标是将自然语言转换成可以在数据库执行的查询语言。当前我们感兴趣的是从自然语言问题生成Cypher查询，大多数大语言模型都知道Cypher是什么，并且了解该语言的基本语法。此过程中的主要挑战是生成一个既正确又与所提问题相关的查询。这需要理解问题的语义，以及被查询知识图谱的模式（schema）。如果我们不提供知识图谱的模式，大语言模型只能假设节点、关系和属性的名称。当提供了模式时，它就充当了用户问题语义与所使用的图模型之间的映射——节点上使用的标签、存在的关系类型、可用的属性，以及节点之间通过哪些关系类型连接。

从自然语言问题生成Cypher查询的工作流程可以分解为以下步骤（图4.1）：

• 从用户处获取问题。
• 获取知识图谱的模式。
• 定义其他有用的信息，例如术语映射、格式说明和少量示例（few-shot examples）。
• 为大语言模型生成提示（prompt）。
• 将提示传递给大语言模型以生成Cypher查询。

向量相似性搜索的一个局限性在于它们能够回答的问题类型是受限的，如问题 “列出史蒂文·斯皮尔伯格执导的评分最高的三部电影及其平均得分。” 这个问题永远无法通过向量相似性搜索来回答，因为需要执行特定查询。相关代码如下：

import random
from neo4j import GraphDatabase

NEO4J_URI = "bolt://localhost:7687"      # Neo4j Bolt 连接地址
NEO4J_USER = "neo4j"                          # 数据库用户名
NEO4J_PASSWORD = "password"              # 数据库密码

# 初始化 Neo4j 驱动
driver = GraphDatabase.driver(NEO4J_URI, auth=(NEO4J_USER, NEO4J_PASSWORD))


def clear_database():
    """
    清空 Neo4j 数据库中所有节点和关系
    使用 DETACH DELETE 确保即使有关系也能彻底删除节点
    """
    with driver.session() as session:
        session.run("MATCH (n) DETACH DELETE n")
    print("数据库已清空")


def create_constraints():
    """
    创建唯一性约束，防止重复创建相同的节点
    约束将应用于 Director、Movie 和 Reviewer 的 name/title 字段
    """
    with driver.session() as session:
        session.run("CREATE CONSTRAINT IF NOT EXISTS FOR (d:Director) REQUIRE d.name IS UNIQUE")
        session.run("CREATE CONSTRAINT IF NOT EXISTS FOR (m:Movie) REQUIRE m.title IS UNIQUE")
        session.run("CREATE CONSTRAINT IF NOT EXISTS FOR (r:Reviewer) REQUIRE r.name IS UNIQUE")
    print("唯一性约束已创建")


def populate_movie_graph():
    """
    向数据库中插入电影图谱数据，包括：
    - 导演（Director）
    - 电影（Movie）
    - 评分者（Reviewer）
    - 评分关系（REVIEWED，包含 score 属性）
    - 导演与电影的关系（DIRECTED）
    """
    # 斯皮尔伯格执导的电影列表（真实作品）
    spielberg_movies = [
        "Schindler's List",
        "Saving Private Ryan",
        "Jurassic Park",
        "E.T. the Extra-Terrestrial",
        "Catch Me If You Can",
        "Minority Report",
        "Jaws",
        "Lincoln",
        "The Post",
        "Bridge of Spies"
    ]

    # 其他导演及其电影（用于丰富图谱）
    other_movies = [
        ("Christopher Nolan", "Inception"),
        ("Christopher Nolan", "Interstellar"),
        ("Martin Scorsese", "The Departed"),
        ("Martin Scorsese", "Goodfellas"),
        ("Quentin Tarantino", "Pulp Fiction"),
        ("James Cameron", "Avatar"),
        ("James Cameron", "Titanic")
    ]

    # 评分者名称列表
    reviewers = ["Alice", "Bob", "Charlie", "Diana", "Eve", "Frank", "Grace", "Henry"]

    with driver.session() as session:
        # 1. 创建导演：Steven Spielberg
        session.run("MERGE (:Director {name: 'Steven Spielberg'})")

        # 2. 创建斯皮尔伯格的电影，并建立 DIRECTED 关系
        for title in spielberg_movies:
            session.run("MERGE (:Movie {title: $title})", title=title)
            session.run("""
                MATCH (d:Director {name: 'Steven Spielberg'})
                MATCH (m:Movie {title: $title})
                MERGE (d)-[:DIRECTED]->(m)
            """, title=title)

        # 3. 创建其他导演及其电影
        for director_name, title in other_movies:
            session.run("MERGE (:Director {name: $name})", name=director_name)
            session.run("MERGE (:Movie {title: $title})", title=title)
            session.run("""
                MATCH (d:Director {name: $director})
                MATCH (m:Movie {title: $title})
                MERGE (d)-[:DIRECTED]->(m)
            """, director=director_name, title=title)

        # 4. 创建评分者，并为电影生成评分
        for reviewer_name in reviewers:
            # 创建评分者节点
            session.run("MERGE (:Reviewer {name: $name})", name=reviewer_name)

            # 随机选择 5-8 部电影进行评分
            all_movies = spielberg_movies + [t for _, t in other_movies]
            sampled_movies = random.sample(all_movies, k=random.randint(5, 8))

            for title in sampled_movies:
                # 斯皮尔伯格的电影基础分更高
                base_score = 8.0 if title in spielberg_movies else 6.5
                # 添加随机扰动，模拟真实评分分布
                score = round(random.gauss(base_score, 1.2), 1)
                score = max(1.0, min(10.0, score))  # 限制在 1.0 到 10.0 之间

                # 创建 REVIEWED 关系并设置 score 属性
                session.run("""
                    MATCH (r:Reviewer {name: $reviewer})
                    MATCH (m:Movie {title: $title})
                    MERGE (r)-[rev:REVIEWED]->(m)
                    SET rev.score = $score
                """, reviewer=reviewer_name, title=title, score=score)

    print("电影图谱数据已写入")


def run_spielberg_top_movies_query():
    """
    执行 Cypher 查询：找出斯皮尔伯格执导的电影中，平均评分最高的前三部
    原始查询中的 'AS avg rating' 存在空格语法错误，已修正为 'AS avg_rating'
    """
    query = """
    MATCH (:Reviewer)-[r:REVIEWED]->(m:Movie)<-[:DIRECTED]-(:Director {name: 'Steven Spielberg'})
    RETURN m.title AS title, AVG(r.score) AS avg_rating
    ORDER BY avg_rating DESC
    LIMIT 3
    """
    with driver.session() as session:
        result = session.run(query)
        records = [dict(record) for record in result]

    if records:
        print("斯皮尔伯格执导的评分最高的 3 部电影：")
        print("-" * 50)
        print(f"{'排名':<4} {'电影名称':<30} {'平均分':<8}")
        print("-" * 50)
        for index, record in enumerate(records, start=1):
            print(f"{index:<4} {record['title']:<30} {record['avg_rating']:>6.2f}")
        print("-" * 50)
    else:
        print("查询未返回结果，请检查数据是否正确写入")


def main():
    """
    主函数：按顺序执行以下操作
    1. 清空数据库
    2. 创建约束
    3. 写入电影图谱数据
    4. 执行 Top 3 查询
    """
    try:
        clear_database()
        create_constraints()
        populate_movie_graph()
        run_spielberg_top_movies_query()
        print("脚本执行完成")
    except Exception as e:
        print(f"执行过程中发生错误: {e}")
    finally:
        driver.close()


if __name__ == "__main__":
    main()

斯皮尔伯格执导的评分最高的 3 部电影：
--------------------------------------------------
排名   电影名称                           平均分     
--------------------------------------------------
1    Catch Me If You Can              9.30
2    Minority Report                  8.70
3    Lincoln                          8.37
--------------------------------------------------

要实现这个操作不仅仅需要进行相似性计算，还需要实现一些聚合查询。所以我们需要实Text2cypher。

4.2 查询语言生成

**从自然语言问题生成Cypher查询时，要确保生成的查询是正确且相关的。**简单来说就是如果大模型对你的数据库一点都不了解，其生成的内容大概率也没什么用。特别是当你输入的问题很复杂，或者你的数据库元素命名不是很规范的时候。

4.2.1 少量示例

少量示例（few-shot examples）通过在提示词中添加少量样本提升大语言模型在text2cypher任务中性能。但通常需要为不用的知识图谱定制提示词，需要具体问题具体分析。

例如我们假设国家是当前数据库中一些独立的节点，而不是电影节点的属性。这个时候我们可以使用如下的提示词:

以下是几个正确的示例：

问题：《头号玩家》（Ready Player One）是在哪个国家制作的？
Cypher：MATCH (m:Movie {title: 'Ready Player One'})-[:PRODUCED_IN]->(c:Country) RETURN c.name

问题：列出由克里斯托弗·诺兰执导的所有电影。
Cypher：MATCH (d:Director {name: 'Christopher Nolan'})-[:DIRECTED]->(m:Movie) RETURN m.title

问题：电影《阿凡达》的评分是多少？
Cypher：MATCH (:Reviewer)-[r:REVIEWED]->(m:Movie {title: 'Avatar'}) RETURN AVG(r.score) AS average_rating

4.2.2 展示知识图谱结构

大语言模型（LLM）需要知道图中的节点标签（Node Labels）、关系类型（Relationship Types）以及属性（Properties，我们同样是需要将其整理成结构化数据发送。如下提供了基于APOC库的实现。

import requests
from neo4j import GraphDatabase
from typing import Dict, List
import re

NEO4J_URI = "bolt://localhost:7687"      # Neo4j Bolt 连接地址
NEO4J_USER = "neo4j"                          # 数据库用户名
NEO4J_PASSWORD = "password"              # 数据库密码
# Ollama 服务地址和使用的模型
OLLAMA_BASE_URL = "http://localhost:11434"
LLM_MODEL = "qwen3:14b"
# 创建连接到 Neo4j 的驱动
driver = GraphDatabase.driver(NEO4J_URI, auth=(NEO4J_USER, NEO4J_PASSWORD))

# 步骤 1：定义 APOC 查询语句
# 查询所有节点标签及其属性（例如：Movie 有 title 属性）
NODE_PROPERTIES_QUERY = """
CALL apoc.meta.data()
YIELD label, elementType, type, property
WHERE elementType = "node" AND type <> "RELATIONSHIP"
WITH label, collect({property: property, type: type}) AS properties
RETURN {labels: label, properties: properties} AS output
"""

# 查询所有关系类型及其属性（例如：REVIEWED 有 score 属性）
REL_PROPERTIES_QUERY = """
CALL apoc.meta.data()
YIELD label, elementType, type, property
WHERE elementType = "relationship" AND type <> "RELATIONSHIP"
WITH label, collect({property: property, type: type}) AS properties
RETURN {type: label, properties: properties} AS output
"""

# 查询图中所有关系的结构（例如：(:Director)-[:DIRECTED]->(:Movie)）
REL_QUERY = """
CALL apoc.meta.data()
YIELD label, elementType, type, property, other
WHERE type = "RELATIONSHIP" AND elementType = "node"
UNWIND other AS end_label
RETURN {start: label, type: property, end: toString(end_label)} AS output
"""


# 步骤 2：从数据库获取结构化模式
def get_structured_schema() -> Dict:
    """
    使用 APOC 的元数据功能，自动读取当前图数据库的结构。
    返回一个字典，包含：
    - node_props: 每个节点标签有哪些属性
    - rel_props: 每种关系类型有哪些属性
    - relationships: 节点之间如何连接（三元组）
    """
    with driver.session() as session:
        # 获取节点属性
        node_result = session.run(NODE_PROPERTIES_QUERY)
        node_properties = [record["output"] for record in node_result]

        # 获取关系属性
        rel_prop_result = session.run(REL_PROPERTIES_QUERY)
        rel_properties = [record["output"] for record in rel_prop_result]

        # 获取关系结构（起点 -> 关系 -> 终点）
        rel_result = session.run(REL_QUERY)
        relationships = [record["output"] for record in rel_result]

    # 整理成字典返回
    return {
        "node_props": {item["labels"]: item["properties"] for item in node_properties},
        "rel_props": {item["type"]: item["properties"] for item in rel_properties},
        "relationships": relationships,
    }


# 步骤 3：将结构格式化为易读的文本
def format_props(props: List[Dict]) -> str:
    """
    将属性列表格式化为 '属性名: 类型' 的字符串，例如：
    name: STRING, age: INTEGER
    """
    return ", ".join([f"{p['property']}: {p['type']}" for p in props])


def get_schema_string(schema: Dict) -> str:
    """
    将结构化模式转换为一段自然可读的文本，便于 LLM 理解图的结构。
    输出格式包括三部分：
    1. 节点标签和属性
    2. 关系类型和属性
    3. 节点之间的连接方式
    """
    # 格式化节点属性
    node_lines = [
        f"{label} {{{format_props(props)}}}"
        for label, props in schema["node_props"].items()
    ]

    # 格式化关系属性
    rel_lines = [
        f"{rtype} {{{format_props(props)}}}"
        for rtype, props in schema["rel_props"].items()
    ]

    # 格式化连接关系
    rel_structures = [
        f"(:{item['start']})-[:{item['type']}]->(:{item['end']})"
        for item in schema["relationships"]
    ]

    # 组合成完整字符串
    return "\n".join([
        "Node labels and properties:",
        "\n".join(node_lines) if node_lines else "None",
        "\nRelationship types and properties:",
        "\n".join(rel_lines) if rel_lines else "None",
        "\nThe relationships:",
        "\n".join(rel_structures) if rel_structures else "None"
    ])


# 步骤 4：生成发送给大语言模型的提示
def generate_prompt(question: str, schema_str: str) -> str:
    """
    构造一个清晰的提示（prompt），告诉 LLM：
    - 它的任务是生成 Cypher 查询
    - 图数据库有哪些结构
    - 当前问题是哪个
    """
    return f"""
你是一个专业的 Cypher 查询生成器，请根据以下图数据库结构，将自然语言问题转换为正确的 Cypher 查询。

请遵守规则：
- 只使用下面列出的标签、关系和属性
- 不要假设节点有未列出的属性
- 使用 MATCH 遍历关系，而不是直接访问不存在的字段
- 返回结果时使用别名提高可读性
- 电影标题使用英文，例如："Jurassic Park"
- 请仅返回一行 Cypher 查询，不要添加任何解释、注释或 <think> 标签
- 查询应以分号结尾
- 查询必须以有效的 Cypher 关键字开头，如 MATCH, CREATE, MERGE 等

图数据库结构如下：

{schema_str}

问题：{question}
Cypher：
""".strip()


# 步骤 5：调用本地大语言模型生成 Cypher
def call_llm(prompt: str) -> str:
    """
    将提示发送给 Ollama 中运行的 LLM（如 qwen3:14b）
    获取生成的 Cypher 查询文本
    """
    payload = {
        "model": LLM_MODEL,
        "prompt": prompt,
        "stream": False
    }
    response = requests.post(f"{OLLAMA_BASE_URL}/api/generate", json=payload)
    response.raise_for_status()
    return response.json()["response"].strip()

# 步骤 6：执行 Cypher 查询并返回结果
def execute_query(cypher: str) -> List[Dict]:
    """
    在 Neo4j 中执行生成的 Cypher 查询
    返回结果列表，每项是一个字典
    """
    with driver.session() as session:
        result = session.run(cypher)
        return [dict(record) for record in result]


# 步骤 7：创建电影图数据
def create_test_data():
    """
    创建一个简单的电影知识图谱用于测试，包括：
    - 导演、电影、国家、评分者
    - DIRECTED, PRODUCED_IN, REVIEWED 等关系
    """
    with driver.session() as session:
        # 创建国家节点
        session.run("MERGE (:Country {name: 'United States'})")

        # 创建斯皮尔伯格和他的电影
        session.run("""
        MERGE (d:Director {name: 'Steven Spielberg'})
        WITH d
        UNWIND [
            'Jurassic Park',
            'Schindler\\'s List',
            'Saving Private Ryan'
        ] AS title
        MERGE (m:Movie {title: title})
        MERGE (d)-[:DIRECTED]->(m)
        MERGE (m)-[:PRODUCED_IN]->(:Country {name: 'United States'})
        """)

        # 创建评分数据 - 简化版本
        session.run("""
        MERGE (r1:Reviewer {name: 'Alice'})
        MERGE (r2:Reviewer {name: 'Bob'})
        """)

        # 为《侏罗纪公园》添加评分
        session.run("""
        MATCH (m:Movie {title: 'Jurassic Park'})
        MATCH (r:Reviewer {name: 'Alice'})
        MERGE (r)-[:REVIEWED {score: 8.8}]->(m)
        """)

        # 为《辛德勒的名单》添加评分
        session.run("""
        MATCH (m:Movie {title: 'Schindler\\'s List'})
        MATCH (r:Reviewer {name: 'Alice'})
        MERGE (r)-[:REVIEWED {score: 9.0}]->(m)
        """)


# 步骤 8：提取纯净的 Cypher 查询(可以进行简化 但大模型的输出有时会出现问题)
def extract_cypher_query(response_text: str) -> str:
    """
    从 LLM 的响应中提取纯净的 Cypher 查询语句
    去除 <think> 标签内容和其他无关信息
    """
    # 去除 <think> 标签内容
    cleaned = re.sub(r'<think>.*?</think>', '', response_text, flags=re.DOTALL)

    # 查找包含完整 Cypher 查询的行
    lines = cleaned.strip().split('\n')

    # 寻找以有效 Cypher 关键字开头的行
    cypher_keywords = {
        'MATCH', 'CREATE', 'MERGE', 'DELETE', 'DETACH', 'WITH', 'UNWIND',
        'CALL', 'RETURN', 'SET', 'REMOVE', 'FOREACH', 'LOAD', 'START'
    }

    # 从后往前查找，因为 Cypher 查询通常在最后
    for i in range(len(lines) - 1, -1, -1):
        line = lines[i].strip()
        if not line:
            continue

        # 检查行是否以 Cypher 关键字开头
        first_word = line.split()[0].upper() if line.split() else ""
        if first_word in cypher_keywords:
            # 找到有效的 Cypher 查询行
            # 继续向后查找直到遇到空行或结束，收集完整的查询
            cypher_lines = []
            for j in range(i, len(lines)):
                current_line = lines[j].strip()
                if not current_line:  # 遇到空行则停止
                    break
                cypher_lines.append(current_line)

            if cypher_lines:
                # 合并所有 Cypher 行
                full_query = ' '.join(cypher_lines)
                # 清理可能的前缀（如 "Cypher:"）
                full_query = re.sub(r'^[Cc]ypher:\s*', '', full_query).strip()
                # 确保以分号结尾
                if not full_query.endswith(';'):
                    full_query += ';'
                return full_query

    # 如果没有找到以关键字开头的行，尝试其他方法
    # 查找包含 MATCH 的行作为备选
    for line in lines:
        line = line.strip()
        if line.upper().startswith('MATCH') and ('->' in line or '-[' in line):
            # 这很可能是一个 Cypher 查询
            if not line.endswith(';'):
                line += ';'
            return line

    # 最后的备选方案：返回清理后的文本，确保格式正确
    cleaned_final = cleaned.strip()
    if cleaned_final and not cleaned_final.endswith(';'):
        cleaned_final += ';'
    return cleaned_final


def main():
    # 1. 清空数据库
    with driver.session() as session:
        session.run("MATCH (n) DETACH DELETE n")
    print("[OK] 数据库已清空")
    # 2. 创建测试数据
    create_test_data()
    print("[OK] 测试数据已创建")
    # 3. 获取图结构
    print("\n[INFO] 正在自动推断图结构...")
    schema = get_structured_schema()
    schema_text = get_schema_string(schema)
    print("推断出的结构：")
    print(schema_text)
    # 4. 用户提问 - 使用英文电影名
    question = "Jurassic Park was produced in which country?"
    # 5. 生成提示
    prompt = generate_prompt(question, schema_text)
    print(f"\n[INFO] 发送给大语言模型的提示：\n{prompt}")
    # 6. 调用 LLM 生成 Cypher
    print("\n[INFO] 正在调用大语言模型...")
    try:
        raw_output = call_llm(prompt)
        print(f"LLM 原始响应：\n{raw_output}")
        # 提取纯净的 Cypher 查询
        cypher_query = extract_cypher_query(raw_output)
        print(f"提取的查询：\n{cypher_query}")
    except Exception as e:
        print(f"调用 LLM 失败：{e}")
        return
    # 7. 执行查询
    try:
        result = execute_query(cypher_query)
        print("\n[RESULT] 查询结果：")
        if result:
            for row in result:
                print(row)
        else:
            print("没有找到结果。")
    except Exception as e:
        print(f"执行 Cypher 失败：{e}")
    # 8. 关闭数据库连接
    driver.close()


# 运行主程序
if __name__ == "__main__":
    main()

[INFO] 正在自动推断图结构...
推断出的结构：
Node labels and properties:
Director {name: STRING}
Movie {title: STRING}
Reviewer {name: STRING}
Country {name: STRING}

Relationship types and properties:
REVIEWED {score: FLOAT}

The relationships:
(:Director)-[:DIRECTED]->(:Movie)
(:Movie)-[:PRODUCED_IN]->(:Country)
(:Reviewer)-[:REVIEWED]->(:Movie)

[INFO] 发送给大语言模型的提示：
你是一个专业的 Cypher 查询生成器，请根据以下图数据库结构，将自然语言问题转换为正确的 Cypher 查询。

请遵守规则：
- 只使用下面列出的标签、关系和属性
- 不要假设节点有未列出的属性
- 使用 MATCH 遍历关系，而不是直接访问不存在的字段
- 返回结果时使用别名提高可读性
- 电影标题使用英文，例如："Jurassic Park"
- 请仅返回一行 Cypher 查询，不要添加任何解释、注释或 <think> 标签
- 查询应以分号结尾
- 查询必须以有效的 Cypher 关键字开头，如 MATCH, CREATE, MERGE 等

图数据库结构如下：

Node labels and properties:
Director {name: STRING}
Movie {title: STRING}
Reviewer {name: STRING}
Country {name: STRING}

Relationship types and properties:
REVIEWED {score: FLOAT}

The relationships:
(:Director)-[:DIRECTED]->(:Movie)
(:Movie)-[:PRODUCED_IN]->(:Country)
(:Reviewer)-[:REVIEWED]->(:Movie)

问题：Jurassic Park was produced in which country?
Cypher：

[INFO] 正在调用大语言模型...
LLM 原始响应：
<think>
Okay, let's see. The user is asking which country produced Jurassic Park. So, the movie title is Jurassic Park, and I need to find the country it was produced in.

Looking at the graph structure, the Movie node has a PRODUCED_IN relationship to the Country node. So the steps would be to match the Movie with the title "Jurassic Park", then follow the PRODUCED_IN relationship to get the country's name.

I need to use the MATCH clause. The Movie node's property is title, so I'll write MATCH (m:Movie {title: "Jurassic Park"}). Then, the produced_in relationship is directed from Movie to Country, so I'll do m-[:PRODUCED_IN]->(c:Country). Then, return the country's name with c.name AS country.

Wait, the user said to use aliases for readability. So the Country node can be aliased as c, and then return c.name. Also, make sure the query ends with a semicolon. Let me check the rules again: only use the specified labels and relationships. The PRODUCED_IN is correct here. No other properties are involved except the title and the name of the country. That should cover it. So the final query should be:

MATCH (m:Movie {title: "Jurassic Park"})-[:PRODUCED_IN]->(c:Country) RETURN c.name AS country;
</think>

MATCH (m:Movie {title: "Jurassic Park"})-[:PRODUCED_IN]->(c:Country) RETURN c.name AS country;
提取的查询：
MATCH (m:Movie {title: "Jurassic Park"})-[:PRODUCED_IN]->(c:Country) RETURN c.name AS country;

[RESULT] 查询结果：
{'country': 'United States'}

4.2.3 术语映射

大语言模型需要知道如何将问题中使用的术语映射到模式中使用的术语。一个设计良好的图模式通常使用名词和动词作为标签和关系类型，使用形容词和名词作为属性。这些映射是特殊的，基于特定的图谱，需要具体问题具体分析。

import requests
from neo4j import GraphDatabase
import re

NEO4J_URI = "bolt://localhost:7687"      # Neo4j Bolt 连接地址
NEO4J_USER = "neo4j"                          # 数据库用户名
NEO4J_PASSWORD = "password"              # 数据库密码
# Ollama 服务地址和使用的模型
OLLAMA_BASE_URL = "http://localhost:11434"
LLM_MODEL = "qwen3:14b"

# 创建连接到 Neo4j 的驱动
driver = GraphDatabase.driver(NEO4J_URI, auth=(NEO4J_USER, NEO4J_PASSWORD))

# 步骤 1：定义术语映射
TERM_MAPPING = """
术语映射说明：
- “导演” 指的是 `Director` 节点，其姓名存储在 `name` 属性中。
- “电影” 指的是 `Movie` 节点，其名称存储在 `title` 属性中（英文标题，如 'Jurassic Park'）。
- “国家” 指的是 `Country` 节点，其名称存储在 `name` 属性中。
- “制作国家” 或 “在哪个国家制作” 对应关系 `[:PRODUCED_IN]`。
- “执导”、“导演了” 对应关系 `[:DIRECTED]`。
- “评分”、“打分” 对应关系 `[:REVIEWED]`，其属性 `score` 表示分数。
- “影评人” 对应 `Reviewer` 节点。
- 中文电影名如 “侏罗纪公园” 应映射为英文标题 'Jurassic Park'。
- “辛德勒的名单” → 'Schindler\\'s List'
- “拯救大兵瑞恩” → 'Saving Private Ryan'
"""


# 步骤 2：生成提示
def generate_prompt(question: str,mapping_str: str) -> str:
    return f"""
你是一个专业的 Cypher 查询生成器。请根据以下图数据库结构和术语映射规则，将自然语言问题转换为正确的 Cypher 查询。

规则：
- 仅使用下面列出的标签、关系和属性。
- 不要假设节点有未列出的属性。
- 使用 MATCH 遍历关系，而不是直接访问字段。
- 返回结果时使用别名提高可读性。
- 电影标题使用英文，如 'Jurassic Park'。
- 请只返回一行 Cypher 查询语句，不要添加任何解释、注释或 <think> 标签。
- 查询必须以 Cypher 关键字开头（如 MATCH），并以分号结尾。

图数据库结构：
{mapping_str}

问题：{question}
Cypher：
""".strip()


# 步骤 3：调用本地大语言模型
def call_llm(prompt: str) -> str:
    payload = {
        "model": LLM_MODEL,
        "prompt": prompt,
        "stream": False
    }
    response = requests.post(f"{OLLAMA_BASE_URL}/api/generate", json=payload)
    response.raise_for_status()
    return response.json()["response"].strip()


# 步骤 4：提取纯净 Cypher 查询
def extract_cypher_query(response_text: str) -> str:
    # 去除 <think>...</think> 和 ```cypher...```
    cleaned = re.sub(r'<think>.*?</think>', '', response_text, flags=re.DOTALL)
    cleaned = re.sub(r'```cypher\s*|\s*```', '', cleaned, flags=re.DOTALL)

    lines = [line.strip() for line in cleaned.split('\n') if line.strip()]
    cypher_keywords = {
        'MATCH', 'CREATE', 'MERGE', 'DELETE', 'DETACH', 'WITH', 'UNWIND',
        'CALL', 'RETURN', 'SET', 'REMOVE', 'FOREACH', 'LOAD', 'START'
    }

    # 从后往前找第一个有效 Cypher 行
    for i in range(len(lines) - 1, -1, -1):
        first_word = lines[i].split()[0].upper()
        if first_word in cypher_keywords:
            query = ' '.join(lines[i:])
            query = re.sub(r'^[Cc]ypher:\s*', '', query).strip()
            if not query.endswith(';'):
                query += ';'
            return query

    # 备选：找包含 MATCH 和关系模式的行
    for line in lines:
        if line.upper().startswith('MATCH') and ('->' in line or '-[' in line):
            if not line.endswith(';'):
                line += ';'
            return line

    # 最后兜底
    cleaned_final = cleaned.strip()
    if cleaned_final and not cleaned_final.endswith(';'):
        cleaned_final += ';'
    return cleaned_final if cleaned_final else "RETURN 'No valid query generated';"


# 步骤 5：执行 Cypher 查询=
def execute_query(cypher: str) -> list:
    with driver.session() as session:
        result = session.run(cypher)
        return [dict(record) for record in result]


# 步骤 6：创建测试数据
def create_test_data():
    with driver.session() as session:
        session.run("MERGE (:Country {name: 'United States'})")

        session.run("""
        MERGE (d:Director {name: 'Steven Spielberg'})
        WITH d
        UNWIND ['Jurassic Park', 'Schindler\\'s List', 'Saving Private Ryan'] AS title
        MERGE (m:Movie {title: title})
        MERGE (d)-[:DIRECTED]->(m)
        MERGE (m)-[:PRODUCED_IN]->(:Country {name: 'United States'})
        """)

        session.run("""
        MERGE (r1:Reviewer {name: 'Alice'})
        MERGE (r2:Reviewer {name: 'Bob'})
        """)

        session.run("""
        MATCH (m:Movie {title: 'Jurassic Park'})
        MATCH (r:Reviewer {name: 'Alice'})
        MERGE (r)-[:REVIEWED {score: 8.8}]->(m)
        """)

        session.run("""
        MATCH (m:Movie {title: 'Schindler\\'s List'})
        MATCH (r:Reviewer {name: 'Alice'})
        MERGE (r)-[:REVIEWED {score: 9.0}]->(m)
        """)


# 主程序
def main():
    # 清空数据库
    with driver.session() as session:
        session.run("MATCH (n) DETACH DELETE n")
    print("[OK] 数据库已清空")

    # 创建测试数据
    create_test_data()
    print("[OK] 测试数据已创建")

    # 用户提问
    question = "侏罗纪公园是在哪个国家制作的？"

    # 生成提示
    prompt = generate_prompt(question,TERM_MAPPING)
    print(f"\n[INFO] 发送给大语言模型的提示：\n{prompt}")

    # 调用 LLM
    print("\n[INFO] 正在调用大语言模型...")
    try:
        raw_output = call_llm(prompt)
        print(f"LLM 原始响应：\n{raw_output}")

        cypher_query = extract_cypher_query(raw_output)
        print(f"提取的 Cypher 查询：\n{cypher_query}")
    except Exception as e:
        print(f"调用 LLM 失败：{e}")
        return

    # 执行查询
    try:
        result = execute_query(cypher_query)
        print("\n[RESULT] 查询结果：")
        if result:
            for row in result:
                print(row)
        else:
            print("没有找到结果。")
    except Exception as e:
        print(f"执行 Cypher 失败：{e}")

    # 关闭连接
    driver.close()


if __name__ == "__main__":
    main()

[OK] 数据库已清空
[OK] 测试数据已创建

[INFO] 发送给大语言模型的提示：
你是一个专业的 Cypher 查询生成器。请根据以下图数据库结构和术语映射规则，将自然语言问题转换为正确的 Cypher 查询。

规则：
- 仅使用下面列出的标签、关系和属性。
- 不要假设节点有未列出的属性。
- 使用 MATCH 遍历关系，而不是直接访问字段。
- 返回结果时使用别名提高可读性。
- 电影标题使用英文，如 'Jurassic Park'。
- 请只返回一行 Cypher 查询语句，不要添加任何解释、注释或 <think> 标签。
- 查询必须以 Cypher 关键字开头（如 MATCH），并以分号结尾。

图数据库结构：

术语映射说明：
- “导演” 指的是 `Director` 节点，其姓名存储在 `name` 属性中。
- “电影” 指的是 `Movie` 节点，其名称存储在 `title` 属性中（英文标题，如 'Jurassic Park'）。
- “国家” 指的是 `Country` 节点，其名称存储在 `name` 属性中。
- “制作国家” 或 “在哪个国家制作” 对应关系 `[:PRODUCED_IN]`。
- “执导”、“导演了” 对应关系 `[:DIRECTED]`。
- “评分”、“打分” 对应关系 `[:REVIEWED]`，其属性 `score` 表示分数。
- “影评人” 对应 `Reviewer` 节点。
- 中文电影名如 “侏罗纪公园” 应映射为英文标题 'Jurassic Park'。
- “辛德勒的名单” → 'Schindler\'s List'
- “拯救大兵瑞恩” → 'Saving Private Ryan'


问题：侏罗纪公园是在哪个国家制作的？
Cypher：

[INFO] 正在调用大语言模型...
LLM 原始响应：
<think>
好的，我现在需要处理用户的问题：“侏罗纪公园是在哪个国家制作的？”。首先，根据用户提供的规则，我需要将中文电影名转换为对应的英文标题，这里“侏罗纪公园”应该映射为'Jurassic Park'。

接下来，我要确定图数据库中的相关结构。根据结构说明，电影是`Movie`节点，国家是`Country`节点，而制作国家对应的关系是`[:PRODUCED_IN]`。因此，我需要找到`Movie`节点与`Country`节点之间的`PRODUCED_IN`关系。

用户的问题是要找出电影制作的国家，所以Cypher查询应该从`Movie`节点出发，通过`PRODUCED_IN`关系连接到`Country`节点。需要确保使用正确的标签和关系名称，并且使用别名来提高可读性。

根据规则，不能直接访问字段，而是通过遍历关系来获取信息。因此，正确的MATCH语句应该是匹配`Movie`节点和`Country`节点之间的关系，然后返回国家的名称。

另外，注意查询必须以Cypher关键字开头，如MATCH，并以分号结尾。同时，返回结果时使用别名，比如用`country.name AS countryName`来别名。

现在，把这些组合起来，正确的Cypher查询应该是：

MATCH (m:Movie {title: 'Jurassic Park'})-[:PRODUCED_IN]->(c:Country) RETURN c.name AS countryName;

检查一下是否所有规则都满足：标签和关系正确，使用了别名，电影标题正确转换为英文，没有添加额外的解释，只返回查询语句。确认无误后，生成最终答案。
</think>

MATCH (m:Movie {title: 'Jurassic Park'})-[:PRODUCED_IN]->(c:Country) RETURN c.name AS countryName;
提取的 Cypher 查询：
MATCH (m:Movie {title: 'Jurassic Park'})-[:PRODUCED_IN]->(c:Country) RETURN c.name AS countryName;

[RESULT] 查询结果：
{'countryName': 'United States'}