From ed40019a14d31557c56ddb76117101ced6cdf2f6 Mon Sep 17 00:00:00 2001
From: yuetsh <517252939@qq.com>
Date: Mon, 5 Jan 2026 10:23:02 +0800
Subject: [PATCH] =?UTF-8?q?=E8=87=AA=E5=8A=A8=E7=94=9F=E6=88=90=E6=B5=81?=
 =?UTF-8?q?=E7=A8=8B=E5=9B=BE?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 ai/views/oj.py         | 10 ++--------
 flowchart/tasks.py     | 10 +++-------
 problem/urls/admin.py  |  2 ++
 problem/views/admin.py | 25 +++++++++++++++++++++++++
 utils/openai.py        | 13 +++++++++++++
 5 files changed, 45 insertions(+), 15 deletions(-)
 create mode 100644 utils/openai.py

diff --git a/ai/views/oj.py b/ai/views/oj.py
index cb101a1..c824116 100644
--- a/ai/views/oj.py
+++ b/ai/views/oj.py
@@ -9,10 +9,9 @@ from django.db.models import Min, Count
 from django.db.models.functions import TruncDate
 from django.http import StreamingHttpResponse
 from django.utils import timezone
-from openai import OpenAI
 
 from utils.api import APIView
-from utils.shortcuts import get_env
+from utils.openai import get_ai_client
 
 from account.models import User
 from problem.models import Problem
@@ -528,12 +527,7 @@ class AIAnalysisAPI(APIView):
         details = request.data.get("details")
         duration = request.data.get("duration")
 
-        api_key = get_env("AI_KEY")
-
-        if not api_key:
-            return self.error("API_KEY is not set")
-
-        client = OpenAI(api_key=api_key, base_url="https://api.deepseek.com")
+        client = get_ai_client()
 
         system_prompt = "你是一个风趣的编程老师，学生使用判题狗平台进行编程练习。请根据学生提供的详细数据和每周数据，给出用户的学习建议，最后写一句鼓励学生的话。请使用 markdown 格式输出，不要在代码块中输出。"
         user_prompt = f"这段时间内的详细数据: {details}\n(其中部分字段含义是 flowcharts:流程图的提交,solved:代码的提交)\n每周或每月的数据: {duration}"
diff --git a/flowchart/tasks.py b/flowchart/tasks.py
index fdf230f..faf5346 100644
--- a/flowchart/tasks.py
+++ b/flowchart/tasks.py
@@ -1,10 +1,10 @@
 import dramatiq
 import json
 import time
-from openai import OpenAI
 from django.db import transaction
 from django.utils import timezone
-from utils.shortcuts import get_env, DRAMATIQ_WORKER_ARGS
+from utils.openai import get_ai_client
+from utils.shortcuts import DRAMATIQ_WORKER_ARGS
 from .models import FlowchartSubmission, FlowchartSubmissionStatus
 
 @dramatiq.actor(**DRAMATIQ_WORKER_ARGS(max_retries=3))
@@ -49,11 +49,7 @@ def evaluate_flowchart_task(submission_id):
         """
         
         # 调用AI进行评分
-        api_key = get_env("AI_KEY")
-        if not api_key:
-            raise Exception("AI_KEY is not set")
-            
-        client = OpenAI(api_key=api_key, base_url="https://api.deepseek.com")
+        client = get_ai_client()
         
         response = client.chat.completions.create(
             model="deepseek-chat",
diff --git a/problem/urls/admin.py b/problem/urls/admin.py
index 2538080..aaf58fa 100644
--- a/problem/urls/admin.py
+++ b/problem/urls/admin.py
@@ -3,6 +3,7 @@ from django.urls import path
 from ..views.admin import (
     ContestProblemAPI,
     ProblemAPI,
+    ProblemFlowchartAIGen,
     TestCaseAPI,
     MakeContestProblemPublicAPIView,
     AddContestProblemAPI,
@@ -13,6 +14,7 @@ urlpatterns = [
     path("test_case", TestCaseAPI.as_view()),
     path("problem", ProblemAPI.as_view()),
     path("problem/visible", ProblemVisibleAPI.as_view()),
+    path("problem/flowchart", ProblemFlowchartAIGen.as_view()),
     path("contest/problem", ContestProblemAPI.as_view()),
     path("contest_problem/make_public", MakeContestProblemPublicAPIView.as_view()),
     path("contest/add_problem_from_public", AddContestProblemAPI.as_view()),
diff --git a/problem/views/admin.py b/problem/views/admin.py
index 8552b46..3ec9e80 100644
--- a/problem/views/admin.py
+++ b/problem/views/admin.py
@@ -15,6 +15,7 @@ from contest.models import Contest, ContestStatus
 from submission.models import Submission
 from utils.api import APIView, CSRFExemptAPIView, validate_serializer, APIError
 from utils.shortcuts import rand_str, natural_sort_key
+from utils.openai import get_ai_client
 from ..models import Problem, ProblemRuleType, ProblemTag
 from ..serializers import (
     CreateContestProblemSerializer,
@@ -484,3 +485,27 @@ class ProblemVisibleAPI(APIView):
         problem.visible = not problem.visible
         problem.save()
         return self.success()
+
+
+class ProblemFlowchartAIGen(APIView):
+    @problem_permission_required
+    def post(self, request):
+        python_code = request.data.get("python", "")
+        client = get_ai_client()
+        response = client.chat.completions.create(
+            model="deepseek-chat",
+            messages=[
+                {
+                    "role": "system",
+                    "content": """你是一个可以将Python代码转换为mermaid的助手。
+                    请将用户提供的Python代码转换为 Mermaid 纯文本。
+                    注意括号内的内容用引号包裹，如果本身就有引号，请注意双引号和单引号的问题。
+                    请只返回 mermaid 代码，连 ``` 都不需要。""",
+                },
+                {"role": "user", "content": python_code},
+            ],
+            temperature=1.0,
+        )
+
+        mermaid_code = response.choices[0].message.content
+        return self.success({"flowchart": mermaid_code})
diff --git a/utils/openai.py b/utils/openai.py
new file mode 100644
index 0000000..c945d40
--- /dev/null
+++ b/utils/openai.py
@@ -0,0 +1,13 @@
+from openai import OpenAI
+
+from utils.shortcuts import get_env
+
+BASE_URL = "https://api.deepseek.com"
+
+
+def get_ai_client() -> OpenAI:
+    key = get_env("AI_KEY")
+    if not key:
+        raise Exception("缺少 AI_KEY")
+
+    return OpenAI(api_key=key, base_url=BASE_URL)