From ddf77656fb5d0e89059c22fc25059677cce7fa4c Mon Sep 17 00:00:00 2001 From: Marc Abramowitz Date: Wed, 15 May 2024 17:36:11 -0700 Subject: [PATCH 1/7] Allow non-admins to use OpenAI routes The problem we were having is non-admin users trying to use `/engines/{model}/chat/completions` were getting an HTTP 401 error. ```shell $ curl -sSL 'http://localhost:4000/engines/gpt-35-turbo-0125/chat/completions' \ --header "Authorization: Bearer ${LITELLM_KEY}" \ --header 'Content-Type: application/json' \ --data ' { "model": "gpt-35-turbo-0125", "messages": [ { "role": "user", "content": "Write a poem about LiteLLM" } ] }' \ | jq '.' { "error": { "message": "Authentication Error, Only proxy admin can be used to generate, delete, update info for new keys/users/teams. Route=/engines/gpt-35-turbo-0125/chat/completions. Your role=unknown. Your user_id=someone@company.com", "type": "auth_error", "param": "None", "code": 401 } } ``` This seems to be related to code in `user_api_key_auth` that checks that the URL matches a list of routes that are allowed for non-admin users, where the list of routes is in `LiteLLMRoutes.openai_routes.value`. The problem is that the route `/engines/{model}/chat/completions` is not in that list and furthermore, that wouldn't even work if it were, because the comparison is done with `request.url.path` and that will have the actual model name in it (e.g.: `gpt-35-turbo-0125`), rather than `{model}`. I added a new list `LiteLLMRoutes.openai_route_names` and added the route **names** to that list. Then I added a check in `user_api_key_auth` to see if the route name is in the list of route names. --- litellm/proxy/_types.py | 10 ++++++++++ litellm/proxy/proxy_server.py | 2 ++ 2 files changed, 12 insertions(+) diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index b1af153e81..204405cbfd 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -77,6 +77,16 @@ class LiteLLM_UpperboundKeyGenerateParams(LiteLLMBase): class LiteLLMRoutes(enum.Enum): + openai_route_names: List = [ + "chat_completion", + "completion", + "embeddings", + "image_generation", + "audio_transcriptions", + "moderations", + "model_info_v1", + "model_info_v2", + ] openai_routes: List = [ # chat completions "/openai/deployments/{model}/chat/completions", diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index cc5327b758..edfe69fa6e 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -1073,6 +1073,8 @@ async def user_api_key_auth( if not _is_user_proxy_admin(user_id_information): # if non-admin if route in LiteLLMRoutes.openai_routes.value: pass + elif request['route'].name in LiteLLMRoutes.openai_route_names.value: + pass elif ( route in LiteLLMRoutes.info_routes.value ): # check if user allowed to call an info route From b079f4cb79d3da739d2a6a17d501bf7255d8055b Mon Sep 17 00:00:00 2001 From: Marc Abramowitz Date: Thu, 16 May 2024 07:54:09 -0700 Subject: [PATCH 2/7] model_info_v{1,2} (LiteLLM) => model_list (OpenAI) --- litellm/proxy/_types.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 204405cbfd..d4027f704e 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -84,8 +84,7 @@ class LiteLLMRoutes(enum.Enum): "image_generation", "audio_transcriptions", "moderations", - "model_info_v1", - "model_info_v2", + "model_list", # OpenAI /v1/models route ] openai_routes: List = [ # chat completions From d5b2e8e7e8dd645df34a860d73948306e6993774 Mon Sep 17 00:00:00 2001 From: Marc Abramowitz Date: Thu, 16 May 2024 09:54:10 -0700 Subject: [PATCH 3/7] Make test_generate_and_call_with_valid_key parametrized This allows us to test the same code with different routes. For example, it lets us test the `/engines/{model}/chat/completions` route, which https://github.com/BerriAI/litellm/pull/3663 fixes. --- litellm/tests/test_key_generate_prisma.py | 18 +++++++++++++++--- 1 file changed, 15 insertions(+), 3 deletions(-) diff --git a/litellm/tests/test_key_generate_prisma.py b/litellm/tests/test_key_generate_prisma.py index 2eb693cf45..75379a58ab 100644 --- a/litellm/tests/test_key_generate_prisma.py +++ b/litellm/tests/test_key_generate_prisma.py @@ -23,6 +23,7 @@ import sys, os import traceback from dotenv import load_dotenv from fastapi import Request +from fastapi.routing import APIRoute from datetime import datetime load_dotenv() @@ -51,6 +52,7 @@ from litellm.proxy.proxy_server import ( user_info, info_key_fn, new_team, + chat_completion, ) from litellm.proxy.utils import PrismaClient, ProxyLogging, hash_token, update_spend from litellm._logging import verbose_proxy_logger @@ -146,7 +148,13 @@ async def test_new_user_response(prisma_client): pytest.fail(f"Got exception {e}") -def test_generate_and_call_with_valid_key(prisma_client): +@pytest.mark.parametrize( + "api_route", [ + APIRoute(path="/chat/completions", endpoint=chat_completion), + APIRoute(path="/engines/gpt-35-turbo-0125/chat/completions", endpoint=chat_completion), + ], +) +def test_generate_and_call_with_valid_key(prisma_client, api_route): # 1. Generate a Key, and use it to make a call print("prisma client=", prisma_client) @@ -181,8 +189,12 @@ def test_generate_and_call_with_valid_key(prisma_client): ) print("token from prisma", value_from_prisma) - request = Request(scope={"type": "http"}) - request._url = URL(url="/chat/completions") + request = Request({ + "type": "http", + "route": api_route, + "path": api_route.path, + "headers": [("Authorization", bearer_token)] + }) # use generated key to auth in result = await user_api_key_auth(request=request, api_key=bearer_token) From c427ea3781ef2d748163a338fa4f0bb01b7b152e Mon Sep 17 00:00:00 2001 From: Marc Abramowitz Date: Thu, 16 May 2024 10:03:23 -0700 Subject: [PATCH 4/7] Add "/engines/{model}/chat/completions" to openai_routes I don't think that this helps with the issue that I'm seeing, but I think it might be nice to have this model listed in the openai_routes list so that it's documented that it's a valid chat_completion route. --- litellm/proxy/_types.py | 1 + 1 file changed, 1 insertion(+) diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index d4027f704e..6ec954c8bc 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -88,6 +88,7 @@ class LiteLLMRoutes(enum.Enum): ] openai_routes: List = [ # chat completions + "/engines/{model}/chat/completions", "/openai/deployments/{model}/chat/completions", "/chat/completions", "/v1/chat/completions", From dc52c83b8880c2d41c0b96abfbb6b073bb0a3bc7 Mon Sep 17 00:00:00 2001 From: Marc Abramowitz Date: Thu, 16 May 2024 10:05:35 -0700 Subject: [PATCH 5/7] Add more routes to test_generate_and_call_with_valid_key --- litellm/tests/test_key_generate_prisma.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/litellm/tests/test_key_generate_prisma.py b/litellm/tests/test_key_generate_prisma.py index 75379a58ab..34913ba205 100644 --- a/litellm/tests/test_key_generate_prisma.py +++ b/litellm/tests/test_key_generate_prisma.py @@ -150,8 +150,10 @@ async def test_new_user_response(prisma_client): @pytest.mark.parametrize( "api_route", [ + APIRoute(path="/engines/{model}/chat/completions", endpoint=chat_completion), + APIRoute(path="/openai/deployments/{model}/chat/completions", endpoint=chat_completion), APIRoute(path="/chat/completions", endpoint=chat_completion), - APIRoute(path="/engines/gpt-35-turbo-0125/chat/completions", endpoint=chat_completion), + APIRoute(path="/v1/chat/completions", endpoint=chat_completion), ], ) def test_generate_and_call_with_valid_key(prisma_client, api_route): From cf71857354df693a0591234847428fc322765da9 Mon Sep 17 00:00:00 2001 From: Marc Abramowitz Date: Thu, 16 May 2024 10:44:36 -0700 Subject: [PATCH 6/7] Add more routes to test_generate_and_call_with_valid_key --- litellm/tests/test_key_generate_prisma.py | 28 +++++++++++++++++++++++ 1 file changed, 28 insertions(+) diff --git a/litellm/tests/test_key_generate_prisma.py b/litellm/tests/test_key_generate_prisma.py index 34913ba205..051136256f 100644 --- a/litellm/tests/test_key_generate_prisma.py +++ b/litellm/tests/test_key_generate_prisma.py @@ -53,6 +53,12 @@ from litellm.proxy.proxy_server import ( info_key_fn, new_team, chat_completion, + completion, + embeddings, + image_generation, + audio_transcriptions, + moderations, + model_list, ) from litellm.proxy.utils import PrismaClient, ProxyLogging, hash_token, update_spend from litellm._logging import verbose_proxy_logger @@ -150,10 +156,32 @@ async def test_new_user_response(prisma_client): @pytest.mark.parametrize( "api_route", [ + # chat_completion APIRoute(path="/engines/{model}/chat/completions", endpoint=chat_completion), APIRoute(path="/openai/deployments/{model}/chat/completions", endpoint=chat_completion), APIRoute(path="/chat/completions", endpoint=chat_completion), APIRoute(path="/v1/chat/completions", endpoint=chat_completion), + # completion + APIRoute(path="/completions", endpoint=completion), + APIRoute(path="/v1/completions", endpoint=completion), + APIRoute(path="/engines/{model}/completions", endpoint=completion), + APIRoute(path="/openai/deployments/{model}/completions", endpoint=completion), + # embeddings + APIRoute(path="/v1/embeddings", endpoint=embeddings), + APIRoute(path="/embeddings", endpoint=embeddings), + APIRoute(path="/openai/deployments/{model}/embeddings", endpoint=embeddings), + # image generation + APIRoute(path="/v1/images/generations", endpoint=image_generation), + APIRoute(path="/images/generations", endpoint=image_generation), + # audio transcriptions + APIRoute(path="/v1/audio/transcriptions", endpoint=audio_transcriptions), + APIRoute(path="/audio/transcriptions", endpoint=audio_transcriptions), + # moderations + APIRoute(path="/v1/moderations", endpoint=moderations), + APIRoute(path="/moderations", endpoint=moderations), + # model_list + APIRoute(path= "/v1/models", endpoint=model_list), + APIRoute(path= "/models", endpoint=model_list), ], ) def test_generate_and_call_with_valid_key(prisma_client, api_route): From 4194bafae044b431934a2d4027cdb40a06eb88b1 Mon Sep 17 00:00:00 2001 From: Marc Abramowitz Date: Thu, 16 May 2024 11:34:22 -0700 Subject: [PATCH 7/7] Add nicer test ids when using pytest -v Replace: ``` test_key_generate_prisma.py::test_generate_and_call_with_valid_key[api_route0] PASSED test_key_generate_prisma.py::test_generate_and_call_with_valid_key[api_route10] PASSED test_key_generate_prisma.py::test_generate_and_call_with_valid_key[api_route11] PASSED test_key_generate_prisma.py::test_generate_and_call_with_valid_key[api_route12] PASSED test_key_generate_prisma.py::test_generate_and_call_with_valid_key[api_route13] PASSED test_key_generate_prisma.py::test_generate_and_call_with_valid_key[api_route14] PASSED ```` with: ``` litellm/tests/test_key_generate_prisma.py::test_generate_and_call_with_valid_key[{'route': 'audio_transcriptions', 'path': '/audio/transcriptions'}] PASSED litellm/tests/test_key_generate_prisma.py::test_generate_and_call_with_valid_key[{'route': 'audio_transcriptions', 'path': '/v1/audio/transcriptions'}] PASSED litellm/tests/test_key_generate_prisma.py::test_generate_and_call_with_valid_key[{'route': 'chat_completion', 'path': '/chat/completions'}] PASSED litellm/tests/test_key_generate_prisma.py::test_generate_and_call_with_valid_key[{'route': 'chat_completion', 'path': '/engines/{model}/chat/completions'}] PASSED litellm/tests/test_key_generate_prisma.py::test_generate_and_call_with_valid_key[{'route': 'chat_completion', 'path': '/openai/deployments/{model}/chat/completions'}] PASSED litellm/tests/test_key_generate_prisma.py::test_generate_and_call_with_valid_key[{'route': 'chat_completion', 'path': '/v1/chat/completions'}] PASSED ``` --- litellm/tests/test_key_generate_prisma.py | 1 + 1 file changed, 1 insertion(+) diff --git a/litellm/tests/test_key_generate_prisma.py b/litellm/tests/test_key_generate_prisma.py index 051136256f..1d3a487da5 100644 --- a/litellm/tests/test_key_generate_prisma.py +++ b/litellm/tests/test_key_generate_prisma.py @@ -183,6 +183,7 @@ async def test_new_user_response(prisma_client): APIRoute(path= "/v1/models", endpoint=model_list), APIRoute(path= "/models", endpoint=model_list), ], + ids=lambda route: str(dict(route=route.endpoint.__name__, path=route.path)), ) def test_generate_and_call_with_valid_key(prisma_client, api_route): # 1. Generate a Key, and use it to make a call