diff --git a/backend/server/src/AIWorker.ts b/backend/server/src/AIWorker.ts
new file mode 100644
index 0000000..aed1aab
--- /dev/null
+++ b/backend/server/src/AIWorker.ts
@@ -0,0 +1,77 @@
+// AIWorker class for handling AI-related operations
+export class AIWorker {
+  private aiWorkerUrl: string
+  private cfAiKey: string
+  private databaseWorkerUrl: string
+  private workersKey: string
+
+  // Constructor to initialize AIWorker with necessary URLs and keys
+  constructor(
+    aiWorkerUrl: string,
+    cfAiKey: string,
+    databaseWorkerUrl: string,
+    workersKey: string
+  ) {
+    this.aiWorkerUrl = aiWorkerUrl
+    this.cfAiKey = cfAiKey
+    this.databaseWorkerUrl = databaseWorkerUrl
+    this.workersKey = workersKey
+  }
+
+  // Method to generate code based on user input
+  async generateCode(
+    userId: string,
+    fileName: string,
+    code: string,
+    line: number,
+    instructions: string
+  ): Promise<{ response: string; success: boolean }> {
+    try {
+      // Fetch request to the database worker
+      const fetchPromise = fetch(
+        `${this.databaseWorkerUrl}/api/sandbox/generate`,
+        {
+          method: "POST",
+          headers: {
+            "Content-Type": "application/json",
+            Authorization: `${this.workersKey}`,
+          },
+          body: JSON.stringify({
+            userId: userId,
+          }),
+        }
+      )
+
+      // Fetch request to the AI worker for code generation
+      const generateCodePromise = fetch(
+        `${this.aiWorkerUrl}/api?fileName=${encodeURIComponent(
+          fileName
+        )}&code=${encodeURIComponent(code)}&line=${encodeURIComponent(
+          line
+        )}&instructions=${encodeURIComponent(instructions)}`,
+        {
+          headers: {
+            "Content-Type": "application/json",
+            Authorization: `${this.cfAiKey}`,
+          },
+        }
+      )
+
+      // Wait for both fetch requests to complete
+      const [fetchResponse, generateCodeResponse] = await Promise.all([
+        fetchPromise,
+        generateCodePromise,
+      ])
+
+      // Parse the response from the AI worker
+      const json = await generateCodeResponse.json()
+
+      // Return the generated code response
+      return { response: json.response, success: true }
+    } catch (e: any) {
+      // Log and throw an error if code generation fails
+      console.error("Error generating code:", e)
+      throw new Error(`Error: code generation. ${e.message ?? e}`)
+    }
+  }
+}
diff --git a/backend/server/src/index.ts b/backend/server/src/index.ts
index 9d82c4e..2d9d42b 100644
--- a/backend/server/src/index.ts
+++ b/backend/server/src/index.ts
@@ -6,6 +6,7 @@ import fs from "fs"
 import { createServer } from "http"
 import { Server } from "socket.io"
 import { z } from "zod"
+import { AIWorker } from "./AIWorker"
 import { DokkuClient } from "./DokkuClient"
 import { FileManager, SandboxFiles } from "./FileManager"
 import {
@@ -161,6 +162,14 @@ const git =
       )
     : null
 
+// Add this near the top of the file, after other initializations
+const aiWorker = new AIWorker(
+  process.env.AI_WORKER_URL!,
+  process.env.CF_AI_KEY!,
+  process.env.DATABASE_WORKER_URL!,
+  process.env.WORKERS_KEY!
+)
+
 // Handle socket connections
 io.on("connection", async (socket) => {
   try {
@@ -470,43 +479,14 @@ io.on("connection", async (socket) => {
         callback
       ) => {
         try {
-          const fetchPromise = fetch(
-            `${process.env.DATABASE_WORKER_URL}/api/sandbox/generate`,
-            {
-              method: "POST",
-              headers: {
-                "Content-Type": "application/json",
-                Authorization: `${process.env.WORKERS_KEY}`,
-              },
-              body: JSON.stringify({
-                userId: data.userId,
-              }),
-            }
+          const result = await aiWorker.generateCode(
+            data.userId,
+            fileName,
+            code,
+            line,
+            instructions
           )
-
-          // Generate code from Cloudflare Workers AI
-          const generateCodePromise = fetch(
-            `${process.env.AI_WORKER_URL}/api?fileName=${encodeURIComponent(
-              fileName
-            )}&code=${encodeURIComponent(code)}&line=${encodeURIComponent(
-              line
-            )}&instructions=${encodeURIComponent(instructions)}`,
-            {
-              headers: {
-                "Content-Type": "application/json",
-                Authorization: `${process.env.CF_AI_KEY}`,
-              },
-            }
-          )
-
-          const [fetchResponse, generateCodeResponse] = await Promise.all([
-            fetchPromise,
-            generateCodePromise,
-          ])
-
-          const json = await generateCodeResponse.json()
-
-          callback({ response: json.response, success: true })
+          callback(result)
         } catch (e: any) {
           console.error("Error generating code:", e)
           io.emit("error", `Error: code generation. ${e.message ?? e}`)