scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-25 01:02:15 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 80

Avoid leaking Lemonade metadata auth

Browse source
This commit is contained in:
Graham Neubig 2026-05-21 12:40:26 -04:00
parent 82a4961afc
commit 8829cbaf9d
2 changed files with 134 additions and 11 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

24
litellm/llms/lemonade/chat/transformation.py
View file

@ -164,12 +164,17 @@ class LemonadeChatConfig(OpenAILikeChatConfig):
"max_output_tokens": None,
}
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:
if model.startswith("lemonade/"):
model = model.split("/", 1)[1]
api_base, api_key = self._get_openai_compatible_provider_info(
api_base=api_base, api_key=None
api_base=api_base, api_key=api_key
)
encoded_model = quote(model, safe="")
@ -205,17 +210,20 @@ class LemonadeChatConfig(OpenAILikeChatConfig):
self, api_base: Optional[str], api_key: Optional[str]
) -> Tuple[Optional[str], Optional[str]]:
# lemonade is openai compatible, we just need to set this to custom_openai and have the api_base be lemonade's endpoint
passed_api_base = api_base
api_base = (
api_base
or get_secret_str("LEMONADE_API_BASE")
or "http://localhost:8000/api/v1"
) # type: ignore
key = (
api_key
or litellm.lemonade_key
or get_secret_str("LEMONADE_API_KEY")
or self._DEFAULT_API_KEY
)
key = self._DEFAULT_API_KEY
if api_key is not None or passed_api_base is None:
key = (
api_key
or litellm.lemonade_key
or get_secret_str("LEMONADE_API_KEY")
or self._DEFAULT_API_KEY
)
return api_base, key
def _get_auth_headers(self, api_key: Optional[str]) -> dict:

121
tests/test_litellm/llms/lemonade/test_lemonade.py
View file

