Honor include block when loading config from bucket

litellm/proxy/common_utils/load_config_utils.py

@@ -1,9 +1,11 @@
+from typing import Callable
+
 import yaml
 
 from litellm._logging import verbose_proxy_logger
 
 
-def get_file_contents_from_s3(bucket_name, object_key):
+async def get_file_contents_from_s3(bucket_name, object_key):
     try:
         # v0 rely on boto3 for authentication - allowing boto3 to handle IAM credentials etc
         import tempfile
@@ -40,6 +42,13 @@ def get_file_contents_from_s3(bucket_name, object_key):
         with open(temp_file_path, "r") as yaml_file:
             config = yaml.safe_load(yaml_file)
 
+        # include file config
+        config = await process_includes_from_bucket(
+            config=config,
+            get_file_method=get_file_contents_from_s3,
+            bucket_name=bucket_name,
+        )
+
         return config
     except ImportError as e:
         # this is most likely if a user is not using the litellm docker container
@@ -64,13 +73,57 @@ async def get_config_file_contents_from_gcs(bucket_name, object_key):
         file_contents = file_contents.decode("utf-8")
         # convert to yaml
         config = yaml.safe_load(file_contents)
+        # include file config
+        config = await process_includes_from_bucket(
+            config=config,
+            get_file_method=get_config_file_contents_from_gcs,
+            bucket_name=bucket_name,
+        )
         return config
 
     except Exception as e:
         verbose_proxy_logger.error(f"Error retrieving file contents: {str(e)}")
         return None
 
 
+async def process_includes_from_bucket(
+    config: dict, get_file_method: Callable, bucket_name: str
+) -> dict:
+    """
+    Process includes by appending their contents to the main config
+
+    Handles nested config.yamls with `include` section
+
+    Example config: This will get the contents from files in `include` and append it
+    ```yaml
+    include:
+        - /path/to/key/model_config.yaml
+
+    litellm_settings:
+        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"]:
+        included_config = await get_file_method(bucket_name, include_file)
+        # 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
+
+
 # # Example usage
 # bucket_name = 'litellm-proxy'
 # object_key = 'litellm_proxy_config.yaml'

litellm/proxy/proxy_server.py

@@ -1489,7 +1489,7 @@ async def get_config(self, config_file_path: Optional[str] = None) -> dict:
                     bucket_name=bucket_name, object_key=object_key
                 )
             else:
-                config = get_file_contents_from_s3(
+                config = await get_file_contents_from_s3(
                     bucket_name=bucket_name, object_key=object_key
                 )
 

pyproject.toml

@@ -117,6 +117,7 @@ types-PyYAML = "*"
 opentelemetry-api = "1.25.0"
 opentelemetry-sdk = "1.25.0"
 opentelemetry-exporter-otlp = "1.25.0"
+moto = "5.0.24"
 
 [tool.poetry.group.proxy-dev.dependencies]
 prisma = "0.11.0"