From 82a4961afc384b74ef8e63025578823bea1213bf Mon Sep 17 00:00:00 2001 From: Graham Neubig <398875+neubig@users.noreply.github.com> Date: Thu, 21 May 2026 10:36:00 -0400 Subject: [PATCH] Avoid leaking Ollama metadata auth --- litellm/llms/ollama/common_utils.py | 31 ++++-- .../llms/ollama/completion/transformation.py | 11 +- .../llms/ollama/test_ollama_model_info.py | 100 ++++++++++++++++++ 3 files changed, 134 insertions(+), 8 deletions(-) diff --git a/litellm/llms/ollama/common_utils.py b/litellm/llms/ollama/common_utils.py index 4787e532fae..df54d1b1e79 100644 --- a/litellm/llms/ollama/common_utils.py +++ b/litellm/llms/ollama/common_utils.py @@ -65,7 +65,8 @@ class OllamaModelInfo(BaseLLMModelInfo): from litellm.secret_managers.main import get_secret_str return ( - os.environ.get("OLLAMA_API_KEY") + api_key + or os.environ.get("OLLAMA_API_KEY") or litellm.api_key or litellm.openai_key or get_secret_str("OLLAMA_API_KEY") @@ -98,8 +99,13 @@ class OllamaModelInfo(BaseLLMModelInfo): List all models available on the Ollama server via /api/tags endpoint. """ + passed_api_base = api_base base = self.get_server_api_base(api_base) - api_key = self.get_api_key() + api_key = ( + self.get_api_key(api_key) + if api_key is not None or passed_api_base is None + else None + ) headers = {"Authorization": f"Bearer {api_key}"} if api_key else {} names: set[str] = set() @@ -176,13 +182,21 @@ class OllamaModelInfo(BaseLLMModelInfo): return None def get_runtime_model_info( - self, model: str, api_base: Optional[str] = None + self, + model: str, + api_base: Optional[str] = None, + api_key: Optional[str] = None, ) -> dict[str, Any]: from litellm import module_level_client model = self._strip_ollama_model_prefix(model) + passed_api_base = api_base api_base = self.get_server_api_base(api_base) - api_key = self.get_api_key() + api_key = ( + self.get_api_key(api_key) + if api_key is not None or passed_api_base is None + else None + ) headers = {"Authorization": f"Bearer {api_key}"} if api_key else {} try: @@ -220,11 +234,16 @@ class OllamaModelInfo(BaseLLMModelInfo): } def get_model_info( - self, model: str, api_base: Optional[str] = None + self, + model: str, + api_base: Optional[str] = None, + api_key: Optional[str] = None, ) -> Optional[dict[str, Any]]: if self._is_static_ollama_model(model): return None - return self.get_runtime_model_info(model=model, api_base=api_base) + return self.get_runtime_model_info( + model=model, api_base=api_base, api_key=api_key + ) def validate_environment( self, diff --git a/litellm/llms/ollama/completion/transformation.py b/litellm/llms/ollama/completion/transformation.py index d3d3304e607..72abc2e79bb 100644 --- a/litellm/llms/ollama/completion/transformation.py +++ b/litellm/llms/ollama/completion/transformation.py @@ -221,13 +221,20 @@ class OllamaConfig(BaseConfig): or get_secret_str("OLLAMA_API_KEY") ) - def get_model_info(self, model: str, api_base: Optional[str] = None) -> Any: + def get_model_info( + self, + model: str, + api_base: Optional[str] = None, + api_key: Optional[str] = None, + ) -> Any: """ curl http://localhost:11434/api/show -d '{ "name": "mistral" }' """ - return OllamaModelInfo().get_runtime_model_info(model=model, api_base=api_base) + return OllamaModelInfo().get_runtime_model_info( + model=model, api_base=api_base, api_key=api_key + ) def get_error_class( self, error_message: str, status_code: int, headers: Union[dict, Headers] 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 7d92b6a19c4..d19055c55ae 100644 --- a/tests/test_litellm/llms/ollama/test_ollama_model_info.py +++ b/tests/test_litellm/llms/ollama/test_ollama_model_info.py @@ -105,6 +105,47 @@ class TestOllamaModelInfo: "Authorization": "Bearer test_api_key" } + def test_get_models_does_not_leak_server_key_to_provided_api_base( + self, monkeypatch + ): + """Model discovery should not send server-side keys to caller-supplied bases.""" + call_headers = [] + + def mock_get(url, headers): + call_headers.append(headers) + return DummyResponse({"models": []}, status_code=200) + + monkeypatch.setenv("OLLAMA_API_KEY", "server-side-ollama-key") + monkeypatch.setattr(litellm, "api_key", "global-provider-key") + monkeypatch.setattr(litellm, "openai_key", "global-openai-key") + monkeypatch.setattr(httpx, "get", mock_get) + + info = OllamaModelInfo() + models = info.get_models(api_base="https://attacker.example") + + assert models == [] + assert call_headers[0] == {} + + def test_get_models_uses_explicit_api_key_for_provided_api_base(self, monkeypatch): + """Model discovery should send an explicitly supplied key to the provided base.""" + call_headers = [] + + def mock_get(url, headers): + call_headers.append(headers) + return DummyResponse({"models": []}, status_code=200) + + monkeypatch.setenv("OLLAMA_API_KEY", "server-side-ollama-key") + monkeypatch.setattr(httpx, "get", mock_get) + + info = OllamaModelInfo() + models = info.get_models( + api_base="https://ollama.example", + api_key="explicit-api-key", + ) + + assert models == [] + assert call_headers[0] == {"Authorization": "Bearer explicit-api-key"} + def test_get_models_from_list_response(self, monkeypatch): """ When the /api/tags endpoint returns a list of dicts, @@ -201,18 +242,77 @@ class TestOllamaGetModelInfo: from litellm.llms.ollama.completion.transformation import OllamaConfig captured_urls = [] + captured_headers = [] def mock_post(url, json, headers=None): captured_urls.append(url) + captured_headers.append(headers) return DummyResponse({"template": "", "model_info": {}}, status_code=200) monkeypatch.setattr("litellm.module_level_client.post", mock_post) monkeypatch.setenv("OLLAMA_API_BASE", "http://env-server:11434") + monkeypatch.setenv("OLLAMA_API_KEY", "env-api-key") config = OllamaConfig() config.get_model_info("llama3") assert captured_urls[0] == "http://env-server:11434/api/show" + assert captured_headers[0] == {"Authorization": "Bearer env-api-key"} + + def test_get_model_info_uses_explicit_api_key_for_provided_api_base( + self, monkeypatch + ): + """When api_key is explicit, model info should send it to the provided api_base.""" + from litellm.llms.ollama.completion.transformation import OllamaConfig + + captured_headers = [] + + def mock_post(url, json, headers=None): + captured_headers.append(headers) + return DummyResponse({"template": "", "model_info": {}}, status_code=200) + + monkeypatch.setattr("litellm.module_level_client.post", mock_post) + + config = OllamaConfig() + config.get_model_info( + "llama3", + api_base="http://my-remote-server:11434", + api_key="explicit-api-key", + ) + + assert captured_headers[0] == {"Authorization": "Bearer explicit-api-key"} + + def test_litellm_get_model_info_does_not_leak_server_key_to_provided_api_base( + self, monkeypatch + ): + """Global model info should not send server-side keys to caller-supplied bases.""" + captured_headers = [] + + def mock_post(url, json, headers=None): + captured_headers.append(headers) + return DummyResponse( + { + "template": "{{ .System }} tools {{ .Prompt }}", + "model_info": {"llama.context_length": 32768}, + }, + status_code=200, + ) + + litellm.get_model_info.cache_clear() + monkeypatch.setattr("litellm.module_level_client.post", mock_post) + monkeypatch.setenv("OLLAMA_API_KEY", "server-side-ollama-key") + monkeypatch.setattr(litellm, "api_key", "global-provider-key") + monkeypatch.setattr(litellm, "openai_key", "global-openai-key") + try: + model_info = litellm.get_model_info( + "ollama/unknown-model", + api_base="https://attacker.example", + ) + finally: + litellm.get_model_info.cache_clear() + + assert model_info["max_input_tokens"] == 32768 + assert captured_headers[0] == {} def test_get_model_info_normalizes_generate_api_base(self, monkeypatch): """When completion passes the final generate URL, model info should use the server base."""