feat: adds oauth credential helper, better streaming support

Author: snus-kin

README.md

@@ -3,6 +3,65 @@
 This is a Python library for interacting with the Athena API (Resolver Unknown
 CSAM Detection).
 
+## Authentication
+
+The Athena client supports two authentication methods:
+
+### Static Token Authentication
+```python
+from athena_client.client.channel import create_channel
+
+# Use a pre-existing authentication token
+channel = create_channel(host="your-host", auth_token="your-token")
+```
+
+### OAuth Credential Helper (Recommended)
+The credential helper automatically handles OAuth token acquisition and refresh:
+
+```python
+import asyncio
+from athena_client.client.channel import CredentialHelper, create_channel_with_credentials
+
+async def main():
+    # Create credential helper with OAuth settings
+    credential_helper = CredentialHelper(
+        client_id="your-oauth-client-id",
+        client_secret="your-oauth-client-secret",
+        auth_url="https://crispthinking.auth0.com/oauth/token",  # Optional, this is default
+        audience="crisp-athena-dev"  # Optional, this is default
+    )
+
+    # Create channel with automatic OAuth handling
+    channel = await create_channel_with_credentials(
+        host="your-host",
+        credential_helper=credential_helper
+    )
+
+asyncio.run(main())
+```
+
+#### Environment Variables
+For the OAuth example to work, set these environment variables:
+```bash
+export OAUTH_CLIENT_ID="your-client-id"
+export OAUTH_CLIENT_SECRET="your-client-secret"
+export ATHENA_HOST="your-athena-host"
+```
+
+#### OAuth Features
+- **Automatic token refresh**: Tokens are automatically refreshed when they expire
+- **Thread-safe**: Multiple concurrent requests will safely share cached tokens
+- **Error handling**: Comprehensive error handling for OAuth failures
+- **Configurable**: Custom OAuth endpoints and audiences supported
+
+See `examples/oauth_example.py` for a complete working example.
+
+## Examples
+
+- `examples/example.py` - Basic classification example with static token
+- `examples/oauth_example.py` - OAuth authentication with credential helper
+- `examples/create_image.py` - Image generation utilities
+
 ## TODO
 
 ### Async pipelines

examples/create_image.py

@@ -56,18 +56,26 @@ def create_random_image(width: int = 640, height: int = 480) -> bytes:
     return img_byte_arr.getvalue()
 
 
-async def iter_images(max_images: int | None = None) -> AsyncIterator[bytes]:
+async def iter_images(
+    max_images: int | None = None, counter: list[int] | None = None
+) -> AsyncIterator[bytes]:
     """Generate random test images.
 
     Args:
         max_images: Maximum number of images to generate. If None, generates
-        infinitely.
+            infinitely.
+        counter: Optional list with single integer to track number of images
+            sent.
+            The first element will be incremented for each image generated.
 
     Yields:
         JPEG image bytes with random content
 
     """
     count = 0
     while max_images is None or count < max_images:
-        yield create_random_image()
+        img = create_random_image()
+        if counter is not None:
+            counter[0] = counter[0] + 1
+        yield img
         count += 1

examples/example.py

@@ -21,80 +21,106 @@ async def run_example(
     options: AthenaOptions,
     auth_token: str,
     max_test_images: int | None = None,
-) -> None:
-    """Run an example classification with the given options.
+) -> tuple[int, int]:
+    """Run example classification with the given options.
 
     Args:
         logger: Logger instance for output
         options: Configuration options for the Athena client
         auth_token: Authentication token for the API
         max_test_images: Maximum number of test images to generate
-            (None for infinite)
+
+    Returns:
+        Number of requests sent and responses received
 
     """
     channel = create_channel(options.host, auth_token)
 
+    sent_counter = [0]  # Use list to allow mutation in closure
+    received_count = 0
+
     async with AthenaClient(channel, options) as client:
-        # Start image classification
-        results = client.classify_images(iter_images(max_test_images))
-        logger.info("Starting classification...")
+        results = client.classify_images(
+            iter_images(max_test_images, sent_counter)
+        )
 
-        # Initialize rate tracking
         start_time = time.time()
-        response_count = 0
 
         try:
             async for result in results:
-                response_count += 1
-                current_time = time.time()
-                elapsed = current_time - start_time
+                received_count += len(result.outputs)
 
-                if response_count % 10 == 0:
-                    rate = response_count / elapsed
+                if received_count % 10 == 0:
+                    elapsed = time.time() - start_time
+                    rate = received_count / elapsed if elapsed > 0 else 0
                     logger.info(
-                        "Processed %d responses (%.1f/sec)",
-                        response_count,
+                        "Sent %d requests, received %d responses (%.1f/sec)",
+                        sent_counter[0],
+                        received_count,
                         rate,
                     )
 
-                logger.debug("Got result: %s", result)
+                for output in result.outputs:
+                    classifications = {
+                        c.label: round(c.weight, 3)
+                        for c in output.classifications
+                    }
+                    logger.debug(
+                        "Result [%s]: %s",
+                        output.correlation_id[:8],
+                        classifications,
+                    )
+
         except Exception:
             logger.exception("Error during classification")
-            if response_count == 0:
-                # Re-raise if we didn't process any responses
+            if received_count == 0:
                 raise
         finally:
-            # Calculate final statistics after loop completes
-            if response_count > 0:
-                duration = time.time() - start_time
-                avg_rate = response_count / duration
+            duration = time.time() - start_time
+            if received_count > 0:
+                avg_rate = received_count / duration if duration > 0 else 0
                 logger.info(
-                    "\nFinal: %d responses in %.1f seconds (%.1f/sec)",
-                    response_count,
+                    "Completed: sent=%d received=%d in %.1fs (%.1f/sec)",
+                    sent_counter[0],
+                    received_count,
                     duration,
                     avg_rate,
                 )
 
+                if options.timeout and duration >= options.timeout * 0.95:
+                    logger.info(
+                        "Stream reached maximum duration: %.1fs (limit: %.1fs)",
+                        duration,
+                        options.timeout,
+                    )
+                elif options.timeout:
+                    logger.info(
+                        "Stream completed naturally in %.1fs (max: %.1fs)",
+                        duration,
+                        options.timeout,
+                    )
+
+    return (sent_counter[0], received_count)
+
 
 async def main() -> int:
-    """Run examples showing both authenticated and unauthenticated usage."""
+    """Run the classification example."""
     logger = logging.getLogger(__name__)
     load_dotenv()
 
-    # Set maximum number of test images to generate (None for infinite)
-    max_test_images = None
+    # Configuration
+    max_test_images = 100
 
-    # Example with authenticated channel
     auth_token = os.getenv("ATHENA_AUTH_TOKEN")
     if not auth_token:
-        logger.error("ATHENA_AUTH_TOKEN not set.")
+        logger.error("ATHENA_AUTH_TOKEN not set")
         return 1
 
-    logger.info("Running example with secure authenticated channel")
     host = os.getenv("ATHENA_HOST", "localhost")
-    channel = create_channel(host, auth_token)
+    logger.info("Connecting to %s", host)
 
-    # Get the first available deployment using context manager
+    # Get available deployment
+    channel = create_channel(host, auth_token)
     async with DeploymentSelector(channel) as deployment_selector:
         deployments = await deployment_selector.list_deployments()
 
@@ -105,15 +131,25 @@ async def main() -> int:
     deployment_id = deployments.deployments[0].deployment_id
     logger.info("Using deployment: %s", deployment_id)
 
-    athena_options = AthenaOptions(
+    # Run classification with persistent connection
+    options = AthenaOptions(
         host=host,
         resize_images=True,
         deployment_id=deployment_id,
         compress_images=True,
+        timeout=120.0,  # Maximum duration, not forced timeout
+        keepalive_interval=30.0,  # Longer intervals for persistent streams
+    )
+
+    sent, received = await run_example(
+        logger, options, auth_token, max_test_images
     )
-    await run_example(logger, athena_options, auth_token, max_test_images)
 
-    return 0
+    if sent == received:
+        logger.info("Success: %d requests processed", sent)
+        return 0
+    logger.warning("Incomplete: %d sent, %d received", sent, received)
+    return 1
 
 
 if __name__ == "__main__":