diff --git a/litellm/integrations/azure_sentinel/azure_sentinel.py b/litellm/integrations/azure_sentinel/azure_sentinel.py index 182a2e185ef..5f8afe58cb0 100644 --- a/litellm/integrations/azure_sentinel/azure_sentinel.py +++ b/litellm/integrations/azure_sentinel/azure_sentinel.py @@ -61,13 +61,15 @@ def __init__( 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. + If not provided, will use AZURE_SENTINEL_AUDIT_STREAM_NAME env var or the standard stream name. """ self.async_httpx_client = get_async_httpx_client(llm_provider=httpxSpecialProvider.LoggingCallback) resolved_dcr_immutable_id = dcr_immutable_id or os.getenv("AZURE_SENTINEL_DCR_IMMUTABLE_ID") resolved_stream_name = stream_name or os.getenv("AZURE_SENTINEL_STREAM_NAME") or "Custom-LiteLLM" - resolved_audit_stream_name = audit_stream_name or resolved_stream_name + resolved_audit_stream_name = ( + audit_stream_name or os.getenv("AZURE_SENTINEL_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") resolved_client_id = client_id or os.getenv("AZURE_SENTINEL_CLIENT_ID") or os.getenv("AZURE_CLIENT_ID") diff --git a/tests/test_litellm/integrations/test_azure_sentinel.py b/tests/test_litellm/integrations/test_azure_sentinel.py index 30b246202fc..55e462c82b8 100644 --- a/tests/test_litellm/integrations/test_azure_sentinel.py +++ b/tests/test_litellm/integrations/test_azure_sentinel.py @@ -263,3 +263,36 @@ async def mock_post(*args, **kwargs): ] assert "Custom-LiteLLM-Audit" in audit_call.kwargs["url"] assert json.loads(audit_call.kwargs["data"].decode("utf-8")) == [audit_log] + + +@pytest.mark.asyncio +async def test_azure_sentinel_audit_stream_name_from_env_var(monkeypatch): + """Audit stream resolves from AZURE_SENTINEL_AUDIT_STREAM_NAME when the string + callback constructs the logger with no audit_stream_name argument.""" + monkeypatch.setenv("AZURE_SENTINEL_STREAM_NAME", "Custom-LiteLLM-Standard") + monkeypatch.setenv("AZURE_SENTINEL_AUDIT_STREAM_NAME", "Custom-LiteLLM-Audit") + + 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", + ) + + assert logger.audit_stream_name == "Custom-LiteLLM-Audit" + assert "streams/Custom-LiteLLM-Audit" in logger.audit_api_endpoint + assert "streams/Custom-LiteLLM-Standard" in logger.api_endpoint + + with patch("asyncio.create_task", side_effect=_close_periodic_flush_task): + explicit_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_stream_name="Custom-LiteLLM-Explicit", + ) + + assert explicit_logger.audit_stream_name == "Custom-LiteLLM-Explicit"