From 1890fde3f377b0896e7ca5908953ef948ec1c8a7 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Tue, 19 Nov 2024 07:02:12 -0800 Subject: [PATCH] (Proxy) add support for DOCS_URL and REDOC_URL (#6806) * add support for DOCS_URL and REDOC_URL * document env vars * add unit tests for docs url and redocs url --- docs/my-website/docs/proxy/configs.md | 2 + litellm/proxy/proxy_server.py | 7 ++-- litellm/proxy/utils.py | 29 ++++++++++++++ tests/proxy_unit_tests/test_proxy_utils.py | 44 ++++++++++++++++++++++ 4 files changed, 79 insertions(+), 3 deletions(-) diff --git a/docs/my-website/docs/proxy/configs.md b/docs/my-website/docs/proxy/configs.md index 1609e16aefc3..3b6b336d6897 100644 --- a/docs/my-website/docs/proxy/configs.md +++ b/docs/my-website/docs/proxy/configs.md @@ -934,6 +934,7 @@ router_settings: | DOCS_DESCRIPTION | Description text for documentation pages | DOCS_FILTERED | Flag indicating filtered documentation | DOCS_TITLE | Title of the documentation pages +| DOCS_URL | The path to the Swagger API documentation. **By default this is "/"** | EMAIL_SUPPORT_CONTACT | Support contact email address | GCS_BUCKET_NAME | Name of the Google Cloud Storage bucket | GCS_PATH_SERVICE_ACCOUNT | Path to the Google Cloud service account JSON file @@ -1041,6 +1042,7 @@ router_settings: | REDIS_HOST | Hostname for Redis server | REDIS_PASSWORD | Password for Redis service | REDIS_PORT | Port number for Redis server +| REDOC_URL | The path to the Redoc Fast API documentation. **By default this is "/redoc"** | SERVER_ROOT_PATH | Root path for the server application | SET_VERBOSE | Flag to enable verbose logging | SLACK_DAILY_REPORT_FREQUENCY | Frequency of daily Slack reports (e.g., daily, weekly) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 2ece9705ae9f..4d4c6a1a27e8 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -222,7 +222,9 @@ def generate_feedback_box(): PrismaClient, ProxyLogging, _cache_user_row, + _get_docs_url, _get_projected_spend_over_limit, + _get_redoc_url, _is_projected_spend_over_limit, _is_valid_team_configs, get_error_message_str, @@ -344,7 +346,6 @@ def generate_feedback_box(): custom_swagger_message = "[**Customize Swagger Docs**](https://docs.litellm.ai/docs/proxy/enterprise#swagger-docs---custom-routes--branding)" ### CUSTOM BRANDING [ENTERPRISE FEATURE] ### -_docs_url = None if os.getenv("NO_DOCS", "False") == "True" else "/" _title = os.getenv("DOCS_TITLE", "LiteLLM API") if premium_user else "LiteLLM API" _description = ( os.getenv( @@ -355,9 +356,9 @@ def generate_feedback_box(): else f"Proxy Server to call 100+ LLMs in the OpenAI format. {custom_swagger_message}\n\n{ui_message}" ) - app = FastAPI( - docs_url=_docs_url, + docs_url=_get_docs_url(), + redoc_url=_get_redoc_url(), title=_title, description=_description, version=version, diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index bcd3ce16c33d..e495f3490dd0 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -3099,6 +3099,34 @@ def get_error_message_str(e: Exception) -> str: return error_message +def _get_redoc_url() -> str: + """ + Get the redoc URL from the environment variables. + + - If REDOC_URL is set, return it. + - Otherwise, default to "/redoc". + """ + return os.getenv("REDOC_URL", "/redoc") + + +def _get_docs_url() -> Optional[str]: + """ + Get the docs URL from the environment variables. + + - If DOCS_URL is set, return it. + - If NO_DOCS is True, return None. + - Otherwise, default to "/". + """ + docs_url = os.getenv("DOCS_URL", None) + if docs_url: + return docs_url + + if os.getenv("NO_DOCS", "False") == "True": + return None + + # default to "/" + return "/" + def handle_exception_on_proxy(e: Exception) -> ProxyException: """ Returns an Exception as ProxyException, this ensures all exceptions are OpenAI API compatible @@ -3120,3 +3148,4 @@ def handle_exception_on_proxy(e: Exception) -> ProxyException: param=getattr(e, "param", "None"), code=status.HTTP_500_INTERNAL_SERVER_ERROR, ) + diff --git a/tests/proxy_unit_tests/test_proxy_utils.py b/tests/proxy_unit_tests/test_proxy_utils.py index 2e857808d5c7..607e54225656 100644 --- a/tests/proxy_unit_tests/test_proxy_utils.py +++ b/tests/proxy_unit_tests/test_proxy_utils.py @@ -2,6 +2,7 @@ import os import sys from unittest.mock import Mock +from litellm.proxy.utils import _get_redoc_url, _get_docs_url import pytest from fastapi import Request @@ -530,3 +531,46 @@ def test_prepare_key_update_data(): data = UpdateKeyRequest(key="test_key", metadata=None) updated_data = prepare_key_update_data(data, existing_key_row) assert updated_data["metadata"] == None + + +@pytest.mark.parametrize( + "env_value, expected_url", + [ + (None, "/redoc"), # default case + ("/custom-redoc", "/custom-redoc"), # custom URL + ("https://example.com/redoc", "https://example.com/redoc"), # full URL + ], +) +def test_get_redoc_url(env_value, expected_url): + if env_value is not None: + os.environ["REDOC_URL"] = env_value + else: + os.environ.pop("REDOC_URL", None) # ensure env var is not set + + result = _get_redoc_url() + assert result == expected_url + + +@pytest.mark.parametrize( + "env_vars, expected_url", + [ + ({}, "/"), # default case + ({"DOCS_URL": "/custom-docs"}, "/custom-docs"), # custom URL + ( + {"DOCS_URL": "https://example.com/docs"}, + "https://example.com/docs", + ), # full URL + ({"NO_DOCS": "True"}, None), # docs disabled + ], +) +def test_get_docs_url(env_vars, expected_url): + # Clear relevant environment variables + for key in ["DOCS_URL", "NO_DOCS"]: + os.environ.pop(key, None) + + # Set test environment variables + for key, value in env_vars.items(): + os.environ[key] = value + + result = _get_docs_url() + assert result == expected_url