tiennm99/litellm
1
0
Fork 0
mirror of https://github.com/tiennm99/litellm.git synced 2026-08-21 08:26:34 +00:00
Code Issues Packages Projects Releases Wiki Activity

Add Azure Sentinel audit log support (#27280)

Browse Source 
* Add Azure Sentinel audit log callback support

Co-authored-by: ishaan-berri <ishaan-berri@users.noreply.github.com>

* Fix Azure Sentinel audit log batching

Co-authored-by: ishaan-berri <ishaan-berri@users.noreply.github.com>

* Fix Azure Sentinel CI checks

Co-authored-by: ishaan-berri <ishaan-berri@users.noreply.github.com>

---------

Co-authored-by: oss-agent-shin <279349115+oss-agent-shin@users.noreply.github.com>
Co-authored-by: ishaan-berri <ishaan-berri@users.noreply.github.com>
This commit is contained in:
oss-agent-shin
2026-05-06 15:50:06 -07:00
committed by GitHub
co-authored by oss-agent-shin ishaan-berri
parent aba131d3cf
commit b318231fe9
4 changed files with 305 additions and 35 deletions
Download Patch File Download Diff File Expand all files Collapse all files
litellm/integrations/azure_sentinel/azure_sentinel.py
+111 -28
View File
@@ -14,16 +14,18 @@ For batching specific details see CustomBatchLogger class
import asyncio
import os
import time
import traceback
from typing import List, Optional
from typing import List, Optional, Union
from litellm._logging import verbose_logger
from litellm.integrations.custom_batch_logger import CustomBatchLogger
from litellm.litellm_core_utils.safe_json_dumps import safe_dumps
from litellm.llms.custom_httpx.http_handler import (
get_async_httpx_client,
httpxSpecialProvider,
)
from litellm.types.utils import StandardLoggingPayload
from litellm.types.utils import StandardAuditLogPayload, StandardLoggingPayload
class AzureSentinelLogger(CustomBatchLogger):
@@ -39,6 +41,7 @@ class AzureSentinelLogger(CustomBatchLogger):
tenant_id: Optional[str] = None,
client_id: Optional[str] = None,
client_secret: Optional[str] = None,
audit_stream_name: Optional[str] = None,
**kwargs,
):
"""
@@ -57,57 +60,77 @@ class AzureSentinelLogger(CustomBatchLogger):
If not provided, will use AZURE_SENTINEL_CLIENT_ID or AZURE_CLIENT_ID env var.
client_secret (str, optional): Azure Client Secret for OAuth2 authentication.
If not provided, will use AZURE_SENTINEL_CLIENT_SECRET or AZURE_CLIENT_SECRET env var.
audit_stream_name (str, optional): Stream name from DCR for audit logs.
If not provided, audit logs use the standard stream name.
"""
self.async_httpx_client = get_async_httpx_client(
llm_provider=httpxSpecialProvider.LoggingCallback
)
self.dcr_immutable_id = dcr_immutable_id or os.getenv(
resolved_dcr_immutable_id = dcr_immutable_id or os.getenv(
"AZURE_SENTINEL_DCR_IMMUTABLE_ID"
)
self.stream_name = stream_name or os.getenv(
"AZURE_SENTINEL_STREAM_NAME", "Custom-LiteLLM"
resolved_stream_name = (
stream_name or os.getenv("AZURE_SENTINEL_STREAM_NAME") or "Custom-LiteLLM"
)
self.endpoint = endpoint or os.getenv("AZURE_SENTINEL_ENDPOINT")
self.tenant_id = (
resolved_audit_stream_name = audit_stream_name or resolved_stream_name
resolved_endpoint = endpoint or os.getenv("AZURE_SENTINEL_ENDPOINT")
resolved_tenant_id = (
tenant_id
or os.getenv("AZURE_SENTINEL_TENANT_ID")
or os.getenv("AZURE_TENANT_ID")
)
self.client_id = (
resolved_client_id = (
client_id
or os.getenv("AZURE_SENTINEL_CLIENT_ID")
or os.getenv("AZURE_CLIENT_ID")
)
self.client_secret = (
resolved_client_secret = (
client_secret
or os.getenv("AZURE_SENTINEL_CLIENT_SECRET")
or os.getenv("AZURE_CLIENT_SECRET")
)
if not self.dcr_immutable_id:
if not resolved_dcr_immutable_id:
raise ValueError(
"AZURE_SENTINEL_DCR_IMMUTABLE_ID is required. Set it as an environment variable or pass dcr_immutable_id parameter."
)
if not self.endpoint:
if not resolved_endpoint:
raise ValueError(
"AZURE_SENTINEL_ENDPOINT is required. Set it as an environment variable or pass endpoint parameter."
)
if not self.tenant_id:
if not resolved_tenant_id:
raise ValueError(
"AZURE_SENTINEL_TENANT_ID or AZURE_TENANT_ID is required. Set it as an environment variable or pass tenant_id parameter."
)
if not self.client_id:
if not resolved_client_id:
raise ValueError(
"AZURE_SENTINEL_CLIENT_ID or AZURE_CLIENT_ID is required. Set it as an environment variable or pass client_id parameter."
)
if not self.client_secret:
if not resolved_client_secret:
raise ValueError(
"AZURE_SENTINEL_CLIENT_SECRET or AZURE_CLIENT_SECRET is required. Set it as an environment variable or pass client_secret parameter."
)
self.dcr_immutable_id = resolved_dcr_immutable_id
self.stream_name = resolved_stream_name
self.audit_stream_name = resolved_audit_stream_name
self.endpoint = resolved_endpoint
self.tenant_id = resolved_tenant_id
self.client_id = resolved_client_id
self.client_secret = resolved_client_secret
# Build API endpoint: {Endpoint}/dataCollectionRules/{DCR Immutable ID}/streams/{Stream Name}?api-version=2023-01-01
self.api_endpoint = f"{self.endpoint.rstrip('/')}/dataCollectionRules/{self.dcr_immutable_id}/streams/{self.stream_name}?api-version=2023-01-01"
self.api_endpoint = self._build_api_endpoint(
endpoint=resolved_endpoint,
dcr_immutable_id=resolved_dcr_immutable_id,
stream_name=resolved_stream_name,
)
self.audit_api_endpoint = self._build_api_endpoint(
endpoint=resolved_endpoint,
dcr_immutable_id=resolved_dcr_immutable_id,
stream_name=resolved_audit_stream_name,
)
# OAuth2 scope for Azure Monitor
self.oauth_scope = "https://monitor.azure.com/.default"
@@ -118,6 +141,13 @@ class AzureSentinelLogger(CustomBatchLogger):
super().__init__(**kwargs, flush_lock=self.flush_lock)
asyncio.create_task(self.periodic_flush())
self.log_queue: List[StandardLoggingPayload] = []
self.audit_log_queue: List[StandardAuditLogPayload] = []
@staticmethod
def _build_api_endpoint(
endpoint: str, dcr_immutable_id: str, stream_name: str
) -> str:
return f"{endpoint.rstrip('/')}/dataCollectionRules/{dcr_immutable_id}/streams/{stream_name}?api-version=2023-01-01"
async def _get_oauth_token(self) -> str:
"""
@@ -126,9 +156,6 @@ class AzureSentinelLogger(CustomBatchLogger):
Returns:
Bearer token string
"""
# Check if we have a valid cached token
import time
if (
self.oauth_token
and self.oauth_token_expires_at
@@ -170,9 +197,6 @@ class AzureSentinelLogger(CustomBatchLogger):
if not self.oauth_token:
raise Exception("OAuth2 token response did not contain access_token")
# Cache token expiry time
import time
self.oauth_token_expires_at = time.time() + expires_in
return self.oauth_token
@@ -246,6 +270,34 @@ class AzureSentinelLogger(CustomBatchLogger):
)
pass
async def async_log_audit_log_event(
self, audit_log: StandardAuditLogPayload
) -> None:
"""
Async log LiteLLM audit log events to Azure Sentinel.
Audit logs are queued separately from standard LLM logs so mixed callback
usage never sends schema-mismatched records in the same ingestion batch.
"""
try:
verbose_logger.debug(
"Azure Sentinel: Logging audit event id=%s action=%s table=%s",
audit_log.get("id"),
audit_log.get("action"),
audit_log.get("table_name"),
)
self.audit_log_queue.append(audit_log)
if len(self.audit_log_queue) >= self.batch_size:
await self.async_send_audit_batch()
except Exception as e:
verbose_logger.exception(
f"Azure Sentinel Audit Log Layer Error - {str(e)}\n{traceback.format_exc()}"
)
pass
async def async_send_batch(self):
"""
Sends the batch of logs to Azure Monitor Logs Ingestion API
@@ -253,22 +305,42 @@ class AzureSentinelLogger(CustomBatchLogger):
Raises:
Raises a NON Blocking verbose_logger.exception if an error occurs
"""
await self._async_send_batch_to_api(
log_queue=self.log_queue,
api_endpoint=self.api_endpoint,
log_type="logs",
)
async def async_send_audit_batch(self):
"""
Sends the batch of audit logs to Azure Monitor Logs Ingestion API
"""
await self._async_send_batch_to_api(
log_queue=self.audit_log_queue,
api_endpoint=self.audit_api_endpoint,
log_type="audit logs",
)
async def _async_send_batch_to_api(
self,
log_queue: List[Union[StandardLoggingPayload, StandardAuditLogPayload]],
api_endpoint: str,
log_type: str,
) -> None:
try:
if not self.log_queue:
if not log_queue:
return
verbose_logger.debug(
"Azure Sentinel - about to flush %s events", len(self.log_queue)
"Azure Sentinel - about to flush %s %s", len(log_queue), log_type
)
from litellm.litellm_core_utils.safe_json_dumps import safe_dumps
# Get OAuth2 token
bearer_token = await self._get_oauth_token()
# Convert log queue to JSON array format expected by Logs Ingestion API
# Each log entry should be a JSON object in the array
body = safe_dumps(self.log_queue)
body = safe_dumps(log_queue)
# Set headers for Logs Ingestion API
headers = {
@@ -278,7 +350,7 @@ class AzureSentinelLogger(CustomBatchLogger):
# Send the request
response = await self.async_httpx_client.post(
url=self.api_endpoint, data=body.encode("utf-8"), headers=headers
url=api_endpoint, data=body.encode("utf-8"), headers=headers
)
if response.status_code not in [200, 204]:
@@ -301,4 +373,15 @@ class AzureSentinelLogger(CustomBatchLogger):
f"Azure Sentinel Error sending batch API - {str(e)}\n{traceback.format_exc()}"
)
finally:
self.log_queue.clear()
log_queue.clear()
async def flush_queue(self):
if self.flush_lock is None:
return
async with self.flush_lock:
if self.log_queue:
await self.async_send_batch()
if self.audit_log_queue:
await self.async_send_audit_batch()
self.last_flush_time = time.time()
litellm/litellm_core_utils/custom_logger_registry.py
+2
View File
@@ -14,6 +14,7 @@ from litellm import _custom_logger_compatible_callbacks_literal
from litellm.integrations.agentops import AgentOps
from litellm.integrations.anthropic_cache_control_hook import AnthropicCacheControlHook
from litellm.integrations.argilla import ArgillaLogger
from litellm.integrations.azure_sentinel.azure_sentinel import AzureSentinelLogger
from litellm.integrations.azure_storage.azure_storage import AzureBlobStorageLogger
from litellm.integrations.bitbucket import BitBucketPromptManager
from litellm.integrations.braintrust_logging import BraintrustLogger
@@ -73,6 +74,7 @@ class CustomLoggerRegistry:
"opik": OpikLogger,
"argilla": ArgillaLogger,
"opentelemetry": OpenTelemetry,
"azure_sentinel": AzureSentinelLogger,
"azure_storage": AzureBlobStorageLogger,
"humanloop": HumanloopLogger,
# OTEL compatible loggers
litellm/proxy/_types.py
+13
View File
@@ -3348,6 +3348,19 @@ class AllCallbacks(LiteLLMPydanticObjectBase):
],
)
azure_sentinel: CallbackOnUI = CallbackOnUI(
litellm_callback_name="azure_sentinel",
ui_callback_name="Azure Sentinel",
litellm_callback_params=[
"AZURE_SENTINEL_DCR_IMMUTABLE_ID",
"AZURE_SENTINEL_ENDPOINT",
"AZURE_SENTINEL_TENANT_ID",
"AZURE_SENTINEL_CLIENT_ID",
"AZURE_SENTINEL_CLIENT_SECRET",
"AZURE_SENTINEL_STREAM_NAME",
],
)
openmeter: CallbackOnUI = CallbackOnUI(
litellm_callback_name="openmeter",
ui_callback_name="OpenMeter",
tests/test_litellm/integrations/test_azure_sentinel.py
+179 -7
View File
@@ -2,13 +2,18 @@
Test Azure Sentinel logging integration
"""
import datetime
from unittest.mock import AsyncMock, patch
import json
from unittest.mock import AsyncMock, MagicMock, patch
import pytest
from litellm.integrations.azure_sentinel.azure_sentinel import AzureSentinelLogger
from litellm.types.utils import StandardLoggingPayload
from litellm.types.utils import StandardAuditLogPayload, StandardLoggingPayload
def _close_periodic_flush_task(coro):
coro.close()
return None
@pytest.mark.asyncio
@@ -20,7 +25,7 @@ async def test_azure_sentinel_oauth_and_send_batch():
test_client_id = "test-client-id"
test_client_secret = "test-client-secret"
with patch("asyncio.create_task"):
with patch("asyncio.create_task", side_effect=_close_periodic_flush_task):
logger = AzureSentinelLogger(
dcr_immutable_id=test_dcr_id,
endpoint=test_endpoint,
@@ -42,9 +47,6 @@ async def test_azure_sentinel_oauth_and_send_batch():
# Add to queue
logger.log_queue.append(standard_payload)
# Mock OAuth token response
from unittest.mock import MagicMock
mock_token_response = MagicMock()
mock_token_response.status_code = 200
mock_token_response.json = MagicMock(
@@ -91,3 +93,173 @@ async def test_azure_sentinel_oauth_and_send_batch():
# Verify queue is cleared
assert len(logger.log_queue) == 0
@pytest.mark.asyncio
async def test_azure_sentinel_queues_audit_log_event():
"""Test that Azure Sentinel supports direct audit log callbacks"""
with patch("asyncio.create_task", side_effect=_close_periodic_flush_task):
logger = AzureSentinelLogger(
dcr_immutable_id="dcr-test123456789",
endpoint="https://test-dce.eastus-1.ingest.monitor.azure.com",
tenant_id="test-tenant-id",
client_id="test-client-id",
client_secret="test-client-secret",
)
logger.batch_size = 2
logger.async_send_audit_batch = AsyncMock()
audit_log = StandardAuditLogPayload(
id="audit-123",
updated_at="2026-05-06T04:39:00+00:00",
changed_by="user-1",
changed_by_api_key="sk-test",
action="created",
table_name="LiteLLM_TeamTable",
object_id="team-1",
before_value=None,
updated_values='{"team_alias": "sentinel-demo"}',
)
await logger.async_log_audit_log_event(audit_log)
assert logger.audit_log_queue == [audit_log]
logger.async_send_audit_batch.assert_not_called()
await logger.async_log_audit_log_event(audit_log)
assert logger.audit_log_queue == [audit_log, audit_log]
logger.async_send_audit_batch.assert_awaited_once()
@pytest.mark.asyncio
async def test_azure_sentinel_sends_audit_log_payload_to_ingestion_api():
"""Test that queued audit logs are sent to Azure Monitor Logs Ingestion"""
with patch("asyncio.create_task", side_effect=_close_periodic_flush_task):
logger = AzureSentinelLogger(
dcr_immutable_id="dcr-test123456789",
endpoint="https://test-dce.eastus-1.ingest.monitor.azure.com",
tenant_id="test-tenant-id",
client_id="test-client-id",
client_secret="test-client-secret",
)
audit_log = StandardAuditLogPayload(
id="audit-123",
updated_at="2026-05-06T04:39:00+00:00",
changed_by="user-1",
changed_by_api_key="sk-test",
action="created",
table_name="LiteLLM_TeamTable",
object_id="team-1",
before_value=None,
updated_values='{"team_alias": "sentinel-demo"}',
)
await logger.async_log_audit_log_event(audit_log)
mock_token_response = MagicMock()
mock_token_response.status_code = 200
mock_token_response.json = MagicMock(
return_value={
"access_token": "test-bearer-token",
"expires_in": 3600,
}
)
mock_token_response.text = "Success"
mock_api_response = MagicMock()
mock_api_response.status_code = 204
mock_api_response.text = "Success"
async def mock_post(*args, **kwargs):
if "oauth2/v2.0/token" in kwargs.get("url", ""):
return mock_token_response
return mock_api_response
logger.async_httpx_client.post = AsyncMock(side_effect=mock_post)
await logger.flush_queue()
api_call_args = logger.async_httpx_client.post.call_args_list[-1]
body = json.loads(api_call_args.kwargs["data"].decode("utf-8"))
assert body == [audit_log]
assert "dcr-test123456789" in api_call_args.kwargs["url"]
assert "Custom-LiteLLM" in api_call_args.kwargs["url"]
assert len(logger.audit_log_queue) == 0
@pytest.mark.asyncio
async def test_azure_sentinel_flushes_standard_and_audit_logs_separately():
"""Test mixed callback roles do not send schema-mismatched batches."""
with patch("asyncio.create_task", side_effect=_close_periodic_flush_task):
logger = AzureSentinelLogger(
dcr_immutable_id="dcr-test123456789",
stream_name="Custom-LiteLLM-Standard",
audit_stream_name="Custom-LiteLLM-Audit",
endpoint="https://test-dce.eastus-1.ingest.monitor.azure.com",
tenant_id="test-tenant-id",
client_id="test-client-id",
client_secret="test-client-secret",
)
standard_payload = StandardLoggingPayload(
id="standard-123",
call_type="completion",
model="gpt-3.5-turbo",
status="success",
messages=[{"role": "user", "content": "Hello"}],
response={"choices": [{"message": {"content": "Hi"}}]},
)
audit_log = StandardAuditLogPayload(
id="audit-123",
updated_at="2026-05-06T04:39:00+00:00",
changed_by="user-1",
changed_by_api_key="sk-test",
action="created",
table_name="LiteLLM_TeamTable",
object_id="team-1",
before_value=None,
updated_values='{"team_alias": "sentinel-demo"}',
)
logger.log_queue.append(standard_payload)
await logger.async_log_audit_log_event(audit_log)
mock_token_response = MagicMock()
mock_token_response.status_code = 200
mock_token_response.json = MagicMock(
return_value={
"access_token": "test-bearer-token",
"expires_in": 3600,
}
)
mock_token_response.text = "Success"
mock_api_response = MagicMock()
mock_api_response.status_code = 204
mock_api_response.text = "Success"
async def mock_post(*args, **kwargs):
if "oauth2/v2.0/token" in kwargs.get("url", ""):
return mock_token_response
return mock_api_response
logger.async_httpx_client.post = AsyncMock(side_effect=mock_post)
await logger.flush_queue()
ingestion_calls = [
call
for call in logger.async_httpx_client.post.call_args_list
if "dataCollectionRules" in call.kwargs["url"]
]
assert len(ingestion_calls) == 2
standard_call, audit_call = ingestion_calls
assert "Custom-LiteLLM-Standard" in standard_call.kwargs["url"]
assert json.loads(standard_call.kwargs["data"].decode("utf-8")) == [
standard_payload
]
assert "Custom-LiteLLM-Audit" in audit_call.kwargs["url"]
assert json.loads(audit_call.kwargs["data"].decode("utf-8")) == [audit_log]