diff --git a/litellm/main.py b/litellm/main.py index aa52795f43..b7af4e8d39 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -5298,11 +5298,15 @@ def moderation( or get_secret_str("OPENAI_API_KEY") ) + # Extract api_base from kwargs + api_base = kwargs.get("api_base", None) + openai_client = kwargs.get("client", None) if openai_client is None: - openai_client = openai.OpenAI( - api_key=api_key, - ) + if api_base is not None: + openai_client = openai.OpenAI(api_key=api_key, base_url=api_base) + else: + openai_client = openai.OpenAI(api_key=api_key) if model is not None: response = openai_client.moderations.create(input=input, model=model) @@ -5332,21 +5336,11 @@ async def amoderation( or litellm.openai_key or get_secret_str("OPENAI_API_KEY") ) - openai_client = kwargs.get("client", None) - if openai_client is None or not isinstance(openai_client, AsyncOpenAI): - # call helper to get OpenAI client - # _get_openai_client maintains in-memory caching logic for OpenAI clients - _openai_client: AsyncOpenAI = openai_chat_completions._get_openai_client( # type: ignore - is_async=True, - api_key=api_key, - ) - else: - _openai_client = openai_client - optional_params = GenericLiteLLMParams(**kwargs) litellm_logging_obj: Optional[LiteLLMLoggingObj] = kwargs.get( "litellm_logging_obj", None ) + _dynamic_api_base = None try: ( model, @@ -5363,6 +5357,18 @@ async def amoderation( # `model` is optional field for moderation - get_llm_provider will throw BadRequestError if model is not set / not recognized pass + openai_client = kwargs.get("client", None) + if openai_client is None or not isinstance(openai_client, AsyncOpenAI): + # call helper to get OpenAI client + # _get_openai_client maintains in-memory caching logic for OpenAI clients + _openai_client: AsyncOpenAI = openai_chat_completions._get_openai_client( # type: ignore + is_async=True, + api_key=api_key, + api_base=optional_params.api_base or _dynamic_api_base, + ) + else: + _openai_client = openai_client + # update litellm_logging_obj with environment variables custom_llm_provider = custom_llm_provider or litellm.LlmProviders.OPENAI.value if litellm_logging_obj is not None: diff --git a/tests/router_unit_tests/test_router_endpoints.py b/tests/router_unit_tests/test_router_endpoints.py index 87af688461..3a6a955f80 100644 --- a/tests/router_unit_tests/test_router_endpoints.py +++ b/tests/router_unit_tests/test_router_endpoints.py @@ -260,6 +260,61 @@ async def test_moderation_endpoint(model): print("moderation response: ", response) +@pytest.mark.asyncio() +async def test_moderation_endpoint_with_api_base(): + """ + Test that the moderation endpoint respects api_base configuration + """ + from unittest.mock import AsyncMock, MagicMock, patch + + custom_api_base = "https://us.api.openai.com/v1" + + router = Router( + model_list=[ + { + "model_name": "openai/omni-moderation-latest", + "litellm_params": { + "model": "openai/omni-moderation-latest", + "api_base": custom_api_base, + "api_key": "test-key" + }, + }, + ] + ) + + # Mock the OpenAI client to verify api_base is passed + with patch("litellm.main.openai_chat_completions._get_openai_client") as mock_get_client: + mock_client = AsyncMock() + mock_response = MagicMock() + mock_response.model_dump.return_value = { + "id": "modr-123", + "model": "omni-moderation-latest", + "results": [ + { + "flagged": False, + "categories": {}, + "category_scores": {}, + "category_applied_input_types": {} + } + ] + } + mock_client.moderations.create = AsyncMock(return_value=mock_response) + mock_get_client.return_value = mock_client + + response = await router.amoderation( + model="openai/omni-moderation-latest", + input="hello this is a test" + ) + + # Verify that _get_openai_client was called with the custom api_base + mock_get_client.assert_called() + call_kwargs = mock_get_client.call_args.kwargs + assert call_kwargs.get("api_base") == custom_api_base, \ + f"Expected api_base to be {custom_api_base}, but got {call_kwargs.get('api_base')}" + + print(f"✓ Moderation endpoint correctly uses api_base: {custom_api_base}") + + @pytest.mark.parametrize("sync_mode", [True, False]) @pytest.mark.asyncio async def test_aaaaatext_completion_endpoint(model_list, sync_mode):