app.py

from flask import Flask, request, jsonify, render_template, Response
import os
from dotenv import load_dotenv
import json
import requests
from PIL import Image
import base64
from io import BytesIO
import logging
import time

# 加载环境变量
load_dotenv()

app = Flask(__name__)
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

def create_headers():
    """创建API请求头"""
    return {
        "Content-Type": "application/json",
        "Authorization": f"Bearer {os.getenv('OPENAI_API_KEY')}"
    }

def image_to_base64(image):
    """将PIL图像转换为base64"""
    buffered = BytesIO()
    image.save(buffered, format="PNG")
    return base64.b64encode(buffered.getvalue()).decode()

def process_image_with_vision(image_base64):
    """使用Vision API处理图像"""
    url = f"{os.getenv('OPENAI_API_BASE')}/v1/chat/completions"
    
    messages = [
        {
            "role": "user",
            "content": [
                {
                    "type": "text",
                    "text": """请仔细分析这张图片中的题目，要求：

1. 识别内容类型：
   - 判断是否为数学题、物理题、化学题等
   - 识别是选择题、填空题还是解答题
   - 判断是否包含图表或特殊符号

2. 内容提取：
   - 提取所有文字内容，包括题干、选项（如果有）
   - 识别所有数学公式、化学方程式或特殊符号
   - 保留原有的排版格式（如换行、缩进等）

3. 公式处理：
   - 将所有数学公式转换为LaTeX格式
   - 使用[formula_n]作为占位符，其中n为公式编号
   - 保持公式的完整性和准确性

请使用以下JSON格式返回：
{
    "type": "题目类型（如：数学/物理/化学）",
    "format": "题目格式（如：选择题/填空题/解答题）",
    "text": "包含[formula_1], [formula_2]等占位符的完整文本",
    "formulas": ["latex公式1", "latex公式2"],
    "options": ["A. xxx", "B. xxx"] // 如果是选择题则包含此字段
    "notes": ["可能存在的问题说明1", "问题说明2"]
}"""
                },
                {
                    "type": "image_url",
                    "image_url": {
                        "url": f"data:image/png;base64,{image_base64}"
                    }
                }
            ]
        }
    ]
    
    payload = {
        "model": os.getenv('OPENAI_VISION_MODEL'),
        "messages": messages,
        "max_tokens": 1500,
        "temperature": 0.2
    }
    
    try:
        response = requests.post(url, headers=create_headers(), json=payload)
        response.raise_for_status()
        result = response.json()
        
        content = result['choices'][0]['message']['content']
        try:
            return json.loads(content)
        except json.JSONDecodeError:
            return {
                "type": "unknown",
                "format": "unknown",
                "text": content,
                "formulas": [],
                "notes": ["无法解析为JSON格式"]
            }
            
    except Exception as e:
        logger.error(f"Vision API调用错误: {str(e)}")
        return {"error": str(e)}

def handle_sse_response(raw_data):
    """处理SSE响应数据"""
    if raw_data:
        try:
            data = json.loads(raw_data)
            if len(data['choices']) > 0:
                delta = data['choices'][0].get('delta', {})
                content = delta.get('content', '')
                print(f"Stream content: {content}", flush=True)  # 直接打印流式内容
                return content
        except json.JSONDecodeError:
            pass
    return ''

def stream_solve(problem_text, formulas):
    url = f"{os.getenv('OPENAI_API_BASE')}/v1/chat/completions"
    
    full_problem = problem_text
    for i, formula in enumerate(formulas, 1):
        full_problem = full_problem.replace(f"[formula_{i}]", f"$${formula}$$")
    
    print(f"\n问题内容: {full_problem}")
    
    messages = [
        {
            "role": "system",
            "content": """你是一位经验丰富的教师。请按照以下格式回答问题：

[开始解答]
这是答案：（写出明确的答案，如"答案为：2.5米"，"答案选 A"等）
[解答结束]

[开始解析]
这是解析：（写出详细的解析过程）
[解析结束]"""
        },
        {
            "role": "user",
            "content": full_problem
        }
    ]
    
    payload = {
        "model": os.getenv('OPENAI_CHAT_MODEL'),
        "messages": messages,
        "stream": True,
        "temperature": 0.3
    }
    
    try:
        response = requests.post(
            url,
            headers=create_headers(),
            json=payload,
            stream=True
        )
        response.raise_for_status()

        complete_response = ""
        
        for line in response.iter_lines():
            if not line or not line.startswith(b'data: '):
                continue
                
            try:
                data = json.loads(line[6:].decode('utf-8'))
                content = data['choices'][0].get('delta', {}).get('content', '')
                complete_response += content
                
            except json.JSONDecodeError:
                continue
                    
    except Exception as e:
        print(f"\n错误: {str(e)}")
        yield f"data: {json.dumps({'error': str(e)})}\n\n"
        return

    # 提取答案和解析
    try:
        if '[开始解答]' in complete_response and '[开始解析]' in complete_response:
            parts = complete_response.split('[开始解答]')
            answer_part = parts[1].split('[解答结束]')[0].strip()
            
            parts = complete_response.split('[开始解析]')
            analysis_part = parts[1].split('[解析结束]')[0].strip()
            
            # 控制台打印
            print("\n完整回答:")
            print(f"答案: {answer_part}")
            print(f"解析: {analysis_part}")
            
            # 返回结果给前端
            yield f"data: {json.dumps({'type': 'answer', 'content': answer_part})}\n\n"
            yield f"data: {json.dumps({'type': 'analysis', 'content': analysis_part})}\n\n"
            
        else:
            print("\n响应格式不符合预期")
            yield f"data: {json.dumps({'error': '响应格式不符合预期'})}\n\n"
            
    except Exception as e:
        print(f"\n解析错误: {str(e)}")
        yield f"data: {json.dumps({'error': f'解析响应时出错: {str(e)}'})}\n\n"

@app.route('/')
def index():
    return render_template('index.html')

@app.route('/process', methods=['POST'])
def process():
    if 'file' not in request.files:
        return jsonify({'error': '没有文件上传'}), 400
    
    file = request.files['file']
    if file.filename == '':
        return jsonify({'error': '没有选择文件'}), 400
    
    try:
        image = Image.open(file)
        image_base64 = image_to_base64(image)
        
        result = process_image_with_vision(image_base64)
        
        if 'error' in result:
            return jsonify({'error': result['error']}), 500
        
        return jsonify({
            'original_image': image_base64,
            'result': result
        })
        
    except Exception as e:
        logger.error(f"处理图像错误: {str(e)}")
        return jsonify({'error': str(e)}), 500

@app.route('/solve', methods=['POST'])
def solve():
    data = request.json
    if not data or 'text' not in data or 'formulas' not in data:
        return jsonify({'error': '无效的请求数据'}), 400
    
    return Response(
        stream_solve(data['text'], data['formulas']),
        content_type='text/event-stream'
    )

if __name__ == '__main__':
    app.run(debug=True)