Merge remote-tracking branch 'origin/litellm_internal_staging' into litellm_/terraform-provider-dep-bump-5feb4a

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)
 )

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

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,

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"
+    )

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"