add SkyvernHTTPException base exception to automate known error handling in agent service (#317)

This commit is contained in:
Shuchang Zheng 2024-05-15 13:35:45 -07:00 committed by GitHub
parent 164a4da03a
commit c17b4ed287
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 27 additions and 10 deletions

View file

@ -1,9 +1,18 @@
from fastapi import status
class SkyvernException(Exception): class SkyvernException(Exception):
def __init__(self, message: str | None = None): def __init__(self, message: str | None = None):
self.message = message self.message = message
super().__init__(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): class InvalidOpenAIResponseFormat(SkyvernException):
def __init__(self, message: str | None = None): def __init__(self, message: str | None = None):
super().__init__(f"Invalid response format: {message}") super().__init__(f"Invalid response format: {message}")
@ -22,9 +31,9 @@ class ProxyLocationNotSupportedError(SkyvernException):
super().__init__(f"Unknown proxy location: {proxy_location}") super().__init__(f"Unknown proxy location: {proxy_location}")
class TaskNotFound(SkyvernException): class TaskNotFound(SkyvernHTTPException):
def __init__(self, task_id: str | None = None): 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): class ScriptNotFound(SkyvernException):
@ -97,9 +106,9 @@ class UnknownBlockType(SkyvernException):
super().__init__(f"Unknown block type {block_type}") super().__init__(f"Unknown block type {block_type}")
class WorkflowNotFound(SkyvernException): class WorkflowNotFound(SkyvernHTTPException):
def __init__(self, workflow_id: str) -> None: 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): class WorkflowRunNotFound(SkyvernException):
@ -119,9 +128,9 @@ class MissingValueForParameter(SkyvernException):
) )
class WorkflowParameterNotFound(SkyvernException): class WorkflowParameterNotFound(SkyvernHTTPException):
def __init__(self, workflow_parameter_id: str) -> None: 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): class FailedToNavigateToUrl(SkyvernException):
@ -163,14 +172,17 @@ class BrowserStateMissingPage(SkyvernException):
super().__init__("BrowserState is missing the main page") super().__init__("BrowserState is missing the main page")
class OrganizationNotFound(SkyvernException): class OrganizationNotFound(SkyvernHTTPException):
def __init__(self, organization_id: str) -> None: 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: 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): class FailedToTakeScreenshot(SkyvernException):

View file

@ -10,6 +10,7 @@ from starlette.requests import HTTPConnection, Request
from starlette_context.middleware import RawContextMiddleware from starlette_context.middleware import RawContextMiddleware
from starlette_context.plugins.base import Plugin from starlette_context.plugins.base import Plugin
from skyvern.exceptions import SkyvernHTTPException
from skyvern.forge import app as forge_app from skyvern.forge import app as forge_app
from skyvern.forge.sdk.core import skyvern_context from skyvern.forge.sdk.core import skyvern_context
from skyvern.forge.sdk.core.skyvern_context import SkyvernContext from skyvern.forge.sdk.core.skyvern_context import SkyvernContext
@ -73,6 +74,10 @@ def get_agent_app(router: APIRouter = base_router) -> FastAPI:
LOG.info("Server startup complete. Skyvern is now online") 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) @app.exception_handler(Exception)
async def unexpected_exception(request: Request, exc: Exception) -> JSONResponse: async def unexpected_exception(request: Request, exc: Exception) -> JSONResponse:
LOG.exception("Unexpected error in agent server.", exc_info=exc) LOG.exception("Unexpected error in agent server.", exc_info=exc)