feat(js/plugins/googleai): added support for imagen and veo models to googleai plugin (#3024)

Author: pavelgj

.gitignore

@@ -42,6 +42,7 @@ js/testapps/personal-testing
 # Test files
 last_recording.mp4
 js/testapps/flow-simple-ai/*.wav
+js/testapps/flow-simple-ai/meme-of-the-day.mp4
 
 # auto-generated
 /js/core/src/__codegen

genkit-tools/common/src/types/model.ts

@@ -122,6 +122,8 @@ export const ModelInfoSchema = z.object({
       constrained: z.enum(['none', 'all', 'no-tools']).optional(),
       /** Model supports controlling tool choice, e.g. forced tool calling. */
       toolChoice: z.boolean().optional(),
+      /** Model supports long running operation interface. */
+      longRunning: z.boolean().optional(),
     })
     .optional(),
   /** At which stage of development this model is.
@@ -194,6 +196,40 @@ export const OutputConfigSchema = z.object({
   contentType: z.string().optional(),
 });
 
+/** Model response finish reason enum. */
+export const FinishReasonSchema = z.enum([
+  'stop',
+  'length',
+  'blocked',
+  'interrupted',
+  'pending',
+  'other',
+  'unknown',
+]);
+
+/**
+ * Zod schema of a long running operation.
+ */
+export const ModelOperationSchema = z
+  .object({
+    name: z.string(),
+    done: z.boolean().optional(),
+    request: z
+      .object({
+        model: z.string(),
+        config: z.record(z.string(), z.any()).optional(),
+      })
+      .optional(),
+    response: z
+      .object({
+        message: MessageSchema.optional(),
+        finishReason: FinishReasonSchema,
+        raw: z.unknown(),
+      })
+      .optional(),
+  })
+  .passthrough();
+
 /**
  * Output config.
  */
@@ -207,7 +243,9 @@ export const ModelRequestSchema = z.object({
   toolChoice: z.enum(['auto', 'required', 'none']).optional(),
   output: OutputConfigSchema.optional(),
   docs: z.array(DocumentDataSchema).optional(),
+  operation: ModelOperationSchema.optional(),
 });
+
 /** ModelRequest represents the parameters that are passed to a model when generating content. */
 export interface ModelRequest<
   CustomOptionsSchema extends z.ZodTypeAny = z.ZodTypeAny,
