mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-04 07:27:51 +00:00
fix(proxy): let a config license unlock enterprise settings and spell "environment" right in license errors (#49745)
* fix(proxy): spell "environment" right in the allowed_ips and AWS KMS V2 license errors * fix(proxy): read general_settings.litellm_license before the allowed_ips and enforced_params gates * fix(proxy): read the config license before callbacks and gates, keeping a valid env license * fix(proxy): keep an invalid config license reported and sync license data with the license in use A config license that fails no longer overwrites the license string when the env license is not valid either, so /health/license still reports has_license. A failed boot-time env license is not re-checked. Both license writes go through one helper that also refreshes premium_user_data, so /health/license and the seat limits follow the license in use. * fix(proxy): keep the env license's boot result when a config license fails, without a second license check --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
This commit is contained in:
paveln83
GitHub
parent
28b4300be2
commit
7bd342d7ea
5 changed files with 261 additions and 11 deletions
|
|
@ -1185,6 +1185,27 @@ server_root_path: Final = get_server_root_path()
|
|||
_license_check = LicenseCheck()
|
||||
premium_user: bool = _license_check.is_premium()
|
||||
premium_user_data: Optional["EnterpriseLicenseData"] = _license_check.airgapped_license_data
|
||||
|
||||
|
||||
def _set_premium_user(is_premium: bool) -> None:
|
||||
global premium_user, premium_user_data
|
||||
premium_user = is_premium # rebind-ok: config load can switch the license in use
|
||||
premium_user_data = _license_check.airgapped_license_data # rebind-ok: tracks the license in use
|
||||
|
||||
|
||||
def _use_config_license(config_license: str, env_license_is_premium: bool) -> bool:
|
||||
env_license: Final = _license_check.license_str
|
||||
env_license_data: Final = _license_check.airgapped_license_data
|
||||
_license_check.license_str = config_license
|
||||
if _license_check.is_premium():
|
||||
return True
|
||||
if not env_license_is_premium:
|
||||
return False
|
||||
_license_check.license_str = env_license
|
||||
_license_check.airgapped_license_data = env_license_data
|
||||
return True
|
||||
|
||||
|
||||
global_max_parallel_request_retries_env: Final[str | None] = os.getenv("LITELLM_GLOBAL_MAX_PARALLEL_REQUEST_RETRIES")
|
||||
proxy_state: Final = ProxyState()
|
||||
SENSITIVE_DATA_MASKER: Final = SensitiveDataMasker()
|
||||
|
|
@ -6686,7 +6707,6 @@ class ProxyConfig:
|
|||
|
||||
def _load_environment_variables(self, config: dict):
|
||||
## ENVIRONMENT VARIABLES
|
||||
global premium_user
|
||||
environment_variables: Final = config.get("environment_variables", None)
|
||||
if environment_variables:
|
||||
for key, value in environment_variables.items():
|
||||
|
|
@ -6726,7 +6746,7 @@ class ProxyConfig:
|
|||
# check if litellm_license in general_settings
|
||||
if "LITELLM_LICENSE" in environment_variables:
|
||||
_license_check.license_str = os.getenv("LITELLM_LICENSE", None)
|
||||
premium_user = _license_check.is_premium()
|
||||
_set_premium_user(_license_check.is_premium())
|
||||
|
||||
def _set_config_environment_variable(self, key: str, value: str) -> None:
|
||||
if key in os.environ and os.environ[key] != value:
|
||||
|
|
@ -6778,7 +6798,6 @@ class ProxyConfig:
|
|||
prompt_injection_detection_obj, \
|
||||
redis_usage_cache, \
|
||||
store_model_in_db, \
|
||||
premium_user, \
|
||||
open_telemetry_logger, \
|
||||
health_check_details, \
|
||||
proxy_batch_polling_interval, \
|
||||
|
|
@ -6791,6 +6810,12 @@ class ProxyConfig:
|
|||
|
||||
self._load_environment_variables(config=config)
|
||||
|
||||
config_license: Final = _GENERAL_SETTINGS_VIEW.validate_python(
|
||||
_GENERAL_SETTINGS_VIEW.validate_python(config).get("general_settings") or {}
|
||||
).get("litellm_license")
|
||||
if isinstance(config_license, str):
|
||||
_set_premium_user(_use_config_license(config_license, env_license_is_premium=premium_user))
|
||||
|
||||
## Coordination Redis (before cache init, so the explicit block wins)
|
||||
coordination_redis_cache: Final = self._init_coordination_redis(config=config)
|
||||
if coordination_redis_cache is not None:
|
||||
|
|
@ -7367,7 +7392,7 @@ class ProxyConfig:
|
|||
allowed_ips: Final = general_settings.get("allowed_ips", None)
|
||||
if allowed_ips is not None and premium_user is False:
|
||||
raise ValueError(
|
||||
"allowed_ips is an Enterprise Feature. Please add a valid LITELLM_LICENSE to your envionment."
|
||||
"allowed_ips is an Enterprise Feature. Please add a valid LITELLM_LICENSE to your environment."
|
||||
)
|
||||
## BUDGET RESCHEDULER ##
|
||||
proxy_budget_rescheduler_min_time = general_settings.get(
|
||||
|
|
@ -7423,11 +7448,6 @@ class ProxyConfig:
|
|||
if general_settings.get("enforced_params") is not None and premium_user is not True:
|
||||
raise ValueError("Trying to use `enforced_params`" + CommonProxyErrors.not_premium_user.value)
|
||||
|
||||
# check if litellm_license in general_settings
|
||||
if "litellm_license" in general_settings:
|
||||
_license_check.license_str = general_settings["litellm_license"]
|
||||
premium_user = _license_check.is_premium()
|
||||
|
||||
router_params: Final[dict] = {
|
||||
"cache_responses": litellm.cache is not None, # cache if user passed in cache values
|
||||
}
|
||||
|
|
|
|||
|
|
@ -70,7 +70,7 @@ class AWSKeyManagementService_V2:
|
|||
is_litellm_license_in_env = True
|
||||
if is_litellm_license_in_env is False:
|
||||
raise ValueError(
|
||||
"AWSKeyManagementService V2 is an Enterprise Feature. Please add a valid LITELLM_LICENSE to your envionment."
|
||||
"AWSKeyManagementService V2 is an Enterprise Feature. Please add a valid LITELLM_LICENSE to your environment."
|
||||
)
|
||||
|
||||
def load_aws_kms(self, use_aws_kms: bool | None):
|
||||
|
|
|
|||
|
|
@ -9,6 +9,7 @@ Pins covered:
|
|||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import base64
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
|
|
@ -25,11 +26,15 @@ from unittest.mock import AsyncMock, MagicMock
|
|||
|
||||
import httpx
|
||||
import pytest
|
||||
from cryptography.hazmat.primitives import hashes
|
||||
from cryptography.hazmat.primitives.asymmetric import padding, rsa
|
||||
from cryptography.hazmat.primitives.asymmetric.rsa import RSAPrivateKey, RSAPublicKey
|
||||
from prisma.errors import DataError
|
||||
from pydantic import JsonValue, TypeAdapter, ValidationError
|
||||
|
||||
import litellm
|
||||
from litellm._logging import turn_on_json
|
||||
from litellm.llms.custom_httpx.http_handler import HTTPHandler
|
||||
from litellm.proxy._types import CommonProxyErrors, ConfigGeneralSettings
|
||||
from litellm.proxy.common_utils.encrypt_decrypt_utils import encrypt_value_helper
|
||||
from litellm.proxy.config_resolvers.settings_store import SettingsStore
|
||||
|
|
@ -2671,6 +2676,213 @@ async def test_load_config_logs_disabled_budget_reservation_once(tmp_path, monke
|
|||
assert [record.levelno for record in records] == ([logging.INFO] if setting == "true" else [])
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_load_config_refuses_allowed_ips_without_a_license(tmp_path: Path, monkeypatch: pytest.MonkeyPatch):
|
||||
config_file: Final = tmp_path / "allowed_ips.yaml"
|
||||
config_file.write_text('model_list: []\nlitellm_settings: {}\ngeneral_settings:\n allowed_ips: ["127.0.0.1"]\n')
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", None)
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.store_model_in_db", False)
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.premium_user", False)
|
||||
monkeypatch.delenv("LITELLM_CONFIG_BUCKET_NAME", raising=False)
|
||||
|
||||
expected_message: Final = (
|
||||
"allowed_ips is an Enterprise Feature. Please add a valid LITELLM_LICENSE to your environment."
|
||||
)
|
||||
with pytest.raises(ValueError, match=f"^{re.escape(expected_message)}$"):
|
||||
await ProxyConfig().load_config(router=None, config_file_path=str(config_file))
|
||||
|
||||
|
||||
def _signed_test_license(private_key: RSAPrivateKey, max_users: int) -> str:
|
||||
message: Final = json.dumps({"expiration_date": "9999-12-31", "user_id": "test", "max_users": max_users}).encode()
|
||||
signature: Final = private_key.sign(
|
||||
message,
|
||||
padding.PSS(mgf=padding.MGF1(hashes.SHA256()), salt_length=padding.PSS.MAX_LENGTH),
|
||||
hashes.SHA256(),
|
||||
)
|
||||
return base64.b64encode(message + b"." + signature).decode()
|
||||
|
||||
|
||||
def _write_config(path: Path, general_settings: dict[str, object], litellm_settings: dict[str, object]) -> str:
|
||||
path.write_text(
|
||||
json.dumps({"model_list": [], "litellm_settings": litellm_settings, "general_settings": general_settings})
|
||||
)
|
||||
return str(path)
|
||||
|
||||
|
||||
def _boot_proxy(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
public_key: RSAPublicKey,
|
||||
env_license: str | None,
|
||||
license_server_urls: list[str],
|
||||
server_accepted_license: str | None = None,
|
||||
) -> None:
|
||||
from litellm.proxy import proxy_server
|
||||
|
||||
def _license_server(request: httpx.Request) -> httpx.Response:
|
||||
license_server_urls.append(str(request.url))
|
||||
return httpx.Response(
|
||||
200, json={"verify": server_accepted_license is not None and server_accepted_license in str(request.url)}
|
||||
)
|
||||
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", None)
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.store_model_in_db", False)
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server._license_check.public_key", public_key)
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server._license_check.airgapped_license_data", None)
|
||||
monkeypatch.setattr(
|
||||
"litellm.proxy.proxy_server._license_check.http_handler",
|
||||
HTTPHandler(client=httpx.Client(transport=httpx.MockTransport(_license_server))),
|
||||
)
|
||||
monkeypatch.delenv("LITELLM_CONFIG_BUCKET_NAME", raising=False)
|
||||
if env_license is None:
|
||||
monkeypatch.delenv("LITELLM_LICENSE", raising=False)
|
||||
else:
|
||||
monkeypatch.setenv("LITELLM_LICENSE", env_license)
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server._license_check.license_str", env_license)
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.premium_user", proxy_server._license_check.is_premium())
|
||||
monkeypatch.setattr(
|
||||
"litellm.proxy.proxy_server.premium_user_data", proxy_server._license_check.airgapped_license_data
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
"gated_config",
|
||||
[
|
||||
{"general_settings": {"allowed_ips": ["127.0.0.1"]}},
|
||||
{"general_settings": {"enforced_params": ["user"]}},
|
||||
{"litellm_settings": {"callbacks": ["openai_moderations"]}},
|
||||
],
|
||||
)
|
||||
async def test_load_config_lets_a_config_license_unlock_its_enterprise_gates(
|
||||
tmp_path: Path, monkeypatch: pytest.MonkeyPatch, gated_config: dict[str, dict[str, list[str]]]
|
||||
):
|
||||
from litellm.proxy import proxy_server
|
||||
|
||||
private_key: Final = rsa.generate_private_key(public_exponent=65537, key_size=2048)
|
||||
license_key: Final = _signed_test_license(private_key, max_users=10)
|
||||
config_file: Final = _write_config(
|
||||
tmp_path / "licensed.yaml",
|
||||
{"litellm_license": license_key, **gated_config.get("general_settings", {})},
|
||||
dict(gated_config.get("litellm_settings", {})),
|
||||
)
|
||||
_boot_proxy(monkeypatch, private_key.public_key(), env_license=None, license_server_urls=[])
|
||||
|
||||
await ProxyConfig().load_config(router=None, config_file_path=config_file)
|
||||
|
||||
assert proxy_server.premium_user is True
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_load_config_keeps_a_valid_env_license_when_the_config_license_is_invalid(
|
||||
tmp_path: Path, monkeypatch: pytest.MonkeyPatch
|
||||
):
|
||||
from litellm.proxy import proxy_server
|
||||
|
||||
private_key: Final = rsa.generate_private_key(public_exponent=65537, key_size=2048)
|
||||
config_file: Final = _write_config(
|
||||
tmp_path / "stale_license.yaml", {"litellm_license": "not-a-valid-license", "allowed_ips": ["127.0.0.1"]}, {}
|
||||
)
|
||||
_boot_proxy(
|
||||
monkeypatch,
|
||||
private_key.public_key(),
|
||||
env_license=_signed_test_license(private_key, max_users=5),
|
||||
license_server_urls=[],
|
||||
)
|
||||
|
||||
await ProxyConfig().load_config(router=None, config_file_path=config_file)
|
||||
|
||||
assert proxy_server.premium_user is True
|
||||
assert proxy_server._license_check.user_seat_limit() == 5
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_load_config_prefers_the_config_license_when_both_licenses_are_valid(
|
||||
tmp_path: Path, monkeypatch: pytest.MonkeyPatch
|
||||
):
|
||||
from litellm.proxy import proxy_server
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
from litellm.proxy.health_endpoints._health_endpoints import health_license_endpoint
|
||||
|
||||
private_key: Final = rsa.generate_private_key(public_exponent=65537, key_size=2048)
|
||||
config_file: Final = _write_config(
|
||||
tmp_path / "licensed.yaml", {"litellm_license": _signed_test_license(private_key, max_users=100)}, {}
|
||||
)
|
||||
_boot_proxy(
|
||||
monkeypatch,
|
||||
private_key.public_key(),
|
||||
env_license=_signed_test_license(private_key, max_users=5),
|
||||
license_server_urls=[],
|
||||
)
|
||||
|
||||
await ProxyConfig().load_config(router=None, config_file_path=config_file)
|
||||
|
||||
assert proxy_server._license_check.user_seat_limit() == 100
|
||||
assert (await health_license_endpoint(user_api_key_dict=UserAPIKeyAuth()))["limits"]["max_users"] == 100
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_load_config_still_reports_an_invalid_config_license_as_configured(
|
||||
tmp_path: Path, monkeypatch: pytest.MonkeyPatch
|
||||
):
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
from litellm.proxy.health_endpoints._health_endpoints import health_license_endpoint
|
||||
|
||||
private_key: Final = rsa.generate_private_key(public_exponent=65537, key_size=2048)
|
||||
config_file: Final = _write_config(tmp_path / "stale_license.yaml", {"litellm_license": "not-a-valid-license"}, {})
|
||||
_boot_proxy(monkeypatch, private_key.public_key(), env_license=None, license_server_urls=[])
|
||||
|
||||
await ProxyConfig().load_config(router=None, config_file_path=config_file)
|
||||
|
||||
license_status: Final = await health_license_endpoint(user_api_key_dict=UserAPIKeyAuth())
|
||||
assert (license_status["has_license"], license_status["license_type"]) == (True, "community")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_load_config_does_not_recheck_an_env_license_that_failed_at_boot(
|
||||
tmp_path: Path, monkeypatch: pytest.MonkeyPatch
|
||||
):
|
||||
from litellm.proxy import proxy_server
|
||||
|
||||
private_key: Final = rsa.generate_private_key(public_exponent=65537, key_size=2048)
|
||||
config_file: Final = _write_config(tmp_path / "stale_license.yaml", {"litellm_license": "not-a-valid-license"}, {})
|
||||
license_server_urls: Final[list[str]] = []
|
||||
_boot_proxy(monkeypatch, private_key.public_key(), "also-not-a-valid-license", license_server_urls)
|
||||
asked_at_boot: Final = len(license_server_urls)
|
||||
|
||||
await ProxyConfig().load_config(router=None, config_file_path=config_file)
|
||||
|
||||
assert proxy_server.premium_user is False
|
||||
assert [url for url in license_server_urls[asked_at_boot:] if "also-not-a-valid-license" in url] == []
|
||||
assert any("not-a-valid-license" in url for url in license_server_urls[asked_at_boot:])
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_load_config_keeps_the_boot_verdict_of_a_server_checked_env_license(
|
||||
tmp_path: Path, monkeypatch: pytest.MonkeyPatch
|
||||
):
|
||||
from litellm.proxy import proxy_server
|
||||
|
||||
private_key: Final = rsa.generate_private_key(public_exponent=65537, key_size=2048)
|
||||
config_file: Final = _write_config(
|
||||
tmp_path / "stale_license.yaml", {"litellm_license": "not-a-valid-license", "allowed_ips": ["127.0.0.1"]}, {}
|
||||
)
|
||||
license_server_urls: Final[list[str]] = []
|
||||
_boot_proxy(
|
||||
monkeypatch,
|
||||
private_key.public_key(),
|
||||
"server-checked-license",
|
||||
license_server_urls,
|
||||
server_accepted_license="server-checked-license",
|
||||
)
|
||||
asked_at_boot: Final = len(license_server_urls)
|
||||
|
||||
await ProxyConfig().load_config(router=None, config_file_path=config_file)
|
||||
|
||||
assert proxy_server.premium_user is True
|
||||
assert proxy_server._license_check.license_str == "server-checked-license"
|
||||
assert [url for url in license_server_urls[asked_at_boot:] if "server-checked-license" in url] == []
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
"setting", ["vector_store_deny_by_default", "search_tool_deny_by_default", "claude_code_marketplace_auth_required"]
|
||||
|
|
|
|||
|
|
@ -4117,7 +4117,11 @@ async def test_load_environment_variables_litellm_license_and_edge_cases():
|
|||
mock_license_check = MagicMock()
|
||||
mock_license_check.is_premium.return_value = True
|
||||
|
||||
with patch("litellm.proxy.proxy_server._license_check", mock_license_check):
|
||||
with (
|
||||
patch("litellm.proxy.proxy_server._license_check", mock_license_check),
|
||||
patch("litellm.proxy.proxy_server.premium_user", False),
|
||||
patch("litellm.proxy.proxy_server.premium_user_data", None),
|
||||
):
|
||||
with patch.dict(os.environ, {}, clear=False):
|
||||
# Call the method under test
|
||||
proxy_config._load_environment_variables(test_config_with_license)
|
||||
|
|
|
|||
|
|
@ -1,4 +1,6 @@
|
|||
import base64
|
||||
import re
|
||||
from typing import Final
|
||||
|
||||
import pytest
|
||||
from botocore.stub import Stubber
|
||||
|
|
@ -34,3 +36,15 @@ def test_decrypt_value_turns_only_boolean_literals_into_bools(monkeypatch, plain
|
|||
|
||||
assert decrypted == expected
|
||||
assert type(decrypted) is type(expected)
|
||||
|
||||
|
||||
def test_kms_v2_without_a_license_tells_the_user_to_set_one(monkeypatch: pytest.MonkeyPatch):
|
||||
monkeypatch.setenv("AWS_REGION_NAME", "us-west-2")
|
||||
monkeypatch.delenv("LITELLM_LICENSE", raising=False)
|
||||
monkeypatch.delenv("LITELLM_SECRET_AWS_KMS_LITELLM_LICENSE", raising=False)
|
||||
|
||||
expected_message: Final = (
|
||||
"AWSKeyManagementService V2 is an Enterprise Feature. Please add a valid LITELLM_LICENSE to your environment."
|
||||
)
|
||||
with pytest.raises(ValueError, match=f"^{re.escape(expected_message)}$"):
|
||||
AWSKeyManagementService_V2()
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue