diff --git a/litellm/litellm_core_utils/get_llm_provider_logic.py b/litellm/litellm_core_utils/get_llm_provider_logic.py index 164e2a73e6..b753e9fa8b 100644 --- a/litellm/litellm_core_utils/get_llm_provider_logic.py +++ b/litellm/litellm_core_utils/get_llm_provider_logic.py @@ -229,10 +229,10 @@ def get_llm_provider( # noqa: PLR0915 elif endpoint == "https://api.ai21.com/studio/v1": custom_llm_provider = "ai21_chat" dynamic_api_key = get_secret_str("AI21_API_KEY") - elif endpoint == "https://codestral.mistral.ai/v1": + elif endpoint == "codestral.mistral.ai/v1/chat/completions": custom_llm_provider = "codestral" dynamic_api_key = get_secret_str("CODESTRAL_API_KEY") - elif endpoint == "https://codestral.mistral.ai/v1": + elif endpoint == "codestral.mistral.ai/v1/fim/completions": custom_llm_provider = "text-completion-codestral" dynamic_api_key = get_secret_str("CODESTRAL_API_KEY") elif endpoint == "app.empower.dev/api/v1": diff --git a/tests/test_litellm/litellm_core_utils/test_codestral_provider_routing.py b/tests/test_litellm/litellm_core_utils/test_codestral_provider_routing.py new file mode 100644 index 0000000000..1a6ed51afd --- /dev/null +++ b/tests/test_litellm/litellm_core_utils/test_codestral_provider_routing.py @@ -0,0 +1,69 @@ +""" +Unit tests for codestral provider routing. + +These tests verify that the chat and FIM endpoints for codestral +are correctly routed to different providers: +- Chat endpoint -> codestral provider +- FIM endpoint -> text-completion-codestral provider + +Related issue: https://github.com/BerriAI/litellm/issues/18464 +""" +import pytest + +import litellm + + +class TestCodestralProviderRouting: + """Tests for codestral endpoint routing in get_llm_provider""" + + def test_codestral_chat_endpoint_routes_to_codestral_provider(self): + """ + Test that the codestral chat endpoint routes to the 'codestral' provider. + + The chat/completions endpoint should be handled by the codestral provider. + """ + model, custom_llm_provider, _, api_base = litellm.get_llm_provider( + model="codestral-latest", + api_base="https://codestral.mistral.ai/v1/chat/completions", + ) + + assert custom_llm_provider == "codestral" + + def test_codestral_fim_endpoint_routes_to_text_completion_provider(self): + """ + Test that the codestral FIM endpoint routes to 'text-completion-codestral'. + + The fim/completions endpoint should be handled by the + text-completion-codestral provider for fill-in-the-middle completions. + """ + model, custom_llm_provider, _, api_base = litellm.get_llm_provider( + model="codestral-latest", + api_base="https://codestral.mistral.ai/v1/fim/completions", + ) + + assert custom_llm_provider == "text-completion-codestral" + + def test_codestral_endpoints_are_different_providers(self): + """ + Test that chat and FIM endpoints route to different providers. + + This is the core fix for issue #18464 - previously both endpoints + would route to 'codestral' due to duplicate conditions. + """ + _, chat_provider, _, _ = litellm.get_llm_provider( + model="codestral-latest", + api_base="https://codestral.mistral.ai/v1/chat/completions", + ) + + _, fim_provider, _, _ = litellm.get_llm_provider( + model="codestral-latest", + api_base="https://codestral.mistral.ai/v1/fim/completions", + ) + + assert chat_provider != fim_provider + assert chat_provider == "codestral" + assert fim_provider == "text-completion-codestral" + + +if __name__ == "__main__": + pytest.main([__file__, "-v"])