use include block when loading config from bucket

Author: Bharat Sinha

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
@@ -20,26 +22,33 @@ 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(
+        verbose_proxy_logger.error(
             f"Retrieving {object_key} from S3 bucket: {bucket_name}"
         )
         response = s3_client.get_object(Bucket=bucket_name, Key=object_key)
-        verbose_proxy_logger.debug(f"Response: {response}")
+        verbose_proxy_logger.error(f"Response: {response}")
 
         # Read the file contents
         file_contents = response["Body"].read().decode("utf-8")
-        verbose_proxy_logger.debug("File contents retrieved from S3")
+        verbose_proxy_logger.error("File contents retrieved from S3")
 
         # Create a temporary file with YAML extension
         with tempfile.NamedTemporaryFile(delete=False, suffix=".yaml") as temp_file:
             temp_file.write(file_contents.encode("utf-8"))
             temp_file_path = temp_file.name
-            verbose_proxy_logger.debug(f"File stored temporarily at: {temp_file_path}")
+            verbose_proxy_logger.error(f"File stored temporarily at: {temp_file_path}")
 
         # Load the YAML file content
         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"

tests/proxy_unit_tests/test_proxy_config_unit_test.py

@@ -5,6 +5,8 @@
 import pytest
 
 from dotenv import load_dotenv
+from moto import mock_aws
+import boto3
 
 import litellm.proxy
 import litellm.proxy.proxy_server
@@ -264,3 +266,68 @@ def test_add_callbacks_invalid_input():
     # Cleanup
     litellm.success_callback = []
     litellm.failure_callback = []
+
+
+@pytest.mark.asyncio
+async def test_reading_configs_with_includes_from_s3():
+    """
+    Test that the config is read correctly from the S3 and read includes
+    """
+    BUCKET_NAME = "config-bucket"
+    config_files = ("config_with_multiple_includes.yaml", "models_file_1.yaml", "models_file_2.yaml")
+    os.environ["LITELLM_CONFIG_BUCKET_NAME"] = BUCKET_NAME
+    os.environ["LITELLM_CONFIG_BUCKET_OBJECT_KEY"] = config_files[0]
+    os.environ["LITELLM_CONFIG_BUCKET_TYPE"] = "S3"
+    with mock_aws():
+        # setup s3 bucket and put config files
+        current_path = os.path.dirname(os.path.abspath(__file__))
+        s3_client = boto3.client("s3")
+        s3_client.create_bucket(Bucket=BUCKET_NAME)
+        for file_name in config_files:
+            config_path = os.path.join(
+                current_path, "example_config_yaml", file_name,
+            )
+            s3_client.put_object(
+                Bucket="config-bucket", Key=file_name, Body=open(config_path, 'rb')
+            )
+        proxy_config_instance = ProxyConfig()    
+        config = await proxy_config_instance.get_config()
+
+        assert config == {
+            'model_list': [
+                {'model_name': 'included-model-1', 'litellm_params': {'model': 'gpt-4'}},
+                {'model_name': 'included-model-2', 'litellm_params': {'model': 'gpt-3.5-turbo'}}
+            ],
+            'litellm_settings': {'callbacks': ['prometheus']},
+        }
+
+    # unset the env variable to avoid side-effects on other tests
+    del os.environ["LITELLM_CONFIG_BUCKET_NAME"]
+    del os.environ["LITELLM_CONFIG_BUCKET_OBJECT_KEY"]
+    del os.environ["LITELLM_CONFIG_BUCKET_TYPE"]
+
+
+@pytest.mark.asyncio
+async def test_reading_configs_from_s3_file_not_found():
+    """
+    Test that the config is not present S3
+    """
+    BUCKET_NAME = "config-bucket"
+    config_files = ("config_with_multiple_includes.yaml", "models_file_1.yaml", "models_file_2.yaml")
+    os.environ["LITELLM_CONFIG_BUCKET_NAME"] = BUCKET_NAME
+    os.environ["LITELLM_CONFIG_BUCKET_OBJECT_KEY"] = config_files[0]
+    os.environ["LITELLM_CONFIG_BUCKET_TYPE"] = "S3"
+    with mock_aws():
+        # setup s3 bucket but do not put file
+        s3_client = boto3.client("s3")
+        s3_client.create_bucket(Bucket=BUCKET_NAME)
+        proxy_config_instance = ProxyConfig()
+        with pytest.raises(Exception) as ex:
+            await proxy_config_instance.get_config()
+        
+        assert str(ex.value) == "Unable to load config from given source."
+
+    # unset the env variable to avoid side-effects on other tests
+    del os.environ["LITELLM_CONFIG_BUCKET_NAME"]
+    del os.environ["LITELLM_CONFIG_BUCKET_OBJECT_KEY"]
+    del os.environ["LITELLM_CONFIG_BUCKET_TYPE"]