diff --git a/backend/apps/litellm/main.py b/backend/apps/litellm/main.py index 9bc08598f..40619be2f 100644 --- a/backend/apps/litellm/main.py +++ b/backend/apps/litellm/main.py @@ -9,6 +9,7 @@ from fastapi.responses import JSONResponse from starlette.middleware.base import BaseHTTPMiddleware, RequestResponseEndpoint from starlette.responses import StreamingResponse import json +import time import requests from pydantic import BaseModel @@ -16,7 +17,7 @@ from typing import Optional, List from utils.utils import get_verified_user, get_current_user, get_admin_user from config import SRC_LOG_LEVELS, ENV -from constants import ERROR_MESSAGES +from constants import MESSAGES log = logging.getLogger(__name__) log.setLevel(SRC_LOG_LEVELS["LITELLM"]) @@ -201,6 +202,7 @@ async def get_models(user=Depends(get_current_user)): return data except Exception as e: + log.exception(e) error_detail = "Open WebUI: Server Connection Error" if r is not None: @@ -211,10 +213,18 @@ async def get_models(user=Depends(get_current_user)): except: error_detail = f"External: {e}" - raise HTTPException( - status_code=r.status_code if r else 500, - detail=error_detail, - ) + return { + "data": [ + { + "id": model["model_name"], + "object": "model", + "created": int(time.time()), + "owned_by": "openai", + } + for model in app.state.CONFIG["model_list"] + ], + "object": "list", + } @app.get("/model/info") @@ -231,6 +241,8 @@ class AddLiteLLMModelForm(BaseModel): async def add_model_to_config( form_data: AddLiteLLMModelForm, user=Depends(get_admin_user) ): + # TODO: Validate model form + app.state.CONFIG["model_list"].append(form_data.model_dump()) with open(LITELLM_CONFIG_DIR, "w") as file: @@ -238,7 +250,7 @@ async def add_model_to_config( await restart_litellm() - return {"message": "model added"} + return {"message": MESSAGES.MODEL_ADDED(form_data.model_name)} class DeleteLiteLLMModelForm(BaseModel): @@ -260,7 +272,7 @@ async def delete_model_from_config( await restart_litellm() - return {"message": "model deleted"} + return {"message": MESSAGES.MODEL_DELETED(form_data.id)} @app.api_route("/{path:path}", methods=["GET", "POST", "PUT", "DELETE"]) diff --git a/backend/constants.py b/backend/constants.py index da1ee0b3f..310c13311 100644 --- a/backend/constants.py +++ b/backend/constants.py @@ -3,6 +3,10 @@ from enum import Enum class MESSAGES(str, Enum): DEFAULT = lambda msg="": f"{msg if msg else ''}" + MODEL_ADDED = lambda model="": f"The model '{model}' has been added successfully." + MODEL_DELETED = ( + lambda model="": f"The model '{model}' has been deleted successfully." + ) class WEBHOOK_MESSAGES(str, Enum):