From 36bb10707cb7cc5fb9fc38205e178e4a6b9739cc Mon Sep 17 00:00:00 2001 From: Akhil Nair Date: Fri, 1 May 2026 14:05:23 +0100 Subject: [PATCH 1/2] feat: add Azure DefaultAzureCredential support for blob storage --- pyproject.toml | 1 + .../extra/workflows/encoding/config.py | 1 + .../workflows/encoding/storage/_azure.py | 28 ++++++++++++++++--- .../encoding/storage/blob_storage.py | 9 ++---- 4 files changed, 29 insertions(+), 10 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 674de675..5862bbeb 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -32,6 +32,7 @@ realtime = [ workflow_payload_offloading_azure = [ "azure-storage-blob[aio]>=12.28.0,<13.0.0", + "azure-identity[aio]>=1.25.0,<2.0.0", ] workflow_payload_offloading_gcs = [ "gcloud-aio-storage>=9.3.0,<10.0.0", diff --git a/src/mistralai/extra/workflows/encoding/config.py b/src/mistralai/extra/workflows/encoding/config.py index ff849fd7..af02f64a 100644 --- a/src/mistralai/extra/workflows/encoding/config.py +++ b/src/mistralai/extra/workflows/encoding/config.py @@ -16,6 +16,7 @@ class BlobStorageConfig(BaseModel): # Azure settings container_name: Optional[str] = None azure_connection_string: Optional[SecretStr] = None + azure_storage_account_url: Optional[str] = None # GCS settings bucket_id: Optional[str] = None diff --git a/src/mistralai/extra/workflows/encoding/storage/_azure.py b/src/mistralai/extra/workflows/encoding/storage/_azure.py index e62d9926..fa822ee6 100644 --- a/src/mistralai/extra/workflows/encoding/storage/_azure.py +++ b/src/mistralai/extra/workflows/encoding/storage/_azure.py @@ -3,6 +3,7 @@ from typing import Any, cast from azure.core.exceptions import ResourceNotFoundError +from azure.identity.aio import DefaultAzureCredential from azure.storage.blob.aio import BlobServiceClient from .blob_storage import BlobNotFoundError, BlobStorage @@ -11,14 +12,25 @@ class AzureBlobStorage(BlobStorage): def __init__( self, container_name: str, - azure_connection_string: str, + azure_connection_string: str | None = None, prefix: str | None = None, + azure_storage_account_url: str | None = None, ): + if azure_connection_string and azure_storage_account_url: + raise ValueError( + "azure_connection_string and azure_storage_account_url are mutually exclusive" + ) + if not azure_connection_string and not azure_storage_account_url: + raise ValueError( + "Either azure_connection_string or azure_storage_account_url must be provided" + ) self.container_name = container_name self.connection_string = azure_connection_string + self.account_url = azure_storage_account_url self.prefix = prefix or "" self._service_client: BlobServiceClient | None = None self._container_client: Any = None + self._credential: Any = None def _get_full_key(self, key: str) -> str: if not self.prefix: @@ -28,9 +40,15 @@ def _get_full_key(self, key: str) -> str: return f"{self.prefix}/{key}" async def __aenter__(self) -> "AzureBlobStorage": - self._service_client = BlobServiceClient.from_connection_string( - self.connection_string - ) + if self.connection_string: + self._service_client = BlobServiceClient.from_connection_string( + self.connection_string + ) + else: + self._credential = DefaultAzureCredential() + self._service_client = BlobServiceClient( + self.account_url, credential=self._credential + ) assert self._service_client is not None self._container_client = self._service_client.get_container_client( self.container_name @@ -40,6 +58,8 @@ async def __aenter__(self) -> "AzureBlobStorage": async def __aexit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None: if self._service_client: await self._service_client.close() + if self._credential: + await self._credential.close() async def upload_blob(self, key: str, content: bytes) -> str: full_key = self._get_full_key(key) diff --git a/src/mistralai/extra/workflows/encoding/storage/blob_storage.py b/src/mistralai/extra/workflows/encoding/storage/blob_storage.py index ce488421..e267b979 100644 --- a/src/mistralai/extra/workflows/encoding/storage/blob_storage.py +++ b/src/mistralai/extra/workflows/encoding/storage/blob_storage.py @@ -65,8 +65,8 @@ async def get_blob_storage( from ._azure import AzureBlobStorage # type: ignore[import-untyped] except ImportError as e: raise ImportError( - "Azure Blob Storage support requires azure-storage-blob. " - "Install it with: pip install 'mistralai[workflow_payload_offloading_azure]'" + "Azure Blob Storage support requires azure-storage-blob and azure-identity. " + "Install with: pip install 'mistralai[workflow_payload_offloading_azure]'" ) from e if not blob_storage_config.container_name: @@ -78,14 +78,11 @@ async def get_blob_storage( if blob_storage_config.azure_connection_string else None ) - if not azure_conn_str: - raise WorkflowPayloadOffloadingException( - "azure_connection_string is required for Azure blob storage" - ) storage = AzureBlobStorage( container_name=blob_storage_config.container_name, azure_connection_string=azure_conn_str, prefix=prefix, + azure_storage_account_url=blob_storage_config.azure_storage_account_url, ) elif blob_storage_config.storage_provider == StorageProvider.GCS: From 520972d822f5d4cfa6dbe9e7147bc811731a06e7 Mon Sep 17 00:00:00 2001 From: Akhil Nair Date: Fri, 1 May 2026 14:22:21 +0100 Subject: [PATCH 2/2] fix: narrow account_url type for mypy/pyright --- .../workflows/encoding/storage/_azure.py | 1 + uv.lock | 45 +++++++++++++++++++ 2 files changed, 46 insertions(+) diff --git a/src/mistralai/extra/workflows/encoding/storage/_azure.py b/src/mistralai/extra/workflows/encoding/storage/_azure.py index fa822ee6..8865634c 100644 --- a/src/mistralai/extra/workflows/encoding/storage/_azure.py +++ b/src/mistralai/extra/workflows/encoding/storage/_azure.py @@ -45,6 +45,7 @@ async def __aenter__(self) -> "AzureBlobStorage": self.connection_string ) else: + assert self.account_url is not None self._credential = DefaultAzureCredential() self._service_client = BlobServiceClient( self.account_url, credential=self._credential diff --git a/uv.lock b/uv.lock index 8f07ac6b..d0ba0c4a 100644 --- a/uv.lock +++ b/uv.lock @@ -282,6 +282,22 @@ aio = [ { name = "aiohttp" }, ] +[[package]] +name = "azure-identity" +version = "1.25.3" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "azure-core" }, + { name = "cryptography" }, + { name = "msal" }, + { name = "msal-extensions" }, + { name = "typing-extensions" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/c5/0e/3a63efb48aa4a5ae2cfca61ee152fbcb668092134d3eb8bfda472dd5c617/azure_identity-1.25.3.tar.gz", hash = "sha256:ab23c0d63015f50b630ef6c6cf395e7262f439ce06e5d07a64e874c724f8d9e6", size = 286304, upload-time = "2026-03-13T01:12:20.892Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/49/9a/417b3a533e01953a7c618884df2cb05a71e7b68bdbce4fbdb62349d2a2e8/azure_identity-1.25.3-py3-none-any.whl", hash = "sha256:f4d0b956a8146f30333e071374171f3cfa7bdb8073adb8c3814b65567aa7447c", size = 192138, upload-time = "2026-03-13T01:12:22.951Z" }, +] + [[package]] name = "azure-storage-blob" version = "12.28.0" @@ -1046,10 +1062,12 @@ workflow-payload-encryption = [ ] workflow-payload-offloading = [ { name = "aioboto3" }, + { name = "azure-identity" }, { name = "azure-storage-blob", extra = ["aio"] }, { name = "gcloud-aio-storage" }, ] workflow-payload-offloading-azure = [ + { name = "azure-identity" }, { name = "azure-storage-blob", extra = ["aio"] }, ] workflow-payload-offloading-gcs = [ @@ -1086,6 +1104,7 @@ lint = [ requires-dist = [ { name = "aioboto3", marker = "extra == 'workflow-payload-offloading-s3'", specifier = ">=12.4.0,<13.0.0" }, { name = "authlib", marker = "extra == 'agents'", specifier = ">=1.5.2,<2.0" }, + { name = "azure-identity", extras = ["aio"], marker = "extra == 'workflow-payload-offloading-azure'", specifier = ">=1.25.0,<2.0.0" }, { name = "azure-storage-blob", extras = ["aio"], marker = "extra == 'workflow-payload-offloading-azure'", specifier = ">=12.28.0,<13.0.0" }, { name = "cryptography", marker = "extra == 'workflow-payload-encryption'", specifier = ">=41.0.0,<47.0.0" }, { name = "eval-type-backport", specifier = ">=0.2.0" }, @@ -1131,6 +1150,32 @@ lint = [ { name = "ruff", specifier = ">=0.11.10,<0.12" }, ] +[[package]] +name = "msal" +version = "1.36.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "cryptography" }, + { name = "pyjwt", extra = ["crypto"] }, + { name = "requests" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/de/cb/b02b0f748ac668922364ccb3c3bff5b71628a05f5adfec2ba2a5c3031483/msal-1.36.0.tar.gz", hash = "sha256:3f6a4af2b036b476a4215111c4297b4e6e236ed186cd804faefba23e4990978b", size = 174217, upload-time = "2026-04-09T10:20:33.525Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/2a/d3/414d1f0a5f6f4fe5313c2b002c54e78a3332970feb3f5fed14237aa17064/msal-1.36.0-py3-none-any.whl", hash = "sha256:36ecac30e2ff4322d956029aabce3c82301c29f0acb1ad89b94edcabb0e58ec4", size = 121547, upload-time = "2026-04-09T10:20:32.336Z" }, +] + +[[package]] +name = "msal-extensions" +version = "1.3.1" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "msal" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/01/99/5d239b6156eddf761a636bded1118414d161bd6b7b37a9335549ed159396/msal_extensions-1.3.1.tar.gz", hash = "sha256:c5b0fd10f65ef62b5f1d62f4251d51cbcaf003fcedae8c91b040a488614be1a4", size = 23315, upload-time = "2025-03-14T23:51:03.902Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/5e/75/bd9b7bb966668920f06b200e84454c8f3566b102183bc55c5473d96cb2b9/msal_extensions-1.3.1-py3-none-any.whl", hash = "sha256:96d3de4d034504e969ac5e85bae8106c8373b5c6568e4c8fa7af2eca9dbe6bca", size = 20583, upload-time = "2025-03-14T23:51:03.016Z" }, +] + [[package]] name = "multidict" version = "6.7.1"