(feat) proxy - support s3_callback_params

This commit is contained in:
ishaan-jaff 2024-01-11 09:57:47 +05:30
parent 53cee85bc2
commit 4a47b17ba2
3 changed files with 54 additions and 18 deletions

View file

@ -135,6 +135,7 @@ model_fallbacks: Optional[List] = None # Deprecated for 'litellm.fallbacks'
model_cost_map_url: str = "https://raw.githubusercontent.com/BerriAI/litellm/main/model_prices_and_context_window.json" model_cost_map_url: str = "https://raw.githubusercontent.com/BerriAI/litellm/main/model_prices_and_context_window.json"
suppress_debug_info = False suppress_debug_info = False
dynamodb_table_name: Optional[str] = None dynamodb_table_name: Optional[str] = None
s3_callback_params: Optional[Dict] = None
#### RELIABILITY #### #### RELIABILITY ####
request_timeout: Optional[float] = 6000 request_timeout: Optional[float] = 6000
num_retries: Optional[int] = None # per model endpoint num_retries: Optional[int] = None # per model endpoint

View file

@ -15,7 +15,7 @@ class S3Logger:
# Class variables or attributes # Class variables or attributes
def __init__( def __init__(
self, self,
s3_bucket_name="litellm-logs", s3_bucket_name=None,
s3_region_name=None, s3_region_name=None,
s3_api_version=None, s3_api_version=None,
s3_use_ssl=True, s3_use_ssl=True,
@ -29,21 +29,51 @@ class S3Logger:
): ):
import boto3 import boto3
self.bucket_name = s3_bucket_name try:
# Create an S3 client with custom endpoint URL print_verbose("in init s3 logger")
self.s3_client = boto3.client(
"s3", if litellm.s3_callback_params is not None:
region_name=s3_region_name, # read in .env variables - example os.environ/AWS_BUCKET_NAME
endpoint_url=s3_endpoint_url, for key, value in litellm.s3_callback_params.items():
api_version=s3_api_version, if type(value) is str and value.startswith("os.environ/"):
use_ssl=s3_use_ssl, litellm.s3_callback_params[key] = litellm.get_secret(value)
verify=s3_verify, # now set s3 params from litellm.s3_logger_params
aws_access_key_id=s3_aws_access_key_id, s3_bucket_name = litellm.s3_callback_params.get("s3_bucket_name")
aws_secret_access_key=s3_aws_secret_access_key, s3_region_name = litellm.s3_callback_params.get("s3_region_name")
aws_session_token=s3_aws_session_token, s3_api_version = litellm.s3_callback_params.get("s3_api_version")
config=s3_config, s3_use_ssl = litellm.s3_callback_params.get("s3_use_ssl")
**kwargs, s3_verify = litellm.s3_callback_params.get("s3_verify")
) s3_endpoint_url = litellm.s3_callback_params.get("s3_endpoint_url")
s3_aws_access_key_id = litellm.s3_callback_params.get(
"s3_aws_access_key_id"
)
s3_aws_secret_access_key = litellm.s3_callback_params.get(
"s3_aws_secret_access_key"
)
s3_aws_session_token = litellm.s3_callback_params.get(
"s3_aws_session_token"
)
s3_config = litellm.s3_callback_params.get("s3_config")
# done reading litellm.s3_callback_params
self.bucket_name = s3_bucket_name
# Create an S3 client with custom endpoint URL
self.s3_client = boto3.client(
"s3",
region_name=s3_region_name,
endpoint_url=s3_endpoint_url,
api_version=s3_api_version,
use_ssl=s3_use_ssl,
verify=s3_verify,
aws_access_key_id=s3_aws_access_key_id,
aws_secret_access_key=s3_aws_secret_access_key,
aws_session_token=s3_aws_session_token,
config=s3_config,
**kwargs,
)
except Exception as e:
print_verbose(f"Got exception on init s3 client {str(e)}")
raise e
async def _async_log_event( async def _async_log_event(
self, kwargs, response_obj, start_time, end_time, print_verbose self, kwargs, response_obj, start_time, end_time, print_verbose
@ -109,7 +139,7 @@ class S3Logger:
print_verbose(f"s3 Layer Logging - final response object: {response_obj}") print_verbose(f"s3 Layer Logging - final response object: {response_obj}")
return response return response
except: except Exception as e:
traceback.print_exc() traceback.print_exc()
print_verbose(f"s3 Layer Error - {traceback.format_exc()}") print_verbose(f"s3 Layer Error - {str(e)}\n{traceback.format_exc()}")
pass pass

View file

@ -24,6 +24,11 @@ def test_s3_logging():
# redirect stdout to log_file # redirect stdout to log_file
litellm.success_callback = ["s3"] litellm.success_callback = ["s3"]
litellm.s3_callback_params = {
"s3_bucket_name": "litellm-logs",
"s3_aws_secret_access_key": "os.environ/AWS_SECRET_ACCESS_KEY",
"s3_aws_access_key_id": "os.environ/AWS_ACCESS_KEY_ID",
}
litellm.set_verbose = True litellm.set_verbose = True
print("Testing async s3 logging") print("Testing async s3 logging")