refactor(js): introduced background actions and reimplemented veo using them (#3070)

Co-authored-by: Michael Bleigh <[email protected]>

Author: pavelgj

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

@@ -49,6 +49,18 @@ export {
 // IMPORTANT: Keep this file in sync with genkit/ai/src/model.ts!
 //
 
+/**
+ * Zod schema of an opration representing a background task.
+ */
+export const OperationSchema = z.object({
+  action: z.string().optional(),
+  id: z.string(),
+  done: z.boolean().optional(),
+  output: z.any().optional(),
+  error: z.object({ message: z.string() }).passthrough().optional(),
+  metadata: z.record(z.string(), z.any()).optional(),
+});
+
 /**
  * Zod schema of message part.
  */
@@ -122,8 +134,6 @@ 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.
@@ -196,40 +206,6 @@ 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.
  */
@@ -243,9 +219,7 @@ 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,
@@ -299,6 +273,16 @@ 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(),
@@ -333,7 +317,7 @@ export const ModelResponseSchema = z.object({
   custom: z.unknown(),
   raw: z.unknown(),
   request: GenerateRequestSchema.optional(),
-  operation: ModelOperationSchema.optional(),
+  operation: OperationSchema.optional(),
 });
 
 /**

genkit-tools/genkit-schema.json

@@ -602,7 +602,6 @@
         "length",
         "blocked",
         "interrupted",
-        "pending",
         "other",
         "unknown"
       ]
@@ -733,9 +732,6 @@
             "$ref": "#/$defs/DocumentData"
           }
         },
-        "operation": {
-          "$ref": "#/$defs/ModelOperation"
-        },
         "candidates": {
           "type": "number"
         }
@@ -794,7 +790,7 @@
           "$ref": "#/$defs/GenerateRequest"
         },
         "operation": {
-          "$ref": "#/$defs/ModelOperation"
+          "$ref": "#/$defs/Operation"
         },
         "candidates": {
           "type": "array",
@@ -962,9 +958,6 @@
             },
             "toolChoice": {
               "type": "boolean"
-            },
-            "longRunning": {
-              "type": "boolean"
             }
           },
           "additionalProperties": false
@@ -982,53 +975,6 @@
       },
       "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": {
@@ -1049,9 +995,6 @@
         },
         "docs": {
           "$ref": "#/$defs/GenerateRequest/properties/docs"
-        },
-        "operation": {
-          "$ref": "#/$defs/GenerateRequest/properties/operation"
         }
       },
       "required": [
@@ -1119,6 +1062,41 @@
       ],
       "additionalProperties": false
     },
+    "Operation": {
+      "type": "object",
+      "properties": {
+        "action": {
+          "type": "string"
+        },
+        "id": {
+          "type": "string"
+        },
+        "done": {
+          "type": "boolean"
+        },
+        "output": {},
+        "error": {
+          "type": "object",
+          "properties": {
+            "message": {
+              "type": "string"
+            }
+          },
+          "required": [
+            "message"
+          ],
+          "additionalProperties": true
+        },
+        "metadata": {
+          "type": "object",
+          "additionalProperties": {}
+        }
+      },
+      "required": [
+        "id"
+      ],
+      "additionalProperties": false
+    },
     "OutputConfig": {
       "type": "object",
       "properties": {

js/ai/src/check-operation.ts

@@ -14,42 +14,27 @@
  * limitations under the License.
  */
 
-import { GenkitError } from '@genkit-ai/core';
+import { GenkitError, Operation } from '@genkit-ai/core';
 import { Registry } from '@genkit-ai/core/registry';
-import { GenerateRequest, ModelAction, ModelOperation } from './model';
 
-export async function checkOperation(
+export async function checkOperation<T = unknown>(
   registry: Registry,
-  operation: ModelOperation
-): Promise<ModelOperation> {
-  if (!operation.request?.model) {
+  operation: Operation<T>
+): Promise<Operation<T>> {
+  if (!operation.action) {
     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) {
+  const backgroundAction = await registry.lookupBackgroundAction(
+    operation.action
+  );
+  if (!backgroundAction) {
     throw new GenkitError({
       status: 'INVALID_ARGUMENT',
-      message: `Failed to resolve model from original request: ${operation.request?.model}`,
+      message: `Failed to resolve background action from original request: ${operation.action}`,
     });
   }
-  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,
-  };
+  return await backgroundAction.check(operation);
 }

js/ai/src/generate.ts

@@ -19,6 +19,7 @@ import {
   GenkitError,
   isAction,
   isDetachedAction,
+  Operation,
   runWithContext,
   runWithStreamingCallback,
   sentinelNoopStreamingCallback,
@@ -44,7 +45,7 @@ import { GenerateResponseChunk } from './generate/chunk.js';
 import { GenerateResponse } from './generate/response.js';
 import { Message } from './message.js';
 import {
-  ModelOperation,
+  GenerateResponseData,
   resolveModel,
   type GenerateActionOptions,
   type GenerateRequest,
@@ -342,9 +343,6 @@ 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(
@@ -365,7 +363,6 @@ 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,
       });
@@ -381,7 +378,7 @@ export async function generateOperation<
   options:
     | GenerateOptions<O, CustomOptions>
     | PromiseLike<GenerateOptions<O, CustomOptions>>
-): Promise<ModelOperation> {
+): Promise<Operation<GenerateResponseData>> {
   assertUnstable(registry, 'beta', 'generateOperation is a beta feature.');
 
   options = await options;

js/ai/src/generate/action.ts

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