diff --git a/litellm/proxy/common_utils/config_includes.py b/litellm/proxy/common_utils/config_includes.py new file mode 100644 index 00000000000..d6939db2e95 --- /dev/null +++ b/litellm/proxy/common_utils/config_includes.py @@ -0,0 +1,68 @@ +from collections.abc import Awaitable, Mapping +from types import MappingProxyType +from typing import Final, Protocol + +INCLUDE_KEY: Final = "include" + + +class ConfigLoader(Protocol): + def __call__(self, include_entry: str, /) -> Awaitable[Mapping[str, object]]: ... + + +def _merged_value(base_value: object, included_value: object) -> object: + if isinstance(included_value, list) and isinstance(base_value, list): + return [*base_value, *included_value] # mutable-ok: a merged config value stays the plain list the proxy loads + return included_value + + +def _merged_entry(base: Mapping[str, object], included: Mapping[str, object], key: str) -> object: + if key not in included: + return base[key] + return _merged_value(base.get(key), included[key]) + + +def _merged(base: Mapping[str, object], included: Mapping[str, object]) -> Mapping[str, object]: + return MappingProxyType({key: _merged_entry(base, included, key) for key in (*base, *included)}) + + +def _without_include(config: Mapping[str, object]) -> Mapping[str, object]: + return MappingProxyType({key: value for key, value in config.items() if key != INCLUDE_KEY}) + + +def include_entries(config: Mapping[str, object]) -> tuple[str, ...]: + if INCLUDE_KEY not in config: + return () + + entries: Final = config[INCLUDE_KEY] + if not isinstance(entries, list): + raise ValueError("'include' must be a list of file paths") + + paths: Final = tuple(entry for entry in entries if isinstance(entry, str)) + if len(paths) != len(entries): + raise ValueError("'include' must be a list of file paths") + + return paths + + +async def _resolve(config: Mapping[str, object], pending: tuple[str, ...], load: ConfigLoader) -> Mapping[str, object]: + if not pending: + return _without_include(config) + + included: Final = await load(pending[0]) + return await _resolve( + _merged(config, _without_include(included)), + (*pending[1:], *include_entries(included)), + load, + ) + + +async def resolve_includes(config: Mapping[str, object], load: ConfigLoader) -> dict[str, object]: + """ + Merge every config named by the `include` directive into the config that declares it. + + List values are extended and every other value is overridden, an included config may declare + further includes, and `load` decides where an entry is read from, so the same merge applies to + configs on disk and to configs hosted in a bucket. + """ + merged: Final = await _resolve(config, include_entries(config), load) + return dict(merged) # mutable-ok: the proxy mutates the config it loads diff --git a/litellm/proxy/common_utils/load_config_utils.py b/litellm/proxy/common_utils/load_config_utils.py index 62649ad6ca1..927deb68826 100644 --- a/litellm/proxy/common_utils/load_config_utils.py +++ b/litellm/proxy/common_utils/load_config_utils.py @@ -1,9 +1,19 @@ import os -from typing import Final +import posixpath +from collections.abc import Awaitable, Mapping +from typing import Final, Protocol import yaml +from pydantic import TypeAdapter, ValidationError from litellm._logging import verbose_proxy_logger +from litellm.proxy.common_utils.config_includes import resolve_includes + +_BUCKET_CONFIG_ADAPTER: Final = TypeAdapter(dict[str, object]) + + +class BucketObjectFetcher(Protocol): + def __call__(self, object_key: str, /) -> Awaitable[Mapping[str, object] | None]: ... def get_file_contents_from_s3(bucket_name, object_key): @@ -62,6 +72,60 @@ async def get_config_file_contents_from_gcs(bucket_name, object_key): return None +def resolve_include_object_key(config_object_key: str, include_entry: str) -> str: + """ + Resolve one `include` entry to the object key it names, relative to the config object's prefix. + + A leading "/" means the bucket root, mirroring how an absolute path on disk ignores the + directory the including config sits in. + """ + if include_entry.startswith("/"): + return posixpath.normpath(include_entry).lstrip("/") + return posixpath.normpath(posixpath.join(posixpath.dirname(config_object_key), include_entry)) + + +async def resolve_bucket_includes( + *, + config: Mapping[str, object], + object_key: str, + fetch: BucketObjectFetcher, +) -> dict[str, object]: + async def load(include_entry: str) -> Mapping[str, object]: + include_key: Final = resolve_include_object_key(object_key, include_entry) + included: Final = await fetch(include_key) + if included is None: + raise FileNotFoundError(f"Included config could not be read from bucket: {include_key}") + return included + + return await resolve_includes(config=config, load=load) + + +async def get_config_from_bucket( + *, + bucket_type: str | None, + bucket_name: str, + object_key: str, +) -> dict[str, object] | None: + async def fetch(key: str) -> Mapping[str, object] | None: + raw: Final = ( + await get_config_file_contents_from_gcs(bucket_name=bucket_name, object_key=key) + if bucket_type == "gcs" + else get_file_contents_from_s3(bucket_name=bucket_name, object_key=key) + ) + if raw is None: + return None + try: + return _BUCKET_CONFIG_ADAPTER.validate_python(raw) + except ValidationError as e: + raise ValueError(f"Config object in bucket is not a YAML mapping: {key}") from e + + config: Final = await fetch(object_key) + if config is None: + return None + + return await resolve_bucket_includes(config=config, object_key=object_key, fetch=fetch) + + def download_python_file_from_s3( bucket_name: str, object_key: str, diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index ba5714fe950..a49df166dca 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -341,6 +341,7 @@ from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import ( AuthCacheInvalidationSubscriber, ) from litellm.proxy.common_utils.callback_utils import initialize_callbacks_on_proxy +from litellm.proxy.common_utils.config_includes import resolve_includes from litellm.proxy.common_utils.config_sync_pubsub import ConfigSyncSubscriber from litellm.proxy.common_utils.debug_utils import init_verbose_loggers from litellm.proxy.common_utils.debug_utils import router as debugging_endpoints_router @@ -359,10 +360,7 @@ from litellm.proxy.common_utils.http_parsing_utils import ( check_file_size_under_limit, get_form_data, ) -from litellm.proxy.common_utils.load_config_utils import ( - get_config_file_contents_from_gcs, - get_file_contents_from_s3, -) +from litellm.proxy.common_utils.load_config_utils import get_config_from_bucket from litellm.proxy.common_utils.model_deprecation import collect_model_deprecations from litellm.proxy.common_utils.model_listing_utils import ( TeamModelNameTranslator, @@ -4538,12 +4536,12 @@ class ProxyConfig: if config is None: raise Exception("Config cannot be None or Empty.") # Process includes - config = self._process_includes(config=config, base_dir=os.path.dirname(os.path.abspath(file_path or ""))) + config = await self._process_includes(config=config, base_dir=os.path.dirname(os.path.abspath(file_path or ""))) # verbose_proxy_logger.debug(f"loaded config={json.dumps(config, indent=4)}") return config - def _process_includes(self, config: dict, base_dir: str) -> dict: + async def _process_includes(self, config: dict, base_dir: str) -> dict: """ Process includes by appending their contents to the main config @@ -4558,29 +4556,14 @@ class ProxyConfig: callbacks: ["prometheus"] ``` """ - if "include" not in config: - return config - if not isinstance(config["include"], list): - raise ValueError("'include' must be a list of file paths") - - # Load and append all included files - for include_file in config["include"]: - file_path = os.path.join(base_dir, include_file) + async def load_included(include_file: str) -> Mapping[str, object]: + file_path: Final = os.path.join(base_dir, include_file) if not os.path.exists(file_path): raise FileNotFoundError(f"Included file not found: {file_path}") + return self._load_yaml_file(file_path) - included_config = self._load_yaml_file(file_path) - # Simply update/extend the main config with included config - for key, value in included_config.items(): - if isinstance(value, list) and key in config: - config[key].extend(value) - else: - config[key] = value - - # Remove the include directive - del config["include"] - return config + return await resolve_includes(config=config, load=load_included) async def save_config(self, new_config: dict, include_env_vars: bool = False): global prisma_client, general_settings, user_config_file_path, store_model_in_db @@ -4936,15 +4919,19 @@ class ProxyConfig: global prisma_client, store_model_in_db # Load existing config - if os.environ.get("LITELLM_CONFIG_BUCKET_NAME") is not None: - bucket_name: Final = os.environ.get("LITELLM_CONFIG_BUCKET_NAME") + bucket_name: Final = os.environ.get("LITELLM_CONFIG_BUCKET_NAME") + if bucket_name is not None: object_key: Final = os.environ.get("LITELLM_CONFIG_BUCKET_OBJECT_KEY") bucket_type: Final = os.environ.get("LITELLM_CONFIG_BUCKET_TYPE") verbose_proxy_logger.debug("bucket_name: %s, object_key: %s", bucket_name, object_key) - if bucket_type == "gcs": - config = await get_config_file_contents_from_gcs(bucket_name=bucket_name, object_key=object_key) - else: - config = get_file_contents_from_s3(bucket_name=bucket_name, object_key=object_key) + if object_key is None: + raise Exception("LITELLM_CONFIG_BUCKET_OBJECT_KEY must be set to load the config from a bucket.") + + config = await get_config_from_bucket( + bucket_type=bucket_type, + bucket_name=bucket_name, + object_key=object_key, + ) if config is None: raise Exception("Unable to load config from given source.") diff --git a/tests/test_litellm/proxy/common_utils/test_load_config_utils.py b/tests/test_litellm/proxy/common_utils/test_load_config_utils.py index 524c260e94a..c5d84460d39 100644 --- a/tests/test_litellm/proxy/common_utils/test_load_config_utils.py +++ b/tests/test_litellm/proxy/common_utils/test_load_config_utils.py @@ -1,9 +1,14 @@ +import re from unittest.mock import MagicMock, mock_open, patch import pytest import yaml -from litellm.proxy.common_utils.load_config_utils import get_file_contents_from_s3 +from litellm.proxy.common_utils.load_config_utils import ( + get_config_from_bucket, + get_file_contents_from_s3, + resolve_bucket_includes, +) class TestGetFileContentsFromS3: @@ -83,3 +88,176 @@ class TestGetFileContentsFromS3: # Verify yaml.safe_load was called with the decoded content mock_yaml_load.assert_called_once_with(yaml_content) + + +class TestBucketConfigIncludes: + """`include:` directives in a bucket-hosted config.yaml (LIT-6982). + + They used to be dropped silently: the proxy booted with the root config applied and everything + the included objects declared missing, with nothing logged. + """ + + @staticmethod + def _bucket(objects): + async def fetch(object_key): + return objects.get(object_key) + + return fetch + + @pytest.mark.asyncio + async def test_include_resolves_against_the_config_objects_prefix(self): + merged = await resolve_bucket_includes( + config={"include": ["model_config.yaml"], "general_settings": {"master_key": "sk-1234"}}, + object_key="configs/prod/config.yaml", + fetch=self._bucket( + {"configs/prod/model_config.yaml": {"model_list": [{"model_name": "gpt-4o-mini"}]}} + ), + ) + + assert merged == { + "general_settings": {"master_key": "sk-1234"}, + "model_list": [{"model_name": "gpt-4o-mini"}], + } + + @pytest.mark.asyncio + async def test_include_with_a_leading_slash_reads_from_the_bucket_root(self): + merged = await resolve_bucket_includes( + config={"include": ["/shared/models.yaml"]}, + object_key="configs/prod/config.yaml", + fetch=self._bucket({"shared/models.yaml": {"model_list": [{"model_name": "shared"}]}}), + ) + + assert merged == {"model_list": [{"model_name": "shared"}]} + + @pytest.mark.asyncio + async def test_include_walks_out_of_the_prefix_with_dot_dot(self): + merged = await resolve_bucket_includes( + config={"include": ["../shared/models.yaml"]}, + object_key="configs/prod/config.yaml", + fetch=self._bucket({"configs/shared/models.yaml": {"model_list": [{"model_name": "shared"}]}}), + ) + + assert merged == {"model_list": [{"model_name": "shared"}]} + + @pytest.mark.asyncio + async def test_included_configs_may_declare_further_includes(self): + merged = await resolve_bucket_includes( + config={"include": ["models.yaml"]}, + object_key="configs/config.yaml", + fetch=self._bucket( + { + "configs/models.yaml": { + "include": ["extra/more_models.yaml"], + "model_list": [{"model_name": "first"}], + }, + "configs/extra/more_models.yaml": {"model_list": [{"model_name": "second"}]}, + } + ), + ) + + assert merged == {"model_list": [{"model_name": "first"}, {"model_name": "second"}]} + + @pytest.mark.asyncio + async def test_list_values_are_extended_and_other_values_are_overridden(self): + merged = await resolve_bucket_includes( + config={ + "include": ["models.yaml"], + "model_list": [{"model_name": "from-root"}], + "litellm_settings": {"drop_params": True}, + }, + object_key="config.yaml", + fetch=self._bucket( + { + "models.yaml": { + "model_list": [{"model_name": "from-include"}], + "litellm_settings": {"num_retries": 3}, + } + } + ), + ) + + assert merged == { + "model_list": [{"model_name": "from-root"}, {"model_name": "from-include"}], + "litellm_settings": {"num_retries": 3}, + } + + @pytest.mark.asyncio + async def test_a_missing_included_object_fails_loudly_with_its_key(self): + with pytest.raises(FileNotFoundError, match=re.escape("configs/prod/model_config.yaml")): + await resolve_bucket_includes( + config={"include": ["model_config.yaml"]}, + object_key="configs/prod/config.yaml", + fetch=self._bucket({}), + ) + + @pytest.mark.asyncio + async def test_a_non_list_include_fails_loudly(self): + with pytest.raises(ValueError, match="'include' must be a list of file paths"): + await resolve_bucket_includes( + config={"include": "model_config.yaml"}, + object_key="config.yaml", + fetch=self._bucket({}), + ) + + @pytest.mark.asyncio + async def test_get_config_from_bucket_merges_includes_over_s3(self, monkeypatch): + objects = { + "lit6982/config.yaml": { + "include": ["model_config.yaml"], + "general_settings": {"master_key": "sk-1234"}, + }, + "lit6982/model_config.yaml": {"model_list": [{"model_name": "included-model"}]}, + } + monkeypatch.setattr( + "litellm.proxy.common_utils.load_config_utils.get_file_contents_from_s3", + lambda bucket_name, object_key: objects.get(object_key), + ) + + config = await get_config_from_bucket( + bucket_type="s3", bucket_name="litellm-configs", object_key="lit6982/config.yaml" + ) + + assert config == { + "general_settings": {"master_key": "sk-1234"}, + "model_list": [{"model_name": "included-model"}], + } + + @pytest.mark.asyncio + async def test_get_config_from_bucket_merges_includes_over_gcs(self, monkeypatch): + objects = { + "lit6982/config.yaml": { + "include": ["model_config.yaml"], + "general_settings": {"master_key": "sk-1234"}, + }, + "lit6982/model_config.yaml": {"model_list": [{"model_name": "included-model"}]}, + } + + async def fake_gcs(bucket_name, object_key): + return objects.get(object_key) + + monkeypatch.setattr( + "litellm.proxy.common_utils.load_config_utils.get_config_file_contents_from_gcs", fake_gcs + ) + + config = await get_config_from_bucket( + bucket_type="gcs", bucket_name="litellm-configs", object_key="lit6982/config.yaml" + ) + + assert config == { + "general_settings": {"master_key": "sk-1234"}, + "model_list": [{"model_name": "included-model"}], + } + + @pytest.mark.asyncio + async def test_get_config_from_bucket_returns_none_when_the_root_object_is_missing(self, monkeypatch): + monkeypatch.setattr( + "litellm.proxy.common_utils.load_config_utils.get_file_contents_from_s3", + lambda bucket_name, object_key: None, + ) + + assert ( + await get_config_from_bucket( + bucket_type="s3", bucket_name="litellm-configs", object_key="missing.yaml" + ) + is None + ) diff --git a/tests/test_litellm/proxy/proxy_server/test_proxy_config.py b/tests/test_litellm/proxy/proxy_server/test_proxy_config.py index 2babfe432f3..9b9c4af9e42 100644 --- a/tests/test_litellm/proxy/proxy_server/test_proxy_config.py +++ b/tests/test_litellm/proxy/proxy_server/test_proxy_config.py @@ -710,22 +710,32 @@ async def test_ProxyConfig__get_config_from_file_missing_path_raises(): # --------------------------------------------------------------------------- -def test_ProxyConfig__process_includes_merges_files(tmp_path): +@pytest.mark.asyncio +async def test_ProxyConfig__process_includes_merges_files(tmp_path): inc = tmp_path / "models.yaml" inc.write_text("model_list:\n - model_name: gpt-4\n") pc = ProxyConfig() cfg = {"include": ["models.yaml"], "model_list": [], "litellm_settings": {}} - result = pc._process_includes(cfg, base_dir=str(tmp_path)) + result = await pc._process_includes(cfg, base_dir=str(tmp_path)) assert result == { "model_list": [{"model_name": "gpt-4"}], "litellm_settings": {}, } -def test_ProxyConfig__process_includes_missing_file_raises(tmp_path): +@pytest.mark.asyncio +async def test_ProxyConfig__process_includes_missing_file_raises(tmp_path): pc = ProxyConfig() with pytest.raises(FileNotFoundError): - pc._process_includes({"include": ["nope.yaml"]}, base_dir=str(tmp_path)) + await pc._process_includes({"include": ["nope.yaml"]}, base_dir=str(tmp_path)) + + +@pytest.mark.asyncio +async def test_ProxyConfig__process_includes_follows_nested_includes(tmp_path): + (tmp_path / "models.yaml").write_text("include:\n - more_models.yaml\nmodel_list:\n - model_name: first\n") + (tmp_path / "more_models.yaml").write_text("model_list:\n - model_name: second\n") + result = await ProxyConfig()._process_includes({"include": ["models.yaml"]}, base_dir=str(tmp_path)) + assert result == {"model_list": [{"model_name": "first"}, {"model_name": "second"}]} # --------------------------------------------------------------------------- @@ -1042,6 +1052,32 @@ async def test_ProxyConfig_get_config_loads_from_file(tmp_path, monkeypatch): } +@pytest.mark.asyncio +async def test_ProxyConfig_get_config_from_a_bucket_merges_includes(monkeypatch): + """A bucket-hosted config.yaml used to drop its `include:` entries silently (LIT-6982).""" + objects = { + "lit6982/config.yaml": { + "include": ["model_config.yaml"], + "general_settings": {"master_key": "sk-1234"}, + }, + "lit6982/model_config.yaml": {"model_list": [{"model_name": "included-model"}]}, + } + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", None) + monkeypatch.setattr("litellm.proxy.proxy_server.store_model_in_db", False) + monkeypatch.setattr( + "litellm.proxy.common_utils.load_config_utils.get_file_contents_from_s3", + lambda bucket_name, object_key: objects.get(object_key), + ) + monkeypatch.setenv("LITELLM_CONFIG_BUCKET_NAME", "litellm-configs") + monkeypatch.setenv("LITELLM_CONFIG_BUCKET_OBJECT_KEY", "lit6982/config.yaml") + monkeypatch.setenv("LITELLM_CONFIG_BUCKET_TYPE", "s3") + + cfg = await ProxyConfig().get_config() + + assert cfg["model_list"] == [{"model_name": "included-model"}] + assert "include" not in cfg + + @pytest.mark.asyncio async def test_ProxyConfig_get_config_missing_file_raises(monkeypatch): monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", None)