From 8f989235eabf24e8a14acf5d67ecd6fe03e7a963 Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Mon, 12 Feb 2024 09:30:59 -0800 Subject: [PATCH] fix(proxy_server.py): restrict proxy to just listed models --- litellm/proxy/proxy_server.py | 21 ++++++++++++++++++--- 1 file changed, 18 insertions(+), 3 deletions(-) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 4927f3db4c..12d272966e 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -2372,8 +2372,13 @@ async def chat_completion( llm_router is not None and data["model"] in llm_router.deployment_names ): # model in router deployments, calling a specific deployment on the router response = await llm_router.acompletion(**data, specific_deployment=True) - else: # router is not set + elif user_model is not None: # `litellm --model ` response = await litellm.acompletion(**data) + else: + raise HTTPException( + status_code=status.HTTP_422_UNPROCESSABLE_ENTITY, + detail={"message": "Invalid model name passed in"}, + ) # Post Call Processing data["litellm_status"] = "success" # used for alerting @@ -2584,8 +2589,13 @@ async def embeddings( llm_router is not None and data["model"] in llm_router.deployment_names ): # model in router deployments, calling a specific deployment on the router response = await llm_router.aembedding(**data, specific_deployment=True) - else: + elif user_model is not None: # `litellm --model ` response = await litellm.aembedding(**data) + else: + raise HTTPException( + status_code=status.HTTP_422_UNPROCESSABLE_ENTITY, + detail={"message": "Invalid model name passed in"}, + ) ### ALERTING ### data["litellm_status"] = "success" # used for alerting @@ -2719,8 +2729,13 @@ async def image_generation( response = await llm_router.aimage_generation( **data ) # ensure this goes the llm_router, router will do the correct alias mapping - else: + elif user_model is not None: # `litellm --model ` response = await litellm.aimage_generation(**data) + else: + raise HTTPException( + status_code=status.HTTP_422_UNPROCESSABLE_ENTITY, + detail={"message": "Invalid model name passed in"}, + ) ### ALERTING ### data["litellm_status"] = "success" # used for alerting