Script generation (#3157)

skyvern/forge/sdk/db/client.py

@@ -3742,31 +3742,71 @@ class AgentDB:
         mime_type: str | None = None,
         encoding: str = "utf-8",
         artifact_id: str | None = None,
-    ) -> None:
-        """Create a script file record."""
-        try:
-            async with self.Session() as session:
-                script_file = ScriptFileModel(
-                    script_revision_id=script_revision_id,
-                    script_id=script_id,
-                    organization_id=organization_id,
-                    file_path=file_path,
-                    file_name=file_name,
-                    file_type=file_type,
-                    content_hash=content_hash,
-                    file_size=file_size,
-                    mime_type=mime_type,
-                    encoding=encoding,
-                    artifact_id=artifact_id,
+    ) -> ScriptFile:
+        """Create a script file."""
+        async with self.Session() as session:
+            script_file = ScriptFileModel(
+                script_revision_id=script_revision_id,
+                script_id=script_id,
+                organization_id=organization_id,
+                file_path=file_path,
+                file_name=file_name,
+                file_type=file_type,
+                content_hash=content_hash,
+                file_size=file_size,
+                mime_type=mime_type,
+                encoding=encoding,
+                artifact_id=artifact_id,
+            )
+            session.add(script_file)
+            await session.commit()
+            await session.refresh(script_file)
+            return convert_to_script_file(script_file)
+
+    async def create_script_block(
+        self,
+        script_revision_id: str,
+        script_id: str,
+        organization_id: str,
+        script_block_label: str,
+        script_file_id: str | None = None,
+    ) -> ScriptBlock:
+        """Create a script block."""
+        async with self.Session() as session:
+            script_block = ScriptBlockModel(
+                script_revision_id=script_revision_id,
+                script_id=script_id,
+                organization_id=organization_id,
+                script_block_label=script_block_label,
+                script_file_id=script_file_id,
+            )
+            session.add(script_block)
+            await session.commit()
+            await session.refresh(script_block)
+            return convert_to_script_block(script_block)
+
+    async def update_script_block(
+        self,
+        script_block_id: str,
+        organization_id: str,
+        script_file_id: str | None = None,
+    ) -> ScriptBlock:
+        async with self.Session() as session:
+            script_block = (
+                await session.scalars(
+                    select(ScriptBlockModel)
+                    .filter_by(script_block_id=script_block_id)
+                    .filter_by(organization_id=organization_id)
                 )
-                session.add(script_file)
+            ).first()
+            if script_block:
+                if script_file_id:
+                    script_block.script_file_id = script_file_id
                 await session.commit()
-        except SQLAlchemyError:
-            LOG.error("SQLAlchemyError", exc_info=True)
-            raise
-        except Exception:
-            LOG.error("UnexpectedError", exc_info=True)
-            raise
+                await session.refresh(script_block)
+                return convert_to_script_block(script_block)
+            else:
+                raise NotFoundError("Script block not found")
 
     async def get_script_files(self, script_revision_id: str, organization_id: str) -> list[ScriptFile]:
         async with self.Session() as session:

skyvern/forge/sdk/routes/scripts.py

@@ -35,44 +35,12 @@ async def create_script(
     current_org: Organization = Depends(org_auth_service.get_current_org),
 ) -> CreateScriptResponse:
     """Create a new script with optional files and metadata."""
-    organization_id = current_org.organization_id
-    LOG.info(
-        "Creating script",
-        organization_id=organization_id,
-        file_count=len(data.files) if data.files else 0,
+    return await script_service.create_script(
+        organization_id=current_org.organization_id,
+        workflow_id=data.workflow_id,
+        run_id=data.run_id,
+        files=data.files,
     )
-    if data.run_id:
-        if not await app.DATABASE.get_run(run_id=data.run_id, organization_id=organization_id):
-            raise HTTPException(status_code=404, detail=f"Run_id {data.run_id} not found")
-    try:
-        # Create the script in the database
-        script = await app.DATABASE.create_script(
-            organization_id=organization_id,
-            run_id=data.run_id,
-        )
-        # Process files if provided
-        file_tree = {}
-        file_count = 0
-        if data.files:
-            file_tree = await script_service.build_file_tree(
-                data.files,
-                organization_id=organization_id,
-                script_id=script.script_id,
-                script_version=script.version,
-                script_revision_id=script.script_revision_id,
-            )
-            file_count = len(data.files)
-        return CreateScriptResponse(
-            script_id=script.script_id,
-            version=script.version,
-            run_id=script.run_id,
-            file_count=file_count,
-            created_at=script.created_at,
-            file_tree=file_tree,
-        )
-    except Exception as e:
-        LOG.error("Failed to create script", error=str(e), exc_info=True)
-        raise HTTPException(status_code=500, detail="Failed to create script")
 
 
 @base_router.get(

skyvern/forge/sdk/workflow/service.py

@@ -11,8 +11,8 @@ from jinja2.sandbox import SandboxedEnvironment
 from skyvern import analytics
 from skyvern.config import settings
 from skyvern.constants import GET_DOWNLOADED_FILES_TIMEOUT, SAVE_DOWNLOADED_FILES_TIMEOUT
-from skyvern.core.code_generations.generate_code import generate_workflow_script as generate_python_workflow_script
-from skyvern.core.code_generations.transform_workflow_run import transform_workflow_run_to_code_gen_input
+from skyvern.core.script_generations.generate_script import generate_workflow_script as generate_python_workflow_script
+from skyvern.core.script_generations.transform_workflow_run import transform_workflow_run_to_code_gen_input
 from skyvern.exceptions import (
     BlockNotFound,
     BrowserSessionNotFound,
@@ -2287,7 +2287,7 @@ class WorkflowService:
                 workflow_run_id=workflow_run.workflow_run_id,
                 organization_id=workflow.organization_id,
             )
-            python_src = generate_python_workflow_script(
+            python_src = await generate_python_workflow_script(
                 file_name=codegen_input.file_name,
                 workflow_run_request=codegen_input.workflow_run,
                 workflow=codegen_input.workflow,