Add workflow payload compression encoder

Author: mistral-vibe

pyproject.toml

@@ -48,6 +48,9 @@ workflow_payload_offloading = [
 workflow_payload_encryption = [
     "cryptography>=41.0.0,<47.0.0",
 ]
+workflow_payload_compression = [
+    "zstandard>=0.25.0,<0.26",
+]
 
 
 [project.urls]
@@ -69,6 +72,7 @@ dev = [
     "griffe>=1.7.3,<2",
     "authlib>=1.5.2,<2",
     "websockets >=13.0",
+    "zstandard>=0.25.0,<0.26",
 ]
 lint = [
     "ruff>=0.11.10,<0.12",

src/mistralai/extra/exceptions.py

@@ -28,6 +28,10 @@ class WorkflowPayloadEncryptionException(MistralClientException):
     """Workflow payload encryption exception"""
 
 
+class WorkflowPayloadCompressionException(MistralClientException):
+    """Workflow payload compression exception"""
+
+
 class RunException(MistralClientException):
     """Conversation run errors."""
 

src/mistralai/extra/tests/fixtures/__init__.py

@@ -0,0 +1,26 @@
+from __future__ import annotations
+
+from typing import Any
+
+
+class InMemoryBlobStorage:
+    def __init__(self) -> None:
+        self.blobs: dict[str, bytes] = {}
+
+    async def __aenter__(self) -> "InMemoryBlobStorage":
+        return self
+
+    async def __aexit__(self, *_args: Any) -> None:
+        pass
+
+    async def upload_blob(self, key: str, content: bytes) -> str:
+        self.blobs[key] = content
+        return key
+
+    async def get_blob(self, key: str) -> bytes:
+        return self.blobs[key]
+
+    async def get_blob_properties(self, key: str) -> dict[str, Any] | None:
+        if key not in self.blobs:
+            return None
+        return {"size": len(self.blobs[key]), "last_modified": "test"}

src/mistralai/extra/tests/fixtures/workflow_encoding.py

@@ -1,6 +1,7 @@
 """Tests for workflow encoding configuration lifecycle."""
 
 import gc
+import json
 
 import pytest
 from pydantic import SecretStr
@@ -12,10 +13,22 @@
     configure_workflow_encoding,
 )
 from mistralai.extra.workflows import (
-    WorkflowEncodingConfig,
+    BlobStorageConfig,
+    EncryptedStrField,
+    PayloadCompressionConfig,
     PayloadEncryptionConfig,
     PayloadEncryptionMode,
+    PayloadOffloadingConfig,
+    StorageProvider,
+    WorkflowEncodingConfig,
+    ZstdCompressionConfig,
+)
+from mistralai.extra.workflows.encoding.models import (
+    EncodedPayloadOptions,
+    WorkflowContext,
 )
+from mistralai.extra.workflows.encoding.payload_encoder import PayloadEncoder
+from mistralai.extra.tests.fixtures.workflow_encoding import InMemoryBlobStorage
 
 
 @pytest.fixture
@@ -29,7 +42,9 @@ def encryption_config() -> WorkflowEncodingConfig:
     )
 
 
-def test_payload_encoder_cleanup_on_client_gc(encryption_config: WorkflowEncodingConfig):
+def test_payload_encoder_cleanup_on_client_gc(
+    encryption_config: WorkflowEncodingConfig,
+):
     """Test that PayloadEncoder is cleaned up when client is garbage collected."""
     initial_config_count = len(_workflow_configs)
 
