feat(go/genkit): added background action and model support (#3262)

Co-authored-by: Sahdev Garg <sahdevgarg@google.com>
Co-authored-by: Alex Pascal <apascal07@gmail.com>

Author: 3 people

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

@@ -147,6 +147,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 operations. */
+      longRunning: z.boolean().optional(),
     })
     .optional(),
   /** At which stage of development this model is.

genkit-tools/genkit-schema.json

@@ -1027,6 +1027,9 @@
             },
             "toolChoice": {
               "type": "boolean"
+            },
+            "longRunning": {
+              "type": "boolean"
             }
           },
           "additionalProperties": false

go/ai/background_model.go

@@ -0,0 +1,230 @@
+// Copyright 2025 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.
+//
+// SPDX-License-Identifier: Apache-2.0
+
+package ai
+
+import (
+	"context"
+	"errors"
+
+	"github.com/firebase/genkit/go/core"
+	"github.com/firebase/genkit/go/core/api"
+	"github.com/firebase/genkit/go/internal/registry"
+)
+
+// BackgroundModel represents a model that can run operations in the background.
+type BackgroundModel interface {
+	// Name returns the registry name of the background model.
+	Name() string
+	// Register registers the model with the given registry.
+	Register(r api.Registry)
+	// Start starts a background operation.
+	Start(ctx context.Context, req *ModelRequest) (*ModelOperation, error)
+	// Check checks the status of a background operation.
+	Check(ctx context.Context, op *ModelOperation) (*ModelOperation, error)
+	// Cancel cancels a background operation.
+	Cancel(ctx context.Context, op *ModelOperation) (*ModelOperation, error)
+	// SupportsCancel returns whether the background action supports cancellation.
+	SupportsCancel() bool
+}
+
+// backgroundModel is the concrete implementation of BackgroundModel interface.
+type backgroundModel struct {
+	core.BackgroundActionDef[*ModelRequest, *ModelResponse]
+}
+
+// ModelOperation is a background operation for a model.
+type ModelOperation = core.Operation[*ModelResponse]
+
+// StartModelOpFunc starts a background model operation.
+type StartModelOpFunc = func(ctx context.Context, req *ModelRequest) (*ModelOperation, error)
+
+// CheckOperationFunc checks the status of a background model operation.
+type CheckModelOpFunc = func(ctx context.Context, op *ModelOperation) (*ModelOperation, error)
+
+// CancelOperationFunc cancels a background model operation.
+type CancelModelOpFunc = func(ctx context.Context, op *ModelOperation) (*ModelOperation, error)
+
+// BackgroundModelOptions holds configuration for defining a background model
+type BackgroundModelOptions struct {
+	ModelOptions
+	Cancel   CancelModelOpFunc // Function that cancels a background model operation.
+	Metadata map[string]any    // Additional metadata.
+}
+
+// LookupBackgroundModel looks up a BackgroundAction registered by [DefineBackgroundModel].
+// It returns nil if the background model was not found.
+func LookupBackgroundModel(r api.Registry, name string) BackgroundModel {
+	key := api.KeyFromName(api.ActionTypeBackgroundModel, name)
+	action := core.LookupBackgroundAction[*ModelRequest, *ModelResponse](r, key)
+	if action == nil {
+		return nil
+	}
+	return &backgroundModel{*action}
+}
+
+// NewBackgroundModel defines a new model that runs in the background.
+func NewBackgroundModel(name string, opts *BackgroundModelOptions, startFn StartModelOpFunc, checkFn CheckModelOpFunc) BackgroundModel {
+	if name == "" {
+		panic("ai.NewBackgroundModel: name is required")
+	}
+	if startFn == nil {
+		panic("ai.NewBackgroundModel: startFn is required")
+	}
+	if checkFn == nil {
+		panic("ai.NewBackgroundModel: checkFn is required")
+	}
+
+	if opts == nil {
+		opts = &BackgroundModelOptions{}
+	}
+	if opts.Label == "" {
+		opts.Label = name
+	}
+	if opts.Supports == nil {
+		opts.Supports = &ModelSupports{}
+	}
+
+	metadata := map[string]any{
+		"type": api.ActionTypeBackgroundModel,
+		"model": map[string]any{
+			"label": opts.Label,
+			"supports": map[string]any{
+				"media":       opts.Supports.Media,
+				"context":     opts.Supports.Context,
+				"multiturn":   opts.Supports.Multiturn,
+				"systemRole":  opts.Supports.SystemRole,
+				"tools":       opts.Supports.Tools,
+				"toolChoice":  opts.Supports.ToolChoice,
+				"constrained": opts.Supports.Constrained,
+				"output":      opts.Supports.Output,
+				"contentType": opts.Supports.ContentType,
+				"longRunning": opts.Supports.LongRunning,
+			},
+			"versions":      opts.Versions,
+			"stage":         opts.Stage,
+			"customOptions": opts.ConfigSchema,
+		},
+	}
+
+	inputSchema := core.InferSchemaMap(ModelRequest{})
+	if inputSchema != nil && opts.ConfigSchema != nil {
+		if props, ok := inputSchema["properties"].(map[string]any); ok {
+			props["config"] = opts.ConfigSchema
+		}
+	}
+
+	mws := []ModelMiddleware{
+		simulateSystemPrompt(&opts.ModelOptions, nil),
+		augmentWithContext(&opts.ModelOptions, nil),
+		validateSupport(name, &opts.ModelOptions),
+		addAutomaticTelemetry(),
+	}
+	fn := core.ChainMiddleware(mws...)(backgroundModelToModelFn(startFn))
+
+	wrappedFn := func(ctx context.Context, req *ModelRequest) (*ModelOperation, error) {
+		resp, err := fn(ctx, req, nil)
+		if err != nil {
+			return nil, err
+		}
+
+		return modelOpFromResponse(resp)
+	}
+
+	return &backgroundModel{*core.NewBackgroundAction(name, api.ActionTypeBackgroundModel, metadata, wrappedFn, checkFn, opts.Cancel)}
+}
+
+// DefineBackgroundModel defines and registers a new model that runs in the background.
+func DefineBackgroundModel(r *registry.Registry, name string, opts *BackgroundModelOptions, fn StartModelOpFunc, checkFn CheckModelOpFunc) BackgroundModel {
+	m := NewBackgroundModel(name, opts, fn, checkFn)
+	m.Register(r)
+	return m
+}
+
+// GenerateOperation generates a model response as a long-running operation based on the provided options.
+func GenerateOperation(ctx context.Context, r *registry.Registry, opts ...GenerateOption) (*ModelOperation, error) {
+	resp, err := Generate(ctx, r, opts...)
+	if err != nil {
+		return nil, err
+	}
+
+	return modelOpFromResponse(resp)
+}
+
+// CheckModelOperation checks the status of a background model operation by looking up the model and calling its Check method.
+func CheckModelOperation(ctx context.Context, r api.Registry, op *ModelOperation) (*ModelOperation, error) {
+	return core.CheckOperation[*ModelRequest](ctx, r, op)
+}
+
+// backgroundModelToModelFn wraps a background model start function into a [ModelFunc] for middleware compatibility.
+func backgroundModelToModelFn(startFn StartModelOpFunc) ModelFunc {
+	return func(ctx context.Context, req *ModelRequest, cb ModelStreamCallback) (*ModelResponse, error) {
+		op, err := startFn(ctx, req)
+		if err != nil {
+			return nil, err
+		}
+
+		var opError *OperationError
+		if op.Error != nil {
+			opError = &OperationError{Message: op.Error.Error()}
+		}
+
+		metadata := op.Metadata
+		if metadata == nil {
+			metadata = make(map[string]any)
+		}
+
+		return &ModelResponse{
+			Operation: &Operation{
+				Action:   op.Action,
+				Id:       op.ID,
+				Done:     op.Done,
+				Output:   op.Output,
+				Error:    opError,
+				Metadata: metadata,
+			},
+			Request: req,
+		}, nil
+	}
+}
+
+// modelOpFromResponse extracts a [ModelOperation] from a [ModelResponse].
+func modelOpFromResponse(resp *ModelResponse) (*ModelOperation, error) {
+	if resp.Operation == nil {
+		return nil, core.NewError(core.FAILED_PRECONDITION, "background model did not return an operation")
+	}
+
+	op := &ModelOperation{
+		Action:   resp.Operation.Action,
+		ID:       resp.Operation.Id,
+		Done:     resp.Operation.Done,
+		Metadata: resp.Operation.Metadata,
+	}
+
+	if resp.Operation.Error != nil {
+		op.Error = errors.New(resp.Operation.Error.Message)
+	}
+
+	if resp.Operation.Output != nil {
+		if modelResp, ok := resp.Operation.Output.(*ModelResponse); ok {
+			op.Output = modelResp
+		} else {
+			return nil, core.NewError(core.INTERNAL, "operation output is not a model response")
+		}
+	}
+
+	return op, nil
+}

go/ai/gen.go

@@ -223,6 +223,7 @@ type ModelSupports struct {
 	Constrained ConstrainedSupport `json:"constrained,omitempty"`
 	ContentType []string           `json:"contentType,omitempty"`
 	Context     bool               `json:"context,omitempty"`
+	LongRunning bool               `json:"longRunning,omitempty"`
 	Media       bool               `json:"media,omitempty"`
 	Multiturn   bool               `json:"multiturn,omitempty"`
 	Output      []string           `json:"output,omitempty"`
@@ -257,8 +258,10 @@ type ModelResponse struct {
 	FinishMessage string       `json:"finishMessage,omitempty"`
 	FinishReason  FinishReason `json:"finishReason,omitempty"`
 	// LatencyMs is the time the request took in milliseconds.
-	LatencyMs float64  `json:"latencyMs,omitempty"`
-	Message   *Message `json:"message,omitempty"`
+	LatencyMs float64    `json:"latencyMs,omitempty"`
+	Message   *Message   `json:"message,omitempty"`
+	Operation *Operation `json:"operation,omitempty"`
+	Raw       any        `json:"raw,omitempty"`
 	// Request is the [ModelRequest] struct used to trigger this response.
 	Request *ModelRequest `json:"request,omitempty"`
 	// Usage describes how many resources were used by this generation request.
@@ -275,6 +278,19 @@ type ModelResponseChunk struct {
 	Role       Role    `json:"role,omitempty"`
 }
 
+type Operation struct {
+	Action   string          `json:"action,omitempty"`
+	Done     bool            `json:"done,omitempty"`
+	Error    *OperationError `json:"error,omitempty"`
+	Id       string          `json:"id,omitempty"`
+	Metadata map[string]any  `json:"metadata,omitempty"`
+	Output   any             `json:"output,omitempty"`
+}
+
+type OperationError struct {
+	Message string `json:"message,omitempty"`
+}
+
 // OutputConfig describes the structure that the model's output
 // should conform to. If Format is [OutputFormatJSON], then Schema
 // can describe the desired form of the generated JSON.