task generation (#450)

Co-authored-by: Shuchang Zheng <wintonzheng0325@gmail.com>

skyvern/forge/sdk/db/client.py

@@ -18,6 +18,7 @@ from skyvern.forge.sdk.db.models import (
     OrganizationModel,
     OutputParameterModel,
     StepModel,
+    TaskGenerationModel,
     TaskModel,
     WorkflowModel,
     WorkflowParameterModel,
@@ -42,6 +43,7 @@ from skyvern.forge.sdk.db.utils import (
     convert_to_workflow_run_parameter,
 )
 from skyvern.forge.sdk.models import Organization, OrganizationAuthToken, Step, StepStatus
+from skyvern.forge.sdk.schemas.task_generations import TaskGeneration
 from skyvern.forge.sdk.schemas.tasks import ProxyLocation, Task, TaskStatus
 from skyvern.forge.sdk.workflow.models.parameter import (
     AWSSecretParameter,
@@ -1236,3 +1238,34 @@ class AgentDB:
             )
             await session.execute(stmt)
             await session.commit()
+
+    async def create_task_generation(
+        self,
+        organization_id: str,
+        user_prompt: str,
+        url: str | None = None,
+        navigation_goal: str | None = None,
+        navigation_payload: dict[str, Any] | None = None,
+        data_extraction_goal: str | None = None,
+        extracted_information_schema: dict[str, Any] | None = None,
+        llm: str | None = None,
+        llm_prompt: str | None = None,
+        llm_response: str | None = None,
+    ) -> TaskGeneration:
+        async with self.Session() as session:
+            new_task_generation = TaskGenerationModel(
+                organization_id=organization_id,
+                user_prompt=user_prompt,
+                url=url,
+                navigation_goal=navigation_goal,
+                navigation_payload=navigation_payload,
+                data_extraction_goal=data_extraction_goal,
+                extracted_information_schema=extracted_information_schema,
+                llm=llm,
+                llm_prompt=llm_prompt,
+                llm_response=llm_response,
+            )
+            session.add(new_task_generation)
+            await session.commit()
+            await session.refresh(new_task_generation)
+            return TaskGeneration.model_validate(new_task_generation)

skyvern/forge/sdk/db/id.py

@@ -40,6 +40,7 @@ WORKFLOW_PARAMETER_PREFIX = "wp"
 AWS_SECRET_PARAMETER_PREFIX = "asp"
 OUTPUT_PARAMETER_PREFIX = "op"
 BITWARDEN_LOGIN_CREDENTIAL_PARAMETER_PREFIX = "blc"
+TASK_GENERATION_PREFIX = "tg"
 
 
 def generate_workflow_id() -> str:
@@ -107,6 +108,11 @@ def generate_user_id() -> str:
     return f"{USER_PREFIX}_{int_id}"
 
 
+def generate_task_generation_id() -> str:
+    int_id = generate_id()
+    return f"{TASK_GENERATION_PREFIX}_{int_id}"
+
+
 def generate_id() -> int:
     """
     generate a 64-bit int ID

skyvern/forge/sdk/db/models.py

@@ -26,6 +26,7 @@ from skyvern.forge.sdk.db.id import (
     generate_organization_auth_token_id,
     generate_output_parameter_id,
     generate_step_id,
+    generate_task_generation_id,
     generate_task_id,
     generate_workflow_id,
     generate_workflow_parameter_id,
@@ -325,3 +326,27 @@ class WorkflowRunOutputParameterModel(Base):
     )
     value = Column(JSON, nullable=False)
     created_at = Column(DateTime, default=datetime.datetime.utcnow, nullable=False)
+
+
+class TaskGenerationModel(Base):
+    """
+    Generate a task based on the prompt (natural language description of the task) from the user
+    """
+
+    __tablename__ = "task_generations"
+
+    task_generation_id = Column(String, primary_key=True, default=generate_task_generation_id)
+    organization_id = Column(String, ForeignKey("organizations.organization_id"), nullable=False)
+    user_prompt = Column(String, nullable=False, index=True)  # The prompt from the user
+    url = Column(String)
+    navigation_goal = Column(String)
+    navigation_payload = Column(JSON)
+    data_extraction_goal = Column(String)
+    extracted_information_schema = Column(JSON)
+
+    llm = Column(String)  # language model to use
+    llm_prompt = Column(String)  # The prompt sent to the language model
+    llm_response = Column(String)  # The response from the language model
+
+    created_at = Column(DateTime, default=datetime.datetime.utcnow, nullable=False)
+    modified_at = Column(DateTime, default=datetime.datetime.utcnow, onupdate=datetime.datetime.utcnow, nullable=False)