feat(js): add support for streaming json output

CONTRIBUTING.md

@@ -131,7 +131,7 @@ genkit eval:flow pdfQA '"What's a brief description of MapReduce?"'
 
 FYI: `js` and `genkit-tools` are in two separate workspaces.
 
-As you make changes you may want to build an test things by running test apps.
+As you make changes you may want to build and test things by running test apps.
 You can reduce the scope of what you're building by running a specific build command:
 
 ```

js/ai/package.json

@@ -15,7 +15,8 @@
     "build:clean": "rm -rf ./lib",
     "build": "npm-run-all build:clean check compile",
     "build:watch": "tsup-node --watch",
-    "test": "node --import tsx --test ./tests/**/*_test.ts"
+    "test": "node --import tsx --test ./tests/**/*_test.ts",
+    "test:single": "node --import tsx --test"
   },
   "repository": {
     "type": "git",
@@ -30,6 +31,7 @@
     "@types/node": "^20.11.19",
     "json5": "^2.2.3",
     "node-fetch": "^3.3.2",
+    "untruncate-json": "^0.0.1",
     "zod": "^3.22.4"
   },
   "devDependencies": {

js/ai/src/generate.ts

@@ -23,6 +23,7 @@ import {
 } from '@genkit-ai/core';
 import { lookupAction } from '@genkit-ai/core/registry';
 import { toJsonSchema, validateSchema } from '@genkit-ai/core/schema';
+import untruncateJson from 'untruncate-json';
 import { z } from 'zod';
 import { DocumentData } from './document.js';
 import { extractJson } from './extract.js';
@@ -357,11 +358,17 @@ export class GenerateResponseChunk<T = unknown>
   content: Part[];
   /** Custom model-specific data for this chunk. */
   custom?: unknown;
+  /** Accumulated chunks for partial output extraction. */
+  accumulatedChunks?: GenerateResponseChunkData[];
 
-  constructor(data: GenerateResponseChunkData) {
+  constructor(
+    data: GenerateResponseChunkData,
+    accumulatedChunks?: GenerateResponseChunkData[]
+  ) {
     this.index = data.index;
     this.content = data.content || [];
     this.custom = data.custom;
+    this.accumulatedChunks = accumulatedChunks;
   }
 
   /**
@@ -399,6 +406,18 @@ export class GenerateResponseChunk<T = unknown>
     ) as ToolRequestPart[];
   }
 
+  /**
+   * Attempts to extract the longest valid JSON substring from the accumulated chunks.
+   * @returns The longest valid JSON substring found in the accumulated chunks.
+   */
+  output(): T | null {
+    if (!this.accumulatedChunks) return null;
+    const accumulatedText = this.accumulatedChunks
+      .map((chunk) => chunk.content.map((part) => part.text || '').join(''))
+      .join('');
+    return extractJson(untruncateJson(accumulatedText));
+  }
+
   toJSON(): GenerateResponseChunkData {
     return { index: this.index, content: this.content, custom: this.custom };
   }
@@ -583,6 +602,7 @@ export class NoValidCandidatesError extends GenkitError {
  * @param options The options for this generation request.
  * @returns The generated response based on the provided parameters.
  */
+
 export async function generate<
   O extends z.ZodTypeAny = z.ZodTypeAny,
   CustomOptions extends z.ZodTypeAny = typeof GenerationCommonConfigSchema,
@@ -614,10 +634,20 @@ export async function generate<
     resolvedOptions,
     request
   );
+
+  const accumulatedChunks: GenerateResponseChunkData[] = [];
+
   const response = await runWithStreamingCallback(
     resolvedOptions.streamingCallback
-      ? (chunk: GenerateResponseChunkData) =>
-          resolvedOptions.streamingCallback!(new GenerateResponseChunk(chunk))
+      ? (chunk: GenerateResponseChunkData) => {
+          // Store accumulated chunk data
+          accumulatedChunks.push(chunk);
+          if (resolvedOptions.streamingCallback) {
+            resolvedOptions.streamingCallback!(
+              new GenerateResponseChunk(chunk, accumulatedChunks)
+            );
+          }
+        }
       : undefined,
     async () => new GenerateResponse<z.infer<O>>(await model(request), request)
   );