QingLong/AI/A1_GenerateMarkdown.py

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

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

from A3_MarkdownToDocx import *
from Config import *


class MarkdownGenerator:
    """Markdown教学大纲生成器"""

    # 固定配置项
    MODEL_R1 = "deepseek-r1"
    MODEL_V3 = "deepseek-v3"
    API_KEY = "sk-01d13a39e09844038322108ecdbd1bbc"

    def __init__(
            self,
            course_name: str,
            output_path: Optional[Path] = None,
            template_path: Optional[Path] = None
    ):
        """
        初始化生成器

        :param course_name: 课程名称（如：小学数学内角和）
        :param output_path: 输出文件路径（可选）
        :param template_path: 模板文件路径（可选）
        """
        self.course_name = course_name
        self.template_path = template_path or DEFAULT_TEMPLATE
        self.output_path = output_path or DEFAULT_OUTPUT_DIR / f"{course_name}.md"

        self._validate_paths()

    def _validate_paths(self):
        """路径验证"""
        if not self.template_path.exists():
            raise FileNotFoundError(f"模板文件不存在: {self.template_path}")

        self.output_path.parent.mkdir(parents=True, exist_ok=True)
        if not self.output_path.parent.is_dir():
            raise NotADirectoryError(f"无效的输出目录: {self.output_path.parent}")

    def _load_template(self) -> str:
        """加载模板内容"""
        return self.template_path.read_text(encoding='utf-8')

    def _generate_stream(self) -> Iterator[DashScopeAPIResponse]:
        """流式生成内容"""
        system_prompt = (
            "请严格按照以下Markdown格式生成教学大纲：\n"
            f"{self._load_template()}\n"
            "注意：\n"
            "1. 使用#/##标记标题\n"
            "2. 每页内容以>开头\n"
            "3. 列表项使用-标记\n"
            "4. 内容中不要包含图片\n"
        )

        return Generation.call(
            model=self.MODEL_R1,
            api_key=self.API_KEY,
            messages=[
                {"role": "system", "content": system_prompt},
                {"role": "user", "content": f"请生成《{self.course_name}》教学大纲"}
            ],
            result_format='message',
            stream=True,
            incremental_output=True
        )

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

        try:
            print(f"🚀 开始生成【{self.course_name}】教学大纲")
            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:
                self.output_path.write_text(''.join(content_buffer), encoding='utf-8')
                print(f"\n\n✅ 生成成功！耗时 {int(time.time() - start_time)}秒")
                print(f"📂 文件已保存至：{self.output_path}")
                return True
            return False

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


def generate_document(
        course_name: str,
        output_path: Optional[str] = None,
        template_path: Optional[str] = None
):
    """
    生成文档入口函数

    :param course_name: 课程名称
    :param output_path: 输出路径（可选）
    :param template_path: 模板路径（可选）
    """
    try:
        generator = MarkdownGenerator(
            course_name=course_name,
            output_path=Path(output_path) if output_path else None,
            template_path=Path(template_path) if template_path else None
        )
        return generator.run()
    except Exception as e:
        print(f"❌ 初始化失败：{str(e)}")
        return False