feat(generate): handle moderation chunks (#58)

src/api-types.ts

@@ -84,6 +84,26 @@ export const GenerateStopReasonSchema = z.enum([
 ]);
 export type GenerateStopReason = z.infer<typeof GenerateStopReasonSchema>;
 
+const GenerateModerationSchema = z
+  .object({
+    hap: z.optional(
+      z.array(
+        z
+          .object({
+            success: z.boolean(),
+            flagged: z.boolean(),
+            score: z.number().min(0).max(1),
+            position: z.object({
+              start: z.number().int().min(0),
+              stop: z.number().int().min(0),
+            }),
+          })
+          .passthrough(),
+      ),
+    ),
+  })
+  .passthrough();
+
 export const GenerateResultSchema = z
   .object({
     generated_text: z.string(),
@@ -99,6 +119,7 @@ export const GenerateOutputSchema = z
     model_id: z.string(),
     created_at: z.coerce.date(),
     results: z.array(GenerateResultSchema),
+    moderation: GenerateModerationSchema.optional(),
   })
   .passthrough();
 export type GenerateOutput = z.infer<typeof GenerateOutputSchema>;

src/client/client.ts

@@ -462,13 +462,18 @@ export class Client {
               stop_reason = null,
               input_token_count = 0,
               generated_token_count = 0,
-            } = chunk.results[0];
+              ...props
+            } = (chunk.results || [{}])[0];
 
             callback(null, {
               generated_text,
               stop_reason,
               input_token_count,
               generated_token_count,
+              ...(chunk.moderation && {
+                moderation: chunk.moderation,
+              }),
+              ...props,
             } as GenerateOutput);
           } catch (e) {
             const err = (chunk || e) as unknown as Error;

src/client/types.ts

@@ -43,7 +43,9 @@ export const GenerateInputSchema = z.union([
   }),
 ]);
 export type GenerateInput = z.infer<typeof GenerateInputSchema>;
-export type GenerateOutput = ApiTypes.GenerateOutput['results'][number];
+export type GenerateOutput = ApiTypes.GenerateOutput['results'][number] & {
+  moderation?: ApiTypes.GenerateOutput['moderation'];
+};
 
 export const GenerateConfigInputSchema = ApiTypes.GenerateConfigInputSchema;
 export type GenerateConfigInput = z.input<typeof GenerateConfigInputSchema>;

src/tests/e2e/client.test.ts

@@ -55,13 +55,14 @@ describe('client', () => {
     }, 15_000);
 
     describe('streaming', () => {
-      const makeValidStream = () =>
+      const makeValidStream = (parameters: Record<string, any> = {}) =>
         client.generate(
           {
             model_id: 'google/ul2',
             input: 'Hello, World',
             parameters: {
               max_new_tokens: 10,
+              ...parameters,
             },
           },
           {
@@ -77,6 +78,10 @@ describe('client', () => {
         expect(chunk.generated_token_count).not.toBeNegative();
         expect(chunk.input_token_count).not.toBeNegative();
         expect(chunk.stop_reason).toSatisfy(isNumberOrNull);
+        expect(chunk.moderation).toBeOneOf([
+          undefined,
+          expect.objectContaining({ hap: expect.any(Array) }),
+        ]);
       };
 
       test('should throw for multiple inputs', () => {
@@ -99,6 +104,27 @@ describe('client', () => {
         ).toThrowError('Cannot do streaming for more than one input!');
       });
 
+      test('should correctly process moderation chunks during streaming', async () => {
+        const stream = makeValidStream({
+          moderations: {
+            min_new_tokens: 1,
+            max_new_tokens: 5,
+            hap: {
+              input: true,
+              threshold: 0.01,
+            },
+          },
+        });
+
+        for await (const chunk of stream) {
+          validateStreamChunk(chunk);
+          if (chunk.moderation) {
+            return;
+          }
+        }
+        throw Error('No moderation chunks has been retrieved from the API');
+      });
+
       test('should return valid stream for a single input', async () => {
         const stream = makeValidStream();