自动生成流程图

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}"

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",

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()),

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})

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)