@@ -56,7 +71,9 @@ def test_payload_encoder_cleanup_on_client_gc(encryption_config: WorkflowEncodin
     assert len(_workflow_configs) == initial_config_count
 
 
-def test_multiple_clients_independent_configs(encryption_config: WorkflowEncodingConfig):
+def test_multiple_clients_independent_configs(
+    encryption_config: WorkflowEncodingConfig,
+):
     """Test that multiple clients have independent configs."""
     initial_config_count = len(_workflow_configs)
 
@@ -132,3 +149,189 @@ def test_reconfigure_same_client(encryption_config: WorkflowEncodingConfig):
     del client
     gc.collect()
     assert config_id not in _workflow_configs
+
+
+@pytest.mark.asyncio
+async def test_payload_encoder_compresses_network_inputs():
+    config = WorkflowEncodingConfig(
+        payload_compression=PayloadCompressionConfig(
+            min_size_bytes=1, algorithm_config=ZstdCompressionConfig(level=3)
+        )
+    )
+    encoder = PayloadEncoder(encoding_config=config)
+    payload = {"data": "x" * 20_000}
+
+    encoded = await encoder.encode_network_input(
+        payload, WorkflowContext(namespace="test", execution_id="exec")
+    )
+
+    assert encoded.encoding_options == [EncodedPayloadOptions.COMPRESSED]
+    assert encoded.encoding_metadata == {}
+
+    compressed_payload = json.loads(encoded.get_payload())
+    assert compressed_payload["algorithm_config"] == {"algorithm": "zstd", "level": 3}
+
+    decoded = await encoder.decode_network_result(encoded.model_dump(mode="json"))
+    assert decoded == payload
+
+
+@pytest.mark.asyncio
+async def test_payload_encoder_partially_encrypts_before_offloading(monkeypatch):
+    storage = InMemoryBlobStorage()
+    monkeypatch.setattr(
+        "mistralai.extra.workflows.encoding.payload_encoder.get_blob_storage",
+        lambda _: storage,
+    )
+    config = WorkflowEncodingConfig(
+        payload_encryption=PayloadEncryptionConfig(
+            mode=PayloadEncryptionMode.PARTIAL,
+            main_key=SecretStr("0" * 64),
+        ),
+        payload_offloading=PayloadOffloadingConfig(
+            min_size_bytes=1,
+            storage_config=BlobStorageConfig(
+                storage_provider=StorageProvider.S3,
+                bucket_name="test-bucket",
+            ),
+        ),
+    )
+    encoder = PayloadEncoder(encoding_config=config)
+    payload = {
+        "data": "plain value",
+        "secret": EncryptedStrField(data="secret value").model_dump(),
+    }
+
+    encoded = await encoder.encode_network_input(
+        payload, WorkflowContext(namespace="test", execution_id="exec")
+    )
+    offloaded_payload = json.loads(encoded.get_payload())
+    offloaded_bytes = storage.blobs[offloaded_payload["key"]]
+
+    assert encoded.encoding_options == [
+        EncodedPayloadOptions.PARTIALLY_ENCRYPTED,
+        EncodedPayloadOptions.OFFLOADED,
+    ]
+    assert b"plain value" in offloaded_bytes
+    assert b"secret value" not in offloaded_bytes
+
+    decoded = await encoder.decode_network_result(encoded.model_dump(mode="json"))
+    assert decoded == payload
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    ("encryption_mode", "expected_options"),
+    [
+        (
+            PayloadEncryptionMode.PARTIAL,
+            [
+                EncodedPayloadOptions.PARTIALLY_ENCRYPTED,
+                EncodedPayloadOptions.COMPRESSED,
+                EncodedPayloadOptions.OFFLOADED,
+            ],
+        ),
+        (
+            PayloadEncryptionMode.FULL,
+            [
+                EncodedPayloadOptions.COMPRESSED,
+                EncodedPayloadOptions.OFFLOADED,
+                EncodedPayloadOptions.ENCRYPTED,
+            ],
+        ),
+    ],
+)
+async def test_payload_encoder_compression_offloading_encryption_roundtrip(
+    monkeypatch,
+    encryption_mode: PayloadEncryptionMode,
+    expected_options: list[EncodedPayloadOptions],
+):
+    storage = InMemoryBlobStorage()
+    monkeypatch.setattr(
+        "mistralai.extra.workflows.encoding.payload_encoder.get_blob_storage",
+        lambda _: storage,
+    )
+    config = WorkflowEncodingConfig(
+        payload_encryption=PayloadEncryptionConfig(
+            mode=encryption_mode,
+            main_key=SecretStr("0" * 64),
+        ),
+        payload_compression=PayloadCompressionConfig(
+            min_size_bytes=1, algorithm_config=ZstdCompressionConfig(level=3)
+        ),
+        payload_offloading=PayloadOffloadingConfig(
+            min_size_bytes=1,
+            storage_config=BlobStorageConfig(
+                storage_provider=StorageProvider.S3,
+                bucket_name="test-bucket",
+            ),
+        ),
+    )
+    encoder = PayloadEncoder(encoding_config=config)
+    payload = {
+        "data": "x" * 20_000,
+        "secret": EncryptedStrField(data="secret value").model_dump(),
+    }
+
+    encoded = await encoder.encode_network_input(
+        payload, WorkflowContext(namespace="test", execution_id="exec")
+    )
+
+    assert encoded.encoding_options == expected_options
+    assert encoded.encoding_metadata == {}
+    assert len(storage.blobs) == 1
+    decoded = await encoder.decode_network_result(encoded.model_dump(mode="json"))
+    assert decoded == payload
+
+
+@pytest.mark.asyncio
+async def test_payload_encoder_does_not_partially_encrypt_when_no_marked_fields():
+    config = WorkflowEncodingConfig(
+        payload_encryption=PayloadEncryptionConfig(
+            mode=PayloadEncryptionMode.PARTIAL,
+            main_key=SecretStr("0" * 64),
+        ),
+        payload_compression=PayloadCompressionConfig(
+            min_size_bytes=1, algorithm_config=ZstdCompressionConfig(level=3)
+        ),
+    )
+    encoder = PayloadEncoder(encoding_config=config)
+    payload = {"data": "x" * 20_000}
+
+    encoded = await encoder.encode_network_input(
+        payload, WorkflowContext(namespace="test", execution_id="exec")
+    )
+
+    assert encoded.encoding_options == [EncodedPayloadOptions.COMPRESSED]
+    assert encoded.encoding_metadata == {}
+    decoded = await encoder.decode_network_result(encoded.model_dump(mode="json"))
+    assert decoded == payload
+
+
+@pytest.mark.asyncio
+async def test_payload_encoder_event_payload_orders_compression_before_full_encryption():
+    config = WorkflowEncodingConfig(
+        payload_encryption=PayloadEncryptionConfig(
+            mode=PayloadEncryptionMode.FULL,
+            main_key=SecretStr("0" * 64),
+        ),
+        payload_compression=PayloadCompressionConfig(
+            min_size_bytes=1, algorithm_config=ZstdCompressionConfig(level=3)
+        ),
+    )
+    encoder = PayloadEncoder(encoding_config=config)
+    payload = json.dumps({"data": "x" * 20_000}).encode()
+
+    (
+        encoded,
+        encoding_options,
+        encoding_metadata,
+    ) = await encoder.encode_event_payload_content(payload)
+    decoded = await encoder.decode_payload_content(
+        encoded, encoding_options, encoding_metadata
+    )
+
+    assert encoding_options == [
+        EncodedPayloadOptions.COMPRESSED,
+        EncodedPayloadOptions.ENCRYPTED,
+    ]
+    assert decoded == payload

