QingLong/AI/Neo4j/InputShiTi.py

# -*- coding: utf-8 -*-
import re
import hashlib
from py2neo import Graph
from openai import OpenAI
from Config import *


class KnowledgeGraph:
    def __init__(self, content: str):
        self.content = content
        self.question_id = hashlib.md5(content.encode()).hexdigest()[:8]
        self.graph = Graph(NEO4J_URI, auth=NEO4J_AUTH)
        self.knowledge_points = self._get_knowledge_points()
        self.client = OpenAI(api_key=MODEL_API_KEY, base_url=MODEL_API_URL)

    def _get_knowledge_points(self) -> dict:
        """修复字段名称错误"""
        try:
            # 确保返回字段与节点属性名称匹配
            return {row['n.id'].lower(): row['n.name']  # 修改为n.id
                    for row in self.graph.run("MATCH (n:KnowledgePoint) RETURN n.id, n.name")}
        except Exception as e:
            print(f"获取知识点失败：", str(e))
            return {}

    def _make_prompt(self) -> str:
        """生成知识点识别专用提示词"""
        example_ids = list(self.knowledge_points.keys())[:5]
        example_names = [self.knowledge_points[k] for k in example_ids]

        return f"""你是一个数学专家，请分析题目考查的知识点，严格：
1. 只使用以下存在的知识点（格式：ID:名称）：
{", ".join([f"{k}:{v}" for k, v in zip(example_ids, example_names)])}...
共{len(self.knowledge_points)}个可用知识点

2. 按此格式生成Cypher：
MERGE (q:Question {{id: "{self.question_id}"}})
SET q.content = "题目内容"
WITH q
MATCH (kp:KnowledgePoint {{id: "知识点ID"}})
MERGE (q)-[:TESTS_KNOWLEDGE]->(kp)"""

    def _clean_cypher(self, code: str) -> str:
        """修复WITH语句顺序问题"""
        safe = []
        cypher_block = re.findall(r"```(?:cypher)?\n(.*?)```", code, re.DOTALL)
        if not cypher_block:
            return ""

        # 强制先创建Question节点
        has_question = False
        for line in cypher_block[0].split('\n'):
            line = line.split('//')[0].strip()
            if not line:
                continue

            # 确保Question节点最先创建
            if 'MERGE (q:Question' in line:
                has_question = True
                safe.insert(0, line)  # 确保这行在最前面
                continue

            # 安全过滤
            if 'CREATE' in line.upper():
                continue

            # 自动补全WITH语句（仅在Question创建之后）
            if has_question and 'MERGE (q)-[:TESTS_KNOWLEDGE]' in line and not any('WITH q' in l for l in safe):
                safe.append("WITH q")

            # ID存在性验证
            if 'MATCH (kp:KnowledgePoint' in line:
                kp_id = re.findall(r"id: ['\"](.*?)['\"]", line)
                if kp_id and kp_id[0] not in self.knowledge_points:
                    continue

            safe.append(line)

        # 补充必要WITH语句
        if has_question and not any(line.startswith('WITH q') for line in safe):
            safe.insert(1, "WITH q")  # 在创建Question之后立即添加

        return '\n'.join([line for line in safe if line])

    def run(self) -> str:
        """执行知识点关联流程"""
        try:
            response = self.client.chat.completions.create(
                model=MODEL_NAME,
                messages=[
                    {
                        "role": "system",
                        "content": self._make_prompt()
                    },
                    {
                        "role": "user",
                        "content": f"题目内容：{self.content}\n请分析考查的知识点，只返回Cypher代码"
                    }
                ]
            )

            raw_cypher = response.choices[0].message.content
            cleaned_cypher = self._clean_cypher(raw_cypher)

            if cleaned_cypher:
                print("验证通过的Cypher：\n", cleaned_cypher)
                return cleaned_cypher
            return ""

        except Exception as e:
            print("知识点分析失败：", str(e))
            return ""

def query_related_knowledge(self):
    """查询题目关联的知识点"""
    cypher = f"""
    MATCH (q:Question {{id: "{self.question_id}"}})-[:TESTS_KNOWLEDGE]->(kp)
    RETURN kp.id AS knowledge_id, kp.name AS knowledge_name
    """
    try:
        result = self.graph.run(cypher).data()
        if result:
            print(f"题目关联的知识点（{self.question_id}）：")
            for row in result:
                print(f"- {row['knowledge_name']} (ID: {row['knowledge_id']})")
        else:
            print("该题目尚未关联知识点")
        return result
    except Exception as e:
        print("查询失败：", str(e))
        return []

# 测试用例
if __name__ == '__main__':
    test_case = """【时间问题】甲乙两车从相距240公里的两地同时出发相向而行，甲车时速60公里，乙车时速40公里，几小时后相遇？"""

    test_case = """【时间问题】甲乙两车从相距240公里的两地同时出发相向而行..."""

    kg = KnowledgeGraph(test_case)
    cypher = kg.run()

    if cypher:
        kg.graph.run(cypher)
        print("执行成功！关联知识点：")
        kg.query_related_knowledge()  # 新增查询
    else:
        print("未生成有效Cypher")