diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 9981c049c18..504b19ab0e6 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -496,6 +496,7 @@ from litellm.proxy.utils import ( _get_openapi_url, _get_projected_spend_over_limit, _get_redoc_url, + _get_openapi_url, _is_projected_spend_over_limit, _is_valid_team_configs, get_custom_url, diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index a6f81986a6f..e822b3ff241 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -5367,6 +5367,22 @@ def _get_docs_url() -> Optional[str]: return "/" +def _get_openapi_url() -> Optional[str]: + """ + Get the OpenAPI JSON URL from the environment variables. + + - If OPENAPI_URL is set, return it. + - If NO_OPENAPI is True, return None. + - Otherwise, default to "/openapi.json". + """ + if openapi_url := os.getenv("OPENAPI_URL"): + return openapi_url + + if str_to_bool(os.getenv("NO_OPENAPI")) is True: + return None + + return "/openapi.json" + def handle_exception_on_proxy(e: Exception) -> ProxyException: """ diff --git a/tests/proxy_unit_tests/test_proxy_utils.py b/tests/proxy_unit_tests/test_proxy_utils.py index 09f6a85938d..9f5f14457e8 100644 --- a/tests/proxy_unit_tests/test_proxy_utils.py +++ b/tests/proxy_unit_tests/test_proxy_utils.py @@ -10,7 +10,7 @@ import pytest from fastapi import Request from starlette.datastructures import State -from litellm.proxy.utils import _get_docs_url, _get_redoc_url +from litellm.proxy.utils import _get_docs_url, _get_openapi_url, _get_redoc_url sys.path.insert( 0, os.path.abspath("../..") @@ -735,6 +735,30 @@ def test_get_docs_url(env_vars, expected_url): result = _get_docs_url() assert result == expected_url +@pytest.mark.parametrize( + "env_vars, expected_url", + [ + ({}, "/openapi.json"), # default case + ({"OPENAPI_URL": "/custom-openapi.json"}, "/custom-openapi.json"), # custom URL + ( + {"OPENAPI_URL": "https://example.com/openapi.json"}, + "https://example.com/openapi.json", + ), # full URL + ({"NO_OPENAPI": "True"}, None), # openapi disabled + ], +) +def test_get_openapi_url(env_vars, expected_url): + # Clear relevant environment variables + for key in ["OPENAPI_URL", "NO_OPENAPI"]: + os.environ.pop(key, None) + + # Set test environment variables + for key, value in env_vars.items(): + os.environ[key] = value + + result = _get_openapi_url() + assert result == expected_url + @pytest.mark.parametrize( "request_tags, tags_to_add, expected_tags",