diff --git a/src/actions/ai-improve.ts b/src/actions/ai-improve.ts
index 7e65716..85fa02d 100644
--- a/src/actions/ai-improve.ts
+++ b/src/actions/ai-improve.ts
@@ -1,34 +1,41 @@
-"use server";
-
 import {
-    AnalyzeComplexityResponse,
-    AnalyzeComplexityResponseSchema,
-    Complexity,
-} from "@/types/complexity";
+    OptimizeCodeInput,
+    OptimizeCodeOutput,
+    OptimizeCodeOutputSchema
+} from "@/types/ai-improve";
 import { openai } from "@/lib/ai";
 import { CoreMessage, generateText } from "ai";
 
-export const analyzeComplexity = async (
-    content: string
-): Promise<AnalyzeComplexityResponse> => {
+export const optimizeCode = async (
+    input: OptimizeCodeInput
+): Promise<OptimizeCodeOutput> => {
     const model = openai("gpt-4o-mini");
 
     const prompt = `
-  Analyze the time and space complexity of the following programming code snippet.
-  Determine the Big O notation from this list: ${Complexity.options.join(", ")}.
-  Provide your response as a JSON object with one key:
-  1. "time": A string representing the time complexity (e.g., "O(N)", "O(N^2)").
-  2. "space": A string representing the space complexity (e.g., "O(1)", "O(N)").
+    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
 
-  Code to analyze:
-  \`\`\`
-  ${content}
-  \`\`\`
+    Original code:
+    \`\`\`
+    ${input.code}
+    \`\`\`
 
-  Respond ONLY with the JSON object. Do not include any other text or markdown formatting like \`\`\`json before or after the object.
-  `;
+    Error message (if any): ${input.error || "No error message provided"}
 
-    const messages: CoreMessage[] = [{role: "user", content: prompt}];
+    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"]
+    }
+    `;
+
+    const messages: CoreMessage[] = [{ role: "user", content: prompt }];
 
     let text;
     try {
@@ -53,12 +60,12 @@ export const analyzeComplexity = async (
     }
 
     const validationResult =
-        AnalyzeComplexityResponseSchema.safeParse(llmResponseJson);
+        OptimizeCodeOutputSchema.safeParse(llmResponseJson);
 
     if (!validationResult.success) {
         console.error("Zod validation failed:", validationResult.error.format());
         throw new Error("Response validation failed");
     }
 
-    return validationResult;
-}
\ No newline at end of file
+    return validationResult.data;
+};
\ No newline at end of file
diff --git a/src/lib/ai.ts b/src/lib/ai.ts
new file mode 100644
index 0000000..95e527c
--- /dev/null
+++ b/src/lib/ai.ts
@@ -0,0 +1,8 @@
+import "server-only";
+
+import { createOpenAI } from "@ai-sdk/openai";
+
+export const openai = createOpenAI({
+    apiKey: process.env.OPENAI_API_KEY || "",
+    baseURL: process.env.OPENAI_BASE_URL || "",
+});
\ No newline at end of file
diff --git a/src/types/ai-improve.ts b/src/types/ai-improve.ts
new file mode 100644
index 0000000..6615eba
--- /dev/null
+++ b/src/types/ai-improve.ts
@@ -0,0 +1,17 @@
+import { z } from "zod";
+
+export const OptimizeCodeInputSchema = z.object({
+    code: z.string(),
+    error: z.string().optional(),
+});
+
+export type OptimizeCodeInput = z.infer<typeof OptimizeCodeInputSchema>;
+
+export const OptimizeCodeOutputSchema = z.object({
+    optimizedCode: z.string(),
+    explanation: z.string(),
+    issuesFixed: z.array(z.string()).optional(),
+});
+
+export type OptimizeCodeOutput = z.infer<typeof OptimizeCodeOutputSchema>;
+