(Proxy) add support for DOCS_URL and REDOC_URL (BerriAI#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

@@ -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)

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,

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,
     )
+

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