feat: Add async transformers and batching middleware

- Implement async transformers for image resizing, compression, and
classification input conversion - Add request batching middleware for
ClassifyRequest streaming - Add tests for async transformers and
batching logic - Update dependencies: add anyio, brotli, numpy,
pytest-cov - Update .gitignore for .env files - Add coverage reporting
to CI workflow - Document async pipeline and correlation ID TODOs in
README

Author: snus-kin

.github/workflows/build_and_test.yml

@@ -72,7 +72,12 @@ jobs:
 
       - name: Run tests
         run: |
-          uv run pytest
+          uv run pytest --cov-report xml:cov.xml --cov=myproj tests/
+
+      - name: Pytest Coverage Comment
+        uses: MishaKav/pytest-coverage-comment@v1.1.56
+        with:
+          junitxml-path: ./cov.xml
 
       - name: Build package
         run: |

.gitignore

@@ -8,3 +8,6 @@ wheels/
 
 # Virtual environments
 .venv
+
+# Environment variables
+.env

README.md

@@ -3,6 +3,16 @@
 This is a Python library for interacting with the Athena API (Resolver Unknown
 CSAM Detection).
 
+## TODO
+
+### Async pipelines
+Make pipeline style invocation of the async interators such that we can
+
+async read file -> async transform -> async classify -> async results
+
+### Generate correlation IDs based on <something>
+Make do this with strategy pattern?
+
 
 ## Development
 This package uses [uv](https://docs.astral.sh/uv/) to manage its packages.

pyproject.toml

@@ -6,10 +6,11 @@ readme = "README.md"
 authors = [{ name = "Resolver Global", email = "opensource@kroll.com" }]
 requires-python = ">=3.10"
 dependencies = [
+    "anyio>=4.10.0",
+    "brotli>=1.1.0",
     "grpcio-tools>=1.74.0",
+    "numpy>=2.2.6",
     "pillow>=11.3.0",
-    "pydantic>=2.11.7",
-    "types-protobuf>=6.30.2.20250703",
 ]
 
 [dependency-groups]
@@ -19,7 +20,9 @@ dev = [
     "pyright>=1.1.403",
     "pytest>=8.4.1",
     "pytest-asyncio>=1.1.0",
+    "pytest-cov>=6.2.1",
     "ruff>=0.12.7",
+    "types-protobuf>=6.30.2.20250703",
 ]
 
 [build-system]
@@ -40,3 +43,5 @@ ignore = ["COM812", "D213", "D211", "D203", "S101"]
 
 [tool.pyright]
 exclude = ["src/athena_client/generated/*", ".venv"]
+venvPath = "."
+venv = ".venv"

src/athena_client/client/__init__.py

@@ -0,0 +1,4 @@
+"""Athena Client.
+
+This module provides a client for interacting with the Athena API.
+"""

src/athena_client/client/athena_client.py

@@ -0,0 +1,122 @@
+"""The Athena Client Class."""
+
+import types
+from collections.abc import AsyncIterator
+
+import grpc
+
+from athena_client.client.athena_options import AthenaOptions
+from athena_client.client.exceptions import AthenaError
+from athena_client.client.transformers.brotli_compressor import BrotliCompressor
+from athena_client.client.transformers.classification_input import (
+    ClassificationInputTransformer,
+)
+from athena_client.client.transformers.image_resizer import ImageResizer
+from athena_client.client.transformers.request_batcher import RequestBatcher
+from athena_client.generated.athena.athena_pb2 import (
+    ClassifyRequest,
+    ClassifyResponse,
+    RequestEncoding,
+)
+from athena_client.grpc_wrappers.classifier_service import (
+    ClassifierServiceClient,
+)
+
+
+class AthenaClient:
+    """The Athena Client Class.
+
+    This class provides coroutine methods for interacting with the
+    Athena service.
+
+    Attributes:
+        options (AthenaOptions): Configuration options for the Athena client.
+
+    """
+
+    def __init__(
+        self, channel: grpc.aio.Channel, options: AthenaOptions
+    ) -> None:
+        """Initialize the Athena Client.
+
+        Args:
+            channel (grpc.aio.Channel): The gRPC channel to use for
+                communication.
+            options (AthenaOptions): Configuration options for the
+                Athena client.
+
+        """
+        self.options = options
+        self.channel = channel
+        self.classifier = ClassifierServiceClient(self.channel)
+        self.deployment_id = options.deployment_id
+        self.max_batch_size = options.max_batch_size
+
+    async def classify_images(
+        self, images: AsyncIterator[bytes]
+    ) -> AsyncIterator[ClassifyResponse]:
+        """Classify an image using the Athena service.
+
+        Args:
+            images (AsyncIterator[Image]): The images to classify in an
+                asynchronous iterator.
+
+        Yields:
+            ClassifyResponse: The classification responses.
+
+        """
+        # Build middleware pipeline
+        image_stream = images
+
+        if self.options.resize_images:
+            image_stream = ImageResizer(image_stream)
+
+        if self.options.compress_images:
+            image_stream = BrotliCompressor(image_stream)
+
+        input_transformer = ClassificationInputTransformer(
+            image_stream,
+            deployment_id=self.options.deployment_id,
+            affiliate=self.options.affiliate,
+            request_encoding=RequestEncoding.REQUEST_ENCODING_BROTLI,
+        )
+
+        request_batcher = RequestBatcher(
+            input_transformer,
+            deployment_id=self.deployment_id,
+            max_batch_size=self.max_batch_size,
+        )
+
+        async for response in self.classifier.classify(request_batcher):
+            yield response
+
+    async def close(self) -> None:
+        """Close the client and GRPC channel."""
+        await self.channel.close()
+
+    async def __aenter__(self) -> "AthenaClient":
+        """Context manager entrypoint.
+
+        Registers the client with the server by sending a
+        classification request. With a deployment ID.
+        """
+
+        # Register the client with the server, just send one deployment ID.
+        async def init_request() -> AsyncIterator[ClassifyRequest]:
+            yield ClassifyRequest(deployment_id=self.deployment_id)
+
+        responses = self.classifier.classify(init_request())
+        async for response in responses:
+            if response.global_error:
+                raise AthenaError(response.global_error.message)
+
+        return self
+
+    async def __aexit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc_val: BaseException | None,
+        exc_tb: types.TracebackType | None,
+    ) -> None:
+        """Context manager exit."""
+        await self.close()

