mirror of
https://github.com/tiennm99/litellm.git
synced 2026-08-14 16:25:29 +00:00
Merge pull request #23167 from BerriAI/litellm_audit_log_s3_export
[Feature] Add Audit Log Export to External Callbacks
This commit is contained in:
ryan-crabbe
@@ -97,6 +97,7 @@ input_callback: List[CALLBACK_TYPES] = []
|
||||
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",
|
||||
|
||||
@@ -27,6 +27,7 @@ from litellm.types.utils import (
|
||||
LLMResponseTypes,
|
||||
ModelResponse,
|
||||
ModelResponseStream,
|
||||
StandardAuditLogPayload,
|
||||
StandardCallbackDynamicParams,
|
||||
StandardLoggingPayload,
|
||||
)
|
||||
@@ -177,6 +178,10 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac
|
||||
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(
|
||||
|
||||
@@ -22,7 +22,7 @@ from litellm.llms.custom_httpx.http_handler import (
|
||||
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 @@ class S3Logger(CustomBatchLogger, BaseAWSLLM):
|
||||
)
|
||||
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(
|
||||
|
||||
@@ -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 @@ from litellm.proxy._types import (
|
||||
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
|
||||
|
||||
@@ -2961,6 +2961,29 @@ class ProxyConfig:
|
||||
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
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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
|
||||
Reference in New Issue
Block a user