diff --git a/skyvern/forge/sdk/schemas/task_generations.py b/skyvern/forge/sdk/schemas/task_generations.py
index 7e1ff4a8..141023f1 100644
--- a/skyvern/forge/sdk/schemas/task_generations.py
+++ b/skyvern/forge/sdk/schemas/task_generations.py
@@ -1,14 +1,9 @@
 from datetime import datetime
-from enum import StrEnum
 from typing import Any
 
 from pydantic import BaseModel, ConfigDict
 
 
-class LLMType(StrEnum):
-    OPENAI_GPT4O = "OPENAI_GPT4O"
-
-
 class TaskGenerationBase(BaseModel):
     model_config = ConfigDict(from_attributes=True)
 
@@ -19,7 +14,7 @@ class TaskGenerationBase(BaseModel):
     navigation_payload: dict[str, Any] | None = None
     data_extraction_goal: str | None = None
     extracted_information_schema: dict[str, Any] | None = None
-    llm: LLMType | None = None
+    llm: str | None = None
     llm_prompt: str | None = None
     llm_response: str | None = None