From c703332acc3c978fec6732847b3d2cfd6d3b1330 Mon Sep 17 00:00:00 2001 From: whiterabbit1983 Date: Fri, 12 Jul 2024 16:44:15 +0300 Subject: [PATCH] fix: Add exception handler for litellm API error (#432) --- agents-api/agents_api/web.py | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/agents-api/agents_api/web.py b/agents-api/agents_api/web.py index 28980eb26..7c2864f0f 100644 --- a/agents-api/agents_api/web.py +++ b/agents-api/agents_api/web.py @@ -14,6 +14,7 @@ from agents_api.exceptions import PromptTooBigError from pycozo.client import QueryException from temporalio.service import RPCError +from litellm.exceptions import APIError from agents_api.dependencies.auth import get_api_key from agents_api.env import sentry_dsn @@ -120,6 +121,14 @@ async def prompt_too_big_error(request: Request, exc: PromptTooBigError): ) +@app.exception_handler(APIError) +async def litellm_api_error(request: Request, exc: APIError): + return JSONResponse( + status_code=status.HTTP_502_BAD_GATEWAY, + content={"error": {"message": str(exc)}}, + ) + + def main( host="127.0.0.1", port=8000,