QingLong/AI/Neo4j/KnowledgeGraph.py

# -*- coding: utf-8 -*-
import time
from typing import Iterator, Tuple
from openai import OpenAI
from openai.types.chat import ChatCompletionChunk
from Config import API_KEY, MODEL_R1, MODEL_URL


class KnowledgeGraph:
    def __init__(self, content: str):
        self.content = content
        self.client = OpenAI(api_key=API_KEY, base_url=MODEL_URL)

    def _generate_stream(self) -> Iterator[ChatCompletionChunk]:
        system_prompt = '''回答以下内容:
1. 这道题目有哪些知识点，哪些能力点
2. 生成Neo4j 5.26.2的插入语句'''
        return self.client.chat.completions.create(
            model=MODEL_R1,
            messages=[
                {"role": "system", "content": system_prompt},
                {"role": "user", "content": self.content}
            ],
            stream=True,
            timeout=300
        )

    def run(self) -> Tuple[bool, str]:
        """执行生成流程（返回状态和完整内容）"""
        start_time = time.time()
        spinner = ['⠋', '⠙', '⠹', '⠸', '⠼', '⠴', '⠦', '⠧', '⠇', '⠏']
        content_buffer = []

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

            for idx, chunk in enumerate(stream):
                print(f"\r{spinner[idx % 10]} 生成中（{int(time.time() - start_time)}秒）", end="")

                if chunk.choices and chunk.choices[0].delta.content:
                    content_chunk = chunk.choices[0].delta.content
                    content_buffer.append(content_chunk)

                    if len(content_buffer) == 1:
                        print("\n\n📝 内容生成开始：")
                    print(content_chunk, end="", flush=True)

            if content_buffer:
                full_content = ''.join(content_buffer)
                print(f"\n\n✅ 生成成功！耗时 {int(time.time() - start_time)}秒")
                print("\n================ 完整结果 ================")
                print(full_content)
                print("========================================")
                return True, full_content
            return False, ""

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


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

    # 如果需要进一步处理结果
    if success:
        with open("result.txt", "w", encoding="utf-8") as f:
            f.write(result)