Update app/main.py

app/main.py

@@ -10,13 +10,15 @@ from datetime import datetime
 from io import BytesIO
 from urllib.parse import quote_plus
 from jinja2 import Template
-import openai
+from openai import OpenAI
 
 # —— 环境变量读取 —— #
 BAIDU_API_KEY    = os.getenv("BAIDU_API_KEY", "")
 BAIDU_SECRET_KEY = os.getenv("BAIDU_SECRET_KEY", "")
 OPENAI_API_KEY   = os.getenv("OPENAI_API_KEY", "")
-openai.api_key   = OPENAI_API_KEY
+
+# —— 初始化 OpenAI v1 客户端 —— #
+client = OpenAI(api_key=OPENAI_API_KEY)
 
 # —— Baidu OCR —— #
 def get_access_token(api_key, secret_key):
@@ -32,9 +34,8 @@ def get_access_token(api_key, secret_key):
     return resp.json().get("access_token")
 
 def ocr_image(image_bytes: bytes, token: str):
-    # Base64 编码后 URL 编码
-    img_b64 = base64.b64encode(image_bytes).decode()
-    img_encoded = quote_plus(img_b64)
+    img_b64     = base64.b64encode(image_bytes).decode()
+    img_encoded = quote_plus(img_b64)  # URL 编码
     resp = requests.post(
         f"https://aip.baidubce.com/rest/2.0/ocr/v1/handwriting?access_token={token}",
         headers={"Content-Type": "application/x-www-form-urlencoded"},
@@ -44,102 +45,95 @@ def ocr_image(image_bytes: bytes, token: str):
         }
     )
     resp.raise_for_status()
-    data = resp.json()
-    return data.get("words_result", [])
+    return resp.json().get("words_result", [])
 
 # —— 文本高亮 —— #
 def highlight_brackets(text: str) -> str:
-    # 中括号绿色；小括号红色
     text = re.sub(r'\[([^\[\]]+)\]', r'<span class="highlight-bracket-green">\1</span>', text)
     text = re.sub(r'\(([^\(\)]+)\)', r'<span class="highlight-bracket">\1</span>', text)
     return text.replace("\n", "<br>")
 
 # —— 主处理函数 —— #
 def process(image_pil):
-    # 1. PIL 转 bytes
+    # —— PIL 转 bytes —— #
     buf = BytesIO()
     image_pil.save(buf, format="PNG")
     image_bytes = buf.getvalue()
 
-    # 2. OCR
-    token = get_access_token(BAIDU_API_KEY, BAIDU_SECRET_KEY)
-    words = ocr_image(image_bytes, token)
+    # —— OCR 提取原文 —— #
+    token      = get_access_token(BAIDU_API_KEY, BAIDU_SECRET_KEY)
+    words      = ocr_image(image_bytes, token)
     essay_text = "\n".join([w["words"] for w in words])
 
-    # 3. 格式化原文
+    # —— 1. 格式化原文 —— #
     fmt_prompt = (
         "请帮我整理下面的英语作文文本格式，只整理英文正文部分，"
         "保证原汁原味（明显错误空格换行、乱码、非常用字符改正），"
         "拼写错误保留：\n\n" + essay_text
     )
