diff --git a/.circleci/config.yml b/.circleci/config.yml index 30738443d51..94fe9b2bfd8 100644 --- a/.circleci/config.yml +++ b/.circleci/config.yml @@ -1158,6 +1158,7 @@ jobs: pip install "google-cloud-aiplatform==1.43.0" pip install "mlflow==2.17.2" pip install "anthropic==0.52.0" + pip install "blockbuster==1.5.24" # Run pytest and generate JUnit XML report - setup_litellm_enterprise_pip - run: diff --git a/docs/my-website/docs/proxy/config_settings.md b/docs/my-website/docs/proxy/config_settings.md index 701200163c8..03406c9bfa4 100644 --- a/docs/my-website/docs/proxy/config_settings.md +++ b/docs/my-website/docs/proxy/config_settings.md @@ -407,6 +407,8 @@ router_settings: | DEFAULT_REPLICATE_GPU_PRICE_PER_SECOND | Default price per second for Replicate GPU. Default is 0.001400 | DEFAULT_REPLICATE_POLLING_DELAY_SECONDS | Default delay in seconds for Replicate polling. Default is 1 | DEFAULT_REPLICATE_POLLING_RETRIES | Default number of retries for Replicate polling. Default is 5 +| DEFAULT_S3_BATCH_SIZE | Default batch size for S3 logging. Default is 512 +| DEFAULT_S3_FLUSH_INTERVAL_SECONDS | Default flush interval for S3 logging. Default is 10 | DEFAULT_SLACK_ALERTING_THRESHOLD | Default threshold for Slack alerting. Default is 300 | DEFAULT_SOFT_BUDGET | Default soft budget for LiteLLM proxy keys. Default is 50.0 | DEFAULT_TRIM_RATIO | Default ratio of tokens to trim from prompt end. Default is 0.75 diff --git a/litellm/__init__.py b/litellm/__init__.py index 841711f6e67..1cdeb0d6ab4 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -119,6 +119,7 @@ _custom_logger_compatible_callbacks_literal = Literal[ "resend_email", "smtp_email", "deepeval", + "s3_v2", ] logged_real_time_event_types: Optional[Union[List[str], Literal["*"]]] = None _known_custom_logger_compatible_callbacks: List = list( diff --git a/litellm/constants.py b/litellm/constants.py index b896a16e6cf..1d69302f5d7 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -4,6 +4,10 @@ from typing import List, Literal ROUTER_MAX_FALLBACKS = int(os.getenv("ROUTER_MAX_FALLBACKS", 5)) DEFAULT_BATCH_SIZE = int(os.getenv("DEFAULT_BATCH_SIZE", 512)) DEFAULT_FLUSH_INTERVAL_SECONDS = int(os.getenv("DEFAULT_FLUSH_INTERVAL_SECONDS", 5)) +DEFAULT_S3_FLUSH_INTERVAL_SECONDS = int( + os.getenv("DEFAULT_S3_FLUSH_INTERVAL_SECONDS", 10) +) +DEFAULT_S3_BATCH_SIZE = int(os.getenv("DEFAULT_S3_BATCH_SIZE", 512)) DEFAULT_MAX_RETRIES = int(os.getenv("DEFAULT_MAX_RETRIES", 2)) DEFAULT_MAX_RECURSE_DEPTH = int(os.getenv("DEFAULT_MAX_RECURSE_DEPTH", 100)) DEFAULT_MAX_RECURSE_DEPTH_SENSITIVE_DATA_MASKER = int( diff --git a/litellm/integrations/s3_v2.py b/litellm/integrations/s3_v2.py new file mode 100644 index 00000000000..121a491cfcf --- /dev/null +++ b/litellm/integrations/s3_v2.py @@ -0,0 +1,438 @@ +""" +s3 Bucket Logging Integration + +async_log_success_event: Processes the event, stores it in memory for DEFAULT_S3_FLUSH_INTERVAL_SECONDS seconds or until DEFAULT_S3_BATCH_SIZE and then flushes to s3 + +NOTE 1: S3 does not provide a BATCH PUT API endpoint, so we create tasks to upload each element individually +""" + +import asyncio +import json +from datetime import datetime +from typing import List, Optional, cast + +import litellm +from litellm._logging import print_verbose, verbose_logger +from litellm.constants import DEFAULT_S3_BATCH_SIZE, DEFAULT_S3_FLUSH_INTERVAL_SECONDS +from litellm.integrations.s3 import get_s3_object_key +from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM +from litellm.llms.custom_httpx.http_handler import ( + _get_httpx_client, + get_async_httpx_client, + httpxSpecialProvider, +) +from litellm.types.integrations.s3_v2 import s3BatchLoggingElement +from litellm.types.utils import StandardLoggingPayload + +from .custom_batch_logger import CustomBatchLogger + + +class S3Logger(CustomBatchLogger, BaseAWSLLM): + def __init__( + self, + s3_bucket_name: Optional[str] = None, + s3_path: Optional[str] = None, + s3_region_name: Optional[str] = None, + s3_api_version: Optional[str] = None, + s3_use_ssl: bool = True, + s3_verify: Optional[bool] = None, + s3_endpoint_url: Optional[str] = None, + s3_aws_access_key_id: Optional[str] = None, + s3_aws_secret_access_key: Optional[str] = None, + s3_aws_session_token: Optional[str] = None, + s3_aws_session_name: Optional[str] = None, + s3_aws_profile_name: Optional[str] = None, + s3_aws_role_name: Optional[str] = None, + s3_aws_web_identity_token: Optional[str] = None, + s3_aws_sts_endpoint: Optional[str] = None, + s3_flush_interval: Optional[int] = DEFAULT_S3_FLUSH_INTERVAL_SECONDS, + s3_batch_size: Optional[int] = DEFAULT_S3_BATCH_SIZE, + s3_config=None, + s3_use_team_prefix: bool = False, + **kwargs, + ): + try: + verbose_logger.debug( + f"in init s3 logger - s3_callback_params {litellm.s3_callback_params}" + ) + + # IMPORTANT: We use a concurrent limit of 1 to upload to s3 + # Files should get uploaded BUT they should not impact latency of LLM calling logic + self.async_httpx_client = get_async_httpx_client( + llm_provider=httpxSpecialProvider.LoggingCallback, + ) + + self._init_s3_params( + s3_bucket_name=s3_bucket_name, + s3_region_name=s3_region_name, + s3_api_version=s3_api_version, + s3_use_ssl=s3_use_ssl, + s3_verify=s3_verify, + s3_endpoint_url=s3_endpoint_url, + s3_aws_access_key_id=s3_aws_access_key_id, + s3_aws_secret_access_key=s3_aws_secret_access_key, + s3_aws_session_token=s3_aws_session_token, + s3_aws_session_name=s3_aws_session_name, + s3_aws_profile_name=s3_aws_profile_name, + s3_aws_role_name=s3_aws_role_name, + s3_aws_web_identity_token=s3_aws_web_identity_token, + s3_aws_sts_endpoint=s3_aws_sts_endpoint, + s3_config=s3_config, + s3_path=s3_path, + s3_use_team_prefix=s3_use_team_prefix, + ) + verbose_logger.debug(f"s3 logger using endpoint url {s3_endpoint_url}") + + asyncio.create_task(self.periodic_flush()) + self.flush_lock = asyncio.Lock() + + verbose_logger.debug( + f"s3 flush interval: {s3_flush_interval}, s3 batch size: {s3_batch_size}" + ) + # Call CustomLogger's __init__ + CustomBatchLogger.__init__( + self, + flush_lock=self.flush_lock, + flush_interval=s3_flush_interval, + batch_size=s3_batch_size, + ) + self.log_queue: List[s3BatchLoggingElement] = [] + + # Call BaseAWSLLM's __init__ + BaseAWSLLM.__init__(self) + + except Exception as e: + print_verbose(f"Got exception on init s3 client {str(e)}") + raise e + + def _init_s3_params( + self, + s3_bucket_name: Optional[str] = None, + s3_region_name: Optional[str] = None, + s3_api_version: Optional[str] = None, + s3_use_ssl: bool = True, + s3_verify: Optional[bool] = None, + s3_endpoint_url: Optional[str] = None, + s3_aws_access_key_id: Optional[str] = None, + s3_aws_secret_access_key: Optional[str] = None, + s3_aws_session_token: Optional[str] = None, + s3_aws_session_name: Optional[str] = None, + s3_aws_profile_name: Optional[str] = None, + s3_aws_role_name: Optional[str] = None, + s3_aws_web_identity_token: Optional[str] = None, + s3_aws_sts_endpoint: Optional[str] = None, + s3_config=None, + s3_path: Optional[str] = None, + s3_use_team_prefix: bool = False, + ): + """ + Initialize the s3 params for this logging callback + """ + litellm.s3_callback_params = litellm.s3_callback_params or {} + # read in .env variables - example os.environ/AWS_BUCKET_NAME + for key, value in litellm.s3_callback_params.items(): + if isinstance(value, str) and value.startswith("os.environ/"): + litellm.s3_callback_params[key] = litellm.get_secret(value) + + self.s3_bucket_name = ( + litellm.s3_callback_params.get("s3_bucket_name") or s3_bucket_name + ) + self.s3_region_name = ( + litellm.s3_callback_params.get("s3_region_name") or s3_region_name + ) + self.s3_api_version = ( + litellm.s3_callback_params.get("s3_api_version") or s3_api_version + ) + self.s3_use_ssl = ( + litellm.s3_callback_params.get("s3_use_ssl", True) or s3_use_ssl + ) + self.s3_verify = litellm.s3_callback_params.get("s3_verify") or s3_verify + self.s3_endpoint_url = ( + litellm.s3_callback_params.get("s3_endpoint_url") or s3_endpoint_url + ) + self.s3_aws_access_key_id = ( + litellm.s3_callback_params.get("s3_aws_access_key_id") + or s3_aws_access_key_id + ) + + self.s3_aws_secret_access_key = ( + litellm.s3_callback_params.get("s3_aws_secret_access_key") + or s3_aws_secret_access_key + ) + + self.s3_aws_session_token = ( + litellm.s3_callback_params.get("s3_aws_session_token") + or s3_aws_session_token + ) + + self.s3_aws_session_name = ( + litellm.s3_callback_params.get("s3_aws_session_name") or s3_aws_session_name + ) + + self.s3_aws_profile_name = ( + litellm.s3_callback_params.get("s3_aws_profile_name") or s3_aws_profile_name + ) + + self.s3_aws_role_name = ( + litellm.s3_callback_params.get("s3_aws_role_name") or s3_aws_role_name + ) + + self.s3_aws_web_identity_token = ( + litellm.s3_callback_params.get("s3_aws_web_identity_token") + or s3_aws_web_identity_token + ) + + self.s3_aws_sts_endpoint = ( + litellm.s3_callback_params.get("s3_aws_sts_endpoint") or s3_aws_sts_endpoint + ) + + self.s3_config = litellm.s3_callback_params.get("s3_config") or s3_config + self.s3_path = litellm.s3_callback_params.get("s3_path") or s3_path + # done reading litellm.s3_callback_params + self.s3_use_team_prefix = ( + bool(litellm.s3_callback_params.get("s3_use_team_prefix", False)) + or s3_use_team_prefix + ) + + return + + async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): + try: + verbose_logger.debug( + f"s3 Logging - Enters logging function for model {kwargs}" + ) + + s3_batch_logging_element = self.create_s3_batch_logging_element( + start_time=start_time, + standard_logging_payload=kwargs.get("standard_logging_object", None), + ) + + if s3_batch_logging_element is None: + raise ValueError("s3_batch_logging_element is None") + + verbose_logger.debug( + "\ns3 Logger - Logging payload = %s", s3_batch_logging_element + ) + + self.log_queue.append(s3_batch_logging_element) + verbose_logger.debug( + "s3 logging: queue length %s, batch size %s", + len(self.log_queue), + self.batch_size, + ) + except Exception as e: + verbose_logger.exception(f"s3 Layer Error - {str(e)}") + pass + + async def async_upload_data_to_s3( + self, batch_logging_element: s3BatchLoggingElement + ): + try: + import hashlib + + import requests + from botocore.auth import SigV4Auth + from botocore.awsrequest import AWSRequest + except ImportError: + raise ImportError("Missing boto3 to call bedrock. Run 'pip install boto3'.") + try: + from litellm.litellm_core_utils.asyncify import asyncify + + asyncified_get_credentials = asyncify(self.get_credentials) + credentials = await asyncified_get_credentials( + aws_access_key_id=self.s3_aws_access_key_id, + aws_secret_access_key=self.s3_aws_secret_access_key, + aws_session_token=self.s3_aws_session_token, + aws_region_name=self.s3_region_name, + aws_session_name=self.s3_aws_session_name, + aws_profile_name=self.s3_aws_profile_name, + aws_role_name=self.s3_aws_role_name, + aws_web_identity_token=self.s3_aws_web_identity_token, + aws_sts_endpoint=self.s3_aws_sts_endpoint, + ) + + verbose_logger.debug( + f"s3_v2 logger - uploading data to s3 - {batch_logging_element.s3_object_key}" + ) + + # Prepare the URL + url = f"https://{self.s3_bucket_name}.s3.{self.s3_region_name}.amazonaws.com/{batch_logging_element.s3_object_key}" + + if self.s3_endpoint_url: + url = self.s3_endpoint_url + "/" + batch_logging_element.s3_object_key + + # Convert JSON to string + json_string = json.dumps(batch_logging_element.payload) + + # Calculate SHA256 hash of the content + content_hash = hashlib.sha256(json_string.encode("utf-8")).hexdigest() + + # Prepare the request + headers = { + "Content-Type": "application/json", + "x-amz-content-sha256": content_hash, + "Content-Language": "en", + "Content-Disposition": f'inline; filename="{batch_logging_element.s3_object_download_filename}"', + "Cache-Control": "private, immutable, max-age=31536000, s-maxage=0", + } + req = requests.Request("PUT", url, data=json_string, headers=headers) + prepped = req.prepare() + + # Sign the request + aws_request = AWSRequest( + method=prepped.method, + url=prepped.url, + data=prepped.body, + headers=prepped.headers, + ) + SigV4Auth(credentials, "s3", self.s3_region_name).add_auth(aws_request) + + # Prepare the signed headers + signed_headers = dict(aws_request.headers.items()) + + # Make the request + response = await self.async_httpx_client.put( + url, data=json_string, headers=signed_headers + ) + response.raise_for_status() + except Exception as e: + verbose_logger.exception(f"Error uploading to s3: {str(e)}") + + async def async_send_batch(self): + """ + + Sends runs from self.log_queue + + Returns: None + + Raises: Does not raise an exception, will only verbose_logger.exception() + """ + verbose_logger.debug(f"s3_v2 logger - sending batch of {len(self.log_queue)}") + if not self.log_queue: + return + + ######################################################### + # Flush the log queue to s3 + # the log queue can be bounded by DEFAULT_S3_BATCH_SIZE + # see custom_batch_logger.py which triggers the flush + ######################################################### + for payload in self.log_queue: + asyncio.create_task(self.async_upload_data_to_s3(payload)) + + def create_s3_batch_logging_element( + self, + start_time: datetime, + standard_logging_payload: Optional[StandardLoggingPayload], + ) -> Optional[s3BatchLoggingElement]: + """ + Helper function to create an s3BatchLoggingElement. + + Args: + start_time (datetime): The start time of the logging event. + standard_logging_payload (Optional[StandardLoggingPayload]): The payload to be logged. + s3_path (Optional[str]): The S3 path prefix. + + Returns: + Optional[s3BatchLoggingElement]: The created s3BatchLoggingElement, or None if payload is None. + """ + if standard_logging_payload is None: + return None + + team_alias = standard_logging_payload["metadata"].get("user_api_key_team_alias") + + team_alias_prefix = "" + if ( + litellm.enable_preview_features + and self.s3_use_team_prefix + and team_alias is not None + ): + team_alias_prefix = f"{team_alias}/" + + s3_file_name = ( + litellm.utils.get_logging_id(start_time, standard_logging_payload) or "" + ) + s3_object_key = get_s3_object_key( + s3_path=cast(Optional[str], self.s3_path) or "", + team_alias_prefix=team_alias_prefix, + start_time=start_time, + s3_file_name=s3_file_name, + ) + + s3_object_download_filename = ( + "time-" + + start_time.strftime("%Y-%m-%dT%H-%M-%S-%f") + + "_" + + standard_logging_payload["id"] + + ".json" + ) + + s3_object_download_filename = f"time-{start_time.strftime('%Y-%m-%dT%H-%M-%S-%f')}_{standard_logging_payload['id']}.json" + + return s3BatchLoggingElement( + payload=dict(standard_logging_payload), + s3_object_key=s3_object_key, + s3_object_download_filename=s3_object_download_filename, + ) + + def upload_data_to_s3(self, batch_logging_element: s3BatchLoggingElement): + try: + import hashlib + + import requests + from botocore.auth import SigV4Auth + from botocore.awsrequest import AWSRequest + from botocore.credentials import Credentials + except ImportError: + raise ImportError("Missing boto3 to call bedrock. Run 'pip install boto3'.") + try: + verbose_logger.debug( + f"s3_v2 logger - uploading data to s3 - {batch_logging_element.s3_object_key}" + ) + credentials: Credentials = self.get_credentials( + aws_access_key_id=self.s3_aws_access_key_id, + aws_secret_access_key=self.s3_aws_secret_access_key, + aws_session_token=self.s3_aws_session_token, + aws_region_name=self.s3_region_name, + ) + + # Prepare the URL + url = f"https://{self.s3_bucket_name}.s3.{self.s3_region_name}.amazonaws.com/{batch_logging_element.s3_object_key}" + + if self.s3_endpoint_url: + url = self.s3_endpoint_url + "/" + batch_logging_element.s3_object_key + + # Convert JSON to string + json_string = json.dumps(batch_logging_element.payload) + + # Calculate SHA256 hash of the content + content_hash = hashlib.sha256(json_string.encode("utf-8")).hexdigest() + + # Prepare the request + headers = { + "Content-Type": "application/json", + "x-amz-content-sha256": content_hash, + "Content-Language": "en", + "Content-Disposition": f'inline; filename="{batch_logging_element.s3_object_download_filename}"', + "Cache-Control": "private, immutable, max-age=31536000, s-maxage=0", + } + req = requests.Request("PUT", url, data=json_string, headers=headers) + prepped = req.prepare() + + # Sign the request + aws_request = AWSRequest( + method=prepped.method, + url=prepped.url, + data=prepped.body, + headers=prepped.headers, + ) + SigV4Auth(credentials, "s3", self.s3_region_name).add_auth(aws_request) + + # Prepare the signed headers + signed_headers = dict(aws_request.headers.items()) + + httpx_client = _get_httpx_client() + # Make the request + response = httpx_client.put(url, data=json_string, headers=signed_headers) + response.raise_for_status() + except Exception as e: + verbose_logger.exception(f"Error uploading to s3: {str(e)}") diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index cb5673f6228..a4f2dcb5586 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -135,6 +135,7 @@ from ..integrations.opik.opik import OpikLogger from ..integrations.prometheus import PrometheusLogger from ..integrations.prompt_layer import PromptLayerLogger from ..integrations.s3 import S3Logger +from ..integrations.s3_v2 import S3Logger as S3V2Logger from ..integrations.supabase import Supabase from ..integrations.traceloop import TraceloopLogger from ..integrations.weights_biases import WeightsBiasesLogger @@ -2699,7 +2700,9 @@ def set_callbacks(callback_list, function_id=None): # noqa: PLR0915 sentry_sdk_instance.init( dsn=os.environ.get("SENTRY_DSN"), traces_sample_rate=float(sentry_trace_rate), # type: ignore - sample_rate=float(sentry_sample_rate), + sample_rate=float( + sentry_sample_rate if sentry_sample_rate else 1.0 + ), ) capture_exception = sentry_sdk_instance.capture_exception add_breadcrumb = sentry_sdk_instance.add_breadcrumb @@ -2867,6 +2870,14 @@ def _init_custom_logger_compatible_class( # noqa: PLR0915 _gcs_bucket_logger = GCSBucketLogger() _in_memory_loggers.append(_gcs_bucket_logger) return _gcs_bucket_logger # type: ignore + elif logging_integration == "s3_v2": + for callback in _in_memory_loggers: + if isinstance(callback, S3V2Logger): + return callback # type: ignore + + _s3_v2_logger = S3V2Logger() + _in_memory_loggers.append(_s3_v2_logger) + return _s3_v2_logger # type: ignore elif logging_integration == "azure_storage": for callback in _in_memory_loggers: if isinstance(callback, AzureBlobStorageLogger): @@ -2962,7 +2973,7 @@ def _init_custom_logger_compatible_class( # noqa: PLR0915 galileo_logger = GalileoObserve() _in_memory_loggers.append(galileo_logger) return galileo_logger # type: ignore - + elif logging_integration == "deepeval": for callback in _in_memory_loggers: if isinstance(callback, DeepEvalLogger): @@ -2970,7 +2981,7 @@ def _init_custom_logger_compatible_class( # noqa: PLR0915 deepeval_logger = DeepEvalLogger() _in_memory_loggers.append(deepeval_logger) return deepeval_logger # type: ignore - + elif logging_integration == "logfire": if "LOGFIRE_TOKEN" not in os.environ: raise ValueError("LOGFIRE_TOKEN not found in environment variables") @@ -3172,6 +3183,10 @@ def get_custom_logger_compatible_class( # noqa: PLR0915 for callback in _in_memory_loggers: if isinstance(callback, GCSBucketLogger): return callback + elif logging_integration == "s3_v2": + for callback in _in_memory_loggers: + if isinstance(callback, S3V2Logger): + return callback elif logging_integration == "azure_storage": for callback in _in_memory_loggers: if isinstance(callback, AzureBlobStorageLogger): diff --git a/litellm/types/integrations/s3_v2.py b/litellm/types/integrations/s3_v2.py new file mode 100644 index 00000000000..43b917e6200 --- /dev/null +++ b/litellm/types/integrations/s3_v2.py @@ -0,0 +1,13 @@ +from typing import Dict + +from pydantic import BaseModel + + +class s3BatchLoggingElement(BaseModel): + """ + Type of element stored in self.log_queue in S3Logger + """ + + payload: Dict + s3_object_key: str + s3_object_download_filename: str diff --git a/tests/logging_callback_tests/test_amazing_s3_logs.py b/tests/logging_callback_tests/test_amazing_s3_logs.py index 915041e7149..92993d0d76d 100644 --- a/tests/logging_callback_tests/test_amazing_s3_logs.py +++ b/tests/logging_callback_tests/test_amazing_s3_logs.py @@ -74,6 +74,59 @@ async def test_basic_s3_logging(sync_mode, streaming): s3.delete_object(Bucket="load-testing-oct", Key=key) + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "streaming", [(True)] +) +async def test_basic_s3_v2_logging(streaming): + from blockbuster import BlockBuster + from litellm.integrations.s3_v2 import S3Logger + s3_v2_logger = S3Logger(s3_flush_interval=1) + litellm.callbacks = [s3_v2_logger] + blockbuster = BlockBuster() + blockbuster.activate() + + litellm._turn_on_debug() + litellm.callbacks = ["s3_v2"] + litellm.s3_callback_params = { + "s3_bucket_name": "load-testing-oct", + "s3_aws_secret_access_key": "os.environ/AWS_SECRET_ACCESS_KEY", + "s3_aws_access_key_id": "os.environ/AWS_ACCESS_KEY_ID", + "s3_region_name": "us-west-2", + } + litellm.set_verbose = True + response_id = None + response = await litellm.acompletion( + model="gpt-4o-mini", + messages=[{"role": "user", "content": "This is a test"}], + stream=streaming, + ) + if streaming: + async for chunk in response: + print(chunk) + response_id = chunk.id + else: + response_id = response.id + + await asyncio.sleep(30) + print(f"response: {response}") + + # stop blockbuster + blockbuster.deactivate() + + total_objects, all_s3_keys = list_all_s3_objects("load-testing-oct") + + print(f"all_s3_keys: {all_s3_keys}") + + #assert that atlest one key has response.id in it + assert any(response_id in key for key in all_s3_keys) + s3 = boto3.client("s3") + # delete all objects + for key in all_s3_keys: + s3.delete_object(Bucket="load-testing-oct", Key=key) + + def list_all_s3_objects(bucket_name): s3 = boto3.client("s3") diff --git a/tests/logging_callback_tests/test_unit_tests_init_callbacks.py b/tests/logging_callback_tests/test_unit_tests_init_callbacks.py index e4915861b76..d779b4ea944 100644 --- a/tests/logging_callback_tests/test_unit_tests_init_callbacks.py +++ b/tests/logging_callback_tests/test_unit_tests_init_callbacks.py @@ -34,6 +34,7 @@ from litellm.integrations.opentelemetry import OpenTelemetry from litellm.integrations.mlflow import MlflowLogger from litellm.integrations.argilla import ArgillaLogger from litellm.integrations.deepeval.deepeval import DeepEvalLogger +from litellm.integrations.s3_v2 import S3Logger from litellm.integrations.anthropic_cache_control_hook import AnthropicCacheControlHook from litellm.integrations.vector_stores.bedrock_vector_store import BedrockVectorStore from litellm.integrations.langfuse.langfuse_prompt_management import ( @@ -88,6 +89,7 @@ callback_class_str_to_classType = { "resend_email": ResendEmailLogger, "smtp_email": SMTPEmailLogger, "deepeval": DeepEvalLogger, + "s3_v2": S3Logger, } expected_env_vars = { @@ -113,6 +115,10 @@ expected_env_vars = { "GCS_PUBSUB_PROJECT_ID": "gcs_pubsub_project_id", "CONFIDENT_API_KEY": "confident_api_key", "LITELM_ENVIRONMENT": "development", + "AWS_BUCKET_NAME": "aws_bucket_name", + "AWS_SECRET_ACCESS_KEY": "aws_secret_access_key", + "AWS_ACCESS_KEY_ID": "aws_access_key_id", + "AWS_REGION": "aws_region", }