scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-26 01:12:21 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 145

Avoid leaking Ollama metadata auth

Browse source
This commit is contained in:
Graham Neubig 2026-05-21 10:36:00 -04:00
parent aec84020e5
commit 82a4961afc
3 changed files with 134 additions and 8 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

31
litellm/llms/ollama/common_utils.py
View file

@ -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,

11
litellm/llms/ollama/completion/transformation.py
View file

@ -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]

100
tests/test_litellm/llms/ollama/test_ollama_model_info.py
View file

@ -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."""