scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-12 03:48:56 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 118

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

Browse source 
Co-authored-by: yassin <yassin@berri.ai>
Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
devin-ai-integration[bot] 2026-10-10 14:48:56 -07:00 • committed by GitHub
parent ef9713f4de
commit 60f35e35ad
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 121 additions and 43 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

65
litellm/integrations/azure_storage/azure_storage.py
View file

@ -2,7 +2,7 @@ import asyncio
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.secret_managers.get_azure_ad_token_provider import (
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 @@ class AzureBlobStorageLogger(CustomBatchLogger):
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 @@ class AzureBlobStorageLogger(CustomBatchLogger):
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)

99
tests/unit/integrations/azure_storage/test_azure_storage.py
View file

@ -4,18 +4,20 @@ import json
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)