知微
fa45d8aa5f
- health_checklist.json: 192.168.1.122→node122
- ocr_client.py: docstring IP→node122
- docs/market-data-requirements.md: IP→node122
- 所有API调用通过ProxyHandler({})绕过系统代理
Privoxy对node122:18003返回500，直连正常
757 lines
30 KiB
Python
757 lines
30 KiB
Python
"""
|
|
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
|
|
async_log_failure_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 time
|
|
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.litellm_core_utils.safe_json_dumps import safe_dumps
|
|
from litellm.litellm_core_utils.sensitive_data_masker import SensitiveDataMasker
|
|
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 StandardAuditLogPayload, 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,
|
|
s3_strip_base64_files: bool = False,
|
|
s3_use_key_prefix: bool = False,
|
|
s3_use_virtual_hosted_style: bool = False,
|
|
s3_callback_params_override: Optional[dict] = None,
|
|
**kwargs,
|
|
):
|
|
try:
|
|
_masker = SensitiveDataMasker()
|
|
if s3_callback_params_override is not None:
|
|
verbose_logger.debug(
|
|
f"in init s3 logger (audit override) - "
|
|
f"{_masker.mask_dict(dict(s3_callback_params_override))}"
|
|
)
|
|
else:
|
|
verbose_logger.debug(
|
|
f"in init s3 logger - s3_callback_params "
|
|
f"{_masker.mask_dict(dict(litellm.s3_callback_params or {}))}"
|
|
)
|
|
|
|
# Initialize S3 params first to get the correct s3_verify value
|
|
self._init_s3_params(
|
|
params_source=s3_callback_params_override,
|
|
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,
|
|
s3_strip_base64_files=s3_strip_base64_files,
|
|
s3_use_key_prefix=s3_use_key_prefix,
|
|
s3_use_virtual_hosted_style=s3_use_virtual_hosted_style,
|
|
)
|
|
verbose_logger.debug(f"s3 logger using endpoint url {s3_endpoint_url}")
|
|
|
|
# IMPORTANT
|
|
# Create httpx client AFTER _init_s3_params so we have the correct s3_verify value
|
|
verbose_logger.debug(
|
|
f"s3_v2 logger creating async httpx client with s3_verify={self.s3_verify}"
|
|
)
|
|
self.async_httpx_client = get_async_httpx_client(
|
|
llm_provider=httpxSpecialProvider.LoggingCallback,
|
|
params={"ssl_verify": self.s3_verify},
|
|
)
|
|
|
|
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,
|
|
s3_strip_base64_files: bool = False,
|
|
s3_use_key_prefix: bool = False,
|
|
s3_use_virtual_hosted_style: bool = False,
|
|
params_source: Optional[dict] = None,
|
|
):
|
|
"""
|
|
Initialize the s3 params for this logging callback. Reads from
|
|
`params_source` if given (e.g. `s3_audit_callback_params` for the
|
|
audit-log instance), otherwise falls back to `litellm.s3_callback_params`.
|
|
Resolves `os.environ/X` markers into a local dict; never mutates the source.
|
|
"""
|
|
if params_source is None:
|
|
params_source = litellm.s3_callback_params or {}
|
|
params: dict = {
|
|
key: (
|
|
litellm.get_secret(value)
|
|
if isinstance(value, str) and value.startswith("os.environ/")
|
|
else value
|
|
)
|
|
for key, value in params_source.items()
|
|
}
|
|
|
|
self.s3_bucket_name = params.get("s3_bucket_name") or s3_bucket_name
|
|
self.s3_region_name = params.get("s3_region_name") or s3_region_name
|
|
self.s3_api_version = params.get("s3_api_version") or s3_api_version
|
|
self.s3_use_ssl = (
|
|
params.get("s3_use_ssl", True)
|
|
if params.get("s3_use_ssl") is not None
|
|
else s3_use_ssl
|
|
)
|
|
self.s3_verify = (
|
|
params.get("s3_verify")
|
|
if params.get("s3_verify") is not None
|
|
else s3_verify
|
|
)
|
|
self.s3_endpoint_url = params.get("s3_endpoint_url") or s3_endpoint_url
|
|
self.s3_aws_access_key_id = (
|
|
params.get("s3_aws_access_key_id") or s3_aws_access_key_id
|
|
)
|
|
|
|
self.s3_aws_secret_access_key = (
|
|
params.get("s3_aws_secret_access_key") or s3_aws_secret_access_key
|
|
)
|
|
|
|
self.s3_aws_session_token = (
|
|
params.get("s3_aws_session_token") or s3_aws_session_token
|
|
)
|
|
|
|
self.s3_aws_session_name = (
|
|
params.get("s3_aws_session_name") or s3_aws_session_name
|
|
)
|
|
|
|
self.s3_aws_profile_name = (
|
|
params.get("s3_aws_profile_name") or s3_aws_profile_name
|
|
)
|
|
|
|
self.s3_aws_role_name = params.get("s3_aws_role_name") or s3_aws_role_name
|
|
|
|
self.s3_aws_web_identity_token = (
|
|
params.get("s3_aws_web_identity_token") or s3_aws_web_identity_token
|
|
)
|
|
|
|
self.s3_aws_sts_endpoint = (
|
|
params.get("s3_aws_sts_endpoint") or s3_aws_sts_endpoint
|
|
)
|
|
|
|
self.s3_config = params.get("s3_config") or s3_config
|
|
self.s3_path = params.get("s3_path") or s3_path
|
|
self.s3_use_team_prefix = (
|
|
bool(params.get("s3_use_team_prefix", False)) or s3_use_team_prefix
|
|
)
|
|
|
|
self.s3_use_key_prefix = (
|
|
bool(params.get("s3_use_key_prefix", False)) or s3_use_key_prefix
|
|
)
|
|
|
|
self.s3_strip_base64_files = (
|
|
bool(params.get("s3_strip_base64_files", False)) or s3_strip_base64_files
|
|
)
|
|
|
|
self.s3_use_virtual_hosted_style = (
|
|
bool(params.get("s3_use_virtual_hosted_style", False))
|
|
or s3_use_virtual_hosted_style
|
|
)
|
|
|
|
return
|
|
|
|
async def async_log_success_event(self, kwargs, response_obj, start_time, end_time):
|
|
await self._async_log_event_base(
|
|
kwargs=kwargs,
|
|
response_obj=response_obj,
|
|
start_time=start_time,
|
|
end_time=end_time,
|
|
)
|
|
|
|
async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time):
|
|
await self._async_log_event_base(
|
|
kwargs=kwargs,
|
|
response_obj=response_obj,
|
|
start_time=start_time,
|
|
end_time=end_time,
|
|
)
|
|
pass
|
|
|
|
async def async_log_audit_log_event(
|
|
self, audit_log: StandardAuditLogPayload
|
|
) -> None:
|
|
"""Batch audit logs and upload to S3 under audit_logs/ prefix."""
|
|
try:
|
|
from datetime import timezone
|
|
|
|
now = datetime.now(timezone.utc)
|
|
audit_log_id = audit_log.get("id", "unknown")
|
|
|
|
s3_path = cast(Optional[str], self.s3_path) or ""
|
|
s3_path = s3_path.rstrip("/") + "/" if s3_path else ""
|
|
|
|
s3_object_key = (
|
|
f"{s3_path}audit_logs/"
|
|
f"{now.strftime('%Y-%m-%d')}/"
|
|
f"{now.strftime('%H-%M-%S')}_{audit_log_id}.json"
|
|
)
|
|
|
|
element = s3BatchLoggingElement(
|
|
payload=dict(audit_log),
|
|
s3_object_key=s3_object_key,
|
|
s3_object_download_filename=f"audit-{audit_log_id}.json",
|
|
)
|
|
|
|
self.log_queue.append(element)
|
|
|
|
if len(self.log_queue) >= self.batch_size:
|
|
await self.flush_queue()
|
|
except Exception as e:
|
|
verbose_logger.exception("S3 audit log error: %s", e)
|
|
|
|
async def _async_log_event_base(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),
|
|
)
|
|
|
|
# afile_delete and other non-model call types never produce a standard_logging_object,
|
|
# so s3_batch_logging_element is None. Skip gracefully instead of raising ValueError.
|
|
if s3_batch_logging_element is None:
|
|
verbose_logger.debug(
|
|
"s3 Logging - skipping event, no standard_logging_object for call_type=%s",
|
|
kwargs.get("call_type", "unknown"),
|
|
)
|
|
return
|
|
|
|
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)}")
|
|
self.handle_callback_failure(callback_name="S3Logger")
|
|
|
|
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}"
|
|
)
|
|
verbose_logger.debug(f"s3_v2 logger - s3_verify setting: {self.s3_verify}")
|
|
|
|
# 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 and self.s3_bucket_name:
|
|
if self.s3_use_virtual_hosted_style:
|
|
# Virtual-hosted-style: bucket.endpoint/key
|
|
endpoint_host = self.s3_endpoint_url.replace(
|
|
"https://", ""
|
|
).replace("http://", "")
|
|
protocol = (
|
|
"https://"
|
|
if self.s3_endpoint_url.startswith("https://")
|
|
else "http://"
|
|
)
|
|
url = f"{protocol}{self.s3_bucket_name}.{endpoint_host}/{batch_logging_element.s3_object_key}"
|
|
else:
|
|
# Path-style: endpoint/bucket/key
|
|
url = (
|
|
self.s3_endpoint_url
|
|
+ "/"
|
|
+ self.s3_bucket_name
|
|
+ "/"
|
|
+ batch_logging_element.s3_object_key
|
|
)
|
|
|
|
# Convert JSON to string
|
|
json_string = safe_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,
|
|
)
|
|
aws_region_name = self.get_aws_region_name_for_non_llm_api_calls(
|
|
aws_region_name=self.s3_region_name
|
|
)
|
|
SigV4Auth(credentials, "s3", aws_region_name).add_auth(aws_request)
|
|
|
|
# Prepare the signed headers
|
|
signed_headers = dict(aws_request.headers.items())
|
|
|
|
# Use prepared URL so path segments match SigV4 canonical request (e.g. %20 for spaces).
|
|
request_url = prepped.url or url
|
|
|
|
# Make the request with retry for transient S3 errors (500/503)
|
|
max_retries = 3
|
|
for attempt in range(max_retries):
|
|
response = await self.async_httpx_client.put(
|
|
request_url, data=json_string, headers=signed_headers
|
|
)
|
|
if response.status_code in (500, 503) and attempt < max_retries - 1:
|
|
wait_time = 2**attempt # 1s, 2s
|
|
verbose_logger.warning(
|
|
f"S3 upload returned {response.status_code}, retrying in {wait_time}s "
|
|
f"(attempt {attempt + 1}/{max_retries}) "
|
|
f"key={batch_logging_element.s3_object_key}"
|
|
)
|
|
await asyncio.sleep(wait_time)
|
|
continue
|
|
response.raise_for_status()
|
|
break
|
|
except Exception as e:
|
|
verbose_logger.exception(f"Error uploading to s3: {str(e)}")
|
|
self.handle_callback_failure(callback_name="S3Logger")
|
|
|
|
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
|
|
|
|
if self.s3_strip_base64_files:
|
|
standard_logging_payload = self._strip_base64_from_messages_sync(
|
|
standard_logging_payload
|
|
)
|
|
|
|
# Base prefix (default empty)
|
|
prefix_components = []
|
|
if self.s3_use_team_prefix:
|
|
team_alias = standard_logging_payload.get("metadata", {}).get(
|
|
"user_api_key_team_alias", None
|
|
)
|
|
if team_alias:
|
|
prefix_components.append(team_alias)
|
|
if self.s3_use_key_prefix:
|
|
user_api_key_alias = standard_logging_payload.get("metadata", {}).get(
|
|
"user_api_key_alias", None
|
|
)
|
|
if user_api_key_alias:
|
|
prefix_components.append(user_api_key_alias)
|
|
|
|
# Construct full prefix path
|
|
prefix_path = "/".join(prefix_components)
|
|
if prefix_path:
|
|
prefix_path += "/"
|
|
|
|
s3_file_name = (
|
|
litellm.utils.get_logging_id(start_time, standard_logging_payload) or ""
|
|
)
|
|
verbose_logger.debug(
|
|
f"Creating s3 file with prefix_components={prefix_components},prefix_path={prefix_path} and {s3_file_name}"
|
|
)
|
|
s3_object_key = get_s3_object_key(
|
|
s3_path=cast(Optional[str], self.s3_path) or "",
|
|
prefix=prefix_path,
|
|
start_time=start_time,
|
|
s3_file_name=s3_file_name,
|
|
)
|
|
verbose_logger.debug(f"s3_object_key={s3_object_key}")
|
|
|
|
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 and self.s3_bucket_name:
|
|
if self.s3_use_virtual_hosted_style:
|
|
# Virtual-hosted-style: bucket.endpoint/key
|
|
endpoint_host = self.s3_endpoint_url.replace(
|
|
"https://", ""
|
|
).replace("http://", "")
|
|
protocol = (
|
|
"https://"
|
|
if self.s3_endpoint_url.startswith("https://")
|
|
else "http://"
|
|
)
|
|
url = f"{protocol}{self.s3_bucket_name}.{endpoint_host}/{batch_logging_element.s3_object_key}"
|
|
else:
|
|
# Path-style: endpoint/bucket/key
|
|
url = (
|
|
self.s3_endpoint_url
|
|
+ "/"
|
|
+ self.s3_bucket_name
|
|
+ "/"
|
|
+ batch_logging_element.s3_object_key
|
|
)
|
|
|
|
# Convert JSON to string
|
|
json_string = safe_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,
|
|
)
|
|
aws_region_name = self.get_aws_region_name_for_non_llm_api_calls(
|
|
aws_region_name=self.s3_region_name
|
|
)
|
|
SigV4Auth(credentials, "s3", aws_region_name).add_auth(aws_request)
|
|
|
|
# Prepare the signed headers
|
|
signed_headers = dict(aws_request.headers.items())
|
|
|
|
# Use prepared URL so path segments match SigV4 canonical request (e.g. %20 for spaces).
|
|
request_url = prepped.url or url
|
|
|
|
httpx_client = _get_httpx_client(
|
|
params=(
|
|
{"ssl_verify": self.s3_verify}
|
|
if self.s3_verify is not None
|
|
else None
|
|
)
|
|
)
|
|
# Make the request with retry for transient S3 errors (500/503)
|
|
max_retries = 3
|
|
for attempt in range(max_retries):
|
|
response = httpx_client.put(
|
|
request_url, data=json_string, headers=signed_headers
|
|
)
|
|
if response.status_code in (500, 503) and attempt < max_retries - 1:
|
|
wait_time = 2**attempt # 1s, 2s
|
|
verbose_logger.warning(
|
|
f"S3 upload returned {response.status_code}, retrying in {wait_time}s "
|
|
f"(attempt {attempt + 1}/{max_retries}) "
|
|
f"key={batch_logging_element.s3_object_key}"
|
|
)
|
|
time.sleep(wait_time)
|
|
continue
|
|
response.raise_for_status()
|
|
break
|
|
except Exception as e:
|
|
verbose_logger.exception(f"Error uploading to s3: {str(e)}")
|
|
self.handle_callback_failure(callback_name="S3Logger")
|
|
|
|
async def _download_object_from_s3(self, s3_object_key: str) -> Optional[dict]:
|
|
"""
|
|
Download and parse JSON object from S3.
|
|
|
|
Args:
|
|
s3_object_key: The S3 object key to download
|
|
|
|
Returns:
|
|
Optional[dict]: The parsed JSON object or None if not found/error
|
|
"""
|
|
try:
|
|
import hashlib
|
|
|
|
import requests
|
|
from botocore.auth import SigV4Auth
|
|
from botocore.awsrequest import AWSRequest
|
|
except ImportError:
|
|
raise ImportError("Missing boto3 to call S3. Run 'pip install boto3'.")
|
|
|
|
try:
|
|
from litellm.litellm_core_utils.asyncify import asyncify
|
|
|
|
# Get AWS credentials
|
|
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 - downloading data from s3 - {s3_object_key}"
|
|
)
|
|
|
|
# Prepare the URL
|
|
url = f"https://{self.s3_bucket_name}.s3.{self.s3_region_name}.amazonaws.com/{s3_object_key}"
|
|
|
|
if self.s3_endpoint_url and self.s3_bucket_name:
|
|
if self.s3_use_virtual_hosted_style:
|
|
# Virtual-hosted-style: bucket.endpoint/key
|
|
endpoint_host = self.s3_endpoint_url.replace(
|
|
"https://", ""
|
|
).replace("http://", "")
|
|
protocol = (
|
|
"https://"
|
|
if self.s3_endpoint_url.startswith("https://")
|
|
else "http://"
|
|
)
|
|
url = f"{protocol}{self.s3_bucket_name}.{endpoint_host}/{s3_object_key}"
|
|
else:
|
|
# Path-style: endpoint/bucket/key
|
|
url = (
|
|
self.s3_endpoint_url
|
|
+ "/"
|
|
+ self.s3_bucket_name
|
|
+ "/"
|
|
+ s3_object_key
|
|
)
|
|
|
|
# Prepare the request for GET operation
|
|
# For GET requests, we need x-amz-content-sha256 with hash of empty string
|
|
empty_string_hash = hashlib.sha256(b"").hexdigest()
|
|
headers = {
|
|
"x-amz-content-sha256": empty_string_hash,
|
|
}
|
|
req = requests.Request("GET", url, headers=headers)
|
|
prepped = req.prepare()
|
|
|
|
# Sign the request
|
|
aws_request = AWSRequest(
|
|
method=prepped.method,
|
|
url=prepped.url,
|
|
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())
|
|
|
|
request_url = prepped.url or url
|
|
response = await self.async_httpx_client.get(
|
|
request_url, headers=signed_headers
|
|
)
|
|
|
|
if response.status_code != 200:
|
|
verbose_logger.exception(
|
|
"S3 object not found, saw response=", response.text
|
|
)
|
|
return None
|
|
|
|
# Parse JSON response
|
|
return response.json()
|
|
|
|
except Exception as e:
|
|
verbose_logger.exception(f"Error downloading from S3: {str(e)}")
|
|
return None
|
|
|
|
async def get_proxy_server_request_from_cold_storage_with_object_key(
|
|
self,
|
|
object_key: str,
|
|
) -> Optional[dict]:
|
|
"""
|
|
Get the proxy server request from cold storage
|
|
|
|
Allows fetching a dict of the proxy server request from s3 or GCS bucket.
|
|
|
|
Args:
|
|
request_id: The unique request ID to search for
|
|
start_time: The start time of the request (datetime or ISO string)
|
|
|
|
Returns:
|
|
Optional[dict]: The request data dictionary or None if not found
|
|
"""
|
|
try:
|
|
# Download and return the object from S3
|
|
downloaded_object = await self._download_object_from_s3(object_key)
|
|
return downloaded_object
|
|
except Exception as e:
|
|
verbose_logger.exception(
|
|
f"Error retrieving object {object_key} from cold storage: {str(e)}"
|
|
)
|
|
return None
|