feat(azure_storage): upload audit logs to Azure Storage via audit_log_callbacks

litellm/integrations/azure_storage/azure_storage.py

@@ -2,7 +2,7 @@
 import os
 import time
 from collections.abc import Callable
-from datetime import datetime, timedelta
+from datetime import datetime, timedelta, timezone
 from functools import cache
 from typing import Final
 from urllib.parse import unquote
@@ -28,7 +28,7 @@
 from litellm.types.secret_managers.get_azure_ad_token_provider import (
     AzureCredentialType,
 )
-from litellm.types.utils import StandardLoggingPayload
+from litellm.types.utils import StandardAuditLogPayload, StandardLoggingPayload
 
 AZURE_STORAGE_TOKEN_SCOPE: Final = "https://storage.azure.com/.default"
 _ADLS_SAFE_NAME: Final = str.maketrans("/", "_", "=")
@@ -160,6 +160,24 @@ async def async_log_failure_event(self, kwargs, response_obj, start_time, end_ti
         except Exception as e:
             verbose_logger.exception("AzureBlobStorageLogger Layer Error - %s", e)
 
+    async def async_log_audit_log_event(self, audit_log: StandardAuditLogPayload) -> None:
+        now: Final = datetime.now(timezone.utc)
+        file_path: Final = f"audit_logs/{now.strftime('%Y-%m-%d')}/{now.strftime('%H-%M-%S')}_{audit_log['id']}.json"
+        await self._upload_json_to_file_path(file_path=file_path, json_payload=safe_dumps(audit_log))
+
+    async def _upload_json_to_file_path(self, file_path: str, json_payload: str) -> None:
+        payload_bytes: Final = json_payload.encode("utf-8")
+        if self.azure_storage_account_key:
+            await self._upload_bytes_with_account_key(file_path=file_path, content=payload_bytes)
+            return
+
+        await self.set_valid_azure_ad_token()
+        async_client: Final = get_async_httpx_client(llm_provider=httpxSpecialProvider.LoggingCallback)
+        base_url: Final = f"{self.azure_storage_dfs_endpoint}/{self.azure_storage_file_system}/{file_path}"
+        await self._create_file(async_client, base_url)
+        await self._append_data(async_client, base_url, json_payload)
+        await self._flush_data(async_client, base_url, len(payload_bytes))
+
     async def async_send_batch(self):
         """
         Sends the in memory logs queue to Azure Blob Storage
@@ -366,41 +384,22 @@ async def upload_to_azure_data_lake_with_azure_account_key(self, payload: Standa
         This is used when Azure Storage Account Key is set - Azure Storage Account Key does not work directly with Azure Rest API
         """
 
-        # Create an async service client
-
-        service_client: Final = await self.get_service_client()
-        # Get file system client
-        file_system_client: Final = service_client.get_file_system_client(file_system=self.azure_storage_file_system)
-
         try:
-            # Create directory with today's date
-            from datetime import datetime
-
             today: Final = datetime.now().strftime("%Y-%m-%d")
-            directory_client: Final = file_system_client.get_directory_client(today)
-
-            # check if the directory exists
-            if not await directory_client.exists():
-                await directory_client.create_directory()
-                verbose_logger.debug("Created directory: %s", today)
-
-            # Create a file client
             file_name: Final = adls_safe_file_name(payload.get("id"))
-            file_client: Final = directory_client.get_file_client(file_name)
-
-            # Create the file
-            await file_client.create_file()
-
-            # Content to append
-            content: Final = safe_dumps(payload).encode("utf-8")
-
-            # Append content to the file
-            await file_client.append_data(data=content, offset=0, length=len(content))
-
-            # Flush the content to finalize the file
-            await file_client.flush_data(position=len(content), offset=0)
-
+            await self._upload_bytes_with_account_key(
+                file_path=f"{today}/{file_name}", content=safe_dumps(payload).encode("utf-8")
+            )
             verbose_logger.debug("Successfully uploaded and wrote to %s/%s", today, file_name)
 
         except Exception as e:
             verbose_logger.exception("Error occurred: %s", e)
+
+    async def _upload_bytes_with_account_key(self, file_path: str, content: bytes) -> None:
+        service_client: Final = await self.get_service_client()
+        file_client: Final = service_client.get_file_system_client(
+            file_system=self.azure_storage_file_system
+        ).get_file_client(file_path)
+        await file_client.create_file()
+        await file_client.append_data(data=content, offset=0, length=len(content))
+        await file_client.flush_data(position=len(content), offset=0)

tests/unit/integrations/azure_storage/test_azure_storage.py

@@ -4,18 +4,20 @@
 import re
 import sys
 import threading
+from datetime import datetime, timedelta, timezone
 from unittest.mock import AsyncMock, MagicMock, patch
 
 import pytest
 
 from litellm.constants import _DEFAULT_TTL_FOR_HTTPX_CLIENTS
+from litellm.integrations.azure_storage import azure_storage as azure_storage_module
 from litellm.integrations.azure_storage.azure_storage import (
     AzureBlobStorageLogger,
     _cached_credential_chain_token_provider,
     adls_safe_file_name,
 )
 from litellm.types.secret_managers.get_azure_ad_token_provider import AzureCredentialType
-from litellm.types.utils import StandardLoggingPayload
+from litellm.types.utils import StandardAuditLogPayload, StandardLoggingPayload
 
 
 @pytest.fixture
@@ -336,11 +338,8 @@ async def test_account_key_auth_never_requests_a_token(workload_identity_env_var
     file_client.create_file = AsyncMock()
     file_client.append_data = AsyncMock()
     file_client.flush_data = AsyncMock()
-    directory_client = MagicMock()
-    directory_client.exists = AsyncMock(return_value=True)
-    directory_client.get_file_client = MagicMock(return_value=file_client)
     file_system_client = MagicMock()
-    file_system_client.get_directory_client = MagicMock(return_value=directory_client)
+    file_system_client.get_file_client = MagicMock(return_value=file_client)
     service_client = MagicMock()
     service_client.get_file_system_client = MagicMock(return_value=file_system_client)
     fake_aio_module = MagicMock()
@@ -523,11 +522,8 @@ async def test_account_key_upload_names_the_file_adls_safe_and_keeps_the_origina
     file_client.create_file = AsyncMock()
     file_client.append_data = AsyncMock()
     file_client.flush_data = AsyncMock()
-    directory_client = MagicMock()
-    directory_client.exists = AsyncMock(return_value=True)
-    directory_client.get_file_client = MagicMock(return_value=file_client)
     file_system_client = MagicMock()
-    file_system_client.get_directory_client = MagicMock(return_value=directory_client)
+    file_system_client.get_file_client = MagicMock(return_value=file_client)
     service_client = MagicMock()
     service_client.get_file_system_client = MagicMock(return_value=file_system_client)
     fake_aio_module = MagicMock()
@@ -537,7 +533,9 @@ async def test_account_key_upload_names_the_file_adls_safe_and_keeps_the_origina
         logger = AzureBlobStorageLogger()
         await logger.async_upload_payload_to_azure_blob_storage({"id": "resp_YWJjZA=="})
 
-    directory_client.get_file_client.assert_called_once_with("resp_YWJjZA.json")
+    file_system_client.get_file_client.assert_called_once_with(
+        f"{datetime.now().strftime('%Y-%m-%d')}/resp_YWJjZA.json"
+    )
     body = json.loads(file_client.append_data.call_args.kwargs["data"])
     assert body["id"] == "resp_YWJjZA==", "the stored payload must keep the original id byte for byte"
 
@@ -569,3 +567,84 @@ async def test_entra_upload_names_the_file_adls_safe_and_keeps_the_original_id(m
     ), f"the Entra path must be the rewritten name, got {put_call_args[0][0]!r}"
     append_call = mock_http_client.patch.call_args_list[0]
     assert "resp_YWJjZA==" in append_call[1]["data"], "the stored payload must keep the original id byte for byte"
+
+
+class _FrozenDatetime(datetime):
+    @classmethod
+    def now(cls, tz=None):
+        instant = datetime(2026, 9, 30, 3, 5, 9, tzinfo=timezone.utc)
+        if tz is None:
+            return instant.astimezone(timezone(timedelta(hours=-7))).replace(tzinfo=None)
+        return instant.astimezone(tz)
+
+
+def _key_created_audit_log(audit_log_id: str) -> StandardAuditLogPayload:
+    return StandardAuditLogPayload(
+        id=audit_log_id,
+        updated_at="2026-09-30T03:05:09+00:00",
+        changed_by="admin-user",
+        changed_by_api_key="hashed-admin-key",
+        action="created",
+        table_name="LiteLLM_VerificationToken",
+        object_id="hashed-new-key",
+        before_value=None,
+        updated_values=json.dumps({"key_alias": "audit-key"}),
+    )
+
+
+@pytest.mark.asyncio
+async def test_audit_log_is_uploaded_to_a_dated_audit_logs_path_with_a_service_principal(mock_env_vars, monkeypatch):
+    monkeypatch.setattr(azure_storage_module, "datetime", _FrozenDatetime)
+    audit_log = _key_created_audit_log("audit-id-sp")
+    with patch(  # test-quality-ok: REST client is created inside the method; assert emitted requests
+        "litellm.integrations.azure_storage.azure_storage.get_async_httpx_client"
+    ) as mock_get_client:
+        mock_http_client = AsyncMock()
+        mock_http_client.put.return_value = MagicMock()
+        mock_http_client.patch.return_value = MagicMock()
+        mock_get_client.return_value = mock_http_client
+
+        logger = AzureBlobStorageLogger()
+        logger.azure_auth_token = "sp-token"
+        logger.token_expiry = datetime(2026, 9, 30, 9, 30)
+        await logger.async_log_audit_log_event(audit_log)
+
+    file_url = (
+        "https://test-account.dfs.core.windows.net/test-container/audit_logs/2026-09-30/03-05-09_audit-id-sp.json"
+    )
+    assert [call.args[0] for call in mock_http_client.put.call_args_list] == [f"{file_url}?resource=file"]
+    append_call, flush_call = mock_http_client.patch.call_args_list
+    assert append_call.args[0] == f"{file_url}?action=append&position=0"
+    assert json.loads(append_call.kwargs["data"]) == audit_log
+    assert flush_call.args[0] == f"{file_url}?action=flush&position={len(append_call.kwargs['data'].encode())}"
+    assert logger.log_queue == [], "audit logs must not be mixed into the request log queue"
+
+
+@pytest.mark.asyncio
+async def test_audit_log_is_uploaded_to_a_dated_audit_logs_path_with_an_account_key(mock_env_vars, monkeypatch):
+    monkeypatch.setenv("AZURE_STORAGE_ACCOUNT_KEY", "dGVzdC1rZXk=")
+    monkeypatch.setattr(azure_storage_module, "datetime", _FrozenDatetime)
+    audit_log = _key_created_audit_log("audit-id-key")
+
+    file_client = MagicMock()
+    file_client.create_file = AsyncMock()
+    file_client.append_data = AsyncMock()
+    file_client.flush_data = AsyncMock()
+    file_system_client = MagicMock()
+    file_system_client.get_file_client = MagicMock(return_value=file_client)
+    service_client = MagicMock()
+    service_client.get_file_system_client = MagicMock(return_value=file_system_client)
+    fake_aio_module = MagicMock()
+    fake_aio_module.DataLakeServiceClient = MagicMock(return_value=service_client)
+
+    with patch.dict(sys.modules, {"azure.storage.filedatalake.aio": fake_aio_module}):
+        logger = AzureBlobStorageLogger()
+        await logger.async_log_audit_log_event(audit_log)
+
+    service_client.get_file_system_client.assert_called_once_with(file_system="test-container")
+    file_system_client.get_file_client.assert_called_once_with("audit_logs/2026-09-30/03-05-09_audit-id-key.json")
+    file_client.create_file.assert_awaited_once_with()
+    uploaded = file_client.append_data.await_args.kwargs["data"]
+    assert json.loads(uploaded) == audit_log
+    assert file_client.append_data.await_args.kwargs == {"data": uploaded, "offset": 0, "length": len(uploaded)}
+    file_client.flush_data.assert_awaited_once_with(position=len(uploaded), offset=0)