From 24e7dc359d1b573012e48c195cdd9315d4415dc3 Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Mon, 1 Jan 2024 13:42:36 +0530 Subject: [PATCH] feat(proxy_server.py): introduces new `/user/auth` endpoint for handling user email auth decouples streamlit ui from proxy server. this then requires the proxy to handle user auth separately. --- litellm/proxy/proxy_server.py | 185 +++++++++++++--------------------- pyproject.toml | 8 +- 2 files changed, 73 insertions(+), 120 deletions(-) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 5410b8ef6..2183b277a 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -101,7 +101,7 @@ from typing import Union app = FastAPI( docs_url="/", title="LiteLLM API", - description="Proxy Server to call 100+ LLMs in the OpenAI format\n\nAdmin Panel on `/admin` endpoint", + description="Proxy Server to call 100+ LLMs in the OpenAI format\n\nAdmin Panel on `https://dashboard.litellm.ai/admin`", ) router = APIRouter() origins = ["*"] @@ -199,18 +199,27 @@ async def user_api_key_auth( if user_custom_auth: response = await user_custom_auth(request=request, api_key=api_key) return UserAPIKeyAuth.model_validate(response) - + ### LITELLM-DEFINED AUTH FUNCTION ### if master_key is None: if isinstance(api_key, str): return UserAPIKeyAuth(api_key=api_key) else: return UserAPIKeyAuth() + route: str = request.url.path + print(f"route: {route}") + if route == "/user/auth": + if general_settings.get("allow_user_auth", False) == True: + return UserAPIKeyAuth() + else: + raise HTTPException( + status_code=status.HTTP_403_FORBIDDEN, + detail="'allow_user_auth' not set or set to False", + ) + if api_key is None: # only require api key if master key is set raise Exception(f"No api key passed in.") - route: str = request.url.path - # note: never string compare api keys, this is vulenerable to a time attack. Use secrets.compare_digest instead is_master_key_valid = secrets.compare_digest(api_key, master_key) if is_master_key_valid: @@ -347,36 +356,6 @@ def load_from_azure_key_vault(use_azure_key_vault: bool = False): ) -async def run_streamlit_ui(): - # Save the current working directory - original_dir = os.getcwd() - # set the working directory to where this script is - abspath = os.path.abspath(__file__) - dname = os.path.dirname(abspath) - os.chdir(dname) - try: - # Start Streamlit without opening the browser automatically - process = subprocess.Popen( - [ - "streamlit", - "run", - "admin_ui.py", - "--server.headless=true", - "--browser.serverAddress=0.0.0.0", - "--server.enableCORS=false", - ] - ) - # Wait for the server to start before exiting the context manager - await asyncio.sleep(1) - print("Streamlit UI server has started successfully.") - os.chdir(original_dir) - # Keep the background task running - while True: - await asyncio.sleep(3600) - except Exception as e: - print_verbose(f"Admin UI - Streamlit. An error occurred: {e}") - - def cost_tracking(): global prisma_client if prisma_client is not None: @@ -1678,8 +1657,6 @@ async def info_key_fn( #### USER MANAGEMENT #### - - @router.post( "/user/new", tags=["user management"], @@ -1719,6 +1696,57 @@ async def new_user(data: NewUserRequest): ) +@router.post( + "/user/auth", tags=["user management"], dependencies=[Depends(user_api_key_auth)] +) +async def user_auth(request: Request): + """ + Allows UI ("https://dashboard.litellm.ai/", or self-hosted - os.getenv("LITELLM_HOSTED_UI")) to request a magic link to be sent to user email, for auth to proxy. + + Only allows emails from accepted email subdomains. + + Rate limit: 1 request every 60s. + + Only works, if you enable 'allow_user_auth' in general settings: + e.g.: + ```yaml + general_settings: + allow_user_auth: true + ``` + + Requirements: + This uses [Resend](https://resend.com/) for sending emails. Needs these 2 keys in your .env: + ```env + RESEND_API_KEY = "my-resend-api-key" + RESEND_API_EMAIL = "my-sending-email" + ``` + """ + data = await request.json() # type: ignore + user_email = data["user_email"] + import os + import resend + + ## [TODO]: Check if user exists, if so - use an existing key, if not - create new user -> return new key + response = await generate_key_helper_fn( + **{"duration": "1hr", "models": [], "aliases": {}, "config": {}, "spend": 0} # type: ignore + ) + + base_url = os.getenv("LITELLM_HOSTED_UI", "https://dashboard.litellm.ai/") + + resend.api_key = os.getenv("RESEND_API_KEY") + + params = { + "from": f"LiteLLM Proxy <{os.getenv('RESEND_API_EMAIL')}>", + "to": [user_email], + "subject": "Your Magic Link", + "html": f" Follow this link, to login:\n\n{base_url}user/?token={response['token']}&user_id={response['user_id']}", + } + + email = resend.Emails.send(params) + print(email) + return "Email sent!" + + @router.post( "/user/info", tags=["user management"], dependencies=[Depends(user_api_key_auth)] ) @@ -1798,6 +1826,12 @@ async def add_new_model(model_params: ModelParams): tags=["model management"], dependencies=[Depends(user_api_key_auth)], ) +@router.get( + "/v1/model/info", + description="Provides more info about each model in /models, including config.yaml descriptions (except api key and api base)", + tags=["model management"], + dependencies=[Depends(user_api_key_auth)], +) async def model_info_v1(request: Request): global llm_model_list, general_settings, user_config_file_path # Load existing config @@ -1822,55 +1856,6 @@ async def model_info_v1(request: Request): return {"data": all_models} -#### [BETA] - This is a beta endpoint, format might change based on user feedback. - https://github.com/BerriAI/litellm/issues/933 -@router.get( - "/v1/model/info", - description="Provides more info about each model in /models, including config.yaml descriptions (except api key and api base)", - tags=["model management"], - dependencies=[Depends(user_api_key_auth)], -) -async def model_info(request: Request): - global llm_model_list, general_settings, user_config_file_path - # Load existing config - with open(f"{user_config_file_path}", "r") as config_file: - config = yaml.safe_load(config_file) - all_models = config["model_list"] - - for model in all_models: - # get the model cost map info - ## make an api call - data = copy.deepcopy(model["litellm_params"]) - data["messages"] = [{"role": "user", "content": "Hey, how's it going?"}] - data["max_tokens"] = 10 - print(f"data going to litellm acompletion: {data}") - response = await litellm.acompletion(**data) - response_model = response["model"] - print(f"response model: {response_model}; response - {response}") - litellm_model_info = litellm.get_model_info(response_model) - model_info = model.get("model_info", {}) - for k, v in litellm_model_info.items(): - if k not in model_info: - model_info[k] = v - model["model_info"] = model_info - # don't return the api key - model["litellm_params"].pop("api_key", None) - - # all_models = list(set([m["model_name"] for m in llm_model_list])) - print_verbose(f"all_models: {all_models}") - return dict( - data=[ - { - "id": model, - "object": "model", - "created": 1677610602, - "owned_by": "openai", - } - for model in all_models - ], - object="list", - ) - - #### [BETA] - This is a beta endpoint, format might change based on user feedback. - https://github.com/BerriAI/litellm/issues/964 @router.post( "/model/delete", @@ -2083,36 +2068,6 @@ async def retrieve_server_log(request: Request): return FileResponse(filepath) -#### ADMIN UI ENDPOINTS #### - - -@router.get("/admin") -async def admin_page(request: Request): - from fastapi.responses import HTMLResponse - - # Assuming your Streamlit app is running on localhost port 8501 - html_content = """ - - - Admin Page - - - - - - - """ - return HTMLResponse(content=html_content) - - #### BASIC ENDPOINTS #### diff --git a/pyproject.toml b/pyproject.toml index aa5229f27..02046cc58 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -36,15 +36,13 @@ proxy = [ "orjson", ] -proxy-ui = [ - "streamlit" -] - extra_proxy = [ "prisma", "azure-identity", "azure-keyvault-secrets", - "google-cloud-kms" + "google-cloud-kms", + "streamlit", + "resend" ]