diff --git a/skyvern/exceptions.py b/skyvern/exceptions.py index b195ee281..431951eb8 100644 --- a/skyvern/exceptions.py +++ b/skyvern/exceptions.py @@ -1,9 +1,18 @@ +from fastapi import status + + class SkyvernException(Exception): def __init__(self, message: str | None = None): self.message = message super().__init__(message) +class SkyvernHTTPException(SkyvernException): + def __init__(self, message: str | None = None, status_code: int = status.HTTP_400_BAD_REQUEST): + self.status_code = status_code + super().__init__(message) + + class InvalidOpenAIResponseFormat(SkyvernException): def __init__(self, message: str | None = None): super().__init__(f"Invalid response format: {message}") @@ -22,9 +31,9 @@ def __init__(self, proxy_location: str | None = None): super().__init__(f"Unknown proxy location: {proxy_location}") -class TaskNotFound(SkyvernException): +class TaskNotFound(SkyvernHTTPException): def __init__(self, task_id: str | None = None): - super().__init__(f"Task {task_id} not found") + super().__init__(f"Task {task_id} not found", status_code=status.HTTP_404_NOT_FOUND) class ScriptNotFound(SkyvernException): @@ -97,9 +106,9 @@ def __init__(self, block_type: str) -> None: super().__init__(f"Unknown block type {block_type}") -class WorkflowNotFound(SkyvernException): +class WorkflowNotFound(SkyvernHTTPException): def __init__(self, workflow_id: str) -> None: - super().__init__(f"Workflow {workflow_id} not found") + super().__init__(f"Workflow {workflow_id} not found", status_code=status.HTTP_404_NOT_FOUND) class WorkflowRunNotFound(SkyvernException): @@ -119,9 +128,9 @@ def __init__(self, parameter_key: str, workflow_id: str, workflow_run_id: str) - ) -class WorkflowParameterNotFound(SkyvernException): +class WorkflowParameterNotFound(SkyvernHTTPException): def __init__(self, workflow_parameter_id: str) -> None: - super().__init__(f"Workflow parameter {workflow_parameter_id} not found") + super().__init__(f"Workflow parameter {workflow_parameter_id} not found", status_code=status.HTTP_404_NOT_FOUND) class FailedToNavigateToUrl(SkyvernException): @@ -163,14 +172,17 @@ def __init__(self) -> None: super().__init__("BrowserState is missing the main page") -class OrganizationNotFound(SkyvernException): +class OrganizationNotFound(SkyvernHTTPException): def __init__(self, organization_id: str) -> None: - super().__init__(f"Organization {organization_id} not found") + super().__init__(f"Organization {organization_id} not found", status_code=status.HTTP_404_NOT_FOUND) -class StepNotFound(SkyvernException): +class StepNotFound(SkyvernHTTPException): def __init__(self, organization_id: str, task_id: str, step_id: str | None = None) -> None: - super().__init__(f"Step {step_id or 'latest'} not found. organization_id={organization_id} task_id={task_id}") + super().__init__( + f"Step {step_id or 'latest'} not found. organization_id={organization_id} task_id={task_id}", + status_code=status.HTTP_404_NOT_FOUND, + ) class FailedToTakeScreenshot(SkyvernException): diff --git a/skyvern/forge/api_app.py b/skyvern/forge/api_app.py index e0175b6c4..744d46830 100644 --- a/skyvern/forge/api_app.py +++ b/skyvern/forge/api_app.py @@ -10,6 +10,7 @@ from starlette_context.middleware import RawContextMiddleware from starlette_context.plugins.base import Plugin +from skyvern.exceptions import SkyvernHTTPException from skyvern.forge import app as forge_app from skyvern.forge.sdk.core import skyvern_context from skyvern.forge.sdk.core.skyvern_context import SkyvernContext @@ -73,6 +74,10 @@ def start_scheduler() -> None: LOG.info("Server startup complete. Skyvern is now online") + @app.exception_handler(SkyvernHTTPException) + async def handle_skyvern_http_exception(request: Request, exc: SkyvernHTTPException) -> JSONResponse: + return JSONResponse(status_code=exc.status_code, content={"detail": exc.message}) + @app.exception_handler(Exception) async def unexpected_exception(request: Request, exc: Exception) -> JSONResponse: LOG.exception("Unexpected error in agent server.", exc_info=exc)