diff --git a/src/actions/ai-improve.ts b/src/actions/ai-improve.ts
index 85fa02d..4411d49 100644
--- a/src/actions/ai-improve.ts
+++ b/src/actions/ai-improve.ts
@@ -1,42 +1,75 @@
+"use server";
+
 import {
     OptimizeCodeInput,
     OptimizeCodeOutput,
-    OptimizeCodeOutputSchema
+    OptimizeCodeOutputSchema,
 } from "@/types/ai-improve";
 import { openai } from "@/lib/ai";
 import { CoreMessage, generateText } from "ai";
+import { prisma } from "@/lib/prisma"; // Prisma客户端
 
+/**
+ * 调用AI优化代码
+ * @param input 包含代码、错误信息、题目ID的输入
+ * @returns 优化后的代码和说明
+ */
 export const optimizeCode = async (
     input: OptimizeCodeInput
 ): Promise<OptimizeCodeOutput> => {
     const model = openai("gpt-4o-mini");
 
-    const prompt = `
-    Analyze the following programming code for potential errors, inefficiencies or code style issues.
-    Provide an optimized version of the code with explanations. Focus on:
-    1. Fixing any syntax errors
-    2. Improving performance
-    3. Enhancing code readability
-    4. Following best practices
-
-    Original code:
-    \`\`\`
-    ${input.code}
-    \`\`\`
-
-    Error message (if any): ${input.error || "No error message provided"}
-
-    Respond ONLY with the JSON object containing the optimized code and explanations.
-    Format:
-    {
-        "optimizedCode": "optimized code here",
-        "explanation": "explanation of changes made",
-        "issuesFixed": ["list of issues fixed"]
+    // 获取题目详情（如果提供了problemId）
+    let problemDetails = "";
+    if (input.problemId) {
+        try {
+            const problem = await prisma.problem.findUnique({
+                where: { problemId: input.problemId },
+            });
+            if (problem) {
+                problemDetails = `
+Problem Requirements:
+-------------------
+Description: ${problem.description}
+Input: ${problem.inputSpec}
+Output: ${problem.outputSpec}
+Test Cases: ${JSON.stringify(problem.testCases)}
+        `;
+            }
+        } catch (error) {
+            console.error("Failed to fetch problem details:", error);
+        }
     }
-    `;
 
+    // 构建AI提示词
+    const prompt = `
+Analyze the following programming code for potential errors, inefficiencies or code style issues.
+Provide an optimized version of the code with explanations. Focus on:
+1. Fixing any syntax errors
+2. Improving performance
+3. Enhancing code readability 
+4. Following best practices
+
+Original code:
+\`\`\`
+${input.code}
+\`\`\`
+
+Error message (if any): ${input.error || "No error message provided"}
+
+${problemDetails}
+
+Respond ONLY with the JSON object containing the optimized code and explanations.
+Format:
+{
+  "optimizedCode": "optimized code here",
+  "explanation": "explanation of changes made",
+  "issuesFixed": ["list of issues fixed"]
+}
+`;
+
+    // 发送请求给OpenAI
     const messages: CoreMessage[] = [{ role: "user", content: prompt }];
-
     let text;
     try {
         const response = await generateText({
@@ -49,6 +82,7 @@ export const optimizeCode = async (
         throw new Error("Failed to generate response from OpenAI");
     }
 
+    // 解析LLM响应
     let llmResponseJson;
     try {
         const cleanedText = text.trim();
@@ -59,13 +93,12 @@ export const optimizeCode = async (
         throw new Error("Invalid JSON response from LLM");
     }
 
-    const validationResult =
-        OptimizeCodeOutputSchema.safeParse(llmResponseJson);
-
+    // 验证响应格式
+    const validationResult = OptimizeCodeOutputSchema.safeParse(llmResponseJson);
     if (!validationResult.success) {
         console.error("Zod validation failed:", validationResult.error.format());
         throw new Error("Response validation failed");
     }
 
     return validationResult.data;
-};
\ No newline at end of file
+};
diff --git a/src/types/ai-improve.ts b/src/types/ai-improve.ts
index 6615eba..f7461b4 100644
--- a/src/types/ai-improve.ts
+++ b/src/types/ai-improve.ts
@@ -1,17 +1,19 @@
 import { z } from "zod";
 
+// 优化代码的输入类型
 export const OptimizeCodeInputSchema = z.object({
-    code: z.string(),
-    error: z.string().optional(),
+    code: z.string(),         // 用户输入的代码
+    error: z.string().optional(), // 可选的错误信息
+    problemId: z.string().optional(), // 可选的题目ID
 });
 
 export type OptimizeCodeInput = z.infer<typeof OptimizeCodeInputSchema>;
 
+// 优化代码的输出类型
 export const OptimizeCodeOutputSchema = z.object({
-    optimizedCode: z.string(),
-    explanation: z.string(),
-    issuesFixed: z.array(z.string()).optional(),
+    optimizedCode: z.string(),    // 优化后的代码
+    explanation: z.string(),      // 优化说明
+    issuesFixed: z.array(z.string()).optional(), // 修复的问题列表
 });
 
 export type OptimizeCodeOutput = z.infer<typeof OptimizeCodeOutputSchema>;
-