Merge branch 'main' into dependabot/uv/mypy-protobuf-5.0.0

Author: iwillspeak

examples/classify_single_example.py

@@ -69,9 +69,8 @@ async def classify_single_image_example(
             logger.info("Classifying single image...")
             correlation_id = uuid.uuid4().hex[:MAX_DEPLOYMENT_ID_LENGTH]
             logger.info("Correlation ID: %s", correlation_id)
-            result = await client.classify_single(
-                image_data, correlation_id=correlation_id
-            )
+            image_data.correlation_id = correlation_id  # Optional
+            result = await client.classify_single(image_data)
 
             # Process the result
             logger.info("Classification completed successfully!")

src/resolver_athena_client/client/athena_client.py

@@ -121,7 +121,7 @@ async def image_stream():
             yield response
 
     async def classify_single(
-        self, image_data: ImageData, correlation_id: str | None = None
+        self, image_data: ImageData
     ) -> ClassificationOutput:
         """Classify a single image synchronously without deployment context.
 
@@ -169,18 +169,19 @@ async def classify_single(
                               f"Weight: {classification.weight}")
 
         """
-        if correlation_id is None:
-            correlation_id = str(uuid.uuid4())
-
         processed_image = image_data
 
         # Apply image resizing if enabled
         if self.options.resize_images:
-            processed_image = await resize_image(processed_image)
+            processed_image = await resize_image(
+                processed_image, self.options.resampling_algorithm
+            )
 
         # Apply compression if enabled
         if self.options.compress_images:
-            processed_image = compress_image(processed_image)
+            processed_image = compress_image(
+                processed_image, self.options.compression_quality
+            )
 
         request_encoding = (
             RequestEncoding.REQUEST_ENCODING_BROTLI
@@ -196,7 +197,7 @@ async def classify_single(
 
         classification_input = ClassificationInput(
             affiliate=self.options.affiliate,
-            correlation_id=correlation_id,
+            correlation_id=processed_image.correlation_id or str(uuid.uuid4()),
             encoding=request_encoding,
             data=processed_image.data,
             format=image_format,
@@ -240,13 +241,17 @@ async def transform_image(image_data: ImageData) -> ClassificationInput:
             """Transform a single image through the full pipeline."""
             # Apply image resizing if enabled
             if self.options.resize_images:
-                resized_image = await resize_image(image_data)
+                resized_image = await resize_image(
+                    image_data, self.options.resampling_algorithm
+                )
             else:
                 resized_image = image_data
 
             # Apply compression if enabled
             if self.options.compress_images:
-                compressed_image = compress_image(resized_image)
+                compressed_image = compress_image(
+                    resized_image, self.options.compression_quality
+                )
             else:
                 compressed_image = resized_image
 
@@ -257,19 +262,23 @@ async def transform_image(image_data: ImageData) -> ClassificationInput:
                 else RequestEncoding.REQUEST_ENCODING_UNCOMPRESSED
             )
 
-            # Create classification input directly
-            correlation_provider = self.options.correlation_provider()
-
             # Ensure we never send UNSPECIFIED format over the API
             image_format = compressed_image.image_format
             if image_format == ImageFormat.IMAGE_FORMAT_UNSPECIFIED:
                 image_format = ImageFormat.IMAGE_FORMAT_RAW_UINT8_BGR
 
+            if compressed_image.correlation_id:
+                correlation_id = compressed_image.correlation_id
+            else:
+                # Create classification input directly
+                correlation_provider = self.options.correlation_provider()
+                correlation_id = correlation_provider.get_correlation_id(
+                    compressed_image.data
+                )
+
             return ClassificationInput(
                 affiliate=self.options.affiliate,
-                correlation_id=correlation_provider.get_correlation_id(
-                    compressed_image.data
-                ),
+                correlation_id=correlation_id,
                 data=compressed_image.data,
                 encoding=request_encoding,
                 format=image_format,

src/resolver_athena_client/client/athena_options.py

@@ -2,6 +2,8 @@
 
 from dataclasses import dataclass
 
+from PIL.Image import Resampling
+
 from resolver_athena_client.client.correlation import (
     CorrelationProvider,
     HashCorrelationProvider,
@@ -66,3 +68,5 @@ class AthenaOptions:
     correlation_provider: type[CorrelationProvider] = HashCorrelationProvider
     timeout: float | None = 120.0
     keepalive_interval: float | None = None
+    compression_quality: int = 11  # Brotli quality level (0-11)
+    resampling_algorithm: Resampling = Resampling.LANCZOS

src/resolver_athena_client/client/models/input_model.py

@@ -67,12 +67,16 @@ async def image_stream():
 
     """
 
-    def __init__(self, image_bytes: bytes) -> None:
+    def __init__(
+        self, image_bytes: bytes, correlation_id: None | str = None
+    ) -> None:
         """Initialize ImageData with bytes and calculate hashes.
 
         Args:
         ----
             image_bytes: The raw bytes of the image.
+            correlation_id: Optional correlation ID to associate with this
+            image data, if not provided, it will be generated by the client.
 
         """
         self.data: bytes = image_bytes
@@ -83,6 +87,7 @@ def __init__(self, image_bytes: bytes) -> None:
             hashlib.sha256(image_bytes).hexdigest()
         ]
         self.md5_hashes: list[str] = [hashlib.md5(image_bytes).hexdigest()]
+        self.correlation_id: None | str = correlation_id
 
     def add_transformation_hashes(self) -> None:
         """Add new hashes for the current data to track transformations.

src/resolver_athena_client/client/transformers/core.py

@@ -25,12 +25,17 @@ def _is_raw_bgr_expected_size(data: bytes) -> bool:
     return len(data) == _expected_raw_size
 
 
-async def resize_image(image_data: ImageData) -> ImageData:
+async def resize_image(
+    image_data: ImageData,
+    sampling_algorithm: Image.Resampling = Image.Resampling.LANCZOS,
+) -> ImageData:
     """Resize an image to expected dimensions.
 
     Args:
     ----
         image_data: The ImageData object to resize
+        sampling_algorithm: The resampling algorithm to use for resizing.
+            Defaults to LANCZOS.
 
     Returns:
     -------
@@ -53,7 +58,7 @@ def process_image() -> tuple[bytes, bool]:
             # Resize if needed
             if rgb_image.size != _target_size:
                 resized_image = rgb_image.resize(
-                    _target_size, Image.Resampling.LANCZOS
+                    _target_size, sampling_algorithm
                 )
             else:
                 resized_image = rgb_image
@@ -82,19 +87,21 @@ def process_image() -> tuple[bytes, bool]:
     return image_data
 
 
-def compress_image(image_data: ImageData) -> ImageData:
+def compress_image(image_data: ImageData, quality: int = 11) -> ImageData:
     """Compress image data using Brotli compression.
 
     Args:
     ----
         image_data: The ImageData object to compress
+        quality: Compression quality level (0-11), higher is better compression
+            but slower. Default is 11 for maximum compression.
 
     Returns:
     -------
         The same ImageData object with compressed data (modified in-place)
 
     """
-    compressed_bytes = brotli.compress(image_data.data)
+    compressed_bytes = brotli.compress(image_data.data, quality=quality)
     # Modify existing ImageData with compressed bytes but preserve hashes
     # since compression doesn't change image content
     image_data.data = compressed_bytes

tests/functional/conftest.py

@@ -92,10 +92,12 @@ def valid_formatted_image(
     tmp_path_factory: pytest.TempPathFactory,
 ) -> bytes:
     image_format = request.param
-    if (magick_path := shutil.which("magick")) is None:
+    if (magick_path := shutil.which("magick")) is None and (
+        magick_path := shutil.which("convert")
+    ) is None:
         pytest.fail(
-            "ImageMagick 'magick' command not found - cannot run "
-            "multi-format test"
+            "ImageMagick 'magick' or 'convert' command not found - cannot "
+            "run multi-format test"
         )
 
     image_dir = tmp_path_factory.mktemp("images")
@@ -119,9 +121,9 @@ def valid_formatted_image(
 
     image_path = image_dir / f"test_image.{image_format}"
     if not image_path.exists():
-        cmd = f'magick "{base_image_path}" "{image_path}"'
+        cmd = [magick_path, str(base_image_path), str(image_path)]
         _ = subprocess.run(  # noqa: S603 - false positive :(
-            [magick_path, str(base_image_path), str(image_path)],
+            cmd,
             check=True,
             shell=False,
         )

tests/test_classify_single.py

@@ -122,11 +122,13 @@ async def test_classify_single_with_correlation_id(
         return_value=mock_output
     )
 
-    # Call classify_single with custom correlation ID
-    _ = await athena_client.classify_single(
-        sample_image_data, correlation_id=custom_correlation_id
+    copied_image_data = ImageData(
+        sample_image_data.data, correlation_id=custom_correlation_id
     )
 
+    # Call classify_single with custom correlation ID
+    _ = await athena_client.classify_single(copied_image_data)
+
     # Verify correlation ID was used
     call_args = athena_client.classifier.classify_single.call_args[0][0]
     assert call_args.correlation_id == custom_correlation_id