feat: add protobuf update instructions to README, enhance classifyImage tests, and improve image processing in hashing

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

@@ -123,6 +123,83 @@ describe('classifyImage function', () => {
     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';
+    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 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);
+
+      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,
+      resize: true,
+    };
+    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();
+  }, 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.
@@ -170,7 +247,8 @@ describe('classifyImage function', () => {
     const imageStream = fs.createReadStream(imagePath);
     const options: ClassifyImageInput = {
       imageStream,
-      correlationId
+      correlationId,
+      format: ImageFormat.JPEG
     };
     try {
       await sdk.sendClassifyRequest(options);

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

@@ -3,6 +3,7 @@ import crypto from 'crypto';
 import sharp from 'sharp';
 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,17 +14,25 @@ import brotli from 'brotli';
 export async function computeHashesFromStream(
   stream: Readable,
   encoding: RequestEncoding = RequestEncoding.UNCOMPRESSED,
+  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).png();
+  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));
@@ -33,6 +42,6 @@ export async function computeHashesFromStream(
     md5: md5.digest('hex'),
     sha1: sha1.digest('hex'),
     data,
-    format: ImageFormat.PNG,
+    format: imageFormat,
   };
 }

src/index.ts

@@ -8,6 +8,7 @@ import {
   ImageHash,
   ClassifyResponse,
   Deployment,
+  ImageFormat,
 } from './athena/athena';
 import * as grpc from '@grpc/grpc-js';
 import {
@@ -31,12 +32,20 @@ import { computeHashesFromStream } from './hashing';
  * @property encoding Optional encoding type for the image data.
  * @property format Optional image format.
  */
-export interface ClassifyImageInput {
+export type ClassifyImageInput = {
   affiliate?: string;
   correlationId?: string;
   imageStream: Readable;
   encoding?: ClassifyRequest['inputs'][number]['encoding'];
-}
+} & (ResizeImageInput | RawImageInput);
+
+export type ResizeImageInput = {
+  resize: true;
+};
+
+export type RawImageInput = {
+  format: ClassifyRequest['inputs'][number]['format'];
+};
 
 /**
  * Options for initializing the ClassifierSdk.
@@ -217,6 +226,7 @@ export class ClassifierSdk extends (EventEmitter as new () => TypedEmitter<Class
       : [request];
 
     const processedInputs: ClassificationInput[] = [];
+
     for (const options of requests) {
       const {
         affiliate = this.options.affiliate,
@@ -225,9 +235,17 @@ export class ClassifierSdk extends (EventEmitter as new () => TypedEmitter<Class
         encoding = RequestEncoding.UNCOMPRESSED,
       } = options;
 
+      let inputFormat: ImageFormat = ImageFormat.UNSPECIFIED;
+
+      if ('resize' in options === false) {
+        inputFormat = options.format;
+      }
+
       const { md5, sha1, data, format } = await computeHashesFromStream(
         imageStream,
         encoding,
+        inputFormat,
+        'resize' in options,
       );
       const hashes: ImageHash[] = [
         { value: md5, type: HashType.MD5 },