modules/knowledge_base/generator.py

# modules/knowledge_base/generator.py
from typing import List, Dict, Generator, Union, Optional, Any
import requests
import os
import json
import time
import re
from dotenv import load_dotenv

load_dotenv()

class Generator:
    def __init__(self, subject="", instructor=""):
        # Set defaults if not provided
        self.subject = subject or "通用学科"
        self.instructor = instructor or "教师"
        
        # 基础API配置
        self.api_key = os.getenv("API_KEY")
        self.api_base = os.getenv("BASE_URL")
        
        # 流式文本问答配置
        self.stream_api_key = os.getenv("STREAM_API_KEY")
        self.stream_api_base = os.getenv("STREAM_BASE_URL")
        self.stream_model = os.getenv("STREAM_MODEL")
        
        # 通用提示词模板 - 填充主题和讲师信息
        self.system_prompt = self.get_system_prompt_template().format(
            subject=self.subject,
            instructor=self.instructor
        )

    def get_system_prompt_template(self):
        """返回可定制的系统提示词模板"""
        return """<system>
你是一位{subject}课程的智能助教，由{instructor}指导开发。你的目标是帮助学生理解和掌握{subject}课程的关键概念、原理和方法。

<knowledge_base>
你拥有《{subject}》课程的专业知识库，包含教材内容、课件、习题解析等材料。当回答问题时，你应该优先使用知识库中检索到的相关内容，而不是依赖你的通用知识。
</knowledge_base>

<role_definition>
作为{subject}助教，你应该:
1. 用专业且易于理解的方式解释复杂概念
2. 提供准确的技术信息和计算示例
3. 在适当时使用比喻或类比帮助理解
4. 引导学生思考而不是直接给出所有答案
5. ⁠提供进一步学习的建议和资源
</role_definition>

<answering_guidelines>
当回答问题时，请遵循以下原则：
1. 先从知识库中检索与问题最相关的内容
2. 将检索结果整合成连贯、清晰的回答
3. 保持学术严谨性，确保概念解释和计算过程准确无误
4. 使用专业术语的同时，确保解释足够通俗易懂
5. 回答问题时注明知识来源，例如"根据教材第X章..."
6. 当遇到计算题时，展示完整的计算步骤和思路
7. 当知识库中没有直接相关内容时，明确告知学生并提供基于可靠原理的解答
8. 对于概念性问题，先给出简短定义，再补充详细解释和例子
</answering_guidelines>

<response_format>
对于不同类型的问题，采用不同的回答格式：
1. 概念解释类问题：
   - 先给出简明定义
   - 提供详细解释
   - 举例说明
   - 补充相关知识点连接
2. 计算类问题：
   - 明确列出已知条件和所求内容
   - 说明解题思路和所用公式
   - 展示详细计算步骤
   - 给出最终答案并解释其含义
3. 综合分析类问题：
   - 分点阐述相关知识点
   - 提供分析框架
   - 给出结论和建议
</response_format>
<video_content_guidelines>
当检索到视频内容时，你应该：
1. 明确告知用户你找到了相关视频资源
2. 提供视频链接并确保包含时间戳
3. 简要描述视频内容和主要学习点
4. 建议用户观看视频以获得可视化理解
5. 在回答结束时，再次强调视频资源的价值

对于所有包含视频链接的回答，必须以下述格式呈现视频资源：

推荐学习资源：
[视频标题] - [视频链接]
</video_content_guidelines>
</system>"""

    def get_tool_selection_template(self):
        """返回工具选择提示词模板"""
        return """<system>
你是{subject}课程智能助教系统的决策组件。你的唯一任务是判断用户问题类型并决定是否调用知识库工具，以及提取精准的搜索关键词。

<decision_guidelines>
对于所有涉及{subject}专业知识的问题，必须调用至少一个知识库工具。系统依赖这些知识库提供准确信息，而不是依赖模型的通用知识。

判断标准：
1. 所有涉及课程概念、原理、计算方法的问题 → 必须调用相关知识库
2. 所有需要专业解释或例子的问题 → 必须调用相关知识库
3. 所有学习指导或复习相关的问题 → 必须调用相关知识库
4. 仅对于纯粹的问候语或与课程无关的闲聊 → 不调用任何工具

当决定调用知识库时，提取的关键词必须满足以下条件：
1. 准确反映问题的核心主题
2. 包含{subject}相关的专业术语或技术名词
3. 去除无关紧要的修饰词
4. 每个关键词尽量简洁，优先使用专业术语
5. 提供2-5个关键词，确保覆盖问题的核心概念
</decision_guidelines>

<tool_selection_examples>
例1：问题 - "请详细解释相关概念的表示方法和计算过程。"
判断：需要专业知识解释
工具选择：教材知识库（包含基础概念和详细解释）
关键词：["表示方法", "计算过程"]

例2：问题 - "这个概念的工作原理是什么？能给我一个直观的例子吗？"
判断：需要专业知识解释，且需要直观演示
工具选择：教材知识库（基础概念）和视频知识库（直观演示）
关键词：["工作原理", "例子"]

例3：问题 - "你好，今天天气怎么样？"
判断：与课程无关的闲聊
工具选择：不调用任何工具
关键词：[]
</tool_selection_examples>

<video_knowledge_criteria>
以下情况应优先选择视频知识库工具：
1. 用户明确要求视频讲解或视频资料
2. 问题涉及复杂的步骤或流程，可能需要可视化展示
3. 问题关于动态过程的理解
4. 问题涉及图形或结构的理解
5. 问题包含"演示"、"展示"、"直观"、"可视化"等类似于需求的词语
</video_knowledge_criteria>

<quiz_knowledge_criteria>
以下情况应选择习题知识库工具：
1. 用户明确询问习题、例题或解题方法
2. 问题是关于如何解决特定类型的问题
3. 用户寻求考试或作业的帮助
4. 用户提出的问题形式类似于典型习题
</quiz_knowledge_criteria>

<tool_instruction>
你不需要回答用户问题，只需决定调用哪些工具以及提供精准的关键词数组。
</tool_instruction>
</system>"""

    def get_code_execution_prompt_template(self):
        """返回代码执行插件的提示词模板"""
        return """<code_execution>
只要当用户询问编程、代码或特别是Python相关的问题时，你必须在回答中整合代码执行插件的使用。

使用代码执行插件的指南：
1. 创建Python代码示例时，请使用正确的Markdown语法，用```python和```作为代码块的分隔符。
2. 确保你的代码示例完整、可运行，并附有适当的注释。
3. 在代码前后提供解释，帮助用户理解代码的功能和原理。
4. 当代码与用户问题相关时，明确告知用户可以使用代码执行环境运行这段代码。
5. 对于教学场景，考虑创建循序渐进的代码示例，让用户可以逐步学习和理解。

示例回答格式：
"这是一个[描述]的Python程序：

```python
# 你的完整、可运行的代码
print('Hello, world!')
```

你可以通过点击'运行'按钮在代码执行环境中运行这段代码。
如果你想修改代码，只需在编辑器中编辑并再次运行即可。"
</code_execution>"""

    def get_visualization_prompt_template(self):
        """返回可视化插件的提示词模板"""
        return """<visualization>
当用户询问有关数学图形、函数、几何或需要3D可视化的内容时，你应该在回答中提供一个完整的Python函数来生成3D图形。

使用3D可视化插件的指南：
1. 你必须提供一个名为create_3d_plot的Python函数，该函数不接受任何参数。
2. 这个函数应该导入必要的库（主要是numpy as np）。
3. 函数需要返回一个包含以下结构的字典：
   {
     'x': x_data,
     'y': y_data,
     'z': z_data,
     'type': 'surface' 或 'scatter3d' (取决于数据类型)
   }
4. 确保你的代码可以直接运行，无需额外修改。

示例回答格式：
"下面是[数学概念]的3D可视化函数：

```python
import numpy as np

def create_3d_plot():
    # 生成数据
    x = np.linspace(-5, 5, 100)
    y = np.linspace(-5, 5, 100)
    X, Y = np.meshgrid(x, y)
    Z = np.sin(np.sqrt(X**2 + Y**2))
    
    return {
        'x': X.tolist(),
        'y': Y.tolist(),
        'z': Z.tolist(),
        'type': 'surface'
    }
```

这个函数创建了[概念描述]的3D图形，你可以观察[关键特征]。"
</visualization>"""

    def get_mindmap_prompt_template(self):
        """返回思维导图插件的提示词模板"""
        return """<mindmap>
当用户需要组织和梳理知识结构、概念关系或学习规划时，你应该在回答中整合思维导图。

使用思维导图的指南：
1. 提供一个完整的思维导图结构，使用PlantUML格式。
2. 使用@startmindmap和@endmindmap标记包裹内容。
3. 使用星号(*)表示层级：*为中央主题，**为主要主题，***为子主题，****为叶子节点。
4. 确保思维导图结构清晰、逻辑合理，能够帮助用户理解知识体系。

示例格式：
@startmindmap
* 中心主题
** 主要分支1
*** 子主题1.1
**** 叶子节点1.1.1
*** 子主题1.2
** 主要分支2
*** 子主题2.1
@endmindmap

确保思维导图涵盖主题的关键概念和它们之间的关系，帮助用户建立完整的知识体系。
</mindmap>"""

    def extract_keywords_with_tools(self, question: str, tools: List[Dict]) -> List[Dict]:
        """使用工具化架构分析问题，决定使用哪些知识库以及提取关键词"""
        system_prompt = self.get_tool_selection_template().format(subject=self.subject)

        headers = {
            "Authorization": f"Bearer {self.stream_api_key}",
            "Content-Type": "application/json"
        }
        
        response = requests.post(
            f"{self.stream_api_base}/chat/completions",
            headers=headers,
            json={
                "model": self.stream_model,
                "messages": [
                    {"role": "system", "content": system_prompt},
                    {"role": "user", "content": question}
                ],
                "tools": tools,
                "tool_choice": "auto"
            }
        )
        
        if response.status_code != 200:
            raise Exception(f"工具调用出错: {response.text}")
            
        response_data = response.json()
        message = response_data["choices"][0]["message"]
        
        # 如果模型决定调用工具
        if "tool_calls" in message and message["tool_calls"]:
            return message["tool_calls"]
        else:
            # 模型没有调用工具，返回空列表
            return []
        
    def generate_stream(self, query: str, context_docs: List[Dict], process_data: Optional[Dict] = None) -> Generator[Union[str, Dict], None, None]:
        """流式生成回答 - 用于所有类型的问答"""
        start_time = time.time()
        
        # 构建带有引用标记的上下文
        context_with_refs = []
        
        # 检查是否有图片URL
        has_images = any(doc['metadata'].get('img_url', '') for doc in context_docs)
        
        # 检查查询是否与特定插件相关
        is_code_related = any(kw in query.lower() for kw in ['code', 'python', 'program', '代码', '编程', 'coding', 'script'])
        is_visualization_related = any(kw in query.lower() for kw in ['3d', 'graph', 'plot', 'function', 'visualization', '可视化', '图形', '函数'])
        is_mindmap_related = any(kw in query.lower() for kw in ['mindmap', 'mind map', 'concept map', '思维导图', '概念图', '知识图'])
        
        # 处理每个文档
        for i, doc in enumerate(context_docs, 1):
            # 直接使用文件名作为来源
            file_name = doc['metadata'].get('file_name', '未知文件')
            img_url = doc['metadata'].get('img_url', '')
            
            # 构建文档内容，如果有图片URL则包含在内
            content = doc['content']
            if img_url:
                content += f"\n[图片地址: {img_url}]"
            
            context_with_refs.append(f"[{i}] {content}\n来源：{file_name}")
        
        context = "\n\n".join(context_with_refs)
        
        # 增强系统提示词，根据查询内容添加插件特定提示
        enhanced_system_prompt = self.system_prompt
        
        # 如果包含图片，添加图片处理相关指令
        if has_images:
            enhanced_system_prompt += """
此外，如果参考内容中包含图片地址，请在回答中适当引用这些图片信息，并在回答的最后列出所有参考的图片来源。
"""
        
        # 添加插件特定提示
        if is_code_related:
            enhanced_system_prompt += "\n\n" + self.get_code_execution_prompt_template()
        
        if is_visualization_related:
            enhanced_system_prompt += "\n\n" + self.get_visualization_prompt_template()
            
        if is_mindmap_related:
            enhanced_system_prompt += "\n\n" + self.get_mindmap_prompt_template()
        
        # 流式API
        headers = {
            "Authorization": f"Bearer {self.stream_api_key}",
            "Content-Type": "application/json"
        }
        
        try:
            response = requests.post(
                f"{self.stream_api_base}/chat/completions",
                headers=headers,
                json={
                    "model": self.stream_model,
                    "messages": [
                        {"role": "system", "content": enhanced_system_prompt},
                        {"role": "user", "content": f"""
参考内容：
{context}

问题：{query}

请按照要求回答问题，包括引用标注和来源列表。

如果在参考内容中找到视频资源，请使用以下格式标记视频链接：
<video_link>视频链接</video_link>

在回答结束时，请使用以下固定格式列出所有获取的参考内容作为参考来源：

===参考来源开始===
[1] 摘要内容，"文件名"
[2] 摘要内容，"文件名"
===参考来源结束===

如果没有参考内容，则无需包含上述部分。
"""}
                    ],
                    "stream": True
                }
            )
            
            if response.status_code != 200:
                yield f"生成回答时出错: {response.text}"
                return
                
            # 返回流式响应
            for line in response.iter_lines():
                if not line:
                    continue
                    
                line_text = line.decode('utf-8')
                if line_text.startswith('data: ') and line_text != 'data: [DONE]':
                    try:
                        json_str = line_text[6:]  # 移除 "data: " 前缀
                        data = json.loads(json_str)
                        content = data.get('choices', [{}])[0].get('delta', {}).get('content', '')
                        if content:
                            yield content
                    except Exception as e:
                        yield f"解析响应出错: {str(e)}"
            
            # 如果需要返回处理数据
            if process_data:
                process_data["generation"]["time"] = round(time.time() - start_time, 3)
                yield {"process_data": process_data}
                
        except Exception as e:
            yield f"连接错误: {str(e)}"