-    fm = openai.ChatCompletion.create(
+    fm_resp = client.chat.completions.create(
         model="gpt-4o-mini",
         messages=[{"role": "user", "content": fmt_prompt}]
     )
-    revised = fm.choices[0].message.content
+    revised = fm_resp.choices[0].message.content
 
-    # 4. 批改
+    # —— 2. 批改 —— #
     corr_prompt = (
         "请帮我把下面的英语作文的语法错误改正，输出改正后的文章，"
         "原文错误用()括起来，修改部分用[]括起来：\n\n" + revised
     )
-    cm = openai.ChatCompletion.create(
+    cm_resp = client.chat.completions.create(
         model="gpt-4o-mini",
         messages=[{"role": "user", "content": corr_prompt}]
     )
-    corrected = cm.choices[0].message.content
+    corrected = cm_resp.choices[0].message.content
 
-    # 5. 批改意见
+    # —— 3. 批改意见 —— #
     review_prompt = (
         "下面是一份已经批改过的英语作文，请根据批注给出逐条批改意见：\n\n" + corrected
     )
-    rm = openai.ChatCompletion.create(
+    rm_resp = client.chat.completions.create(
         model="gpt-4o-mini",
         messages=[{"role": "user", "content": review_prompt}]
     )
-    review = rm.choices[0].message.content
+    review = rm_resp.choices[0].message.content
 
-    # 6. 评分
+    # —— 4. 评分 —— #
     rate_prompt = (
         "请按照 IELTS/CEFR 写作评价体系，从语言通顺度、连贯度、词汇与语法三维度打分并给出原因：\n\n" + revised
     )
-    rr = openai.ChatCompletion.create(
+    rr_resp = client.chat.completions.create(
         model="gpt-4o-mini",
         messages=[{"role": "user", "content": rate_prompt}]
     )
-    rating = rr.choices[0].message.content
+    rating = rr_resp.choices[0].message.content
 
-    # 7. 优秀范文
+    # —— 5. 优秀范文 —— #
     rewrite_prompt = (
         "请使用优秀表达重写下面这篇作文，加粗可供学习的部分：\n\n" + revised
     )
-    wm = openai.ChatCompletion.create(
+    wm_resp = client.chat.completions.create(
         model="gpt-4o-mini",
         messages=[{"role": "user", "content": rewrite_prompt}]
     )
-    perfect = wm.choices[0].message.content
+    perfect = wm_resp.choices[0].message.content
 
-    # 8. 渲染 HTML
-    code = f"{random.randint(0,9999):04}-{datetime.now().strftime('%Y%m%d%H%M%S')}"
+    # —— 6. 渲染 HTML —— #
+    code     = f"{random.randint(0,9999):04}-{datetime.now().strftime('%Y%m%d%H%M%S')}"
     tpl_path = os.path.join("app", "templates", "base.html")
     with open(tpl_path, encoding="utf-8") as f:
         tpl = Template(f.read())
 
     html_content = (
-        "<h2>原文格式化</h2>"
-        + markdown.markdown(revised)
-        + "<h2>批改结果</h2>"
-        + highlight_brackets(corrected)
-        + "<h2>批改意见</h2>"
-        + markdown.markdown(review)
-        + "<h2>评分</h2>"
-        + markdown.markdown(rating)
-        + "<h2>优秀范文</h2>"
-        + markdown.markdown(perfect)
+        "<h2>原文格式化</h2>"   + markdown.markdown(revised) +
+        "<h2>批改结果</h2>"     + highlight_brackets(corrected) +
+        "<h2>批改意见</h2>"     + markdown.markdown(review) +
+        "<h2>评分</h2>"         + markdown.markdown(rating) +
+        "<h2>优秀范文</h2>"     + markdown.markdown(perfect)
     )
     full_html = tpl.render(code=code, content=html_content)
 
-    # 9. 写文件并返回
+    # —— 7. 写文件并返回 —— #
     output_dir = os.path.join("app", "output")
     os.makedirs(output_dir, exist_ok=True)
     html_path = os.path.join(output_dir, f"{code}.html")
@@ -148,19 +142,21 @@ def process(image_pil):
     with open(html_path, "w", encoding="utf-8") as f:
         f.write(full_html)
 
-    pdfkit.from_string(full_html, pdf_path,
-                       options={"enable-local-file-access": ""})
+    pdfkit.from_string(
+        full_html, pdf_path,
+        options={"enable-local-file-access": ""}
+    )
 
     return full_html, html_path, pdf_path
 
 # —— Gradio 接口 —— #
 with gr.Blocks(title="英语作文批改") as demo:
     gr.Markdown("## 上传英语作文照片，等待批改完成后下载 HTML 或 PDF")
-    image_in  = gr.Image(type="pil", label="上传照片")
+    image_in    = gr.Image(type="pil", label="上传照片")
     output_html = gr.HTML()
-    btn = gr.Button("开始批改")
-    file_html = gr.File(label="下载 HTML")
-    file_pdf  = gr.File(label="下载 PDF")
+    btn         = gr.Button("开始批改")
+    file_html   = gr.File(label="下载 HTML")
+    file_pdf    = gr.File(label="下载 PDF")
 
     btn.click(
         fn=process,