feat: correlationid can be specified on ImageData

Author: anna-singleton-resolver

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,9 +169,6 @@ 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
@@ -196,7 +193,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,
@@ -257,19 +254,24 @@ 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/models/input_model.py

@@ -67,12 +67,14 @@ 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 +85,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.

tests/test_classify_single.py

@@ -122,10 +122,10 @@ async def test_classify_single_with_correlation_id(
         return_value=mock_output
     )
 
+    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(
-        sample_image_data, correlation_id=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]