diff --git a/litellm/__init__.py b/litellm/__init__.py index 7f72e0b0e89..e45d926e8db 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -97,6 +97,7 @@ success_callback: List[CALLBACK_TYPES] = [] failure_callback: List[CALLBACK_TYPES] = [] service_callback: List[CALLBACK_TYPES] = [] +audit_log_callbacks: List[CALLBACK_TYPES] = [] # logging_callback_manager is lazy-loaded via __getattr__ _custom_logger_compatible_callbacks_literal = Literal[ "lago", diff --git a/litellm/integrations/custom_logger.py b/litellm/integrations/custom_logger.py index 06ba9675ca2..cccabf53e51 100644 --- a/litellm/integrations/custom_logger.py +++ b/litellm/integrations/custom_logger.py @@ -27,6 +27,7 @@ LLMResponseTypes, ModelResponse, ModelResponseStream, + StandardAuditLogPayload, StandardCallbackDynamicParams, StandardLoggingPayload, ) @@ -177,6 +178,10 @@ async def async_log_success_event(self, kwargs, response_obj, start_time, end_ti async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time): pass + async def async_log_audit_log_event(self, audit_log: "StandardAuditLogPayload"): + """Called when an audit log is created. Override in subclasses to handle.""" + pass + #### PROMPT MANAGEMENT HOOKS #### async def async_get_chat_completion_prompt( diff --git a/litellm/integrations/s3_v2.py b/litellm/integrations/s3_v2.py index c8db4be7cea..405bf9698cc 100644 --- a/litellm/integrations/s3_v2.py +++ b/litellm/integrations/s3_v2.py @@ -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" + ) + + 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( diff --git a/litellm/proxy/management_helpers/audit_logs.py b/litellm/proxy/management_helpers/audit_logs.py index ea082f468ae..08d92c3e6a6 100644 --- a/litellm/proxy/management_helpers/audit_logs.py +++ b/litellm/proxy/management_helpers/audit_logs.py @@ -2,12 +2,15 @@ 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, @@ -15,6 +18,95 @@ 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 + ) 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 diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index e982c934aa6..61fb12a855d 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -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) + + _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." + ) elif key == "cache_params": # this is set in the cache branch # see usage here: https://docs.litellm.ai/docs/proxy/caching diff --git a/litellm/types/utils.py b/litellm/types/utils.py index 38425c7ac4a..e7f0cd77143 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -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" + 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) diff --git a/tests/test_litellm/proxy/management_helpers/test_audit_log_callbacks.py b/tests/test_litellm/proxy/management_helpers/test_audit_log_callbacks.py new file mode 100644 index 00000000000..85eda4368cd --- /dev/null +++ b/tests/test_litellm/proxy/management_helpers/test_audit_log_callbacks.py @@ -0,0 +1,345 @@ +""" +Tests for audit log callback dispatch. + +Tests the flow: create_audit_log_for_update -> _dispatch_audit_log_to_callbacks -> CustomLogger.async_log_audit_log_event +""" + +import asyncio +import json +from datetime import datetime, timezone +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest + +import litellm +from litellm.integrations.custom_logger import CustomLogger +from litellm.proxy._types import LiteLLM_AuditLogs, LitellmTableNames +from litellm.proxy.management_helpers.audit_logs import ( + _audit_log_task_done_callback, + _build_audit_log_payload, + _dispatch_audit_log_to_callbacks, + create_audit_log_for_update, +) +from litellm.types.utils import StandardAuditLogPayload + + +@pytest.fixture(autouse=True) +def reset_audit_log_callbacks(): + """Reset audit_log_callbacks before and after each test.""" + original = litellm.audit_log_callbacks + litellm.audit_log_callbacks = [] + yield + litellm.audit_log_callbacks = original + + +def _make_audit_log( + action: str = "created", + table_name: LitellmTableNames = LitellmTableNames.TEAM_TABLE_NAME, +) -> LiteLLM_AuditLogs: + return LiteLLM_AuditLogs( + id="test-audit-id", + updated_at=datetime(2026, 3, 9, 12, 0, 0, tzinfo=timezone.utc), + changed_by="user-123", + changed_by_api_key="sk-abc", + action=action, + table_name=table_name, + object_id="team-456", + updated_values=json.dumps({"name": "new-team"}), + before_value=json.dumps({"name": "old-team"}), + ) + + +class TestBuildAuditLogPayload: + def test_builds_correct_payload(self): + audit_log = _make_audit_log() + payload = _build_audit_log_payload(audit_log) + + assert payload["id"] == "test-audit-id" + assert payload["updated_at"] == "2026-03-09T12:00:00+00:00" + assert payload["changed_by"] == "user-123" + assert payload["changed_by_api_key"] == "sk-abc" + assert payload["action"] == "created" + assert payload["table_name"] == "LiteLLM_TeamTable" + assert payload["object_id"] == "team-456" + assert payload["updated_values"] == json.dumps({"name": "new-team"}) + assert payload["before_value"] == json.dumps({"name": "old-team"}) + + def test_handles_none_values(self): + audit_log = LiteLLM_AuditLogs( + id="test-id", + updated_at=datetime(2026, 1, 1, tzinfo=timezone.utc), + changed_by=None, + changed_by_api_key=None, + action="deleted", + table_name=LitellmTableNames.KEY_TABLE_NAME, + object_id="key-789", + updated_values=None, + before_value=None, + ) + payload = _build_audit_log_payload(audit_log) + + assert payload["changed_by"] == "" + assert payload["changed_by_api_key"] == "" + assert payload["before_value"] is None + assert payload["updated_values"] is None + + +class TestDispatchAuditLogToCallbacks: + @pytest.mark.asyncio + async def test_dispatches_to_custom_logger_instance(self): + mock_logger = MagicMock(spec=CustomLogger) + mock_logger.async_log_audit_log_event = AsyncMock() + litellm.audit_log_callbacks = [mock_logger] + + audit_log = _make_audit_log() + await _dispatch_audit_log_to_callbacks(audit_log) + + # Let asyncio.create_task run + await asyncio.sleep(0.1) + + mock_logger.async_log_audit_log_event.assert_called_once() + payload = mock_logger.async_log_audit_log_event.call_args[0][0] + assert payload["id"] == "test-audit-id" + assert payload["action"] == "created" + + @pytest.mark.asyncio + async def test_no_dispatch_when_callbacks_empty(self): + litellm.audit_log_callbacks = [] + audit_log = _make_audit_log() + # Should return immediately without error + await _dispatch_audit_log_to_callbacks(audit_log) + + @pytest.mark.asyncio + async def test_resolves_string_callback(self): + mock_logger = MagicMock(spec=CustomLogger) + mock_logger.async_log_audit_log_event = AsyncMock() + + litellm.audit_log_callbacks = ["s3_v2"] + + with patch( + "litellm.proxy.management_helpers.audit_logs._resolve_audit_log_callback", + return_value=mock_logger, + ): + audit_log = _make_audit_log() + await _dispatch_audit_log_to_callbacks(audit_log) + await asyncio.sleep(0.1) + + mock_logger.async_log_audit_log_event.assert_called_once() + + @pytest.mark.asyncio + async def test_nonblocking_on_callback_failure(self): + """Callback errors should not propagate.""" + mock_logger = MagicMock(spec=CustomLogger) + mock_logger.async_log_audit_log_event = AsyncMock( + side_effect=RuntimeError("boom") + ) + litellm.audit_log_callbacks = [mock_logger] + + audit_log = _make_audit_log() + # Should not raise + await _dispatch_audit_log_to_callbacks(audit_log) + await asyncio.sleep(0.1) + + @pytest.mark.asyncio + async def test_skips_unresolvable_string_callback(self): + litellm.audit_log_callbacks = ["nonexistent_callback"] + + with patch( + "litellm.proxy.management_helpers.audit_logs._resolve_audit_log_callback", + return_value=None, + ): + audit_log = _make_audit_log() + # Should not raise + await _dispatch_audit_log_to_callbacks(audit_log) + + +class TestCreateAuditLogForUpdateWithCallbacks: + @pytest.mark.asyncio + async def test_dispatches_to_callbacks_after_db_write(self): + mock_logger = MagicMock(spec=CustomLogger) + mock_logger.async_log_audit_log_event = AsyncMock() + litellm.audit_log_callbacks = [mock_logger] + + with patch("litellm.proxy.proxy_server.premium_user", True), patch( + "litellm.store_audit_logs", True + ), patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma: + mock_prisma.db.litellm_auditlog.create = AsyncMock() + + audit_log = _make_audit_log() + await create_audit_log_for_update(audit_log) + await asyncio.sleep(0.1) + + # DB write should happen + mock_prisma.db.litellm_auditlog.create.assert_called_once() + # Callback should also be called + mock_logger.async_log_audit_log_event.assert_called_once() + + @pytest.mark.asyncio + async def test_no_dispatch_when_not_premium(self): + mock_logger = MagicMock(spec=CustomLogger) + mock_logger.async_log_audit_log_event = AsyncMock() + litellm.audit_log_callbacks = [mock_logger] + + with patch("litellm.proxy.proxy_server.premium_user", False), patch( + "litellm.store_audit_logs", True + ): + audit_log = _make_audit_log() + await create_audit_log_for_update(audit_log) + await asyncio.sleep(0.1) + + mock_logger.async_log_audit_log_event.assert_not_called() + + @pytest.mark.asyncio + async def test_no_dispatch_when_store_audit_logs_false(self): + mock_logger = MagicMock(spec=CustomLogger) + mock_logger.async_log_audit_log_event = AsyncMock() + litellm.audit_log_callbacks = [mock_logger] + + with patch("litellm.store_audit_logs", False): + audit_log = _make_audit_log() + await create_audit_log_for_update(audit_log) + await asyncio.sleep(0.1) + + mock_logger.async_log_audit_log_event.assert_not_called() + + @pytest.mark.asyncio + async def test_dispatches_even_when_prisma_client_is_none(self): + """Callbacks should fire even if DB is unavailable.""" + mock_logger = MagicMock(spec=CustomLogger) + mock_logger.async_log_audit_log_event = AsyncMock() + litellm.audit_log_callbacks = [mock_logger] + + with patch("litellm.proxy.proxy_server.premium_user", True), patch( + "litellm.store_audit_logs", True + ), patch("litellm.proxy.proxy_server.prisma_client", None): + audit_log = _make_audit_log() + await create_audit_log_for_update(audit_log) + await asyncio.sleep(0.1) + + # Callback should still be called despite no DB + mock_logger.async_log_audit_log_event.assert_called_once() + + @pytest.mark.asyncio + async def test_dispatches_even_when_db_write_fails(self): + """Callbacks should fire even if the DB write raises.""" + mock_logger = MagicMock(spec=CustomLogger) + mock_logger.async_log_audit_log_event = AsyncMock() + litellm.audit_log_callbacks = [mock_logger] + + with patch("litellm.proxy.proxy_server.premium_user", True), patch( + "litellm.store_audit_logs", True + ), patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma: + mock_prisma.db.litellm_auditlog.create = AsyncMock( + side_effect=RuntimeError("DB connection lost") + ) + + audit_log = _make_audit_log() + await create_audit_log_for_update(audit_log) + await asyncio.sleep(0.1) + + # Callback should still be called despite DB failure + mock_logger.async_log_audit_log_event.assert_called_once() + + +class TestAuditLogTaskDoneCallback: + def test_logs_exception_from_failed_task(self): + """Done callback should log task exceptions.""" + mock_task = MagicMock(spec=asyncio.Task) + mock_task.exception.return_value = RuntimeError("callback failed") + + with patch( + "litellm.proxy.management_helpers.audit_logs.verbose_proxy_logger" + ) as mock_logger: + _audit_log_task_done_callback(mock_task) + mock_logger.error.assert_called_once() + assert "callback failed" in str(mock_logger.error.call_args) + + def test_no_log_on_success(self): + """Done callback should not log when task succeeds.""" + mock_task = MagicMock(spec=asyncio.Task) + mock_task.exception.return_value = None + + with patch( + "litellm.proxy.management_helpers.audit_logs.verbose_proxy_logger" + ) as mock_logger: + _audit_log_task_done_callback(mock_task) + mock_logger.error.assert_not_called() + + def test_handles_cancelled_task(self): + """Done callback should handle cancelled tasks gracefully.""" + mock_task = MagicMock(spec=asyncio.Task) + mock_task.exception.side_effect = asyncio.CancelledError() + + with patch( + "litellm.proxy.management_helpers.audit_logs.verbose_proxy_logger" + ) as mock_logger: + _audit_log_task_done_callback(mock_task) + mock_logger.error.assert_not_called() + + +class TestS3LoggerAuditLogEvent: + @pytest.mark.asyncio + async def test_queues_audit_log_with_correct_s3_key(self): + with patch( + "litellm.integrations.s3_v2.S3Logger.__init__", return_value=None + ): + from litellm.integrations.s3_v2 import S3Logger + + logger = S3Logger() + logger.s3_path = "my-prefix" + logger.log_queue = [] + logger.batch_size = 100 + + audit_log = StandardAuditLogPayload( + id="audit-123", + updated_at="2026-03-09T12:00:00+00:00", + changed_by="user-1", + changed_by_api_key="sk-abc", + action="created", + table_name="LiteLLM_TeamTable", + object_id="team-1", + before_value=None, + updated_values='{"name": "new"}', + ) + + await logger.async_log_audit_log_event(audit_log) + + assert len(logger.log_queue) == 1 + element = logger.log_queue[0] + assert element.s3_object_key.startswith("my-prefix/audit_logs/") + assert "audit-123" in element.s3_object_key + assert element.s3_object_key.endswith(".json") + assert element.s3_object_download_filename == "audit-audit-123.json" + assert element.payload["id"] == "audit-123" + assert element.payload["action"] == "created" + + @pytest.mark.asyncio + async def test_s3_key_format_no_path(self): + with patch( + "litellm.integrations.s3_v2.S3Logger.__init__", return_value=None + ): + from litellm.integrations.s3_v2 import S3Logger + + logger = S3Logger() + logger.s3_path = None + logger.log_queue = [] + logger.batch_size = 100 + + audit_log = StandardAuditLogPayload( + id="audit-456", + updated_at="2026-03-09T12:00:00+00:00", + changed_by="user-1", + changed_by_api_key="sk-abc", + action="deleted", + table_name="LiteLLM_VerificationToken", + object_id="key-1", + before_value=None, + updated_values=None, + ) + + await logger.async_log_audit_log_event(audit_log) + + assert len(logger.log_queue) == 1 + element = logger.log_queue[0] + assert element.s3_object_key.startswith("audit_logs/") + assert "audit-456" in element.s3_object_key