(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
This commit is contained in:
Ishaan Jaff 2024-11-19 07:02:12 -08:00 committed by GitHub
parent 7550aba474
commit 1890fde3f3
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
4 changed files with 79 additions and 3 deletions

View file

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

View file

@ -222,7 +222,9 @@ from litellm.proxy.utils import (
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 @@ ui_message += "\n\n💸 [```LiteLLM Model Cost Map```](https://models.litellm.ai
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 @@ _description = (
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,

View file

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

View file

@ -2,6 +2,7 @@ import asyncio
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