From fc0429a888a9d93cdc956c295d279817939cb867 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Thu, 17 Oct 2024 12:59:22 +0530 Subject: [PATCH 01/11] add check for os.environ vars when readin config.yaml --- litellm/proxy/proxy_config.yaml | 5 +++-- litellm/proxy/proxy_server.py | 32 +++++++++++++++++++++++++++++++- 2 files changed, 34 insertions(+), 3 deletions(-) diff --git a/litellm/proxy/proxy_config.yaml b/litellm/proxy/proxy_config.yaml index 7c70332fd25b..36c3935271d2 100644 --- a/litellm/proxy/proxy_config.yaml +++ b/litellm/proxy/proxy_config.yaml @@ -5,6 +5,7 @@ model_list: api_key: fake-key api_base: https://exampleopenaiendpoint-production.up.railway.app/ -litellm_settings: - callbacks: ["arize"] +litellm_settings: + default_internal_user_params: + user_role: os.environ/DEFAULT_USER_ROLE \ No newline at end of file diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 18d51fe41f9d..f1cb223be919 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -1425,7 +1425,7 @@ async def get_config(self, config_file_path: Optional[str] = None) -> dict: else: # if it's not in the config - then add it config[param_name] = param_value - + config = self._check_for_os_environ_vars(config=config) return config async def save_config(self, new_config: dict): @@ -1492,6 +1492,36 @@ def _init_cache( ## INIT PROXY REDIS USAGE CLIENT ## redis_usage_cache = litellm.cache.cache + def _check_for_os_environ_vars( + self, config: dict, depth: int = 0, max_depth: int = 10 + ) -> dict: + """ + Check for os.environ/ variables in the config and replace them with the actual values. + Includes a depth limit to prevent infinite recursion. + + Args: + config (dict): The configuration dictionary to process. + depth (int): Current recursion depth. + max_depth (int): Maximum allowed recursion depth. + + Returns: + dict: Processed configuration dictionary. + """ + if depth > max_depth: + verbose_proxy_logger.warning( + f"Maximum recursion depth ({max_depth}) reached while processing config." + ) + return config + + for key, value in config.items(): + if isinstance(value, dict): + config[key] = self._check_for_os_environ_vars( + config=value, depth=depth + 1, max_depth=max_depth + ) + elif isinstance(value, str) and value.startswith("os.environ/"): + config[key] = get_secret_str(value) + return config + async def load_config( self, router: Optional[litellm.Router], config_file_path: str ): From 00163bcec3237d7fb285f3e23af5be19fe6e38d8 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Thu, 17 Oct 2024 13:48:19 +0530 Subject: [PATCH 02/11] use base class for reading from config.yaml --- .../proxy/common_utils/base_config_class.py | 280 ++++++++++++++++++ .../proxy/common_utils/load_config_utils.py | 77 ----- litellm/proxy/proxy_server.py | 139 +-------- 3 files changed, 283 insertions(+), 213 deletions(-) create mode 100644 litellm/proxy/common_utils/base_config_class.py delete mode 100644 litellm/proxy/common_utils/load_config_utils.py diff --git a/litellm/proxy/common_utils/base_config_class.py b/litellm/proxy/common_utils/base_config_class.py new file mode 100644 index 000000000000..f44ada45604f --- /dev/null +++ b/litellm/proxy/common_utils/base_config_class.py @@ -0,0 +1,280 @@ +import asyncio +import os +from typing import TYPE_CHECKING, Any, Dict, Optional + +import yaml + +from litellm._logging import verbose_proxy_logger +from litellm.secret_managers.main import get_secret_str + +if TYPE_CHECKING: + from litellm.proxy.utils import PrismaClient +else: + PrismaClient = Any + + +class BaseProxyConfig: + """ + Base class for proxy config.yaml + + Handles loading litellm config.yaml from a file, S3 bucket, GCS bucket or DB + """ + + def is_yaml(self, config_file_path: str) -> bool: + if not os.path.isfile(config_file_path): + return False + + _, file_extension = os.path.splitext(config_file_path) + return file_extension.lower() == ".yaml" or file_extension.lower() == ".yml" + + async def get_config( + self, + config_file_path: Optional[str] = None, + ) -> Dict: + """ + Get the config.yaml file contents from a File, S3 bucket, GCS bucket or DB + + - Read from s3/GCS bucket when `LITELLM_CONFIG_BUCKET_NAME` is set in .env + - Default to reading from config file path + """ + # Load existing config + if os.environ.get("LITELLM_CONFIG_BUCKET_NAME") is not None: + bucket_name = os.environ.get("LITELLM_CONFIG_BUCKET_NAME") + object_key = os.environ.get("LITELLM_CONFIG_BUCKET_OBJECT_KEY") + bucket_type = 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 self._get_config_file_contents_from_gcs( + bucket_name=bucket_name, object_key=object_key + ) + else: + config = self._get_file_contents_from_s3( + bucket_name=bucket_name, object_key=object_key + ) + + if config is None: + raise Exception("Unable to load config from given source.") + else: + # default to file + config = await self._get_config_from_file(config_file_path=config_file_path) + + config = self._check_for_os_environ_vars(config=config) + return config + + async def _get_config_from_file( + self, config_file_path: Optional[str] = None + ) -> Dict: + from litellm.proxy.proxy_server import ( + general_settings, + prisma_client, + store_model_in_db, + user_config_file_path, + ) + + file_path = config_file_path or user_config_file_path + if config_file_path is not None: + user_config_file_path = config_file_path + # Load existing config + ## Yaml + if os.path.exists(f"{file_path}"): + with open(f"{file_path}", "r") as config_file: + config = yaml.safe_load(config_file) + else: + config = { + "model_list": [], + "general_settings": {}, + "router_settings": {}, + "litellm_settings": {}, + } + + ## DB + if prisma_client is not None and ( + general_settings.get("store_model_in_db", False) is True + or store_model_in_db is True + ): + config = await self._update_config_from_db( + config=config, prisma_client=prisma_client + ) + return config + + async def _update_config_from_db( + self, config: Dict, prisma_client: PrismaClient + ) -> Dict: + _tasks = [] + keys = [ + "general_settings", + "router_settings", + "litellm_settings", + "environment_variables", + ] + for k in keys: + response = prisma_client.get_generic_data( + key="param_name", value=k, table_name="config" + ) + _tasks.append(response) + + responses = await asyncio.gather(*_tasks) + for response in responses: + if response is not None: + param_name = getattr(response, "param_name", None) + param_value = getattr(response, "param_value", None) + if param_name is not None and param_value is not None: + # check if param_name is already in the config + if param_name in config: + if isinstance(config[param_name], dict): + config[param_name].update(param_value) + else: + config[param_name] = param_value + else: + # if it's not in the config - then add it + config[param_name] = param_value + return config + + def _get_file_contents_from_s3( + self, bucket_name: Optional[str] = None, object_key: Optional[str] = None + ) -> Optional[Dict]: + """ + Get the config.yaml file contents from a S3 bucket + """ + try: + # v0 rely on boto3 for authentication - allowing boto3 to handle IAM credentials etc + import tempfile + + import boto3 + from botocore.config import Config + from botocore.credentials import Credentials + + from litellm.main import bedrock_converse_chat_completion + + credentials: Credentials = ( + bedrock_converse_chat_completion.get_credentials() + ) + s3_client = boto3.client( + "s3", + aws_access_key_id=credentials.access_key, + aws_secret_access_key=credentials.secret_key, + aws_session_token=credentials.token, # Optional, if using temporary credentials + ) + verbose_proxy_logger.debug( + 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}") + + # Read the file contents + file_contents = response["Body"].read().decode("utf-8") + verbose_proxy_logger.debug("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}" + ) + + # Load the YAML file content + with open(temp_file_path, "r") as yaml_file: + config = yaml.safe_load(yaml_file) + + return config + except ImportError as e: + # this is most likely if a user is not using the litellm docker container + verbose_proxy_logger.error(f"ImportError: {str(e)}") + pass + except Exception as e: + verbose_proxy_logger.error(f"Error retrieving file contents: {str(e)}") + return None + + async def _get_config_file_contents_from_gcs( + self, bucket_name: Optional[str] = None, object_key: Optional[str] = None + ) -> Optional[Dict]: + """ + Get the config.yaml file contents from a GCS bucket + """ + try: + from litellm.integrations.gcs_bucket.gcs_bucket import GCSBucketLogger + + gcs_bucket = GCSBucketLogger( + bucket_name=bucket_name, + ) + if object_key is None: + raise Exception(f"Object key is None for {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}") + # 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 = yaml.safe_load(file_contents) + return config + + except Exception as e: + verbose_proxy_logger.error(f"Error retrieving file contents: {str(e)}") + return None + + # # Example usage + # bucket_name = 'litellm-proxy' + # object_key = 'litellm_proxy_config.yaml' + + def _check_for_os_environ_vars( + self, config: dict, depth: int = 0, max_depth: int = 10 + ) -> dict: + """ + Check for os.environ/ variables in the config and replace them with the actual values. + Includes a depth limit to prevent infinite recursion. + + Args: + config (dict): The configuration dictionary to process. + depth (int): Current recursion depth. + max_depth (int): Maximum allowed recursion depth. + + Returns: + dict: Processed configuration dictionary. + """ + if depth > max_depth: + verbose_proxy_logger.warning( + f"Maximum recursion depth ({max_depth}) reached while processing config." + ) + return config + + for key, value in config.items(): + if isinstance(value, dict): + config[key] = self._check_for_os_environ_vars( + config=value, depth=depth + 1, max_depth=max_depth + ) + elif isinstance(value, str) and value.startswith("os.environ/"): + config[key] = get_secret_str(value) + return config + + async def save_config(self, new_config: dict): + """ + Save the config.yaml contents to the DB (if user has opted in) or save in file + """ + from litellm.proxy.proxy_server import ( + general_settings, + prisma_client, + store_model_in_db, + user_config_file_path, + ) + + # Load existing config + ## DB - writes valid config to db + """ + - Do not write restricted params like 'api_key' to the database + - if api_key is passed, save that to the local environment or connected secret manage (maybe expose `litellm.save_secret()`) + """ + if prisma_client is not None and ( + general_settings.get("store_model_in_db", False) is True + or store_model_in_db + ): + # if using - db for config - models are in ModelTable + new_config.pop("model_list", None) + await prisma_client.insert_data(data=new_config, table_name="config") + else: + # Save the updated config - if user is not using a dB + ## YAML + with open(f"{user_config_file_path}", "w") as config_file: + yaml.dump(new_config, config_file, default_flow_style=False) diff --git a/litellm/proxy/common_utils/load_config_utils.py b/litellm/proxy/common_utils/load_config_utils.py deleted file mode 100644 index f262837d9226..000000000000 --- a/litellm/proxy/common_utils/load_config_utils.py +++ /dev/null @@ -1,77 +0,0 @@ -import yaml - -from litellm._logging import verbose_proxy_logger - - -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 - - import boto3 - from botocore.config import Config - from botocore.credentials import Credentials - - from litellm.main import bedrock_converse_chat_completion - - credentials: Credentials = bedrock_converse_chat_completion.get_credentials() - s3_client = boto3.client( - "s3", - aws_access_key_id=credentials.access_key, - aws_secret_access_key=credentials.secret_key, - aws_session_token=credentials.token, # Optional, if using temporary credentials - ) - verbose_proxy_logger.debug( - 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}") - - # Read the file contents - file_contents = response["Body"].read().decode("utf-8") - verbose_proxy_logger.debug("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}") - - # Load the YAML file content - with open(temp_file_path, "r") as yaml_file: - config = yaml.safe_load(yaml_file) - - return config - except ImportError as e: - # this is most likely if a user is not using the litellm docker container - verbose_proxy_logger.error(f"ImportError: {str(e)}") - pass - except Exception as e: - verbose_proxy_logger.error(f"Error retrieving file contents: {str(e)}") - return None - - -async def get_config_file_contents_from_gcs(bucket_name, object_key): - try: - from litellm.integrations.gcs_bucket.gcs_bucket import GCSBucketLogger - - gcs_bucket = GCSBucketLogger( - 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}") - # 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 = yaml.safe_load(file_contents) - return config - - except Exception as e: - verbose_proxy_logger.error(f"Error retrieving file contents: {str(e)}") - return None - - -# # Example usage -# bucket_name = 'litellm-proxy' -# object_key = 'litellm_proxy_config.yaml' diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index f1cb223be919..532baf1bfbdc 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -139,6 +139,7 @@ def generate_feedback_box(): ## Import All Misc routes here ## from litellm.proxy.caching_routes import router as caching_router from litellm.proxy.common_utils.admin_ui_utils import html_form +from litellm.proxy.common_utils.base_config_class import BaseProxyConfig from litellm.proxy.common_utils.callback_utils import ( get_logging_caching_headers, get_remaining_tokens_and_requests_from_request_data, @@ -1358,7 +1359,7 @@ async def _run_background_health_check(): await asyncio.sleep(health_check_interval) -class ProxyConfig: +class ProxyConfig(BaseProxyConfig): """ Abstraction class on top of config loading/updating logic. Gives us one place to control all config updating logic. """ @@ -1366,89 +1367,6 @@ class ProxyConfig: def __init__(self) -> None: pass - def is_yaml(self, config_file_path: str) -> bool: - if not os.path.isfile(config_file_path): - return False - - _, file_extension = os.path.splitext(config_file_path) - return file_extension.lower() == ".yaml" or file_extension.lower() == ".yml" - - async def get_config(self, config_file_path: Optional[str] = None) -> dict: - global prisma_client, user_config_file_path - - file_path = config_file_path or user_config_file_path - if config_file_path is not None: - user_config_file_path = config_file_path - # Load existing config - ## Yaml - if os.path.exists(f"{file_path}"): - with open(f"{file_path}", "r") as config_file: - config = yaml.safe_load(config_file) - else: - config = { - "model_list": [], - "general_settings": {}, - "router_settings": {}, - "litellm_settings": {}, - } - - ## DB - if prisma_client is not None and ( - general_settings.get("store_model_in_db", False) is True - or store_model_in_db is True - ): - _tasks = [] - keys = [ - "general_settings", - "router_settings", - "litellm_settings", - "environment_variables", - ] - for k in keys: - response = prisma_client.get_generic_data( - key="param_name", value=k, table_name="config" - ) - _tasks.append(response) - - responses = await asyncio.gather(*_tasks) - for response in responses: - if response is not None: - param_name = getattr(response, "param_name", None) - param_value = getattr(response, "param_value", None) - if param_name is not None and param_value is not None: - # check if param_name is already in the config - if param_name in config: - if isinstance(config[param_name], dict): - config[param_name].update(param_value) - else: - config[param_name] = param_value - else: - # if it's not in the config - then add it - config[param_name] = param_value - config = self._check_for_os_environ_vars(config=config) - return config - - async def save_config(self, new_config: dict): - global prisma_client, general_settings, user_config_file_path, store_model_in_db - # Load existing config - ## DB - writes valid config to db - """ - - Do not write restricted params like 'api_key' to the database - - if api_key is passed, save that to the local environment or connected secret manage (maybe expose `litellm.save_secret()`) - """ - if prisma_client is not None and ( - general_settings.get("store_model_in_db", False) is True - or store_model_in_db - ): - # if using - db for config - models are in ModelTable - new_config.pop("model_list", None) - await prisma_client.insert_data(data=new_config, table_name="config") - else: - # Save the updated config - if user is not using a dB - ## YAML - with open(f"{user_config_file_path}", "w") as config_file: - yaml.dump(new_config, config_file, default_flow_style=False) - async def load_team_config(self, team_id: str): """ - for a given team id @@ -1492,36 +1410,6 @@ def _init_cache( ## INIT PROXY REDIS USAGE CLIENT ## redis_usage_cache = litellm.cache.cache - def _check_for_os_environ_vars( - self, config: dict, depth: int = 0, max_depth: int = 10 - ) -> dict: - """ - Check for os.environ/ variables in the config and replace them with the actual values. - Includes a depth limit to prevent infinite recursion. - - Args: - config (dict): The configuration dictionary to process. - depth (int): Current recursion depth. - max_depth (int): Maximum allowed recursion depth. - - Returns: - dict: Processed configuration dictionary. - """ - if depth > max_depth: - verbose_proxy_logger.warning( - f"Maximum recursion depth ({max_depth}) reached while processing config." - ) - return config - - for key, value in config.items(): - if isinstance(value, dict): - config[key] = self._check_for_os_environ_vars( - config=value, depth=depth + 1, max_depth=max_depth - ) - elif isinstance(value, str) and value.startswith("os.environ/"): - config[key] = get_secret_str(value) - return config - async def load_config( self, router: Optional[litellm.Router], config_file_path: str ): @@ -1530,28 +1418,7 @@ async def load_config( """ global master_key, user_config_file_path, otel_logging, user_custom_auth, user_custom_auth_path, user_custom_key_generate, user_custom_sso, use_background_health_checks, health_check_interval, use_queue, proxy_budget_rescheduler_max_time, proxy_budget_rescheduler_min_time, ui_access_mode, litellm_master_key_hash, proxy_batch_write_at, disable_spend_logs, prompt_injection_detection_obj, redis_usage_cache, store_model_in_db, premium_user, open_telemetry_logger, health_check_details, callback_settings - # Load existing config - if os.environ.get("LITELLM_CONFIG_BUCKET_NAME") is not None: - bucket_name = os.environ.get("LITELLM_CONFIG_BUCKET_NAME") - object_key = os.environ.get("LITELLM_CONFIG_BUCKET_OBJECT_KEY") - bucket_type = 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 config is None: - raise Exception("Unable to load config from given source.") - else: - # default to file - config = await self.get_config(config_file_path=config_file_path) + config: dict = await self.get_config(config_file_path=config_file_path) ## PRINT YAML FOR CONFIRMING IT WORKS printed_yaml = copy.deepcopy(config) printed_yaml.pop("environment_variables", None) From 33d25d12838df4e9784ec008e902763f2de238f6 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Thu, 17 Oct 2024 13:51:36 +0530 Subject: [PATCH 03/11] fix import --- litellm/proxy/proxy_server.py | 4 ---- 1 file changed, 4 deletions(-) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 532baf1bfbdc..6225b1873579 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -155,10 +155,6 @@ def generate_feedback_box(): _read_request_body, check_file_size_under_limit, ) -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.openai_endpoint_utils import ( remove_sensitive_info_from_deployment, ) From 752c03395812d483320c0653875163f6c90c1cb3 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Thu, 17 Oct 2024 13:54:37 +0530 Subject: [PATCH 04/11] fix linting --- litellm/proxy/common_utils/base_config_class.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/litellm/proxy/common_utils/base_config_class.py b/litellm/proxy/common_utils/base_config_class.py index f44ada45604f..74a259957356 100644 --- a/litellm/proxy/common_utils/base_config_class.py +++ b/litellm/proxy/common_utils/base_config_class.py @@ -183,7 +183,7 @@ def _get_file_contents_from_s3( except ImportError as e: # this is most likely if a user is not using the litellm docker container verbose_proxy_logger.error(f"ImportError: {str(e)}") - pass + return None except Exception as e: verbose_proxy_logger.error(f"Error retrieving file contents: {str(e)}") return None From e0ebf51c9098a8e11e0e34b5d4b7da97567ff0d5 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Thu, 17 Oct 2024 14:26:01 +0530 Subject: [PATCH 05/11] add unit tests for base config class --- .../proxy/common_utils/base_config_class.py | 13 +- litellm/proxy/proxy_config.yaml | 51 ++++++-- .../config_with_env_vars.yaml | 48 ++++++++ .../test_proxy_base_config_unit_test.py | 116 ++++++++++++++++++ 4 files changed, 219 insertions(+), 9 deletions(-) create mode 100644 tests/local_testing/example_config_yaml/config_with_env_vars.yaml create mode 100644 tests/local_testing/test_proxy_base_config_unit_test.py diff --git a/litellm/proxy/common_utils/base_config_class.py b/litellm/proxy/common_utils/base_config_class.py index 74a259957356..3a7f615512fb 100644 --- a/litellm/proxy/common_utils/base_config_class.py +++ b/litellm/proxy/common_utils/base_config_class.py @@ -5,7 +5,7 @@ import yaml from litellm._logging import verbose_proxy_logger -from litellm.secret_managers.main import get_secret_str +from litellm.secret_managers.main import get_secret, get_secret_str if TYPE_CHECKING: from litellm.proxy.utils import PrismaClient @@ -81,6 +81,8 @@ async def _get_config_from_file( if os.path.exists(f"{file_path}"): with open(f"{file_path}", "r") as config_file: config = yaml.safe_load(config_file) + elif config_file_path is not None: + raise Exception(f"Config file not found at {config_file_path}") else: config = { "model_list": [], @@ -245,8 +247,15 @@ def _check_for_os_environ_vars( config[key] = self._check_for_os_environ_vars( config=value, depth=depth + 1, max_depth=max_depth ) + elif isinstance(value, list): + for item in value: + if isinstance(item, dict): + item = self._check_for_os_environ_vars( + config=item, depth=depth + 1, max_depth=max_depth + ) + # if the value is a string and starts with "os.environ/" - then it's an environment variable elif isinstance(value, str) and value.startswith("os.environ/"): - config[key] = get_secret_str(value) + config[key] = get_secret(value) return config async def save_config(self, new_config: dict): diff --git a/litellm/proxy/proxy_config.yaml b/litellm/proxy/proxy_config.yaml index 36c3935271d2..bae738c735d7 100644 --- a/litellm/proxy/proxy_config.yaml +++ b/litellm/proxy/proxy_config.yaml @@ -1,11 +1,48 @@ model_list: - - model_name: fake-openai-endpoint + ################################################################################ + # Azure + - model_name: gpt-4o-mini litellm_params: - model: openai/fake - api_key: fake-key - api_base: https://exampleopenaiendpoint-production.up.railway.app/ - + model: azure/gpt-4o-mini + api_base: https://amazin-prod.openai.azure.com + api_key: "os.environ/AZURE_GPT_4O" + deployment_id: gpt-4o-mini + - model_name: gpt-4o + litellm_params: + model: azure/gpt-4o + api_base: https://very-cool-prod.openai.azure.com + api_key: "os.environ/AZURE_GPT_4O" + deployment_id: gpt-4o -litellm_settings: + ################################################################################ + # Fireworks + - model_name: fireworks-llama-v3p1-405b-instruct + litellm_params: + model: fireworks_ai/accounts/fireworks/models/llama-v3p1-405b-instruct + api_key: "os.environ/FIREWORKS" + - model_name: fireworks-llama-v3p1-70b-instruct + litellm_params: + model: fireworks_ai/accounts/fireworks/models/llama-v3p1-70b-instruct + api_key: "os.environ/FIREWORKS" + +general_settings: + alerting_threshold: 300 # sends alerts if requests hang for 5min+ and responses take 5min+ +litellm_settings: # module level litellm settings - https://github.com/BerriAI/litellm/blob/main/litellm/__init__.py + success_callback: ["prometheus"] + service_callback: ["prometheus_system"] + drop_params: False # Raise an exception if the openai param being passed in isn't supported. + cache: false default_internal_user_params: - user_role: os.environ/DEFAULT_USER_ROLE \ No newline at end of file + user_role: os.environ/DEFAULT_USER_ROLE + + success_callback: ["s3"] + s3_callback_params: + s3_bucket_name: logs-bucket-litellm # AWS Bucket Name for S3 + s3_region_name: us-west-2 # AWS Region Name for S3 + s3_aws_access_key_id: os.environ/AWS_ACCESS_KEY_ID # us os.environ/ to pass environment variables. This is AWS Access Key ID for S3 + s3_aws_secret_access_key: os.environ/AWS_SECRET_ACCESS_KEY # AWS Secret Access Key for S3 + s3_path: my-test-path # [OPTIONAL] set path in bucket you want to write logs to + s3_endpoint_url: https://s3.amazonaws.com # [OPTIONAL] S3 endpoint URL, if you want to use Backblaze/cloudflare s3 buckets + +router_settings: + routing_strategy: simple-shuffle # "simple-shuffle" shown to result in highest throughput. https://docs.litellm.ai/docs/proxy/configs#load-balancing diff --git a/tests/local_testing/example_config_yaml/config_with_env_vars.yaml b/tests/local_testing/example_config_yaml/config_with_env_vars.yaml new file mode 100644 index 000000000000..bae738c735d7 --- /dev/null +++ b/tests/local_testing/example_config_yaml/config_with_env_vars.yaml @@ -0,0 +1,48 @@ +model_list: + ################################################################################ + # Azure + - model_name: gpt-4o-mini + litellm_params: + model: azure/gpt-4o-mini + api_base: https://amazin-prod.openai.azure.com + api_key: "os.environ/AZURE_GPT_4O" + deployment_id: gpt-4o-mini + - model_name: gpt-4o + litellm_params: + model: azure/gpt-4o + api_base: https://very-cool-prod.openai.azure.com + api_key: "os.environ/AZURE_GPT_4O" + deployment_id: gpt-4o + + ################################################################################ + # Fireworks + - model_name: fireworks-llama-v3p1-405b-instruct + litellm_params: + model: fireworks_ai/accounts/fireworks/models/llama-v3p1-405b-instruct + api_key: "os.environ/FIREWORKS" + - model_name: fireworks-llama-v3p1-70b-instruct + litellm_params: + model: fireworks_ai/accounts/fireworks/models/llama-v3p1-70b-instruct + api_key: "os.environ/FIREWORKS" + +general_settings: + alerting_threshold: 300 # sends alerts if requests hang for 5min+ and responses take 5min+ +litellm_settings: # module level litellm settings - https://github.com/BerriAI/litellm/blob/main/litellm/__init__.py + success_callback: ["prometheus"] + service_callback: ["prometheus_system"] + drop_params: False # Raise an exception if the openai param being passed in isn't supported. + cache: false + default_internal_user_params: + user_role: os.environ/DEFAULT_USER_ROLE + + success_callback: ["s3"] + s3_callback_params: + s3_bucket_name: logs-bucket-litellm # AWS Bucket Name for S3 + s3_region_name: us-west-2 # AWS Region Name for S3 + s3_aws_access_key_id: os.environ/AWS_ACCESS_KEY_ID # us os.environ/ to pass environment variables. This is AWS Access Key ID for S3 + s3_aws_secret_access_key: os.environ/AWS_SECRET_ACCESS_KEY # AWS Secret Access Key for S3 + s3_path: my-test-path # [OPTIONAL] set path in bucket you want to write logs to + s3_endpoint_url: https://s3.amazonaws.com # [OPTIONAL] S3 endpoint URL, if you want to use Backblaze/cloudflare s3 buckets + +router_settings: + routing_strategy: simple-shuffle # "simple-shuffle" shown to result in highest throughput. https://docs.litellm.ai/docs/proxy/configs#load-balancing diff --git a/tests/local_testing/test_proxy_base_config_unit_test.py b/tests/local_testing/test_proxy_base_config_unit_test.py new file mode 100644 index 000000000000..a579b57bc27e --- /dev/null +++ b/tests/local_testing/test_proxy_base_config_unit_test.py @@ -0,0 +1,116 @@ +import os +import sys +import traceback +from unittest import mock +import pytest + +from dotenv import load_dotenv + +import litellm.proxy +import litellm.proxy.proxy_server + +load_dotenv() +import io +import os + +# this file is to test litellm/proxy + +sys.path.insert( + 0, os.path.abspath("../..") +) # Adds the parent directory to the system path +import asyncio +import logging + +from litellm.proxy.common_utils.base_config_class import BaseProxyConfig + +# get all the files from example_config_yaml +files = os.listdir("example_config_yaml") +print(files) + + +@pytest.mark.asyncio +async def test_basic_reading_configs_from_files(): + """ + Test that the config is read correctly from the files in the example_config_yaml folder + """ + _base_config = BaseProxyConfig() + current_path = os.path.dirname(os.path.abspath(__file__)) + + for file in files: + config_path = os.path.join(current_path, "example_config_yaml", file) + config = await _base_config.get_config(config_file_path=config_path) + print(config) + + +@pytest.mark.asyncio +async def test_read_config_from_bad_file_path(): + """ + Raise an exception if the file path is not valid + """ + _base_config = BaseProxyConfig() + config_path = "non-existent-file.yaml" + with pytest.raises(Exception): + config = await _base_config.get_config(config_file_path=config_path) + + +@pytest.mark.asyncio +async def test_read_config_file_with_os_environ_vars(): + """ + Ensures os.environ variables are read correctly from config.yaml + Following vars are set as os.environ variables in the config.yaml file + - DEFAULT_USER_ROLE + - AWS_ACCESS_KEY_ID + - AWS_SECRET_ACCESS_KEY + - AZURE_GPT_4O + - FIREWORKS + """ + + _env_vars_for_testing = { + "DEFAULT_USER_ROLE": "admin", + "AWS_ACCESS_KEY_ID": "1234567890", + "AWS_SECRET_ACCESS_KEY": "1234567890", + "AZURE_GPT_4O": "1234567890", + "FIREWORKS": "1234567890", + } + + _old_env_vars = {} + for key, value in _env_vars_for_testing.items(): + if key in os.environ: + _old_env_vars[key] = os.environ.get(key) + os.environ[key] = value + + # Read config + _base_config = BaseProxyConfig() + current_path = os.path.dirname(os.path.abspath(__file__)) + config_path = os.path.join( + current_path, "example_config_yaml", "config_with_env_vars.yaml" + ) + config = await _base_config.get_config(config_file_path=config_path) + print(config) + + # Add assertions + assert ( + config["litellm_settings"]["default_internal_user_params"]["user_role"] + == "admin" + ) + assert ( + config["litellm_settings"]["s3_callback_params"]["s3_aws_access_key_id"] + == "1234567890" + ) + assert ( + config["litellm_settings"]["s3_callback_params"]["s3_aws_secret_access_key"] + == "1234567890" + ) + + for model in config["model_list"]: + if "azure" in model["litellm_params"]["model"]: + assert model["litellm_params"]["api_key"] == "1234567890" + elif "fireworks" in model["litellm_params"]["model"]: + assert model["litellm_params"]["api_key"] == "1234567890" + + # cleanup + for key, value in _env_vars_for_testing.items(): + if key in _old_env_vars: + os.environ[key] = _old_env_vars[key] + else: + del os.environ[key] From 5a431a447182af444cc15a1f7951f2d363a5d2a9 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Thu, 17 Oct 2024 14:30:13 +0530 Subject: [PATCH 06/11] fix order of reading elements from config.yaml --- litellm/proxy/common_utils/base_config_class.py | 10 ++++++++++ litellm/proxy/proxy_server.py | 8 -------- 2 files changed, 10 insertions(+), 8 deletions(-) diff --git a/litellm/proxy/common_utils/base_config_class.py b/litellm/proxy/common_utils/base_config_class.py index 3a7f615512fb..33ff1b9e382d 100644 --- a/litellm/proxy/common_utils/base_config_class.py +++ b/litellm/proxy/common_utils/base_config_class.py @@ -1,4 +1,6 @@ import asyncio +import copy +import json import os from typing import TYPE_CHECKING, Any, Dict, Optional @@ -60,6 +62,14 @@ async def get_config( # default to file config = await self._get_config_from_file(config_file_path=config_file_path) + ## PRINT YAML FOR CONFIRMING IT WORKS + printed_yaml = copy.deepcopy(config) + printed_yaml.pop("environment_variables", None) + + verbose_proxy_logger.debug( + f"Loaded config YAML (api_key and environment_variables are not shown):\n{json.dumps(printed_yaml, indent=2)}" + ) + config = self._check_for_os_environ_vars(config=config) return config diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 6225b1873579..a778eff02117 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -1415,14 +1415,6 @@ async def load_config( global master_key, user_config_file_path, otel_logging, user_custom_auth, user_custom_auth_path, user_custom_key_generate, user_custom_sso, use_background_health_checks, health_check_interval, use_queue, proxy_budget_rescheduler_max_time, proxy_budget_rescheduler_min_time, ui_access_mode, litellm_master_key_hash, proxy_batch_write_at, disable_spend_logs, prompt_injection_detection_obj, redis_usage_cache, store_model_in_db, premium_user, open_telemetry_logger, health_check_details, callback_settings config: dict = await self.get_config(config_file_path=config_file_path) - ## PRINT YAML FOR CONFIRMING IT WORKS - printed_yaml = copy.deepcopy(config) - printed_yaml.pop("environment_variables", None) - - verbose_proxy_logger.debug( - f"Loaded config YAML (api_key and environment_variables are not shown):\n{json.dumps(printed_yaml, indent=2)}" - ) - ## ENVIRONMENT VARIABLES environment_variables = config.get("environment_variables", None) if environment_variables: From 37e455adef884233217cf8dfe5e71c024754c1ae Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Thu, 17 Oct 2024 14:58:14 +0530 Subject: [PATCH 07/11] unit tests for reading configs from files --- .../local_testing/test_proxy_base_config_unit_test.py | 11 ++++++----- 1 file changed, 6 insertions(+), 5 deletions(-) diff --git a/tests/local_testing/test_proxy_base_config_unit_test.py b/tests/local_testing/test_proxy_base_config_unit_test.py index a579b57bc27e..972f0d67dfab 100644 --- a/tests/local_testing/test_proxy_base_config_unit_test.py +++ b/tests/local_testing/test_proxy_base_config_unit_test.py @@ -23,10 +23,6 @@ from litellm.proxy.common_utils.base_config_class import BaseProxyConfig -# get all the files from example_config_yaml -files = os.listdir("example_config_yaml") -print(files) - @pytest.mark.asyncio async def test_basic_reading_configs_from_files(): @@ -35,9 +31,14 @@ async def test_basic_reading_configs_from_files(): """ _base_config = BaseProxyConfig() current_path = os.path.dirname(os.path.abspath(__file__)) + example_config_yaml_path = os.path.join(current_path, "example_config_yaml") + + # get all the files from example_config_yaml + files = os.listdir(example_config_yaml_path) + print(files) for file in files: - config_path = os.path.join(current_path, "example_config_yaml", file) + config_path = os.path.join(example_config_yaml_path, file) config = await _base_config.get_config(config_file_path=config_path) print(config) From c7242d815165bb0d0adde5fc4135eb9474d79e4e Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Thu, 17 Oct 2024 15:14:06 +0530 Subject: [PATCH 08/11] fix user_config_file_path --- litellm/proxy/proxy_server.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index a778eff02117..5474e402461d 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -461,7 +461,7 @@ async def redirect_ui_middleware(request: Request, call_next): user_telemetry = True user_config = None user_headers = None -user_config_file_path = f"config_{int(time.time())}.yaml" +user_config_file_path: Optional[str] = None local_logging = True # writes logs to a local api_log.json file for debugging experimental = False #### GLOBAL VARIABLES #### From 2993fae515620a9945a1c118f841c419e0969bca Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Fri, 18 Oct 2024 08:48:07 +0530 Subject: [PATCH 09/11] use simpler implementation --- .../proxy/common_utils/base_config_class.py | 299 ------------------ .../proxy/common_utils/load_config_utils.py | 77 +++++ litellm/proxy/proxy_server.py | 160 +++++++++- 3 files changed, 234 insertions(+), 302 deletions(-) delete mode 100644 litellm/proxy/common_utils/base_config_class.py create mode 100644 litellm/proxy/common_utils/load_config_utils.py diff --git a/litellm/proxy/common_utils/base_config_class.py b/litellm/proxy/common_utils/base_config_class.py deleted file mode 100644 index 33ff1b9e382d..000000000000 --- a/litellm/proxy/common_utils/base_config_class.py +++ /dev/null @@ -1,299 +0,0 @@ -import asyncio -import copy -import json -import os -from typing import TYPE_CHECKING, Any, Dict, Optional - -import yaml - -from litellm._logging import verbose_proxy_logger -from litellm.secret_managers.main import get_secret, get_secret_str - -if TYPE_CHECKING: - from litellm.proxy.utils import PrismaClient -else: - PrismaClient = Any - - -class BaseProxyConfig: - """ - Base class for proxy config.yaml - - Handles loading litellm config.yaml from a file, S3 bucket, GCS bucket or DB - """ - - def is_yaml(self, config_file_path: str) -> bool: - if not os.path.isfile(config_file_path): - return False - - _, file_extension = os.path.splitext(config_file_path) - return file_extension.lower() == ".yaml" or file_extension.lower() == ".yml" - - async def get_config( - self, - config_file_path: Optional[str] = None, - ) -> Dict: - """ - Get the config.yaml file contents from a File, S3 bucket, GCS bucket or DB - - - Read from s3/GCS bucket when `LITELLM_CONFIG_BUCKET_NAME` is set in .env - - Default to reading from config file path - """ - # Load existing config - if os.environ.get("LITELLM_CONFIG_BUCKET_NAME") is not None: - bucket_name = os.environ.get("LITELLM_CONFIG_BUCKET_NAME") - object_key = os.environ.get("LITELLM_CONFIG_BUCKET_OBJECT_KEY") - bucket_type = 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 self._get_config_file_contents_from_gcs( - bucket_name=bucket_name, object_key=object_key - ) - else: - config = self._get_file_contents_from_s3( - bucket_name=bucket_name, object_key=object_key - ) - - if config is None: - raise Exception("Unable to load config from given source.") - else: - # default to file - config = await self._get_config_from_file(config_file_path=config_file_path) - - ## PRINT YAML FOR CONFIRMING IT WORKS - printed_yaml = copy.deepcopy(config) - printed_yaml.pop("environment_variables", None) - - verbose_proxy_logger.debug( - f"Loaded config YAML (api_key and environment_variables are not shown):\n{json.dumps(printed_yaml, indent=2)}" - ) - - config = self._check_for_os_environ_vars(config=config) - return config - - async def _get_config_from_file( - self, config_file_path: Optional[str] = None - ) -> Dict: - from litellm.proxy.proxy_server import ( - general_settings, - prisma_client, - store_model_in_db, - user_config_file_path, - ) - - file_path = config_file_path or user_config_file_path - if config_file_path is not None: - user_config_file_path = config_file_path - # Load existing config - ## Yaml - if os.path.exists(f"{file_path}"): - with open(f"{file_path}", "r") as config_file: - config = yaml.safe_load(config_file) - elif config_file_path is not None: - raise Exception(f"Config file not found at {config_file_path}") - else: - config = { - "model_list": [], - "general_settings": {}, - "router_settings": {}, - "litellm_settings": {}, - } - - ## DB - if prisma_client is not None and ( - general_settings.get("store_model_in_db", False) is True - or store_model_in_db is True - ): - config = await self._update_config_from_db( - config=config, prisma_client=prisma_client - ) - return config - - async def _update_config_from_db( - self, config: Dict, prisma_client: PrismaClient - ) -> Dict: - _tasks = [] - keys = [ - "general_settings", - "router_settings", - "litellm_settings", - "environment_variables", - ] - for k in keys: - response = prisma_client.get_generic_data( - key="param_name", value=k, table_name="config" - ) - _tasks.append(response) - - responses = await asyncio.gather(*_tasks) - for response in responses: - if response is not None: - param_name = getattr(response, "param_name", None) - param_value = getattr(response, "param_value", None) - if param_name is not None and param_value is not None: - # check if param_name is already in the config - if param_name in config: - if isinstance(config[param_name], dict): - config[param_name].update(param_value) - else: - config[param_name] = param_value - else: - # if it's not in the config - then add it - config[param_name] = param_value - return config - - def _get_file_contents_from_s3( - self, bucket_name: Optional[str] = None, object_key: Optional[str] = None - ) -> Optional[Dict]: - """ - Get the config.yaml file contents from a S3 bucket - """ - try: - # v0 rely on boto3 for authentication - allowing boto3 to handle IAM credentials etc - import tempfile - - import boto3 - from botocore.config import Config - from botocore.credentials import Credentials - - from litellm.main import bedrock_converse_chat_completion - - credentials: Credentials = ( - bedrock_converse_chat_completion.get_credentials() - ) - s3_client = boto3.client( - "s3", - aws_access_key_id=credentials.access_key, - aws_secret_access_key=credentials.secret_key, - aws_session_token=credentials.token, # Optional, if using temporary credentials - ) - verbose_proxy_logger.debug( - 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}") - - # Read the file contents - file_contents = response["Body"].read().decode("utf-8") - verbose_proxy_logger.debug("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}" - ) - - # Load the YAML file content - with open(temp_file_path, "r") as yaml_file: - config = yaml.safe_load(yaml_file) - - return config - except ImportError as e: - # this is most likely if a user is not using the litellm docker container - verbose_proxy_logger.error(f"ImportError: {str(e)}") - return None - except Exception as e: - verbose_proxy_logger.error(f"Error retrieving file contents: {str(e)}") - return None - - async def _get_config_file_contents_from_gcs( - self, bucket_name: Optional[str] = None, object_key: Optional[str] = None - ) -> Optional[Dict]: - """ - Get the config.yaml file contents from a GCS bucket - """ - try: - from litellm.integrations.gcs_bucket.gcs_bucket import GCSBucketLogger - - gcs_bucket = GCSBucketLogger( - bucket_name=bucket_name, - ) - if object_key is None: - raise Exception(f"Object key is None for {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}") - # 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 = yaml.safe_load(file_contents) - return config - - except Exception as e: - verbose_proxy_logger.error(f"Error retrieving file contents: {str(e)}") - return None - - # # Example usage - # bucket_name = 'litellm-proxy' - # object_key = 'litellm_proxy_config.yaml' - - def _check_for_os_environ_vars( - self, config: dict, depth: int = 0, max_depth: int = 10 - ) -> dict: - """ - Check for os.environ/ variables in the config and replace them with the actual values. - Includes a depth limit to prevent infinite recursion. - - Args: - config (dict): The configuration dictionary to process. - depth (int): Current recursion depth. - max_depth (int): Maximum allowed recursion depth. - - Returns: - dict: Processed configuration dictionary. - """ - if depth > max_depth: - verbose_proxy_logger.warning( - f"Maximum recursion depth ({max_depth}) reached while processing config." - ) - return config - - for key, value in config.items(): - if isinstance(value, dict): - config[key] = self._check_for_os_environ_vars( - config=value, depth=depth + 1, max_depth=max_depth - ) - elif isinstance(value, list): - for item in value: - if isinstance(item, dict): - item = self._check_for_os_environ_vars( - config=item, depth=depth + 1, max_depth=max_depth - ) - # if the value is a string and starts with "os.environ/" - then it's an environment variable - elif isinstance(value, str) and value.startswith("os.environ/"): - config[key] = get_secret(value) - return config - - async def save_config(self, new_config: dict): - """ - Save the config.yaml contents to the DB (if user has opted in) or save in file - """ - from litellm.proxy.proxy_server import ( - general_settings, - prisma_client, - store_model_in_db, - user_config_file_path, - ) - - # Load existing config - ## DB - writes valid config to db - """ - - Do not write restricted params like 'api_key' to the database - - if api_key is passed, save that to the local environment or connected secret manage (maybe expose `litellm.save_secret()`) - """ - if prisma_client is not None and ( - general_settings.get("store_model_in_db", False) is True - or store_model_in_db - ): - # if using - db for config - models are in ModelTable - new_config.pop("model_list", None) - await prisma_client.insert_data(data=new_config, table_name="config") - else: - # Save the updated config - if user is not using a dB - ## YAML - with open(f"{user_config_file_path}", "w") as config_file: - yaml.dump(new_config, config_file, default_flow_style=False) diff --git a/litellm/proxy/common_utils/load_config_utils.py b/litellm/proxy/common_utils/load_config_utils.py new file mode 100644 index 000000000000..f262837d9226 --- /dev/null +++ b/litellm/proxy/common_utils/load_config_utils.py @@ -0,0 +1,77 @@ +import yaml + +from litellm._logging import verbose_proxy_logger + + +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 + + import boto3 + from botocore.config import Config + from botocore.credentials import Credentials + + from litellm.main import bedrock_converse_chat_completion + + credentials: Credentials = bedrock_converse_chat_completion.get_credentials() + s3_client = boto3.client( + "s3", + aws_access_key_id=credentials.access_key, + aws_secret_access_key=credentials.secret_key, + aws_session_token=credentials.token, # Optional, if using temporary credentials + ) + verbose_proxy_logger.debug( + 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}") + + # Read the file contents + file_contents = response["Body"].read().decode("utf-8") + verbose_proxy_logger.debug("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}") + + # Load the YAML file content + with open(temp_file_path, "r") as yaml_file: + config = yaml.safe_load(yaml_file) + + return config + except ImportError as e: + # this is most likely if a user is not using the litellm docker container + verbose_proxy_logger.error(f"ImportError: {str(e)}") + pass + except Exception as e: + verbose_proxy_logger.error(f"Error retrieving file contents: {str(e)}") + return None + + +async def get_config_file_contents_from_gcs(bucket_name, object_key): + try: + from litellm.integrations.gcs_bucket.gcs_bucket import GCSBucketLogger + + gcs_bucket = GCSBucketLogger( + 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}") + # 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 = yaml.safe_load(file_contents) + return config + + except Exception as e: + verbose_proxy_logger.error(f"Error retrieving file contents: {str(e)}") + return None + + +# # Example usage +# bucket_name = 'litellm-proxy' +# object_key = 'litellm_proxy_config.yaml' diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 5474e402461d..31b269e7eca5 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -139,7 +139,6 @@ def generate_feedback_box(): ## Import All Misc routes here ## from litellm.proxy.caching_routes import router as caching_router from litellm.proxy.common_utils.admin_ui_utils import html_form -from litellm.proxy.common_utils.base_config_class import BaseProxyConfig from litellm.proxy.common_utils.callback_utils import ( get_logging_caching_headers, get_remaining_tokens_and_requests_from_request_data, @@ -155,6 +154,10 @@ def generate_feedback_box(): _read_request_body, check_file_size_under_limit, ) +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.openai_endpoint_utils import ( remove_sensitive_info_from_deployment, ) @@ -1355,7 +1358,7 @@ async def _run_background_health_check(): await asyncio.sleep(health_check_interval) -class ProxyConfig(BaseProxyConfig): +class ProxyConfig: """ Abstraction class on top of config loading/updating logic. Gives us one place to control all config updating logic. """ @@ -1363,6 +1366,126 @@ class ProxyConfig(BaseProxyConfig): def __init__(self) -> None: pass + def is_yaml(self, config_file_path: str) -> bool: + if not os.path.isfile(config_file_path): + return False + + _, file_extension = os.path.splitext(config_file_path) + return file_extension.lower() == ".yaml" or file_extension.lower() == ".yml" + + async def get_config(self, config_file_path: Optional[str] = None) -> dict: + global prisma_client, user_config_file_path + + file_path = config_file_path or user_config_file_path + if config_file_path is not None: + user_config_file_path = config_file_path + # Load existing config + ## Yaml + if os.path.exists(f"{file_path}"): + with open(f"{file_path}", "r") as config_file: + config = yaml.safe_load(config_file) + else: + config = { + "model_list": [], + "general_settings": {}, + "router_settings": {}, + "litellm_settings": {}, + } + + ## DB + if prisma_client is not None and ( + general_settings.get("store_model_in_db", False) is True + or store_model_in_db is True + ): + _tasks = [] + keys = [ + "general_settings", + "router_settings", + "litellm_settings", + "environment_variables", + ] + for k in keys: + response = prisma_client.get_generic_data( + key="param_name", value=k, table_name="config" + ) + _tasks.append(response) + + responses = await asyncio.gather(*_tasks) + for response in responses: + if response is not None: + param_name = getattr(response, "param_name", None) + param_value = getattr(response, "param_value", None) + if param_name is not None and param_value is not None: + # check if param_name is already in the config + if param_name in config: + if isinstance(config[param_name], dict): + config[param_name].update(param_value) + else: + config[param_name] = param_value + else: + # if it's not in the config - then add it + config[param_name] = param_value + + return config + + async def save_config(self, new_config: dict): + global prisma_client, general_settings, user_config_file_path, store_model_in_db + # Load existing config + ## DB - writes valid config to db + """ + - Do not write restricted params like 'api_key' to the database + - if api_key is passed, save that to the local environment or connected secret manage (maybe expose `litellm.save_secret()`) + """ + if prisma_client is not None and ( + general_settings.get("store_model_in_db", False) is True + or store_model_in_db + ): + # if using - db for config - models are in ModelTable + new_config.pop("model_list", None) + await prisma_client.insert_data(data=new_config, table_name="config") + else: + # Save the updated config - if user is not using a dB + ## YAML + with open(f"{user_config_file_path}", "w") as config_file: + yaml.dump(new_config, config_file, default_flow_style=False) + + def _check_for_os_environ_vars( + self, config: dict, depth: int = 0, max_depth: int = 10 + ) -> dict: + """ + Check for os.environ/ variables in the config and replace them with the actual values. + Includes a depth limit to prevent infinite recursion. + + Args: + config (dict): The configuration dictionary to process. + depth (int): Current recursion depth. + max_depth (int): Maximum allowed recursion depth. + + Returns: + dict: Processed configuration dictionary. + """ + if depth > max_depth: + verbose_proxy_logger.warning( + f"Maximum recursion depth ({max_depth}) reached while processing config." + ) + return config + + for key, value in config.items(): + if isinstance(value, dict): + config[key] = self._check_for_os_environ_vars( + config=value, depth=depth + 1, max_depth=max_depth + ) + elif isinstance(value, list): + for item in value: + if isinstance(item, dict): + item = self._check_for_os_environ_vars( + config=item, depth=depth + 1, max_depth=max_depth + ) + # if the value is a string and starts with "os.environ/" - then it's an environment variable + elif isinstance(value, str) and value.startswith("os.environ/"): + config[key] = get_secret(value) + return config + async def load_team_config(self, team_id: str): """ - for a given team id @@ -1414,7 +1537,38 @@ async def load_config( """ global master_key, user_config_file_path, otel_logging, user_custom_auth, user_custom_auth_path, user_custom_key_generate, user_custom_sso, use_background_health_checks, health_check_interval, use_queue, proxy_budget_rescheduler_max_time, proxy_budget_rescheduler_min_time, ui_access_mode, litellm_master_key_hash, proxy_batch_write_at, disable_spend_logs, prompt_injection_detection_obj, redis_usage_cache, store_model_in_db, premium_user, open_telemetry_logger, health_check_details, callback_settings - config: dict = await self.get_config(config_file_path=config_file_path) + # Load existing config + if os.environ.get("LITELLM_CONFIG_BUCKET_NAME") is not None: + bucket_name = os.environ.get("LITELLM_CONFIG_BUCKET_NAME") + object_key = os.environ.get("LITELLM_CONFIG_BUCKET_OBJECT_KEY") + bucket_type = 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 config is None: + raise Exception("Unable to load config from given source.") + else: + # default to file + config = await self.get_config(config_file_path=config_file_path) + ## PRINT YAML FOR CONFIRMING IT WORKS + printed_yaml = copy.deepcopy(config) + printed_yaml.pop("environment_variables", None) + + verbose_proxy_logger.debug( + f"Loaded config YAML (api_key and environment_variables are not shown):\n{json.dumps(printed_yaml, indent=2)}" + ) + + config = self._check_for_os_environ_vars(config=config) + ## ENVIRONMENT VARIABLES environment_variables = config.get("environment_variables", None) if environment_variables: From d2e6ea05e4b9676daf00fb31fb0eed51463c20a3 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Fri, 18 Oct 2024 08:58:30 +0530 Subject: [PATCH 10/11] use helper to get_config --- litellm/proxy/proxy_server.py | 47 +++++++++++++++---- ...test.py => test_proxy_config_unit_test.py} | 8 ++-- 2 files changed, 43 insertions(+), 12 deletions(-) rename tests/local_testing/{test_proxy_base_config_unit_test.py => test_proxy_config_unit_test.py} (94%) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 31b269e7eca5..569fa9b56001 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -1373,7 +1373,19 @@ def is_yaml(self, config_file_path: str) -> bool: _, file_extension = os.path.splitext(config_file_path) return file_extension.lower() == ".yaml" or file_extension.lower() == ".yml" - async def get_config(self, config_file_path: Optional[str] = None) -> dict: + async def _get_config_from_file( + self, config_file_path: Optional[str] = None + ) -> dict: + """ + Given a config file path, load the config from the file. + + If `store_model_in_db` is True, then read the DB and update the config with the DB values. + + Args: + config_file_path (str): path to the config file + Returns: + dict: config + """ global prisma_client, user_config_file_path file_path = config_file_path or user_config_file_path @@ -1529,14 +1541,21 @@ def _init_cache( ## INIT PROXY REDIS USAGE CLIENT ## redis_usage_cache = litellm.cache.cache - async def load_config( - self, router: Optional[litellm.Router], config_file_path: str - ): - """ - Load config values into proxy global state + async def get_config(self, config_file_path: Optional[str] = None) -> dict: """ - global master_key, user_config_file_path, otel_logging, user_custom_auth, user_custom_auth_path, user_custom_key_generate, user_custom_sso, use_background_health_checks, health_check_interval, use_queue, proxy_budget_rescheduler_max_time, proxy_budget_rescheduler_min_time, ui_access_mode, litellm_master_key_hash, proxy_batch_write_at, disable_spend_logs, prompt_injection_detection_obj, redis_usage_cache, store_model_in_db, premium_user, open_telemetry_logger, health_check_details, callback_settings + Load config file + Supports reading from: + - .yaml file paths + - LiteLLM connected DB + - GCS + - S3 + Args: + config_file_path (str): path to the config file + Returns: + dict: config + + """ # Load existing config if os.environ.get("LITELLM_CONFIG_BUCKET_NAME") is not None: bucket_name = os.environ.get("LITELLM_CONFIG_BUCKET_NAME") @@ -1558,7 +1577,7 @@ async def load_config( raise Exception("Unable to load config from given source.") else: # default to file - config = await self.get_config(config_file_path=config_file_path) + config = await self._get_config_from_file(config_file_path=config_file_path) ## PRINT YAML FOR CONFIRMING IT WORKS printed_yaml = copy.deepcopy(config) printed_yaml.pop("environment_variables", None) @@ -1569,6 +1588,18 @@ async def load_config( config = self._check_for_os_environ_vars(config=config) + return config + + async def load_config( + self, router: Optional[litellm.Router], config_file_path: str + ): + """ + Load config values into proxy global state + """ + global master_key, user_config_file_path, otel_logging, user_custom_auth, user_custom_auth_path, user_custom_key_generate, user_custom_sso, use_background_health_checks, health_check_interval, use_queue, proxy_budget_rescheduler_max_time, proxy_budget_rescheduler_min_time, ui_access_mode, litellm_master_key_hash, proxy_batch_write_at, disable_spend_logs, prompt_injection_detection_obj, redis_usage_cache, store_model_in_db, premium_user, open_telemetry_logger, health_check_details, callback_settings + + config: dict = await self.get_config(config_file_path=config_file_path) + ## ENVIRONMENT VARIABLES environment_variables = config.get("environment_variables", None) if environment_variables: diff --git a/tests/local_testing/test_proxy_base_config_unit_test.py b/tests/local_testing/test_proxy_config_unit_test.py similarity index 94% rename from tests/local_testing/test_proxy_base_config_unit_test.py rename to tests/local_testing/test_proxy_config_unit_test.py index 972f0d67dfab..efe51fe6a7b8 100644 --- a/tests/local_testing/test_proxy_base_config_unit_test.py +++ b/tests/local_testing/test_proxy_config_unit_test.py @@ -21,7 +21,7 @@ import asyncio import logging -from litellm.proxy.common_utils.base_config_class import BaseProxyConfig +from litellm.proxy.proxy_server import ProxyConfig @pytest.mark.asyncio @@ -29,7 +29,7 @@ async def test_basic_reading_configs_from_files(): """ Test that the config is read correctly from the files in the example_config_yaml folder """ - _base_config = BaseProxyConfig() + _base_config = ProxyConfig() current_path = os.path.dirname(os.path.abspath(__file__)) example_config_yaml_path = os.path.join(current_path, "example_config_yaml") @@ -48,7 +48,7 @@ async def test_read_config_from_bad_file_path(): """ Raise an exception if the file path is not valid """ - _base_config = BaseProxyConfig() + _base_config = ProxyConfig() config_path = "non-existent-file.yaml" with pytest.raises(Exception): config = await _base_config.get_config(config_file_path=config_path) @@ -81,7 +81,7 @@ async def test_read_config_file_with_os_environ_vars(): os.environ[key] = value # Read config - _base_config = BaseProxyConfig() + _base_config = ProxyConfig() current_path = os.path.dirname(os.path.abspath(__file__)) config_path = os.path.join( current_path, "example_config_yaml", "config_with_env_vars.yaml" From 64ebe6193f667fd7794df79fc24460058e5aacea Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Fri, 18 Oct 2024 09:00:52 +0530 Subject: [PATCH 11/11] working unit tests for reading configs --- litellm/proxy/proxy_server.py | 2 ++ tests/local_testing/test_proxy_config_unit_test.py | 12 ++++++------ 2 files changed, 8 insertions(+), 6 deletions(-) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 569fa9b56001..5cbc629915c7 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -1396,6 +1396,8 @@ async def _get_config_from_file( if os.path.exists(f"{file_path}"): with open(f"{file_path}", "r") as config_file: config = yaml.safe_load(config_file) + elif file_path is not None: + raise Exception(f"Config file not found: {file_path}") else: config = { "model_list": [], diff --git a/tests/local_testing/test_proxy_config_unit_test.py b/tests/local_testing/test_proxy_config_unit_test.py index efe51fe6a7b8..bb51ce7268fe 100644 --- a/tests/local_testing/test_proxy_config_unit_test.py +++ b/tests/local_testing/test_proxy_config_unit_test.py @@ -29,7 +29,7 @@ async def test_basic_reading_configs_from_files(): """ Test that the config is read correctly from the files in the example_config_yaml folder """ - _base_config = ProxyConfig() + proxy_config_instance = ProxyConfig() current_path = os.path.dirname(os.path.abspath(__file__)) example_config_yaml_path = os.path.join(current_path, "example_config_yaml") @@ -39,7 +39,7 @@ async def test_basic_reading_configs_from_files(): for file in files: config_path = os.path.join(example_config_yaml_path, file) - config = await _base_config.get_config(config_file_path=config_path) + config = await proxy_config_instance.get_config(config_file_path=config_path) print(config) @@ -48,10 +48,10 @@ async def test_read_config_from_bad_file_path(): """ Raise an exception if the file path is not valid """ - _base_config = ProxyConfig() + proxy_config_instance = ProxyConfig() config_path = "non-existent-file.yaml" with pytest.raises(Exception): - config = await _base_config.get_config(config_file_path=config_path) + config = await proxy_config_instance.get_config(config_file_path=config_path) @pytest.mark.asyncio @@ -81,12 +81,12 @@ async def test_read_config_file_with_os_environ_vars(): os.environ[key] = value # Read config - _base_config = ProxyConfig() + proxy_config_instance = ProxyConfig() current_path = os.path.dirname(os.path.abspath(__file__)) config_path = os.path.join( current_path, "example_config_yaml", "config_with_env_vars.yaml" ) - config = await _base_config.get_config(config_file_path=config_path) + config = await proxy_config_instance.get_config(config_file_path=config_path) print(config) # Add assertions