src/mistralai/extra/tests/test_workflow_encoding.py

@@ -9,13 +9,16 @@
     WorkflowExtensions,
 )
 from .encoding import (
-    WorkflowEncodingConfig,
-    PayloadOffloadingConfig,
+    AlgorithmConfig,
+    BlobStorageConfig,
+    EncryptedStrField,
+    PayloadCompressionConfig,
     PayloadEncryptionConfig,
     PayloadEncryptionMode,
-    BlobStorageConfig,
+    PayloadOffloadingConfig,
     StorageProvider,
-    EncryptedStrField,
+    WorkflowEncodingConfig,
+    ZstdCompressionConfig,
     configure_workflow_encoding,
     generate_two_part_id,
 )
@@ -27,10 +30,13 @@
     "ConnectorSlot",
     "WorkflowExtensions",
     "execute_with_connector_auth_async",
+    "AlgorithmConfig",
     "WorkflowEncodingConfig",
     "PayloadOffloadingConfig",
     "PayloadEncryptionConfig",
     "PayloadEncryptionMode",
+    "PayloadCompressionConfig",
+    "ZstdCompressionConfig",
     "BlobStorageConfig",
     "StorageProvider",
     "EncryptedStrField",

src/mistralai/extra/workflows/__init__.py

@@ -1,20 +1,26 @@
 from .config import (
-    WorkflowEncodingConfig,
-    PayloadOffloadingConfig,
+    AlgorithmConfig,
+    BlobStorageConfig,
+    PayloadCompressionConfig,
     PayloadEncryptionConfig,
     PayloadEncryptionMode,
-    BlobStorageConfig,
+    PayloadOffloadingConfig,
     StorageProvider,
+    WorkflowEncodingConfig,
+    ZstdCompressionConfig,
 )
 from .models import EncryptedStrField
 from .payload_encoder import PayloadEncoder
 from .helpers import configure_workflow_encoding, generate_two_part_id
 
 __all__ = [
+    "AlgorithmConfig",
     "WorkflowEncodingConfig",
     "PayloadOffloadingConfig",
     "PayloadEncryptionConfig",
     "PayloadEncryptionMode",
+    "PayloadCompressionConfig",
+    "ZstdCompressionConfig",
     "BlobStorageConfig",
     "StorageProvider",
     "EncryptedStrField",

src/mistralai/extra/workflows/encoding/__init__.py

@@ -1,6 +1,7 @@
 from enum import Enum
-from pydantic import SecretStr, BaseModel
-from typing import Optional
+from typing import Annotated, Literal, Optional, Union
+
+from pydantic import BaseModel, Field, SecretStr
 
 
 class StorageProvider(str, Enum):
@@ -47,6 +48,22 @@ class PayloadEncryptionConfig(BaseModel):
     secondary_key: Optional[SecretStr] = None
 
 
+class ZstdCompressionConfig(BaseModel):
+    algorithm: Literal["zstd"] = "zstd"
+    level: int = Field(default=3, ge=1, le=22)
+
+
+AlgorithmConfig = Annotated[
+    Union[ZstdCompressionConfig], Field(discriminator="algorithm")
+]
+
+
+class PayloadCompressionConfig(BaseModel):
+    min_size_bytes: int = 1024 * 1024  # 1MB
+    algorithm_config: AlgorithmConfig = Field(default_factory=ZstdCompressionConfig)
+
+
 class WorkflowEncodingConfig(BaseModel):
     payload_offloading: PayloadOffloadingConfig | None = None
     payload_encryption: PayloadEncryptionConfig | None = None
+    payload_compression: PayloadCompressionConfig | None = None