diff --git a/litellm/batches/batch_line_items.py b/litellm/batches/batch_line_items.py index 8f1e3efaaf3..a6522214dba 100644 --- a/litellm/batches/batch_line_items.py +++ b/litellm/batches/batch_line_items.py @@ -24,7 +24,7 @@ from litellm.types.utils import ( ) if TYPE_CHECKING: - from litellm.integrations.s3_v3 import S3Logger + from litellm.integrations.s3_v2 import S3Logger BatchFileProvider: TypeAlias = Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic", "mistral"] @@ -242,7 +242,7 @@ async def _deliver( def _resolve_logger(callback: "litellm.CALLBACK_TYPES | S3Logger") -> "CustomLogger | S3Logger | None": - from litellm.integrations.s3_v3 import S3Logger + from litellm.integrations.s3_v2 import S3Logger match callback: case CustomLogger() | S3Logger(): diff --git a/litellm/integrations/s3_v2.py b/litellm/integrations/s3_v2.py index 98892a7fdf4..45a6f887650 100644 --- a/litellm/integrations/s3_v2.py +++ b/litellm/integrations/s3_v2.py @@ -1,3 +1,481 @@ -from litellm.integrations.s3_v3 import S3BatchUploadError, S3Logger +import asyncio +import copy +import json +import os +from collections.abc import Callable, Collection, Mapping +from typing import TYPE_CHECKING, Final -__all__ = ["S3BatchUploadError", "S3Logger"] +from typing_extensions import Unpack + +import litellm +from litellm._logging import verbose_logger +from litellm.constants import ( + COMPLETION_HTTP_FALLBACK_SECONDS, + DEFAULT_S3_BATCH_SIZE, + DEFAULT_S3_FLUSH_INTERVAL_SECONDS, + DEFAULT_S3_MAX_ADAPTIVE_CONCURRENCY, + DEFAULT_S3_MAX_CONCURRENT_UPLOADS, + HTTP_HANDLER_CONNECT_TIMEOUT_SECONDS, + S3_UPLOAD_ATTEMPT_TIMEOUT_SECONDS, +) +from litellm.integrations.batch_line_item_objects import batch_line_items_ndjson, batch_line_items_object_name +from litellm.integrations.custom_batch_logger import CustomBatchLogger +from litellm.integrations.custom_logger import CustomLogger +from litellm.integrations.custom_logger_v2.call import CallContext +from litellm.integrations.custom_logger_v2.derivations import ( + cold_storage_object_key, + log_id, + logged_key_identity, + mapping, + payload_copy, +) +from litellm.integrations.custom_logger_v2.enums import Field, Hook +from litellm.integrations.custom_logger_v2.field_sources import HookArgs +from litellm.integrations.custom_logger_v2.logger import CustomLoggerV2, on +from litellm.integrations.custom_logger_v2.outcomes import Abort, Patch +from litellm.integrations.custom_logger_v2.rust import RustLogger +from litellm.integrations.custom_logger_v2.selection import Selection +from litellm.integrations.s3 import resolve_s3_max_queue_size, resolve_s3_max_retry_age_seconds +from litellm.llms.bedrock.web_identity_session_policy import S3_LOGGING_SESSION_POLICY, session_policy_document +from litellm.secret_managers.main import get_secret, get_secret_str +from litellm.types.integrations.s3_v2 import S3LoggerSettings, S3PartitionGranularity +from litellm.types.utils import StandardAuditLogPayload, StandardBatchLineItemsPayload, StandardLoggingPayload + +if TYPE_CHECKING: + from litellm.rust_bridge._native import NativeS3Log + +CALLBACK_NAME: Final = "s3_v2" +STRIP_CALLBACK_NAME: Final = "s3_v2_strip_base64_files" +_WEB_IDENTITY_TOKEN_MIN_LENGTH: Final = 4 +_DEFAULT_MAX_RETRY_AGE_SECONDS: Final = 3600 + + +class S3BatchUploadError(Exception): + def __init__(self, failed: int, total: int) -> None: + self.failed: Final = failed + self.total: Final = total + super().__init__(f"{failed} of {total} S3 uploads failed; transient failures kept in queue for the next flush") + + +def _raise_batch_upload_error(failed: int, total: int) -> None: + raise S3BatchUploadError(failed=failed, total=total) + + +def batch_upload_failed(failure: tuple[int, int, int] | None) -> None: + if failure is None: + return + queued, failed, total = failure + try: + _raise_batch_upload_error(failed, total) + except S3BatchUploadError: + verbose_logger.exception( + "CustomLogger: async_send_batch raised; preserving %s events in queue for retry", queued + ) + + +def _cold_storage_is_s3_v2() -> bool: + return litellm.cold_storage_custom_logger == "s3_v2" + + +def _region_from_environment() -> str | None: + region_name: Final = get_secret_str("AWS_REGION_NAME", None) + region: Final = get_secret_str("AWS_REGION", None) + return region if region is not None else region_name + + +class _Host: + def __init__(self) -> None: + self.failures: Final = CustomLogger() + + def callback_failed(self) -> None: + self.failures.handle_callback_failure(callback_name="S3Logger") + + def web_identity_token(self, reference: str) -> str: + if reference.startswith("os.environ/") or reference in os.environ: + raise ValueError("Invalid web identity token reference.") + match _read_secret(get_secret, reference): + case None: + raise ValueError("OIDC token could not be retrieved from secret manager.") + case str() as token if len(token) < _WEB_IDENTITY_TOKEN_MIN_LENGTH: + raise ValueError( + "Parameter validation failed:\nInvalid length for parameter WebIdentityToken, " + f"value: {len(token)}, valid min length: {_WEB_IDENTITY_TOKEN_MIN_LENGTH}" + ) + case str() as token: + return token + case other: + raise ValueError( + "Parameter validation failed:\nInvalid type for parameter WebIdentityToken, " + f"value: {other}, type: {type(other)}, valid types: {str}" + ) + + +def _read_secret(read: Callable[[str], object], reference: str) -> object: + return read(reference) + + +def _resolved_params(params_source: Mapping[str, object] | None) -> dict[str, object]: # mutable-ok: handed to Rust + configured: Final[Mapping[object, object] | None] = getattr(litellm, "s3_callback_params", None) + source: Final[Mapping[str, object] | Mapping[object, object]] = ( + params_source if params_source is not None else (configured or {}) + ) + return { + key: (litellm.get_secret(value) if isinstance(value, str) and value.startswith("os.environ/") else value) + for key, value in source.items() + if isinstance(key, str) + } + + +def _constants(batch_size: int | None = DEFAULT_S3_BATCH_SIZE) -> Mapping[str, object]: + return { + "batch_size": batch_size or litellm.DEFAULT_BATCH_SIZE, + "max_concurrent_uploads": DEFAULT_S3_MAX_CONCURRENT_UPLOADS, + "max_adaptive_concurrency": DEFAULT_S3_MAX_ADAPTIVE_CONCURRENCY, + "upload_attempt_timeout_seconds": S3_UPLOAD_ATTEMPT_TIMEOUT_SECONDS, + "download_read_timeout_seconds": COMPLETION_HTTP_FALLBACK_SECONDS, + "download_connect_timeout_seconds": HTTP_HANDLER_CONNECT_TIMEOUT_SECONDS, + "web_identity_session_policy": session_policy_document(S3_LOGGING_SESSION_POLICY), + } + + +def _flush_interval(flush_interval: float | None = DEFAULT_S3_FLUSH_INTERVAL_SECONDS) -> float: + return flush_interval or litellm.DEFAULT_FLUSH_INTERVAL_SECONDS + + +_CONFIGURED_FIRST: Final = frozenset( + { + "s3_bucket_name", + "s3_path", + "s3_region_name", + "s3_api_version", + "s3_endpoint_url", + "s3_aws_access_key_id", + "s3_aws_secret_access_key", + "s3_aws_session_token", + "s3_aws_session_name", + "s3_aws_profile_name", + "s3_aws_role_name", + "s3_aws_web_identity_token", + "s3_aws_sts_endpoint", + "s3_config", + "s3_server_side_encryption", + "s3_sse_kms_key_id", + } +) +_CONSTRUCTOR_ENABLES: Final = frozenset( + { + "s3_use_team_prefix", + "s3_strip_base64_files", + "s3_use_key_prefix", + "s3_use_virtual_hosted_style", + "s3_adaptive_concurrency", + "s3_batch_file_upload", + } +) +_CONFIGURED_UNLESS_NONE: Final = frozenset({"s3_use_ssl", "s3_verify", "s3_drop_on_terminal_error"}) +_CONSTRUCTOR_UNLESS_NONE: Final = frozenset({"s3_log_prompts_only", "s3_partition_granularity"}) +_CONFIGURED_UNLESS_BLANK: Final = frozenset({"s3_max_concurrent_uploads", "s3_max_adaptive_concurrency"}) +_CONSTRUCTOR_DEFAULTS: Final[Mapping[str, object]] = { + "s3_use_ssl": True, + "s3_drop_on_terminal_error": True, + "s3_max_concurrent_uploads": DEFAULT_S3_MAX_CONCURRENT_UPLOADS, + "s3_max_retry_age_seconds": _DEFAULT_MAX_RETRY_AGE_SECONDS, + "s3_flush_interval": DEFAULT_S3_FLUSH_INTERVAL_SECONDS, + "s3_batch_size": DEFAULT_S3_BATCH_SIZE, +} + + +def _blank(value: object) -> bool: + return value is None or value == "" + + +def _constructor_wins(key: str, configured: object, value: object) -> bool: + return ( + (key in _CONFIGURED_FIRST and not configured) + or (key in _CONFIGURED_UNLESS_NONE and configured is None) + or (key in _CONFIGURED_UNLESS_BLANK and _blank(configured)) + or (key in _CONSTRUCTOR_UNLESS_NONE and value is not None) + ) + + +_KEEP_CONFIGURED: Final = object() + + +def _constructor_value(key: str, configured: object, value: object) -> object: + if _constructor_wins(key, configured, value): + return value + if key in _CONSTRUCTOR_ENABLES and value: + return True + if key == "s3_max_queue_size" and value is not None: + return resolve_s3_max_queue_size( + configured, resolve_s3_max_queue_size(value, CustomBatchLogger.DEFAULT_MAX_QUEUE_SIZE) + ) + if key == "s3_max_retry_age_seconds": + fallback: Final = resolve_s3_max_retry_age_seconds(value, _DEFAULT_MAX_RETRY_AGE_SECONDS) + resolved: Final = fallback if _blank(configured) else resolve_s3_max_retry_age_seconds(configured, fallback) + return 0 if resolved is None else resolved + return _KEEP_CONFIGURED + + +def _constructed_params( + params_source: Mapping[str, object] | None, settings: S3LoggerSettings +) -> dict[str, object]: # mutable-ok: handed to Rust + params: Final = _resolved_params(params_source) + overrides: Final = { + key: replaced + for key, value in settings.items() + if value != _CONSTRUCTOR_DEFAULTS.get(key, False if key in _CONSTRUCTOR_ENABLES else None) + and (replaced := _constructor_value(key, params.get(key), value)) is not _KEEP_CONFIGURED + } + return {**params, **overrides} + + +class _Base64Stripper(CustomLogger): + def strip(self, payload: StandardLoggingPayload) -> None: + _ = self._strip_base64_from_messages_sync(payload) + + def stripped(self, messages: list[object]) -> list[dict[str, object]]: + return self._process_messages(messages=[_shallow_copy(message) for message in messages]) + + +def _shallow_copy(message: object) -> object: + entries: Final = mapping(message) + return copy.copy(entries) if entries is not None and isinstance(message, dict) else message + + +_STRIPPER: Final = _Base64Stripper() + + +def legacy_should_log(kwargs: Mapping[str, object]) -> bool: + from litellm.litellm_core_utils import litellm_logging + + if litellm.global_disable_no_log_param: + return True + litellm_params: Final = kwargs.get("litellm_params", {}) + lookup: Final = getattr(litellm_params, "get", None) + if callable(lookup) and lookup("no-log", False) is True: + return False + return not litellm_logging.callback_disabled_dynamically( + callback=CALLBACK_NAME, + litellm_params=litellm_params, + standard_callback_dynamic_params=kwargs.get("standard_callback_dynamic_params") or {}, # pyright: ignore[reportArgumentType] # the dynamic params dict the legacy path hands the same check + ) + + +def _payload_messages(payload: Mapping[str, object]) -> list[object]: + raw: Final = payload.get("messages") + return [*raw] if isinstance(raw, list) else [] + + +class S3StripBase64Files(CustomLoggerV2): + selection = Selection(name=STRIP_CALLBACK_NAME) + + def __init__(self, s3_logger_key: str) -> None: + self.s3_logger_key: Final = s3_logger_key + + @on( + Hook.log_prepare, + reads={Field.deprecated_kwargs, Field.standard_logging_payload}, + writes={Field.standard_logging_payload}, + ) + async def on_log_prepare(self, view: object, ctx: CallContext[object]) -> Patch | None: + from litellm.litellm_core_utils.litellm_logging import V2_ASYNC_LOGGING_LINE + + kwargs: Final = mapping(getattr(view, Field.deprecated_kwargs.value, None)) + payload: Final = mapping(getattr(view, Field.standard_logging_payload.value, None)) + if not V2_ASYNC_LOGGING_LINE.get() or kwargs is None or payload is None or not legacy_should_log(kwargs): + return None + messages: Final = _payload_messages(payload) + if not messages: + return None + return Patch().merge(Field.standard_logging_payload, {"messages": _STRIPPER.stripped(messages)}) + + +async def _periodic_flush(sink: "NativeS3Log", interval: float) -> None: + while True: + await asyncio.sleep(interval) + try: + batch_upload_failed(await sink.flush()) + except (asyncio.CancelledError, KeyboardInterrupt, SystemExit): + raise + except BaseException as e: # noqa: BLE001 # a failed flush keeps the queue and the next interval tries again, as flush_queue does + verbose_logger.exception("CustomLogger: periodic flush failed: %s", e) + + +def _excluding_fields(kwargs: object) -> Mapping[str, object] | None: + excluded: Final[object] = getattr(litellm, "standard_logging_payload_excluded_fields", None) + details: Final = mapping(kwargs) + logged: Final = None if details is None else mapping(details.get("standard_logging_object")) + if details is None or logged is None or not isinstance(excluded, Collection) or not excluded: + return None + if logged.get("status") == "failure": + return None + return {**details, "standard_logging_object": {key: value for key, value in logged.items() if key not in excluded}} + + +def _logged_payload(kwargs: object, details: Mapping[str, object] | None) -> object: + source: Final = mapping(kwargs) if details is None else details + return None if source is None else source.get("standard_logging_object") + + +def _legacy_kwargs(view: object) -> object: + return getattr(view, Field.deprecated_kwargs.value, None) + + +_read_attribute: Final[Callable[[object, str], object]] = getattr + + +class _ExcludedView: + def __init__(self, view: object, details: Mapping[str, object]) -> None: + args: Final = HookArgs(request=details) + self.view: Final = view + self.fields: Final[Mapping[str, object]] = { + Field.deprecated_kwargs.value: details, + Field.standard_logging_payload_ref.value: payload_copy(args), + Field.log_id.value: log_id(args), + Field.cold_storage_object_key.value: cold_storage_object_key(args), + Field.logged_key_identity.value: logged_key_identity(args), + } + + def __getattr__(self, name: str) -> object: + if name not in self.fields: + return _read_attribute(self.view, name) + value: Final = self.fields[name] + if value is None: + raise AttributeError(name) + return value + + +_ATTEMPT_ENDED_READS: Final = frozenset( + { + Field.deprecated_kwargs, + Field.attempt_timing, + Field.log_id, + Field.cold_storage_object_key, + Field.standard_logging_payload_ref, + Field.logged_key_identity, + } +) + + +class S3Logger(RustLogger): + selection = Selection(name=CALLBACK_NAME, honors_no_log=True) + + def __init__( + self, + s3_callback_params_override: Mapping[str, object] | None = None, + **settings: Unpack[S3LoggerSettings], # kwargs-ok: the S3Logger constructor keywords + ) -> None: + from litellm.rust_bridge._native import NativeS3Log + + host: Final = _Host() + params: Final = _constructed_params(s3_callback_params_override, settings) + sink: Final = NativeS3Log( + params, + _constants(settings.get("s3_batch_size", DEFAULT_S3_BATCH_SIZE)), + host, + _cold_storage_is_s3_v2(), + None if params.get("s3_region_name") else _region_from_environment(), + ) + self.s3_region_name: Final[object] = params.get("s3_region_name") or settings.get("s3_region_name") + self.sink: Final = sink + self.s3_host: Final = host + self.periodic_flush_task: Final = asyncio.create_task( + _periodic_flush(sink, _flush_interval(settings.get("s3_flush_interval", DEFAULT_S3_FLUSH_INTERVAL_SECONDS))) + ) + super().__init__(sink.callback(), {Field.deprecated_kwargs: lambda _kwargs: None}) + self.v2_companion: Final = ( + S3StripBase64Files(litellm.logging_callback_manager.get_custom_logger_key(self)) + if sink.strip_base64_files + else None + ) + + @on(Hook.attempt_ended, reads=_ATTEMPT_ENDED_READS, awaited=True) + async def on_attempt_ended(self, view: object, ctx: CallContext[object]) -> Patch | Abort | None: + return await self.run_rust(Hook.attempt_ended, _ATTEMPT_ENDED_READS, view, ctx) + + async def run_rust( + self, hook: Hook, reads: frozenset[Field], view: object, ctx: CallContext[object] + ) -> Patch | Abort | None: + from litellm.litellm_core_utils.litellm_logging import V2_ASYNC_LOGGING_LINE + + kwargs: Final = getattr(view, Field.deprecated_kwargs.value, None) + if not V2_ASYNC_LOGGING_LINE.get() or not isinstance(kwargs, Mapping) or not legacy_should_log(kwargs): # pyright: ignore[reportUnknownArgumentType] # the legacy kwargs dict + return None + payload: Final[object] = kwargs.get("standard_logging_object") # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # the legacy kwargs dict + if payload is not None and not isinstance(payload, Mapping): + return self._failed(AttributeError(f"'{type(payload).__name__}' object has no attribute 'get'")) # pyright: ignore[reportUnknownArgumentType] # any non-mapping payload + try: + raw: Final = _legacy_kwargs(view) + details: Final = _excluding_fields(raw) + logged: Final = _logged_payload(raw, details) + if self.sink.strip_base64_files and isinstance(logged, Mapping): + _STRIPPER.strip(logged) # pyright: ignore[reportArgumentType] # the shared StandardLoggingPayload, stripped in place as S3Logger does + return await super().run_rust(hook, reads, view if details is None else _ExcludedView(view, details), ctx) + except Exception as e: # noqa: BLE001 # S3Logger counts a failure to build the element and moves on + return self._failed(e) + + def _failed(self, error: Exception) -> None: + verbose_logger.error("s3 Layer Error - %s", error, exc_info=error) + self.s3_host.callback_failed() + + @property + def s3_path(self) -> str | None: + return self.sink.s3_path + + @property + def s3_bucket_name(self) -> str | None: + return self.sink.s3_bucket_name + + def resolve_partition_granularity(self) -> S3PartitionGranularity: + return "hour" if self.sink.partition_granularity() == "hour" else "day" + + async def flush_queue(self) -> None: + batch_upload_failed(await self.sink.flush()) + + async def async_log_audit_log_event(self, audit_log: StandardAuditLogPayload) -> None: + try: + audit_log_id: Final = audit_log.get("id", "unknown") + batch_upload_failed(await self.sink.submit_audit(dict(audit_log), f"{audit_log_id}")) + except Exception as e: # noqa: BLE001 # an audit log failure is logged and swallowed, as in S3Logger + verbose_logger.exception("S3 audit log error: %s", e) + + async def async_log_batch_line_items(self, payload: StandardBatchLineItemsPayload) -> None: + object_name: Final = batch_line_items_object_name(payload) + s3_path: Final = self.s3_path + s3_object_key: Final = (f"{s3_path.rstrip('/')}/" if s3_path else "") + object_name + body: Final = await asyncio.to_thread(batch_line_items_ndjson, payload) + delivered: Final = await self.sink.upload_line_items( + s3_object_key, body, "-".join(object_name.rsplit("/", 2)[-2:]) + ) + if not delivered: + verbose_logger.error( + "s3 logging: the %d line items of part %d of batch %s were not stored at %s", + len(payload["lines"]), + payload["part"], + payload["batch_id"], + s3_object_key, + ) + + async def _download_object_from_s3(self, s3_object_key: str) -> object: + try: + downloaded: Final = await self.sink.download(s3_object_key) + return None if downloaded is None else json.loads(downloaded) + except Exception as e: # noqa: BLE001 # a failed read is None, as in S3Logger + verbose_logger.exception("Error downloading from S3: %s", e) + return None + + async def get_proxy_server_request_from_cold_storage_with_object_key(self, object_key: str) -> object: + try: + return await self._download_object_from_s3(object_key) + except Exception as e: # noqa: BLE001 # a failed read is None, as in S3Logger + verbose_logger.exception("Error retrieving object %s from cold storage: %s", object_key, e) + return None + + +__all__ = [ + "S3BatchUploadError", + "S3Logger", + "S3StripBase64Files", + "legacy_should_log", +] diff --git a/litellm/integrations/s3_v3.py b/litellm/integrations/s3_v3.py deleted file mode 100644 index 45a6f887650..00000000000 --- a/litellm/integrations/s3_v3.py +++ /dev/null @@ -1,481 +0,0 @@ -import asyncio -import copy -import json -import os -from collections.abc import Callable, Collection, Mapping -from typing import TYPE_CHECKING, Final - -from typing_extensions import Unpack - -import litellm -from litellm._logging import verbose_logger -from litellm.constants import ( - COMPLETION_HTTP_FALLBACK_SECONDS, - DEFAULT_S3_BATCH_SIZE, - DEFAULT_S3_FLUSH_INTERVAL_SECONDS, - DEFAULT_S3_MAX_ADAPTIVE_CONCURRENCY, - DEFAULT_S3_MAX_CONCURRENT_UPLOADS, - HTTP_HANDLER_CONNECT_TIMEOUT_SECONDS, - S3_UPLOAD_ATTEMPT_TIMEOUT_SECONDS, -) -from litellm.integrations.batch_line_item_objects import batch_line_items_ndjson, batch_line_items_object_name -from litellm.integrations.custom_batch_logger import CustomBatchLogger -from litellm.integrations.custom_logger import CustomLogger -from litellm.integrations.custom_logger_v2.call import CallContext -from litellm.integrations.custom_logger_v2.derivations import ( - cold_storage_object_key, - log_id, - logged_key_identity, - mapping, - payload_copy, -) -from litellm.integrations.custom_logger_v2.enums import Field, Hook -from litellm.integrations.custom_logger_v2.field_sources import HookArgs -from litellm.integrations.custom_logger_v2.logger import CustomLoggerV2, on -from litellm.integrations.custom_logger_v2.outcomes import Abort, Patch -from litellm.integrations.custom_logger_v2.rust import RustLogger -from litellm.integrations.custom_logger_v2.selection import Selection -from litellm.integrations.s3 import resolve_s3_max_queue_size, resolve_s3_max_retry_age_seconds -from litellm.llms.bedrock.web_identity_session_policy import S3_LOGGING_SESSION_POLICY, session_policy_document -from litellm.secret_managers.main import get_secret, get_secret_str -from litellm.types.integrations.s3_v2 import S3LoggerSettings, S3PartitionGranularity -from litellm.types.utils import StandardAuditLogPayload, StandardBatchLineItemsPayload, StandardLoggingPayload - -if TYPE_CHECKING: - from litellm.rust_bridge._native import NativeS3Log - -CALLBACK_NAME: Final = "s3_v2" -STRIP_CALLBACK_NAME: Final = "s3_v2_strip_base64_files" -_WEB_IDENTITY_TOKEN_MIN_LENGTH: Final = 4 -_DEFAULT_MAX_RETRY_AGE_SECONDS: Final = 3600 - - -class S3BatchUploadError(Exception): - def __init__(self, failed: int, total: int) -> None: - self.failed: Final = failed - self.total: Final = total - super().__init__(f"{failed} of {total} S3 uploads failed; transient failures kept in queue for the next flush") - - -def _raise_batch_upload_error(failed: int, total: int) -> None: - raise S3BatchUploadError(failed=failed, total=total) - - -def batch_upload_failed(failure: tuple[int, int, int] | None) -> None: - if failure is None: - return - queued, failed, total = failure - try: - _raise_batch_upload_error(failed, total) - except S3BatchUploadError: - verbose_logger.exception( - "CustomLogger: async_send_batch raised; preserving %s events in queue for retry", queued - ) - - -def _cold_storage_is_s3_v2() -> bool: - return litellm.cold_storage_custom_logger == "s3_v2" - - -def _region_from_environment() -> str | None: - region_name: Final = get_secret_str("AWS_REGION_NAME", None) - region: Final = get_secret_str("AWS_REGION", None) - return region if region is not None else region_name - - -class _Host: - def __init__(self) -> None: - self.failures: Final = CustomLogger() - - def callback_failed(self) -> None: - self.failures.handle_callback_failure(callback_name="S3Logger") - - def web_identity_token(self, reference: str) -> str: - if reference.startswith("os.environ/") or reference in os.environ: - raise ValueError("Invalid web identity token reference.") - match _read_secret(get_secret, reference): - case None: - raise ValueError("OIDC token could not be retrieved from secret manager.") - case str() as token if len(token) < _WEB_IDENTITY_TOKEN_MIN_LENGTH: - raise ValueError( - "Parameter validation failed:\nInvalid length for parameter WebIdentityToken, " - f"value: {len(token)}, valid min length: {_WEB_IDENTITY_TOKEN_MIN_LENGTH}" - ) - case str() as token: - return token - case other: - raise ValueError( - "Parameter validation failed:\nInvalid type for parameter WebIdentityToken, " - f"value: {other}, type: {type(other)}, valid types: {str}" - ) - - -def _read_secret(read: Callable[[str], object], reference: str) -> object: - return read(reference) - - -def _resolved_params(params_source: Mapping[str, object] | None) -> dict[str, object]: # mutable-ok: handed to Rust - configured: Final[Mapping[object, object] | None] = getattr(litellm, "s3_callback_params", None) - source: Final[Mapping[str, object] | Mapping[object, object]] = ( - params_source if params_source is not None else (configured or {}) - ) - return { - key: (litellm.get_secret(value) if isinstance(value, str) and value.startswith("os.environ/") else value) - for key, value in source.items() - if isinstance(key, str) - } - - -def _constants(batch_size: int | None = DEFAULT_S3_BATCH_SIZE) -> Mapping[str, object]: - return { - "batch_size": batch_size or litellm.DEFAULT_BATCH_SIZE, - "max_concurrent_uploads": DEFAULT_S3_MAX_CONCURRENT_UPLOADS, - "max_adaptive_concurrency": DEFAULT_S3_MAX_ADAPTIVE_CONCURRENCY, - "upload_attempt_timeout_seconds": S3_UPLOAD_ATTEMPT_TIMEOUT_SECONDS, - "download_read_timeout_seconds": COMPLETION_HTTP_FALLBACK_SECONDS, - "download_connect_timeout_seconds": HTTP_HANDLER_CONNECT_TIMEOUT_SECONDS, - "web_identity_session_policy": session_policy_document(S3_LOGGING_SESSION_POLICY), - } - - -def _flush_interval(flush_interval: float | None = DEFAULT_S3_FLUSH_INTERVAL_SECONDS) -> float: - return flush_interval or litellm.DEFAULT_FLUSH_INTERVAL_SECONDS - - -_CONFIGURED_FIRST: Final = frozenset( - { - "s3_bucket_name", - "s3_path", - "s3_region_name", - "s3_api_version", - "s3_endpoint_url", - "s3_aws_access_key_id", - "s3_aws_secret_access_key", - "s3_aws_session_token", - "s3_aws_session_name", - "s3_aws_profile_name", - "s3_aws_role_name", - "s3_aws_web_identity_token", - "s3_aws_sts_endpoint", - "s3_config", - "s3_server_side_encryption", - "s3_sse_kms_key_id", - } -) -_CONSTRUCTOR_ENABLES: Final = frozenset( - { - "s3_use_team_prefix", - "s3_strip_base64_files", - "s3_use_key_prefix", - "s3_use_virtual_hosted_style", - "s3_adaptive_concurrency", - "s3_batch_file_upload", - } -) -_CONFIGURED_UNLESS_NONE: Final = frozenset({"s3_use_ssl", "s3_verify", "s3_drop_on_terminal_error"}) -_CONSTRUCTOR_UNLESS_NONE: Final = frozenset({"s3_log_prompts_only", "s3_partition_granularity"}) -_CONFIGURED_UNLESS_BLANK: Final = frozenset({"s3_max_concurrent_uploads", "s3_max_adaptive_concurrency"}) -_CONSTRUCTOR_DEFAULTS: Final[Mapping[str, object]] = { - "s3_use_ssl": True, - "s3_drop_on_terminal_error": True, - "s3_max_concurrent_uploads": DEFAULT_S3_MAX_CONCURRENT_UPLOADS, - "s3_max_retry_age_seconds": _DEFAULT_MAX_RETRY_AGE_SECONDS, - "s3_flush_interval": DEFAULT_S3_FLUSH_INTERVAL_SECONDS, - "s3_batch_size": DEFAULT_S3_BATCH_SIZE, -} - - -def _blank(value: object) -> bool: - return value is None or value == "" - - -def _constructor_wins(key: str, configured: object, value: object) -> bool: - return ( - (key in _CONFIGURED_FIRST and not configured) - or (key in _CONFIGURED_UNLESS_NONE and configured is None) - or (key in _CONFIGURED_UNLESS_BLANK and _blank(configured)) - or (key in _CONSTRUCTOR_UNLESS_NONE and value is not None) - ) - - -_KEEP_CONFIGURED: Final = object() - - -def _constructor_value(key: str, configured: object, value: object) -> object: - if _constructor_wins(key, configured, value): - return value - if key in _CONSTRUCTOR_ENABLES and value: - return True - if key == "s3_max_queue_size" and value is not None: - return resolve_s3_max_queue_size( - configured, resolve_s3_max_queue_size(value, CustomBatchLogger.DEFAULT_MAX_QUEUE_SIZE) - ) - if key == "s3_max_retry_age_seconds": - fallback: Final = resolve_s3_max_retry_age_seconds(value, _DEFAULT_MAX_RETRY_AGE_SECONDS) - resolved: Final = fallback if _blank(configured) else resolve_s3_max_retry_age_seconds(configured, fallback) - return 0 if resolved is None else resolved - return _KEEP_CONFIGURED - - -def _constructed_params( - params_source: Mapping[str, object] | None, settings: S3LoggerSettings -) -> dict[str, object]: # mutable-ok: handed to Rust - params: Final = _resolved_params(params_source) - overrides: Final = { - key: replaced - for key, value in settings.items() - if value != _CONSTRUCTOR_DEFAULTS.get(key, False if key in _CONSTRUCTOR_ENABLES else None) - and (replaced := _constructor_value(key, params.get(key), value)) is not _KEEP_CONFIGURED - } - return {**params, **overrides} - - -class _Base64Stripper(CustomLogger): - def strip(self, payload: StandardLoggingPayload) -> None: - _ = self._strip_base64_from_messages_sync(payload) - - def stripped(self, messages: list[object]) -> list[dict[str, object]]: - return self._process_messages(messages=[_shallow_copy(message) for message in messages]) - - -def _shallow_copy(message: object) -> object: - entries: Final = mapping(message) - return copy.copy(entries) if entries is not None and isinstance(message, dict) else message - - -_STRIPPER: Final = _Base64Stripper() - - -def legacy_should_log(kwargs: Mapping[str, object]) -> bool: - from litellm.litellm_core_utils import litellm_logging - - if litellm.global_disable_no_log_param: - return True - litellm_params: Final = kwargs.get("litellm_params", {}) - lookup: Final = getattr(litellm_params, "get", None) - if callable(lookup) and lookup("no-log", False) is True: - return False - return not litellm_logging.callback_disabled_dynamically( - callback=CALLBACK_NAME, - litellm_params=litellm_params, - standard_callback_dynamic_params=kwargs.get("standard_callback_dynamic_params") or {}, # pyright: ignore[reportArgumentType] # the dynamic params dict the legacy path hands the same check - ) - - -def _payload_messages(payload: Mapping[str, object]) -> list[object]: - raw: Final = payload.get("messages") - return [*raw] if isinstance(raw, list) else [] - - -class S3StripBase64Files(CustomLoggerV2): - selection = Selection(name=STRIP_CALLBACK_NAME) - - def __init__(self, s3_logger_key: str) -> None: - self.s3_logger_key: Final = s3_logger_key - - @on( - Hook.log_prepare, - reads={Field.deprecated_kwargs, Field.standard_logging_payload}, - writes={Field.standard_logging_payload}, - ) - async def on_log_prepare(self, view: object, ctx: CallContext[object]) -> Patch | None: - from litellm.litellm_core_utils.litellm_logging import V2_ASYNC_LOGGING_LINE - - kwargs: Final = mapping(getattr(view, Field.deprecated_kwargs.value, None)) - payload: Final = mapping(getattr(view, Field.standard_logging_payload.value, None)) - if not V2_ASYNC_LOGGING_LINE.get() or kwargs is None or payload is None or not legacy_should_log(kwargs): - return None - messages: Final = _payload_messages(payload) - if not messages: - return None - return Patch().merge(Field.standard_logging_payload, {"messages": _STRIPPER.stripped(messages)}) - - -async def _periodic_flush(sink: "NativeS3Log", interval: float) -> None: - while True: - await asyncio.sleep(interval) - try: - batch_upload_failed(await sink.flush()) - except (asyncio.CancelledError, KeyboardInterrupt, SystemExit): - raise - except BaseException as e: # noqa: BLE001 # a failed flush keeps the queue and the next interval tries again, as flush_queue does - verbose_logger.exception("CustomLogger: periodic flush failed: %s", e) - - -def _excluding_fields(kwargs: object) -> Mapping[str, object] | None: - excluded: Final[object] = getattr(litellm, "standard_logging_payload_excluded_fields", None) - details: Final = mapping(kwargs) - logged: Final = None if details is None else mapping(details.get("standard_logging_object")) - if details is None or logged is None or not isinstance(excluded, Collection) or not excluded: - return None - if logged.get("status") == "failure": - return None - return {**details, "standard_logging_object": {key: value for key, value in logged.items() if key not in excluded}} - - -def _logged_payload(kwargs: object, details: Mapping[str, object] | None) -> object: - source: Final = mapping(kwargs) if details is None else details - return None if source is None else source.get("standard_logging_object") - - -def _legacy_kwargs(view: object) -> object: - return getattr(view, Field.deprecated_kwargs.value, None) - - -_read_attribute: Final[Callable[[object, str], object]] = getattr - - -class _ExcludedView: - def __init__(self, view: object, details: Mapping[str, object]) -> None: - args: Final = HookArgs(request=details) - self.view: Final = view - self.fields: Final[Mapping[str, object]] = { - Field.deprecated_kwargs.value: details, - Field.standard_logging_payload_ref.value: payload_copy(args), - Field.log_id.value: log_id(args), - Field.cold_storage_object_key.value: cold_storage_object_key(args), - Field.logged_key_identity.value: logged_key_identity(args), - } - - def __getattr__(self, name: str) -> object: - if name not in self.fields: - return _read_attribute(self.view, name) - value: Final = self.fields[name] - if value is None: - raise AttributeError(name) - return value - - -_ATTEMPT_ENDED_READS: Final = frozenset( - { - Field.deprecated_kwargs, - Field.attempt_timing, - Field.log_id, - Field.cold_storage_object_key, - Field.standard_logging_payload_ref, - Field.logged_key_identity, - } -) - - -class S3Logger(RustLogger): - selection = Selection(name=CALLBACK_NAME, honors_no_log=True) - - def __init__( - self, - s3_callback_params_override: Mapping[str, object] | None = None, - **settings: Unpack[S3LoggerSettings], # kwargs-ok: the S3Logger constructor keywords - ) -> None: - from litellm.rust_bridge._native import NativeS3Log - - host: Final = _Host() - params: Final = _constructed_params(s3_callback_params_override, settings) - sink: Final = NativeS3Log( - params, - _constants(settings.get("s3_batch_size", DEFAULT_S3_BATCH_SIZE)), - host, - _cold_storage_is_s3_v2(), - None if params.get("s3_region_name") else _region_from_environment(), - ) - self.s3_region_name: Final[object] = params.get("s3_region_name") or settings.get("s3_region_name") - self.sink: Final = sink - self.s3_host: Final = host - self.periodic_flush_task: Final = asyncio.create_task( - _periodic_flush(sink, _flush_interval(settings.get("s3_flush_interval", DEFAULT_S3_FLUSH_INTERVAL_SECONDS))) - ) - super().__init__(sink.callback(), {Field.deprecated_kwargs: lambda _kwargs: None}) - self.v2_companion: Final = ( - S3StripBase64Files(litellm.logging_callback_manager.get_custom_logger_key(self)) - if sink.strip_base64_files - else None - ) - - @on(Hook.attempt_ended, reads=_ATTEMPT_ENDED_READS, awaited=True) - async def on_attempt_ended(self, view: object, ctx: CallContext[object]) -> Patch | Abort | None: - return await self.run_rust(Hook.attempt_ended, _ATTEMPT_ENDED_READS, view, ctx) - - async def run_rust( - self, hook: Hook, reads: frozenset[Field], view: object, ctx: CallContext[object] - ) -> Patch | Abort | None: - from litellm.litellm_core_utils.litellm_logging import V2_ASYNC_LOGGING_LINE - - kwargs: Final = getattr(view, Field.deprecated_kwargs.value, None) - if not V2_ASYNC_LOGGING_LINE.get() or not isinstance(kwargs, Mapping) or not legacy_should_log(kwargs): # pyright: ignore[reportUnknownArgumentType] # the legacy kwargs dict - return None - payload: Final[object] = kwargs.get("standard_logging_object") # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # the legacy kwargs dict - if payload is not None and not isinstance(payload, Mapping): - return self._failed(AttributeError(f"'{type(payload).__name__}' object has no attribute 'get'")) # pyright: ignore[reportUnknownArgumentType] # any non-mapping payload - try: - raw: Final = _legacy_kwargs(view) - details: Final = _excluding_fields(raw) - logged: Final = _logged_payload(raw, details) - if self.sink.strip_base64_files and isinstance(logged, Mapping): - _STRIPPER.strip(logged) # pyright: ignore[reportArgumentType] # the shared StandardLoggingPayload, stripped in place as S3Logger does - return await super().run_rust(hook, reads, view if details is None else _ExcludedView(view, details), ctx) - except Exception as e: # noqa: BLE001 # S3Logger counts a failure to build the element and moves on - return self._failed(e) - - def _failed(self, error: Exception) -> None: - verbose_logger.error("s3 Layer Error - %s", error, exc_info=error) - self.s3_host.callback_failed() - - @property - def s3_path(self) -> str | None: - return self.sink.s3_path - - @property - def s3_bucket_name(self) -> str | None: - return self.sink.s3_bucket_name - - def resolve_partition_granularity(self) -> S3PartitionGranularity: - return "hour" if self.sink.partition_granularity() == "hour" else "day" - - async def flush_queue(self) -> None: - batch_upload_failed(await self.sink.flush()) - - async def async_log_audit_log_event(self, audit_log: StandardAuditLogPayload) -> None: - try: - audit_log_id: Final = audit_log.get("id", "unknown") - batch_upload_failed(await self.sink.submit_audit(dict(audit_log), f"{audit_log_id}")) - except Exception as e: # noqa: BLE001 # an audit log failure is logged and swallowed, as in S3Logger - verbose_logger.exception("S3 audit log error: %s", e) - - async def async_log_batch_line_items(self, payload: StandardBatchLineItemsPayload) -> None: - object_name: Final = batch_line_items_object_name(payload) - s3_path: Final = self.s3_path - s3_object_key: Final = (f"{s3_path.rstrip('/')}/" if s3_path else "") + object_name - body: Final = await asyncio.to_thread(batch_line_items_ndjson, payload) - delivered: Final = await self.sink.upload_line_items( - s3_object_key, body, "-".join(object_name.rsplit("/", 2)[-2:]) - ) - if not delivered: - verbose_logger.error( - "s3 logging: the %d line items of part %d of batch %s were not stored at %s", - len(payload["lines"]), - payload["part"], - payload["batch_id"], - s3_object_key, - ) - - async def _download_object_from_s3(self, s3_object_key: str) -> object: - try: - downloaded: Final = await self.sink.download(s3_object_key) - return None if downloaded is None else json.loads(downloaded) - except Exception as e: # noqa: BLE001 # a failed read is None, as in S3Logger - verbose_logger.exception("Error downloading from S3: %s", e) - return None - - async def get_proxy_server_request_from_cold_storage_with_object_key(self, object_key: str) -> object: - try: - return await self._download_object_from_s3(object_key) - except Exception as e: # noqa: BLE001 # a failed read is None, as in S3Logger - verbose_logger.exception("Error retrieving object %s from cold storage: %s", object_key, e) - return None - - -__all__ = [ - "S3BatchUploadError", - "S3Logger", - "S3StripBase64Files", - "legacy_should_log", -] diff --git a/litellm/litellm_core_utils/custom_logger_registry.py b/litellm/litellm_core_utils/custom_logger_registry.py index 535fef8fcec..74badc2c8c1 100644 --- a/litellm/litellm_core_utils/custom_logger_registry.py +++ b/litellm/litellm_core_utils/custom_logger_registry.py @@ -46,7 +46,7 @@ from litellm.integrations.opik.opik import OpikLogger from litellm.integrations.pointfive import PointFiveLogger from litellm.integrations.posthog import PostHogLogger from litellm.integrations.prometheus import PrometheusLogger -from litellm.integrations.s3_v3 import S3Logger +from litellm.integrations.s3_v2 import S3Logger from litellm.integrations.sqs import SQSLogger from litellm.integrations.ternary.ternary_logger import TernaryLogger from litellm.integrations.vantage.vantage_logger import VantageLogger diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 15a95e91c52..6ab55772394 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -6536,7 +6536,7 @@ def build_custom_logger_compatible_class( _in_memory_loggers.append(_gcs_bucket_logger) return _gcs_bucket_logger elif logging_integration == "s3_v2": - from litellm.integrations.s3_v3 import S3Logger as S3V2Logger + from litellm.integrations.s3_v2 import S3Logger as S3V2Logger for callback in _in_memory_loggers: if isinstance(callback, S3V2Logger): @@ -7360,7 +7360,7 @@ def get_custom_logger_compatible_class( if isinstance(callback, GCSBucketLogger): return callback elif logging_integration == "s3_v2": - from litellm.integrations.s3_v3 import S3Logger as S3V2Logger + from litellm.integrations.s3_v2 import S3Logger as S3V2Logger for callback in _in_memory_loggers: if isinstance(callback, S3V2Logger): @@ -8067,7 +8067,7 @@ class StandardLoggingPayloadSetup: ) if custom_logger and hasattr(custom_logger, "s3_path") and getattr(custom_logger, "s3_path"): s3_path = getattr(custom_logger, "s3_path") - from litellm.integrations.s3_v3 import S3Logger as S3V2Logger + from litellm.integrations.s3_v2 import S3Logger as S3V2Logger if isinstance(custom_logger, S3V2Logger): partition_granularity = custom_logger.resolve_partition_granularity() diff --git a/litellm/proxy/management_helpers/audit_logs.py b/litellm/proxy/management_helpers/audit_logs.py index 65b64a297c9..c5ee783bea1 100644 --- a/litellm/proxy/management_helpers/audit_logs.py +++ b/litellm/proxy/management_helpers/audit_logs.py @@ -22,7 +22,7 @@ from litellm.repositories.table_repositories import AuditLogRepository from litellm.types.utils import StandardAuditLogPayload if TYPE_CHECKING: - from litellm.integrations.s3_v3 import S3Logger + from litellm.integrations.s3_v2 import S3Logger _audit_log_callback_cache: Final[dict[str, "CustomLogger | S3Logger"]] = {} ALLOW_LITELLM_CHANGED_BY_HEADER_METADATA_KEY: Final = "allow_litellm_changed_by_header" @@ -78,9 +78,9 @@ def _resolve_audit_log_callback(name: str) -> "CustomLogger | S3Logger | None": instance: CustomLogger | S3Logger | None audit_params: Final[Mapping[str, object] | None] = getattr(litellm, "s3_audit_callback_params", None) if name == "s3_v2" and audit_params is not None: - from litellm.integrations import s3_v3 + from litellm.integrations import s3_v2 - instance = s3_v3.S3Logger(s3_callback_params_override=audit_params) + instance = s3_v2.S3Logger(s3_callback_params_override=audit_params) else: from litellm.litellm_core_utils.litellm_logging import ( _init_custom_logger_compatible_class, diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index eb238a7aa9d..bb71779fd70 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -7452,7 +7452,7 @@ class ProxyConfig: if key == "request_timeout": litellm.request_timeout_explicitly_set = True if key in {"s3_audit_callback_params", "s3_callback_params"}: - from litellm.integrations.s3_v3 import S3Logger as S3V2Logger + from litellm.integrations.s3_v2 import S3Logger as S3V2Logger from litellm.litellm_core_utils.litellm_logging import ( _in_memory_loggers, ) diff --git a/tests/e2e/coverage_registry/logging.yaml b/tests/e2e/coverage_registry/logging.yaml index 6eebb3c5173..69cad8eb22e 100644 --- a/tests/e2e/coverage_registry/logging.yaml +++ b/tests/e2e/coverage_registry/logging.yaml @@ -1,7 +1,7 @@ # Logging integration delivery (behavior features). Grounded in litellm/integrations/. -- {id: logging.s3.success.writes_object, module: logging, tier: P0, event: success, assertions: [writes_object], exercised_on: [chat_completions, messages, embeddings], source: "integrations/s3_v3.py", rationale: "Primary audit trail; batch flush no-drop"} -- {id: logging.s3.failure.writes_object, module: logging, tier: P0, event: failure, assertions: [writes_object], exercised_on: [chat_completions, messages], source: "integrations/s3_v3.py", rationale: "Failed calls persisted for compliance"} -- {id: logging.s3.success.partition_layout, module: logging, tier: P1, event: success, assertions: [object_key_layout], exercised_on: [chat_completions], source: "integrations/s3_v3.py / LIT-8985", rationale: "s3_partition_granularity picks the date or date/hour folder every downstream query and lifecycle rule reads"} +- {id: logging.s3.success.writes_object, module: logging, tier: P0, event: success, assertions: [writes_object], exercised_on: [chat_completions, messages, embeddings], source: "integrations/s3_v2.py", rationale: "Primary audit trail; batch flush no-drop"} +- {id: logging.s3.failure.writes_object, module: logging, tier: P0, event: failure, assertions: [writes_object], exercised_on: [chat_completions, messages], source: "integrations/s3_v2.py", rationale: "Failed calls persisted for compliance"} +- {id: logging.s3.success.partition_layout, module: logging, tier: P1, event: success, assertions: [object_key_layout], exercised_on: [chat_completions], source: "integrations/s3_v2.py / LIT-8985", rationale: "s3_partition_granularity picks the date or date/hour folder every downstream query and lifecycle rule reads"} - {id: logging.gcs_bucket.success.writes_object, module: logging, tier: P0, event: success, assertions: [writes_object], exercised_on: [chat_completions, messages, embeddings], source: "integrations/gcs_bucket/gcs_bucket.py", rationale: "GCS parallel to S3"} - {id: logging.datadog.success.exports_metric, module: logging, tier: P0, event: success, assertions: [exports_metric], exercised_on: [chat_completions, messages, responses, embeddings], source: "integrations/datadog/datadog.py", rationale: "Powers dashboards/alerts; cardinality regressions common"} - {id: logging.datadog.stream.exports_metric, module: logging, tier: P0, event: stream, assertions: [exports_metric], exercised_on: [chat_completions, messages, responses], source: "integrations/datadog/datadog.py", rationale: "Streaming aggregates usage after the last chunk; delivery and cost must survive that path"} diff --git a/tests/unit/integrations/s3_v3_rust_support.py b/tests/unit/integrations/s3_v2_rust_support.py similarity index 99% rename from tests/unit/integrations/s3_v3_rust_support.py rename to tests/unit/integrations/s3_v2_rust_support.py index fcd566c98fb..e188c51405c 100644 --- a/tests/unit/integrations/s3_v3_rust_support.py +++ b/tests/unit/integrations/s3_v2_rust_support.py @@ -16,7 +16,7 @@ from litellm.integrations.custom_logger_v2.enforcement import FieldView from litellm.integrations.custom_logger_v2.field_sources import HookArgs, with_sources from litellm.integrations.custom_logger_v2.legacy import KwargsFields from litellm.integrations.custom_logger_v2.rust import WireManifest -from litellm.integrations.s3_v3 import S3Logger +from litellm.integrations.s3_v2 import S3Logger from litellm.litellm_core_utils.litellm_logging import LOGGING_KEYS, V2_ASYNC_LOGGING_LINE from litellm.types.integrations.custom_logger_v2 import Outcome diff --git a/tests/unit/integrations/test_s3_v3_native.py b/tests/unit/integrations/test_s3_v2_native.py similarity index 98% rename from tests/unit/integrations/test_s3_v3_native.py rename to tests/unit/integrations/test_s3_v2_native.py index 86e14d3fea9..50c4724127f 100644 --- a/tests/unit/integrations/test_s3_v3_native.py +++ b/tests/unit/integrations/test_s3_v2_native.py @@ -21,7 +21,7 @@ import pytest from pydantic import BaseModel import litellm -from litellm.integrations import s3_v3 +from litellm.integrations import s3_v2 from litellm.integrations.custom_logger_v2.derivations import DERIVED_ATTEMPT_FIELDS from litellm.integrations.custom_logger_v2.enforcement import FieldView from litellm.integrations.custom_logger_v2.enums import Field, Hook @@ -29,7 +29,7 @@ from litellm.integrations.custom_logger_v2.host import request_scope from litellm.integrations.custom_logger_v2.legacy import KwargsFields from litellm.integrations.custom_logger_v2.field_sources import HookArgs, with_sources from litellm.integrations.custom_logger_v2.selection import RequestSelection -from litellm.integrations.s3_v3 import S3Logger +from litellm.integrations.s3_v2 import S3Logger from litellm.litellm_core_utils.litellm_logging import LOGGING_KEYS, V2_ASYNC_LOGGING_LINE from litellm.llms.bedrock import base_aws_llm from litellm.types.integrations.custom_logger_v2 import Outcome @@ -505,7 +505,7 @@ async def test_both_loggers_skip_and_count_the_same_bad_payloads( golden: Golden, bucket: FakeBucket, monkeypatch: pytest.MonkeyPatch, payload: object ) -> None: logger = await _logger(monkeypatch, bucket, {"s3_use_team_prefix": True}) - with patch.object(s3_v3._Host, "callback_failed") as rust_failures: + with patch.object(s3_v2._Host, "callback_failed") as rust_failures: await _log(logger, payload) await logger.flush_queue() await asyncio.sleep(0.2) @@ -703,7 +703,7 @@ def test_both_loggers_need_a_running_event_loop(golden: Golden, monkeypatch: pyt async def test_the_native_logger_flushes_on_its_own_after_the_interval( golden: Golden, bucket: FakeBucket, monkeypatch: pytest.MonkeyPatch ) -> None: - monkeypatch.setattr(s3_v3, "DEFAULT_S3_FLUSH_INTERVAL_SECONDS", 0.05) + monkeypatch.setattr(s3_v2, "DEFAULT_S3_FLUSH_INTERVAL_SECONDS", 0.05) logger = await _logger(monkeypatch, bucket, {}) await _log_native(logger, {"standard_logging_object": _payload()}) for _ in range(100): @@ -735,7 +735,7 @@ async def test_without_a_region_param_both_sign_with_the_same_environment_region for name, value in environment.items(): monkeypatch.setenv(name, value) logger = await _logger(monkeypatch, bucket, {"s3_region_name": None}) - with patch.object(s3_v3._Host, "callback_failed") as rust_failures: + with patch.object(s3_v2._Host, "callback_failed") as rust_failures: await _log(logger, _payload()) await logger.flush_queue() await asyncio.sleep(0.2) @@ -852,7 +852,7 @@ async def test_a_web_identity_token_botocore_rejects_fails_both_before_sts( return token if name == reference else default_value monkeypatch.setattr(base_aws_llm, "get_secret", secret) - monkeypatch.setattr(s3_v3, "get_secret", secret) + monkeypatch.setattr(s3_v2, "get_secret", secret) identity: Final = { "s3_aws_access_key_id": None, "s3_aws_secret_access_key": None, @@ -863,7 +863,7 @@ async def test_a_web_identity_token_botocore_rejects_fails_both_before_sts( "s3_aws_sts_endpoint": f"http://127.0.0.1:{sts.port}", } logger = await _logger(monkeypatch, bucket, identity) - with patch.object(s3_v3._Host, "callback_failed") as rust_failures: + with patch.object(s3_v2._Host, "callback_failed") as rust_failures: await _log(logger, _payload()) await logger.flush_queue() assert sts.calls == [] @@ -873,7 +873,7 @@ async def test_a_web_identity_token_botocore_rejects_fails_both_before_sts( assert logger.sink.queued_keys() == golden.expected("queue") python_error: Final = golden.expected("error") with pytest.raises(ValueError, match=r"^Parameter validation failed:") as rust_error: - s3_v3._Host().web_identity_token(reference) + s3_v2._Host().web_identity_token(reference) assert str(rust_error.value) == python_error @@ -918,7 +918,7 @@ async def test_the_native_opt_out_matches_should_run_callback( "standard_callback_dynamic_params": logging_obj.standard_callback_dynamic_params, "standard_logging_object": _payload(), } - assert s3_v3.legacy_should_log(kwargs) is expected + assert s3_v2.legacy_should_log(kwargs) is expected await _log_native(logger, kwargs) await logger.flush_queue() assert len(bucket.puts()) == int(expected) @@ -979,7 +979,7 @@ async def test_a_payload_that_changes_while_it_is_rendered_is_kept_for_retry_by_ metadata: Final = racing["metadata"] assert isinstance(metadata, dict) metadata["mutating"] = _Mutating(metadata) # pyright: ignore[reportUnknownArgumentType] # the metadata dict _payload builds - with patch.object(s3_v3._Host, "callback_failed") as rust_failures: + with patch.object(s3_v2._Host, "callback_failed") as rust_failures: await _log_native(logger, {"standard_logging_object": racing}) await _log(logger, _payload("steady")) await logger.flush_queue() @@ -1010,7 +1010,7 @@ class _Panic(BaseException): @pytest.mark.asyncio async def test_the_periodic_flush_keeps_running_after_a_flush_fails() -> None: sink: Final = _FlakySink((RuntimeError("bucket gone"), _Panic("a PanicException is a BaseException too"))) - task: Final = asyncio.create_task(s3_v3._periodic_flush(sink, 0.01)) # pyright: ignore[reportArgumentType] # the loop only calls flush + task: Final = asyncio.create_task(s3_v2._periodic_flush(sink, 0.01)) # pyright: ignore[reportArgumentType] # the loop only calls flush for _ in range(100): if sink.flushed >= 3: break @@ -1200,7 +1200,7 @@ async def test_both_loggers_widen_the_adaptive_upload_limit_alike( async def test_audit_events_flush_at_the_same_batch_size( golden: Golden, bucket: FakeBucket, monkeypatch: pytest.MonkeyPatch, batch_size: int, puts: int ) -> None: - monkeypatch.setattr(s3_v3, "DEFAULT_S3_BATCH_SIZE", batch_size) + monkeypatch.setattr(s3_v2, "DEFAULT_S3_BATCH_SIZE", batch_size) logger = await _logger(monkeypatch, bucket, {}) for number in (1, 2): audit = {"id": f"audit-{number}", "action": "created", "updated_at": "2026-11-13"} @@ -1216,13 +1216,13 @@ async def test_audit_events_flush_at_the_same_batch_size( async def test_both_loggers_keep_uploads_that_fail_in_transport( golden: Golden, bucket: FakeBucket, monkeypatch: pytest.MonkeyPatch, failure: str ) -> None: - monkeypatch.setattr(s3_v3, "S3_UPLOAD_ATTEMPT_TIMEOUT_SECONDS", 0.3) + monkeypatch.setattr(s3_v2, "S3_UPLOAD_ATTEMPT_TIMEOUT_SECONDS", 0.3) if failure == "timeout": bucket.delays.append(("_a\\.json", 1.0)) else: bucket.failures.append(("_a\\.json", 0, "")) logger = await _logger(monkeypatch, bucket, {}) - with patch.object(s3_v3._Host, "callback_failed") as rust_failures: + with patch.object(s3_v2._Host, "callback_failed") as rust_failures: for response_id in ("a", "b"): await _log(logger, _payload(response_id)) await logger.flush_queue() @@ -1262,7 +1262,7 @@ async def test_cold_storage_handles_odd_metadata_like_python( del payload["metadata"] else: payload["metadata"] = metadata - with patch.object(s3_v3._Host, "callback_failed") as rust_failures: + with patch.object(s3_v2._Host, "callback_failed") as rust_failures: await _log(logger, payload) await logger.flush_queue() await asyncio.sleep(0.2) @@ -1345,7 +1345,7 @@ class _AliasLookupFails(dict[str, object]): async def test_the_attempt_ended_hook_only_enqueues_while_s3_is_slow( golden: Golden, bucket: FakeBucket, monkeypatch: pytest.MonkeyPatch ) -> None: - monkeypatch.setattr(s3_v3, "DEFAULT_S3_BATCH_SIZE", 1) + monkeypatch.setattr(s3_v2, "DEFAULT_S3_BATCH_SIZE", 1) logger = await _logger(monkeypatch, bucket, {}) bucket.delays.append((".", 3.0)) started: Final = time.monotonic() @@ -1369,7 +1369,7 @@ async def test_a_payload_whose_alias_lookup_fails_is_counted_and_released( golden: Golden, bucket: FakeBucket, monkeypatch: pytest.MonkeyPatch ) -> None: logger = await _logger(monkeypatch, bucket, {"s3_use_team_prefix": True}) - with patch.object(s3_v3._Host, "callback_failed") as rust_failures: + with patch.object(s3_v2._Host, "callback_failed") as rust_failures: await _log(logger, _payload(metadata=_AliasLookupFails(team="t"))) await logger.flush_queue() await asyncio.sleep(0.2) diff --git a/tests/unit/integrations/test_s3_v2_rust.py b/tests/unit/integrations/test_s3_v2_rust.py index 3c19d63c615..84dc0cf40dd 100644 --- a/tests/unit/integrations/test_s3_v2_rust.py +++ b/tests/unit/integrations/test_s3_v2_rust.py @@ -16,7 +16,7 @@ from urllib.parse import quote, unquote, urlsplit, urlunsplit import pytest import requests -from tests.unit.integrations.s3_v3_rust_support import ( +from tests.unit.integrations.s3_v2_rust_support import ( DROP, Received, ScriptedBucket, @@ -27,8 +27,8 @@ from tests.unit.integrations.s3_v3_rust_support import ( ) import litellm -from litellm.integrations import s3_v3 -from litellm.integrations.s3_v3 import S3Logger +from litellm.integrations import s3_v2 +from litellm.integrations.s3_v2 import S3Logger from litellm.litellm_core_utils.safe_json_dumps import safe_dumps from litellm.rust_bridge import _native from litellm.types.integrations.s3_v2 import S3LoggerSettings @@ -83,7 +83,7 @@ async def _upload(logger: S3Logger, payload: Mapping[str, object], start_time: d def _failures() -> object: - return patch.object(s3_v3._Host, "callback_failed") + return patch.object(s3_v2._Host, "callback_failed") def _uploaded(put: Received) -> dict[str, object]: @@ -2117,14 +2117,14 @@ async def test_batch_line_items_are_serialized_and_hashed_off_the_event_loop( import threading logger: Final = _s3_logger_answering(monkeypatch, bucket, 200, "cold/") - serialize: Final = s3_v3.batch_line_items_ndjson + serialize: Final = s3_v2.batch_line_items_ndjson threads: Final[dict[str, threading.Thread]] = {} def _serialize(payload: object) -> str: threads["serialize"] = threading.current_thread() return serialize(payload) - monkeypatch.setattr(s3_v3, "batch_line_items_ndjson", _serialize) + monkeypatch.setattr(s3_v2, "batch_line_items_ndjson", _serialize) await logger.async_log_batch_line_items(_batch_line_items_payload()) diff --git a/tests/unit/integrations/test_s3_v3_strip_writer.py b/tests/unit/integrations/test_s3_v2_strip_writer.py similarity index 99% rename from tests/unit/integrations/test_s3_v3_strip_writer.py rename to tests/unit/integrations/test_s3_v2_strip_writer.py index 3678624ca09..fdb344d42d3 100644 --- a/tests/unit/integrations/test_s3_v3_strip_writer.py +++ b/tests/unit/integrations/test_s3_v2_strip_writer.py @@ -14,13 +14,13 @@ from litellm.integrations.custom_logger import CustomLogger from litellm.integrations.custom_logger_v2.enums import Field, Hook from litellm.integrations.custom_logger_v2.field_sources import HookArgs from litellm.integrations.custom_logger_v2.manifest import build_manifest -from litellm.integrations.s3_v3 import STRIP_CALLBACK_NAME, S3Logger, S3StripBase64Files +from litellm.integrations.s3_v2 import STRIP_CALLBACK_NAME, S3Logger, S3StripBase64Files from litellm.litellm_core_utils import litellm_logging from litellm.litellm_core_utils.litellm_logging import Logging from litellm.rust_bridge import _native from litellm.types.integrations.custom_logger_v2 import Outcome from litellm.types.utils import ModelResponse -from tests.unit.integrations.s3_v3_rust_support import ScriptedBucket, scripted_bucket +from tests.unit.integrations.s3_v2_rust_support import ScriptedBucket, scripted_bucket pytestmark = pytest.mark.skipif(not hasattr(_native, "NativeS3Log"), reason="the built _native has no NativeS3Log") diff --git a/tests/unit/proxy/management_helpers/test_audit_log_callbacks.py b/tests/unit/proxy/management_helpers/test_audit_log_callbacks.py index e6323bd876b..1468e6ab092 100644 --- a/tests/unit/proxy/management_helpers/test_audit_log_callbacks.py +++ b/tests/unit/proxy/management_helpers/test_audit_log_callbacks.py @@ -22,7 +22,7 @@ from litellm.proxy.management_helpers.audit_logs import ( is_audit_logging_enabled, ) from litellm.types.utils import StandardAuditLogPayload -from tests.unit.integrations.s3_v3_rust_support import scripted_bucket +from tests.unit.integrations.s3_v2_rust_support import scripted_bucket @pytest.fixture(autouse=True) diff --git a/tests/unit/responses/litellm_completion_transformation/test_session_handler_with_cold_storage.py b/tests/unit/responses/litellm_completion_transformation/test_session_handler_with_cold_storage.py index 8fd97ea5bce..522b3258961 100644 --- a/tests/unit/responses/litellm_completion_transformation/test_session_handler_with_cold_storage.py +++ b/tests/unit/responses/litellm_completion_transformation/test_session_handler_with_cold_storage.py @@ -24,7 +24,7 @@ from litellm.responses.litellm_completion_transformation.session_handler import ResponsesSessionHandler, ) from litellm.types.utils import StandardLoggingMetadata -from tests.unit.integrations.s3_v3_rust_support import scripted_bucket +from tests.unit.integrations.s3_v2_rust_support import scripted_bucket class TestColdStorageObjectKeyIntegration: