From 8f989235eabf24e8a14acf5d67ecd6fe03e7a963 Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Mon, 12 Feb 2024 09:30:59 -0800 Subject: [PATCH 1/3] 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 From 1a452057af683d4a6e68dd6d659b9ea1e74a47cd Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Mon, 12 Feb 2024 11:13:02 -0800 Subject: [PATCH 2/3] fix(proxy_server.py): fix exception raising --- litellm/proxy/proxy_server.py | 33 ++++++++++++++----- litellm/tests/test_proxy_exception_mapping.py | 1 - 2 files changed, 24 insertions(+), 10 deletions(-) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 12d272966e..b341adde67 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -2376,8 +2376,8 @@ async def chat_completion( response = await litellm.acompletion(**data) else: raise HTTPException( - status_code=status.HTTP_422_UNPROCESSABLE_ENTITY, - detail={"message": "Invalid model name passed in"}, + status_code=status.HTTP_400_BAD_REQUEST, + detail={"error": "Invalid model name passed in"}, ) # Post Call Processing @@ -2439,7 +2439,12 @@ async def chat_completion( traceback.print_exc() if isinstance(e, HTTPException): - raise e + raise ProxyException( + message=getattr(e, "detail", str(e)), + type=getattr(e, "type", "None"), + param=getattr(e, "param", "None"), + code=getattr(e, "status_code", status.HTTP_400_BAD_REQUEST), + ) else: error_traceback = traceback.format_exc() error_msg = f"{str(e)}\n\n{error_traceback}" @@ -2593,8 +2598,8 @@ async def embeddings( response = await litellm.aembedding(**data) else: raise HTTPException( - status_code=status.HTTP_422_UNPROCESSABLE_ENTITY, - detail={"message": "Invalid model name passed in"}, + status_code=status.HTTP_400_BAD_REQUEST, + detail={"error": "Invalid model name passed in"}, ) ### ALERTING ### @@ -2613,7 +2618,12 @@ async def embeddings( ) traceback.print_exc() if isinstance(e, HTTPException): - raise e + raise ProxyException( + message=getattr(e, "message", str(e)), + type=getattr(e, "type", "None"), + param=getattr(e, "param", "None"), + code=getattr(e, "status_code", status.HTTP_400_BAD_REQUEST), + ) else: error_traceback = traceback.format_exc() error_msg = f"{str(e)}\n\n{error_traceback}" @@ -2733,8 +2743,8 @@ async def image_generation( response = await litellm.aimage_generation(**data) else: raise HTTPException( - status_code=status.HTTP_422_UNPROCESSABLE_ENTITY, - detail={"message": "Invalid model name passed in"}, + status_code=status.HTTP_400_BAD_REQUEST, + detail={"error": "Invalid model name passed in"}, ) ### ALERTING ### @@ -2753,7 +2763,12 @@ async def image_generation( ) traceback.print_exc() if isinstance(e, HTTPException): - raise e + raise ProxyException( + message=getattr(e, "message", str(e)), + type=getattr(e, "type", "None"), + param=getattr(e, "param", "None"), + code=getattr(e, "status_code", status.HTTP_400_BAD_REQUEST), + ) else: error_traceback = traceback.format_exc() error_msg = f"{str(e)}\n\n{error_traceback}" diff --git a/litellm/tests/test_proxy_exception_mapping.py b/litellm/tests/test_proxy_exception_mapping.py index cab26f319e..31829586d8 100644 --- a/litellm/tests/test_proxy_exception_mapping.py +++ b/litellm/tests/test_proxy_exception_mapping.py @@ -160,7 +160,6 @@ def test_chat_completion_exception_any_model(client): response = client.post("/chat/completions", json=test_data) json_response = response.json() - print("keys in json response", json_response.keys()) assert json_response.keys() == {"error"} # make an openai client to call _make_status_error_from_response From 456ce692df013c72101d5dfda29bbe7dbe0ff3d1 Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Mon, 12 Feb 2024 16:25:35 -0800 Subject: [PATCH 3/3] test(test_proxy_exception_mapping.py): fix test --- litellm/tests/test_proxy_exception_mapping.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/litellm/tests/test_proxy_exception_mapping.py b/litellm/tests/test_proxy_exception_mapping.py index 31829586d8..a2c7952586 100644 --- a/litellm/tests/test_proxy_exception_mapping.py +++ b/litellm/tests/test_proxy_exception_mapping.py @@ -139,7 +139,7 @@ def test_exception_openai_bad_model(client): response=response ) print("Type of exception=", type(openai_exception)) - assert isinstance(openai_exception, openai.NotFoundError) + assert isinstance(openai_exception, openai.BadRequestError) except Exception as e: pytest.fail(f"LiteLLM Proxy test failed. Exception {str(e)}")