diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 4fe732f..7738e51 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -52,31 +52,11 @@ jobs: - name: Install the project run: uv sync --locked --all-extras --dev - - name: Ensure no differences in generated code - if: runner.os != 'Windows' + - name: Generate protobuf code shell: bash run: | source .venv/bin/activate - GENERATED_DIR="src/resolver_athena_client/generated" - BACKUP_DIR="src/resolver_athena_client/generated_backup" - - cp -r $GENERATED_DIR $BACKUP_DIR - - ./scripts/compile_proto.sh || (echo "Protobuf compilation failed. Ensure submodules are initialized and the proto file exists." && exit 1) - - # Fix imports in generated files - if [[ -f "$GENERATED_DIR/athena/athena_pb2_grpc.py" && -f "$GENERATED_DIR/athena/athena_pb2.py" ]]; then - sed -i.bak 's/^from athena /from resolver_athena_client.generated.athena /' "$GENERATED_DIR/athena/athena_pb2_grpc.py" - sed -i.bak 's/^from athena /from resolver_athena_client.generated.athena /' "$GENERATED_DIR/athena/athena_pb2.py" - rm -f "$GENERATED_DIR/athena/athena_pb2_grpc.py.bak" "$GENERATED_DIR/athena/athena_pb2.py.bak" - else - echo "Error: Expected files not found in $GENERATED_DIR/athena" - exit 1 - fi - - diff -r $GENERATED_DIR $BACKUP_DIR || (echo "Generated code differs. Please commit the changes after running compile_proto.sh." && exit 1) - - rm -rf $BACKUP_DIR + ./scripts/compile_proto.sh || (echo "Protobuf compilation failed. Ensure submodules are initialized and proto files exist." && exit 1) - name: Run linter run: | @@ -123,6 +103,17 @@ jobs: with: enable-cache: true + - name: Set up Python + run: uv python install + + - name: Install the project + run: uv sync --locked --all-extras --dev + + - name: Generate protobuf code + run: | + source .venv/bin/activate + ./scripts/compile_proto.sh || (echo "Protobuf compilation failed. Ensure submodules are initialized and proto files exist." && exit 1) + - name: Set version from tag run: | if [[ "$GITHUB_REF" == refs/tags/* ]]; then diff --git a/.gitignore b/.gitignore index cb0c4f4..a91126d 100644 --- a/.gitignore +++ b/.gitignore @@ -11,3 +11,6 @@ wheels/ .env docs/_build/ + +# Generated protobuf code +src/resolver_athena_client/generated/ diff --git a/AGENTS.md b/AGENTS.md index c746526..94ab69c 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -8,7 +8,7 @@ - Format code: `ruff format` - Lint code: `ruff check` - Install git hooks: `pre-commit install` -- Compile protobufs: `bash scripts/compile_proto.sh` (run from root) +- Compile protobufs: `bash scripts/compile_proto.sh` (run from root, required for local development) ## Code style - Use Python type hints throughout @@ -30,7 +30,7 @@ ## PR instructions - Title format: [component] Description -- Run `ruff check`, `pyright`, and `pytest` before committing +- Run `bash scripts/compile_proto.sh`, `ruff check`, `pyright`, and `pytest` before committing - Keep PRs focused on a single change - Add tests for new functionality - Update documentation for API changes @@ -40,7 +40,8 @@ - Use `uv` package manager instead of pip - Don't use `uv pip` commands, just the base `uv` commands - Run formatters before committing -- Check generated code in `src/athena_client/generated/` +- Generated code is built automatically in CI, but run `bash scripts/compile_proto.sh` locally for development +- Generated code in `src/resolver_athena_client/generated/` is not committed to the repo - Add error handling at each pipeline stage - Use correlation IDs for request tracing diff --git a/athena-protobufs b/athena-protobufs index ca80312..7f69cdf 160000 --- a/athena-protobufs +++ b/athena-protobufs @@ -1 +1 @@ -Subproject commit ca8031229a4336df41199bd48cc7db15e92e3156 +Subproject commit 7f69cdfb55bc761e18a5e244da034ec361f265dd diff --git a/examples/classify_single_example.py b/examples/classify_single_example.py new file mode 100755 index 0000000..6d432e4 --- /dev/null +++ b/examples/classify_single_example.py @@ -0,0 +1,271 @@ +#!/usr/bin/env python3 +"""Example script demonstrating the classify_single method.""" + +import asyncio +import logging +import os +import sys +import uuid +from pathlib import Path + +from create_image import create_test_image +from dotenv import load_dotenv + +from resolver_athena_client.client.athena_client import AthenaClient +from resolver_athena_client.client.athena_options import AthenaOptions +from resolver_athena_client.client.channel import ( + CredentialHelper, + create_channel_with_credentials, +) +from resolver_athena_client.client.models import ImageData + + +async def classify_single_image_example( + logger: logging.Logger, + options: AthenaOptions, + credential_helper: CredentialHelper, + image_path: str | None = None, +) -> bool: + """Demonstrate single image classification. + + Args: + logger: Logger instance for output + options: Configuration options for the Athena client + credential_helper: OAuth credential helper for authentication + image_path: Path to image file to classify (optional) + + Returns: + True if classification was successful, False otherwise + + """ + # Create gRPC channel with credentials + channel = await create_channel_with_credentials( + options.host, credential_helper + ) + + async with AthenaClient(channel, options) as client: + # Load image data + if image_path and Path(image_path).exists(): + logger.info("Loading image from: %s", image_path) + image_bytes = Path(image_path).read_bytes() + else: + # Create a simple test image if no path provided + logger.info("Creating synthetic test image") + image_bytes = create_test_image() + + # Create ImageData object + image_data = ImageData(image_bytes) + logger.info( + "Image loaded: %d bytes, MD5: %s", + len(image_data.data), + image_data.md5_hashes[0][:8] + "...", + ) + + try: + # Classify the single image + logger.info("Classifying single image...") + correlation_id = uuid.uuid4().hex[:63] + logger.info("Correlation ID: %s", correlation_id) + result = await client.classify_single( + image_data, correlation_id=correlation_id + ) + + # Process the result + logger.info("Classification completed successfully!") + + if result.error.code: + logger.error( + "Classification error: %s (%s)", + result.error.message, + result.error.code, + ) + if result.error.details: + logger.error("Error details: %s", result.error.details) + return False + + if result.classifications: + logger.info( + "Found %d classifications:", len(result.classifications) + ) + for i, classification in enumerate(result.classifications, 1): + logger.info( + " %d. Label: %s, Weight: %.3f", + i, + classification.label, + classification.weight, + ) + else: + logger.info("No classifications found for this image") + + except Exception: + logger.exception("Error during single image classification") + return False + else: + return True + + +async def classify_multiple_single_images_example( + logger: logging.Logger, + options: AthenaOptions, + credential_helper: CredentialHelper, + num_images: int = 3, +) -> int: + """Demonstrate classifying multiple images individually. + + This shows how classify_single can be used for multiple images + when you want individual control over each classification request. + + Args: + logger: Logger instance for output + options: Configuration options for the Athena client + credential_helper: OAuth credential helper for authentication + num_images: Number of test images to classify + + Returns: + Number of successfully classified images + + """ + # Create gRPC channel with credentials + channel = await create_channel_with_credentials( + options.host, credential_helper + ) + + successful_count = 0 + + async with AthenaClient(channel, options) as client: + logger.info("Classifying %d images individually...", num_images) + + for i in range(num_images): + try: + # Create a unique test image for each iteration + image_bytes = create_test_image(seed=i) + image_data = ImageData(image_bytes) + + # Classify with auto-generated correlation ID + result = await client.classify_single(image_data) + + logger.info( + "Image %d/%d - Correlation: %s", + i + 1, + num_images, + result.correlation_id[:8] + "...", + ) + + if result.error.code: + logger.warning( + "Image %d failed: %s", i + 1, result.error.message + ) + elif result.classifications: + top_classification = max( + result.classifications, key=lambda c: c.weight + ) + logger.info( + "Image %d - Top result: %s (%.3f)", + i + 1, + top_classification.label, + top_classification.weight, + ) + successful_count += 1 + else: + logger.info("Image %d - No classifications", i + 1) + successful_count += 1 + + except Exception: # noqa: PERF203 + logger.exception("Failed to classify image %d", i + 1) + + logger.info( + "Completed: %d/%d images classified successfully", + successful_count, + num_images, + ) + return successful_count + + +async def main() -> int: + """Run the classify_single examples.""" + logger = logging.getLogger(__name__) + load_dotenv() + + # OAuth credentials from environment + client_id = os.getenv("OAUTH_CLIENT_ID") + client_secret = os.getenv("OAUTH_CLIENT_SECRET") + auth_url = os.getenv( + "OAUTH_AUTH_URL", "https://crispthinking.auth0.com/oauth/token" + ) + audience = os.getenv("OAUTH_AUDIENCE", "crisp-athena-dev") + + if not client_id or not client_secret: + logger.error("OAUTH_CLIENT_ID and OAUTH_CLIENT_SECRET must be set") + return 1 + + host = os.getenv("ATHENA_HOST", "localhost") + logger.info("Connecting to %s", host) + + # Create credential helper + credential_helper = CredentialHelper( + client_id=client_id, + client_secret=client_secret, + auth_url=auth_url, + audience=audience, + ) + + # Test token acquisition + try: + logger.info("Acquiring OAuth token...") + token = await credential_helper.get_token() + logger.info("Successfully acquired token (length: %d)", len(token)) + except Exception: + logger.exception("Failed to acquire OAuth token") + return 1 + + # Configure client options + options = AthenaOptions( + host=host, + resize_images=True, + compress_images=True, + timeout=30.0, # Shorter timeout for single requests + affiliate="Crisp", + deployment_id="single-example-deployment", # Not used + ) + + try: + # Example 1: Classify a single image + logger.info("\n=== Example 1: Single Image Classification ===") + success = await classify_single_image_example( + logger, + options, + credential_helper, + image_path=os.getenv("TEST_IMAGE_PATH"), # Optional image path + ) + + if not success: + logger.error("Single image classification failed") + return 1 + + # Example 2: Classify multiple images individually + logger.info("\n=== Example 2: Multiple Individual Classifications ===") + successful_count = await classify_multiple_single_images_example( + logger, options, credential_helper, num_images=5 + ) + + if successful_count == 0: + logger.error("No images were successfully classified") + return 1 + + logger.info("\n=== All examples completed successfully! ===") + + except Exception: + logger.exception("Examples failed") + return 1 + else: + return 0 + + +if __name__ == "__main__": + logging.basicConfig( + level=logging.INFO, + format="%(asctime)s.%(msecs)03d %(levelname)s: %(message)s", + datefmt="%H:%M:%S", + ) + + sys.exit(asyncio.run(main())) diff --git a/examples/create_image.py b/examples/create_image.py index 06b75b0..1991ba3 100644 --- a/examples/create_image.py +++ b/examples/create_image.py @@ -139,3 +139,23 @@ async def iter_images( counter[0] += 1 yield ImageData(img_bytes) count += 1 + + +def create_test_image( + width: int = 160, height: int = 120, seed: int | None = None +) -> bytes: + """Create a test image with specified dimensions and optional seed. + + Args: + width: Width of the test image in pixels (default: 160) + height: Height of the test image in pixels (default: 120) + seed: Optional seed for reproducible image generation + + Returns: + PNG image bytes + + """ + if seed is not None: + _rng.seed(seed) + + return create_random_image(width, height) diff --git a/pyproject.toml b/pyproject.toml index fdcfd8f..d517a48 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,5 +1,5 @@ [project] -name = "resolver-athena-client" +name = "resolver_athena_client" version = "0.0.0" description = "Python client library for Athena API - CSAM detection and content classification" readme = "README.md" @@ -58,7 +58,7 @@ build-backend = "uv_build" [tool.ruff] line-length = 80 -exclude = ["src/resolver_athena_client/generated/*"] +exclude = ["src/resolver_athena_client/generated/*", "docs"] [tool.ruff.lint] select = ["ALL"] diff --git a/scripts/compile_proto.sh b/scripts/compile_proto.sh index 1cec6de..f6c9f91 100755 --- a/scripts/compile_proto.sh +++ b/scripts/compile_proto.sh @@ -7,30 +7,50 @@ OUT_DIR="src/resolver_athena_client/generated" # Ensure output directory exists mkdir -p "$OUT_DIR" +# Create __init__.py files for Python packages +touch "$OUT_DIR/__init__.py" +mkdir -p "$OUT_DIR/athena" +touch "$OUT_DIR/athena/__init__.py" + # Compile .proto files -if [ ! -f "$PROTO_DIR/athena/athena.proto" ]; then - echo "Error: Protobuf file not found at $PROTO_DIR/athena/athena.proto. Ensure the submodule is initialized and updated." +if [ ! -d "$PROTO_DIR/athena" ]; then + echo "Error: Athena protobuf directory not found at $PROTO_DIR/athena. Ensure the submodule is initialized and updated." exit 1 fi -python -m grpc_tools.protoc \ +# Find all .proto files in the athena directory +PROTO_FILES=$(find "$PROTO_DIR/athena" -name "*.proto") +if [ -z "$PROTO_FILES" ]; then + echo "Error: No .proto files found in $PROTO_DIR/athena." + exit 1 +fi + +echo "Found proto files: $PROTO_FILES" + +uv run python -m grpc_tools.protoc \ --proto_path="$PROTO_DIR" \ - --proto_path="$(python -c "import grpc_tools; print(grpc_tools.__path__[0])")" \ --python_out="$OUT_DIR" \ --grpc_python_out="$OUT_DIR" \ --mypy_out="$OUT_DIR" \ - "$PROTO_DIR/athena/athena.proto" + $PROTO_FILES # Fix imports in generated files -if [[ "$OSTYPE" == "darwin"* ]]; then - sed -i '' -e 's/^from athena /from resolver_athena_client.generated.athena /' "$OUT_DIR/athena/athena_pb2_grpc.py" - sed -i '' -e 's/^from athena /from resolver_athena_client.generated.athena /' "$OUT_DIR/athena/athena_pb2.py" +GENERATED_FILES=$(find "$OUT_DIR/athena" -name "*_pb2*.py" 2>/dev/null) +if [ -n "$GENERATED_FILES" ]; then + for file in $GENERATED_FILES; do + echo "Fixing imports in $file" + if [[ "$OSTYPE" == "darwin"* ]]; then + sed -i '' -e 's/^from athena /from resolver_athena_client.generated.athena /' "$file" + else + sed -i -e 's/^from athena /from resolver_athena_client.generated.athena /' "$file" + fi + done else - sed -i -e 's/^from athena /from resolver_athena_client.generated.athena /' "$OUT_DIR/athena/athena_pb2_grpc.py" - sed -i -e 's/^from athena /from resolver_athena_client.generated.athena /' "$OUT_DIR/athena/athena_pb2.py" + echo "Warning: No generated Python files found to fix imports." fi -if [ ! -f "$OUT_DIR/athena/athena_pb2.py" ] || [ ! -f "$OUT_DIR/athena/athena_pb2_grpc.py" ]; then +# Verify at least some files were generated +if [ ! -d "$OUT_DIR/athena" ] || [ -z "$(find "$OUT_DIR/athena" -name "*_pb2*.py" 2>/dev/null)" ]; then echo "Error: Protobuf files were not generated successfully in $OUT_DIR." exit 1 fi diff --git a/src/resolver_athena_client/client/athena_client.py b/src/resolver_athena_client/client/athena_client.py index 2c82897..38fc85c 100644 --- a/src/resolver_athena_client/client/athena_client.py +++ b/src/resolver_athena_client/client/athena_client.py @@ -3,6 +3,7 @@ import asyncio import logging import types +import uuid from collections.abc import AsyncGenerator, AsyncIterator import grpc @@ -16,15 +17,24 @@ from resolver_athena_client.client.transformers.classification_input import ( ClassificationInputTransformer, ) +from resolver_athena_client.client.transformers.core import ( + compress_image, + resize_image, +) from resolver_athena_client.client.transformers.image_resizer import ( ImageResizer, ) from resolver_athena_client.client.transformers.request_batcher import ( RequestBatcher, ) -from resolver_athena_client.generated.athena.athena_pb2 import ( +from resolver_athena_client.generated.athena.models_pb2 import ( + ClassificationInput, + ClassificationOutput, ClassifyRequest, ClassifyResponse, + HashType, + ImageFormat, + ImageHash, RequestEncoding, ) from resolver_athena_client.grpc_wrappers.classifier_service import ( @@ -155,6 +165,101 @@ async def image_stream(): max_reconnects, ) + async def classify_single( + self, image_data: ImageData, correlation_id: str | None = None + ) -> ClassificationOutput: + """Classify a single image synchronously without deployment context. + + This method provides immediate, synchronous classification results for + single images without requiring deployment coordination, session + management, or streaming setup. It's ideal for: + + - Low-throughput, low-latency classification scenarios + - Simple one-off image classifications + - Applications where immediate responses are preferred over streaming + - Testing and debugging individual image classifications + + Args: + image_data: ImageData object containing image bytes and metadata. + The image will be processed through the same transformation + pipeline as the streaming classify method (resize, compression) + based on client options. + correlation_id: Optional unique identifier for correlating this + request. If not provided, a UUID will be generated + automatically. + + Returns: + ClassificationOutput containing either classification results or + error information for the single image. + + Raises: + AthenaError: If the service returns an error. + grpc.aio.AioRpcError: For gRPC communication errors. + + Example: + # Create ImageData from raw bytes + image_data = ImageData(image_bytes) + + async with AthenaClient(channel, options) as client: + result = await client.classify_single(image_data) + if result.error: + print(f"Classification error: {result.error.message}") + else: + for classification in result.classifications: + print(f"Label: {classification.label}, " + 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) + + # Apply compression if enabled + if self.options.compress_images: + processed_image = compress_image(processed_image) + + request_encoding = ( + RequestEncoding.REQUEST_ENCODING_BROTLI + if self.options.compress_images + else RequestEncoding.REQUEST_ENCODING_UNCOMPRESSED + ) + + classification_input = ClassificationInput( + affiliate=self.options.affiliate, + correlation_id=correlation_id, + encoding=request_encoding, + data=processed_image.data, + format=ImageFormat.IMAGE_FORMAT_RAW_UINT8, + hashes=[ + ImageHash( + value=hash_value, + type=HashType.HASH_TYPE_MD5, + ) + for hash_value in processed_image.md5_hashes + ], + ) + + try: + result = await self.classifier.classify_single( + classification_input, timeout=self.options.timeout + ) + except grpc.aio.AioRpcError: + self.logger.exception( + "gRPC error in classify_single", + ) + raise + + # Check for errors in the response + if result.error and result.error.message: + self._raise_athena_error(result.error.message) + + return result + def _create_request_pipeline( self, images: AsyncIterator[ImageData] ) -> RequestBatcher: diff --git a/src/resolver_athena_client/client/correlation.py b/src/resolver_athena_client/client/correlation.py index c977568..b874636 100644 --- a/src/resolver_athena_client/client/correlation.py +++ b/src/resolver_athena_client/client/correlation.py @@ -69,7 +69,7 @@ def get_correlation_id(self, input_data: bytes | str | bytearray) -> str: else: data_bytes = str(input_data).encode("utf-8") - return hashlib.sha256(data_bytes).hexdigest() + return hashlib.sha256(data_bytes).hexdigest()[:63] except Exception as e: error_msg = f"Failed to generate correlation ID from input: {e}" raise ValueError(error_msg) from e diff --git a/src/resolver_athena_client/client/deployment_selector.py b/src/resolver_athena_client/client/deployment_selector.py index 020f7bd..7f6a983 100644 --- a/src/resolver_athena_client/client/deployment_selector.py +++ b/src/resolver_athena_client/client/deployment_selector.py @@ -5,7 +5,7 @@ import grpc -from resolver_athena_client.generated.athena.athena_pb2 import ( +from resolver_athena_client.generated.athena.models_pb2 import ( ListDeploymentsResponse, ) from resolver_athena_client.grpc_wrappers.classifier_service import ( diff --git a/src/resolver_athena_client/client/exceptions.py b/src/resolver_athena_client/client/exceptions.py index 0066414..5ee8957 100644 --- a/src/resolver_athena_client/client/exceptions.py +++ b/src/resolver_athena_client/client/exceptions.py @@ -3,7 +3,7 @@ from typing import TYPE_CHECKING if TYPE_CHECKING: - from resolver_athena_client.generated.athena.athena_pb2 import ( + from resolver_athena_client.generated.athena.models_pb2 import ( ClassificationError, ) diff --git a/src/resolver_athena_client/client/transformers/__init__.py b/src/resolver_athena_client/client/transformers/__init__.py index cac3ca3..67ba8d6 100644 --- a/src/resolver_athena_client/client/transformers/__init__.py +++ b/src/resolver_athena_client/client/transformers/__init__.py @@ -9,6 +9,10 @@ from resolver_athena_client.client.transformers.classification_input import ( ClassificationInputTransformer, ) +from resolver_athena_client.client.transformers.core import ( + compress_image, + resize_image, +) from resolver_athena_client.client.transformers.image_resizer import ( ImageResizer, ) @@ -22,4 +26,6 @@ "ClassificationInputTransformer", "ImageResizer", "RequestBatcher", + "compress_image", + "resize_image", ] diff --git a/src/resolver_athena_client/client/transformers/brotli_compressor.py b/src/resolver_athena_client/client/transformers/brotli_compressor.py index 1e12248..14df4fb 100644 --- a/src/resolver_athena_client/client/transformers/brotli_compressor.py +++ b/src/resolver_athena_client/client/transformers/brotli_compressor.py @@ -1,11 +1,10 @@ """Compression middleware for images.""" -import brotli - from resolver_athena_client.client.models import ImageData from resolver_athena_client.client.transformers.async_transformer import ( AsyncTransformer, ) +from resolver_athena_client.client.transformers.core import compress_image class BrotliCompressor(AsyncTransformer[ImageData, ImageData]): @@ -21,8 +20,4 @@ async def transform(self, data: ImageData) -> ImageData: ImageData with compressed bytes but original hashes preserved. """ - compressed_bytes = brotli.compress(data.data) - # Modify existing ImageData with compressed bytes but preserve hashes - # since compression doesn't change image content - data.data = compressed_bytes - return data + return compress_image(data) diff --git a/src/resolver_athena_client/client/transformers/classification_input.py b/src/resolver_athena_client/client/transformers/classification_input.py index 2341b3c..2778803 100644 --- a/src/resolver_athena_client/client/transformers/classification_input.py +++ b/src/resolver_athena_client/client/transformers/classification_input.py @@ -7,7 +7,7 @@ from resolver_athena_client.client.transformers.async_transformer import ( AsyncTransformer, ) -from resolver_athena_client.generated.athena.athena_pb2 import ( +from resolver_athena_client.generated.athena.models_pb2 import ( ClassificationInput, ImageFormat, RequestEncoding, diff --git a/src/resolver_athena_client/client/transformers/core.py b/src/resolver_athena_client/client/transformers/core.py new file mode 100644 index 0000000..e97d64d --- /dev/null +++ b/src/resolver_athena_client/client/transformers/core.py @@ -0,0 +1,86 @@ +"""Core transformation functions that operate on single ImageData objects. + +This module provides the core transformation logic without async iterator +dependencies, making them easier to use for single-item operations and more +composable. +""" + +import asyncio +from io import BytesIO + +import brotli +from PIL import Image + +from resolver_athena_client.client.consts import EXPECTED_HEIGHT, EXPECTED_WIDTH +from resolver_athena_client.client.models import ImageData + +# Global optimization constants +_target_size = (EXPECTED_WIDTH, EXPECTED_HEIGHT) +_expected_raw_size = EXPECTED_WIDTH * EXPECTED_HEIGHT * 3 + + +def _is_raw_rgb_expected_size(data: bytes) -> bool: + """Detect if data is already a raw RGB array of expected size.""" + return len(data) == _expected_raw_size + + +async def resize_image(image_data: ImageData) -> ImageData: + """Resize an image to expected dimensions. + + Args: + image_data: The ImageData object to resize + + Returns: + The same ImageData object with resized data (modified in-place) + + """ + + def process_image() -> tuple[bytes, bool]: + # Fast path for raw RGB arrays of correct size + if _is_raw_rgb_expected_size(image_data.data): + return image_data.data, False # No transformation needed + + # Try to load the image data directly + input_buffer = BytesIO(image_data.data) + + with Image.open(input_buffer) as image: + # Convert to RGB if needed + rgb_image = image.convert("RGB") if image.mode != "RGB" else image + + # Resize if needed + if rgb_image.size != _target_size: + resized_image = rgb_image.resize( + _target_size, Image.Resampling.LANCZOS + ) + else: + resized_image = rgb_image + + # Convert to raw RGB bytes (C-order: height x width x channels) + return resized_image.tobytes(), True # Data was transformed + + # Use thread pool for CPU-intensive processing + resized_bytes, was_transformed = await asyncio.to_thread(process_image) + + # Only modify data and add hashes if transformation occurred + if was_transformed: + image_data.data = resized_bytes + image_data.add_transformation_hashes() + + return image_data + + +def compress_image(image_data: ImageData) -> ImageData: + """Compress image data using Brotli compression. + + Args: + image_data: The ImageData object to compress + + Returns: + The same ImageData object with compressed data (modified in-place) + + """ + compressed_bytes = brotli.compress(image_data.data) + # Modify existing ImageData with compressed bytes but preserve hashes + # since compression doesn't change image content + image_data.data = compressed_bytes + return image_data diff --git a/src/resolver_athena_client/client/transformers/image_resizer.py b/src/resolver_athena_client/client/transformers/image_resizer.py index 999aede..d9cfa84 100644 --- a/src/resolver_athena_client/client/transformers/image_resizer.py +++ b/src/resolver_athena_client/client/transformers/image_resizer.py @@ -1,25 +1,12 @@ """Optimized image resizer that ensures all images match expected dimensions.""" -import asyncio from collections.abc import AsyncIterator -from io import BytesIO -from PIL import Image - -from resolver_athena_client.client.consts import EXPECTED_HEIGHT, EXPECTED_WIDTH from resolver_athena_client.client.models import ImageData from resolver_athena_client.client.transformers.async_transformer import ( AsyncTransformer, ) - -# Global optimization constants -_target_size = (EXPECTED_WIDTH, EXPECTED_HEIGHT) -_expected_raw_size = EXPECTED_WIDTH * EXPECTED_HEIGHT * 3 - - -def _is_raw_rgb_expected_size(data: bytes) -> bool: - """Detect if data is already a raw RGB array of expected size.""" - return len(data) == _expected_raw_size +from resolver_athena_client.client.transformers.core import resize_image class ImageResizer(AsyncTransformer[ImageData, ImageData]): @@ -41,39 +28,4 @@ async def transform(self, data: ImageData) -> ImageData: Returns raw RGB bytes in C-order format (height x width x 3). """ - - def process_image() -> tuple[bytes, bool]: - # Fast path for raw RGB arrays of correct size - if _is_raw_rgb_expected_size(data.data): - return data.data, False # No transformation needed - - # Try to load the image data directly - input_buffer = BytesIO(data.data) - - with Image.open(input_buffer) as image: - # Convert to RGB if needed - if image.mode != "RGB": - rgb_image = image.convert("RGB") - else: - rgb_image = image - - # Resize if needed - if rgb_image.size != _target_size: - resized_image = rgb_image.resize( - _target_size, Image.Resampling.LANCZOS - ) - else: - resized_image = rgb_image - - # Convert to raw RGB bytes (C-order: height x width x channels) - return resized_image.tobytes(), True # Data was transformed - - # Use thread pool for CPU-intensive processing - resized_bytes, was_transformed = await asyncio.to_thread(process_image) - - # Only modify data and add hashes if transformation occurred - if was_transformed: - data.data = resized_bytes - data.add_transformation_hashes() - - return data + return await resize_image(data) diff --git a/src/resolver_athena_client/client/transformers/request_batcher.py b/src/resolver_athena_client/client/transformers/request_batcher.py index df24263..0a382c4 100644 --- a/src/resolver_athena_client/client/transformers/request_batcher.py +++ b/src/resolver_athena_client/client/transformers/request_batcher.py @@ -5,7 +5,7 @@ import time from collections.abc import AsyncIterator -from resolver_athena_client.generated.athena.athena_pb2 import ( +from resolver_athena_client.generated.athena.models_pb2 import ( ClassificationInput, ClassifyRequest, ) diff --git a/src/resolver_athena_client/client/utils.py b/src/resolver_athena_client/client/utils.py index 4bb1d1a..41135d6 100644 --- a/src/resolver_athena_client/client/utils.py +++ b/src/resolver_athena_client/client/utils.py @@ -6,7 +6,7 @@ from resolver_athena_client.client.exceptions import ClassificationOutputError if TYPE_CHECKING: - from resolver_athena_client.generated.athena.athena_pb2 import ( + from resolver_athena_client.generated.athena.models_pb2 import ( ClassificationOutput, ClassifyResponse, ) diff --git a/src/resolver_athena_client/generated/__init__.py b/src/resolver_athena_client/generated/__init__.py deleted file mode 100644 index e26d0e0..0000000 --- a/src/resolver_athena_client/generated/__init__.py +++ /dev/null @@ -1 +0,0 @@ -"""Library methods for Athena Client.""" diff --git a/src/resolver_athena_client/generated/athena/__init__.py b/src/resolver_athena_client/generated/athena/__init__.py deleted file mode 100644 index e69de29..0000000 diff --git a/src/resolver_athena_client/generated/athena/athena_pb2.py b/src/resolver_athena_client/generated/athena/athena_pb2.py deleted file mode 100644 index 2b00f0c..0000000 --- a/src/resolver_athena_client/generated/athena/athena_pb2.py +++ /dev/null @@ -1,60 +0,0 @@ -# -*- coding: utf-8 -*- -# Generated by the protocol buffer compiler. DO NOT EDIT! -# NO CHECKED-IN PROTOBUF GENCODE -# source: athena/athena.proto -# Protobuf Python Version: 6.31.1 -"""Generated protocol buffer code.""" -from google.protobuf import descriptor as _descriptor -from google.protobuf import descriptor_pool as _descriptor_pool -from google.protobuf import runtime_version as _runtime_version -from google.protobuf import symbol_database as _symbol_database -from google.protobuf.internal import builder as _builder -_runtime_version.ValidateProtobufRuntimeVersion( - _runtime_version.Domain.PUBLIC, - 6, - 31, - 1, - '', - 'athena/athena.proto' -) -# @@protoc_insertion_point(imports) - -_sym_db = _symbol_database.Default() - - -from google.protobuf import empty_pb2 as google_dot_protobuf_dot_empty__pb2 - - -DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n\x13\x61thena/athena.proto\x12\x06\x61thena\x1a\x1bgoogle/protobuf/empty.proto\"B\n\x17ListDeploymentsResponse\x12\'\n\x0b\x64\x65ployments\x18\x01 \x03(\x0b\x32\x12.athena.Deployment\"4\n\nDeployment\x12\x15\n\rdeployment_id\x18\x01 \x01(\t\x12\x0f\n\x07\x62\x61\x63klog\x18\x02 \x01(\x05\"U\n\x0f\x43lassifyRequest\x12\x15\n\rdeployment_id\x18\x01 \x01(\t\x12+\n\x06inputs\x18\x02 \x03(\x0b\x32\x1b.athena.ClassificationInput\"\x9e\x01\n\x13\x43lassificationInput\x12\x11\n\taffiliate\x18\x01 \x01(\t\x12\x16\n\x0e\x63orrelation_id\x18\x02 \x01(\t\x12)\n\x08\x65ncoding\x18\x03 \x01(\x0e\x32\x17.athena.RequestEncoding\x12\x0c\n\x04\x64\x61ta\x18\x04 \x01(\x0c\x12#\n\x06\x66ormat\x18\x05 \x01(\x0e\x32\x13.athena.ImageFormat\"t\n\x10\x43lassifyResponse\x12\x31\n\x0cglobal_error\x18\x01 \x01(\x0b\x32\x1b.athena.ClassificationError\x12-\n\x07outputs\x18\x02 \x03(\x0b\x32\x1c.athena.ClassificationOutput\"\x8b\x01\n\x14\x43lassificationOutput\x12\x16\n\x0e\x63orrelation_id\x18\x01 \x01(\t\x12/\n\x0f\x63lassifications\x18\x02 \x03(\x0b\x32\x16.athena.Classification\x12*\n\x05\x65rror\x18\x03 \x01(\x0b\x32\x1b.athena.ClassificationError\"/\n\x0e\x43lassification\x12\r\n\x05label\x18\x01 \x01(\t\x12\x0e\n\x06weight\x18\x02 \x01(\x02\"X\n\x13\x43lassificationError\x12\x1f\n\x04\x63ode\x18\x01 \x01(\x0e\x32\x11.athena.ErrorCode\x12\x0f\n\x07message\x18\x02 \x01(\t\x12\x0f\n\x07\x64\x65tails\x18\x03 \x01(\t*\x8b\x01\n\tErrorCode\x12\x1a\n\x16\x45RROR_CODE_UNSPECIFIED\x10\x00\x12\x1e\n\x1a\x45RROR_CODE_IMAGE_TOO_LARGE\x10\x02\x12\x1a\n\x16\x45RROR_CODE_MODEL_ERROR\x10\x03\x12&\n\"ERROR_CODE_AFFILIATE_NOT_PERMITTED\x10\x04*s\n\x0fRequestEncoding\x12 \n\x1cREQUEST_ENCODING_UNSPECIFIED\x10\x00\x12!\n\x1dREQUEST_ENCODING_UNCOMPRESSED\x10\x01\x12\x1b\n\x17REQUEST_ENCODING_BROTLI\x10\x02*\xbf\x03\n\x0bImageFormat\x12\x1c\n\x18IMAGE_FORMAT_UNSPECIFIED\x10\x00\x12\x14\n\x10IMAGE_FORMAT_GIF\x10\x01\x12\x15\n\x11IMAGE_FORMAT_JPEG\x10\x02\x12\x14\n\x10IMAGE_FORMAT_BMP\x10\x03\x12\x14\n\x10IMAGE_FORMAT_DIB\x10\x04\x12\x14\n\x10IMAGE_FORMAT_PNG\x10\x05\x12\x15\n\x11IMAGE_FORMAT_WEBP\x10\x06\x12\x14\n\x10IMAGE_FORMAT_PBM\x10\x07\x12\x14\n\x10IMAGE_FORMAT_PGM\x10\x08\x12\x14\n\x10IMAGE_FORMAT_PPM\x10\t\x12\x14\n\x10IMAGE_FORMAT_PXM\x10\n\x12\x14\n\x10IMAGE_FORMAT_PNM\x10\x0b\x12\x14\n\x10IMAGE_FORMAT_PFM\x10\x0c\x12\x13\n\x0fIMAGE_FORMAT_SR\x10\r\x12\x14\n\x10IMAGE_FORMAT_RAS\x10\x0e\x12\x15\n\x11IMAGE_FORMAT_TIFF\x10\x0f\x12\x14\n\x10IMAGE_FORMAT_HDR\x10\x10\x12\x14\n\x10IMAGE_FORMAT_PIC\x10\x11\x12\x1a\n\x16IMAGE_FORMAT_RAW_UINT8\x10\x12\x32\xa2\x01\n\x11\x43lassifierService\x12\x41\n\x08\x43lassify\x12\x17.athena.ClassifyRequest\x1a\x18.athena.ClassifyResponse(\x01\x30\x01\x12J\n\x0fListDeployments\x12\x16.google.protobuf.Empty\x1a\x1f.athena.ListDeploymentsResponseBz\n\x18\x63om.resolver.athena.grpcB\x0b\x41thenaProto\xa2\x02\x03RAT\xaa\x02\x14Resolver.Athena.Grpc\xba\x02\x03RAT\xca\x02\x14Resolver\\Athena\\Grpc\xea\x02\x16Resolver::Athena::Grpcb\x06proto3') - -_globals = globals() -_builder.BuildMessageAndEnumDescriptors(DESCRIPTOR, _globals) -_builder.BuildTopDescriptorsAndMessages(DESCRIPTOR, 'athena.athena_pb2', _globals) -if not _descriptor._USE_C_DESCRIPTORS: - _globals['DESCRIPTOR']._loaded_options = None - _globals['DESCRIPTOR']._serialized_options = b'\n\030com.resolver.athena.grpcB\013AthenaProto\242\002\003RAT\252\002\024Resolver.Athena.Grpc\272\002\003RAT\312\002\024Resolver\\Athena\\Grpc\352\002\026Resolver::Athena::Grpc' - _globals['_ERRORCODE']._serialized_start=830 - _globals['_ERRORCODE']._serialized_end=969 - _globals['_REQUESTENCODING']._serialized_start=971 - _globals['_REQUESTENCODING']._serialized_end=1086 - _globals['_IMAGEFORMAT']._serialized_start=1089 - _globals['_IMAGEFORMAT']._serialized_end=1536 - _globals['_LISTDEPLOYMENTSRESPONSE']._serialized_start=60 - _globals['_LISTDEPLOYMENTSRESPONSE']._serialized_end=126 - _globals['_DEPLOYMENT']._serialized_start=128 - _globals['_DEPLOYMENT']._serialized_end=180 - _globals['_CLASSIFYREQUEST']._serialized_start=182 - _globals['_CLASSIFYREQUEST']._serialized_end=267 - _globals['_CLASSIFICATIONINPUT']._serialized_start=270 - _globals['_CLASSIFICATIONINPUT']._serialized_end=428 - _globals['_CLASSIFYRESPONSE']._serialized_start=430 - _globals['_CLASSIFYRESPONSE']._serialized_end=546 - _globals['_CLASSIFICATIONOUTPUT']._serialized_start=549 - _globals['_CLASSIFICATIONOUTPUT']._serialized_end=688 - _globals['_CLASSIFICATION']._serialized_start=690 - _globals['_CLASSIFICATION']._serialized_end=737 - _globals['_CLASSIFICATIONERROR']._serialized_start=739 - _globals['_CLASSIFICATIONERROR']._serialized_end=827 - _globals['_CLASSIFIERSERVICE']._serialized_start=1539 - _globals['_CLASSIFIERSERVICE']._serialized_end=1701 -# @@protoc_insertion_point(module_scope) diff --git a/src/resolver_athena_client/generated/athena/athena_pb2.pyi b/src/resolver_athena_client/generated/athena/athena_pb2.pyi deleted file mode 100644 index 2c00a65..0000000 --- a/src/resolver_athena_client/generated/athena/athena_pb2.pyi +++ /dev/null @@ -1,389 +0,0 @@ -""" -@generated by mypy-protobuf. Do not edit manually! -isort:skip_file -""" - -import builtins -import collections.abc -import google.protobuf.descriptor -import google.protobuf.internal.containers -import google.protobuf.internal.enum_type_wrapper -import google.protobuf.message -import sys -import typing - -if sys.version_info >= (3, 10): - import typing as typing_extensions -else: - import typing_extensions - -DESCRIPTOR: google.protobuf.descriptor.FileDescriptor - -class _ErrorCode: - ValueType = typing.NewType("ValueType", builtins.int) - V: typing_extensions.TypeAlias = ValueType - -class _ErrorCodeEnumTypeWrapper(google.protobuf.internal.enum_type_wrapper._EnumTypeWrapper[_ErrorCode.ValueType], builtins.type): - DESCRIPTOR: google.protobuf.descriptor.EnumDescriptor - ERROR_CODE_UNSPECIFIED: _ErrorCode.ValueType # 0 - """Unknown or unspecified error""" - ERROR_CODE_IMAGE_TOO_LARGE: _ErrorCode.ValueType # 2 - """Image is too large to process""" - ERROR_CODE_MODEL_ERROR: _ErrorCode.ValueType # 3 - """Opaque error from the classifier""" - ERROR_CODE_AFFILIATE_NOT_PERMITTED: _ErrorCode.ValueType # 4 - """Attempt to send data for an affiliate which is not granted to the current - client. - """ - -class ErrorCode(_ErrorCode, metaclass=_ErrorCodeEnumTypeWrapper): - """Enumeration of possible classification error codes.""" - -ERROR_CODE_UNSPECIFIED: ErrorCode.ValueType # 0 -"""Unknown or unspecified error""" -ERROR_CODE_IMAGE_TOO_LARGE: ErrorCode.ValueType # 2 -"""Image is too large to process""" -ERROR_CODE_MODEL_ERROR: ErrorCode.ValueType # 3 -"""Opaque error from the classifier""" -ERROR_CODE_AFFILIATE_NOT_PERMITTED: ErrorCode.ValueType # 4 -"""Attempt to send data for an affiliate which is not granted to the current -client. -""" -global___ErrorCode = ErrorCode - -class _RequestEncoding: - ValueType = typing.NewType("ValueType", builtins.int) - V: typing_extensions.TypeAlias = ValueType - -class _RequestEncodingEnumTypeWrapper(google.protobuf.internal.enum_type_wrapper._EnumTypeWrapper[_RequestEncoding.ValueType], builtins.type): - DESCRIPTOR: google.protobuf.descriptor.EnumDescriptor - REQUEST_ENCODING_UNSPECIFIED: _RequestEncoding.ValueType # 0 - """Unspecified encoding. Assumed to be uncompressed.""" - REQUEST_ENCODING_UNCOMPRESSED: _RequestEncoding.ValueType # 1 - """Uncompressed raw image data (e.g., raw RGB pixels, BMP, etc.)""" - REQUEST_ENCODING_BROTLI: _RequestEncoding.ValueType # 2 - """Brotli-compressed image data for reduced bandwidth usage - Server will decompress using Brotli algorithm before processing - """ - -class RequestEncoding(_RequestEncoding, metaclass=_RequestEncodingEnumTypeWrapper): - """Enumeration of supported image data encoding formats. - Determines how the server should interpret the image data bytes. - """ - -REQUEST_ENCODING_UNSPECIFIED: RequestEncoding.ValueType # 0 -"""Unspecified encoding. Assumed to be uncompressed.""" -REQUEST_ENCODING_UNCOMPRESSED: RequestEncoding.ValueType # 1 -"""Uncompressed raw image data (e.g., raw RGB pixels, BMP, etc.)""" -REQUEST_ENCODING_BROTLI: RequestEncoding.ValueType # 2 -"""Brotli-compressed image data for reduced bandwidth usage -Server will decompress using Brotli algorithm before processing -""" -global___RequestEncoding = RequestEncoding - -class _ImageFormat: - ValueType = typing.NewType("ValueType", builtins.int) - V: typing_extensions.TypeAlias = ValueType - -class _ImageFormatEnumTypeWrapper(google.protobuf.internal.enum_type_wrapper._EnumTypeWrapper[_ImageFormat.ValueType], builtins.type): - DESCRIPTOR: google.protobuf.descriptor.EnumDescriptor - IMAGE_FORMAT_UNSPECIFIED: _ImageFormat.ValueType # 0 - IMAGE_FORMAT_GIF: _ImageFormat.ValueType # 1 - IMAGE_FORMAT_JPEG: _ImageFormat.ValueType # 2 - """Covers .jpeg, .jpg, .jpe extensions""" - IMAGE_FORMAT_BMP: _ImageFormat.ValueType # 3 - IMAGE_FORMAT_DIB: _ImageFormat.ValueType # 4 - IMAGE_FORMAT_PNG: _ImageFormat.ValueType # 5 - IMAGE_FORMAT_WEBP: _ImageFormat.ValueType # 6 - IMAGE_FORMAT_PBM: _ImageFormat.ValueType # 7 - IMAGE_FORMAT_PGM: _ImageFormat.ValueType # 8 - IMAGE_FORMAT_PPM: _ImageFormat.ValueType # 9 - IMAGE_FORMAT_PXM: _ImageFormat.ValueType # 10 - IMAGE_FORMAT_PNM: _ImageFormat.ValueType # 11 - IMAGE_FORMAT_PFM: _ImageFormat.ValueType # 12 - IMAGE_FORMAT_SR: _ImageFormat.ValueType # 13 - IMAGE_FORMAT_RAS: _ImageFormat.ValueType # 14 - IMAGE_FORMAT_TIFF: _ImageFormat.ValueType # 15 - """Covers .tiff, .tif extensions""" - IMAGE_FORMAT_HDR: _ImageFormat.ValueType # 16 - IMAGE_FORMAT_PIC: _ImageFormat.ValueType # 17 - IMAGE_FORMAT_RAW_UINT8: _ImageFormat.ValueType # 18 - """Raw unsigned 8-bit image data (RGB, C order array)""" - -class ImageFormat(_ImageFormat, metaclass=_ImageFormatEnumTypeWrapper): - """Enumeration of supported image file formats. - Specifies the image file format structure for proper parsing. - """ - -IMAGE_FORMAT_UNSPECIFIED: ImageFormat.ValueType # 0 -IMAGE_FORMAT_GIF: ImageFormat.ValueType # 1 -IMAGE_FORMAT_JPEG: ImageFormat.ValueType # 2 -"""Covers .jpeg, .jpg, .jpe extensions""" -IMAGE_FORMAT_BMP: ImageFormat.ValueType # 3 -IMAGE_FORMAT_DIB: ImageFormat.ValueType # 4 -IMAGE_FORMAT_PNG: ImageFormat.ValueType # 5 -IMAGE_FORMAT_WEBP: ImageFormat.ValueType # 6 -IMAGE_FORMAT_PBM: ImageFormat.ValueType # 7 -IMAGE_FORMAT_PGM: ImageFormat.ValueType # 8 -IMAGE_FORMAT_PPM: ImageFormat.ValueType # 9 -IMAGE_FORMAT_PXM: ImageFormat.ValueType # 10 -IMAGE_FORMAT_PNM: ImageFormat.ValueType # 11 -IMAGE_FORMAT_PFM: ImageFormat.ValueType # 12 -IMAGE_FORMAT_SR: ImageFormat.ValueType # 13 -IMAGE_FORMAT_RAS: ImageFormat.ValueType # 14 -IMAGE_FORMAT_TIFF: ImageFormat.ValueType # 15 -"""Covers .tiff, .tif extensions""" -IMAGE_FORMAT_HDR: ImageFormat.ValueType # 16 -IMAGE_FORMAT_PIC: ImageFormat.ValueType # 17 -IMAGE_FORMAT_RAW_UINT8: ImageFormat.ValueType # 18 -"""Raw unsigned 8-bit image data (RGB, C order array)""" -global___ImageFormat = ImageFormat - -@typing.final -class ListDeploymentsResponse(google.protobuf.message.Message): - """Response message for ListDeployments RPC - Contains the list of active deployments and their details. - """ - - DESCRIPTOR: google.protobuf.descriptor.Descriptor - - DEPLOYMENTS_FIELD_NUMBER: builtins.int - @property - def deployments(self) -> google.protobuf.internal.containers.RepeatedCompositeFieldContainer[global___Deployment]: - """List of active deployments with their backlog information.""" - - def __init__( - self, - *, - deployments: collections.abc.Iterable[global___Deployment] | None = ..., - ) -> None: ... - def ClearField(self, field_name: typing.Literal["deployments", b"deployments"]) -> None: ... - -global___ListDeploymentsResponse = ListDeploymentsResponse - -@typing.final -class Deployment(google.protobuf.message.Message): - """A single active deployment part of a `ListDeployments` response""" - - DESCRIPTOR: google.protobuf.descriptor.Descriptor - - DEPLOYMENT_ID_FIELD_NUMBER: builtins.int - BACKLOG_FIELD_NUMBER: builtins.int - deployment_id: builtins.str - """active deployment identifier""" - backlog: builtins.int - """Backlog of classification responses in this deployment""" - def __init__( - self, - *, - deployment_id: builtins.str = ..., - backlog: builtins.int = ..., - ) -> None: ... - def ClearField(self, field_name: typing.Literal["backlog", b"backlog", "deployment_id", b"deployment_id"]) -> None: ... - -global___Deployment = Deployment - -@typing.final -class ClassifyRequest(google.protobuf.message.Message): - """The request message containing the image data to classify. - Each request represents a batch of images that should be processed within - the same deployment context. - """ - - DESCRIPTOR: google.protobuf.descriptor.Descriptor - - DEPLOYMENT_ID_FIELD_NUMBER: builtins.int - INPUTS_FIELD_NUMBER: builtins.int - deployment_id: builtins.str - """Client's unique identifier for this deployment. Responses returned will be - sent to a client with a matching deployment_id. - """ - @property - def inputs(self) -> google.protobuf.internal.containers.RepeatedCompositeFieldContainer[global___ClassificationInput]: - """Array of images to be classified in this request batch - Allows sending multiple images in a single request for efficiency. - """ - - def __init__( - self, - *, - deployment_id: builtins.str = ..., - inputs: collections.abc.Iterable[global___ClassificationInput] | None = ..., - ) -> None: ... - def ClearField(self, field_name: typing.Literal["deployment_id", b"deployment_id", "inputs", b"inputs"]) -> None: ... - -global___ClassifyRequest = ClassifyRequest - -@typing.final -class ClassificationInput(google.protobuf.message.Message): - """A single image within a classification request batch. - Contains all necessary metadata and data for classifying one image. - """ - - DESCRIPTOR: google.protobuf.descriptor.Descriptor - - AFFILIATE_FIELD_NUMBER: builtins.int - CORRELATION_ID_FIELD_NUMBER: builtins.int - ENCODING_FIELD_NUMBER: builtins.int - DATA_FIELD_NUMBER: builtins.int - FORMAT_FIELD_NUMBER: builtins.int - affiliate: builtins.str - """The affiliate or source system that provided this image - Used for tracking, analytics, and routing purposes. - """ - correlation_id: builtins.str - """Unique identifier for correlating this input with its response - Must be unique within the deployment to properly match responses - """ - encoding: global___RequestEncoding.ValueType - """Specifies the encoding/compression format of the image data - Allows the server to properly decode the image before classification - """ - data: builtins.bytes - """The raw image data bytes in the format specified by encoding - Can be compressed or uncompressed based on the encoding field - """ - format: global___ImageFormat.ValueType - """The image file format of the data bytes""" - def __init__( - self, - *, - affiliate: builtins.str = ..., - correlation_id: builtins.str = ..., - encoding: global___RequestEncoding.ValueType = ..., - data: builtins.bytes = ..., - format: global___ImageFormat.ValueType = ..., - ) -> None: ... - def ClearField(self, field_name: typing.Literal["affiliate", b"affiliate", "correlation_id", b"correlation_id", "data", b"data", "encoding", b"encoding", "format", b"format"]) -> None: ... - -global___ClassificationInput = ClassificationInput - -@typing.final -class ClassifyResponse(google.protobuf.message.Message): - """The response message containing the classification results. - Sent back to clients for each processed batch, containing either - a global error or individual results for each image in the batch. - """ - - DESCRIPTOR: google.protobuf.descriptor.Descriptor - - GLOBAL_ERROR_FIELD_NUMBER: builtins.int - OUTPUTS_FIELD_NUMBER: builtins.int - @property - def global_error(self) -> global___ClassificationError: - """Global error affecting the entire batch/request - If present, indicates that the entire request failed and no individual - image results will be provided - """ - - @property - def outputs(self) -> google.protobuf.internal.containers.RepeatedCompositeFieldContainer[global___ClassificationOutput]: - """Array of classification results, one for each input image - Will be empty if global_error is present - """ - - def __init__( - self, - *, - global_error: global___ClassificationError | None = ..., - outputs: collections.abc.Iterable[global___ClassificationOutput] | None = ..., - ) -> None: ... - def HasField(self, field_name: typing.Literal["global_error", b"global_error"]) -> builtins.bool: ... - def ClearField(self, field_name: typing.Literal["global_error", b"global_error", "outputs", b"outputs"]) -> None: ... - -global___ClassifyResponse = ClassifyResponse - -@typing.final -class ClassificationOutput(google.protobuf.message.Message): - """Individual classification result for a single image. - Contains the correlation ID and classification results for one image. - """ - - DESCRIPTOR: google.protobuf.descriptor.Descriptor - - CORRELATION_ID_FIELD_NUMBER: builtins.int - CLASSIFICATIONS_FIELD_NUMBER: builtins.int - ERROR_FIELD_NUMBER: builtins.int - correlation_id: builtins.str - """Matches the correlationId from the corresponding ClassificationInput - Allows clients to match responses with their original requests - """ - @property - def classifications(self) -> google.protobuf.internal.containers.RepeatedCompositeFieldContainer[global___Classification]: - """Array of all classifications detected for this image - Multiple classifications may be returned with different confidence levels - """ - - @property - def error(self) -> global___ClassificationError: - """Error information if this specific image classification failed - If present, indicates that this particular image could not be processed - """ - - def __init__( - self, - *, - correlation_id: builtins.str = ..., - classifications: collections.abc.Iterable[global___Classification] | None = ..., - error: global___ClassificationError | None = ..., - ) -> None: ... - def HasField(self, field_name: typing.Literal["error", b"error"]) -> builtins.bool: ... - def ClearField(self, field_name: typing.Literal["classifications", b"classifications", "correlation_id", b"correlation_id", "error", b"error"]) -> None: ... - -global___ClassificationOutput = ClassificationOutput - -@typing.final -class Classification(google.protobuf.message.Message): - """A single classification result for an image. - Represents one detected category with its confidence score. - """ - - DESCRIPTOR: google.protobuf.descriptor.Descriptor - - LABEL_FIELD_NUMBER: builtins.int - WEIGHT_FIELD_NUMBER: builtins.int - label: builtins.str - """Human-readable label describing what was classified - Examples: "CatA", "CatB", "Indicitive", "Distraction", etc. - """ - weight: builtins.float - """Confidence score between 0.0 and 1.0 indicating certainty - Higher values indicate greater confidence in the classification - """ - def __init__( - self, - *, - label: builtins.str = ..., - weight: builtins.float = ..., - ) -> None: ... - def ClearField(self, field_name: typing.Literal["label", b"label", "weight", b"weight"]) -> None: ... - -global___Classification = Classification - -@typing.final -class ClassificationError(google.protobuf.message.Message): - """Error information for failed classification attempts. - Provides details about why a classification could not be completed. - """ - - DESCRIPTOR: google.protobuf.descriptor.Descriptor - - CODE_FIELD_NUMBER: builtins.int - MESSAGE_FIELD_NUMBER: builtins.int - DETAILS_FIELD_NUMBER: builtins.int - code: global___ErrorCode.ValueType - """Error code indicating the type of failure""" - message: builtins.str - """Human-readable error message providing details about the failure""" - details: builtins.str - """Additional context or details about the error (optional)""" - def __init__( - self, - *, - code: global___ErrorCode.ValueType = ..., - message: builtins.str = ..., - details: builtins.str = ..., - ) -> None: ... - def ClearField(self, field_name: typing.Literal["code", b"code", "details", b"details", "message", b"message"]) -> None: ... - -global___ClassificationError = ClassificationError diff --git a/src/resolver_athena_client/generated/athena/athena_pb2_grpc.py b/src/resolver_athena_client/generated/athena/athena_pb2_grpc.py deleted file mode 100644 index 4dc4e03..0000000 --- a/src/resolver_athena_client/generated/athena/athena_pb2_grpc.py +++ /dev/null @@ -1,156 +0,0 @@ -# Generated by the gRPC Python protocol compiler plugin. DO NOT EDIT! -"""Client and server classes corresponding to protobuf-defined services.""" -import grpc -import warnings - -from resolver_athena_client.generated.athena import athena_pb2 as athena_dot_athena__pb2 -from google.protobuf import empty_pb2 as google_dot_protobuf_dot_empty__pb2 - -GRPC_GENERATED_VERSION = '1.74.0' -GRPC_VERSION = grpc.__version__ -_version_not_supported = False - -try: - from grpc._utilities import first_version_is_lower - _version_not_supported = first_version_is_lower(GRPC_VERSION, GRPC_GENERATED_VERSION) -except ImportError: - _version_not_supported = True - -if _version_not_supported: - raise RuntimeError( - f'The grpc package installed is at version {GRPC_VERSION},' - + f' but the generated code in athena/athena_pb2_grpc.py depends on' - + f' grpcio>={GRPC_GENERATED_VERSION}.' - + f' Please upgrade your grpc module to grpcio>={GRPC_GENERATED_VERSION}' - + f' or downgrade your generated code using grpcio-tools<={GRPC_VERSION}.' - ) - - -class ClassifierServiceStub(object): - """The classifier service definition. - Provides image classification capabilities with session-based streaming - and client management functionality. - """ - - def __init__(self, channel): - """Constructor. - - Args: - channel: A grpc.Channel. - """ - self.Classify = channel.stream_stream( - '/athena.ClassifierService/Classify', - request_serializer=athena_dot_athena__pb2.ClassifyRequest.SerializeToString, - response_deserializer=athena_dot_athena__pb2.ClassifyResponse.FromString, - _registered_method=True) - self.ListDeployments = channel.unary_unary( - '/athena.ClassifierService/ListDeployments', - request_serializer=google_dot_protobuf_dot_empty__pb2.Empty.SerializeToString, - response_deserializer=athena_dot_athena__pb2.ListDeploymentsResponse.FromString, - _registered_method=True) - - -class ClassifierServiceServicer(object): - """The classifier service definition. - Provides image classification capabilities with session-based streaming - and client management functionality. - """ - - def Classify(self, request_iterator, context): - """Classify images in a deployment-based streaming context - Multiple affiliates can join the same deployment to share responses - Supports bidirectional streaming for real-time classification - """ - context.set_code(grpc.StatusCode.UNIMPLEMENTED) - context.set_details('Method not implemented!') - raise NotImplementedError('Method not implemented!') - - def ListDeployments(self, request, context): - """Retrieves a list of all active deployment IDs - Returns the active deployment_id values that can be used in Classify requests - Useful for monitoring and debugging active connections - """ - context.set_code(grpc.StatusCode.UNIMPLEMENTED) - context.set_details('Method not implemented!') - raise NotImplementedError('Method not implemented!') - - -def add_ClassifierServiceServicer_to_server(servicer, server): - rpc_method_handlers = { - 'Classify': grpc.stream_stream_rpc_method_handler( - servicer.Classify, - request_deserializer=athena_dot_athena__pb2.ClassifyRequest.FromString, - response_serializer=athena_dot_athena__pb2.ClassifyResponse.SerializeToString, - ), - 'ListDeployments': grpc.unary_unary_rpc_method_handler( - servicer.ListDeployments, - request_deserializer=google_dot_protobuf_dot_empty__pb2.Empty.FromString, - response_serializer=athena_dot_athena__pb2.ListDeploymentsResponse.SerializeToString, - ), - } - generic_handler = grpc.method_handlers_generic_handler( - 'athena.ClassifierService', rpc_method_handlers) - server.add_generic_rpc_handlers((generic_handler,)) - server.add_registered_method_handlers('athena.ClassifierService', rpc_method_handlers) - - - # This class is part of an EXPERIMENTAL API. -class ClassifierService(object): - """The classifier service definition. - Provides image classification capabilities with session-based streaming - and client management functionality. - """ - - @staticmethod - def Classify(request_iterator, - target, - options=(), - channel_credentials=None, - call_credentials=None, - insecure=False, - compression=None, - wait_for_ready=None, - timeout=None, - metadata=None): - return grpc.experimental.stream_stream( - request_iterator, - target, - '/athena.ClassifierService/Classify', - athena_dot_athena__pb2.ClassifyRequest.SerializeToString, - athena_dot_athena__pb2.ClassifyResponse.FromString, - options, - channel_credentials, - insecure, - call_credentials, - compression, - wait_for_ready, - timeout, - metadata, - _registered_method=True) - - @staticmethod - def ListDeployments(request, - target, - options=(), - channel_credentials=None, - call_credentials=None, - insecure=False, - compression=None, - wait_for_ready=None, - timeout=None, - metadata=None): - return grpc.experimental.unary_unary( - request, - target, - '/athena.ClassifierService/ListDeployments', - google_dot_protobuf_dot_empty__pb2.Empty.SerializeToString, - athena_dot_athena__pb2.ListDeploymentsResponse.FromString, - options, - channel_credentials, - insecure, - call_credentials, - compression, - wait_for_ready, - timeout, - metadata, - _registered_method=True) diff --git a/src/resolver_athena_client/grpc_wrappers/classifier_service.py b/src/resolver_athena_client/grpc_wrappers/classifier_service.py index f3ddc20..f21e852 100644 --- a/src/resolver_athena_client/grpc_wrappers/classifier_service.py +++ b/src/resolver_athena_client/grpc_wrappers/classifier_service.py @@ -6,14 +6,16 @@ from google.protobuf.empty_pb2 import Empty from grpc import aio -from resolver_athena_client.generated.athena.athena_pb2 import ( +from resolver_athena_client.generated.athena.athena_pb2_grpc import ( + ClassifierServiceStub, +) +from resolver_athena_client.generated.athena.models_pb2 import ( + ClassificationInput, + ClassificationOutput, ClassifyRequest, ClassifyResponse, ListDeploymentsResponse, ) -from resolver_athena_client.generated.athena.athena_pb2_grpc import ( - ClassifierServiceStub, -) if TYPE_CHECKING: from grpc.aio import StreamStreamCall @@ -65,3 +67,26 @@ async def list_deployments(self) -> ListDeploymentsResponse: """ return await self.stub.ListDeployments(Empty()) + + async def classify_single( + self, + request: ClassificationInput, + timeout: float | None = None, + ) -> ClassificationOutput: + """Classify a single image synchronously without deployment context. + + Args: + request (ClassificationInput): The classification input containing + the image data and metadata to be classified. + timeout (float | None): RPC timeout in seconds. None for no timeout. + + Returns: + ClassificationOutput: The classification result containing either + classifications or error information. + + """ + return await self.stub.ClassifySingle( + request, + timeout=timeout, + wait_for_ready=True, + ) diff --git a/src/resolver_athena_client/version.py b/src/resolver_athena_client/version.py index 3edca84..178edaf 100644 --- a/src/resolver_athena_client/version.py +++ b/src/resolver_athena_client/version.py @@ -2,4 +2,4 @@ import importlib.metadata -__version__ = importlib.metadata.version("athena_client") +__version__ = importlib.metadata.version("resolver-athena-client") diff --git a/tests/client/test_athena_client.py b/tests/client/test_athena_client.py index 8c6e49c..b11169d 100644 --- a/tests/client/test_athena_client.py +++ b/tests/client/test_athena_client.py @@ -9,7 +9,7 @@ from resolver_athena_client.client.athena_options import AthenaOptions from resolver_athena_client.client.exceptions import AthenaError from resolver_athena_client.client.models import ImageData -from resolver_athena_client.generated.athena.athena_pb2 import ( +from resolver_athena_client.generated.athena.models_pb2 import ( ClassificationError, ClassificationOutput, ClassifyResponse, diff --git a/tests/client/test_deployment_selector.py b/tests/client/test_deployment_selector.py index f622ec1..e4ba6a6 100644 --- a/tests/client/test_deployment_selector.py +++ b/tests/client/test_deployment_selector.py @@ -6,7 +6,7 @@ from grpc import aio from resolver_athena_client.client.deployment_selector import DeploymentSelector -from resolver_athena_client.generated.athena.athena_pb2 import ( +from resolver_athena_client.generated.athena.models_pb2 import ( Deployment, ListDeploymentsResponse, ) diff --git a/tests/client/test_output_error_handling.py b/tests/client/test_output_error_handling.py index 0078277..0f395d0 100644 --- a/tests/client/test_output_error_handling.py +++ b/tests/client/test_output_error_handling.py @@ -10,7 +10,7 @@ log_output_errors, process_classification_outputs, ) -from resolver_athena_client.generated.athena.athena_pb2 import ( +from resolver_athena_client.generated.athena.models_pb2 import ( Classification, ClassificationError, ClassificationOutput, diff --git a/tests/client/test_timeout_behavior.py b/tests/client/test_timeout_behavior.py index d09907b..1ab229f 100644 --- a/tests/client/test_timeout_behavior.py +++ b/tests/client/test_timeout_behavior.py @@ -13,7 +13,7 @@ from resolver_athena_client.client.athena_client import AthenaClient from resolver_athena_client.client.athena_options import AthenaOptions from resolver_athena_client.client.models import ImageData -from resolver_athena_client.generated.athena.athena_pb2 import ( +from resolver_athena_client.generated.athena.models_pb2 import ( ClassificationOutput, ClassifyResponse, ) diff --git a/tests/client/transformers/test_classification_input.py b/tests/client/transformers/test_classification_input.py index 4092a92..2ad0ba9 100644 --- a/tests/client/transformers/test_classification_input.py +++ b/tests/client/transformers/test_classification_input.py @@ -7,7 +7,7 @@ from resolver_athena_client.client.transformers.classification_input import ( ClassificationInputTransformer, ) -from resolver_athena_client.generated.athena.athena_pb2 import ( +from resolver_athena_client.generated.athena.models_pb2 import ( ImageFormat, RequestEncoding, ) diff --git a/tests/client/transformers/test_core.py b/tests/client/transformers/test_core.py new file mode 100644 index 0000000..bcf2338 --- /dev/null +++ b/tests/client/transformers/test_core.py @@ -0,0 +1,221 @@ +"""Test core transformation functions.""" + +from io import BytesIO + +import pytest +from PIL import Image + +from resolver_athena_client.client.consts import ( + EXPECTED_HEIGHT, + EXPECTED_WIDTH, +) +from resolver_athena_client.client.models import ImageData +from resolver_athena_client.client.transformers.core import ( + compress_image, + resize_image, +) + + +def create_test_image( + width: int = 100, height: int = 100, mode: str = "RGB" +) -> bytes: + """Create a test image with specified dimensions.""" + img = Image.new(mode, (width, height), color="red") + img_bytes = BytesIO() + img.save(img_bytes, format="PNG") + return img_bytes.getvalue() + + +@pytest.mark.asyncio +async def test_resize_image_basic() -> None: + """Test basic image resizing functionality.""" + # Create a test image that needs resizing + test_image_bytes = create_test_image(200, 150) + image_data = ImageData(test_image_bytes) + + original_hash_count = len(image_data.md5_hashes) + + # Resize the image + result = await resize_image(image_data) + + # Should be the same object + assert result is image_data + + # Should have added a new hash since transformation occurred + assert len(image_data.md5_hashes) == original_hash_count + 1 + + # Data should be different (resized) + assert image_data.data != test_image_bytes + + +@pytest.mark.asyncio +async def test_resize_image_already_correct_size() -> None: + """Test resizing when image is already correct size but needs conversion.""" + # Create image with expected dimensions (but still PNG format) + test_image_bytes = create_test_image(EXPECTED_WIDTH, EXPECTED_HEIGHT) + image_data = ImageData(test_image_bytes) + + original_hash_count = len(image_data.md5_hashes) + + # Resize the image + result = await resize_image(image_data) + + # Should be the same object + assert result is image_data + + # Should have added new hash since format conversion occurred + assert len(image_data.md5_hashes) == original_hash_count + 1 + + # Data should be different (converted to raw RGB) + assert image_data.data != test_image_bytes + + +@pytest.mark.asyncio +async def test_resize_image_grayscale_conversion() -> None: + """Test resizing with grayscale to RGB conversion.""" + # Create a grayscale test image + test_image_bytes = create_test_image(100, 100, mode="L") + image_data = ImageData(test_image_bytes) + + original_hash_count = len(image_data.md5_hashes) + + # Resize the image + result = await resize_image(image_data) + + # Should be the same object + assert result is image_data + + # Should have added a new hash since transformation occurred + assert len(image_data.md5_hashes) == original_hash_count + 1 + + # Data should be different (converted and resized) + assert image_data.data != test_image_bytes + + +@pytest.mark.asyncio +async def test_resize_image_raw_rgb_fast_path() -> None: + """Test the fast path for raw RGB data of correct size.""" + # Create raw RGB data of expected size + expected_size = EXPECTED_WIDTH * EXPECTED_HEIGHT * 3 + raw_rgb_data = b"\x00" * expected_size + image_data = ImageData(raw_rgb_data) + + original_hash_count = len(image_data.md5_hashes) + original_data = image_data.data + + # Resize the image + result = await resize_image(image_data) + + # Should be the same object + assert result is image_data + + # Should not have added new hash (fast path, no transformation) + assert len(image_data.md5_hashes) == original_hash_count + + # Data should be unchanged + assert image_data.data == original_data + + +def test_compress_image_basic() -> None: + """Test basic image compression functionality.""" + # Create test image data + test_data = b"This is some test image data that should be compressed" + image_data = ImageData(test_data) + + original_hash_count = len(image_data.md5_hashes) + + # Compress the image + result = compress_image(image_data) + + # Should be the same object + assert result is image_data + + # Should not have added new hash (compression preserves hashes) + assert len(image_data.md5_hashes) == original_hash_count + + # Data should be different (compressed) + assert image_data.data != test_data + assert len(image_data.data) < len(test_data) # Should be smaller + + +def test_compress_image_empty_data() -> None: + """Test compression with empty data.""" + image_data = ImageData(b"") + + original_hash_count = len(image_data.md5_hashes) + + # Compress the image + result = compress_image(image_data) + + # Should be the same object + assert result is image_data + + # Should not have added new hash + assert len(image_data.md5_hashes) == original_hash_count + + # Data should be compressed (but still small) + assert ( + len(image_data.data) > 0 + ) # Brotli adds some overhead even for empty data + + +def test_compress_image_preserves_hashes() -> None: + """Test that compression preserves the original hash list.""" + # Create image data and add some transformation hashes + image_data = ImageData(b"test data") + image_data.add_transformation_hashes() # Simulate a previous transformation + + original_hashes = image_data.md5_hashes.copy() + + # Compress the image + compress_image(image_data) + + # Hashes should be unchanged + assert image_data.md5_hashes == original_hashes + + +@pytest.mark.asyncio +async def test_combined_transformations() -> None: + """Test using both resize and compress transformations together.""" + # Create a test image that needs resizing + test_image_bytes = create_test_image(200, 150) + image_data = ImageData(test_image_bytes) + + original_data = image_data.data + original_hash_count = len(image_data.md5_hashes) + + # Apply resize transformation + await resize_image(image_data) + + # Should have new hash from resizing + assert len(image_data.md5_hashes) == original_hash_count + 1 + resized_data = image_data.data + assert resized_data != original_data + + # Apply compression transformation + compress_image(image_data) + + # Should still have the same number of hashes (compression preserves) + assert len(image_data.md5_hashes) == original_hash_count + 1 + compressed_data = image_data.data + assert compressed_data != resized_data + assert compressed_data != original_data + + +@pytest.mark.asyncio +async def test_transformations_modify_in_place() -> None: + """Test that transformations modify the ImageData object in-place.""" + test_image_bytes = create_test_image(100, 100) + image_data = ImageData(test_image_bytes) + + original_id = id(image_data) + + # Apply transformations + result1 = await resize_image(image_data) + result2 = compress_image(image_data) + + # All results should be the same object + assert id(result1) == original_id + assert id(result2) == original_id + assert result1 is image_data + assert result2 is image_data diff --git a/tests/client/transformers/test_request_batcher.py b/tests/client/transformers/test_request_batcher.py index 5ddb9d3..71f8caa 100644 --- a/tests/client/transformers/test_request_batcher.py +++ b/tests/client/transformers/test_request_batcher.py @@ -9,7 +9,7 @@ from resolver_athena_client.client.transformers.request_batcher import ( RequestBatcher, ) -from resolver_athena_client.generated.athena.athena_pb2 import ( +from resolver_athena_client.generated.athena.models_pb2 import ( ClassificationInput, ClassifyRequest, RequestEncoding, diff --git a/tests/grpc_wrappers/classifier_service_test.py b/tests/grpc_wrappers/classifier_service_test.py index d0e21cf..12f166a 100644 --- a/tests/grpc_wrappers/classifier_service_test.py +++ b/tests/grpc_wrappers/classifier_service_test.py @@ -3,7 +3,7 @@ import pytest from google.protobuf.empty_pb2 import Empty -from resolver_athena_client.generated.athena.athena_pb2 import ( +from resolver_athena_client.generated.athena.models_pb2 import ( ClassifyResponse, ListDeploymentsResponse, ) diff --git a/tests/test_classify_single.py b/tests/test_classify_single.py new file mode 100644 index 0000000..edadd1b --- /dev/null +++ b/tests/test_classify_single.py @@ -0,0 +1,378 @@ +"""Tests for the classify_single method in AthenaClient.""" + +import io +import uuid +from unittest.mock import AsyncMock, Mock + +import grpc.aio +import pytest +from PIL import Image + +from resolver_athena_client.client.athena_client import AthenaClient +from resolver_athena_client.client.athena_options import AthenaOptions +from resolver_athena_client.client.exceptions import AthenaError +from resolver_athena_client.client.models import ImageData +from resolver_athena_client.generated.athena.models_pb2 import ( + Classification, + ClassificationError, + ClassificationInput, + ClassificationOutput, + ErrorCode, + HashType, + ImageFormat, + ImageHash, + RequestEncoding, +) + + +@pytest.fixture +def athena_options() -> AthenaOptions: + """Create test AthenaOptions.""" + return AthenaOptions( + host="localhost:8080", + affiliate="test-affiliate", + deployment_id="test-deployment", + resize_images=False, + compress_images=False, + timeout=30.0, + ) + + +@pytest.fixture +def mock_channel() -> Mock: + """Create a mock gRPC channel.""" + return Mock() + + +@pytest.fixture +def mock_classifier() -> Mock: + """Create a mock classifier service client.""" + return Mock() + + +@pytest.fixture +def athena_client( + mock_channel: Mock, athena_options: AthenaOptions, mock_classifier: Mock +) -> AthenaClient: + """Create AthenaClient with mocked dependencies.""" + client = AthenaClient(mock_channel, athena_options) + client.classifier = mock_classifier + return client + + +@pytest.fixture +def sample_image_data() -> ImageData: + """Create sample image data for testing.""" + return ImageData(b"fake_image_bytes") + + +@pytest.mark.asyncio +async def test_classify_single_success( + athena_client: AthenaClient, sample_image_data: ImageData +) -> None: + """Test successful single image classification.""" + # Setup mock response + expected_classification = Classification(label="test_label", weight=0.95) + mock_output = ClassificationOutput( + correlation_id="test-correlation", + classifications=[expected_classification], + ) + athena_client.classifier.classify_single = AsyncMock( + return_value=mock_output + ) + + # Call classify_single + result = await athena_client.classify_single(sample_image_data) + + # Verify result + assert result == mock_output + assert len(result.classifications) == 1 + assert result.classifications[0].label == "test_label" + tolerance = 0.001 + assert abs(result.classifications[0].weight - 0.95) < tolerance + assert not result.HasField("error") + + # Verify the call was made with correct parameters + athena_client.classifier.classify_single.assert_called_once() + call_args = athena_client.classifier.classify_single.call_args[0][0] + + assert isinstance(call_args, ClassificationInput) + assert call_args.affiliate == "test-affiliate" + assert call_args.data == b"fake_image_bytes" + assert call_args.encoding == RequestEncoding.REQUEST_ENCODING_UNCOMPRESSED + assert call_args.format == ImageFormat.IMAGE_FORMAT_RAW_UINT8 + assert len(call_args.hashes) == 1 + assert call_args.hashes[0].type == HashType.HASH_TYPE_MD5 + assert len(call_args.hashes[0].value) > 0 # MD5 hash should be generated + + +@pytest.mark.asyncio +async def test_classify_single_with_correlation_id( + athena_client: AthenaClient, sample_image_data: ImageData +) -> None: + """Test classify_single with custom correlation ID.""" + custom_correlation_id = "custom-correlation-123" + + # Setup mock response + mock_output = ClassificationOutput( + correlation_id=custom_correlation_id, + classifications=[], + ) + athena_client.classifier.classify_single = AsyncMock( + return_value=mock_output + ) + + # Call classify_single with custom correlation ID + await athena_client.classify_single( + sample_image_data, correlation_id=custom_correlation_id + ) + + # Verify correlation ID was used + call_args = athena_client.classifier.classify_single.call_args[0][0] + assert call_args.correlation_id == custom_correlation_id + + +@pytest.mark.asyncio +async def test_classify_single_auto_correlation_id( + athena_client: AthenaClient, sample_image_data: ImageData +) -> None: + """Test classify_single generates correlation ID when not provided.""" + # Setup mock response + mock_output = ClassificationOutput( + correlation_id="auto-generated", + classifications=[], + ) + athena_client.classifier.classify_single = AsyncMock( + return_value=mock_output + ) + + # Call classify_single without correlation ID + await athena_client.classify_single(sample_image_data) + + # Verify a correlation ID was generated + call_args = athena_client.classifier.classify_single.call_args[0][0] + assert call_args.correlation_id is not None + assert len(call_args.correlation_id) > 0 + # Should be a valid UUID format + uuid.UUID(call_args.correlation_id) # This will raise if not a valid UUID + + +@pytest.mark.asyncio +async def test_classify_single_with_compression( + athena_client: AthenaClient, sample_image_data: ImageData +) -> None: + """Test classify_single with compression enabled.""" + # Enable compression + athena_client.options.compress_images = True + + # Setup mock response + mock_output = ClassificationOutput( + correlation_id="test-correlation", + classifications=[], + ) + athena_client.classifier.classify_single = AsyncMock( + return_value=mock_output + ) + + # Call classify_single + await athena_client.classify_single(sample_image_data) + + # Verify compression settings were applied + call_args = athena_client.classifier.classify_single.call_args[0][0] + assert call_args.encoding == RequestEncoding.REQUEST_ENCODING_BROTLI + # Data should be compressed - check it's the same as the modified image data + assert call_args.data == sample_image_data.data + + +@pytest.mark.asyncio +async def test_classify_single_error_handling( + athena_client: AthenaClient, +) -> None: + """Test classify_single with image resizing enabled.""" + # Create a simple valid image for testing + + # Create a simple 1x1 pixel image + img = Image.new("RGB", (1, 1), color="red") + img_bytes = io.BytesIO() + img.save(img_bytes, format="PNG") + valid_image_data = ImageData(img_bytes.getvalue()) + + # Enable resizing + athena_client.options.resize_images = True + + # Setup mock response + mock_output = ClassificationOutput( + correlation_id="test-correlation", + classifications=[], + ) + athena_client.classifier.classify_single = AsyncMock( + return_value=mock_output + ) + + # Call classify_single with valid image + await athena_client.classify_single(valid_image_data) + + # Verify resizing was processed (encoding should be uncompressed) + call_args = athena_client.classifier.classify_single.call_args[0][0] + assert call_args.encoding == RequestEncoding.REQUEST_ENCODING_UNCOMPRESSED + + +@pytest.mark.asyncio +async def test_classify_single_with_error_response( + athena_client: AthenaClient, sample_image_data: ImageData +) -> None: + """Test classify_single handling error in response.""" + # Setup mock response with error + error = ClassificationError( + code=ErrorCode.ERROR_CODE_IMAGE_TOO_LARGE, + message="Image is too large", + details="Max size is 10MB", + ) + mock_output = ClassificationOutput( + correlation_id="test-correlation", + classifications=[], + error=error, + ) + athena_client.classifier.classify_single = AsyncMock( + return_value=mock_output + ) + + # Call classify_single and expect AthenaError + with pytest.raises(AthenaError, match="Image is too large"): + await athena_client.classify_single(sample_image_data) + + +@pytest.mark.asyncio +async def test_classify_single_timeout_handling( + athena_client: AthenaClient, sample_image_data: ImageData +) -> None: + """Test classify_single handling gRPC errors.""" + # Setup mock to raise gRPC error + grpc_error = grpc.aio.AioRpcError( + code=grpc.StatusCode.UNAVAILABLE, + initial_metadata=grpc.aio.Metadata(), + trailing_metadata=grpc.aio.Metadata(), + details="Service unavailable", + ) + athena_client.classifier.classify_single = AsyncMock(side_effect=grpc_error) + + # Call classify_single and expect gRPC error to be re-raised + with pytest.raises(grpc.aio.AioRpcError): + await athena_client.classify_single(sample_image_data) + + +@pytest.mark.asyncio +async def test_classify_single_timeout_parameter( + athena_client: AthenaClient, sample_image_data: ImageData +) -> None: + """Test that timeout is passed to the gRPC call.""" + # Setup mock response + mock_output = ClassificationOutput( + correlation_id="test-correlation", + classifications=[], + ) + athena_client.classifier.classify_single = AsyncMock( + return_value=mock_output + ) + + # Call classify_single + await athena_client.classify_single(sample_image_data) + + # Verify timeout was passed + call_kwargs = athena_client.classifier.classify_single.call_args[1] + expected_timeout = 30.0 # From the fixture options + assert call_kwargs["timeout"] == expected_timeout + + +@pytest.mark.asyncio +async def test_classify_single_multiple_hashes( + athena_client: AthenaClient, +) -> None: + """Test classify_single with multiple transformation hashes.""" + # Create image with multiple transformations + image_data = ImageData(b"original_image") + image_data.add_transformation_hashes() # Simulate a transformation + image_data.data = b"transformed_image" + image_data.add_transformation_hashes() # Another transformation + + # Setup mock response + mock_output = ClassificationOutput( + correlation_id="test-correlation", + classifications=[], + error=None, + ) + athena_client.classifier.classify_single = AsyncMock( + return_value=mock_output + ) + + # Call classify_single + await athena_client.classify_single(image_data) + + # Verify all hashes were included + call_args = athena_client.classifier.classify_single.call_args[0][0] + expected_hash_count = 3 # Original + 2 transformations + assert len(call_args.hashes) == expected_hash_count + for hash_obj in call_args.hashes: + assert isinstance(hash_obj, ImageHash) + assert hash_obj.type == HashType.HASH_TYPE_MD5 + assert len(hash_obj.value) > 0 + + +@pytest.mark.asyncio +async def test_classify_single_multiple_classifications( + athena_client: AthenaClient, sample_image_data: ImageData +) -> None: + """Test classify_single with multiple classification results.""" + # Setup mock response with multiple classifications + classifications = [ + Classification(label="cat", weight=0.85), + Classification(label="animal", weight=0.92), + Classification(label="pet", weight=0.78), + ] + mock_output = ClassificationOutput( + correlation_id="test-correlation", + classifications=classifications, + ) + athena_client.classifier.classify_single = AsyncMock( + return_value=mock_output + ) + + # Call classify_single + result = await athena_client.classify_single(sample_image_data) + + # Verify all classifications are returned + expected_classification_count = 3 + assert len(result.classifications) == expected_classification_count + labels = [c.label for c in result.classifications] + weights = [c.weight for c in result.classifications] + + assert "cat" in labels + assert "animal" in labels + assert "pet" in labels + tolerance = 0.001 + assert any(abs(w - 0.85) < tolerance for w in weights) + assert any(abs(w - 0.92) < tolerance for w in weights) + assert any(abs(w - 0.78) < tolerance for w in weights) + + +@pytest.mark.asyncio +async def test_classify_single_empty_classifications( + athena_client: AthenaClient, sample_image_data: ImageData +) -> None: + """Test classify_single with no classification results.""" + # Setup mock response with empty classifications + mock_output = ClassificationOutput( + correlation_id="test-correlation", + classifications=[], + ) + athena_client.classifier.classify_single = AsyncMock( + return_value=mock_output + ) + + # Call classify_single + result = await athena_client.classify_single(sample_image_data) + + # Verify empty result is handled correctly + assert len(result.classifications) == 0 + assert not result.HasField("error") # Check protobuf field is not set + assert result.correlation_id == "test-correlation"