app.py

from flask import Flask, request, jsonify, send_from_directory, render_template, redirect, url_for
from flask_cors import CORS
import os
import time
import traceback
import json
import re
import sys
import io
import threading
import queue
import contextlib
import signal
import psutil
from dotenv import load_dotenv

# 导入模块路由
from modules.knowledge_base.routes import knowledge_bp
from modules.code_executor.routes import code_executor_bp
from modules.visualization.routes import visualization_bp
from modules.agent_builder.routes import agent_builder_bp

# 加载环境变量
load_dotenv()

app = Flask(__name__)
CORS(app)

# 注册蓝图
app.register_blueprint(knowledge_bp, url_prefix='/api/knowledge')
app.register_blueprint(code_executor_bp, url_prefix='/api/code')
app.register_blueprint(visualization_bp, url_prefix='/api/visualization')
app.register_blueprint(agent_builder_bp, url_prefix='/api/agent')

# 确保目录存在
os.makedirs('static', exist_ok=True)
os.makedirs('uploads', exist_ok=True)
os.makedirs('agents', exist_ok=True)

# 用于代码执行的上下文
execution_contexts = {}

def get_memory_usage():
    """获取当前进程的内存使用情况"""
    process = psutil.Process(os.getpid())
    return f"{process.memory_info().rss / 1024 / 1024:.1f} MB"

class CustomStdin:
    def __init__(self, input_queue):
        self.input_queue = input_queue
        self.buffer = ""
    
    def readline(self):
        if not self.buffer:
            self.buffer = self.input_queue.get() + "\n"
        
        result = self.buffer
        self.buffer = ""
        return result

class InteractiveExecution:
    """管理Python代码的交互式执行"""
    def __init__(self, code):
        self.code = code
        self.context_id = str(time.time())
        self.is_complete = False
        self.is_waiting_for_input = False
        self.stdout_buffer = io.StringIO()
        self.last_read_position = 0
        self.input_queue = queue.Queue()
        self.error = None
        self.thread = None
        self.should_terminate = False
    
    def run(self):
        """在单独的线程中启动执行"""
        self.thread = threading.Thread(target=self._execute)
        self.thread.daemon = True
        self.thread.start()
        
        # 给执行一点时间开始
        time.sleep(0.1)
        return self.context_id
    
    def _execute(self):
        """执行代码，处理标准输入输出"""
        try:
            # 保存原始的stdin/stdout
            orig_stdin = sys.stdin
            orig_stdout = sys.stdout
            
            # 创建自定义stdin
            custom_stdin = CustomStdin(self.input_queue)
            
            # 重定向stdin和stdout
            sys.stdin = custom_stdin
            sys.stdout = self.stdout_buffer
            
            try:
                # 检查终止的函数
                self._last_check_time = 0
                
                def check_termination():
                    if self.should_terminate:
                        raise KeyboardInterrupt("Execution terminated by user")
                
                # 设置一个模拟__main__模块的命名空间
                shared_namespace = {
                    "__builtins__": __builtins__,
                    "_check_termination": check_termination,
                    "time": time, 
                    "__name__": "__main__"
                }
                
                # 在这个命名空间中执行用户代码
                try:
                    exec(self.code, shared_namespace)
                except KeyboardInterrupt:
                    print("\nExecution terminated by user")
                
            except Exception as e:
                self.error = {
                    "error": str(e),
                    "traceback": traceback.format_exc()
                }
            
            finally:
                # 恢复原始stdin/stdout
                sys.stdin = orig_stdin
                sys.stdout = orig_stdout
                
                # 标记执行完成
                self.is_complete = True
                
        except Exception as e:
            self.error = {
                "error": str(e),
                "traceback": traceback.format_exc()
            }
            self.is_complete = True
    
    def terminate(self):
        """终止执行"""
        self.should_terminate = True
        
        # 如果在等待输入，放入一些内容以解除阻塞
        if self.is_waiting_for_input:
            self.input_queue.put("\n")
        
        # 给执行一点时间终止
        time.sleep(0.2)
        
        # 标记为完成
        self.is_complete = True
        
        return True
    
    def provide_input(self, user_input):
        """为运行的代码提供输入"""
        self.input_queue.put(user_input)
        self.is_waiting_for_input = False
        return True
    
    def get_output(self):
        """获取stdout缓冲区的当前内容"""
        output = self.stdout_buffer.getvalue()
        return output
        
    def get_new_output(self):
        """只获取自上次读取以来的新输出"""
        current_value = self.stdout_buffer.getvalue()
        if self.last_read_position < len(current_value):
            new_output = current_value[self.last_read_position:]
            self.last_read_position = len(current_value)
            return new_output
        return ""