src/athena_client/client/athena_options.py

@@ -0,0 +1,16 @@
+"""Options object for the Athena client."""
+
+from dataclasses import dataclass
+
+
+@dataclass
+class AthenaOptions:
+    """Options for the Athena client."""
+
+    host: str = "localhost"
+    resize_images: bool = False
+    compress_images: bool = True
+    deployment_id: str = "default"
+    affiliate: str = "default"
+    correlation_id: str = "default"
+    max_batch_size: int = 100

src/athena_client/client/channel.py

@@ -0,0 +1,26 @@
+"""Channel creation utilities for the Athena client."""
+
+import grpc
+from grpc.aio import Channel
+
+
+def create_channel(host: str, auth_token: str) -> Channel:
+    """Create a gRPC channel with optional authentication.
+
+    Args:
+        host: The host address to connect to
+        auth_token: Optional authentication token. If provided, creates a secure
+            channel with token-based authentication. If not provided, creates an
+            insecure channel.
+
+    Returns:
+        A gRPC channel (either secure or insecure)
+
+    """
+    # Create credentials with token authentication
+    credentials = grpc.composite_channel_credentials(
+        grpc.ssl_channel_credentials(),
+        grpc.access_token_call_credentials(auth_token),
+    )
+
+    return grpc.aio.secure_channel(host, credentials)

src/athena_client/client/consts.py

@@ -0,0 +1,5 @@
+"""Constants for the athena_client's image properties."""
+
+# Athena's classifier expects images to be 448x448 pixels.
+EXPECTED_WIDTH = 448
+EXPECTED_HEIGHT = 448

src/athena_client/client/deployment_selector.py

@@ -0,0 +1,61 @@
+"""Deployment selector for the Athena client."""
+
+import logging
+
+import grpc
+
+from athena_client.generated.athena.athena_pb2 import ListDeploymentsResponse
+from athena_client.grpc_wrappers.classifier_service import (
+    ClassifierServiceClient,
+)
+
+
+class DeploymentSelector:
+    """A controller for selecting deployments from the Athena service.
+
+    This class provides functionality to list available deployments for use
+    with the Athena client.
+
+    Attributes:
+        classifier (ClassifierServiceClient): The classifier service client used
+            to communicate with the Athena service.
+
+    """
+
+    def __init__(self, channel: grpc.aio.Channel) -> None:
+        """Initialize the deployment selector.
+
+        Args:
+            channel (grpc.aio.Channel): Channel with which to communicate with
+                the Athena service.
+
+        """
+        self.logger = logging.getLogger(__name__)
+        self.classifier = ClassifierServiceClient(channel)
+
+    async def list_deployments(self) -> ListDeploymentsResponse:
+        """Retrieve a list of all active deployments.
+
+        Returns:
+            ListDeploymentsResponse: Response containing the list of
+                deployments.
+
+        """
+        self.logger.debug("Retrieving list of deployments from server")
+        response = await self.classifier.list_deployments()
+
+        if not response.deployments:
+            self.logger.error("No deployments available from server")
+        else:
+            self.logger.debug(
+                "Retrieved %d deployments: %s",
+                len(response.deployments),
+                ", ".join(
+                    [
+                        deployment.deployment_id
+                        for deployment in response.deployments
+                    ]
+                ),
+            )
+
+        return response