@ -71,6 +71,44 @@ def test_get_openai_compatible_provider_info_with_api_key_env(monkeypatch):
assert key == "test-key"
def test_get_openai_compatible_provider_info_skips_env_key_for_custom_base(
monkeypatch,
):
"""Test that caller-supplied bases do not receive server-side Lemonade keys."""
monkeypatch.setenv("LEMONADE_API_KEY", "server-side-lemonade-key")
monkeypatch.setattr(litellm, "lemonade_key", "configured-lemonade-key")
monkeypatch.setattr(litellm, "api_key", None)
config = LemonadeChatConfig()
api_base, key = config._get_openai_compatible_provider_info(
api_base="https://attacker.example/v1", api_key=None
)
assert api_base == "https://attacker.example/v1"
assert key == "lemonade"
assert config._get_auth_headers(key) == {}
def test_get_openai_compatible_provider_info_uses_explicit_key_for_custom_base(
monkeypatch,
):
"""Test that explicitly supplied Lemonade keys are sent to supplied bases."""
monkeypatch.setenv("LEMONADE_API_KEY", "server-side-lemonade-key")
monkeypatch.setattr(litellm, "lemonade_key", "configured-lemonade-key")
monkeypatch.setattr(litellm, "api_key", None)
config = LemonadeChatConfig()
api_base, key = config._get_openai_compatible_provider_info(
api_base="https://lemonade.example/v1", api_key="explicit-lemonade-key"
)
assert api_base == "https://lemonade.example/v1"
assert key == "explicit-lemonade-key"
assert config._get_auth_headers(key) == {
"Authorization": "Bearer explicit-lemonade-key"
}
def test_get_openai_compatible_provider_info_ignores_global_api_key(monkeypatch):
"""Test that Lemonade discovery does not send unrelated global API keys."""
monkeypatch.delenv("LEMONADE_API_KEY", raising=False)
@ -87,6 +125,25 @@ def test_get_openai_compatible_provider_info_ignores_global_api_key(monkeypatch)
assert config._get_auth_headers(key) == {}
def test_get_models_does_not_leak_lemonade_key_to_custom_base(monkeypatch):
"""Test Lemonade discovery does not send server-side keys to supplied bases."""
monkeypatch.setenv("LEMONADE_API_KEY", "server-side-lemonade-key")
monkeypatch.setattr(litellm, "lemonade_key", "configured-lemonade-key")
monkeypatch.setattr(litellm, "api_key", "global-provider-key")
config = LemonadeChatConfig()
response = MagicMock()
response.status_code = 200
response.json.return_value = {"data": []}
with patch.object(
litellm.module_level_client, "get", return_value=response
) as mock_get:
models = config.get_models(api_base="https://attacker.example/v1")
assert models == []
assert mock_get.call_args.kwargs["headers"] == {}
def test_get_model_info_uses_loaded_context_size():
"""Test that Lemonade model info prefers the effective loaded ctx_size."""
config = LemonadeChatConfig()
@ -170,9 +227,33 @@ def test_get_model_info_reads_context_from_provider_specific_entry():
}
def test_get_model_info_sends_lemonade_api_key(monkeypatch):
"""Test that Lemonade model info uses auth for authenticated servers."""
def test_get_model_info_sends_lemonade_api_key_for_configured_base(monkeypatch):
"""Test that Lemonade model info uses auth for configured servers."""
monkeypatch.setenv("LEMONADE_API_KEY", "test-key")
monkeypatch.setenv("LEMONADE_API_BASE", "http://lemonade.test/v1")
monkeypatch.setattr(litellm, "lemonade_key", None)
monkeypatch.setattr(litellm, "api_key", None)
config = LemonadeChatConfig()
response = MagicMock()
response.status_code = 200
response.json.return_value = {
"id": "Qwen3.6-35B-A3B-GGUF",
"recipe_options": {"ctx_size": 65536},
}
with patch.object(
litellm.module_level_client, "get", return_value=response
) as mock_get:
config.get_model_info(
model="lemonade/Qwen3.6-35B-A3B-GGUF",
)
assert mock_get.call_args.kwargs["headers"] == {"Authorization": "Bearer test-key"}
def test_get_model_info_sends_explicit_lemonade_api_key_for_custom_base(monkeypatch):
"""Test that Lemonade model info sends explicitly supplied auth to supplied bases."""
monkeypatch.setenv("LEMONADE_API_KEY", "server-side-key")
monkeypatch.setattr(litellm, "lemonade_key", None)
monkeypatch.setattr(litellm, "api_key", None)
config = LemonadeChatConfig()
@ -189,9 +270,43 @@ def test_get_model_info_sends_lemonade_api_key(monkeypatch):
config.get_model_info(
model="lemonade/Qwen3.6-35B-A3B-GGUF",
api_base="http://lemonade.test/v1",
api_key="explicit-test-key",
)
assert mock_get.call_args.kwargs["headers"] == {"Authorization": "Bearer test-key"}
assert mock_get.call_args.kwargs["headers"] == {
"Authorization": "Bearer explicit-test-key"
}
def test_litellm_get_model_info_does_not_leak_lemonade_key_to_custom_base(
monkeypatch,
):
"""Test top-level model info does not send server-side keys to supplied bases."""
monkeypatch.setenv("LEMONADE_API_KEY", "server-side-lemonade-key")
monkeypatch.setattr(litellm, "lemonade_key", "configured-lemonade-key")
monkeypatch.setattr(litellm, "api_key", "global-provider-key")
response = MagicMock()
response.status_code = 200
response.json.return_value = {
"id": "Qwen3.6-35B-A3B-GGUF",
"max_input_tokens": 65536,
"max_context_window": 262144,
}
litellm.get_model_info.cache_clear()
with patch.object(
litellm.module_level_client, "get", return_value=response
) as mock_get:
try:
model_info = litellm.get_model_info(
model="lemonade/Qwen3.6-35B-A3B-GGUF",
api_base="https://attacker.example/v1",
)
finally:
litellm.get_model_info.cache_clear()
assert model_info["max_input_tokens"] == 65536
assert mock_get.call_args.kwargs["headers"] == {}
def test_litellm_get_model_info_uses_lemonade_api_base():