@app.route('/')
def index():
    """主界面"""
    return render_template('index.html')
@app.route('/code_execution.html')
def index2():
    """主界面"""
    return render_template('code_execution.html')
@app.route('/api/progress/<task_id>', methods=['GET'])
def get_progress(task_id):
    """获取文档处理进度"""
    try:
        # 从知识库模块访问处理任务
        from modules.knowledge_base.routes import processing_tasks
        
        progress_data = processing_tasks.get(task_id, {
            'progress': 0,
            'status': '未找到任务',
            'error': True
        })
        
        return jsonify({"success": True, "data": progress_data})
    except Exception as e:
        traceback.print_exc()
        return jsonify({"success": False, "message": str(e)}), 500

@app.route('/student/<agent_id>')
def student_view(agent_id):
    """学生访问Agent界面"""
    token = request.args.get('token', '')
    
    # 验证Agent存在
    agent_path = os.path.join('agents', f"{agent_id}.json")
    if not os.path.exists(agent_path):
        return render_template('error.html', 
                              message="找不到指定的Agent",
                              error_code=404)
    
    # 加载Agent配置
    with open(agent_path, 'r', encoding='utf-8') as f:
        try:
            agent_config = json.load(f)
        except:
            return render_template('error.html', 
                                  message="Agent配置无效",
                                  error_code=500)
    
    # 验证访问令牌
    if token:
        valid_token = False
        if "distributions" in agent_config:
            for dist in agent_config["distributions"]:
                if dist.get("token") == token:
                    valid_token = True
                    break
        
        if not valid_token:
            return render_template('error.html', 
                                  message="访问令牌无效",
                                  error_code=403)
    
    # 渲染学生页面
    return render_template('student.html', 
                          agent_id=agent_id,
                          agent_name=agent_config.get('name', 'AI学习助手'),
                          agent_description=agent_config.get('description', ''),
                          token=token)

@app.route('/code_execution.html')
def code_execution_page():
    """代码执行页面"""
    return send_from_directory(os.path.dirname(os.path.abspath(__file__)), 'code_execution.html')