@@ -261,16 +299,6 @@ export const GenerationUsageSchema = z.object({
  */
 export type GenerationUsage = z.infer<typeof GenerationUsageSchema>;
 
-/** Model response finish reason enum. */
-export const FinishReasonSchema = z.enum([
-  'stop',
-  'length',
-  'blocked',
-  'interrupted',
-  'other',
-  'unknown',
-]);
-
 /** @deprecated All responses now return a single candidate. Only the first candidate will be used if supplied. */
 export const CandidateSchema = z.object({
   index: z.number(),
@@ -305,6 +333,7 @@ export const ModelResponseSchema = z.object({
   custom: z.unknown(),
   raw: z.unknown(),
   request: GenerateRequestSchema.optional(),
+  operation: ModelOperationSchema.optional(),
 });
 
 /**

genkit-tools/genkit-schema.json

@@ -602,6 +602,7 @@
         "length",
         "blocked",
         "interrupted",
+        "pending",
         "other",
         "unknown"
       ]
@@ -732,6 +733,9 @@
             "$ref": "#/$defs/DocumentData"
           }
         },
+        "operation": {
+          "$ref": "#/$defs/ModelOperation"
+        },
         "candidates": {
           "type": "number"
         }
@@ -789,6 +793,9 @@
         "request": {
           "$ref": "#/$defs/GenerateRequest"
         },
+        "operation": {
+          "$ref": "#/$defs/ModelOperation"
+        },
         "candidates": {
           "type": "array",
           "items": {
@@ -955,6 +962,9 @@
             },
             "toolChoice": {
               "type": "boolean"
+            },
+            "longRunning": {
+              "type": "boolean"
             }
           },
           "additionalProperties": false
@@ -972,6 +982,53 @@
       },
       "additionalProperties": false
     },
+    "ModelOperation": {
+      "type": "object",
+      "properties": {
+        "name": {
+          "type": "string"
+        },
+        "done": {
+          "type": "boolean"
+        },
+        "request": {
+          "type": "object",
+          "properties": {
+            "model": {
+              "type": "string"
+            },
+            "config": {
+              "type": "object",
+              "additionalProperties": {}
+            }
+          },
+          "required": [
+            "model"
+          ],
+          "additionalProperties": false
+        },
+        "response": {
+          "type": "object",
+          "properties": {
+            "message": {
+              "$ref": "#/$defs/Message"
+            },
+            "finishReason": {
+              "$ref": "#/$defs/FinishReason"
+            },
+            "raw": {}
+          },
+          "required": [
+            "finishReason"
+          ],
+          "additionalProperties": false
+        }
+      },
+      "required": [
+        "name"
+      ],
+      "additionalProperties": true
+    },
     "ModelRequest": {
       "type": "object",
       "properties": {
@@ -992,6 +1049,9 @@
         },
         "docs": {
           "$ref": "#/$defs/GenerateRequest/properties/docs"
+        },
+        "operation": {
+          "$ref": "#/$defs/GenerateRequest/properties/operation"
         }
       },
       "required": [
@@ -1049,6 +1109,9 @@
         },
         "request": {
           "$ref": "#/$defs/GenerateResponse/properties/request"
+        },
+        "operation": {
+          "$ref": "#/$defs/GenerateResponse/properties/operation"
         }
       },
       "required": [

js/ai/src/check-operation.ts

@@ -0,0 +1,55 @@
+/**
+ * Copyright 2024 Google LLC
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+import { GenkitError } from '@genkit-ai/core';
+import { Registry } from '@genkit-ai/core/registry';
+import { GenerateRequest, ModelAction, ModelOperation } from './model';
+
+export async function checkOperation(
+  registry: Registry,
+  operation: ModelOperation
+): Promise<ModelOperation> {
+  if (!operation.request?.model) {
+    throw new GenkitError({
+      status: 'INVALID_ARGUMENT',
+      message: 'Provided operation is missing original request information',
+    });
+  }
+  const model = (await registry.lookupAction(
+    `/model/${operation.request?.model}`
+  )) as ModelAction;
+  if (!model) {
+    throw new GenkitError({
+      status: 'INVALID_ARGUMENT',
+      message: `Failed to resolve model from original request: ${operation.request?.model}`,
+    });
+  }
+  const request = {
+    operation,
+    messages: [],
+  } as GenerateRequest;
+  const rawResponse = await model(request);
+  if (!rawResponse.operation) {
+    throw new GenkitError({
+      status: 'FAILED_PRECONDITION',
+      message: `The model did not return expected operation information: ${JSON.stringify(rawResponse)}`,
+    });
+  }
+  return {
+    ...rawResponse.operation!,
+    request: operation.request,
+  };
+}

js/ai/src/generate.ts

@@ -15,6 +15,7 @@
  */
 
 import {
+  assertUnstable,
   GenkitError,
   isAction,
   isDetachedAction,
@@ -43,6 +44,7 @@ import { GenerateResponseChunk } from './generate/chunk.js';
 import { GenerateResponse } from './generate/response.js';
 import { Message } from './message.js';
 import {
+  ModelOperation,
   resolveModel,
   type GenerateActionOptions,
   type GenerateRequest,
@@ -340,6 +342,9 @@ export async function generate<
   registry = maybeRegisterDynamicTools(registry, resolvedOptions);
 
   const params = await toGenerateActionOptions(registry, resolvedOptions);
+  const model = await resolveModel(registry, resolvedOptions.model, {
+    warnDeprecated: true,
+  });
 
   const tools = await toolsToActionRefs(registry, resolvedOptions.tools);
   return await runWithStreamingCallback(
@@ -360,13 +365,46 @@ export async function generate<
         tools,
       });
       return new GenerateResponse<O>(response, {
+        model: model.modelAction.__action.name,
         request: response.request ?? request,
         parser: resolvedFormat?.handler(request.output?.schema).parseMessage,
       });
     }
   );
 }
 
+export async function generateOperation<
+  O extends z.ZodTypeAny = z.ZodTypeAny,
+  CustomOptions extends z.ZodTypeAny = typeof GenerationCommonConfigSchema,
+>(
+  registry: Registry,
+  options:
+    | GenerateOptions<O, CustomOptions>
+    | PromiseLike<GenerateOptions<O, CustomOptions>>
+): Promise<ModelOperation> {
+  assertUnstable(registry, 'beta', 'generateOperation is a beta feature.');
+
+  options = await options;
+  const resolvedModel = await resolveModel(registry, options.model);
+  if (
+    !resolvedModel.modelAction.__action.metadata?.model.supports?.longRunning
+  ) {
+    throw new GenkitError({
+      status: 'INVALID_ARGUMENT',
+      message: `Model '${resolvedModel.modelAction.__action.name}' does not support long running operations.`,
+    });
+  }
+
+  const { operation } = await generate(registry, options);
+  if (!operation) {
+    throw new GenkitError({
+      status: 'FAILED_PRECONDITION',
+      message: `Model '${resolvedModel.modelAction.__action.name}' did not return an operation.`,
+    });
+  }
+  return operation;
+}
+
 function maybeRegisterDynamicTools<
   O extends z.ZodTypeAny = z.ZodTypeAny,
   CustomOptions extends z.ZodTypeAny = typeof GenerationCommonConfigSchema,

js/ai/src/generate/action.ts

@@ -322,7 +322,13 @@ async function generate(
         );
       };
 
-      return new GenerateResponse(await dispatch(0, request), {
+      const rawResponse = await dispatch(0, request);
+      if (!rawResponse.model) {
+        rawResponse.model = model.__action.name;
+      }
+
+      return new GenerateResponse(rawResponse, {
+        model: model.__action.name,
         request,
         parser: format?.handler(request.output?.schema).parseMessage,
       });
@@ -331,6 +337,11 @@ async function generate(
 
   // Throw an error if the response is not usable.
   response.assertValid();
+
+  if (response.operation) {
+    return response.toJSON();
+  }
+
   const generatedMessage = response.message!; // would have thrown if no message
 
   const toolRequests = generatedMessage.content.filter(