diff --git a/litellm/constants.py b/litellm/constants.py index 73f0d1e160e..0c7316455d6 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -1305,6 +1305,7 @@ RESPONSE_FORMAT_TOOL_NAME = "json_tool_call" # default tool name used when conv ########################### Logging Callback Constants ########################### AZURE_STORAGE_MSFT_VERSION: Final = "2019-07-07" +AZURE_STORAGE_DEFAULT_ENDPOINT_SUFFIX: Final = "core.windows.net" PROMETHEUS_BUDGET_METRICS_REFRESH_INTERVAL_MINUTES: Final = int( os.getenv("PROMETHEUS_BUDGET_METRICS_REFRESH_INTERVAL_MINUTES", 5) ) diff --git a/litellm/integrations/azure_storage/azure_storage.py b/litellm/integrations/azure_storage/azure_storage.py index d2181dbbb38..cb7175691df 100644 --- a/litellm/integrations/azure_storage/azure_storage.py +++ b/litellm/integrations/azure_storage/azure_storage.py @@ -6,7 +6,11 @@ from typing import Final from litellm._logging import verbose_logger from litellm._uuid import uuid -from litellm.constants import _DEFAULT_TTL_FOR_HTTPX_CLIENTS, AZURE_STORAGE_MSFT_VERSION +from litellm.constants import ( + _DEFAULT_TTL_FOR_HTTPX_CLIENTS, + AZURE_STORAGE_DEFAULT_ENDPOINT_SUFFIX, + AZURE_STORAGE_MSFT_VERSION, +) from litellm.integrations.custom_batch_logger import CustomBatchLogger from litellm.litellm_core_utils.safe_json_dumps import safe_dumps from litellm.llms.azure.common_utils import get_azure_ad_token_from_entra_id @@ -41,6 +45,9 @@ class AzureBlobStorageLogger(CustomBatchLogger): if not _azure_storage_file_system: raise ValueError("Missing required environment variable: AZURE_STORAGE_FILE_SYSTEM") self.azure_storage_file_system: str = _azure_storage_file_system + self.azure_storage_endpoint_suffix: str = ( + os.getenv("AZURE_STORAGE_ENDPOINT_SUFFIX") or AZURE_STORAGE_DEFAULT_ENDPOINT_SUFFIX + ) self._service_client = None # Time that the azure service client expires, in order to reset the connection pool and keep it fresh self._service_client_timeout: float | None = None @@ -59,6 +66,14 @@ class AzureBlobStorageLogger(CustomBatchLogger): ) raise e + @property + def azure_storage_dfs_endpoint(self) -> str: + return f"https://{self.azure_storage_account_name}.dfs.{self.azure_storage_endpoint_suffix}" + + @property + def azure_storage_blob_endpoint(self) -> str: + return f"https://{self.azure_storage_account_name}.blob.{self.azure_storage_endpoint_suffix}" + async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): """ Async Log success events to Azure Blob Storage @@ -144,7 +159,7 @@ class AzureBlobStorageLogger(CustomBatchLogger): json_payload: Final = safe_dumps(payload) + "\n" # Add newline for each log entry payload_bytes: Final = json_payload.encode("utf-8") filename: Final = f"{payload.get('id') or str(uuid.uuid4())}.json" - base_url = f"https://{self.azure_storage_account_name}.dfs.core.windows.net/{self.azure_storage_file_system}/{filename}" + base_url = f"{self.azure_storage_dfs_endpoint}/{self.azure_storage_file_system}/{filename}" # Execute the 3-step upload process await self._create_file(async_client, base_url) @@ -296,7 +311,7 @@ class AzureBlobStorageLogger(CustomBatchLogger): self._service_client = None if not self._service_client: self._service_client = DataLakeServiceClient( - account_url=f"https://{self.azure_storage_account_name}.dfs.core.windows.net", + account_url=self.azure_storage_dfs_endpoint, credential=self.azure_storage_account_key, ) self._service_client_timeout = time.time() + _DEFAULT_TTL_FOR_HTTPX_CLIENTS diff --git a/litellm/llms/base_llm/files/azure_blob_storage_backend.py b/litellm/llms/base_llm/files/azure_blob_storage_backend.py index b8d27b71636..e22cf528856 100644 --- a/litellm/llms/base_llm/files/azure_blob_storage_backend.py +++ b/litellm/llms/base_llm/files/azure_blob_storage_backend.py @@ -8,7 +8,7 @@ to reuse all authentication and Azure Storage operations. import time from typing import Final -from urllib.parse import quote +from urllib.parse import quote, urlparse from litellm._logging import verbose_logger from litellm._uuid import uuid @@ -47,6 +47,8 @@ class AzureBlobStorageBackend(BaseFileStorageBackend, AzureBlobStorageLogger): - AZURE_STORAGE_TENANT_ID (optional, if using Azure AD) - AZURE_STORAGE_CLIENT_ID (optional, if using Azure AD) - AZURE_STORAGE_CLIENT_SECRET (optional, if using Azure AD) + - AZURE_STORAGE_ENDPOINT_SUFFIX (optional, defaults to core.windows.net; set to + core.usgovcloudapi.net or another sovereign-cloud suffix as needed) Note: We skip periodic_flush since we're not using this as a logger. """ @@ -103,7 +105,7 @@ class AzureBlobStorageBackend(BaseFileStorageBackend, AzureBlobStorageLogger): """ Upload a file to Azure Blob Storage. - Returns the blob URL in format: https://{account}.blob.core.windows.net/{container}/{path} + Returns the blob URL in format: https://{account}.blob.{endpoint_suffix}/{container}/{path} """ try: # Generate file name @@ -172,7 +174,7 @@ class AzureBlobStorageBackend(BaseFileStorageBackend, AzureBlobStorageLogger): await file_client.flush_data(position=len(file_content), offset=0) # Return blob URL (not DFS URL) - blob_url = f"https://{self.azure_storage_account_name}.blob.core.windows.net/{self.azure_storage_file_system}/{full_path}" + blob_url = f"{self.azure_storage_blob_endpoint}/{self.azure_storage_file_system}/{full_path}" return blob_url async def _upload_file_with_azure_ad(self, file_content: bytes, full_path: str) -> str: @@ -188,7 +190,7 @@ class AzureBlobStorageBackend(BaseFileStorageBackend, AzureBlobStorageLogger): async_client: Final = get_async_httpx_client(llm_provider=httpxSpecialProvider.LoggingCallback) # Use DFS endpoint for upload - base_url = f"https://{self.azure_storage_account_name}.dfs.core.windows.net/{self.azure_storage_file_system}/{full_path}" + base_url = f"{self.azure_storage_dfs_endpoint}/{self.azure_storage_file_system}/{full_path}" # Execute 3-step upload process: create, append, flush # Reuse the logger's helper methods @@ -198,7 +200,7 @@ class AzureBlobStorageBackend(BaseFileStorageBackend, AzureBlobStorageLogger): await self._flush_data(async_client, base_url, len(file_content)) # Return blob URL (not DFS URL) - blob_url = f"https://{self.azure_storage_account_name}.blob.core.windows.net/{self.azure_storage_file_system}/{full_path}" + blob_url = f"{self.azure_storage_blob_endpoint}/{self.azure_storage_file_system}/{full_path}" return blob_url async def _append_data_bytes(self, client, base_url: str, file_content: bytes): @@ -222,23 +224,22 @@ class AzureBlobStorageBackend(BaseFileStorageBackend, AzureBlobStorageLogger): Download a file from Azure Blob Storage. Args: - storage_url: Blob URL in format: https://{account}.blob.core.windows.net/{container}/{path} + storage_url: Blob URL in format: https://{account}.blob.{endpoint_suffix}/{container}/{path} Returns: bytes: File content """ try: # Parse blob URL to extract path - # URL format: https://{account}.blob.core.windows.net/{container}/{path} - if ".blob.core.windows.net/" not in storage_url: + # URL format: https://{account}.blob.{endpoint_suffix}/{container}/{path} + parsed_url: Final = urlparse(storage_url) + if ".blob." not in (parsed_url.hostname or ""): raise ValueError(f"Invalid Azure Blob Storage URL: {storage_url}") # Extract path after container name - container_and_path: Final = storage_url.split(".blob.core.windows.net/", 1)[1] - path_parts: Final = container_and_path.split("/", 1) - if len(path_parts) < 2: + _, _, file_path = parsed_url.path.lstrip("/").partition("/") + if not file_path: raise ValueError(f"Invalid Azure Blob Storage URL format: {storage_url}") - file_path: Final = path_parts[1] # Path after container name if self.azure_storage_account_key: # Use Azure SDK (reuse logger's service client) @@ -279,7 +280,7 @@ class AzureBlobStorageBackend(BaseFileStorageBackend, AzureBlobStorageLogger): async_client: Final = get_async_httpx_client(llm_provider=httpxSpecialProvider.LoggingCallback) # Use blob endpoint for download (simpler than DFS) - blob_url = f"https://{self.azure_storage_account_name}.blob.core.windows.net/{self.azure_storage_file_system}/{file_path}" + blob_url = f"{self.azure_storage_blob_endpoint}/{self.azure_storage_file_system}/{file_path}" headers: Final = { "x-ms-version": AZURE_STORAGE_MSFT_VERSION, diff --git a/tests/test_litellm/integrations/azure_storage/test_azure_storage.py b/tests/test_litellm/integrations/azure_storage/test_azure_storage.py index d6c9d7a5c92..5d7c55e81af 100644 --- a/tests/test_litellm/integrations/azure_storage/test_azure_storage.py +++ b/tests/test_litellm/integrations/azure_storage/test_azure_storage.py @@ -20,6 +20,13 @@ def mock_env_vars(monkeypatch): monkeypatch.setenv("AZURE_STORAGE_TENANT_ID", "test-tenant-id") monkeypatch.setenv("AZURE_STORAGE_CLIENT_ID", "test-client-id") monkeypatch.setenv("AZURE_STORAGE_CLIENT_SECRET", "test-client-secret") + monkeypatch.delenv("AZURE_STORAGE_ENDPOINT_SUFFIX", raising=False) + + +@pytest.fixture +def mock_gov_env_vars(mock_env_vars, monkeypatch): + """Point the logger at an Azure Government storage account""" + monkeypatch.setenv("AZURE_STORAGE_ENDPOINT_SUFFIX", "core.usgovcloudapi.net") @pytest.mark.asyncio @@ -99,3 +106,76 @@ async def test_async_upload_payload_to_azure_blob_storage(mock_env_vars): # Verify raise_for_status was called on all responses assert mock_response.raise_for_status.call_count == 3 + + +@pytest.mark.asyncio +async def test_async_upload_payload_uses_configured_endpoint_suffix(mock_gov_env_vars): + """ + AZURE_STORAGE_ENDPOINT_SUFFIX must reach the Entra-ID REST upload path so a + sovereign-cloud account is addressed instead of the commercial dfs host. + """ + with patch( + "litellm.integrations.azure_storage.azure_storage.get_async_httpx_client" + ) as mock_get_client: + mock_http_client = AsyncMock() + mock_response = MagicMock() + mock_http_client.put.return_value = mock_response + mock_http_client.patch.return_value = mock_response + mock_get_client.return_value = mock_http_client + + logger = AzureBlobStorageLogger() + logger.azure_auth_token = "mock-azure-ad-token" + logger.token_expiry = None + + test_payload: StandardLoggingPayload = {"id": "gov-log-id"} + + await logger.async_upload_payload_to_azure_blob_storage(test_payload) + + expected_base_url = ( + "https://test-account.dfs.core.usgovcloudapi.net/test-container/gov-log-id.json" + ) + assert mock_http_client.put.call_args[0][0] == f"{expected_base_url}?resource=file" + assert ( + mock_http_client.patch.call_args_list[0][0][0] + == f"{expected_base_url}?action=append&position=0" + ) + assert mock_http_client.patch.call_args_list[1][0][0].startswith( + f"{expected_base_url}?action=flush" + ) + + +@pytest.mark.asyncio +async def test_service_client_uses_configured_endpoint_suffix(mock_gov_env_vars): + """ + The account key path builds its own account_url; the Azure SDK derives the blob + host from it, so the suffix has to be applied here too. + """ + fake_aio_module = MagicMock() + + with patch.dict( + sys.modules, {"azure.storage.filedatalake.aio": fake_aio_module} + ): + logger = AzureBlobStorageLogger() + await logger.get_service_client() + + assert ( + fake_aio_module.DataLakeServiceClient.call_args.kwargs["account_url"] + == "https://test-account.dfs.core.usgovcloudapi.net" + ) + + +@pytest.mark.asyncio +async def test_service_client_defaults_to_commercial_endpoint(mock_env_vars): + """Unset AZURE_STORAGE_ENDPOINT_SUFFIX keeps the pre-existing commercial host""" + fake_aio_module = MagicMock() + + with patch.dict( + sys.modules, {"azure.storage.filedatalake.aio": fake_aio_module} + ): + logger = AzureBlobStorageLogger() + await logger.get_service_client() + + assert ( + fake_aio_module.DataLakeServiceClient.call_args.kwargs["account_url"] + == "https://test-account.dfs.core.windows.net" + ) diff --git a/tests/test_litellm/llms/base_llm/files/__init__.py b/tests/test_litellm/llms/base_llm/files/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/test_litellm/llms/base_llm/files/test_azure_blob_storage_backend.py b/tests/test_litellm/llms/base_llm/files/test_azure_blob_storage_backend.py new file mode 100644 index 00000000000..b924ea8f93f --- /dev/null +++ b/tests/test_litellm/llms/base_llm/files/test_azure_blob_storage_backend.py @@ -0,0 +1,226 @@ +import sys +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest + +from litellm.llms.base_llm.files.azure_blob_storage_backend import ( + AzureBlobStorageBackend, +) + +GOV_SUFFIX = "core.usgovcloudapi.net" + + +@pytest.fixture +def mock_env_vars(monkeypatch): + """Azure AD (no account key) configuration for the files backend""" + monkeypatch.setenv("AZURE_STORAGE_ACCOUNT_NAME", "test-account") + monkeypatch.setenv("AZURE_STORAGE_FILE_SYSTEM", "test-container") + monkeypatch.setenv("AZURE_STORAGE_TENANT_ID", "test-tenant-id") + monkeypatch.setenv("AZURE_STORAGE_CLIENT_ID", "test-client-id") + monkeypatch.setenv("AZURE_STORAGE_CLIENT_SECRET", "test-client-secret") + monkeypatch.delenv("AZURE_STORAGE_ACCOUNT_KEY", raising=False) + monkeypatch.delenv("AZURE_STORAGE_ENDPOINT_SUFFIX", raising=False) + + +@pytest.fixture +def mock_gov_env_vars(mock_env_vars, monkeypatch): + monkeypatch.setenv("AZURE_STORAGE_ENDPOINT_SUFFIX", GOV_SUFFIX) + + +def _make_backend() -> AzureBlobStorageBackend: + backend = AzureBlobStorageBackend() + backend.azure_auth_token = "mock-azure-ad-token" + backend.token_expiry = None + return backend + + +def _mock_upload_client() -> AsyncMock: + client = AsyncMock() + response = MagicMock() + client.put = AsyncMock(return_value=response) + client.patch = AsyncMock(return_value=response) + return client + + +@pytest.mark.parametrize( + "env_fixture, expected_suffix", + [("mock_env_vars", "core.windows.net"), ("mock_gov_env_vars", GOV_SUFFIX)], +) +@pytest.mark.asyncio +async def test_upload_file_with_azure_ad_honors_endpoint_suffix(request, env_fixture, expected_suffix): + """ + The REST upload targets the dfs host and the returned handle is a blob URL, so both + have to follow AZURE_STORAGE_ENDPOINT_SUFFIX or a sovereign-cloud account is unreachable. + """ + request.getfixturevalue(env_fixture) + client = _mock_upload_client() + + with patch( + "litellm.llms.custom_httpx.http_handler.get_async_httpx_client", + return_value=client, + ): + backend = _make_backend() + storage_url = await backend.upload_file( + file_content=b"hello", + filename="report.json", + content_type="application/json", + path_prefix="logs", + file_naming_strategy="original_filename", + ) + + expected_dfs = f"https://test-account.dfs.{expected_suffix}/test-container/logs/report.json" + assert client.put.call_args[0][0] == f"{expected_dfs}?resource=file" + assert client.patch.call_args_list[0][0][0] == f"{expected_dfs}?action=append&position=0" + assert storage_url == f"https://test-account.blob.{expected_suffix}/test-container/logs/report.json" + + +@pytest.mark.parametrize( + "env_fixture, expected_suffix", + [("mock_env_vars", "core.windows.net"), ("mock_gov_env_vars", GOV_SUFFIX)], +) +@pytest.mark.asyncio +async def test_download_file_honors_endpoint_suffix(request, env_fixture, expected_suffix): + """ + download_file both validates and splits the stored blob URL on the host, so a + sovereign-cloud URL must parse and round-trip back to the same host. + """ + request.getfixturevalue(env_fixture) + response = MagicMock() + response.content = b"file-bytes" + client = AsyncMock() + client.get = AsyncMock(return_value=response) + + storage_url = f"https://test-account.blob.{expected_suffix}/test-container/logs/report.json" + + with patch( + "litellm.llms.custom_httpx.http_handler.get_async_httpx_client", + return_value=client, + ): + backend = _make_backend() + content = await backend.download_file(storage_url) + + assert content == b"file-bytes" + assert client.get.call_args[0][0] == storage_url + + +@pytest.mark.asyncio +async def test_download_file_accepts_url_persisted_before_the_suffix_was_set(mock_gov_env_vars): + """ + storage_url is persisted in the managed files table while the suffix is process config, + so rows written before the suffix was configured must still resolve. Only the path after + the container is taken from the stored URL; the host comes from the current config. + """ + response = MagicMock() + response.content = b"file-bytes" + client = AsyncMock() + client.get = AsyncMock(return_value=response) + + with patch( + "litellm.llms.custom_httpx.http_handler.get_async_httpx_client", + return_value=client, + ): + backend = _make_backend() + content = await backend.download_file( + "https://old-account.blob.core.windows.net/old-container/logs/report.json" + ) + + assert content == b"file-bytes" + assert ( + client.get.call_args[0][0] + == f"https://test-account.blob.{GOV_SUFFIX}/test-container/logs/report.json" + ) + + +@pytest.mark.parametrize( + "storage_url", + [ + "https://example-bucket.s3.amazonaws.com/container/report.json", + "https://example.com/download?u=.blob.core.windows.net/container/report.json", + "mygovacct.blob.core.windows.net/container/report.json", + ], + ids=["other-provider", "blob-host-only-in-query", "no-scheme"], +) +@pytest.mark.asyncio +async def test_download_file_rejects_url_whose_host_is_not_an_azure_blob_host(mock_env_vars, storage_url): + """ + The host is checked on the parsed hostname, so a blob host appearing anywhere else in the + string no longer passes. No first-party producer emits these, and rejecting beats issuing a + request built from a mis-split path. + """ + client = AsyncMock() + + with patch( + "litellm.llms.custom_httpx.http_handler.get_async_httpx_client", + return_value=client, + ): + backend = _make_backend() + + with pytest.raises(ValueError, match="Invalid Azure Blob Storage URL"): + await backend.download_file(storage_url) + + client.get.assert_not_called() + + +@pytest.mark.asyncio +async def test_download_file_drops_query_string_from_the_stored_url(mock_env_vars): + """A query string on the stored URL is not part of the blob path and must not reach the request""" + response = MagicMock() + response.content = b"file-bytes" + client = AsyncMock() + client.get = AsyncMock(return_value=response) + + with patch( + "litellm.llms.custom_httpx.http_handler.get_async_httpx_client", + return_value=client, + ): + backend = _make_backend() + await backend.download_file( + "https://test-account.blob.core.windows.net/test-container/logs/report.json?sig=redacted&se=2026" + ) + + assert ( + client.get.call_args[0][0] + == "https://test-account.blob.core.windows.net/test-container/logs/report.json" + ) + + +@pytest.mark.parametrize( + "env_fixture, expected_suffix", + [("mock_env_vars", "core.windows.net"), ("mock_gov_env_vars", GOV_SUFFIX)], +) +@pytest.mark.asyncio +async def test_upload_file_with_account_key_honors_endpoint_suffix(request, env_fixture, expected_suffix, monkeypatch): + """The account key path returns its own blob URL, built independently of the REST path""" + request.getfixturevalue(env_fixture) + monkeypatch.setenv("AZURE_STORAGE_ACCOUNT_KEY", "dGVzdC1rZXk=") + + file_client = MagicMock() + 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.exists = AsyncMock(return_value=True) + file_system_client.get_directory_client = MagicMock(return_value=directory_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}): + backend = AzureBlobStorageBackend() + storage_url = await backend.upload_file( + file_content=b"hello", + filename="report.json", + content_type="application/json", + path_prefix="logs", + file_naming_strategy="original_filename", + ) + + assert storage_url == f"https://test-account.blob.{expected_suffix}/test-container/logs/report.json"