mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-14 23:21:35 +00:00
Merge pull request #40772 from BerriAI/litellm_lit6982_bucket_config_includes
fix(proxy): resolve config include directives for bucket-hosted configs
This commit is contained in:
Mateo Wang
GitHub
commit
e1bff56f0b
7 changed files with 864 additions and 69 deletions
|
|
@ -29,8 +29,6 @@ class GCSBucketLogger(GCSBucketBase, AdditionalLoggingUtils):
|
|||
def __init__(self, bucket_name: str | None = None) -> None:
|
||||
from litellm.proxy.proxy_server import premium_user
|
||||
|
||||
super().__init__(bucket_name=bucket_name)
|
||||
|
||||
self.batch_size = int(os.getenv("GCS_BATCH_SIZE", GCS_DEFAULT_BATCH_SIZE))
|
||||
self.flush_interval = int(os.getenv("GCS_FLUSH_INTERVAL", GCS_DEFAULT_FLUSH_INTERVAL_SECONDS))
|
||||
self.use_batched_logging = (
|
||||
|
|
@ -38,6 +36,7 @@ class GCSBucketLogger(GCSBucketBase, AdditionalLoggingUtils):
|
|||
)
|
||||
self.flush_lock = asyncio.Lock()
|
||||
super().__init__(
|
||||
bucket_name=bucket_name,
|
||||
flush_lock=self.flush_lock,
|
||||
batch_size=self.batch_size,
|
||||
flush_interval=self.flush_interval,
|
||||
|
|
|
|||
132
litellm/proxy/common_utils/config_includes.py
Normal file
132
litellm/proxy/common_utils/config_includes.py
Normal file
|
|
@ -0,0 +1,132 @@
|
|||
import os
|
||||
from collections.abc import Awaitable, Mapping
|
||||
from types import MappingProxyType
|
||||
from typing import Final, Protocol
|
||||
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
|
||||
INCLUDE_KEY: Final = "include"
|
||||
|
||||
|
||||
def resolve_include_file_path(include_file: str, declared_in: str, root_config_path: str) -> str:
|
||||
"""
|
||||
Resolve one `include` entry to the file it names, next to the config that declares it.
|
||||
|
||||
A config written before nested entries resolved this way can name a file sitting next to the root
|
||||
config instead, so that file is still read, with a warning naming where it was found. When both
|
||||
files exist the one next to the declaring config wins and the other is named in a warning.
|
||||
"""
|
||||
declared_relative: Final = os.path.abspath(os.path.join(os.path.dirname(declared_in), include_file))
|
||||
root_relative: Final = os.path.abspath(os.path.join(os.path.dirname(root_config_path), include_file))
|
||||
if root_relative == declared_relative or not os.path.exists(root_relative):
|
||||
return declared_relative
|
||||
|
||||
if not os.path.exists(declared_relative):
|
||||
verbose_proxy_logger.warning(
|
||||
"Config include '%s' declared in %s was not found next to it, so %s was read instead. "
|
||||
"Move the included file next to the config that declares it.",
|
||||
include_file,
|
||||
declared_in,
|
||||
root_relative,
|
||||
)
|
||||
return root_relative
|
||||
|
||||
verbose_proxy_logger.warning(
|
||||
"Config include '%s' declared in %s matches two files. %s sits next to that config and was read, "
|
||||
"so %s was skipped. Rename one of the two to say which one you meant.",
|
||||
include_file,
|
||||
declared_in,
|
||||
declared_relative,
|
||||
root_relative,
|
||||
)
|
||||
return declared_relative
|
||||
|
||||
|
||||
class IncludeResolver(Protocol):
|
||||
def __call__(self, include_entry: str, declared_in: str, /) -> str: ...
|
||||
|
||||
|
||||
class ConfigReader(Protocol):
|
||||
def __call__(self, location: 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
|
||||
|
||||
|
||||
def _pending_from(config: Mapping[str, object], location: str) -> tuple[tuple[str, str], ...]:
|
||||
return tuple((entry, location) for entry in include_entries(config))
|
||||
|
||||
|
||||
async def _resolve(
|
||||
config: Mapping[str, object],
|
||||
pending: tuple[tuple[str, str], ...],
|
||||
loaded: frozenset[str],
|
||||
resolve: IncludeResolver,
|
||||
read: ConfigReader,
|
||||
) -> Mapping[str, object]:
|
||||
if not pending:
|
||||
return _without_include(config)
|
||||
|
||||
entry, declared_in = pending[0]
|
||||
location: Final = resolve(entry, declared_in)
|
||||
if location in loaded:
|
||||
return await _resolve(config, pending[1:], loaded, resolve, read)
|
||||
|
||||
included: Final = await read(location)
|
||||
return await _resolve(
|
||||
_merged(config, _without_include(included)),
|
||||
(*pending[1:], *_pending_from(included, location)),
|
||||
loaded | frozenset((location,)),
|
||||
resolve,
|
||||
read,
|
||||
)
|
||||
|
||||
|
||||
async def resolve_includes(
|
||||
config: Mapping[str, object],
|
||||
location: str,
|
||||
resolve: IncludeResolver,
|
||||
read: ConfigReader,
|
||||
) -> 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, `resolve` turns each entry into the
|
||||
location it names relative to the config that declares it, a config already pulled in is neither
|
||||
read nor merged a second time, and `read` decides where a location is read from, so the same merge
|
||||
applies to configs on disk and to configs hosted in a bucket.
|
||||
"""
|
||||
merged: Final = await _resolve(config, _pending_from(config, location), frozenset((location,)), resolve, read)
|
||||
return dict(merged) # mutable-ok: the proxy mutates the config it loads
|
||||
|
|
@ -1,12 +1,47 @@
|
|||
import asyncio
|
||||
import os
|
||||
from typing import Final
|
||||
import posixpath
|
||||
from collections.abc import Awaitable, Mapping
|
||||
from types import MappingProxyType
|
||||
from typing import TYPE_CHECKING, 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
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from litellm.integrations.gcs_bucket.gcs_bucket_base import GCSBucketBase
|
||||
|
||||
_BUCKET_CONFIG_ADAPTER: Final = TypeAdapter(dict[str, object])
|
||||
|
||||
|
||||
def get_file_contents_from_s3(bucket_name, object_key):
|
||||
class BucketObjectFetcher(Protocol):
|
||||
def __call__(self, object_key: str, /) -> Awaitable[Mapping[str, object] | None]: ...
|
||||
|
||||
|
||||
class BucketObjectReader(Protocol):
|
||||
def __call__(self, object_key: str, /) -> Awaitable[object | None]: ...
|
||||
|
||||
|
||||
class SyncBucketObjectReader(Protocol):
|
||||
def __call__(self, object_key: str, /) -> object | None: ...
|
||||
|
||||
|
||||
def _parsed_config(object_key: str, file_contents: str) -> object | None:
|
||||
try:
|
||||
parsed: Final = yaml.safe_load(file_contents)
|
||||
except yaml.YAMLError as e:
|
||||
verbose_proxy_logger.error("Config object %s is not valid YAML: %s", object_key, e)
|
||||
return None
|
||||
return MappingProxyType({}) if parsed is None else parsed
|
||||
|
||||
|
||||
def s3_object_reader(bucket_name: str) -> SyncBucketObjectReader:
|
||||
"""
|
||||
Build one reader for a whole config, so an `include` tree costs one S3 client rather than one per object.
|
||||
"""
|
||||
try:
|
||||
# v0 rely on boto3 for authentication - allowing boto3 to handle IAM credentials etc
|
||||
import boto3
|
||||
|
|
@ -21,46 +56,147 @@ def get_file_contents_from_s3(bucket_name, object_key):
|
|||
aws_secret_access_key=credentials.secret_key,
|
||||
aws_session_token=credentials.token, # Optional, if using temporary credentials
|
||||
)
|
||||
verbose_proxy_logger.debug("Retrieving %s from S3 bucket: %s", object_key, bucket_name)
|
||||
response: Final = s3_client.get_object(Bucket=bucket_name, Key=object_key)
|
||||
verbose_proxy_logger.debug("Response: %s", response)
|
||||
|
||||
# Read the file contents and directly parse YAML
|
||||
file_contents: Final = response["Body"].read().decode("utf-8")
|
||||
verbose_proxy_logger.debug("File contents retrieved from S3")
|
||||
|
||||
# Parse YAML directly from string
|
||||
config: Final = yaml.safe_load(file_contents)
|
||||
return config
|
||||
|
||||
except ImportError as e:
|
||||
# this is most likely if a user is not using the litellm docker container
|
||||
verbose_proxy_logger.error("ImportError: %s", e)
|
||||
return lambda object_key: None
|
||||
except Exception as e:
|
||||
verbose_proxy_logger.error("Error retrieving file contents: %s", e)
|
||||
verbose_proxy_logger.error("Error creating the S3 client for bucket %s: %s", bucket_name, e)
|
||||
return lambda object_key: None
|
||||
|
||||
def read(object_key: str) -> object | None:
|
||||
try:
|
||||
verbose_proxy_logger.debug("Retrieving %s from S3 bucket: %s", object_key, bucket_name)
|
||||
response: Final = s3_client.get_object(Bucket=bucket_name, Key=object_key)
|
||||
file_contents: Final = response["Body"].read().decode("utf-8")
|
||||
except Exception as e: # noqa: BLE001 # any boto3 error must read as a missing object
|
||||
verbose_proxy_logger.error("Error retrieving %s from S3 bucket %s: %s", object_key, bucket_name, e)
|
||||
return None
|
||||
|
||||
return _parsed_config(object_key, file_contents)
|
||||
|
||||
return read
|
||||
|
||||
|
||||
def get_file_contents_from_s3(bucket_name: str, object_key: str) -> object | None:
|
||||
return s3_object_reader(bucket_name)(object_key)
|
||||
|
||||
|
||||
def gcs_config_bucket(bucket_name: str) -> "GCSBucketBase | None":
|
||||
"""
|
||||
Build a plain GCS client for reading config objects.
|
||||
|
||||
Reading a config out of a bucket is not GCS logging, so it neither needs the enterprise license
|
||||
that gate covers nor the batching task the logger starts and never stops.
|
||||
"""
|
||||
try:
|
||||
from litellm.integrations.gcs_bucket.gcs_bucket_base import GCSBucketBase
|
||||
|
||||
return GCSBucketBase(bucket_name=bucket_name)
|
||||
except Exception as e: # noqa: BLE001 # an unbuildable client must read as an unreadable bucket
|
||||
verbose_proxy_logger.error("Error creating the GCS client for bucket %s: %s", bucket_name, e)
|
||||
return None
|
||||
|
||||
|
||||
async def get_config_file_contents_from_gcs(bucket_name, object_key):
|
||||
async def get_config_file_contents_from_gcs(
|
||||
bucket_name: str,
|
||||
object_key: str,
|
||||
gcs_bucket: "GCSBucketBase | None" = None,
|
||||
) -> object | None:
|
||||
try:
|
||||
from litellm.integrations.gcs_bucket.gcs_bucket import GCSBucketLogger
|
||||
|
||||
gcs_bucket: Final = GCSBucketLogger(
|
||||
bucket_name=bucket_name,
|
||||
)
|
||||
file_contents = await gcs_bucket.download_gcs_object(object_key)
|
||||
bucket: Final = gcs_config_bucket(bucket_name) if gcs_bucket is None else gcs_bucket
|
||||
if bucket is None:
|
||||
return None
|
||||
file_contents: Final = await bucket.download_gcs_object(object_key)
|
||||
if file_contents is None:
|
||||
raise Exception(f"File contents are None for {object_key}")
|
||||
# file_contentis is a bytes object, so we need to convert it to yaml
|
||||
file_contents = file_contents.decode("utf-8")
|
||||
# convert to yaml
|
||||
config: Final = yaml.safe_load(file_contents)
|
||||
return config
|
||||
decoded: Final = file_contents.decode("utf-8")
|
||||
|
||||
except Exception as e:
|
||||
verbose_proxy_logger.error("Error retrieving file contents: %s", e)
|
||||
verbose_proxy_logger.error("Error retrieving %s from GCS bucket %s: %s", object_key, bucket_name, e)
|
||||
return None
|
||||
|
||||
return _parsed_config(object_key, decoded)
|
||||
|
||||
|
||||
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 read(include_key: str) -> Mapping[str, object]:
|
||||
included: Final = await fetch(include_key)
|
||||
if included is None:
|
||||
raise FileNotFoundError(
|
||||
f"Included config could not be read from bucket: {include_key}. "
|
||||
"The underlying bucket error is logged above."
|
||||
)
|
||||
return included
|
||||
|
||||
def resolve(include_entry: str, declared_in: str) -> str:
|
||||
return resolve_include_object_key(declared_in, include_entry)
|
||||
|
||||
return await resolve_includes(config=config, location=object_key, resolve=resolve, read=read)
|
||||
|
||||
|
||||
async def bucket_object_reader(bucket_type: str | None, bucket_name: str) -> BucketObjectReader:
|
||||
"""
|
||||
Build one reader for a whole config, so an `include` tree costs one bucket client rather than one per object.
|
||||
"""
|
||||
if bucket_type != "gcs":
|
||||
read_object: Final = await asyncio.to_thread(s3_object_reader, bucket_name)
|
||||
|
||||
async def read_from_s3(object_key: str) -> object | None:
|
||||
return await asyncio.to_thread(read_object, object_key)
|
||||
|
||||
return read_from_s3
|
||||
|
||||
gcs_bucket: Final = gcs_config_bucket(bucket_name)
|
||||
|
||||
async def read_from_gcs(object_key: str) -> object | None:
|
||||
if gcs_bucket is None:
|
||||
return None
|
||||
return await get_config_file_contents_from_gcs(bucket_name, object_key, gcs_bucket)
|
||||
|
||||
return read_from_gcs
|
||||
|
||||
|
||||
async def get_config_from_bucket(
|
||||
*,
|
||||
bucket_type: str | None,
|
||||
bucket_name: str,
|
||||
object_key: str,
|
||||
) -> dict[str, object] | None:
|
||||
read: Final = await bucket_object_reader(bucket_type, bucket_name)
|
||||
|
||||
async def fetch(key: str) -> Mapping[str, object] | None:
|
||||
raw: Final = await read(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 not config:
|
||||
return None
|
||||
|
||||
return await resolve_bucket_includes(config=config, object_key=object_key, fetch=fetch)
|
||||
|
||||
|
||||
def download_python_file_from_s3(
|
||||
bucket_name: str,
|
||||
|
|
@ -136,11 +272,9 @@ async def download_python_file_from_gcs(
|
|||
bool: True if successful, False otherwise
|
||||
"""
|
||||
try:
|
||||
from litellm.integrations.gcs_bucket.gcs_bucket import GCSBucketLogger
|
||||
from litellm.integrations.gcs_bucket.gcs_bucket_base import GCSBucketBase
|
||||
|
||||
gcs_bucket: Final = GCSBucketLogger(
|
||||
bucket_name=bucket_name,
|
||||
)
|
||||
gcs_bucket: Final = GCSBucketBase(bucket_name=bucket_name)
|
||||
file_contents = await gcs_bucket.download_gcs_object(object_key)
|
||||
if file_contents is None:
|
||||
raise Exception(f"File contents are None for {object_key}")
|
||||
|
|
|
|||
|
|
@ -353,6 +353,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_include_file_path, 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
|
||||
|
|
@ -371,10 +372,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 (
|
||||
ClaudeCodeRoutingNames,
|
||||
|
|
@ -4806,12 +4804,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, config_file_path=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, config_file_path: str) -> dict:
|
||||
"""
|
||||
Process includes by appending their contents to the main config
|
||||
|
||||
|
|
@ -4826,29 +4824,21 @@ 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")
|
||||
included_config_adapter: Final = TypeAdapter(dict[str, object])
|
||||
|
||||
# Load and append all included files
|
||||
for include_file in config["include"]:
|
||||
file_path = os.path.join(base_dir, include_file)
|
||||
def resolve(include_file: str, declared_in: str) -> str:
|
||||
return resolve_include_file_path(include_file, declared_in, config_file_path)
|
||||
|
||||
async def read_included(file_path: str) -> Mapping[str, object]:
|
||||
if not os.path.exists(file_path):
|
||||
raise FileNotFoundError(f"Included file not found: {file_path}")
|
||||
try:
|
||||
return included_config_adapter.validate_python(self._load_yaml_file(file_path))
|
||||
except ValidationError as e:
|
||||
raise ValueError(f"Included config file is not a YAML mapping: {file_path}") from e
|
||||
|
||||
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, location=config_file_path, resolve=resolve, read=read_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
|
||||
|
|
@ -5206,15 +5196,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.")
|
||||
|
|
|
|||
|
|
@ -128,3 +128,20 @@ class TestGCSBucketBase:
|
|||
assert object_name.endswith("-target_uploadType_media")
|
||||
assert ".." not in object_name
|
||||
assert "?" not in object_name
|
||||
|
||||
|
||||
class TestGCSBucketLoggerBucketName:
|
||||
@pytest.mark.asyncio
|
||||
async def test_the_bucket_name_it_is_constructed_with_survives(self, monkeypatch):
|
||||
"""Reading config.yaml out of a GCS bucket asks for that bucket, not the logging one (LIT-6982)."""
|
||||
monkeypatch.setenv("GCS_BUCKET_NAME", "logging-bucket")
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.premium_user", True)
|
||||
|
||||
assert GCSBucketLogger(bucket_name="config-bucket").BUCKET_NAME == "config-bucket"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_no_bucket_name_still_falls_back_to_the_environment(self, monkeypatch):
|
||||
monkeypatch.setenv("GCS_BUCKET_NAME", "logging-bucket")
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.premium_user", True)
|
||||
|
||||
assert GCSBucketLogger().BUCKET_NAME == "logging-bucket"
|
||||
|
|
|
|||
|
|
@ -1,9 +1,18 @@
|
|||
import asyncio
|
||||
import logging
|
||||
import re
|
||||
import threading
|
||||
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 (
|
||||
gcs_config_bucket,
|
||||
get_config_from_bucket,
|
||||
get_file_contents_from_s3,
|
||||
resolve_bucket_includes,
|
||||
)
|
||||
|
||||
|
||||
class TestGetFileContentsFromS3:
|
||||
|
|
@ -83,3 +92,385 @@ class TestGetFileContentsFromS3:
|
|||
|
||||
# Verify yaml.safe_load was called with the decoded content
|
||||
mock_yaml_load.assert_called_once_with(yaml_content)
|
||||
|
||||
|
||||
class TestBucketConfigIncludes:
|
||||
|
||||
@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_a_nested_include_resolves_against_the_object_that_declares_it(self):
|
||||
merged = await resolve_bucket_includes(
|
||||
config={"include": ["shared/models.yaml"]},
|
||||
object_key="configs/config.yaml",
|
||||
fetch=self._bucket(
|
||||
{
|
||||
"configs/shared/models.yaml": {
|
||||
"include": ["more_models.yaml"],
|
||||
"model_list": [{"model_name": "first"}],
|
||||
},
|
||||
"configs/shared/more_models.yaml": {"model_list": [{"model_name": "second"}]},
|
||||
"configs/more_models.yaml": {"model_list": [{"model_name": "wrong-prefix"}]},
|
||||
}
|
||||
),
|
||||
)
|
||||
|
||||
assert merged == {"model_list": [{"model_name": "first"}, {"model_name": "second"}]}
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_an_object_pulled_in_twice_is_merged_once(self):
|
||||
merged = await resolve_bucket_includes(
|
||||
config={"include": ["a.yaml", "b.yaml"]},
|
||||
object_key="configs/config.yaml",
|
||||
fetch=self._bucket(
|
||||
{
|
||||
"configs/a.yaml": {"include": ["shared.yaml"]},
|
||||
"configs/b.yaml": {"include": ["./shared.yaml"]},
|
||||
"configs/shared.yaml": {"model_list": [{"model_name": "shared"}]},
|
||||
}
|
||||
),
|
||||
)
|
||||
|
||||
assert merged == {"model_list": [{"model_name": "shared"}]}
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_cycle_between_included_objects_terminates(self):
|
||||
merged = await asyncio.wait_for(
|
||||
resolve_bucket_includes(
|
||||
config={"include": ["a.yaml"]},
|
||||
object_key="configs/config.yaml",
|
||||
fetch=self._bucket(
|
||||
{
|
||||
"configs/a.yaml": {"include": ["b.yaml"], "model_list": [{"model_name": "from-a"}]},
|
||||
"configs/b.yaml": {"include": ["a.yaml"], "model_list": [{"model_name": "from-b"}]},
|
||||
}
|
||||
),
|
||||
),
|
||||
timeout=10,
|
||||
)
|
||||
|
||||
assert merged == {"model_list": [{"model_name": "from-a"}, {"model_name": "from-b"}]}
|
||||
|
||||
@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.s3_object_reader",
|
||||
lambda bucket_name: objects.get,
|
||||
)
|
||||
|
||||
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_the_blocking_s3_work_runs_off_the_event_loop_thread(self, monkeypatch):
|
||||
loop_thread = threading.current_thread()
|
||||
threads = []
|
||||
|
||||
def build_reader(bucket_name):
|
||||
threads.append(threading.current_thread())
|
||||
|
||||
def read(object_key):
|
||||
threads.append(threading.current_thread())
|
||||
return {"model_list": [{"model_name": "a-model"}]}
|
||||
|
||||
return read
|
||||
|
||||
monkeypatch.setattr("litellm.proxy.common_utils.load_config_utils.s3_object_reader", build_reader)
|
||||
|
||||
await get_config_from_bucket(bucket_type="s3", bucket_name="litellm-configs", object_key="config.yaml")
|
||||
|
||||
assert len(threads) == 2 and loop_thread not in threads
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_one_s3_client_serves_the_whole_include_tree(self, monkeypatch):
|
||||
objects = {
|
||||
"lit6982/config.yaml": {"include": ["model_config.yaml"]},
|
||||
"lit6982/model_config.yaml": {"model_list": [{"model_name": "included-model"}]},
|
||||
}
|
||||
readers = []
|
||||
|
||||
def build_reader(bucket_name):
|
||||
requested = []
|
||||
readers.append(requested)
|
||||
|
||||
def read(object_key):
|
||||
requested.append(object_key)
|
||||
return objects.get(object_key)
|
||||
|
||||
return read
|
||||
|
||||
monkeypatch.setattr("litellm.proxy.common_utils.load_config_utils.s3_object_reader", build_reader)
|
||||
|
||||
await get_config_from_bucket(
|
||||
bucket_type="s3", bucket_name="litellm-configs", object_key="lit6982/config.yaml"
|
||||
)
|
||||
|
||||
assert readers == [["lit6982/config.yaml", "lit6982/model_config.yaml"]]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_an_empty_included_object_merges_as_an_empty_config(self, monkeypatch):
|
||||
objects = {
|
||||
"lit6982/config.yaml": "include:\n - empty.yaml\nmodel_list:\n - model_name: only-model\n",
|
||||
"lit6982/empty.yaml": "",
|
||||
}
|
||||
|
||||
class FakeGCSBucket:
|
||||
async def download_gcs_object(self, object_key):
|
||||
return objects[object_key].encode("utf-8")
|
||||
|
||||
monkeypatch.setattr(
|
||||
"litellm.proxy.common_utils.load_config_utils.gcs_config_bucket",
|
||||
lambda bucket_name: FakeGCSBucket(),
|
||||
)
|
||||
|
||||
config = await get_config_from_bucket(
|
||||
bucket_type="gcs", bucket_name="litellm-configs", object_key="lit6982/config.yaml"
|
||||
)
|
||||
|
||||
assert config == {"model_list": [{"model_name": "only-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"}]},
|
||||
}
|
||||
|
||||
buckets = []
|
||||
|
||||
class FakeGCSBucket:
|
||||
def __init__(self):
|
||||
self.requested = []
|
||||
buckets.append(self)
|
||||
|
||||
async def download_gcs_object(self, object_key):
|
||||
self.requested.append(object_key)
|
||||
return yaml.dump(objects[object_key]).encode("utf-8")
|
||||
|
||||
monkeypatch.setattr(
|
||||
"litellm.proxy.common_utils.load_config_utils.gcs_config_bucket",
|
||||
lambda bucket_name: FakeGCSBucket(),
|
||||
)
|
||||
|
||||
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"}],
|
||||
}
|
||||
assert [bucket.requested for bucket in buckets] == [
|
||||
["lit6982/config.yaml", "lit6982/model_config.yaml"]
|
||||
]
|
||||
|
||||
@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.s3_object_reader",
|
||||
lambda bucket_name: (lambda object_key: None),
|
||||
)
|
||||
|
||||
assert (
|
||||
await get_config_from_bucket(
|
||||
bucket_type="s3", bucket_name="litellm-configs", object_key="missing.yaml"
|
||||
)
|
||||
is None
|
||||
)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_an_object_pulled_in_twice_is_read_once(self):
|
||||
objects = {
|
||||
"configs/a.yaml": {"include": ["shared.yaml"]},
|
||||
"configs/b.yaml": {"include": ["./shared.yaml"]},
|
||||
"configs/shared.yaml": {"model_list": [{"model_name": "shared"}]},
|
||||
}
|
||||
requested = []
|
||||
|
||||
async def fetch(object_key):
|
||||
requested.append(object_key)
|
||||
return objects.get(object_key)
|
||||
|
||||
await resolve_bucket_includes(
|
||||
config={"include": ["a.yaml", "b.yaml"]},
|
||||
object_key="configs/config.yaml",
|
||||
fetch=fetch,
|
||||
)
|
||||
|
||||
assert requested == ["configs/a.yaml", "configs/b.yaml", "configs/shared.yaml"]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_an_empty_root_object_does_not_boot_an_empty_proxy(self, monkeypatch):
|
||||
class FakeGCSBucket:
|
||||
async def download_gcs_object(self, object_key):
|
||||
return b""
|
||||
|
||||
monkeypatch.setattr(
|
||||
"litellm.proxy.common_utils.load_config_utils.gcs_config_bucket",
|
||||
lambda bucket_name: FakeGCSBucket(),
|
||||
)
|
||||
|
||||
config = await get_config_from_bucket(
|
||||
bucket_type="gcs", bucket_name="litellm-configs", object_key="lit6982/config.yaml"
|
||||
)
|
||||
|
||||
assert config is None
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_an_object_that_is_not_valid_yaml_is_reported_as_a_yaml_error(self, monkeypatch, caplog):
|
||||
class FakeGCSBucket:
|
||||
async def download_gcs_object(self, object_key):
|
||||
return b"model_list: [\n"
|
||||
|
||||
monkeypatch.setattr(
|
||||
"litellm.proxy.common_utils.load_config_utils.gcs_config_bucket",
|
||||
lambda bucket_name: FakeGCSBucket(),
|
||||
)
|
||||
|
||||
with caplog.at_level(logging.ERROR, logger="LiteLLM Proxy"):
|
||||
config = await get_config_from_bucket(
|
||||
bucket_type="gcs", bucket_name="litellm-configs", object_key="lit6982/config.yaml"
|
||||
)
|
||||
|
||||
assert config is None
|
||||
assert [
|
||||
record
|
||||
for record in caplog.records
|
||||
if "not valid YAML" in record.getMessage() and "lit6982/config.yaml" in record.getMessage()
|
||||
]
|
||||
|
||||
|
||||
class TestGCSConfigBucketClient:
|
||||
@pytest.mark.asyncio
|
||||
async def test_reading_a_config_from_gcs_does_not_need_an_enterprise_license(self, monkeypatch):
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.premium_user", False)
|
||||
|
||||
bucket = gcs_config_bucket("litellm-configs")
|
||||
|
||||
assert bucket is not None
|
||||
assert bucket.BUCKET_NAME == "litellm-configs"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_reading_a_config_from_gcs_starts_no_background_task(self, monkeypatch):
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.premium_user", True)
|
||||
running_before = asyncio.all_tasks()
|
||||
|
||||
gcs_config_bucket("litellm-configs")
|
||||
|
||||
assert asyncio.all_tasks() - running_before == set()
|
||||
|
|
|
|||
|
|
@ -8,6 +8,7 @@ Pins covered:
|
|||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
|
|
@ -712,22 +713,124 @@ 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, config_file_path=str(tmp_path / "config.yaml"))
|
||||
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"]}, config_file_path=str(tmp_path / "config.yaml"))
|
||||
|
||||
|
||||
@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"]}, config_file_path=str(tmp_path / "config.yaml")
|
||||
)
|
||||
assert result == {"model_list": [{"model_name": "first"}, {"model_name": "second"}]}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_ProxyConfig__process_includes_resolves_a_nested_include_next_to_its_own_file(tmp_path):
|
||||
(tmp_path / "shared").mkdir()
|
||||
(tmp_path / "shared" / "models.yaml").write_text(
|
||||
"include:\n - more_models.yaml\nmodel_list:\n - model_name: first\n"
|
||||
)
|
||||
(tmp_path / "shared" / "more_models.yaml").write_text("model_list:\n - model_name: second\n")
|
||||
(tmp_path / "more_models.yaml").write_text("model_list:\n - model_name: wrong-directory\n")
|
||||
|
||||
result = await ProxyConfig()._process_includes(
|
||||
{"include": ["shared/models.yaml"]}, config_file_path=str(tmp_path / "config.yaml")
|
||||
)
|
||||
|
||||
assert result == {"model_list": [{"model_name": "first"}, {"model_name": "second"}]}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_ProxyConfig__process_includes_still_reads_a_nested_include_left_beside_the_root_config(tmp_path):
|
||||
(tmp_path / "shared").mkdir()
|
||||
(tmp_path / "shared" / "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": ["shared/models.yaml"]}, config_file_path=str(tmp_path / "config.yaml")
|
||||
)
|
||||
|
||||
assert result == {"model_list": [{"model_name": "first"}, {"model_name": "second"}]}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_ProxyConfig__process_includes_names_both_files_when_a_nested_include_matches_two(tmp_path, caplog):
|
||||
(tmp_path / "shared").mkdir()
|
||||
(tmp_path / "shared" / "models.yaml").write_text(
|
||||
"include:\n - more_models.yaml\nmodel_list:\n - model_name: first\n"
|
||||
)
|
||||
(tmp_path / "shared" / "more_models.yaml").write_text("model_list:\n - model_name: next-to-the-declaring-file\n")
|
||||
(tmp_path / "more_models.yaml").write_text("model_list:\n - model_name: next-to-the-root-config\n")
|
||||
|
||||
with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"):
|
||||
result = await ProxyConfig()._process_includes(
|
||||
{"include": ["shared/models.yaml"]}, config_file_path=str(tmp_path / "config.yaml")
|
||||
)
|
||||
|
||||
assert result == {"model_list": [{"model_name": "first"}, {"model_name": "next-to-the-declaring-file"}]}
|
||||
assert [
|
||||
record
|
||||
for record in caplog.records
|
||||
if str(tmp_path / "shared" / "more_models.yaml") in record.getMessage()
|
||||
and str(tmp_path / "more_models.yaml") in record.getMessage()
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_ProxyConfig__process_includes_merges_a_shared_file_once(tmp_path):
|
||||
(tmp_path / "shared.yaml").write_text("model_list:\n - model_name: shared\n")
|
||||
(tmp_path / "a.yaml").write_text("include:\n - shared.yaml\n")
|
||||
(tmp_path / "b.yaml").write_text("include:\n - ./shared.yaml\n")
|
||||
|
||||
result = await ProxyConfig()._process_includes(
|
||||
{"include": ["a.yaml", "b.yaml"]}, config_file_path=str(tmp_path / "config.yaml")
|
||||
)
|
||||
|
||||
assert result == {"model_list": [{"model_name": "shared"}]}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_ProxyConfig__process_includes_names_the_file_when_it_is_not_a_mapping(tmp_path):
|
||||
(tmp_path / "models.yaml").write_text("- model_name: gpt-4\n")
|
||||
|
||||
with pytest.raises(ValueError, match=re.escape(str(tmp_path / "models.yaml"))):
|
||||
await ProxyConfig()._process_includes(
|
||||
{"include": ["models.yaml"]}, config_file_path=str(tmp_path / "config.yaml")
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_ProxyConfig__process_includes_terminates_on_a_cycle(tmp_path):
|
||||
(tmp_path / "a.yaml").write_text("include:\n - b.yaml\nmodel_list:\n - model_name: from-a\n")
|
||||
(tmp_path / "b.yaml").write_text("include:\n - a.yaml\nmodel_list:\n - model_name: from-b\n")
|
||||
|
||||
result = await asyncio.wait_for(
|
||||
ProxyConfig()._process_includes({"include": ["a.yaml"]}, config_file_path=str(tmp_path / "config.yaml")),
|
||||
timeout=10,
|
||||
)
|
||||
|
||||
assert result == {"model_list": [{"model_name": "from-a"}, {"model_name": "from-b"}]}
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
|
|
@ -1044,6 +1147,31 @@ 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):
|
||||
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.s3_object_reader",
|
||||
lambda bucket_name: objects.get,
|
||||
)
|
||||
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)
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue