Merge pull request #6 from crispthinking/feature/multiple-inputs

Feature/multiple inputs

Author: abbottdev-kroll

README.md

@@ -14,7 +14,14 @@ Athena is a gRPC-based image classification service designed for CSAM (Child Sex
 - **Error Handling**: Comprehensive error codes and detailed error messages
 - **Monitoring**: Active deployment tracking and backlog monitoring
 
-## Contributing
+# Contributing
+
+## Updating the Protobuf definitions
+
+Protobufs are stored in the [@crispthinking/athena-protobuffs](https://github.com/crispthinking/athena-protobufs) repository.
+
+To update the protobuf definitions for client generation, run:
+`git subtree pull --prefix=athena-protobufs https://github.com/crispthinking/athena-protobufs.git <sha> --squash`
 
 ## Regenerating the TypeScript gRPC Client
 

__tests__/unit/main.test.ts

@@ -1,5 +1,5 @@
 import { describe, it,} from 'vitest';
-import { ClassifierSdk, ClassifyImageOptions, ClassifyResponse, ImageFormat } from '../../src/main';
+import { ClassificationOutput, ClassifierSdk, ClassifyImageInput, ImageFormat } from '../../src';
 import fs from 'fs';
 import { randomUUID } from 'crypto';
 
@@ -31,7 +31,99 @@ describe('classifierHelper', () => {
 });
 
 describe('classifyImage function', () => {
-  it('should classify Steamboat-willie.jpg and return responses (integration smoke test)', async ({expect}) => {
+  it('should classify 10 images in a single request and return responses (integration smoke test)', async ({expect, annotate}) => {
+    // This is a smoke test. You must have a running gRPC server at localhost:50051 for this to pass.
+    // You may want to mock the gRPC client for true unit testing.
+    const imagePath = __dirname + '/Steamboat-willie.jpg';
+    const sdk = new ClassifierSdk({
+        deploymentId: process.env.VITE_ATHENA_DEPLOYMENT_ID,
+        affiliate: process.env.VITE_ATHENA_AFFILIATE,
+        authentication: {
+          issuerUrl: process.env.VITE_OAUTH_ISSUER,
+          clientId: process.env.VITE_ATHENA_CLIENT_ID,
+          clientSecret: process.env.VITE_ATHENA_CLIENT_SECRET,
+          scope: 'manage:classify'
+        }
+    });
+
+    // Generate 10 unique correlationIds
+    const correlationIds = Array.from({ length: 10 }, () => randomUUID().toString());
+
+    correlationIds.sort((a, b) => a.localeCompare(b));
+
+    annotate(`Correlation IDs: ${correlationIds.join(', ')}`);
+
+    // Create 10 input objects, each with a new stream and unique correlationId
+    const inputs: ClassifyImageInput[] = correlationIds.map((correlationId) => ({
+      imageStream: fs.createReadStream(imagePath),
+      format: ImageFormat.PNG,
+      correlationId
+    }));
+
+    // Create a promise to wrap the event emitter event 'data'
+    const promise = new Promise<ClassificationOutput[]>((resolve, reject) => {
+      const results:ClassificationOutput[] = [];
+
+      sdk.on('data', (data) => {
+        if (data.globalError)
+        {
+          reject(data.globalError);
+        }
+
+        // Check that all correlationIds are present in the outputs
+        for(const result of data.outputs)
+        {
+          if (correlationIds.includes(result.correlationId)) {
+            results.push(result);
+          }
+        }
+        if (results.length == correlationIds.length) {
+          resolve(results);
+        }
+      });
+      sdk.once('error', (err) => {
+        reject(err);
+      });
+    });
+
+    let error: any = undefined;
+
+    await sdk.open();
+
+    try {
+      await sdk.sendClassifyRequest(inputs);
+    } catch (err) {
+      error = err;
+    }
+
+    // Wait for classifier to process some data....
+    const outputs = await promise;
+    sdk.close();
+
+    expect(error).toBeUndefined();
+
+    outputs.sort((a, b) => a.correlationId.localeCompare(b.correlationId));
+
+    expect(outputs).toBeDefined();
+    // Check that all correlationIds are present in the outputs
+    expect(outputs.length).toBe(correlationIds.length);
+
+    const expectedOutputs = correlationIds.map(id =>(
+      {
+        correlationId: id,
+        classifications: expect.arrayContaining([
+          {
+            label: expect.any(String),
+            weight: expect.any(Number)
+          }
+        ])
+      }
+    ));
+
+    expect(outputs).toMatchObject(expectedOutputs);
+  }, 120000);
+
+  it('should classify Steamboat-willie.jpg with raw uint8 resize return responses (integration smoke test)', async ({expect, annotate}) => {
     // This is a smoke test. You must have a running gRPC server at localhost:50051 for this to pass.
     // You may want to mock the gRPC client for true unit testing.
     const imagePath = __dirname + '/Steamboat-willie.jpg';
@@ -48,16 +140,24 @@ describe('classifyImage function', () => {
 
     const correlationId = randomUUID();
 
+    annotate(`Correlation IDs: ${correlationId}`);
+
     // Create a promise to wrap the event emitter event 'data'
-    const promise = new Promise<ClassifyResponse>((resolve, reject) => {
-      sdk.once('data', (data) => {
+    const promise = new Promise<ClassificationOutput[]>((resolve, reject) => {
+      // Add a timeout to reject the promise if no data is received in 30 seconds
+      const timeout = setTimeout(() => {
+        reject(new Error('Timeout waiting for classification response'));
+      }, 30000);
+
+      sdk.on('data', (data) => {
         const byCorrelationId = data.outputs.filter(o => o.correlationId === correlationId);
         if (byCorrelationId.length > 0) {
-          resolve(data);
+          clearTimeout(timeout);
+          resolve(byCorrelationId);
         }
-        sdk.close();
       });
       sdk.once('error', (err) => {
+        clearTimeout(timeout);
         reject(err);
       });
     });
@@ -68,10 +168,10 @@ describe('classifyImage function', () => {
     await sdk.open();
 
     const imageStream = fs.createReadStream(imagePath);
-    const options: ClassifyImageOptions = {
+    const options: ClassifyImageInput = {
       imageStream,
-      format: ImageFormat.PNG,
-      correlationId
+      correlationId,
+      resize: true,
     };
     try {
       await sdk.sendClassifyRequest(options);
@@ -81,12 +181,97 @@ describe('classifyImage function', () => {
 
     // Wait for classifier to process some data....
     const first = await promise;
+    sdk.close();
 
     expect(first).toBeDefined();
+    expect(first).toMatchObject([
+      {
+        correlationId,
+        classifications: expect.arrayContaining([
+          {
+            label: expect.any(String),
+            weight: expect.any(Number)
+          }
+        ])
+      } as ClassificationOutput
+    ]);
+
+    // Accept either a successful call or a connection error (for CI/dev convenience)
+    expect(error).toBeUndefined();
+  }, 120000);
+
+  it('should classify Steamboat-willie.jpg and return responses (integration smoke test)', async ({expect, annotate}) => {
+    // This is a smoke test. You must have a running gRPC server at localhost:50051 for this to pass.
+    // You may want to mock the gRPC client for true unit testing.
+    const imagePath = __dirname + '/Steamboat-willie.jpg';
+    const sdk = new ClassifierSdk({
+        deploymentId: process.env.VITE_ATHENA_DEPLOYMENT_ID,
+        affiliate: process.env.VITE_ATHENA_AFFILIATE,
+        authentication: {
+          issuerUrl: process.env.VITE_OAUTH_ISSUER,
+          clientId: process.env.VITE_ATHENA_CLIENT_ID,
+          clientSecret: process.env.VITE_ATHENA_CLIENT_SECRET,
+          scope: 'manage:classify'
+        }
+    });
 
-    const byCorrelationId = first.outputs.filter(o => o.correlationId === correlationId);
+    const correlationId = randomUUID();
+
+    annotate(`Correlation IDs: ${correlationId}`);
+
+    // Create a promise to wrap the event emitter event 'data'
+    const promise = new Promise<ClassificationOutput[]>((resolve, reject) => {
+      // Add a timeout to reject the promise if no data is received in 30 seconds
+      const timeout = setTimeout(() => {
+        reject(new Error('Timeout waiting for classification response'));
+      }, 30000);
 
-    expect(byCorrelationId.length).toBeGreaterThan(0);
+      sdk.on('data', (data) => {
+        const byCorrelationId = data.outputs.filter(o => o.correlationId === correlationId);
+        if (byCorrelationId.length > 0) {
+          clearTimeout(timeout);
+          resolve(byCorrelationId);
+        }
+      });
+      sdk.once('error', (err) => {
+        clearTimeout(timeout);
+        reject(err);
+      });
+    });
+
+    // This will fail if no server is running, but will exercise the code path.
+    let error: any = undefined;
+
+    await sdk.open();
+
+    const imageStream = fs.createReadStream(imagePath);
+    const options: ClassifyImageInput = {
+      imageStream,
+      correlationId,
+      format: ImageFormat.JPEG
+    };
+    try {
+      await sdk.sendClassifyRequest(options);
+    } catch (err) {
+      error = err;
+    }
+
+    // Wait for classifier to process some data....
+    const first = await promise;
+    sdk.close();
+
+    expect(first).toBeDefined();
+    expect(first).toMatchObject([
+      {
+        correlationId,
+        classifications: expect.arrayContaining([
+          {
+            label: expect.any(String),
+            weight: expect.any(Number)
+          }
+        ])
+      } as ClassificationOutput
+    ]);
 
     // Accept either a successful call or a connection error (for CI/dev convenience)
     expect(error).toBeUndefined();

package.json

@@ -47,7 +47,7 @@
     "test:watch": "vitest unit"
   },
   "author": "James Abbott <abbottdev@users.noreply.github.com>",
-  "license": "Apache-2.0",
+  "license": "MIT",
   "dependencies": {
     "@bufbuild/protobuf": "^2.7.0",
     "@grpc/grpc-js": "^1.13.4",

src/hashing.ts

@@ -1,8 +1,9 @@
 import { Readable } from 'stream';
 import crypto from 'crypto';
 import sharp from 'sharp';
-import { RequestEncoding } from './main';
+import { ImageFormat, RequestEncoding } from '.';
 import brotli from 'brotli';
+import { buffer } from 'stream/consumers';
 
 /**
  * Computes MD5 and SHA1 hashes from a readable stream and resizes any image data.
@@ -13,16 +14,25 @@ import brotli from 'brotli';
 export async function computeHashesFromStream(
   stream: Readable,
   encoding: RequestEncoding = RequestEncoding.UNCOMPRESSED,
-): Promise<{ md5: string; sha1: string; data: Buffer }> {
+  imageFormat: ImageFormat = ImageFormat.UNSPECIFIED,
+  resize: boolean = false,
+): Promise<{ md5: string; sha1: string; data: Buffer; format: ImageFormat }> {
   const md5 = crypto.createHash('md5');
   const sha1 = crypto.createHash('sha1');
-  const resizer = sharp().resize(448, 448).raw({ depth: 'uint' });
+
+  let data: Buffer<ArrayBufferLike>;
 
   stream.pipe(md5);
   stream.pipe(sha1);
-  stream.pipe(resizer);
 
-  let data = await resizer.toBuffer();
+  if (resize) {
+    const resizer = sharp().resize(448, 448).raw({ depth: 'uint' });
+    stream.pipe(resizer);
+    data = await resizer.toBuffer();
+    imageFormat = ImageFormat.RAW_UINT8;
+  } else {
+    data = await buffer(stream);
+  }
 
   if (encoding === RequestEncoding.BROTLI) {
     data = Buffer.from(await brotli.compress(data));
@@ -32,5 +42,6 @@ export async function computeHashesFromStream(
     md5: md5.digest('hex'),
     sha1: sha1.digest('hex'),
     data,
+    format: imageFormat,
   };
 }