feat: add support for multiple images and add smoke test

Author: abbottdev-kroll

__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,119 @@ 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 =
+    [
+      randomUUID().toString(),
+      randomUUID().toString(),
+      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);
+    expect(outputs).toMatchObject([
+      {
+        correlationId: correlationIds[0],
+        classifications: expect.arrayContaining([
+          {
+            label: expect.any(String),
+            weight: expect.any(Number)
+          }
+        ])
+      } as ClassificationOutput,
+      {
+        correlationId: correlationIds[1],
+        classifications: expect.arrayContaining([
+          {
+            label: expect.any(String),
+            weight: expect.any(Number)
+          }
+        ])
+      } as ClassificationOutput,
+      {
+        correlationId: correlationIds[2],
+        classifications: expect.arrayContaining([
+          {
+            label: expect.any(String),
+            weight: expect.any(Number)
+          }
+        ])
+      } as ClassificationOutput
+    ]);
+  }, 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';
@@ -48,16 +160,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,9 +188,8 @@ describe('classifyImage function', () => {
     await sdk.open();
 
     const imageStream = fs.createReadStream(imagePath);
-    const options: ClassifyImageOptions = {
+    const options: ClassifyImageInput = {
       imageStream,
-      format: ImageFormat.PNG,
       correlationId
     };
     try {
@@ -81,12 +200,20 @@ describe('classifyImage function', () => {
 
     // Wait for classifier to process some data....
     const first = await promise;
+    sdk.close();
 
     expect(first).toBeDefined();
-
-    const byCorrelationId = first.outputs.filter(o => o.correlationId === correlationId);
-
-    expect(byCorrelationId.length).toBeGreaterThan(0);
+    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();

src/hashing.ts

@@ -1,7 +1,7 @@
 import { Readable } from 'stream';
 import crypto from 'crypto';
 import sharp from 'sharp';
-import { RequestEncoding } from './main';
+import { ImageFormat, RequestEncoding } from '.';
 import brotli from 'brotli';
 
 /**
@@ -13,10 +13,11 @@ import brotli from 'brotli';
 export async function computeHashesFromStream(
   stream: Readable,
   encoding: RequestEncoding = RequestEncoding.UNCOMPRESSED,
-): Promise<{ md5: string; sha1: string; data: Buffer }> {
+): 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' });
+
+  const resizer = sharp().resize(448, 448).png();
 
   stream.pipe(md5);
   stream.pipe(sha1);
@@ -32,5 +33,6 @@ export async function computeHashesFromStream(
     md5: md5.digest('hex'),
     sha1: sha1.digest('hex'),
     data,
+    format: ImageFormat.PNG,
   };
 }

src/main.ts src/index.tssrc/main.ts renamed to src/index.ts

@@ -5,11 +5,10 @@ import {
   ClassificationInput,
   HashType,
   RequestEncoding,
-  ImageFormat,
   ImageHash,
   ClassifyResponse,
   Deployment,
-} from './athena/athena.js';
+} from './athena/athena';
 import * as grpc from '@grpc/grpc-js';
 import {
   IClassifierServiceClient,
@@ -32,23 +31,22 @@ import { computeHashesFromStream } from './hashing';
  * @property encoding Optional encoding type for the image data.
  * @property format Optional image format.
  */
-export interface ClassifyImageOptions {
+export interface ClassifyImageInput {
   affiliate?: string;
   correlationId?: string;
   imageStream: Readable;
   encoding?: ClassifyRequest['inputs'][number]['encoding'];
-  format?: ClassifyRequest['inputs'][number]['format'];
 }
 
 /**
- * Options for initializing the ClassifierSdk helper.
+ * Options for initializing the ClassifierSdk.
  * @property keepAliveInterval Optional interval (ms) for keep-alive pings.
  * @property grpcAddress Optional gRPC server address.
  * @property deploymentId Deployment ID to use for classification.
  * @property affiliate Affiliate identifier for requests.
  * @property authentication Authentication options for the SDK.
  */
-export interface ClassifierHelperOptions {
+export interface ClassifierSdkOptions {
   keepAliveInterval?: number;
   grpcAddress?: string;
   deploymentId: string;
@@ -70,7 +68,6 @@ export type ClassifierEvents = {
   open: () => void;
 };
 
-// export the default endpoint of csam-classification-messages.crispdev.com so library consumers may use it.
 export const defaultGrpcAddress =
   'csam-classification-messages.crispdev.com:443';
 
@@ -85,7 +82,7 @@ export class ClassifierSdk extends (EventEmitter as new () => TypedEmitter<Class
     ClassifyRequest,
     ClassifyResponse
   > | null = null;
-  private options: ClassifierHelperOptions;
+  private options: ClassifierSdkOptions;
   private auth: AuthenticationManager;
   private keepAlive: NodeJS.Timeout;
   private metadata?: grpc.Metadata;
@@ -100,7 +97,7 @@ export class ClassifierSdk extends (EventEmitter as new () => TypedEmitter<Class
     deploymentId,
     affiliate,
     authentication,
-  }: ClassifierHelperOptions) {
+  }: ClassifierSdkOptions) {
     super();
     this.grpcAddress = grpcAddress;
     this.client = new ClassifierServiceClient(
@@ -202,46 +199,58 @@ export class ClassifierSdk extends (EventEmitter as new () => TypedEmitter<Class
    * @throws Error if the gRPC stream is not open.
    * @returns Promise that resolves when the request is sent.
    */
+  /**
+   * Sends one or more classify requests for images to the Athena service.
+   * @param request Single input or array of image classification request options.
+   * @throws Error if the gRPC stream is not open.
+   * @returns Promise that resolves when the request is sent.
+   */
   public async sendClassifyRequest(
-    options: ClassifyImageOptions,
+    request: ClassifyImageInput | ClassifyImageInput[],
   ): Promise<void> {
     if (!this.classifierGrpcCall) {
       throw new Error('gRPC stream is not open. Call open() first.');
     }
-    const {
-      affiliate = this.options.affiliate,
-      correlationId = randomUUID(),
-      imageStream,
-      encoding = RequestEncoding.UNCOMPRESSED,
-      format = ImageFormat.UNSPECIFIED,
-    } = options;
 
-    const { md5, sha1, data } = await computeHashesFromStream(
-      imageStream,
-      encoding,
-    );
+    const requests: ClassifyImageInput[] = Array.isArray(request)
+      ? request
+      : [request];
 
-    const hashes: ImageHash[] = [
-      { value: md5, type: HashType.MD5 },
-      { value: sha1, type: HashType.SHA1 },
-    ];
+    const processedInputs: ClassificationInput[] = [];
+    for (const options of requests) {
+      const {
+        affiliate = this.options.affiliate,
+        correlationId = randomUUID(),
+        imageStream,
+        encoding = RequestEncoding.UNCOMPRESSED,
+      } = options;
 
-    const input: ClassificationInput = {
-      affiliate,
-      correlationId,
-      encoding,
-      data,
-      format,
-      hashes,
-    };
+      const { md5, sha1, data, format } = await computeHashesFromStream(
+        imageStream,
+        encoding,
+      );
+      const hashes: ImageHash[] = [
+        { value: md5, type: HashType.MD5 },
+        { value: sha1, type: HashType.SHA1 },
+      ];
+
+      processedInputs.push({
+        affiliate,
+        correlationId,
+        encoding,
+        data,
+        format,
+        hashes,
+      });
+    }
 
-    const request: ClassifyRequest = {
+    const classifyRequest: ClassifyRequest = {
       deploymentId: this.options.deploymentId,
-      inputs: [input],
+      inputs: processedInputs,
     };
 
     await new Promise<void>((resolve) => {
-      if (this.classifierGrpcCall.write(request) == false) {
+      if (this.classifierGrpcCall.write(classifyRequest) == false) {
         this.classifierGrpcCall.once('drain', resolve);
       } else {
         resolve();
@@ -265,6 +274,6 @@ export class ClassifierSdk extends (EventEmitter as new () => TypedEmitter<Class
   }
 }
 
-export * from './athena/athena';
-export * from './athena/athena.grpc-client';
-export * from './hashing';
+export * from './athena/athena.js';
+export * from './athena/athena.grpc-client.js';
+export * from './hashing.js';