From 0a36f89009855d7cc19f1ccab0565dc3dec3adfe Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Mon, 7 Jul 2025 14:42:33 -0700 Subject: [PATCH] [Bug Fix] s3 config.yaml file - ensure yaml safe load is used (#12373) * use yaml safe load * test_get_file_contents_from_s3_no_temp_file_creation --- .../proxy/common_utils/load_config_utils.py | 19 ++-- .../common_utils/test_load_config_utils.py | 90 +++++++++++++++++++ 2 files changed, 95 insertions(+), 14 deletions(-) create mode 100644 tests/test_litellm/proxy/common_utils/test_load_config_utils.py diff --git a/litellm/proxy/common_utils/load_config_utils.py b/litellm/proxy/common_utils/load_config_utils.py index 38e7b3f33b2..34554cbb260 100644 --- a/litellm/proxy/common_utils/load_config_utils.py +++ b/litellm/proxy/common_utils/load_config_utils.py @@ -6,8 +6,6 @@ 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.credentials import Credentials @@ -26,21 +24,14 @@ def get_file_contents_from_s3(bucket_name, object_key): response = s3_client.get_object(Bucket=bucket_name, Key=object_key) verbose_proxy_logger.debug(f"Response: {response}") - # Read the file contents + # Read the file contents and directly parse YAML 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) - + + # Parse YAML directly from string + config = 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(f"ImportError: {str(e)}") diff --git a/tests/test_litellm/proxy/common_utils/test_load_config_utils.py b/tests/test_litellm/proxy/common_utils/test_load_config_utils.py new file mode 100644 index 00000000000..0bb63ad60fd --- /dev/null +++ b/tests/test_litellm/proxy/common_utils/test_load_config_utils.py @@ -0,0 +1,90 @@ +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 + + +class TestGetFileContentsFromS3: + """Test suite for S3 config loading functionality.""" + + @patch('boto3.client') + @patch('litellm.main.bedrock_converse_chat_completion') + @patch('yaml.safe_load') + def test_get_file_contents_from_s3_no_temp_file_creation( + self, mock_yaml_load, mock_bedrock, mock_boto3_client + ): + """ + Test that get_file_contents_from_s3 doesn't create temporary files + and uses yaml.safe_load directly on the S3 response content. + + Note: It's critical that yaml.safe_load is used + + Relevant issue/PR: https://github.com/BerriAI/litellm/pull/12078 + """ + # Mock credentials + mock_credentials = MagicMock() + mock_credentials.access_key = "test_access_key" + mock_credentials.secret_key = "test_secret_key" + mock_credentials.token = "test_token" + mock_bedrock.get_credentials.return_value = mock_credentials + + # Mock S3 client and response + mock_s3_client = MagicMock() + mock_boto3_client.return_value = mock_s3_client + + # Mock S3 response with YAML content + yaml_content = """ + model_list: + - model_name: gpt-3.5-turbo + litellm_params: + model: gpt-3.5-turbo + """ + mock_response_body = MagicMock() + mock_response_body.read.return_value = yaml_content.encode('utf-8') + mock_s3_response = { + 'Body': mock_response_body + } + mock_s3_client.get_object.return_value = mock_s3_response + + # Mock yaml.safe_load to return parsed config + expected_config = { + 'model_list': [{ + 'model_name': 'gpt-3.5-turbo', + 'litellm_params': { + 'model': 'gpt-3.5-turbo' + } + }] + } + mock_yaml_load.return_value = expected_config + + # Call the function + bucket_name = "test-bucket" + object_key = "config.yaml" + result = get_file_contents_from_s3(bucket_name, object_key) + + # Assertions + assert result == expected_config + + # Verify S3 client was created with correct credentials + mock_boto3_client.assert_called_once_with( + "s3", + aws_access_key_id="test_access_key", + aws_secret_access_key="test_secret_key", + aws_session_token="test_token" + ) + + # Verify S3 get_object was called with correct parameters + mock_s3_client.get_object.assert_called_once_with( + Bucket=bucket_name, + Key=object_key + ) + + # Verify the response body was read and decoded + mock_response_body.read.assert_called_once() + + # Verify yaml.safe_load was called with the decoded content + mock_yaml_load.assert_called_once_with(yaml_content) + +