QingLong/AI/Neo4j/KnowledgeGraph.py

# -*- coding: utf-8 -*-
import time
from typing import Iterator, Optional

from dashscope import Generation
from dashscope.api_entities.dashscope_response import DashScopeAPIResponse

from Config import *


class KnowledgeGraph:
    def __init__(
            self,
            shiti_content: str,
    ):
        """
        初始化生成器
        """
        self.shiti_content = shiti_content

    def _generate_stream(self) -> Iterator[DashScopeAPIResponse]:
        """流式生成内容"""
        system_prompt = (
            '''
            回答以下内容:
            1、这道题目有哪些知识点，哪些能力点
            2、我准备把你返回的知识点和能力点，存入到neo4j中去，版本是neo4j-community-5.26.2,生成插入到数据库中的语句
            '''
        )

        return Generation.call(
            model=MODEL_R1,
            api_key=API_KEY,
            messages=[
                {"role": "system", "content": system_prompt},
                {"role": "user", "content": shiti_content}
            ],
            result_format='message',
            stream=True,
            incremental_output=True
        )

    def run(self) -> bool:
        """执行生成流程"""
        start_time = time.time()
        spinner = ['⠋', '⠙', '⠹', '⠸', '⠼', '⠴', '⠦', '⠧', '⠇', '⠏']
        content_buffer = []

        try:
            print(f"🚀 开始生成知识点和能力点的总结和插入语句")
            responses = self._generate_stream()

            for idx, response in enumerate(responses):
                # 显示进度
                print(f"\r{spinner[idx % 10]} 生成中（{int(time.time() - start_time)}秒）", end="")

                if response.status_code == 200 and response.output:
                    if chunk := response.output.choices[0].message.content:
                        content_buffer.append(chunk)
                        if len(content_buffer) == 1:
                            print("\n\n📝 内容生成开始：")
                        print(chunk, end="", flush=True)

            # 保存结果
            if content_buffer:
                print(f"\n\n✅ 生成成功！耗时 {int(time.time() - start_time)}秒")
                return True
            return False

        except Exception as e:
            print(f"\n\n❌ 生成失败：{str(e)}")
            return False


if __name__ == '__main__':
    shiti_content='''
    下面是一道小学三年级的数学题目，巧求周长：
    把7个完全相同的小长方形拼成如图的样子，已知每个小长方形的长是10厘米，则拼成的大长方形的周长是多少厘米？    
    '''
    KnowledgeGraph = KnowledgeGraph(shiti_content)
    KnowledgeGraph.run()