diff --git a/litellm/constants.py b/litellm/constants.py index 9235916dd433..1215b1bd2019 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -519,6 +519,7 @@ "https://api.friendli.ai/serverless/v1", "api.sambanova.ai/v1", "api.x.ai/v1", + "ollama.com", "api.galadriel.ai/v1", "api.llama.com/compat/v1/", "api.featherless.ai/v1", diff --git a/litellm/litellm_core_utils/get_llm_provider_logic.py b/litellm/litellm_core_utils/get_llm_provider_logic.py index 1efea63beb7d..a4e576c7718a 100644 --- a/litellm/litellm_core_utils/get_llm_provider_logic.py +++ b/litellm/litellm_core_utils/get_llm_provider_logic.py @@ -235,6 +235,9 @@ def get_llm_provider( # noqa: PLR0915 elif endpoint == "api.deepseek.com/v1": custom_llm_provider = "deepseek" dynamic_api_key = get_secret_str("DEEPSEEK_API_KEY") + elif endpoint == "ollama.com": + custom_llm_provider = "ollama" + dynamic_api_key = get_secret_str("OLLAMA_API_KEY") elif endpoint == "https://api.friendli.ai/serverless/v1": custom_llm_provider = "friendliai" dynamic_api_key = get_secret_str( @@ -547,6 +550,13 @@ def _get_openai_compatible_provider_info( # noqa: PLR0915 or "https://api.studio.nebius.ai/v1" ) # type: ignore dynamic_api_key = api_key or get_secret_str("NEBIUS_API_KEY") + elif custom_llm_provider == "ollama": + api_base = ( + api_base + or get_secret("OLLAMA_API_BASE") + or "http://127.0.0.1:11434" + ) # type: ignore + dynamic_api_key = api_key or get_secret_str("OLLAMA_API_KEY") elif (custom_llm_provider == "ai21_chat") or ( custom_llm_provider == "ai21" and model in litellm.ai21_chat_models ): diff --git a/litellm/main.py b/litellm/main.py index 3e51785ac637..36e0d52da3f0 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -3439,6 +3439,9 @@ def completion( # type: ignore # noqa: PLR0915 or get_secret("OLLAMA_API_BASE") or "http://localhost:11434" ) + if api_key is not None and "Authorization" not in headers: + headers["Authorization"] = f"Bearer {api_key}" + response = base_llm_http_handler.completion( model=model, stream=stream, @@ -3472,6 +3475,9 @@ def completion( # type: ignore # noqa: PLR0915 or os.environ.get("OLLAMA_API_KEY") or litellm.api_key ) + if api_key is not None and "Authorization" not in headers: + headers["Authorization"] = f"Bearer {api_key}" + response = base_llm_http_handler.completion( model=model, diff --git a/tests/test_litellm/llms/ollama/test_ollama_model_info.py b/tests/test_litellm/llms/ollama/test_ollama_model_info.py index b0aa464bf9df..7eef15cd4d29 100644 --- a/tests/test_litellm/llms/ollama/test_ollama_model_info.py +++ b/tests/test_litellm/llms/ollama/test_ollama_model_info.py @@ -136,3 +136,327 @@ def mock_get(url, headers): models = info.get_models() # Default static ollama_models is ['llama2'], so expect ['ollama/llama2'] assert models == ["ollama/llama2"] + + +class TestOllamaAuthHeaders: + """Tests for Ollama authentication header handling in completion calls.""" + + def test_ollama_completion_with_api_key_adds_auth_header(self, monkeypatch): + """ + Test that when an api_key is provided to ollama completion, + the Authorization header is added with Bearer token format. + + This tests the bug fix where Ollama requests with API keys + were not including the Authorization header. + """ + import litellm + from unittest.mock import MagicMock, patch + + # Track the headers that were passed to the completion call + captured_headers = {} + + def mock_completion(*args, **kwargs): + # Capture the headers that were passed + if 'headers' in kwargs: + captured_headers.update(kwargs['headers']) + # Return a mock response + mock_response = MagicMock() + mock_response.choices = [MagicMock()] + mock_response.choices[0].message = MagicMock() + mock_response.choices[0].message.content = "Test response" + return mock_response + + # Mock the base_llm_http_handler.completion method at the module level + with patch('litellm.main.base_llm_http_handler.completion', side_effect=mock_completion): + try: + # Call completion with ollama provider and api_key + litellm.completion( + model="ollama/llama2", + messages=[{"role": "user", "content": "Hello"}], + api_key="test-api-key-12345", + api_base="http://localhost:11434" + ) + + # Verify that Authorization header was added + assert "Authorization" in captured_headers, \ + "Authorization header should be present when api_key is provided" + assert captured_headers["Authorization"] == "Bearer test-api-key-12345", \ + f"Authorization header should be 'Bearer test-api-key-12345', got {captured_headers.get('Authorization')}" + + except Exception as e: + pytest.fail(f"Ollama completion with api_key failed: {e}") + + def test_ollama_chat_completion_with_api_key_adds_auth_header(self, monkeypatch): + """ + Test that when an api_key is provided to ollama_chat completion, + the Authorization header is added with Bearer token format. + + This tests the bug fix for the ollama_chat provider variant. + """ + import litellm + from unittest.mock import MagicMock, patch + + # Track the headers that were passed to the completion call + captured_headers = {} + + def mock_completion(*args, **kwargs): + # Capture the headers that were passed + if 'headers' in kwargs: + captured_headers.update(kwargs['headers']) + # Return a mock response + mock_response = MagicMock() + mock_response.choices = [MagicMock()] + mock_response.choices[0].message = MagicMock() + mock_response.choices[0].message.content = "Test response" + return mock_response + + # Mock the base_llm_http_handler.completion method at the module level + with patch('litellm.main.base_llm_http_handler.completion', side_effect=mock_completion): + try: + # Call completion with ollama_chat provider and api_key + litellm.completion( + model="ollama_chat/llama2", + messages=[{"role": "user", "content": "Hello"}], + api_key="test-api-key-67890", + api_base="http://localhost:11434" + ) + + # Verify that Authorization header was added + assert "Authorization" in captured_headers, \ + "Authorization header should be present when api_key is provided" + assert captured_headers["Authorization"] == "Bearer test-api-key-67890", \ + f"Authorization header should be 'Bearer test-api-key-67890', got {captured_headers.get('Authorization')}" + + except Exception as e: + pytest.fail(f"Ollama_chat completion with api_key failed: {e}") + + def test_ollama_completion_without_api_key_no_auth_header(self, monkeypatch): + """ + Test that when no api_key is provided to ollama completion, + no Authorization header is added. + """ + import litellm + from unittest.mock import MagicMock, patch + + # Track the headers that were passed to the completion call + captured_headers = {} + + def mock_completion(*args, **kwargs): + # Capture the headers that were passed + if 'headers' in kwargs: + captured_headers.update(kwargs['headers']) + # Return a mock response + mock_response = MagicMock() + mock_response.choices = [MagicMock()] + mock_response.choices[0].message = MagicMock() + mock_response.choices[0].message.content = "Test response" + return mock_response + + # Mock the base_llm_http_handler.completion method at the module level + with patch('litellm.main.base_llm_http_handler.completion', side_effect=mock_completion): + try: + # Call completion without api_key + litellm.completion( + model="ollama/llama2", + messages=[{"role": "user", "content": "Hello"}], + api_base="http://localhost:11434" + ) + + # Verify that Authorization header was NOT added + assert "Authorization" not in captured_headers, \ + "Authorization header should not be present when api_key is not provided" + + except Exception as e: + pytest.fail(f"Ollama completion without api_key failed: {e}") + + def test_ollama_completion_preserves_existing_auth_header(self, monkeypatch): + """ + Test that when an Authorization header is already present in headers, + it is not overwritten even if api_key is provided. + + This ensures the fix respects existing Authorization headers. + """ + import litellm + from unittest.mock import MagicMock, patch + + # Track the headers that were passed to the completion call + captured_headers = {} + + def mock_completion(*args, **kwargs): + # Capture the headers that were passed + if 'headers' in kwargs: + captured_headers.update(kwargs['headers']) + # Return a mock response + mock_response = MagicMock() + mock_response.choices = [MagicMock()] + mock_response.choices[0].message = MagicMock() + mock_response.choices[0].message.content = "Test response" + return mock_response + + # Mock the base_llm_http_handler.completion method at the module level + with patch('litellm.main.base_llm_http_handler.completion', side_effect=mock_completion): + try: + # Call completion with both api_key and existing Authorization header + existing_auth = "Bearer existing-token" + litellm.completion( + model="ollama/llama2", + messages=[{"role": "user", "content": "Hello"}], + api_key="test-api-key-should-not-be-used", + api_base="http://localhost:11434", + headers={"Authorization": existing_auth} + ) + + # Verify that existing Authorization header was preserved + assert "Authorization" in captured_headers, \ + "Authorization header should be present" + assert captured_headers["Authorization"] == existing_auth, \ + f"Existing Authorization header should be preserved, got {captured_headers.get('Authorization')}" + + except Exception as e: + pytest.fail(f"Ollama completion with existing auth header failed: {e}") + + def test_ollama_completion_with_ollama_com_api_base(self, monkeypatch): + """ + Test that when using https://ollama.com as api_base with an api_key, + the Authorization header is correctly added. + + This tests the real-world use case of using Ollama's hosted service. + """ + import litellm + from unittest.mock import MagicMock, patch + + # Track the headers and api_base that were passed to the completion call + captured_headers = {} + captured_api_base = None + + def mock_completion(*args, **kwargs): + nonlocal captured_api_base + # Capture the headers and api_base that were passed + if 'headers' in kwargs: + captured_headers.update(kwargs['headers']) + if 'api_base' in kwargs: + captured_api_base = kwargs['api_base'] + # Return a mock response + mock_response = MagicMock() + mock_response.choices = [MagicMock()] + mock_response.choices[0].message = MagicMock() + mock_response.choices[0].message.content = "Test response" + return mock_response + + # Mock the base_llm_http_handler.completion method at the module level + with patch('litellm.main.base_llm_http_handler.completion', side_effect=mock_completion): + try: + # Call completion with ollama.com as api_base and api_key + litellm.completion( + model="ollama/qwen3-vl:235b-cloud", + messages=[{"role": "user", "content": "Hello"}], + api_key="test-ollama-com-api-key", + api_base="https://ollama.com" + ) + + # Verify that Authorization header was added + assert "Authorization" in captured_headers, \ + "Authorization header should be present when using ollama.com with api_key" + assert captured_headers["Authorization"] == "Bearer test-ollama-com-api-key", \ + f"Authorization header should be 'Bearer test-ollama-com-api-key', got {captured_headers.get('Authorization')}" + + # Verify the api_base was passed correctly + assert captured_api_base == "https://ollama.com", \ + f"API base should be 'https://ollama.com', got {captured_api_base}" + + except Exception as e: + pytest.fail(f"Ollama completion with ollama.com api_base failed: {e}") + + def test_ollama_chat_completion_with_ollama_com_api_base(self, monkeypatch): + """ + Test that when using https://ollama.com as api_base with an api_key + for ollama_chat provider, the Authorization header is correctly added. + + This tests the real-world use case for the ollama_chat variant. + """ + import litellm + from unittest.mock import MagicMock, patch + + # Track the headers and api_base that were passed to the completion call + captured_headers = {} + captured_api_base = None + + def mock_completion(*args, **kwargs): + nonlocal captured_api_base + # Capture the headers and api_base that were passed + if 'headers' in kwargs: + captured_headers.update(kwargs['headers']) + if 'api_base' in kwargs: + captured_api_base = kwargs['api_base'] + # Return a mock response + mock_response = MagicMock() + mock_response.choices = [MagicMock()] + mock_response.choices[0].message = MagicMock() + mock_response.choices[0].message.content = "Test response" + return mock_response + + # Mock the base_llm_http_handler.completion method at the module level + with patch('litellm.main.base_llm_http_handler.completion', side_effect=mock_completion): + try: + # Call completion with ollama.com as api_base and api_key + litellm.completion( + model="ollama_chat/qwen3-vl:235b-cloud", + messages=[{"role": "user", "content": "Hello"}], + api_key="test-ollama-com-chat-key", + api_base="https://ollama.com" + ) + + # Verify that Authorization header was added + assert "Authorization" in captured_headers, \ + "Authorization header should be present when using ollama.com with api_key" + assert captured_headers["Authorization"] == "Bearer test-ollama-com-chat-key", \ + f"Authorization header should be 'Bearer test-ollama-com-chat-key', got {captured_headers.get('Authorization')}" + + # Verify the api_base was passed correctly + assert captured_api_base == "https://ollama.com", \ + f"API base should be 'https://ollama.com', got {captured_api_base}" + + except Exception as e: + pytest.fail(f"Ollama_chat completion with ollama.com api_base failed: {e}") + + def test_ollama_completion_with_ollama_com_without_api_key_fails_gracefully(self, monkeypatch): + """ + Test that when using https://ollama.com as api_base without an api_key, + no Authorization header is added (which would likely fail on the server side, + but we're testing the client behavior). + + This ensures we don't add empty or None Authorization headers. + """ + import litellm + from unittest.mock import MagicMock, patch + + # Track the headers that were passed to the completion call + captured_headers = {} + + def mock_completion(*args, **kwargs): + # Capture the headers that were passed + if 'headers' in kwargs: + captured_headers.update(kwargs['headers']) + # Return a mock response + mock_response = MagicMock() + mock_response.choices = [MagicMock()] + mock_response.choices[0].message = MagicMock() + mock_response.choices[0].message.content = "Test response" + return mock_response + + # Mock the base_llm_http_handler.completion method at the module level + with patch('litellm.main.base_llm_http_handler.completion', side_effect=mock_completion): + try: + # Call completion with ollama.com but no api_key + litellm.completion( + model="ollama/llama2", + messages=[{"role": "user", "content": "Hello"}], + api_base="https://ollama.com" + ) + + # Verify that Authorization header was NOT added + assert "Authorization" not in captured_headers, \ + "Authorization header should not be present when api_key is not provided, even with ollama.com" + + except Exception as e: + pytest.fail(f"Ollama completion with ollama.com without api_key failed: {e}")