-
-
Notifications
You must be signed in to change notification settings - Fork 9.7k
[Feature] Add Audit Log Export to External Callbacks #23167
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Changes from all commits
c1b3640
364f5c6
1da02b6
9c47d50
ae21527
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -22,7 +22,7 @@ | |
| httpxSpecialProvider, | ||
| ) | ||
| from litellm.types.integrations.s3_v2 import s3BatchLoggingElement | ||
| from litellm.types.utils import StandardLoggingPayload | ||
| from litellm.types.utils import StandardAuditLogPayload, StandardLoggingPayload | ||
|
|
||
| from .custom_batch_logger import CustomBatchLogger | ||
|
|
||
|
|
@@ -248,6 +248,38 @@ async def async_log_failure_event(self, kwargs, response_obj, start_time, end_ti | |
| ) | ||
| 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" | ||
| ) | ||
|
Comment on lines
+264
to
+268
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more.
The actual event timestamp is already available in from datetime import datetime, timezone
event_time_str = audit_log.get("updated_at", "")
try:
event_time = datetime.fromisoformat(event_time_str)
except (ValueError, TypeError):
event_time = datetime.now(timezone.utc)
s3_object_key = (
f"{s3_path}audit_logs/"
f"{event_time.strftime('%Y-%m-%d')}/"
f"{event_time.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( | ||
|
|
||
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -2,19 +2,111 @@ | |
| Functions to create audit logs for LiteLLM Proxy | ||
| """ | ||
|
|
||
| import asyncio | ||
| import json | ||
| from litellm._uuid import uuid | ||
| from datetime import datetime, timezone | ||
| from typing import Dict | ||
|
|
||
| import litellm | ||
| from litellm._logging import verbose_proxy_logger | ||
| from litellm._uuid import uuid | ||
| from litellm.integrations.custom_logger import CustomLogger | ||
| from litellm.proxy._types import ( | ||
| AUDIT_ACTIONS, | ||
| LiteLLM_AuditLogs, | ||
| LitellmTableNames, | ||
| Optional, | ||
| UserAPIKeyAuth, | ||
| ) | ||
| from litellm.types.utils import StandardAuditLogPayload | ||
|
|
||
| _audit_log_callback_cache: Dict[str, CustomLogger] = {} | ||
|
|
||
|
|
||
| def _resolve_audit_log_callback(name: str) -> Optional[CustomLogger]: | ||
| """Resolve a string callback name to a CustomLogger instance, with caching.""" | ||
| if name in _audit_log_callback_cache: | ||
| return _audit_log_callback_cache[name] | ||
|
|
||
| from litellm.litellm_core_utils.litellm_logging import ( | ||
| _init_custom_logger_compatible_class, | ||
| ) | ||
|
|
||
| instance = _init_custom_logger_compatible_class( | ||
| logging_integration=name, # type: ignore | ||
| internal_usage_cache=None, | ||
| llm_router=None, | ||
| ) | ||
|
|
||
| if instance is not None: | ||
| _audit_log_callback_cache[name] = instance | ||
| return instance | ||
|
|
||
|
|
||
| def _build_audit_log_payload( | ||
| request_data: LiteLLM_AuditLogs, | ||
| ) -> StandardAuditLogPayload: | ||
| """Convert LiteLLM_AuditLogs to StandardAuditLogPayload for callback dispatch.""" | ||
| updated_at = "" | ||
| if request_data.updated_at is not None: | ||
| updated_at = request_data.updated_at.isoformat() | ||
|
|
||
| table_name = request_data.table_name | ||
| if isinstance(table_name, LitellmTableNames): | ||
| table_name = table_name.value | ||
|
|
||
| return StandardAuditLogPayload( | ||
| id=request_data.id, | ||
| updated_at=updated_at, | ||
| changed_by=request_data.changed_by or "", | ||
| changed_by_api_key=request_data.changed_by_api_key or "", | ||
| action=request_data.action, | ||
| table_name=str(table_name), | ||
| object_id=request_data.object_id, | ||
| before_value=request_data.before_value, | ||
| updated_values=request_data.updated_values, | ||
| ) | ||
|
|
||
|
|
||
| def _audit_log_task_done_callback(task: asyncio.Task) -> None: | ||
| """Log exceptions from audit log callback tasks so they don't slip through silently.""" | ||
| try: | ||
| exc = task.exception() | ||
| except asyncio.CancelledError: | ||
| return | ||
| if exc is not None: | ||
| verbose_proxy_logger.error( | ||
| "Audit log callback task failed: %s", exc, exc_info=exc | ||
| ) | ||
|
|
||
|
|
||
| async def _dispatch_audit_log_to_callbacks( | ||
| request_data: LiteLLM_AuditLogs, | ||
| ) -> None: | ||
| """Dispatch audit log to all registered audit_log_callbacks.""" | ||
| if not litellm.audit_log_callbacks: | ||
| return | ||
|
|
||
| payload = _build_audit_log_payload(request_data) | ||
|
|
||
| for callback in litellm.audit_log_callbacks: | ||
| try: | ||
| resolved = callback | ||
| if isinstance(callback, str): | ||
| resolved = _resolve_audit_log_callback(callback) | ||
| if resolved is None: | ||
| verbose_proxy_logger.warning( | ||
| "Could not resolve audit log callback: %s", callback | ||
| ) | ||
| continue | ||
|
|
||
| if isinstance(resolved, CustomLogger): | ||
| task = asyncio.create_task(resolved.async_log_audit_log_event(payload)) | ||
| task.add_done_callback(_audit_log_task_done_callback) | ||
| except Exception as e: | ||
| verbose_proxy_logger.error( | ||
| "Failed dispatching audit log to callback: %s", e | ||
| ) | ||
|
ryan-crabbe marked this conversation as resolved.
Comment on lines
+92
to
+109
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more.
If Consider adding an if isinstance(resolved, CustomLogger):
task = asyncio.create_task(
resolved.async_log_audit_log_event(payload)
)
task.add_done_callback(_audit_log_task_done_callback)
else:
verbose_proxy_logger.warning(
"Audit log callback %r is not a CustomLogger instance and will be skipped.",
resolved,
) |
||
|
|
||
|
|
||
| async def create_object_audit_log( | ||
|
|
@@ -81,9 +173,6 @@ async def create_audit_log_for_update(request_data: LiteLLM_AuditLogs): | |
| if premium_user is not True: | ||
| return | ||
|
|
||
| if prisma_client is None: | ||
| raise Exception("prisma_client is None, no DB connected") | ||
|
|
||
| verbose_proxy_logger.debug("creating audit log for %s", request_data) | ||
|
|
||
| if isinstance(request_data.updated_values, dict): | ||
|
|
@@ -92,6 +181,15 @@ async def create_audit_log_for_update(request_data: LiteLLM_AuditLogs): | |
| if isinstance(request_data.before_value, dict): | ||
| request_data.before_value = json.dumps(request_data.before_value) | ||
|
|
||
| # Dispatch to external audit log callbacks regardless of DB availability | ||
| await _dispatch_audit_log_to_callbacks(request_data) | ||
|
|
||
| if prisma_client is None: | ||
| verbose_proxy_logger.error( | ||
| "prisma_client is None, cannot write audit log to DB" | ||
| ) | ||
| return | ||
|
|
||
| _request_data = request_data.model_dump(exclude_none=True) | ||
|
|
||
| try: | ||
|
|
@@ -103,5 +201,3 @@ async def create_audit_log_for_update(request_data: LiteLLM_AuditLogs): | |
| except Exception as e: | ||
| # [Non-Blocking Exception. Do not allow blocking LLM API call] | ||
| verbose_proxy_logger.error(f"Failed Creating audit log {e}") | ||
|
|
||
| return | ||
| Original file line number | Diff line number | Diff line change | ||||||||||||||||||||||||||||||||||||||||
|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|
|
|
@@ -2961,6 +2961,29 @@ async def load_config( # noqa: PLR0915 | |||||||||||||||||||||||||||||||||||||||||
| print( # noqa | ||||||||||||||||||||||||||||||||||||||||||
| f"{blue_color_code} Initialized Failure Callbacks - {litellm.failure_callback} {reset_color_code}" | ||||||||||||||||||||||||||||||||||||||||||
| ) # noqa | ||||||||||||||||||||||||||||||||||||||||||
| elif key == "audit_log_callbacks": | ||||||||||||||||||||||||||||||||||||||||||
| litellm.audit_log_callbacks = [] | ||||||||||||||||||||||||||||||||||||||||||
|
|
||||||||||||||||||||||||||||||||||||||||||
| for callback in value: | ||||||||||||||||||||||||||||||||||||||||||
| if "." in callback: | ||||||||||||||||||||||||||||||||||||||||||
| litellm.audit_log_callbacks.append( | ||||||||||||||||||||||||||||||||||||||||||
| get_instance_fn(value=callback) | ||||||||||||||||||||||||||||||||||||||||||
| ) | ||||||||||||||||||||||||||||||||||||||||||
| else: | ||||||||||||||||||||||||||||||||||||||||||
| litellm.audit_log_callbacks.append(callback) | ||||||||||||||||||||||||||||||||||||||||||
|
Comment on lines
+2967
to
+2973
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more.
Every item in The existing
Suggested change
|
||||||||||||||||||||||||||||||||||||||||||
|
|
||||||||||||||||||||||||||||||||||||||||||
| _store_audit_logs = litellm_settings.get( | ||||||||||||||||||||||||||||||||||||||||||
| "store_audit_logs", litellm.store_audit_logs | ||||||||||||||||||||||||||||||||||||||||||
| ) | ||||||||||||||||||||||||||||||||||||||||||
| if _store_audit_logs: | ||||||||||||||||||||||||||||||||||||||||||
| print( # noqa | ||||||||||||||||||||||||||||||||||||||||||
| f"{blue_color_code} Initialized Audit Log Callbacks - {litellm.audit_log_callbacks} {reset_color_code}" | ||||||||||||||||||||||||||||||||||||||||||
| ) # noqa | ||||||||||||||||||||||||||||||||||||||||||
| else: | ||||||||||||||||||||||||||||||||||||||||||
| verbose_proxy_logger.warning( | ||||||||||||||||||||||||||||||||||||||||||
| "'audit_log_callbacks' is configured but 'store_audit_logs' is not enabled. " | ||||||||||||||||||||||||||||||||||||||||||
| "Audit log callbacks will not fire until 'store_audit_logs: true' is added to litellm_settings." | ||||||||||||||||||||||||||||||||||||||||||
| ) | ||||||||||||||||||||||||||||||||||||||||||
|
Comment on lines
+2982
to
+2986
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more.
The warning check reads Additionally, the warning doesn't surface the other gate: audit log callbacks are also blocked for non-premium users ( Consider updating the warning to cover both conditions: _store_audit_logs = litellm_settings.get(
"store_audit_logs", litellm.store_audit_logs
) or get_secret_bool("LITELLM_STORE_AUDIT_LOGS")
if not _store_audit_logs:
verbose_proxy_logger.warning(
"'audit_log_callbacks' is configured but 'store_audit_logs' is not enabled "
"(set 'store_audit_logs: true' in litellm_settings or LITELLM_STORE_AUDIT_LOGS=true). "
"Additionally, a premium license is required for audit log callbacks to fire."
) |
||||||||||||||||||||||||||||||||||||||||||
| elif key == "cache_params": | ||||||||||||||||||||||||||||||||||||||||||
| # this is set in the cache branch | ||||||||||||||||||||||||||||||||||||||||||
| # see usage here: https://docs.litellm.ai/docs/proxy/caching | ||||||||||||||||||||||||||||||||||||||||||
|
|
||||||||||||||||||||||||||||||||||||||||||
| Original file line number | Diff line number | Diff line change | ||||||||
|---|---|---|---|---|---|---|---|---|---|---|
|
|
@@ -2804,6 +2804,20 @@ class StandardLoggingPayloadStatusFields(TypedDict, total=False): | |||||||||
| """ | ||||||||||
|
|
||||||||||
|
|
||||||||||
| class StandardAuditLogPayload(TypedDict): | ||||||||||
| """Payload for audit log events dispatched to external callbacks.""" | ||||||||||
|
|
||||||||||
| id: str | ||||||||||
| updated_at: str # ISO-8601 | ||||||||||
| changed_by: str | ||||||||||
| changed_by_api_key: str | ||||||||||
| action: str # "created" | "updated" | "deleted" | "blocked" | "rotated" | ||||||||||
|
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more.
Suggested change
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more.
from typing import TYPE_CHECKING
if TYPE_CHECKING:
from litellm.proxy._types import AUDIT_ACTIONSand then type the field accordingly, or simply copy the Literal definition into
Suggested change
|
||||||||||
| table_name: str | ||||||||||
| object_id: str | ||||||||||
| before_value: Optional[str] | ||||||||||
| updated_values: Optional[str] | ||||||||||
|
|
||||||||||
|
|
||||||||||
| class StandardLoggingPayload(TypedDict): | ||||||||||
| id: str | ||||||||||
| trace_id: str # Trace multiple LLM calls belonging to same overall request (e.g. fallbacks/retries) | ||||||||||
|
|
||||||||||
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
timezonedatetimeis already imported at the module level (from datetime import datetime). Thefrom datetime import timezoneinside the function body is redundant —timezonecan be accessed directly since it lives in the samedatetimemodule. Move this to the module-level import:and remove the in-function import entirely.