@app.route('/api/student/chat/<agent_id>', methods=['POST'])
def student_chat(agent_id):
    """学生与Agent聊天的API"""
    try:
        data = request.json
        message = data.get('message', '')
        token = data.get('token', '')
        
        if not message:
            return jsonify({"success": False, "message": "消息不能为空"}), 400
        
        # 验证Agent和令牌
        agent_path = os.path.join('agents', f"{agent_id}.json")
        if not os.path.exists(agent_path):
            return jsonify({"success": False, "message": "Agent不存在"}), 404
        
        with open(agent_path, 'r', encoding='utf-8') as f:
            agent_config = json.load(f)
        
        # 验证令牌（如果提供）
        if token and "distributions" in agent_config:
            valid_token = False
            for dist in agent_config["distributions"]:
                if dist.get("token") == token:
                    valid_token = True
                    
                    # 更新使用计数
                    dist["usage_count"] = dist.get("usage_count", 0) + 1
                    break
            
            if not valid_token:
                return jsonify({"success": False, "message": "访问令牌无效"}), 403
                
        # 更新Agent使用统计
        if "stats" not in agent_config:
            agent_config["stats"] = {}
        
        agent_config["stats"]["usage_count"] = agent_config["stats"].get("usage_count", 0) + 1
        agent_config["stats"]["last_used"] = int(time.time())
        
        # 保存更新后的Agent配置
        with open(agent_path, 'w', encoding='utf-8') as f:
            json.dump(agent_config, f, ensure_ascii=False, indent=2)
        
        # 获取Agent关联的知识库和插件
        knowledge_bases = agent_config.get('knowledge_bases', [])
        plugins = agent_config.get('plugins', [])
        
        # 获取学科和指导者信息
        subject = agent_config.get('subject', agent_config.get('name', '通用学科'))
        instructor = agent_config.get('instructor', '教师')
        
        # 创建Generator实例，传入学科和指导者信息
        from modules.knowledge_base.generator import Generator
        generator = Generator(subject=subject, instructor=instructor)
        
        # 检测需要使用的插件
        suggested_plugins = []
            
        # 检测是否需要代码执行插件
        if 'code' in plugins and ('代码' in message or 'python' in message.lower() or '编程' in message or 'code' in message.lower() or 'program' in message.lower()):
            suggested_plugins.append('code')
                
        # 检测是否需要3D可视化插件
        if 'visualization' in plugins and ('3d' in message.lower() or '可视化' in message or '图形' in message):
            suggested_plugins.append('visualization')
                
        # 检测是否需要思维导图插件
        if 'mindmap' in plugins and ('思维导图' in message or 'mindmap' in message.lower()):
            suggested_plugins.append('mindmap')
        
        # 检查是否有配置知识库
        if not knowledge_bases:
            # 没有知识库，直接使用模型进行回答
            print(f"\n=== 处理查询: {message} (无知识库) ===")
            
            # 使用空的文档列表调用生成器进行回答
            final_response = ""
            for chunk in generator.generate_stream(message, []):
                if isinstance(chunk, dict):
                    continue  # 跳过处理数据
                final_response += chunk
            
            # 返回生成的回答
            return jsonify({
                "success": True,
                "message": final_response,
                "tools": suggested_plugins
            })
        
        # 有知识库配置，执行知识库查询流程
        try:
            # 导入RAG系统组件
            from modules.knowledge_base.retriever import Retriever
            from modules.knowledge_base.reranker import Reranker
            
            retriever = Retriever()
            reranker = Reranker()
            
            # 构建工具定义 - 将所有知识库作为工具
            tools = []
            
            # 创建工具名称到索引的映射
            tool_to_index = {}
            
            for i, index in enumerate(knowledge_bases):
                display_name = index[4:] if index.startswith('rag_') else index
                
                # 判断是否是视频知识库
                is_video = "视频" in display_name or "video" in display_name.lower()
                
                # 根据内容类型生成适当的工具名称
                if is_video:
                    tool_name = f"video_knowledge_base_{i+1}"
                    description = f"在'{display_name}'视频知识库中搜索，返回带时间戳的视频链接。适用于需要视频讲解的问题。"
                else:
                    tool_name = f"knowledge_base_{i+1}"
                    description = f"在'{display_name}'知识库中搜索专业知识、概念和原理。适用于需要文本说明的问题。"
                
                # 添加工具名到索引的映射
                tool_to_index[tool_name] = index
                
                tools.append({
                    "type": "function",
                    "function": {
                        "name": tool_name,
                        "description": description,
                        "parameters": {
                            "type": "object",
                            "properties": {
                                "keywords": {
                                    "type": "array",
                                    "items": {"type": "string"},
                                    "description": "搜索的关键词列表"
                                }
                            },
                            "required": ["keywords"],
                            "additionalProperties": False
                        },
                        "strict": True
                    }
                })
            
            # 第一阶段：工具选择决策
            print(f"\n=== 处理查询: {message} ===")
            tool_calls = generator.extract_keywords_with_tools(message, tools)
            
            # 如果不需要调用工具，直接回答
            if not tool_calls:
                print("未检测到需要使用知识库，直接回答")
                final_response = ""
                for chunk in generator.generate_stream(message, []):
                    if isinstance(chunk, dict):
                        continue  # 跳过处理数据
                    final_response += chunk
                
                return jsonify({
                    "success": True,
                    "message": final_response,
                    "tools": suggested_plugins
                })
            
            # 收集来自工具执行的所有文档
            all_docs = []
            
            # 执行每个工具调用
            for tool_call in tool_calls:
                try:
                    tool_name = tool_call["function"]["name"]
                    actual_index = tool_to_index.get(tool_name)
                    
                    if not actual_index:
                        print(f"找不到工具名称 '{tool_name}' 对应的索引")
                        continue
                    
                    print(f"\n执行工具 '{tool_name}' -> 使用索引 '{actual_index}'")
                    
                    arguments = json.loads(tool_call["function"]["arguments"])
                    keywords = " ".join(arguments.get("keywords", []))
                    
                    if not keywords:
                        print("没有提供关键词，跳过检索")
                        continue
                    
                    print(f"检索关键词: {keywords}")
                    
                    # 执行检索
                    retrieved_docs, _ = retriever.retrieve(keywords, specific_index=actual_index)
                    print(f"检索到 {len(retrieved_docs)} 个文档")
                    
                    # 重排序文档
                    reranked_docs = reranker.rerank(message, retrieved_docs, actual_index)
                    print(f"重排序完成，排序后有 {len(reranked_docs)} 个文档")
                    
                    # 添加结果
                    all_docs.extend(reranked_docs)
                    
                except Exception as e:
                    print(f"执行工具 '{tool_call.get('function', {}).get('name', '未知')}' 调用时出错: {str(e)}")
                    import traceback
                    traceback.print_exc()
            
            # 如果没有检索到任何文档，直接回答
            if not all_docs:
                print("未检索到任何相关文档，直接回答")
                final_response = ""
                for chunk in generator.generate_stream(message, []):
                    if isinstance(chunk, dict):
                        continue  # 跳过处理数据
                    final_response += chunk
                
                return jsonify({
                    "success": True,
                    "message": final_response,
                    "tools": suggested_plugins
                })
            
            # 按相关性排序
            all_docs.sort(key=lambda x: x.get('rerank_score', 0), reverse=True)
            print(f"\n最终收集到 {len(all_docs)} 个文档用于生成回答")
            
            # 提取参考信息
            references = []
            for i, doc in enumerate(all_docs[:3], 1):  # 只展示前3个参考来源
                file_name = doc['metadata'].get('file_name', '未知文件')
                content = doc['content']
                
                # 提取大约前100字符作为摘要
                summary = content[:100] + ('...' if len(content) > 100 else '')
                
                references.append({
                    'index': i,
                    'file_name': file_name,
                    'content': content,
                    'summary': summary
                })
            
            # 第二阶段：生成最终答案
            final_response = ""
            for chunk in generator.generate_stream(message, all_docs):
                if isinstance(chunk, dict):
                    continue  # 跳过处理数据
                final_response += chunk
            
            # 构建回复
            return jsonify({
                "success": True,
                "message": final_response,
                "tools": suggested_plugins,
                "references": references
            })
            
        except Exception as e:
            import traceback
            traceback.print_exc()
            return jsonify({
                "success": False, 
                "message": f"处理查询时出错: {str(e)}"
            }), 500
            
    except Exception as e:
        import traceback
        traceback.print_exc()
        return jsonify({"success": False, "message": str(e)}), 500

if __name__ == '__main__':
    app.run(debug=True, host='0.0.0.0', port=5000)