From a8d54e885934cebd59e8ec718a2a00f3318fc7e8 Mon Sep 17 00:00:00 2001 From: Alex Pascal Date: Wed, 29 Jul 2026 15:33:14 -0700 Subject: [PATCH 01/10] feat(go/ai): add typed-config New*WithConfig constructors NewModelWithConfig, NewEmbedderWithConfig, NewEvaluatorWithConfig, NewBatchEvaluatorWithConfig, and NewBackgroundModelWithConfig take a Config type parameter inferred from the provided function's signature. The framework validates the request's raw config against the config schema, deserializes it into the typed value (exact type, *Config, or map[string]any; anything else is INVALID_ARGUMENT), normalizes the request's type-erased config slot to the converted value, and passes the typed value to the function. This removes the per-provider configFromRequest boilerplate once plugins migrate. The config JSON schema is inferred from Config unless ConfigSchema overrides it (kept for types whose reflection is insufficient, e.g. SDK wrapper generics like Opt[float64]). The request input schema's config slot accepts an explicit JSON null so typed-nil configs pass input validation. Config resolution runs outermost in the model's built-in chain; version validation moves there from validateSupport because it must see the raw, pre-conversion config. Batch evaluators now advertise an inferred input schema. The untyped New*/Define* constructors are deprecated and delegate with Config=any, which passes the raw config through unchanged (pinned by TestDeprecatedModelPassesConfigThrough). Retrievers are left untouched since the subsystem is slated for removal. --- go/ai/background_model.go | 66 +++++-- go/ai/config.go | 109 ++++++++++++ go/ai/config_test.go | 312 +++++++++++++++++++++++++++++++++ go/ai/embedder.go | 58 +++++- go/ai/evaluator.go | 153 ++++++++++++---- go/ai/generate.go | 71 ++++++-- go/ai/model_middleware.go | 7 +- go/ai/model_middleware_test.go | 14 +- 8 files changed, 712 insertions(+), 78 deletions(-) create mode 100644 go/ai/config.go create mode 100644 go/ai/config_test.go diff --git a/go/ai/background_model.go b/go/ai/background_model.go index 21e3817b95..d4e914176d 100644 --- a/go/ai/background_model.go +++ b/go/ai/background_model.go @@ -52,6 +52,12 @@ type ModelOperation = core.Operation[*ModelResponse] // StartModelOpFunc starts a background model operation. type StartModelOpFunc = func(ctx context.Context, req *ModelRequest) (*ModelOperation, error) +// StartModelOpFuncWithConfig is a [StartModelOpFunc] that additionally +// receives the request's typed Config: the framework deserializes the +// request's raw config into it before calling the function (see +// [NewBackgroundModelWithConfig]). +type StartModelOpFuncWithConfig[Config any] = func(ctx context.Context, req *ModelRequest, config Config) (*ModelOperation, error) + // CheckModelOpFunc checks the status of a background model operation. type CheckModelOpFunc = func(ctx context.Context, op *ModelOperation) (*ModelOperation, error) @@ -76,16 +82,21 @@ func LookupBackgroundModel(r api.Registry, name string) BackgroundModel { return &backgroundModel{*action} } -// NewBackgroundModel defines a new model that runs in the background. -func NewBackgroundModel(name string, opts *BackgroundModelOptions, startFn StartModelOpFunc, checkFn CheckModelOpFunc) BackgroundModel { +// NewBackgroundModelWithConfig creates a new [BackgroundModel]. Register it +// with [BackgroundModel.Register] to make it resolvable by name. +// +// Config is the model's typed configuration; it is usually inferred from +// startFn's signature. See [NewModelWithConfig] for how the request's config +// is deserialized. +func NewBackgroundModelWithConfig[Config any](name string, opts *BackgroundModelOptions, startFn StartModelOpFuncWithConfig[Config], checkFn CheckModelOpFunc) BackgroundModel { if name == "" { - panic("ai.NewBackgroundModel: name is required") + panic("ai.NewBackgroundModelWithConfig: name is required") } if startFn == nil { - panic("ai.NewBackgroundModel: startFn is required") + panic("ai.NewBackgroundModelWithConfig: startFn is required") } if checkFn == nil { - panic("ai.NewBackgroundModel: checkFn is required") + panic("ai.NewBackgroundModelWithConfig: checkFn is required") } if opts == nil { @@ -98,6 +109,8 @@ func NewBackgroundModel(name string, opts *BackgroundModelOptions, startFn Start opts.Supports = &ModelSupports{} } + configSchema, inputSchema := actionConfigSchemas[Config](opts.ConfigSchema, ModelRequest{}, "config") + metadata := map[string]any{ "type": api.ActionTypeBackgroundModel, "model": map[string]any{ @@ -116,23 +129,28 @@ func NewBackgroundModel(name string, opts *BackgroundModelOptions, startFn Start }, "versions": opts.Versions, "stage": opts.Stage, - "customOptions": opts.ConfigSchema, + "customOptions": configSchema, }, } - inputSchema := core.InferSchemaMap(ModelRequest{}) - if inputSchema != nil && opts.ConfigSchema != nil { - if props, ok := inputSchema["properties"].(map[string]any); ok { - props["config"] = opts.ConfigSchema + typedStartFn := func(ctx context.Context, req *ModelRequest) (*ModelOperation, error) { + // req.Config was normalized to the exact Config type by + // normalizeConfig below, so this hits the fast path. + cfg, err := resolveConfig[Config](req.Config) + if err != nil { + return nil, err } + return startFn(ctx, req, cfg) } - mws := []ModelMiddleware{ + // normalizeConfig runs outermost so that the built-in wrappers and the + // start function all see the typed, converted config on the request. + fn := core.ChainMiddleware( + normalizeConfig[Config](name, opts.Versions), simulateSystemPrompt(&opts.ModelOptions, nil), augmentWithContext(&opts.ModelOptions, nil), validateSupport(name, &opts.ModelOptions), - } - fn := core.ChainMiddleware(mws...)(backgroundModelToModelFn(startFn)) + )(backgroundModelToModelFn(typedStartFn)) wrappedFn := func(ctx context.Context, req *ModelRequest) (*ModelOperation, error) { resp, err := fn(ctx, req, nil) @@ -143,10 +161,30 @@ func NewBackgroundModel(name string, opts *BackgroundModelOptions, startFn Start return modelOpFromResponse(resp) } - return &backgroundModel{*core.NewBackgroundAction(name, api.ActionTypeBackgroundModel, metadata, wrappedFn, checkFn, opts.Cancel)} + return &backgroundModel{*core.NewBackgroundActionOf(api.ActionTypeBackgroundModel, name, &core.BackgroundActionOptions{ + Metadata: metadata, + InputSchema: inputSchema, + }, wrappedFn, checkFn, opts.Cancel)} +} + +// NewBackgroundModel defines a new model that runs in the background. +// +// Deprecated: Use [NewBackgroundModelWithConfig], which passes the request's +// config to startFn as a typed value instead of leaving it type-erased on the +// request. +func NewBackgroundModel(name string, opts *BackgroundModelOptions, startFn StartModelOpFunc, checkFn CheckModelOpFunc) BackgroundModel { + if startFn == nil { + panic("ai.NewBackgroundModel: startFn is required") + } + return NewBackgroundModelWithConfig(name, opts, func(ctx context.Context, req *ModelRequest, _ any) (*ModelOperation, error) { + return startFn(ctx, req) + }, checkFn) } // DefineBackgroundModel defines and registers a new model that runs in the background. +// +// Deprecated: Use [NewBackgroundModelWithConfig] and register the result with +// [BackgroundModel.Register]. func DefineBackgroundModel(r *registry.Registry, name string, opts *BackgroundModelOptions, fn StartModelOpFunc, checkFn CheckModelOpFunc) BackgroundModel { m := NewBackgroundModel(name, opts, fn, checkFn) m.Register(r) diff --git a/go/ai/config.go b/go/ai/config.go new file mode 100644 index 0000000000..6775b2561d --- /dev/null +++ b/go/ai/config.go @@ -0,0 +1,109 @@ +// Copyright 2026 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" + "fmt" + + "github.com/firebase/genkit/go/core" + "github.com/firebase/genkit/go/internal/base" +) + +// This file holds the typed-config plumbing shared by models, embedders, and +// evaluators. Requests carry config as `any` on the wire; the +// Define*/New*WithConfig constructors wrap the user's typed function so that +// the raw value is deserialized into the Config type parameter before the +// function runs, and the request's type-erased config slot is normalized to +// that same converted value so the two views never disagree. + +// nullableConfigSchema wraps a config schema for the request input-schema +// slot so that an explicit JSON null is accepted on the wire: a typed-nil Go +// config marshals to null and resolves to the zero Config value, so it must +// not be rejected by input validation. The advertised config schema (the +// customOptions metadata) stays unwrapped. +func nullableConfigSchema(schema map[string]any) map[string]any { + if schema == nil { + return nil + } + return map[string]any{"anyOf": []any{schema, map[string]any{"type": "null"}}} +} + +// normalizeConfig returns a model middleware that resolves the request's raw +// config into Config and writes the converted value back into the request. +// It runs as the outermost step of a model's built-in chain so that every +// wrapper after it, and the model function itself, sees the typed value; by +// then the config has already been validated against the model's config +// schema at the action boundary. +// +// Version validation runs here, against the raw config, because conversion is +// lossy: a "version" key sent by a JSON caller would be silently dropped when +// deserializing into a Config type that has no such field. +func normalizeConfig[Config any](model string, versions []string) ModelMiddleware { + return func(next ModelFunc) ModelFunc { + return func(ctx context.Context, req *ModelRequest, cb ModelStreamCallback) (*ModelResponse, error) { + if err := validateVersion(model, versions, req.Config); err != nil { + return nil, err + } + cfg, err := resolveConfig[Config](req.Config) + if err != nil { + return nil, err + } + req.Config = cfg + return next(ctx, req, cb) + } + } +} + +// actionConfigSchemas returns the action's effective config schema (the +// explicit override when set, otherwise the schema inferred from Config) and +// the request input schema with its config slot replaced by the null-tolerant +// wrapping of that schema. reqZero is the zero request value to infer the +// input schema from and key is the config slot's wire name ("config" for +// models, "options" for embedders and evaluators). +func actionConfigSchemas[Config any](override map[string]any, reqZero any, key string) (configSchema, inputSchema map[string]any) { + configSchema = override + if configSchema == nil { + configSchema = base.SchemaMapFor[Config]() + } + + inputSchema = core.InferSchemaMap(reqZero) + if inputSchema != nil && configSchema != nil { + if props, ok := inputSchema["properties"].(map[string]any); ok { + props[key] = nullableConfigSchema(configSchema) + } + } + return configSchema, inputSchema +} + +// resolveConfig converts the raw config value carried by a request into the +// typed Config the action was defined with. It accepts the exact Config type +// (or a pointer to it, which is dereferenced), a map[string]any (as sent by +// the Dev UI and other JSON callers, deserialized via a JSON round-trip), or +// nil (yielding the zero value). Any other type is rejected so one provider's +// config cannot be silently passed to another provider's action. +func resolveConfig[Config any](raw any) (Config, error) { + cfg, err := base.ConvertToExact[Config](raw) + if err != nil { + if errors.Is(err, base.ErrTypeMismatch) { + return cfg, core.NewPublicError(core.INVALID_ARGUMENT, fmt.Sprintf("invalid config type %T, want %T or map[string]any", raw, cfg), nil) + } + return cfg, core.NewPublicError(core.INVALID_ARGUMENT, fmt.Sprintf("invalid config for %T; check that field names and value types match: %v", cfg, err), nil) + } + return cfg, nil +} diff --git a/go/ai/config_test.go b/go/ai/config_test.go new file mode 100644 index 0000000000..1e69536b4c --- /dev/null +++ b/go/ai/config_test.go @@ -0,0 +1,312 @@ +// Copyright 2026 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" + "strings" + "testing" + + "github.com/firebase/genkit/go/core/api" + "github.com/firebase/genkit/go/internal/registry" +) + +// testTypedConfig is a provider-style config struct used to exercise the +// typed-config deserialization that New*WithConfig wraps around user +// functions. +type testTypedConfig struct { + Temperature float64 `json:"temperature,omitempty"` + MaxTokens int `json:"maxTokens,omitempty"` +} + +// otherProviderConfig stands in for a different provider's config type. Its +// JSON shape is compatible with testTypedConfig so that requests carrying it +// pass the action's input schema validation and rejection happens on the Go +// type itself. +type otherProviderConfig struct { + Temperature float64 `json:"temperature,omitempty"` +} + +func TestModelTypedConfig(t *testing.T) { + r := registry.New() + + var got testTypedConfig + var gotReqConfig any + m := NewModelWithConfig("test/typed-config", nil, func(ctx context.Context, req *ModelRequest, cfg testTypedConfig, cb ModelStreamCallback) (*ModelResponse, error) { + got = cfg + gotReqConfig = req.Config + return &ModelResponse{Message: NewModelTextMessage("ok"), Request: req}, nil + }) + m.Register(r) + + tests := []struct { + name string + config any + want testTypedConfig + wantErr string + }{ + {name: "nil config yields zero value", config: nil, want: testTypedConfig{}}, + {name: "exact type", config: testTypedConfig{Temperature: 0.5, MaxTokens: 10}, want: testTypedConfig{Temperature: 0.5, MaxTokens: 10}}, + {name: "pointer to exact type", config: &testTypedConfig{Temperature: 0.7}, want: testTypedConfig{Temperature: 0.7}}, + {name: "map is deserialized", config: map[string]any{"temperature": 0.9, "maxTokens": 5}, want: testTypedConfig{Temperature: 0.9, MaxTokens: 5}}, + {name: "mismatched struct type is rejected", config: otherProviderConfig{Temperature: 0.2}, wantErr: "invalid config type"}, + {name: "mismatched pointer type is rejected", config: &otherProviderConfig{Temperature: 0.2}, wantErr: "invalid config type"}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + got = testTypedConfig{} + req := &ModelRequest{ + Messages: []*Message{NewUserTextMessage("hi")}, + Config: tt.config, + } + _, err := m.Generate(context.Background(), req, nil) + if tt.wantErr != "" { + if err == nil || !strings.Contains(err.Error(), tt.wantErr) { + t.Fatalf("Generate() error = %v, want error containing %q", err, tt.wantErr) + } + return + } + if err != nil { + t.Fatalf("Generate() unexpected error: %v", err) + } + if got != tt.want { + t.Errorf("config = %+v, want %+v", got, tt.want) + } + if gotReqConfig != any(tt.want) { + t.Errorf("req.Config = %#v, want normalized to %#v", gotReqConfig, tt.want) + } + }) + } +} + +// TestModelConfigNormalizedBeforeBuiltins pins the boundary ordering: config +// resolution runs outermost in the model's built-in chain, so a mismatched +// config type is reported before capability validation (which would otherwise +// reject this request for containing unsupported media). +func TestModelConfigNormalizedBeforeBuiltins(t *testing.T) { + r := registry.New() + + m := NewModelWithConfig("test/config-order", &ModelOptions{ + Supports: &ModelSupports{Multiturn: true}, // no media support + }, func(ctx context.Context, req *ModelRequest, cfg testTypedConfig, cb ModelStreamCallback) (*ModelResponse, error) { + return &ModelResponse{Message: NewModelTextMessage("ok"), Request: req}, nil + }) + m.Register(r) + + req := &ModelRequest{ + Messages: []*Message{NewUserMessage(NewMediaPart("image/png", "data:image/png;base64,aGk="))}, + Config: otherProviderConfig{Temperature: 0.2}, + } + _, err := m.Generate(context.Background(), req, nil) + if err == nil || !strings.Contains(err.Error(), "invalid config type") { + t.Fatalf("Generate() error = %v, want config type error before media support error", err) + } +} + +func TestBackgroundModelConfigValidation(t *testing.T) { + r := registry.New() + + bm := NewBackgroundModelWithConfig("test/bg-typed-config", nil, + func(ctx context.Context, req *ModelRequest, cfg testTypedConfig) (*ModelOperation, error) { + return &ModelOperation{ID: "op1", Done: false}, nil + }, + func(ctx context.Context, op *ModelOperation) (*ModelOperation, error) { + return op, nil + }) + bm.Register(r) + + // A config that violates the inferred schema is rejected by input + // validation at the action boundary, before conversion. + req := &ModelRequest{ + Messages: []*Message{NewUserTextMessage("hi")}, + Config: map[string]any{"unknownOption": true}, + } + if _, err := bm.Start(context.Background(), req); err == nil || !strings.Contains(err.Error(), "schema") { + t.Fatalf("Start() error = %v, want schema validation error", err) + } + + // A valid map config passes validation and is converted. + req = &ModelRequest{ + Messages: []*Message{NewUserTextMessage("hi")}, + Config: map[string]any{"temperature": 0.5}, + } + if _, err := bm.Start(context.Background(), req); err != nil { + t.Fatalf("Start() unexpected error: %v", err) + } +} + +func TestModelConfigSchemaInference(t *testing.T) { + newFn := func() ModelFuncWithConfig[testTypedConfig] { + return func(ctx context.Context, req *ModelRequest, cfg testTypedConfig, cb ModelStreamCallback) (*ModelResponse, error) { + return &ModelResponse{Message: NewModelTextMessage("ok"), Request: req}, nil + } + } + + configSchemaOf := func(t *testing.T, m Model) any { + t.Helper() + desc := m.(api.Action).Desc() + modelMeta, ok := desc.Metadata["model"].(map[string]any) + if !ok { + t.Fatalf("missing model metadata: %+v", desc.Metadata) + } + return modelMeta["customOptions"] + } + + t.Run("inferred from Config type", func(t *testing.T) { + m := NewModelWithConfig("test/inferred-schema", nil, newFn()) + schema, ok := configSchemaOf(t, m).(map[string]any) + if !ok { + t.Fatalf("customOptions = %v, want inferred schema map", configSchemaOf(t, m)) + } + props, ok := schema["properties"].(map[string]any) + if !ok { + t.Fatalf("schema has no properties: %v", schema) + } + if _, ok := props["temperature"]; !ok { + t.Errorf("inferred schema missing temperature property: %v", props) + } + }) + + t.Run("explicit ConfigSchema wins", func(t *testing.T) { + override := map[string]any{"type": "object", "properties": map[string]any{"custom": map[string]any{"type": "string"}}} + m := NewModelWithConfig("test/override-schema", &ModelOptions{ConfigSchema: override}, newFn()) + schema, ok := configSchemaOf(t, m).(map[string]any) + if !ok { + t.Fatalf("customOptions missing") + } + props := schema["properties"].(map[string]any) + if _, ok := props["custom"]; !ok { + t.Errorf("override schema not used: %v", schema) + } + }) + + t.Run("any config infers no schema", func(t *testing.T) { + m := NewModelWithConfig("test/any-schema", nil, func(ctx context.Context, req *ModelRequest, cfg any, cb ModelStreamCallback) (*ModelResponse, error) { + return &ModelResponse{Message: NewModelTextMessage("ok"), Request: req}, nil + }) + if s, _ := configSchemaOf(t, m).(map[string]any); len(s) != 0 { + t.Errorf("customOptions = %v, want no inferred schema for any config", s) + } + }) + + t.Run("deprecated NewModel behaves like any config", func(t *testing.T) { + m := NewModel("test/legacy-schema", nil, func(ctx context.Context, req *ModelRequest, cb ModelStreamCallback) (*ModelResponse, error) { + return &ModelResponse{Message: NewModelTextMessage("ok"), Request: req}, nil + }) + if s, _ := configSchemaOf(t, m).(map[string]any); len(s) != 0 { + t.Errorf("customOptions = %v, want no inferred schema for legacy models", s) + } + }) +} + +// TestDeprecatedModelPassesConfigThrough pins the compatibility contract of +// the deprecated untyped constructors: any config value, including another +// provider's struct, still reaches the model function unconverted. +func TestDeprecatedModelPassesConfigThrough(t *testing.T) { + r := registry.New() + + var gotConfig any + m := DefineModel(r, "test/legacy-config", nil, func(ctx context.Context, req *ModelRequest, cb ModelStreamCallback) (*ModelResponse, error) { + gotConfig = req.Config + return &ModelResponse{Message: NewModelTextMessage("ok"), Request: req}, nil + }) + + cfg := otherProviderConfig{Temperature: 0.4} + req := &ModelRequest{ + Messages: []*Message{NewUserTextMessage("hi")}, + Config: cfg, + } + if _, err := m.Generate(context.Background(), req, nil); err != nil { + t.Fatalf("Generate() unexpected error: %v", err) + } + if gotConfig != any(cfg) { + t.Errorf("req.Config = %#v, want passed through as %#v", gotConfig, cfg) + } +} + +func TestEmbedderTypedConfig(t *testing.T) { + r := registry.New() + + var got testTypedConfig + var gotReqOptions any + e := NewEmbedderWithConfig("test/typed-config-embedder", nil, func(ctx context.Context, req *EmbedRequest, cfg testTypedConfig) (*EmbedResponse, error) { + got = cfg + gotReqOptions = req.Options + return &EmbedResponse{}, nil + }) + e.Register(r) + + req := &EmbedRequest{Options: map[string]any{"maxTokens": 7}} + if _, err := e.Embed(context.Background(), req); err != nil { + t.Fatalf("Embed() unexpected error: %v", err) + } + if got.MaxTokens != 7 { + t.Errorf("config = %+v, want MaxTokens 7", got) + } + if gotReqOptions != any(got) { + t.Errorf("req.Options = %#v, want normalized to %#v", gotReqOptions, got) + } + + req = &EmbedRequest{Options: otherProviderConfig{Temperature: 0.1}} + if _, err := e.Embed(context.Background(), req); err == nil || !strings.Contains(err.Error(), "invalid config type") { + t.Fatalf("Embed() error = %v, want invalid config type error", err) + } +} + +func TestEvaluatorTypedConfig(t *testing.T) { + r := registry.New() + + var got testTypedConfig + e := NewEvaluatorWithConfig("test/typed-config-evaluator", nil, func(ctx context.Context, req *EvaluatorCallbackRequest, cfg testTypedConfig) (*EvaluatorCallbackResponse, error) { + got = cfg + return &EvaluatorCallbackResponse{ + TestCaseId: req.Input.TestCaseId, + Evaluation: []Score{{Id: "s", Score: 1, Status: ScoreStatusPass.String()}}, + }, nil + }) + e.Register(r) + + req := &EvaluatorRequest{ + Dataset: []*Example{{TestCaseId: "tc1", Input: "in"}}, + EvaluationId: "run1", + Options: &testTypedConfig{Temperature: 0.3}, + } + if _, err := e.Evaluate(context.Background(), req); err != nil { + t.Fatalf("Evaluate() unexpected error: %v", err) + } + if got.Temperature != 0.3 { + t.Errorf("config = %+v, want Temperature 0.3", got) + } + + var gotBatch testTypedConfig + be := NewBatchEvaluatorWithConfig("test/typed-config-batch-evaluator", nil, func(ctx context.Context, req *EvaluatorRequest, cfg testTypedConfig) (*EvaluatorResponse, error) { + gotBatch = cfg + return &EvaluatorResponse{}, nil + }) + be.Register(r) + breq := &EvaluatorRequest{ + Dataset: []*Example{{TestCaseId: "tc1", Input: "in"}}, + EvaluationId: "run2", + Options: map[string]any{"temperature": 0.8}, + } + if _, err := be.Evaluate(context.Background(), breq); err != nil { + t.Fatalf("Evaluate() unexpected error: %v", err) + } + if gotBatch.Temperature != 0.8 { + t.Errorf("config = %+v, want Temperature 0.8", gotBatch) + } +} diff --git a/go/ai/embedder.go b/go/ai/embedder.go index 5801da9592..e43505e68c 100644 --- a/go/ai/embedder.go +++ b/go/ai/embedder.go @@ -27,6 +27,11 @@ import ( // EmbedderFunc is the function type for embedding documents. type EmbedderFunc = func(context.Context, *EmbedRequest) (*EmbedResponse, error) +// EmbedderFuncWithConfig is an [EmbedderFunc] that additionally receives the +// request's typed Config: the framework deserializes the request's raw +// options into it before calling the function (see [NewEmbedderWithConfig]). +type EmbedderFuncWithConfig[Config any] = func(context.Context, *EmbedRequest, Config) (*EmbedResponse, error) + // Embedder represents an embedder that can perform content embedding. type Embedder interface { // Name returns the registry name of the embedder. @@ -88,10 +93,20 @@ type embedder struct { core.Action[*EmbedRequest, *EmbedResponse, struct{}] } -// NewEmbedder creates a new [Embedder]. -func NewEmbedder(name string, opts *EmbedderOptions, fn EmbedderFunc) Embedder { +// NewEmbedderWithConfig creates a new [Embedder]. Register it with +// [Embedder.Register] to make it resolvable by name. +// +// Config is the embedder's typed configuration; it is usually inferred from +// fn's signature. The framework deserializes the request's raw options into +// Config before calling fn: the exact Config type (or a pointer to it) and +// map[string]any (from the Dev UI and other JSON callers) are accepted, and +// mismatched types are rejected. The request's [EmbedRequest.Options] is +// normalized to the converted value, so it always matches the typed +// parameter. The config's JSON schema is inferred from Config unless +// [EmbedderOptions.ConfigSchema] overrides it. +func NewEmbedderWithConfig[Config any](name string, opts *EmbedderOptions, fn EmbedderFuncWithConfig[Config]) Embedder { if name == "" { - panic("ai.NewEmbedder: name is required") + panic("ai.NewEmbedderWithConfig: name is required") } if opts == nil { @@ -103,6 +118,8 @@ func NewEmbedder(name string, opts *EmbedderOptions, fn EmbedderFunc) Embedder { opts.Supports = &EmbedderSupports{} } + configSchema, inputSchema := actionConfigSchemas[Config](opts.ConfigSchema, EmbedRequest{}, "options") + metadata := map[string]any{ "type": api.ActionTypeEmbedder, // TODO: This should be under "embedder" but JS has it as "info". @@ -115,24 +132,47 @@ func NewEmbedder(name string, opts *EmbedderOptions, fn EmbedderFunc) Embedder { }, }, "embedder": map[string]any{ - "customOptions": opts.ConfigSchema, + "customOptions": configSchema, }, } - inputSchema := core.InferSchemaMap(EmbedRequest{}) - if inputSchema != nil && opts.ConfigSchema != nil { - if props, ok := inputSchema["properties"].(map[string]any); ok { - props["options"] = opts.ConfigSchema + rawFn := func(ctx context.Context, req *EmbedRequest) (*EmbedResponse, error) { + cfg, err := resolveConfig[Config](req.Options) + if err != nil { + return nil, err } + // Normalize the request so its type-erased Options always carries the + // same converted value the typed parameter does. + req.Options = cfg + return fn(ctx, req, cfg) } return &embedder{ - Action: *core.NewAction(name, api.ActionTypeEmbedder, metadata, inputSchema, fn), + Action: *core.NewActionOf(api.ActionTypeEmbedder, name, &core.ActionOptions{ + Metadata: metadata, + InputSchema: inputSchema, + }, rawFn), } } +// NewEmbedder creates a new [Embedder]. +// +// Deprecated: Use [NewEmbedderWithConfig], which passes the request's options +// to fn as a typed value instead of leaving them type-erased on the request. +func NewEmbedder(name string, opts *EmbedderOptions, fn EmbedderFunc) Embedder { + if name == "" { + panic("ai.NewEmbedder: name is required") + } + return NewEmbedderWithConfig(name, opts, func(ctx context.Context, req *EmbedRequest, _ any) (*EmbedResponse, error) { + return fn(ctx, req) + }) +} + // DefineEmbedder registers the given embed function as an action, and returns an // [Embedder] that runs it. +// +// Deprecated: Use [NewEmbedderWithConfig] and register the result with +// [Embedder.Register]. func DefineEmbedder(r api.Registry, name string, opts *EmbedderOptions, fn EmbedderFunc) Embedder { e := NewEmbedder(name, opts, fn) e.Register(r) diff --git a/go/ai/evaluator.go b/go/ai/evaluator.go index 1ab7335932..55e051f875 100644 --- a/go/ai/evaluator.go +++ b/go/ai/evaluator.go @@ -31,9 +31,20 @@ import ( // EvaluatorFunc is the function type for evaluator implementations. type EvaluatorFunc = func(context.Context, *EvaluatorCallbackRequest) (*EvaluatorCallbackResponse, error) +// EvaluatorFuncWithConfig is an [EvaluatorFunc] that additionally receives +// the request's typed Config: the framework deserializes the request's raw +// options into it before calling the function (see [NewEvaluatorWithConfig]). +type EvaluatorFuncWithConfig[Config any] = func(context.Context, *EvaluatorCallbackRequest, Config) (*EvaluatorCallbackResponse, error) + // BatchEvaluatorFunc is the function type for batch evaluator implementations. type BatchEvaluatorFunc = func(context.Context, *EvaluatorRequest) (*EvaluatorResponse, error) +// BatchEvaluatorFuncWithConfig is a [BatchEvaluatorFunc] that additionally +// receives the request's typed Config: the framework deserializes the +// request's raw options into it before calling the function (see +// [NewBatchEvaluatorWithConfig]). +type BatchEvaluatorFuncWithConfig[Config any] = func(context.Context, *EvaluatorRequest, Config) (*EvaluatorResponse, error) + // Evaluator represents a evaluator action. type Evaluator interface { // Name returns the name of the evaluator. @@ -161,19 +172,10 @@ type EvaluatorCallbackRequest struct { // EvaluatorCallbackResponse is the result on evaluating a single [Example] type EvaluatorCallbackResponse = EvaluationResult -// NewEvaluator creates a new [Evaluator]. -// This method processes the input dataset one-by-one. -func NewEvaluator(name string, opts *EvaluatorOptions, fn EvaluatorFunc) Evaluator { - if name == "" { - panic("ai.NewEvaluator: evaluator name is required") - } - - if opts == nil { - opts = &EvaluatorOptions{} - } - +// evaluatorMetadata builds the shared action metadata for an evaluator. +func evaluatorMetadata(opts *EvaluatorOptions) map[string]any { // TODO(ssbushi): Set this on `evaluator` key on action metadata - metadata := map[string]any{ + return map[string]any{ "type": api.ActionTypeEvaluator, "evaluator": map[string]any{ "evaluatorIsBilled": opts.IsBilled, @@ -181,16 +183,42 @@ func NewEvaluator(name string, opts *EvaluatorOptions, fn EvaluatorFunc) Evaluat "evaluatorDefinition": opts.Definition, }, } +} - inputSchema := core.InferSchemaMap(EvaluatorRequest{}) - if inputSchema != nil && opts.ConfigSchema != nil { - if props, ok := inputSchema["properties"].(map[string]any); ok { - props["options"] = opts.ConfigSchema - } +// NewEvaluatorWithConfig creates a new [Evaluator]. Register it with +// [Evaluator.Register] to make it resolvable by name. +// This method processes the input dataset one-by-one. +// +// Config is the evaluator's typed configuration; it is usually inferred from +// fn's signature. The framework deserializes the request's raw options into +// Config before calling fn: the exact Config type (or a pointer to it) and +// map[string]any (from the Dev UI and other JSON callers) are accepted, and +// mismatched types are rejected. The config's JSON schema is inferred from +// Config unless [EvaluatorOptions.ConfigSchema] overrides it. +func NewEvaluatorWithConfig[Config any](name string, opts *EvaluatorOptions, fn EvaluatorFuncWithConfig[Config]) Evaluator { + if name == "" { + panic("ai.NewEvaluatorWithConfig: evaluator name is required") } + if opts == nil { + opts = &EvaluatorOptions{} + } + + _, inputSchema := actionConfigSchemas[Config](opts.ConfigSchema, EvaluatorRequest{}, "options") + return &evaluator{ - Action: *core.NewAction(name, api.ActionTypeEvaluator, metadata, inputSchema, func(ctx context.Context, req *EvaluatorRequest) (output *EvaluatorResponse, err error) { + Action: *core.NewActionOf(api.ActionTypeEvaluator, name, &core.ActionOptions{ + Metadata: evaluatorMetadata(opts), + InputSchema: inputSchema, + }, func(ctx context.Context, req *EvaluatorRequest) (output *EvaluatorResponse, err error) { + cfg, err := resolveConfig[Config](req.Options) + if err != nil { + return nil, err + } + // Normalize the request so its type-erased Options always carries + // the same converted value the typed parameter does. + req.Options = cfg + var results []EvaluationResult for _, datapoint := range req.Dataset { if datapoint.TestCaseId == "" { @@ -208,10 +236,10 @@ func NewEvaluator(name string, opts *EvaluatorOptions, fn EvaluatorFunc) Evaluat callbackRequest := EvaluatorCallbackRequest{ Input: *input, - Options: req.Options, + Options: cfg, } - result, err := fn(ctx, &callbackRequest) + result, err := fn(ctx, &callbackRequest, cfg) if err != nil { failedScore := Score{ Status: ScoreStatusFail.String(), @@ -245,8 +273,64 @@ func NewEvaluator(name string, opts *EvaluatorOptions, fn EvaluatorFunc) Evaluat } } +// NewBatchEvaluatorWithConfig creates a new [Evaluator]. Register it with +// [Evaluator.Register] to make it resolvable by name. +// This method provides the full [EvaluatorRequest] to the callback function, +// giving more flexibility to the user for processing the data, such as batching or parallelization. +// +// Config is the evaluator's typed configuration; it is usually inferred from +// fn's signature. See [NewEvaluatorWithConfig] for how the request's options +// are deserialized. +func NewBatchEvaluatorWithConfig[Config any](name string, opts *EvaluatorOptions, fn BatchEvaluatorFuncWithConfig[Config]) Evaluator { + if name == "" { + panic("ai.NewBatchEvaluatorWithConfig: batch evaluator name is required") + } + + if opts == nil { + opts = &EvaluatorOptions{} + } + + _, inputSchema := actionConfigSchemas[Config](opts.ConfigSchema, EvaluatorRequest{}, "options") + + rawFn := func(ctx context.Context, req *EvaluatorRequest) (*EvaluatorResponse, error) { + cfg, err := resolveConfig[Config](req.Options) + if err != nil { + return nil, err + } + // Normalize the request so its type-erased Options always carries the + // same converted value the typed parameter does. + req.Options = cfg + return fn(ctx, req, cfg) + } + + return &evaluator{ + Action: *core.NewActionOf(api.ActionTypeEvaluator, name, &core.ActionOptions{ + Metadata: evaluatorMetadata(opts), + InputSchema: inputSchema, + }, rawFn), + } +} + +// NewEvaluator creates a new [Evaluator]. +// This method processes the input dataset one-by-one. +// +// Deprecated: Use [NewEvaluatorWithConfig], which passes the request's +// options to fn as a typed value instead of leaving them type-erased on the +// request. +func NewEvaluator(name string, opts *EvaluatorOptions, fn EvaluatorFunc) Evaluator { + if name == "" { + panic("ai.NewEvaluator: evaluator name is required") + } + return NewEvaluatorWithConfig(name, opts, func(ctx context.Context, req *EvaluatorCallbackRequest, _ any) (*EvaluatorCallbackResponse, error) { + return fn(ctx, req) + }) +} + // DefineEvaluator creates a new [Evaluator] and registers it. // This method processes the input dataset one-by-one. +// +// Deprecated: Use [NewEvaluatorWithConfig] and register the result with +// [Evaluator.Register]. func DefineEvaluator(r api.Registry, name string, opts *EvaluatorOptions, fn EvaluatorFunc) Evaluator { e := NewEvaluator(name, opts, fn) e.Register(r) @@ -256,35 +340,28 @@ func DefineEvaluator(r api.Registry, name string, opts *EvaluatorOptions, fn Eva // NewBatchEvaluator creates a new [Evaluator]. // This method provides the full [EvaluatorRequest] to the callback function, // giving more flexibility to the user for processing the data, such as batching or parallelization. +// +// Deprecated: Use [NewBatchEvaluatorWithConfig], which passes the request's +// options to fn as a typed value instead of leaving them type-erased on the +// request. func NewBatchEvaluator(name string, opts *EvaluatorOptions, fn BatchEvaluatorFunc) Evaluator { if name == "" { panic("ai.NewBatchEvaluator: batch evaluator name is required") } - - if opts == nil { - opts = &EvaluatorOptions{} - } - - metadata := map[string]any{ - "type": api.ActionTypeEvaluator, - "evaluator": map[string]any{ - "evaluatorIsBilled": opts.IsBilled, - "evaluatorDisplayName": opts.DisplayName, - "evaluatorDefinition": opts.Definition, - }, - } - - return &evaluator{ - Action: *core.NewAction(name, api.ActionTypeEvaluator, metadata, nil, fn), - } + return NewBatchEvaluatorWithConfig(name, opts, func(ctx context.Context, req *EvaluatorRequest, _ any) (*EvaluatorResponse, error) { + return fn(ctx, req) + }) } // DefineBatchEvaluator creates a new [Evaluator] and registers it. // This method provides the full [EvaluatorRequest] to the callback function, // giving more flexibility to the user for processing the data, such as batching or parallelization. +// +// Deprecated: Use [NewBatchEvaluatorWithConfig] and register the result with +// [Evaluator.Register]. func DefineBatchEvaluator(r api.Registry, name string, opts *EvaluatorOptions, fn BatchEvaluatorFunc) Evaluator { e := NewBatchEvaluator(name, opts, fn) - e.(*evaluator).Register(r) + e.Register(r) return e } diff --git a/go/ai/generate.go b/go/ai/generate.go index a535903de5..e53fc85c4b 100644 --- a/go/ai/generate.go +++ b/go/ai/generate.go @@ -69,6 +69,11 @@ type ToolConfig struct { // ModelFunc is a streaming function that takes in a ModelRequest and generates a ModelResponse, optionally streaming ModelResponseChunks. type ModelFunc = core.StreamingFunc[*ModelRequest, *ModelResponse, *ModelResponseChunk] +// ModelFuncWithConfig is a [ModelFunc] that additionally receives the +// request's typed Config: the framework deserializes the request's raw config +// into it before calling the function (see [NewModelWithConfig]). +type ModelFuncWithConfig[Config any] = func(context.Context, *ModelRequest, Config, ModelStreamCallback) (*ModelResponse, error) + // ModelStreamCallback is a stream callback of a ModelAction. type ModelStreamCallback = func(context.Context, *ModelResponseChunk) error @@ -135,10 +140,26 @@ func DefineGenerateAction(ctx context.Context, r api.Registry) *generateAction { return (*generateAction)(a) } -// NewModel creates a new [Model]. -func NewModel(name string, opts *ModelOptions, fn ModelFunc) Model { +// NewModelWithConfig creates a new [Model]. Register it with +// [Model.Register] to make it resolvable by name. +// +// Config is the model's typed configuration; it is usually inferred from fn's +// signature. The framework deserializes the request's raw config into Config +// before calling fn: the exact Config type (or a pointer to it) and +// map[string]any (from the Dev UI and other JSON callers) are accepted, and +// mismatched types are rejected. The request's [ModelRequest.Config] is +// normalized to the converted value, so it always matches the typed +// parameter. The config's JSON schema is inferred from Config unless +// [ModelOptions.ConfigSchema] overrides it. +// +// The config schema is enforced by input validation on every call, so if +// Config's JSON marshaling diverges from its reflected schema (e.g. SDK +// wrapper types like Opt[float64] that marshal to primitives but reflect as +// objects), set [ModelOptions.ConfigSchema] explicitly or requests will be +// rejected at the action boundary. +func NewModelWithConfig[Config any](name string, opts *ModelOptions, fn ModelFuncWithConfig[Config]) Model { if name == "" { - panic("ai.NewModel: name is required") + panic("ai.NewModelWithConfig: name is required") } if opts == nil { @@ -150,6 +171,8 @@ func NewModel(name string, opts *ModelOptions, fn ModelFunc) Model { opts.Supports = &ModelSupports{} } + configSchema, inputSchema := actionConfigSchemas[Config](opts.ConfigSchema, ModelRequest{}, "config") + metadata := map[string]any{ "type": api.ActionTypeModel, "model": map[string]any{ @@ -168,29 +191,53 @@ func NewModel(name string, opts *ModelOptions, fn ModelFunc) Model { }, "versions": opts.Versions, "stage": opts.Stage, - "customOptions": opts.ConfigSchema, + "customOptions": configSchema, }, } - inputSchema := core.InferSchemaMap(ModelRequest{}) - if inputSchema != nil && opts.ConfigSchema != nil { - if props, ok := inputSchema["properties"].(map[string]any); ok { - props["config"] = opts.ConfigSchema + typedFn := func(ctx context.Context, req *ModelRequest, cb ModelStreamCallback) (*ModelResponse, error) { + // req.Config was normalized to the exact Config type by + // normalizeConfig below, so this hits the fast path. + cfg, err := resolveConfig[Config](req.Config) + if err != nil { + return nil, err } + return fn(ctx, req, cfg, cb) } - mws := []ModelMiddleware{ + // normalizeConfig runs outermost so that the built-in wrappers and the + // model function all see the typed, converted config on the request. + rawFn := core.ChainMiddleware( + normalizeConfig[Config](name, opts.Versions), simulateSystemPrompt(opts, nil), augmentWithContext(opts, nil), validateSupport(name, opts), addAutomaticTelemetry(), - } - fn = core.ChainMiddleware(mws...)(fn) + )(typedFn) - return &model{*core.NewStreamingAction(name, api.ActionTypeModel, metadata, inputSchema, fn)} + return &model{*core.NewStreamingActionOf(api.ActionTypeModel, name, &core.ActionOptions{ + Metadata: metadata, + InputSchema: inputSchema, + }, rawFn)} +} + +// NewModel creates a new [Model]. +// +// Deprecated: Use [NewModelWithConfig], which passes the request's config to +// fn as a typed value instead of leaving it type-erased on the request. +func NewModel(name string, opts *ModelOptions, fn ModelFunc) Model { + if name == "" { + panic("ai.NewModel: name is required") + } + return NewModelWithConfig(name, opts, func(ctx context.Context, req *ModelRequest, _ any, cb ModelStreamCallback) (*ModelResponse, error) { + return fn(ctx, req, cb) + }) } // DefineModel creates a new [Model] and registers it. +// +// Deprecated: Use [NewModelWithConfig] and register the result with +// [Model.Register]. func DefineModel(r api.Registry, name string, opts *ModelOptions, fn ModelFunc) Model { m := NewModel(name, opts, fn) m.Register(r) diff --git a/go/ai/model_middleware.go b/go/ai/model_middleware.go index ed5aa9340c..d82bfd5e56 100644 --- a/go/ai/model_middleware.go +++ b/go/ai/model_middleware.go @@ -265,16 +265,15 @@ func validateSupport(model string, opts *ModelOptions) ModelMiddleware { return nil, core.NewError(core.INVALID_ARGUMENT, "model %q does not support native constrained output, but constrained output was requested. Request: %+v", model, input) } - if err := validateVersion(model, opts.Versions, input.Config); err != nil { - return nil, err - } - return next(ctx, input, cb) } } } // validateVersion validates that the requested model version is supported. +// It runs against the raw, pre-conversion config (see [normalizeConfig]) +// because conversion into a Config type without a version field would +// silently drop the key. func validateVersion(model string, versions []string, config any) error { var configMap map[string]any diff --git a/go/ai/model_middleware_test.go b/go/ai/model_middleware_test.go index 4ba4a84577..f4ef0efa5e 100644 --- a/go/ai/model_middleware_test.go +++ b/go/ai/model_middleware_test.go @@ -23,6 +23,8 @@ import ( "net/http/httptest" "reflect" "testing" + + "github.com/firebase/genkit/go/core" ) func TestValidateSupport(t *testing.T) { @@ -254,7 +256,17 @@ func TestValidateSupport(t *testing.T) { for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { - handler := validateSupport("test-model", tt.opts)(mockModelFunc) + // Version validation moved from validateSupport to + // normalizeConfig (it must see the raw, pre-conversion config), + // so chain both here the way model construction does. + var versions []string + if tt.opts != nil { + versions = tt.opts.Versions + } + handler := core.ChainMiddleware( + normalizeConfig[any]("test-model", versions), + validateSupport("test-model", tt.opts), + )(mockModelFunc) _, err := handler(context.Background(), tt.input, nil) if (err != nil) != tt.wantErr { From 883a804b7c6c829cc8d1a5abbe9cba3caa920c6c Mon Sep 17 00:00:00 2001 From: Alex Pascal Date: Wed, 29 Jul 2026 15:33:19 -0700 Subject: [PATCH 02/10] feat(go/genkit): add Define*WithConfig wrappers DefineModelWithConfig, DefineBackgroundModelWithConfig, DefineEmbedderWithConfig, DefineEvaluatorWithConfig, and DefineBatchEvaluatorWithConfig wrap the ai New*WithConfig constructors as New* + Register(g.reg). The untyped Define* wrappers are deprecated in favor of them. --- go/genkit/genkit.go | 129 ++++++++++++++++++++++++++++++++++---------- 1 file changed, 100 insertions(+), 29 deletions(-) diff --git a/go/genkit/genkit.go b/go/genkit/genkit.go index 3f9b3c14da..cdbf09035f 100644 --- a/go/genkit/genkit.go +++ b/go/genkit/genkit.go @@ -498,26 +498,30 @@ func ListTools(g *Genkit) []ai.Tool { return tools } -// DefineModel defines a custom model implementation, registers it as a [core.Action] -// of type Model, and returns an [ai.Model] interface. +// DefineModelWithConfig defines a custom model implementation, registers it +// as a [core.Action] of type Model, and returns an [ai.Model] interface. // // The `name` argument is the unique identifier for the model (e.g., "myProvider/myModel"). // The `opts` argument provides metadata about the model's capabilities ([ai.ModelOptions]). -// The `fn` argument ([ai.ModelFunc]) implements the actual generation logic, handling -// input requests ([ai.ModelRequest]) and producing responses ([ai.ModelResponse]), +// The `fn` argument ([ai.ModelFuncWithConfig]) implements the actual generation logic, +// handling input requests ([ai.ModelRequest]) and producing responses ([ai.ModelResponse]), // potentially streaming chunks ([ai.ModelResponseChunk]) via the callback. // +// Config is the model's typed configuration; it is usually inferred from fn's +// signature. See [ai.NewModelWithConfig] for how the request's config is +// deserialized and validated. +// // For models that don't need to be registered (e.g., for plugin development or testing), -// use [ai.NewModel] instead. +// use [ai.NewModelWithConfig] instead. // // Example: // -// echoModel := genkit.DefineModel(g, "custom/echo", +// echoModel := genkit.DefineModelWithConfig(g, "custom/echo", // &ai.ModelOptions{ // Label: "Echo Model", // Supports: &ai.ModelSupports{Multiturn: true}, // }, -// func(ctx context.Context, req *ai.ModelRequest, cb ai.ModelStreamCallback) (*ai.ModelResponse, error) { +// func(ctx context.Context, req *ai.ModelRequest, config MyConfig, cb ai.ModelStreamCallback) (*ai.ModelResponse, error) { // // Simple echo implementation // resp := &ai.ModelResponse{ // Message: &ai.Message{ @@ -555,16 +559,43 @@ func ListTools(g *Genkit) []ai.Tool { // return resp, nil // }, // ) +func DefineModelWithConfig[Config any](g *Genkit, name string, opts *ai.ModelOptions, fn ai.ModelFuncWithConfig[Config]) ai.Model { + m := ai.NewModelWithConfig(name, opts, fn) + m.Register(g.reg) + return m +} + +// DefineModel defines a custom model implementation, registers it as a [core.Action] +// of type Model, and returns an [ai.Model] interface. +// +// Deprecated: Use [DefineModelWithConfig], which passes the request's config +// to fn as a typed value instead of leaving it type-erased on the request. func DefineModel(g *Genkit, name string, opts *ai.ModelOptions, fn ai.ModelFunc) ai.Model { return ai.DefineModel(g.reg, name, opts, fn) } -// DefineBackgroundModel defines a background model, registers it as a [ai.BackgroundModel], -// and returns an [ai.BackgroundModel]. +// DefineBackgroundModelWithConfig defines a background model, registers it as +// an [ai.BackgroundModel], and returns an [ai.BackgroundModel]. // // The `name` is the identifier the model uses to request the background model. The `opts` // are the options for the background model. The `startFn` is the function that starts the background model. // The `checkFn` is the function that checks the status of the background model. +// +// Config is the model's typed configuration; it is usually inferred from +// startFn's signature. See [ai.NewModelWithConfig] for how the request's +// config is deserialized and validated. +func DefineBackgroundModelWithConfig[Config any](g *Genkit, name string, opts *ai.BackgroundModelOptions, startFn ai.StartModelOpFuncWithConfig[Config], checkFn ai.CheckModelOpFunc) ai.BackgroundModel { + m := ai.NewBackgroundModelWithConfig(name, opts, startFn, checkFn) + m.Register(g.reg) + return m +} + +// DefineBackgroundModel defines a background model, registers it as a [ai.BackgroundModel], +// and returns an [ai.BackgroundModel]. +// +// Deprecated: Use [DefineBackgroundModelWithConfig], which passes the +// request's config to startFn as a typed value instead of leaving it +// type-erased on the request. func DefineBackgroundModel(g *Genkit, name string, opts *ai.BackgroundModelOptions, startFn ai.StartModelOpFunc, checkFn ai.CheckModelOpFunc) ai.BackgroundModel { return ai.DefineBackgroundModel(g.reg, name, opts, startFn, checkFn) } @@ -1339,16 +1370,33 @@ func LookupRetriever(g *Genkit, name string) ai.Retriever { return ai.LookupRetriever(g.reg, name) } -// DefineEmbedder defines a custom text embedding implementation, registers it as a -// [core.Action] of type Embedder, and returns an [ai.Embedder]. -// Embedders convert text documents or queries into numerical vector representations (embeddings). +// DefineEmbedderWithConfig defines a custom text embedding implementation, +// registers it as a [core.Action] of type Embedder, and returns an +// [ai.Embedder]. Embedders convert text documents or queries into numerical +// vector representations (embeddings). // // The `name` is the unique identifier for the embedder. // The `fn` function contains the logic to process an [ai.EmbedRequest] (containing documents or a query) // and return an [ai.EmbedResponse] (containing the corresponding embeddings). // +// Config is the embedder's typed configuration; it is usually inferred from +// fn's signature. See [ai.NewEmbedderWithConfig] for how the request's +// options are deserialized. +// // For embedders that don't need to be registered (e.g., for plugin development), -// use [ai.NewEmbedder] instead. +// use [ai.NewEmbedderWithConfig] instead. +func DefineEmbedderWithConfig[Config any](g *Genkit, name string, opts *ai.EmbedderOptions, fn ai.EmbedderFuncWithConfig[Config]) ai.Embedder { + e := ai.NewEmbedderWithConfig(name, opts, fn) + e.Register(g.reg) + return e +} + +// DefineEmbedder defines a custom text embedding implementation, registers it as a +// [core.Action] of type Embedder, and returns an [ai.Embedder]. +// +// Deprecated: Use [DefineEmbedderWithConfig], which passes the request's +// options to fn as a typed value instead of leaving them type-erased on the +// request. func DefineEmbedder(g *Genkit, name string, opts *ai.EmbedderOptions, fn ai.EmbedderFunc) ai.Embedder { return ai.DefineEmbedder(g.reg, name, opts, fn) } @@ -1370,33 +1418,56 @@ func LookupPlugin(g *Genkit, name string) api.Plugin { return g.reg.LookupPlugin(name) } -// DefineEvaluator defines an evaluator that processes test cases one by one, -// registers it as a [core.Action] of type Evaluator, and returns an [ai.Evaluator]. -// Evaluators are used to assess the quality or performance of AI models or flows -// based on a dataset of test cases. +// DefineEvaluatorWithConfig defines an evaluator that processes test cases +// one by one, registers it as a [core.Action] of type Evaluator, and returns +// an [ai.Evaluator]. Evaluators are used to assess the quality or performance +// of AI models or flows based on a dataset of test cases. // -// This variant calls the provided `eval` function for each individual test case +// This variant calls the provided `fn` function for each individual test case // ([ai.EvaluatorCallbackRequest]) in the evaluation dataset. // -// The `provider` and `name` form the unique identifier. `options` provide -// metadata about the evaluator ([ai.EvaluatorOptions]). The `eval` function -// implements the logic to score a single test case and returns the results -// in an [ai.EvaluatorCallbackResponse]. +// Config is the evaluator's typed configuration; it is usually inferred from +// fn's signature. See [ai.NewEvaluatorWithConfig] for how the request's +// options are deserialized. +func DefineEvaluatorWithConfig[Config any](g *Genkit, name string, opts *ai.EvaluatorOptions, fn ai.EvaluatorFuncWithConfig[Config]) ai.Evaluator { + e := ai.NewEvaluatorWithConfig(name, opts, fn) + e.Register(g.reg) + return e +} + +// DefineEvaluator defines an evaluator that processes test cases one by one, +// registers it as a [core.Action] of type Evaluator, and returns an [ai.Evaluator]. +// +// Deprecated: Use [DefineEvaluatorWithConfig], which passes the request's +// options to fn as a typed value instead of leaving them type-erased on the +// request. func DefineEvaluator(g *Genkit, name string, opts *ai.EvaluatorOptions, fn ai.EvaluatorFunc) ai.Evaluator { return ai.DefineEvaluator(g.reg, name, opts, fn) } -// DefineBatchEvaluator defines an evaluator that processes the entire dataset at once, -// registers it as a [core.Action] of type Evaluator, and returns an [ai.Evaluator]. +// DefineBatchEvaluatorWithConfig defines an evaluator that processes the +// entire dataset at once, registers it as a [core.Action] of type Evaluator, +// and returns an [ai.Evaluator]. // // This variant provides the full evaluation request ([ai.EvaluatorRequest]), including -// the entire dataset, to the `eval` function. This allows for more flexible processing, +// the entire dataset, to the `fn` function. This allows for more flexible processing, // such as batching calls to external services or parallelizing computations. // -// The `provider` and `name` form the unique identifier. `options` provide -// metadata about the evaluator ([ai.EvaluatorOptions]). The `eval` function -// implements the logic to score the dataset and returns the aggregated results -// in an [ai.EvaluatorResponse]. +// Config is the evaluator's typed configuration; it is usually inferred from +// fn's signature. See [ai.NewEvaluatorWithConfig] for how the request's +// options are deserialized. +func DefineBatchEvaluatorWithConfig[Config any](g *Genkit, name string, opts *ai.EvaluatorOptions, fn ai.BatchEvaluatorFuncWithConfig[Config]) ai.Evaluator { + e := ai.NewBatchEvaluatorWithConfig(name, opts, fn) + e.Register(g.reg) + return e +} + +// DefineBatchEvaluator defines an evaluator that processes the entire dataset at once, +// registers it as a [core.Action] of type Evaluator, and returns an [ai.Evaluator]. +// +// Deprecated: Use [DefineBatchEvaluatorWithConfig], which passes the +// request's options to fn as a typed value instead of leaving them +// type-erased on the request. func DefineBatchEvaluator(g *Genkit, name string, opts *ai.EvaluatorOptions, fn ai.BatchEvaluatorFunc) ai.Evaluator { return ai.DefineBatchEvaluator(g.reg, name, opts, fn) } From 001cb0e8b9d5f7a4be39f47c9d0113da501401e8 Mon Sep 17 00:00:00 2001 From: Alex Pascal Date: Wed, 29 Jul 2026 15:45:22 -0700 Subject: [PATCH 03/10] refactor(go): rename New*WithConfig constructors to NewTyped* NewTypedModel, NewTypedEmbedder, NewTypedEvaluator, NewTypedBatchEvaluator, and NewTypedBackgroundModel name what the variant actually changes: the function receives the request's config as a typed value. Func aliases follow (TypedModelFunc[Config] et al.), and the genkit wrappers become DefineTyped*. These names shipped earlier on this branch and were never released, so this is a rename, not a break. --- go/ai/background_model.go | 24 ++++++++-------- go/ai/config.go | 2 +- go/ai/config_test.go | 22 +++++++-------- go/ai/embedder.go | 18 ++++++------ go/ai/evaluator.go | 38 ++++++++++++------------- go/ai/generate.go | 18 ++++++------ go/core/doc.go | 2 +- go/genkit/genkit.go | 58 +++++++++++++++++++-------------------- 8 files changed, 91 insertions(+), 91 deletions(-) diff --git a/go/ai/background_model.go b/go/ai/background_model.go index d4e914176d..0cfc6d4887 100644 --- a/go/ai/background_model.go +++ b/go/ai/background_model.go @@ -52,11 +52,11 @@ type ModelOperation = core.Operation[*ModelResponse] // StartModelOpFunc starts a background model operation. type StartModelOpFunc = func(ctx context.Context, req *ModelRequest) (*ModelOperation, error) -// StartModelOpFuncWithConfig is a [StartModelOpFunc] that additionally +// TypedStartModelOpFunc is a [StartModelOpFunc] that additionally // receives the request's typed Config: the framework deserializes the // request's raw config into it before calling the function (see -// [NewBackgroundModelWithConfig]). -type StartModelOpFuncWithConfig[Config any] = func(ctx context.Context, req *ModelRequest, config Config) (*ModelOperation, error) +// [NewTypedBackgroundModel]). +type TypedStartModelOpFunc[Config any] = func(ctx context.Context, req *ModelRequest, config Config) (*ModelOperation, error) // CheckModelOpFunc checks the status of a background model operation. type CheckModelOpFunc = func(ctx context.Context, op *ModelOperation) (*ModelOperation, error) @@ -82,21 +82,21 @@ func LookupBackgroundModel(r api.Registry, name string) BackgroundModel { return &backgroundModel{*action} } -// NewBackgroundModelWithConfig creates a new [BackgroundModel]. Register it +// NewTypedBackgroundModel creates a new [BackgroundModel]. Register it // with [BackgroundModel.Register] to make it resolvable by name. // // Config is the model's typed configuration; it is usually inferred from -// startFn's signature. See [NewModelWithConfig] for how the request's config +// startFn's signature. See [NewTypedModel] for how the request's config // is deserialized. -func NewBackgroundModelWithConfig[Config any](name string, opts *BackgroundModelOptions, startFn StartModelOpFuncWithConfig[Config], checkFn CheckModelOpFunc) BackgroundModel { +func NewTypedBackgroundModel[Config any](name string, opts *BackgroundModelOptions, startFn TypedStartModelOpFunc[Config], checkFn CheckModelOpFunc) BackgroundModel { if name == "" { - panic("ai.NewBackgroundModelWithConfig: name is required") + panic("ai.NewTypedBackgroundModel: name is required") } if startFn == nil { - panic("ai.NewBackgroundModelWithConfig: startFn is required") + panic("ai.NewTypedBackgroundModel: startFn is required") } if checkFn == nil { - panic("ai.NewBackgroundModelWithConfig: checkFn is required") + panic("ai.NewTypedBackgroundModel: checkFn is required") } if opts == nil { @@ -169,21 +169,21 @@ func NewBackgroundModelWithConfig[Config any](name string, opts *BackgroundModel // NewBackgroundModel defines a new model that runs in the background. // -// Deprecated: Use [NewBackgroundModelWithConfig], which passes the request's +// Deprecated: Use [NewTypedBackgroundModel], which passes the request's // config to startFn as a typed value instead of leaving it type-erased on the // request. func NewBackgroundModel(name string, opts *BackgroundModelOptions, startFn StartModelOpFunc, checkFn CheckModelOpFunc) BackgroundModel { if startFn == nil { panic("ai.NewBackgroundModel: startFn is required") } - return NewBackgroundModelWithConfig(name, opts, func(ctx context.Context, req *ModelRequest, _ any) (*ModelOperation, error) { + return NewTypedBackgroundModel(name, opts, func(ctx context.Context, req *ModelRequest, _ any) (*ModelOperation, error) { return startFn(ctx, req) }, checkFn) } // DefineBackgroundModel defines and registers a new model that runs in the background. // -// Deprecated: Use [NewBackgroundModelWithConfig] and register the result with +// Deprecated: Use [NewTypedBackgroundModel] and register the result with // [BackgroundModel.Register]. func DefineBackgroundModel(r *registry.Registry, name string, opts *BackgroundModelOptions, fn StartModelOpFunc, checkFn CheckModelOpFunc) BackgroundModel { m := NewBackgroundModel(name, opts, fn, checkFn) diff --git a/go/ai/config.go b/go/ai/config.go index 6775b2561d..aeccb05603 100644 --- a/go/ai/config.go +++ b/go/ai/config.go @@ -27,7 +27,7 @@ import ( // This file holds the typed-config plumbing shared by models, embedders, and // evaluators. Requests carry config as `any` on the wire; the -// Define*/New*WithConfig constructors wrap the user's typed function so that +// NewTyped* constructors wrap the user's typed function so that // the raw value is deserialized into the Config type parameter before the // function runs, and the request's type-erased config slot is normalized to // that same converted value so the two views never disagree. diff --git a/go/ai/config_test.go b/go/ai/config_test.go index 1e69536b4c..51055c20d4 100644 --- a/go/ai/config_test.go +++ b/go/ai/config_test.go @@ -26,7 +26,7 @@ import ( ) // testTypedConfig is a provider-style config struct used to exercise the -// typed-config deserialization that New*WithConfig wraps around user +// typed-config deserialization that NewTyped* constructors wrap around user // functions. type testTypedConfig struct { Temperature float64 `json:"temperature,omitempty"` @@ -46,7 +46,7 @@ func TestModelTypedConfig(t *testing.T) { var got testTypedConfig var gotReqConfig any - m := NewModelWithConfig("test/typed-config", nil, func(ctx context.Context, req *ModelRequest, cfg testTypedConfig, cb ModelStreamCallback) (*ModelResponse, error) { + m := NewTypedModel("test/typed-config", nil, func(ctx context.Context, req *ModelRequest, cfg testTypedConfig, cb ModelStreamCallback) (*ModelResponse, error) { got = cfg gotReqConfig = req.Config return &ModelResponse{Message: NewModelTextMessage("ok"), Request: req}, nil @@ -100,7 +100,7 @@ func TestModelTypedConfig(t *testing.T) { func TestModelConfigNormalizedBeforeBuiltins(t *testing.T) { r := registry.New() - m := NewModelWithConfig("test/config-order", &ModelOptions{ + m := NewTypedModel("test/config-order", &ModelOptions{ Supports: &ModelSupports{Multiturn: true}, // no media support }, func(ctx context.Context, req *ModelRequest, cfg testTypedConfig, cb ModelStreamCallback) (*ModelResponse, error) { return &ModelResponse{Message: NewModelTextMessage("ok"), Request: req}, nil @@ -120,7 +120,7 @@ func TestModelConfigNormalizedBeforeBuiltins(t *testing.T) { func TestBackgroundModelConfigValidation(t *testing.T) { r := registry.New() - bm := NewBackgroundModelWithConfig("test/bg-typed-config", nil, + bm := NewTypedBackgroundModel("test/bg-typed-config", nil, func(ctx context.Context, req *ModelRequest, cfg testTypedConfig) (*ModelOperation, error) { return &ModelOperation{ID: "op1", Done: false}, nil }, @@ -150,7 +150,7 @@ func TestBackgroundModelConfigValidation(t *testing.T) { } func TestModelConfigSchemaInference(t *testing.T) { - newFn := func() ModelFuncWithConfig[testTypedConfig] { + newFn := func() TypedModelFunc[testTypedConfig] { return func(ctx context.Context, req *ModelRequest, cfg testTypedConfig, cb ModelStreamCallback) (*ModelResponse, error) { return &ModelResponse{Message: NewModelTextMessage("ok"), Request: req}, nil } @@ -167,7 +167,7 @@ func TestModelConfigSchemaInference(t *testing.T) { } t.Run("inferred from Config type", func(t *testing.T) { - m := NewModelWithConfig("test/inferred-schema", nil, newFn()) + m := NewTypedModel("test/inferred-schema", nil, newFn()) schema, ok := configSchemaOf(t, m).(map[string]any) if !ok { t.Fatalf("customOptions = %v, want inferred schema map", configSchemaOf(t, m)) @@ -183,7 +183,7 @@ func TestModelConfigSchemaInference(t *testing.T) { t.Run("explicit ConfigSchema wins", func(t *testing.T) { override := map[string]any{"type": "object", "properties": map[string]any{"custom": map[string]any{"type": "string"}}} - m := NewModelWithConfig("test/override-schema", &ModelOptions{ConfigSchema: override}, newFn()) + m := NewTypedModel("test/override-schema", &ModelOptions{ConfigSchema: override}, newFn()) schema, ok := configSchemaOf(t, m).(map[string]any) if !ok { t.Fatalf("customOptions missing") @@ -195,7 +195,7 @@ func TestModelConfigSchemaInference(t *testing.T) { }) t.Run("any config infers no schema", func(t *testing.T) { - m := NewModelWithConfig("test/any-schema", nil, func(ctx context.Context, req *ModelRequest, cfg any, cb ModelStreamCallback) (*ModelResponse, error) { + m := NewTypedModel("test/any-schema", nil, func(ctx context.Context, req *ModelRequest, cfg any, cb ModelStreamCallback) (*ModelResponse, error) { return &ModelResponse{Message: NewModelTextMessage("ok"), Request: req}, nil }) if s, _ := configSchemaOf(t, m).(map[string]any); len(s) != 0 { @@ -243,7 +243,7 @@ func TestEmbedderTypedConfig(t *testing.T) { var got testTypedConfig var gotReqOptions any - e := NewEmbedderWithConfig("test/typed-config-embedder", nil, func(ctx context.Context, req *EmbedRequest, cfg testTypedConfig) (*EmbedResponse, error) { + e := NewTypedEmbedder("test/typed-config-embedder", nil, func(ctx context.Context, req *EmbedRequest, cfg testTypedConfig) (*EmbedResponse, error) { got = cfg gotReqOptions = req.Options return &EmbedResponse{}, nil @@ -271,7 +271,7 @@ func TestEvaluatorTypedConfig(t *testing.T) { r := registry.New() var got testTypedConfig - e := NewEvaluatorWithConfig("test/typed-config-evaluator", nil, func(ctx context.Context, req *EvaluatorCallbackRequest, cfg testTypedConfig) (*EvaluatorCallbackResponse, error) { + e := NewTypedEvaluator("test/typed-config-evaluator", nil, func(ctx context.Context, req *EvaluatorCallbackRequest, cfg testTypedConfig) (*EvaluatorCallbackResponse, error) { got = cfg return &EvaluatorCallbackResponse{ TestCaseId: req.Input.TestCaseId, @@ -293,7 +293,7 @@ func TestEvaluatorTypedConfig(t *testing.T) { } var gotBatch testTypedConfig - be := NewBatchEvaluatorWithConfig("test/typed-config-batch-evaluator", nil, func(ctx context.Context, req *EvaluatorRequest, cfg testTypedConfig) (*EvaluatorResponse, error) { + be := NewTypedBatchEvaluator("test/typed-config-batch-evaluator", nil, func(ctx context.Context, req *EvaluatorRequest, cfg testTypedConfig) (*EvaluatorResponse, error) { gotBatch = cfg return &EvaluatorResponse{}, nil }) diff --git a/go/ai/embedder.go b/go/ai/embedder.go index e43505e68c..5976954252 100644 --- a/go/ai/embedder.go +++ b/go/ai/embedder.go @@ -27,10 +27,10 @@ import ( // EmbedderFunc is the function type for embedding documents. type EmbedderFunc = func(context.Context, *EmbedRequest) (*EmbedResponse, error) -// EmbedderFuncWithConfig is an [EmbedderFunc] that additionally receives the +// TypedEmbedderFunc is an [EmbedderFunc] that additionally receives the // request's typed Config: the framework deserializes the request's raw -// options into it before calling the function (see [NewEmbedderWithConfig]). -type EmbedderFuncWithConfig[Config any] = func(context.Context, *EmbedRequest, Config) (*EmbedResponse, error) +// options into it before calling the function (see [NewTypedEmbedder]). +type TypedEmbedderFunc[Config any] = func(context.Context, *EmbedRequest, Config) (*EmbedResponse, error) // Embedder represents an embedder that can perform content embedding. type Embedder interface { @@ -93,7 +93,7 @@ type embedder struct { core.Action[*EmbedRequest, *EmbedResponse, struct{}] } -// NewEmbedderWithConfig creates a new [Embedder]. Register it with +// NewTypedEmbedder creates a new [Embedder]. Register it with // [Embedder.Register] to make it resolvable by name. // // Config is the embedder's typed configuration; it is usually inferred from @@ -104,9 +104,9 @@ type embedder struct { // normalized to the converted value, so it always matches the typed // parameter. The config's JSON schema is inferred from Config unless // [EmbedderOptions.ConfigSchema] overrides it. -func NewEmbedderWithConfig[Config any](name string, opts *EmbedderOptions, fn EmbedderFuncWithConfig[Config]) Embedder { +func NewTypedEmbedder[Config any](name string, opts *EmbedderOptions, fn TypedEmbedderFunc[Config]) Embedder { if name == "" { - panic("ai.NewEmbedderWithConfig: name is required") + panic("ai.NewTypedEmbedder: name is required") } if opts == nil { @@ -157,13 +157,13 @@ func NewEmbedderWithConfig[Config any](name string, opts *EmbedderOptions, fn Em // NewEmbedder creates a new [Embedder]. // -// Deprecated: Use [NewEmbedderWithConfig], which passes the request's options +// Deprecated: Use [NewTypedEmbedder], which passes the request's options // to fn as a typed value instead of leaving them type-erased on the request. func NewEmbedder(name string, opts *EmbedderOptions, fn EmbedderFunc) Embedder { if name == "" { panic("ai.NewEmbedder: name is required") } - return NewEmbedderWithConfig(name, opts, func(ctx context.Context, req *EmbedRequest, _ any) (*EmbedResponse, error) { + return NewTypedEmbedder(name, opts, func(ctx context.Context, req *EmbedRequest, _ any) (*EmbedResponse, error) { return fn(ctx, req) }) } @@ -171,7 +171,7 @@ func NewEmbedder(name string, opts *EmbedderOptions, fn EmbedderFunc) Embedder { // DefineEmbedder registers the given embed function as an action, and returns an // [Embedder] that runs it. // -// Deprecated: Use [NewEmbedderWithConfig] and register the result with +// Deprecated: Use [NewTypedEmbedder] and register the result with // [Embedder.Register]. func DefineEmbedder(r api.Registry, name string, opts *EmbedderOptions, fn EmbedderFunc) Embedder { e := NewEmbedder(name, opts, fn) diff --git a/go/ai/evaluator.go b/go/ai/evaluator.go index 55e051f875..70a5f4cd4c 100644 --- a/go/ai/evaluator.go +++ b/go/ai/evaluator.go @@ -31,19 +31,19 @@ import ( // EvaluatorFunc is the function type for evaluator implementations. type EvaluatorFunc = func(context.Context, *EvaluatorCallbackRequest) (*EvaluatorCallbackResponse, error) -// EvaluatorFuncWithConfig is an [EvaluatorFunc] that additionally receives +// TypedEvaluatorFunc is an [EvaluatorFunc] that additionally receives // the request's typed Config: the framework deserializes the request's raw -// options into it before calling the function (see [NewEvaluatorWithConfig]). -type EvaluatorFuncWithConfig[Config any] = func(context.Context, *EvaluatorCallbackRequest, Config) (*EvaluatorCallbackResponse, error) +// options into it before calling the function (see [NewTypedEvaluator]). +type TypedEvaluatorFunc[Config any] = func(context.Context, *EvaluatorCallbackRequest, Config) (*EvaluatorCallbackResponse, error) // BatchEvaluatorFunc is the function type for batch evaluator implementations. type BatchEvaluatorFunc = func(context.Context, *EvaluatorRequest) (*EvaluatorResponse, error) -// BatchEvaluatorFuncWithConfig is a [BatchEvaluatorFunc] that additionally +// TypedBatchEvaluatorFunc is a [BatchEvaluatorFunc] that additionally // receives the request's typed Config: the framework deserializes the // request's raw options into it before calling the function (see -// [NewBatchEvaluatorWithConfig]). -type BatchEvaluatorFuncWithConfig[Config any] = func(context.Context, *EvaluatorRequest, Config) (*EvaluatorResponse, error) +// [NewTypedBatchEvaluator]). +type TypedBatchEvaluatorFunc[Config any] = func(context.Context, *EvaluatorRequest, Config) (*EvaluatorResponse, error) // Evaluator represents a evaluator action. type Evaluator interface { @@ -185,7 +185,7 @@ func evaluatorMetadata(opts *EvaluatorOptions) map[string]any { } } -// NewEvaluatorWithConfig creates a new [Evaluator]. Register it with +// NewTypedEvaluator creates a new [Evaluator]. Register it with // [Evaluator.Register] to make it resolvable by name. // This method processes the input dataset one-by-one. // @@ -195,9 +195,9 @@ func evaluatorMetadata(opts *EvaluatorOptions) map[string]any { // map[string]any (from the Dev UI and other JSON callers) are accepted, and // mismatched types are rejected. The config's JSON schema is inferred from // Config unless [EvaluatorOptions.ConfigSchema] overrides it. -func NewEvaluatorWithConfig[Config any](name string, opts *EvaluatorOptions, fn EvaluatorFuncWithConfig[Config]) Evaluator { +func NewTypedEvaluator[Config any](name string, opts *EvaluatorOptions, fn TypedEvaluatorFunc[Config]) Evaluator { if name == "" { - panic("ai.NewEvaluatorWithConfig: evaluator name is required") + panic("ai.NewTypedEvaluator: evaluator name is required") } if opts == nil { @@ -273,17 +273,17 @@ func NewEvaluatorWithConfig[Config any](name string, opts *EvaluatorOptions, fn } } -// NewBatchEvaluatorWithConfig creates a new [Evaluator]. Register it with +// NewTypedBatchEvaluator creates a new [Evaluator]. Register it with // [Evaluator.Register] to make it resolvable by name. // This method provides the full [EvaluatorRequest] to the callback function, // giving more flexibility to the user for processing the data, such as batching or parallelization. // // Config is the evaluator's typed configuration; it is usually inferred from -// fn's signature. See [NewEvaluatorWithConfig] for how the request's options +// fn's signature. See [NewTypedEvaluator] for how the request's options // are deserialized. -func NewBatchEvaluatorWithConfig[Config any](name string, opts *EvaluatorOptions, fn BatchEvaluatorFuncWithConfig[Config]) Evaluator { +func NewTypedBatchEvaluator[Config any](name string, opts *EvaluatorOptions, fn TypedBatchEvaluatorFunc[Config]) Evaluator { if name == "" { - panic("ai.NewBatchEvaluatorWithConfig: batch evaluator name is required") + panic("ai.NewTypedBatchEvaluator: batch evaluator name is required") } if opts == nil { @@ -314,14 +314,14 @@ func NewBatchEvaluatorWithConfig[Config any](name string, opts *EvaluatorOptions // NewEvaluator creates a new [Evaluator]. // This method processes the input dataset one-by-one. // -// Deprecated: Use [NewEvaluatorWithConfig], which passes the request's +// Deprecated: Use [NewTypedEvaluator], which passes the request's // options to fn as a typed value instead of leaving them type-erased on the // request. func NewEvaluator(name string, opts *EvaluatorOptions, fn EvaluatorFunc) Evaluator { if name == "" { panic("ai.NewEvaluator: evaluator name is required") } - return NewEvaluatorWithConfig(name, opts, func(ctx context.Context, req *EvaluatorCallbackRequest, _ any) (*EvaluatorCallbackResponse, error) { + return NewTypedEvaluator(name, opts, func(ctx context.Context, req *EvaluatorCallbackRequest, _ any) (*EvaluatorCallbackResponse, error) { return fn(ctx, req) }) } @@ -329,7 +329,7 @@ func NewEvaluator(name string, opts *EvaluatorOptions, fn EvaluatorFunc) Evaluat // DefineEvaluator creates a new [Evaluator] and registers it. // This method processes the input dataset one-by-one. // -// Deprecated: Use [NewEvaluatorWithConfig] and register the result with +// Deprecated: Use [NewTypedEvaluator] and register the result with // [Evaluator.Register]. func DefineEvaluator(r api.Registry, name string, opts *EvaluatorOptions, fn EvaluatorFunc) Evaluator { e := NewEvaluator(name, opts, fn) @@ -341,14 +341,14 @@ func DefineEvaluator(r api.Registry, name string, opts *EvaluatorOptions, fn Eva // This method provides the full [EvaluatorRequest] to the callback function, // giving more flexibility to the user for processing the data, such as batching or parallelization. // -// Deprecated: Use [NewBatchEvaluatorWithConfig], which passes the request's +// Deprecated: Use [NewTypedBatchEvaluator], which passes the request's // options to fn as a typed value instead of leaving them type-erased on the // request. func NewBatchEvaluator(name string, opts *EvaluatorOptions, fn BatchEvaluatorFunc) Evaluator { if name == "" { panic("ai.NewBatchEvaluator: batch evaluator name is required") } - return NewBatchEvaluatorWithConfig(name, opts, func(ctx context.Context, req *EvaluatorRequest, _ any) (*EvaluatorResponse, error) { + return NewTypedBatchEvaluator(name, opts, func(ctx context.Context, req *EvaluatorRequest, _ any) (*EvaluatorResponse, error) { return fn(ctx, req) }) } @@ -357,7 +357,7 @@ func NewBatchEvaluator(name string, opts *EvaluatorOptions, fn BatchEvaluatorFun // This method provides the full [EvaluatorRequest] to the callback function, // giving more flexibility to the user for processing the data, such as batching or parallelization. // -// Deprecated: Use [NewBatchEvaluatorWithConfig] and register the result with +// Deprecated: Use [NewTypedBatchEvaluator] and register the result with // [Evaluator.Register]. func DefineBatchEvaluator(r api.Registry, name string, opts *EvaluatorOptions, fn BatchEvaluatorFunc) Evaluator { e := NewBatchEvaluator(name, opts, fn) diff --git a/go/ai/generate.go b/go/ai/generate.go index e53fc85c4b..f3c5def75b 100644 --- a/go/ai/generate.go +++ b/go/ai/generate.go @@ -69,10 +69,10 @@ type ToolConfig struct { // ModelFunc is a streaming function that takes in a ModelRequest and generates a ModelResponse, optionally streaming ModelResponseChunks. type ModelFunc = core.StreamingFunc[*ModelRequest, *ModelResponse, *ModelResponseChunk] -// ModelFuncWithConfig is a [ModelFunc] that additionally receives the +// TypedModelFunc is a [ModelFunc] that additionally receives the // request's typed Config: the framework deserializes the request's raw config -// into it before calling the function (see [NewModelWithConfig]). -type ModelFuncWithConfig[Config any] = func(context.Context, *ModelRequest, Config, ModelStreamCallback) (*ModelResponse, error) +// into it before calling the function (see [NewTypedModel]). +type TypedModelFunc[Config any] = func(context.Context, *ModelRequest, Config, ModelStreamCallback) (*ModelResponse, error) // ModelStreamCallback is a stream callback of a ModelAction. type ModelStreamCallback = func(context.Context, *ModelResponseChunk) error @@ -140,7 +140,7 @@ func DefineGenerateAction(ctx context.Context, r api.Registry) *generateAction { return (*generateAction)(a) } -// NewModelWithConfig creates a new [Model]. Register it with +// NewTypedModel creates a new [Model]. Register it with // [Model.Register] to make it resolvable by name. // // Config is the model's typed configuration; it is usually inferred from fn's @@ -157,9 +157,9 @@ func DefineGenerateAction(ctx context.Context, r api.Registry) *generateAction { // wrapper types like Opt[float64] that marshal to primitives but reflect as // objects), set [ModelOptions.ConfigSchema] explicitly or requests will be // rejected at the action boundary. -func NewModelWithConfig[Config any](name string, opts *ModelOptions, fn ModelFuncWithConfig[Config]) Model { +func NewTypedModel[Config any](name string, opts *ModelOptions, fn TypedModelFunc[Config]) Model { if name == "" { - panic("ai.NewModelWithConfig: name is required") + panic("ai.NewTypedModel: name is required") } if opts == nil { @@ -223,20 +223,20 @@ func NewModelWithConfig[Config any](name string, opts *ModelOptions, fn ModelFun // NewModel creates a new [Model]. // -// Deprecated: Use [NewModelWithConfig], which passes the request's config to +// Deprecated: Use [NewTypedModel], which passes the request's config to // fn as a typed value instead of leaving it type-erased on the request. func NewModel(name string, opts *ModelOptions, fn ModelFunc) Model { if name == "" { panic("ai.NewModel: name is required") } - return NewModelWithConfig(name, opts, func(ctx context.Context, req *ModelRequest, _ any, cb ModelStreamCallback) (*ModelResponse, error) { + return NewTypedModel(name, opts, func(ctx context.Context, req *ModelRequest, _ any, cb ModelStreamCallback) (*ModelResponse, error) { return fn(ctx, req, cb) }) } // DefineModel creates a new [Model] and registers it. // -// Deprecated: Use [NewModelWithConfig] and register the result with +// Deprecated: Use [NewTypedModel] and register the result with // [Model.Register]. func DefineModel(r api.Registry, name string, opts *ModelOptions, fn ModelFunc) Model { m := NewModel(name, opts, fn) diff --git a/go/core/doc.go b/go/core/doc.go index b283bf5dfa..15cb336356 100644 --- a/go/core/doc.go +++ b/go/core/doc.go @@ -160,7 +160,7 @@ and other capabilities. Implement the [api.Plugin] interface: func (p *MyPlugin) Init(ctx context.Context) []api.Action { // Initialize the plugin and return actions to register - model := ai.NewModel(...) + model := ai.NewTypedModel(...) tool := ai.NewTool(...) return []api.Action{model, tool} } diff --git a/go/genkit/genkit.go b/go/genkit/genkit.go index cdbf09035f..787992122a 100644 --- a/go/genkit/genkit.go +++ b/go/genkit/genkit.go @@ -498,25 +498,25 @@ func ListTools(g *Genkit) []ai.Tool { return tools } -// DefineModelWithConfig defines a custom model implementation, registers it +// DefineTypedModel defines a custom model implementation, registers it // as a [core.Action] of type Model, and returns an [ai.Model] interface. // // The `name` argument is the unique identifier for the model (e.g., "myProvider/myModel"). // The `opts` argument provides metadata about the model's capabilities ([ai.ModelOptions]). -// The `fn` argument ([ai.ModelFuncWithConfig]) implements the actual generation logic, +// The `fn` argument ([ai.TypedModelFunc]) implements the actual generation logic, // handling input requests ([ai.ModelRequest]) and producing responses ([ai.ModelResponse]), // potentially streaming chunks ([ai.ModelResponseChunk]) via the callback. // // Config is the model's typed configuration; it is usually inferred from fn's -// signature. See [ai.NewModelWithConfig] for how the request's config is +// signature. See [ai.NewTypedModel] for how the request's config is // deserialized and validated. // // For models that don't need to be registered (e.g., for plugin development or testing), -// use [ai.NewModelWithConfig] instead. +// use [ai.NewTypedModel] instead. // // Example: // -// echoModel := genkit.DefineModelWithConfig(g, "custom/echo", +// echoModel := genkit.DefineTypedModel(g, "custom/echo", // &ai.ModelOptions{ // Label: "Echo Model", // Supports: &ai.ModelSupports{Multiturn: true}, @@ -559,8 +559,8 @@ func ListTools(g *Genkit) []ai.Tool { // return resp, nil // }, // ) -func DefineModelWithConfig[Config any](g *Genkit, name string, opts *ai.ModelOptions, fn ai.ModelFuncWithConfig[Config]) ai.Model { - m := ai.NewModelWithConfig(name, opts, fn) +func DefineTypedModel[Config any](g *Genkit, name string, opts *ai.ModelOptions, fn ai.TypedModelFunc[Config]) ai.Model { + m := ai.NewTypedModel(name, opts, fn) m.Register(g.reg) return m } @@ -568,13 +568,13 @@ func DefineModelWithConfig[Config any](g *Genkit, name string, opts *ai.ModelOpt // DefineModel defines a custom model implementation, registers it as a [core.Action] // of type Model, and returns an [ai.Model] interface. // -// Deprecated: Use [DefineModelWithConfig], which passes the request's config +// Deprecated: Use [DefineTypedModel], which passes the request's config // to fn as a typed value instead of leaving it type-erased on the request. func DefineModel(g *Genkit, name string, opts *ai.ModelOptions, fn ai.ModelFunc) ai.Model { return ai.DefineModel(g.reg, name, opts, fn) } -// DefineBackgroundModelWithConfig defines a background model, registers it as +// DefineTypedBackgroundModel defines a background model, registers it as // an [ai.BackgroundModel], and returns an [ai.BackgroundModel]. // // The `name` is the identifier the model uses to request the background model. The `opts` @@ -582,10 +582,10 @@ func DefineModel(g *Genkit, name string, opts *ai.ModelOptions, fn ai.ModelFunc) // The `checkFn` is the function that checks the status of the background model. // // Config is the model's typed configuration; it is usually inferred from -// startFn's signature. See [ai.NewModelWithConfig] for how the request's +// startFn's signature. See [ai.NewTypedModel] for how the request's // config is deserialized and validated. -func DefineBackgroundModelWithConfig[Config any](g *Genkit, name string, opts *ai.BackgroundModelOptions, startFn ai.StartModelOpFuncWithConfig[Config], checkFn ai.CheckModelOpFunc) ai.BackgroundModel { - m := ai.NewBackgroundModelWithConfig(name, opts, startFn, checkFn) +func DefineTypedBackgroundModel[Config any](g *Genkit, name string, opts *ai.BackgroundModelOptions, startFn ai.TypedStartModelOpFunc[Config], checkFn ai.CheckModelOpFunc) ai.BackgroundModel { + m := ai.NewTypedBackgroundModel(name, opts, startFn, checkFn) m.Register(g.reg) return m } @@ -593,7 +593,7 @@ func DefineBackgroundModelWithConfig[Config any](g *Genkit, name string, opts *a // DefineBackgroundModel defines a background model, registers it as a [ai.BackgroundModel], // and returns an [ai.BackgroundModel]. // -// Deprecated: Use [DefineBackgroundModelWithConfig], which passes the +// Deprecated: Use [DefineTypedBackgroundModel], which passes the // request's config to startFn as a typed value instead of leaving it // type-erased on the request. func DefineBackgroundModel(g *Genkit, name string, opts *ai.BackgroundModelOptions, startFn ai.StartModelOpFunc, checkFn ai.CheckModelOpFunc) ai.BackgroundModel { @@ -1370,7 +1370,7 @@ func LookupRetriever(g *Genkit, name string) ai.Retriever { return ai.LookupRetriever(g.reg, name) } -// DefineEmbedderWithConfig defines a custom text embedding implementation, +// DefineTypedEmbedder defines a custom text embedding implementation, // registers it as a [core.Action] of type Embedder, and returns an // [ai.Embedder]. Embedders convert text documents or queries into numerical // vector representations (embeddings). @@ -1380,13 +1380,13 @@ func LookupRetriever(g *Genkit, name string) ai.Retriever { // and return an [ai.EmbedResponse] (containing the corresponding embeddings). // // Config is the embedder's typed configuration; it is usually inferred from -// fn's signature. See [ai.NewEmbedderWithConfig] for how the request's +// fn's signature. See [ai.NewTypedEmbedder] for how the request's // options are deserialized. // // For embedders that don't need to be registered (e.g., for plugin development), -// use [ai.NewEmbedderWithConfig] instead. -func DefineEmbedderWithConfig[Config any](g *Genkit, name string, opts *ai.EmbedderOptions, fn ai.EmbedderFuncWithConfig[Config]) ai.Embedder { - e := ai.NewEmbedderWithConfig(name, opts, fn) +// use [ai.NewTypedEmbedder] instead. +func DefineTypedEmbedder[Config any](g *Genkit, name string, opts *ai.EmbedderOptions, fn ai.TypedEmbedderFunc[Config]) ai.Embedder { + e := ai.NewTypedEmbedder(name, opts, fn) e.Register(g.reg) return e } @@ -1394,7 +1394,7 @@ func DefineEmbedderWithConfig[Config any](g *Genkit, name string, opts *ai.Embed // DefineEmbedder defines a custom text embedding implementation, registers it as a // [core.Action] of type Embedder, and returns an [ai.Embedder]. // -// Deprecated: Use [DefineEmbedderWithConfig], which passes the request's +// Deprecated: Use [DefineTypedEmbedder], which passes the request's // options to fn as a typed value instead of leaving them type-erased on the // request. func DefineEmbedder(g *Genkit, name string, opts *ai.EmbedderOptions, fn ai.EmbedderFunc) ai.Embedder { @@ -1418,7 +1418,7 @@ func LookupPlugin(g *Genkit, name string) api.Plugin { return g.reg.LookupPlugin(name) } -// DefineEvaluatorWithConfig defines an evaluator that processes test cases +// DefineTypedEvaluator defines an evaluator that processes test cases // one by one, registers it as a [core.Action] of type Evaluator, and returns // an [ai.Evaluator]. Evaluators are used to assess the quality or performance // of AI models or flows based on a dataset of test cases. @@ -1427,10 +1427,10 @@ func LookupPlugin(g *Genkit, name string) api.Plugin { // ([ai.EvaluatorCallbackRequest]) in the evaluation dataset. // // Config is the evaluator's typed configuration; it is usually inferred from -// fn's signature. See [ai.NewEvaluatorWithConfig] for how the request's +// fn's signature. See [ai.NewTypedEvaluator] for how the request's // options are deserialized. -func DefineEvaluatorWithConfig[Config any](g *Genkit, name string, opts *ai.EvaluatorOptions, fn ai.EvaluatorFuncWithConfig[Config]) ai.Evaluator { - e := ai.NewEvaluatorWithConfig(name, opts, fn) +func DefineTypedEvaluator[Config any](g *Genkit, name string, opts *ai.EvaluatorOptions, fn ai.TypedEvaluatorFunc[Config]) ai.Evaluator { + e := ai.NewTypedEvaluator(name, opts, fn) e.Register(g.reg) return e } @@ -1438,14 +1438,14 @@ func DefineEvaluatorWithConfig[Config any](g *Genkit, name string, opts *ai.Eval // DefineEvaluator defines an evaluator that processes test cases one by one, // registers it as a [core.Action] of type Evaluator, and returns an [ai.Evaluator]. // -// Deprecated: Use [DefineEvaluatorWithConfig], which passes the request's +// Deprecated: Use [DefineTypedEvaluator], which passes the request's // options to fn as a typed value instead of leaving them type-erased on the // request. func DefineEvaluator(g *Genkit, name string, opts *ai.EvaluatorOptions, fn ai.EvaluatorFunc) ai.Evaluator { return ai.DefineEvaluator(g.reg, name, opts, fn) } -// DefineBatchEvaluatorWithConfig defines an evaluator that processes the +// DefineTypedBatchEvaluator defines an evaluator that processes the // entire dataset at once, registers it as a [core.Action] of type Evaluator, // and returns an [ai.Evaluator]. // @@ -1454,10 +1454,10 @@ func DefineEvaluator(g *Genkit, name string, opts *ai.EvaluatorOptions, fn ai.Ev // such as batching calls to external services or parallelizing computations. // // Config is the evaluator's typed configuration; it is usually inferred from -// fn's signature. See [ai.NewEvaluatorWithConfig] for how the request's +// fn's signature. See [ai.NewTypedEvaluator] for how the request's // options are deserialized. -func DefineBatchEvaluatorWithConfig[Config any](g *Genkit, name string, opts *ai.EvaluatorOptions, fn ai.BatchEvaluatorFuncWithConfig[Config]) ai.Evaluator { - e := ai.NewBatchEvaluatorWithConfig(name, opts, fn) +func DefineTypedBatchEvaluator[Config any](g *Genkit, name string, opts *ai.EvaluatorOptions, fn ai.TypedBatchEvaluatorFunc[Config]) ai.Evaluator { + e := ai.NewTypedBatchEvaluator(name, opts, fn) e.Register(g.reg) return e } @@ -1465,7 +1465,7 @@ func DefineBatchEvaluatorWithConfig[Config any](g *Genkit, name string, opts *ai // DefineBatchEvaluator defines an evaluator that processes the entire dataset at once, // registers it as a [core.Action] of type Evaluator, and returns an [ai.Evaluator]. // -// Deprecated: Use [DefineBatchEvaluatorWithConfig], which passes the +// Deprecated: Use [DefineTypedBatchEvaluator], which passes the // request's options to fn as a typed value instead of leaving them // type-erased on the request. func DefineBatchEvaluator(g *Genkit, name string, opts *ai.EvaluatorOptions, fn ai.BatchEvaluatorFunc) ai.Evaluator { From d9a857ae834c97e7e9629f4e2023bf00b8aacf9e Mon Sep 17 00:00:00 2001 From: Alex Pascal Date: Thu, 30 Jul 2026 09:30:58 -0700 Subject: [PATCH 04/10] refactor(go/ai)!: delete the registry-taking Define helpers A registry is unobtainable outside the framework: plugins return actions from Init for the framework to register and applications go through the genkit package, so the registry-taking Define helpers had no reachable external callers. Instead of carrying them through a deprecation cycle, delete them; the genkit wrappers remain the user touchpoints and now compose New* + Register directly. Deleted: ai.DefineModel, ai.DefineEmbedder, ai.DefineEvaluator, ai.DefineBatchEvaluator, ai.DefineBackgroundModel, ai.DefineMiddleware, ai.DefineTool, ai.DefineToolWithInputSchema, ai.DefineMultipartTool, ai.DefineResource, and the experimental exp.DefineTool and exp.DefineInterruptibleTool. Kept, because the registry is a structural input rather than just a registration target: ai.DefinePrompt, ai.DefineDataPrompt, ai.DefineFormat, ai.DefineGenerateAction (init bootstrap), the exp agent constructors, and the frozen retriever surface. Tests use package-local define helpers that compose New* + Register, which also exercises the composition path users are pointed at. --- go/ai/action_test.go | 4 +- go/ai/background_model.go | 12 +-- go/ai/config_test.go | 2 +- go/ai/define_test.go | 70 +++++++++++++++ go/ai/embedder.go | 13 +-- go/ai/embedder_test.go | 22 ++--- go/ai/evaluator.go | 25 +----- go/ai/evaluator_test.go | 18 ++-- go/ai/exp/agent_test.go | 24 ++--- go/ai/exp/agents_conformance_test.go | 23 ++++- go/ai/exp/define_test.go | 49 +++++++++++ go/ai/exp/tools.go | 31 +------ go/ai/exp/tools_test.go | 18 ++-- go/ai/formatter_test.go | 2 +- go/ai/generate.go | 12 +-- go/ai/generate_test.go | 88 +++++++++---------- go/ai/middleware.go | 7 -- go/ai/middleware_test.go | 6 +- go/ai/prompt_test.go | 46 +++++----- go/ai/resource.go | 8 -- go/ai/resource_test.go | 16 ++-- go/ai/testutil_test.go | 6 +- go/ai/tools.go | 70 --------------- go/ai/tools_test.go | 56 ++++++------ go/genkit/exp/tools.go | 8 +- go/genkit/genkit.go | 40 ++++++--- go/internal/fakeembedder/fakeembedder_test.go | 3 +- go/plugins/middleware/define_test.go | 43 +++++++++ go/plugins/middleware/filesystem_test.go | 44 +++++----- go/plugins/middleware/retry_test.go | 20 ++--- go/plugins/middleware/skills_test.go | 14 +-- go/plugins/middleware/tool_approval_test.go | 14 +-- 32 files changed, 424 insertions(+), 390 deletions(-) create mode 100644 go/ai/define_test.go create mode 100644 go/ai/exp/define_test.go create mode 100644 go/plugins/middleware/define_test.go diff --git a/go/ai/action_test.go b/go/ai/action_test.go index d42f5c8a8b..d7509f100a 100644 --- a/go/ai/action_test.go +++ b/go/ai/action_test.go @@ -71,7 +71,7 @@ func defineProgrammableModel(r api.Registry) *programmableModel { Tools: true, Multiturn: true, } - DefineModel(r, "programmableModel", &ModelOptions{Supports: supports}, func(ctx context.Context, req *ModelRequest, cb ModelStreamCallback) (*ModelResponse, error) { + defineModel(r, "programmableModel", &ModelOptions{Supports: supports}, func(ctx context.Context, req *ModelRequest, cb ModelStreamCallback) (*ModelResponse, error) { return pm.Generate(ctx, r, req, &ToolConfig{MaxTurns: 5}, cb) }) return pm @@ -97,7 +97,7 @@ func TestGenerateAction(t *testing.T) { pm := defineProgrammableModel(r) - DefineTool(r, "testTool", "description", + defineTool(r, "testTool", "description", func(ctx *ToolContext, input any) (any, error) { return "tool called", nil }) diff --git a/go/ai/background_model.go b/go/ai/background_model.go index 0cfc6d4887..c982e88d55 100644 --- a/go/ai/background_model.go +++ b/go/ai/background_model.go @@ -71,7 +71,7 @@ type BackgroundModelOptions struct { Metadata map[string]any // Additional metadata. } -// LookupBackgroundModel looks up a BackgroundAction registered by [DefineBackgroundModel]. +// LookupBackgroundModel looks up a registered [BackgroundModel] by name. // It returns nil if the background model was not found. func LookupBackgroundModel(r api.Registry, name string) BackgroundModel { key := api.KeyFromName(api.ActionTypeBackgroundModel, name) @@ -181,16 +181,6 @@ func NewBackgroundModel(name string, opts *BackgroundModelOptions, startFn Start }, checkFn) } -// DefineBackgroundModel defines and registers a new model that runs in the background. -// -// Deprecated: Use [NewTypedBackgroundModel] and register the result with -// [BackgroundModel.Register]. -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...) diff --git a/go/ai/config_test.go b/go/ai/config_test.go index 51055c20d4..abf1ba05a3 100644 --- a/go/ai/config_test.go +++ b/go/ai/config_test.go @@ -220,7 +220,7 @@ func TestDeprecatedModelPassesConfigThrough(t *testing.T) { r := registry.New() var gotConfig any - m := DefineModel(r, "test/legacy-config", nil, func(ctx context.Context, req *ModelRequest, cb ModelStreamCallback) (*ModelResponse, error) { + m := defineModel(r, "test/legacy-config", nil, func(ctx context.Context, req *ModelRequest, cb ModelStreamCallback) (*ModelResponse, error) { gotConfig = req.Config return &ModelResponse{Message: NewModelTextMessage("ok"), Request: req}, nil }) diff --git a/go/ai/define_test.go b/go/ai/define_test.go new file mode 100644 index 0000000000..b15ef42ed0 --- /dev/null +++ b/go/ai/define_test.go @@ -0,0 +1,70 @@ +// Copyright 2026 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 ( + "github.com/firebase/genkit/go/core/api" +) + +// Test-local define helpers: New* + Register in one call, mirroring the +// removed registry-taking Define functions so tests stay concise. + +func defineModel(r api.Registry, name string, opts *ModelOptions, fn ModelFunc) Model { + m := NewModel(name, opts, fn) + m.Register(r) + return m +} + +func defineEmbedder(r api.Registry, name string, opts *EmbedderOptions, fn EmbedderFunc) Embedder { + e := NewEmbedder(name, opts, fn) + e.Register(r) + return e +} + +func defineEvaluator(r api.Registry, name string, opts *EvaluatorOptions, fn EvaluatorFunc) Evaluator { + e := NewEvaluator(name, opts, fn) + e.Register(r) + return e +} + +func defineBatchEvaluator(r api.Registry, name string, opts *EvaluatorOptions, fn BatchEvaluatorFunc) Evaluator { + e := NewBatchEvaluator(name, opts, fn) + e.Register(r) + return e +} + +func defineTool[In, Out any](r api.Registry, name, description string, fn ToolFunc[In, Out], opts ...ToolOption) *ToolDef[In, Out] { + t := NewTool(name, description, fn, opts...) + t.Register(r) + return t +} + +func defineMultipartTool[In any](r api.Registry, name, description string, fn MultipartToolFunc[In], opts ...ToolOption) *ToolDef[In, *MultipartToolResponse] { + t := NewMultipartTool(name, description, fn, opts...) + t.Register(r) + return t +} + +func defineResource(r api.Registry, name string, opts *ResourceOptions, fn ResourceFunc) Resource { + res := NewResource(name, opts, fn) + res.Register(r) + return res +} + +func defineToolWithInputSchema[Out any](r api.Registry, name, description string, inputSchema map[string]any, fn ToolFunc[any, Out]) *ToolDef[any, Out] { + return defineTool(r, name, description, fn, WithInputSchema(inputSchema)) +} diff --git a/go/ai/embedder.go b/go/ai/embedder.go index 5976954252..cc7c9281ec 100644 --- a/go/ai/embedder.go +++ b/go/ai/embedder.go @@ -168,18 +168,7 @@ func NewEmbedder(name string, opts *EmbedderOptions, fn EmbedderFunc) Embedder { }) } -// DefineEmbedder registers the given embed function as an action, and returns an -// [Embedder] that runs it. -// -// Deprecated: Use [NewTypedEmbedder] and register the result with -// [Embedder.Register]. -func DefineEmbedder(r api.Registry, name string, opts *EmbedderOptions, fn EmbedderFunc) Embedder { - e := NewEmbedder(name, opts, fn) - e.Register(r) - return e -} - -// LookupEmbedder looks up an [Embedder] registered by [DefineEmbedder]. +// LookupEmbedder looks up a registered [Embedder] by name. // It will try to resolve the embedder dynamically if the embedder is not found. // It returns nil if the embedder was not resolved. func LookupEmbedder(r api.Registry, name string) Embedder { diff --git a/go/ai/embedder_test.go b/go/ai/embedder_test.go index 43404479f1..c7a4e025a6 100644 --- a/go/ai/embedder_test.go +++ b/go/ai/embedder_test.go @@ -112,7 +112,7 @@ func TestDefineEmbedder(t *testing.T) { r := newTestRegistry(t) called := false - e := DefineEmbedder(r, "test/defineEmbedder", nil, func(ctx context.Context, req *EmbedRequest) (*EmbedResponse, error) { + e := defineEmbedder(r, "test/defineEmbedder", nil, func(ctx context.Context, req *EmbedRequest) (*EmbedResponse, error) { called = true return &EmbedResponse{ Embeddings: []*Embedding{{Embedding: []float32{0.1, 0.2, 0.3}}}, @@ -146,7 +146,7 @@ func TestDefineEmbedder(t *testing.T) { func TestLookupEmbedder(t *testing.T) { t.Run("returns embedder when found", func(t *testing.T) { r := newTestRegistry(t) - DefineEmbedder(r, "test/lookupEmbedder", nil, func(ctx context.Context, req *EmbedRequest) (*EmbedResponse, error) { + defineEmbedder(r, "test/lookupEmbedder", nil, func(ctx context.Context, req *EmbedRequest) (*EmbedResponse, error) { return &EmbedResponse{}, nil }) @@ -171,7 +171,7 @@ func TestEmbedderEmbed(t *testing.T) { r := newTestRegistry(t) var capturedReq *EmbedRequest - e := DefineEmbedder(r, "test/embedDocuments", nil, func(ctx context.Context, req *EmbedRequest) (*EmbedResponse, error) { + e := defineEmbedder(r, "test/embedDocuments", nil, func(ctx context.Context, req *EmbedRequest) (*EmbedResponse, error) { capturedReq = req embeddings := make([]*Embedding, len(req.Input)) for i := range req.Input { @@ -210,7 +210,7 @@ func TestEmbedderEmbed(t *testing.T) { r := newTestRegistry(t) expectedErr := errors.New("embedding failed") - e := DefineEmbedder(r, "test/embedError", nil, func(ctx context.Context, req *EmbedRequest) (*EmbedResponse, error) { + e := defineEmbedder(r, "test/embedError", nil, func(ctx context.Context, req *EmbedRequest) (*EmbedResponse, error) { return nil, expectedErr }) @@ -226,7 +226,7 @@ func TestEmbedderEmbed(t *testing.T) { r := newTestRegistry(t) var capturedOpts any - e := DefineEmbedder(r, "test/embedOpts", nil, func(ctx context.Context, req *EmbedRequest) (*EmbedResponse, error) { + e := defineEmbedder(r, "test/embedOpts", nil, func(ctx context.Context, req *EmbedRequest) (*EmbedResponse, error) { capturedOpts = req.Options return &EmbedResponse{Embeddings: []*Embedding{{Embedding: []float32{0.1}}}}, nil }) @@ -247,7 +247,7 @@ func TestEmbedderEmbed(t *testing.T) { func TestEmbedFunction(t *testing.T) { t.Run("embeds with embedder directly", func(t *testing.T) { r := newTestRegistry(t) - e := DefineEmbedder(r, "test/embedFunc", nil, func(ctx context.Context, req *EmbedRequest) (*EmbedResponse, error) { + e := defineEmbedder(r, "test/embedFunc", nil, func(ctx context.Context, req *EmbedRequest) (*EmbedResponse, error) { return &EmbedResponse{ Embeddings: []*Embedding{{Embedding: []float32{0.1, 0.2, 0.3}}}, }, nil @@ -266,7 +266,7 @@ func TestEmbedFunction(t *testing.T) { t.Run("embeds with embedder ref", func(t *testing.T) { r := newTestRegistry(t) - DefineEmbedder(r, "test/embedFuncRef", nil, func(ctx context.Context, req *EmbedRequest) (*EmbedResponse, error) { + defineEmbedder(r, "test/embedFuncRef", nil, func(ctx context.Context, req *EmbedRequest) (*EmbedResponse, error) { return &EmbedResponse{ Embeddings: []*Embedding{{Embedding: []float32{0.1, 0.2, 0.3}}}, }, nil @@ -286,7 +286,7 @@ func TestEmbedFunction(t *testing.T) { t.Run("embeds with embedder name", func(t *testing.T) { r := newTestRegistry(t) - DefineEmbedder(r, "test/embedFuncName", nil, func(ctx context.Context, req *EmbedRequest) (*EmbedResponse, error) { + defineEmbedder(r, "test/embedFuncName", nil, func(ctx context.Context, req *EmbedRequest) (*EmbedResponse, error) { return &EmbedResponse{ Embeddings: []*Embedding{{Embedding: []float32{0.1, 0.2, 0.3}}}, }, nil @@ -307,7 +307,7 @@ func TestEmbedFunction(t *testing.T) { r := newTestRegistry(t) var capturedOpts any - DefineEmbedder(r, "test/embedRefConfig", nil, func(ctx context.Context, req *EmbedRequest) (*EmbedResponse, error) { + defineEmbedder(r, "test/embedRefConfig", nil, func(ctx context.Context, req *EmbedRequest) (*EmbedResponse, error) { capturedOpts = req.Options return &EmbedResponse{Embeddings: []*Embedding{{Embedding: []float32{0.1}}}}, nil }) @@ -330,7 +330,7 @@ func TestEmbedFunction(t *testing.T) { r := newTestRegistry(t) var capturedOpts any - DefineEmbedder(r, "test/embedOverrideConfig", nil, func(ctx context.Context, req *EmbedRequest) (*EmbedResponse, error) { + defineEmbedder(r, "test/embedOverrideConfig", nil, func(ctx context.Context, req *EmbedRequest) (*EmbedResponse, error) { capturedOpts = req.Options return &EmbedResponse{Embeddings: []*Embedding{{Embedding: []float32{0.1}}}}, nil }) @@ -374,7 +374,7 @@ func TestEmbedFunction(t *testing.T) { r := newTestRegistry(t) var capturedDocs []*Document - DefineEmbedder(r, "test/embedDocs", nil, func(ctx context.Context, req *EmbedRequest) (*EmbedResponse, error) { + defineEmbedder(r, "test/embedDocs", nil, func(ctx context.Context, req *EmbedRequest) (*EmbedResponse, error) { capturedDocs = req.Input embeddings := make([]*Embedding, len(req.Input)) for i := range req.Input { diff --git a/go/ai/evaluator.go b/go/ai/evaluator.go index 70a5f4cd4c..64340457e9 100644 --- a/go/ai/evaluator.go +++ b/go/ai/evaluator.go @@ -326,17 +326,6 @@ func NewEvaluator(name string, opts *EvaluatorOptions, fn EvaluatorFunc) Evaluat }) } -// DefineEvaluator creates a new [Evaluator] and registers it. -// This method processes the input dataset one-by-one. -// -// Deprecated: Use [NewTypedEvaluator] and register the result with -// [Evaluator.Register]. -func DefineEvaluator(r api.Registry, name string, opts *EvaluatorOptions, fn EvaluatorFunc) Evaluator { - e := NewEvaluator(name, opts, fn) - e.Register(r) - return e -} - // NewBatchEvaluator creates a new [Evaluator]. // This method provides the full [EvaluatorRequest] to the callback function, // giving more flexibility to the user for processing the data, such as batching or parallelization. @@ -353,19 +342,7 @@ func NewBatchEvaluator(name string, opts *EvaluatorOptions, fn BatchEvaluatorFun }) } -// DefineBatchEvaluator creates a new [Evaluator] and registers it. -// This method provides the full [EvaluatorRequest] to the callback function, -// giving more flexibility to the user for processing the data, such as batching or parallelization. -// -// Deprecated: Use [NewTypedBatchEvaluator] and register the result with -// [Evaluator.Register]. -func DefineBatchEvaluator(r api.Registry, name string, opts *EvaluatorOptions, fn BatchEvaluatorFunc) Evaluator { - e := NewBatchEvaluator(name, opts, fn) - e.Register(r) - return e -} - -// LookupEvaluator looks up an [Evaluator] registered by [DefineEvaluator]. +// LookupEvaluator looks up a registered [Evaluator] by name. // It returns nil if the evaluator was not defined. func LookupEvaluator(r api.Registry, name string) Evaluator { action := core.ResolveActionFor[*EvaluatorRequest, *EvaluatorResponse, struct{}](r, api.ActionTypeEvaluator, name) diff --git a/go/ai/evaluator_test.go b/go/ai/evaluator_test.go index 6cf5c58953..4b23ea2190 100644 --- a/go/ai/evaluator_test.go +++ b/go/ai/evaluator_test.go @@ -90,7 +90,7 @@ var testRequest = EvaluatorRequest{ func TestSimpleEvaluator(t *testing.T) { r := registry.New() - evaluator := DefineEvaluator(r, "test/testEvaluator", &evalOpts, testEvalFunc) + evaluator := defineEvaluator(r, "test/testEvaluator", &evalOpts, testEvalFunc) resp, err := evaluator.Evaluate(context.Background(), &testRequest) if err != nil { @@ -117,14 +117,14 @@ func TestSimpleEvaluator(t *testing.T) { func TestOptionsRequired(t *testing.T) { r := registry.New() - _ = DefineEvaluator(r, "test/testEvaluator", &evalOpts, testEvalFunc) - _ = DefineBatchEvaluator(r, "test/testBatchEvaluator", &evalOpts, testBatchEvalFunc) + _ = defineEvaluator(r, "test/testEvaluator", &evalOpts, testEvalFunc) + _ = defineBatchEvaluator(r, "test/testBatchEvaluator", &evalOpts, testBatchEvalFunc) } func TestFailingEvaluator(t *testing.T) { r := registry.New() - evalAction := DefineEvaluator(r, "test/testEvaluator", &evalOpts, testFailingEvalFunc) + evalAction := defineEvaluator(r, "test/testEvaluator", &evalOpts, testFailingEvalFunc) resp, err := evalAction.Evaluate(context.Background(), &testRequest) if err != nil { @@ -142,8 +142,8 @@ func TestFailingEvaluator(t *testing.T) { func TestLookupEvaluator(t *testing.T) { r := registry.New() - DefineEvaluator(r, "test/testEvaluator", &evalOpts, testEvalFunc) - DefineBatchEvaluator(r, "test/testBatchEvaluator", &evalOpts, testBatchEvalFunc) + defineEvaluator(r, "test/testEvaluator", &evalOpts, testEvalFunc) + defineBatchEvaluator(r, "test/testBatchEvaluator", &evalOpts, testBatchEvalFunc) if LookupEvaluator(r, "test/testEvaluator") == nil { t.Errorf("LookupEvaluator(r, \"test/testEvaluator\") is nil") @@ -156,7 +156,7 @@ func TestLookupEvaluator(t *testing.T) { func TestEvaluate(t *testing.T) { r := registry.New() - evalAction := DefineEvaluator(r, "test/testEvaluator", &evalOpts, testEvalFunc) + evalAction := defineEvaluator(r, "test/testEvaluator", &evalOpts, testEvalFunc) resp, err := Evaluate(context.Background(), r, WithEvaluator(evalAction), @@ -184,7 +184,7 @@ func TestEvaluate(t *testing.T) { func TestBatchEvaluator(t *testing.T) { r := registry.New() - evalAction := DefineBatchEvaluator(r, "test/testBatchEvaluator", &evalOpts, testBatchEvalFunc) + evalAction := defineBatchEvaluator(r, "test/testBatchEvaluator", &evalOpts, testBatchEvalFunc) resp, err := evalAction.Evaluate(context.Background(), &testRequest) if err != nil { @@ -249,7 +249,7 @@ func TestEvaluatorRefUsedWithEvaluate(t *testing.T) { r := registry.New() // Define evaluator that uses config - DefineEvaluator(r, "test/configEvaluator", &evalOpts, func(ctx context.Context, req *EvaluatorCallbackRequest) (*EvaluatorCallbackResponse, error) { + defineEvaluator(r, "test/configEvaluator", &evalOpts, func(ctx context.Context, req *EvaluatorCallbackRequest) (*EvaluatorCallbackResponse, error) { score := Score{ Id: "configScore", Score: 1, diff --git a/go/ai/exp/agent_test.go b/go/ai/exp/agent_test.go index f60f21aeb2..e0a2d47108 100644 --- a/go/ai/exp/agent_test.go +++ b/go/ai/exp/agent_test.go @@ -1651,7 +1651,7 @@ func setupPromptTestRegistry(t *testing.T) *registry.Registry { ctx := context.Background() ai.ConfigureFormats(reg) - ai.DefineModel(reg, "test/echo", &ai.ModelOptions{Supports: &ai.ModelSupports{Multiturn: true, SystemRole: true}}, + defineTestModel(reg, "test/echo", &ai.ModelOptions{Supports: &ai.ModelSupports{Multiturn: true, SystemRole: true}}, func(ctx context.Context, req *ai.ModelRequest, cb ai.ModelStreamCallback) (*ai.ModelResponse, error) { // Echo back the last user message text. var text string @@ -1701,7 +1701,7 @@ func TestPromptAgent_NamedPromptSharedAcrossAgents(t *testing.T) { var mu sync.Mutex var renderedSystems []string - ai.DefineModel(reg, "test/capture", &ai.ModelOptions{Supports: &ai.ModelSupports{Multiturn: true, SystemRole: true}}, + defineTestModel(reg, "test/capture", &ai.ModelOptions{Supports: &ai.ModelSupports{Multiturn: true, SystemRole: true}}, func(ctx context.Context, req *ai.ModelRequest, cb ai.ModelStreamCallback) (*ai.ModelResponse, error) { mu.Lock() for _, m := range req.Messages { @@ -1765,7 +1765,7 @@ func TestDefinePromptAgent_DefaultAndNamed(t *testing.T) { var mu sync.Mutex var renderedSystems []string - ai.DefineModel(reg, "test/capture", &ai.ModelOptions{Supports: &ai.ModelSupports{Multiturn: true, SystemRole: true}}, + defineTestModel(reg, "test/capture", &ai.ModelOptions{Supports: &ai.ModelSupports{Multiturn: true, SystemRole: true}}, func(ctx context.Context, req *ai.ModelRequest, cb ai.ModelStreamCallback) (*ai.ModelResponse, error) { mu.Lock() for _, m := range req.Messages { @@ -1908,7 +1908,7 @@ func TestPromptAgent_MultiTurnHistory(t *testing.T) { reg := setupPromptTestRegistry(t) // Use a model that echoes all message count so we can verify history grows. - ai.DefineModel(reg, "test/history", &ai.ModelOptions{Supports: &ai.ModelSupports{Multiturn: true, SystemRole: true}}, + defineTestModel(reg, "test/history", &ai.ModelOptions{Supports: &ai.ModelSupports{Multiturn: true, SystemRole: true}}, func(ctx context.Context, req *ai.ModelRequest, cb ai.ModelStreamCallback) (*ai.ModelResponse, error) { // Count total messages received (includes prompt-rendered + history). var parts []string @@ -2055,14 +2055,14 @@ func TestPromptAgent_ToolLoopMessages(t *testing.T) { ai.ConfigureFormats(reg) // Define two tools so the model can call them across multiple rounds. - ai.DefineTool(reg, "greet", "returns a greeting", + defineTestTool(reg, "greet", "returns a greeting", func(ctx *ai.ToolContext, input struct { Name string `json:"name"` }) (string, error) { return "hello " + input.Name, nil }, ) - ai.DefineTool(reg, "farewell", "returns a farewell", + defineTestTool(reg, "farewell", "returns a farewell", func(ctx *ai.ToolContext, input struct { Name string `json:"name"` }) (string, error) { @@ -2074,7 +2074,7 @@ func TestPromptAgent_ToolLoopMessages(t *testing.T) { // Round 1: request "greet" tool // Round 2: after seeing greet response, request "farewell" tool // Round 3: after seeing farewell response, return final text - ai.DefineModel(reg, "test/toolmodel", &ai.ModelOptions{Supports: &ai.ModelSupports{Multiturn: true, SystemRole: true, Tools: true}}, + defineTestModel(reg, "test/toolmodel", &ai.ModelOptions{Supports: &ai.ModelSupports{Multiturn: true, SystemRole: true, Tools: true}}, func(ctx context.Context, req *ai.ModelRequest, cb ai.ModelStreamCallback) (*ai.ModelResponse, error) { // Count tool responses to determine which round we're in. toolResps := 0 @@ -2362,7 +2362,7 @@ func TestPromptAgent_RejectsInvalidInputMessage(t *testing.T) { ctx := context.Background() reg := setupPromptTestRegistry(t) var modelCalls atomic.Int64 - ai.DefineModel(reg, "test/reject", &ai.ModelOptions{Supports: &ai.ModelSupports{Multiturn: true}}, + defineTestModel(reg, "test/reject", &ai.ModelOptions{Supports: &ai.ModelSupports{Multiturn: true}}, func(ctx context.Context, req *ai.ModelRequest, cb ai.ModelStreamCallback) (*ai.ModelResponse, error) { modelCalls.Add(1) return &ai.ModelResponse{Message: ai.NewModelTextMessage("unexpected")}, nil @@ -2539,7 +2539,7 @@ func TestPromptAgent_RejectsResumeForUnrequestedTool(t *testing.T) { ai.ConfigureFormats(reg) var modelCalls atomic.Int32 - ai.DefineModel(reg, "test/plain", &ai.ModelOptions{Supports: &ai.ModelSupports{Multiturn: true, Tools: true}}, + defineTestModel(reg, "test/plain", &ai.ModelOptions{Supports: &ai.ModelSupports{Multiturn: true, Tools: true}}, func(ctx context.Context, req *ai.ModelRequest, cb ai.ModelStreamCallback) (*ai.ModelResponse, error) { modelCalls.Add(1) return &ai.ModelResponse{Request: req, Message: ai.NewModelTextMessage("hello")}, nil @@ -5488,7 +5488,7 @@ func TestPromptAgent_ForwardsFinishReason(t *testing.T) { ctx := context.Background() reg := registry.New() ai.ConfigureFormats(reg) - ai.DefineModel(reg, "test/length", &ai.ModelOptions{Supports: &ai.ModelSupports{Multiturn: true, SystemRole: true}}, + defineTestModel(reg, "test/length", &ai.ModelOptions{Supports: &ai.ModelSupports{Multiturn: true, SystemRole: true}}, func(ctx context.Context, req *ai.ModelRequest, cb ai.ModelStreamCallback) (*ai.ModelResponse, error) { return &ai.ModelResponse{ Request: req, @@ -5880,14 +5880,14 @@ func TestPromptAgent_ForwardsInterruptedFinishReason(t *testing.T) { reg := registry.New() ai.ConfigureFormats(reg) - interruptTool := ai.DefineTool(reg, "interruptor", "always interrupts", + interruptTool := defineTestTool(reg, "interruptor", "always interrupts", func(tc *ai.ToolContext, input any) (any, error) { return nil, tc.Interrupt(&ai.InterruptOptions{ Metadata: map[string]any{"reason": "needs approval"}, }) }, ) - ai.DefineModel(reg, "test/interrupt", &ai.ModelOptions{Supports: &ai.ModelSupports{Multiturn: true, Tools: true}}, + defineTestModel(reg, "test/interrupt", &ai.ModelOptions{Supports: &ai.ModelSupports{Multiturn: true, Tools: true}}, func(ctx context.Context, req *ai.ModelRequest, cb ai.ModelStreamCallback) (*ai.ModelResponse, error) { return &ai.ModelResponse{ Request: req, diff --git a/go/ai/exp/agents_conformance_test.go b/go/ai/exp/agents_conformance_test.go index 12ace7ff2c..de0e331ce4 100644 --- a/go/ai/exp/agents_conformance_test.go +++ b/go/ai/exp/agents_conformance_test.go @@ -48,6 +48,7 @@ import ( "github.com/firebase/genkit/go/ai/exp" "github.com/firebase/genkit/go/ai/exp/localstore" "github.com/firebase/genkit/go/core" + "github.com/firebase/genkit/go/core/api" "github.com/firebase/genkit/go/internal/registry" ) @@ -141,28 +142,28 @@ func setupHarness(t *testing.T) *harness { ai.ConfigureFormats(reg) pm := &programmableModel{} - ai.DefineModel(reg, "programmableModel", + defineTestModel(reg, "programmableModel", &ai.ModelOptions{Supports: &ai.ModelSupports{Multiturn: true, SystemRole: true, Tools: true}}, pm.generate) ai.DefineGenerateAction(context.Background(), reg) // --- Tools --- - testTool := ai.DefineTool(reg, "testTool", "A simple test tool", + testTool := defineTestTool(reg, "testTool", "A simple test tool", func(tc *ai.ToolContext, _ struct{}) (string, error) { return "tool called", nil }) // interruptTool always pauses the turn, returning the tool request to the // client for external resolution (resume.respond). - interruptTool := ai.DefineTool(reg, "interruptTool", "An interrupt tool", + interruptTool := defineTestTool(reg, "interruptTool", "An interrupt tool", func(tc *ai.ToolContext, _ interruptIn) (interruptOut, error) { return interruptOut{}, tc.Interrupt(&ai.InterruptOptions{}) }) // restartTool interrupts on first call and succeeds when restarted with // resumed metadata (resume.restart). - restartTool := ai.DefineTool(reg, "restartTool", "A tool that requires confirmation before executing", + restartTool := defineTestTool(reg, "restartTool", "A tool that requires confirmation before executing", func(tc *ai.ToolContext, in restartIn) (restartOut, error) { if tc.Resumed == nil { return restartOut{}, tc.Interrupt(&ai.InterruptOptions{ @@ -1127,3 +1128,17 @@ func toFloat(v any) float64 { return 0 } } + +// defineTestModel and defineTestTool register actions in one call, mirroring +// the removed registry-taking ai.Define* helpers. +func defineTestModel(r api.Registry, name string, opts *ai.ModelOptions, fn ai.ModelFunc) ai.Model { + m := ai.NewModel(name, opts, fn) + m.Register(r) + return m +} + +func defineTestTool[In, Out any](r api.Registry, name, description string, fn ai.ToolFunc[In, Out], opts ...ai.ToolOption) *ai.ToolDef[In, Out] { + t := ai.NewTool(name, description, fn, opts...) + t.Register(r) + return t +} diff --git a/go/ai/exp/define_test.go b/go/ai/exp/define_test.go new file mode 100644 index 0000000000..3fb2e253f1 --- /dev/null +++ b/go/ai/exp/define_test.go @@ -0,0 +1,49 @@ +// Copyright 2026 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 exp + +import ( + "github.com/firebase/genkit/go/ai" + "github.com/firebase/genkit/go/core/api" +) + +// Test-local define helpers: New* + Register in one call, mirroring the +// removed registry-taking ai.Define* helpers so tests stay concise. + +func defineTestModel(r api.Registry, name string, opts *ai.ModelOptions, fn ai.ModelFunc) ai.Model { + m := ai.NewModel(name, opts, fn) + m.Register(r) + return m +} + +func defineTestTool[In, Out any](r api.Registry, name, description string, fn ai.ToolFunc[In, Out], opts ...ai.ToolOption) *ai.ToolDef[In, Out] { + t := ai.NewTool(name, description, fn, opts...) + t.Register(r) + return t +} + +func defineTestExpTool[In, Out any](r api.Registry, name, description string, fn ToolFunc[In, Out], opts ...ai.ToolOption) *Tool[In, Out] { + t := NewTool(name, description, fn, opts...) + t.Register(r) + return t +} + +func defineTestInterruptibleTool[In, Out, Res any](r api.Registry, name, description string, fn InterruptibleToolFunc[In, Out, Res], opts ...ai.ToolOption) *InterruptibleTool[In, Out, Res] { + t := NewInterruptibleTool(name, description, fn, opts...) + t.Register(r) + return t +} diff --git a/go/ai/exp/tools.go b/go/ai/exp/tools.go index 8f17d3021f..ae633f736e 100644 --- a/go/ai/exp/tools.go +++ b/go/ai/exp/tools.go @@ -29,7 +29,7 @@ import ( "github.com/firebase/genkit/go/internal/base" ) -// ToolFunc is the function signature for tools created with [DefineTool] and [NewTool]. +// ToolFunc is the function signature for tools created with [NewTool]. type ToolFunc[In, Out any] = func(ctx context.Context, input In) (Out, error) // InterruptibleToolFunc is the function signature for tools created with @@ -151,20 +151,6 @@ func (t *InterruptibleTool[In, Out, Resume]) Respond(part *ai.Part, output Out) return tool.Respond(part, output) } -// DefineTool creates a new tool with a simple function signature and registers it. -// The function receives a plain [context.Context] instead of [ai.ToolContext]. -// Use [tool.AttachParts] inside the function to return additional content parts. -func DefineTool[In, Out any]( - r api.Registry, - name, description string, - fn ToolFunc[In, Out], - opts ...ai.ToolOption, -) *Tool[In, Out] { - t := NewTool(name, description, fn, opts...) - t.Register(r) - return t -} - // NewTool creates a new unregistered tool with a simple function signature. // Use [tool.AttachParts] inside the function to return additional content parts. func NewTool[In, Out any]( @@ -177,21 +163,6 @@ func NewTool[In, Out any]( return &Tool[In, Out]{inner: inner} } -// DefineInterruptibleTool creates a new interruptible tool and registers it. -// The resumed parameter is non-nil when the tool is being resumed after an -// interrupt. Use [tool.Interrupt] inside the function to interrupt execution -// and send data to the caller. -func DefineInterruptibleTool[In, Out, Res any]( - r api.Registry, - name, description string, - fn InterruptibleToolFunc[In, Out, Res], - opts ...ai.ToolOption, -) *InterruptibleTool[In, Out, Res] { - t := NewInterruptibleTool(name, description, fn, opts...) - t.Register(r) - return t -} - // NewInterruptibleTool creates a new unregistered interruptible tool. func NewInterruptibleTool[In, Out, Res any]( name, description string, diff --git a/go/ai/exp/tools_test.go b/go/ai/exp/tools_test.go index f33b5e6447..7c4ff9df96 100644 --- a/go/ai/exp/tools_test.go +++ b/go/ai/exp/tools_test.go @@ -42,7 +42,7 @@ func newToolTestRegistry(t *testing.T) *registry.Registry { // reqs (typically tool requests), and once a tool response is in history it // returns the final text "done". This drives a single tool round per Generate. func defineToolThenFinishModel(reg *registry.Registry, reqs ...*ai.Part) { - ai.DefineModel(reg, "test/model", + defineTestModel(reg, "test/model", &ai.ModelOptions{Supports: &ai.ModelSupports{Multiturn: true, Tools: true}}, func(ctx context.Context, req *ai.ModelRequest, cb ai.ModelStreamCallback) (*ai.ModelResponse, error) { for _, m := range req.Messages { @@ -76,7 +76,7 @@ func TestDefineTool_PlainContext(t *testing.T) { })) var gotCity string - weather := DefineTool(reg, "getWeather", "fetches the weather", + weather := defineTestExpTool(reg, "getWeather", "fetches the weather", func(ctx context.Context, in weatherIn) (string, error) { gotCity = in.City return "Sunny", nil @@ -101,7 +101,7 @@ func TestDefineTool_PlainContext(t *testing.T) { // multipart response without changing the function signature. func TestTool_AttachParts(t *testing.T) { reg := newToolTestRegistry(t) - shot := DefineTool(reg, "screenshot", "takes a screenshot", + shot := defineTestExpTool(reg, "screenshot", "takes a screenshot", func(ctx context.Context, _ struct{}) (string, error) { tool.AttachParts(ctx, ai.NewMediaPart("image/png", "pngbytes")) return "captured", nil @@ -174,7 +174,7 @@ func TestTool_OutputSchemaMatchesClassic(t *testing.T) { // when no streaming callback is wired (here, a direct RunRaw). func TestTool_SendPartialNoOpWithoutStreaming(t *testing.T) { reg := newToolTestRegistry(t) - tl := DefineTool(reg, "noop", "streams when it can", + tl := defineTestExpTool(reg, "noop", "streams when it can", func(ctx context.Context, _ struct{}) (string, error) { tool.SendPartial(ctx, map[string]any{"progress": 50}) return "ok", nil @@ -211,7 +211,7 @@ func TestDefineInterruptibleTool_TypedRoundTrip(t *testing.T) { })) var gotResume *confirmation - transfer := DefineInterruptibleTool(reg, "transfer", "transfers money", + transfer := defineTestInterruptibleTool(reg, "transfer", "transfers money", func(ctx context.Context, in transferIn, res *confirmation) (string, error) { if res == nil { return "", tool.Interrupt(transferInterrupt{Reason: "large_amount", Amount: in.Amount}) @@ -302,7 +302,7 @@ func TestResume_NonObjectData_ReturnsClearError(t *testing.T) { // the tool runs. func TestInterrupt_NonObjectData_ReturnsClearError(t *testing.T) { reg := newToolTestRegistry(t) - tl := DefineInterruptibleTool(reg, "bad", "interrupts with a scalar", + tl := defineTestInterruptibleTool(reg, "bad", "interrupts with a scalar", func(ctx context.Context, _ struct{}, _ *struct{}) (string, error) { return "", tool.Interrupt("not an object") }) @@ -323,7 +323,7 @@ func TestSendPartial_StreamsPartialToolResponse(t *testing.T) { reg := newToolTestRegistry(t) defineToolThenFinishModel(reg, ai.NewToolRequestPart(&ai.ToolRequest{Name: "progressTool", Input: map[string]any{}})) - DefineTool(reg, "progressTool", "streams progress", + defineTestExpTool(reg, "progressTool", "streams progress", func(ctx context.Context, _ struct{}) (string, error) { tool.SendPartial(ctx, map[string]any{"progress": 50}) return "complete", nil @@ -381,8 +381,8 @@ func TestConcurrentStreamingTools_NoDataRace(t *testing.T) { } return "ok", nil } - DefineTool(reg, "toolA", "streams", streamer) - DefineTool(reg, "toolB", "streams", streamer) + defineTestExpTool(reg, "toolA", "streams", streamer) + defineTestExpTool(reg, "toolB", "streams", streamer) for _, err := range ai.GenerateStream(context.Background(), reg, ai.WithModelName("test/model"), diff --git a/go/ai/formatter_test.go b/go/ai/formatter_test.go index e52162cfd4..bd8d197a6e 100644 --- a/go/ai/formatter_test.go +++ b/go/ai/formatter_test.go @@ -864,7 +864,7 @@ func TestConstrainedGenerateWithFormats(t *testing.T) { }, } - formatModel := DefineModel(r, "test/format", &modelOpts, func(ctx context.Context, gr *ModelRequest, msc ModelStreamCallback) (*ModelResponse, error) { + formatModel := defineModel(r, "test/format", &modelOpts, func(ctx context.Context, gr *ModelRequest, msc ModelStreamCallback) (*ModelResponse, error) { return &ModelResponse{ Request: gr, Message: NewModelTextMessage(JSONmd), diff --git a/go/ai/generate.go b/go/ai/generate.go index f3c5def75b..dfdc0ccbf7 100644 --- a/go/ai/generate.go +++ b/go/ai/generate.go @@ -234,17 +234,7 @@ func NewModel(name string, opts *ModelOptions, fn ModelFunc) Model { }) } -// DefineModel creates a new [Model] and registers it. -// -// Deprecated: Use [NewTypedModel] and register the result with -// [Model.Register]. -func DefineModel(r api.Registry, name string, opts *ModelOptions, fn ModelFunc) Model { - m := NewModel(name, opts, fn) - m.Register(r) - return m -} - -// LookupModel looks up a [Model] registered by [DefineModel]. +// LookupModel looks up a registered [Model] by name. // It will try to resolve the model dynamically if the model is not found. // It returns nil if the model was not resolved. func LookupModel(r api.Registry, name string) Model { diff --git a/go/ai/generate_test.go b/go/ai/generate_test.go index ecea849456..bada1cbf03 100644 --- a/go/ai/generate_test.go +++ b/go/ai/generate_test.go @@ -61,7 +61,7 @@ var ( Stage: ModelStageDeprecated, } - echoModel = DefineModel(r, "test/"+modelName, &metadata, func(ctx context.Context, gr *ModelRequest, msc ModelStreamCallback) (*ModelResponse, error) { + echoModel = defineModel(r, "test/"+modelName, &metadata, func(ctx context.Context, gr *ModelRequest, msc ModelStreamCallback) (*ModelResponse, error) { if msc != nil { msc(ctx, &ModelResponseChunk{ Content: []*Part{NewTextPart("stream!")}, @@ -81,7 +81,7 @@ var ( ) // with tools -var gablorkenTool = DefineTool(r, "gablorken", "use when need to calculate a gablorken", +var gablorkenTool = defineTool(r, "gablorken", "use when need to calculate a gablorken", func(ctx *ToolContext, input struct { Value float64 Over float64 @@ -96,7 +96,7 @@ func TestStreamingChunksHaveRoleAndIndex(t *testing.T) { ctx := context.Background() - convertTempTool := DefineTool(r, "convertTemp", "converts temperature", + convertTempTool := defineTool(r, "convertTemp", "converts temperature", func(ctx *ToolContext, input struct { From string To string @@ -110,7 +110,7 @@ func TestStreamingChunksHaveRoleAndIndex(t *testing.T) { }, ) - toolModel := DefineModel(r, "test/toolModel", &metadata, func(ctx context.Context, gr *ModelRequest, msc ModelStreamCallback) (*ModelResponse, error) { + toolModel := defineModel(r, "test/toolModel", &metadata, func(ctx context.Context, gr *ModelRequest, msc ModelStreamCallback) (*ModelResponse, error) { hasToolResponse := false for _, msg := range gr.Messages { if msg.Role == RoleTool { @@ -364,7 +364,7 @@ func TestGenerate(t *testing.T) { JSON := "{\"subject\": \"bananas\", \"location\": \"tropics\"}" JSONmd := "```json" + JSON + "```" - bananaModel := DefineModel(r, "test/banana", &metadata, func(ctx context.Context, gr *ModelRequest, msc ModelStreamCallback) (*ModelResponse, error) { + bananaModel := defineModel(r, "test/banana", &metadata, func(ctx context.Context, gr *ModelRequest, msc ModelStreamCallback) (*ModelResponse, error) { if msc != nil { msc(ctx, &ModelResponseChunk{ Content: []*Part{NewTextPart("stream!")}, @@ -478,7 +478,7 @@ func TestGenerate(t *testing.T) { }) t.Run("handles tool interrupts", func(t *testing.T) { - interruptTool := DefineTool(r, "interruptor", "always interrupts", + interruptTool := defineTool(r, "interruptor", "always interrupts", func(ctx *ToolContext, input any) (any, error) { return nil, ctx.Interrupt(&InterruptOptions{ Metadata: map[string]any{ @@ -494,7 +494,7 @@ func TestGenerate(t *testing.T) { Tools: true, }, } - interruptModel := DefineModel(r, "test/interrupt", info, + interruptModel := defineModel(r, "test/interrupt", info, func(ctx context.Context, gr *ModelRequest, msc ModelStreamCallback) (*ModelResponse, error) { return &ModelResponse{ Request: gr, @@ -553,7 +553,7 @@ func TestGenerate(t *testing.T) { Tools: true, }, } - parallelModel := DefineModel(r, "test/parallel", info, + parallelModel := defineModel(r, "test/parallel", info, func(ctx context.Context, gr *ModelRequest, msc ModelStreamCallback) (*ModelResponse, error) { roundCount++ if roundCount == 1 { @@ -618,7 +618,7 @@ func TestGenerate(t *testing.T) { Tools: true, }, } - multiRoundModel := DefineModel(r, "test/multiround", info, + multiRoundModel := defineModel(r, "test/multiround", info, func(ctx context.Context, gr *ModelRequest, msc ModelStreamCallback) (*ModelResponse, error) { roundCount++ if roundCount == 1 { @@ -686,7 +686,7 @@ func TestGenerate(t *testing.T) { Tools: true, }, } - infiniteModel := DefineModel(r, "test/infinite", info, + infiniteModel := defineModel(r, "test/infinite", info, func(ctx context.Context, gr *ModelRequest, msc ModelStreamCallback) (*ModelResponse, error) { return &ModelResponse{ Request: gr, @@ -770,7 +770,7 @@ func TestGenerate(t *testing.T) { Tools: true, }, } - toolCallModel := DefineModel(r, "test/toolcall", info, + toolCallModel := defineModel(r, "test/toolcall", info, func(ctx context.Context, gr *ModelRequest, msc ModelStreamCallback) (*ModelResponse, error) { roundCount++ if roundCount == 1 { @@ -871,7 +871,7 @@ func TestGenerateWithOutputSchemaName(t *testing.T) { ConfigureFormats(r) // Define a model that supports constrained output - model := DefineModel(r, "test/constrained", &ModelOptions{ + model := defineModel(r, "test/constrained", &ModelOptions{ Supports: &ModelSupports{Constrained: ConstrainedSupportAll}, }, func(ctx context.Context, req *ModelRequest, cb ModelStreamCallback) (*ModelResponse, error) { // Mock response @@ -1005,7 +1005,7 @@ type resumableToolInput struct { } func TestToolInterruptsAndResume(t *testing.T) { - conditionalTool := DefineTool(r, "conditional", "tool that may interrupt based on input", + conditionalTool := defineTool(r, "conditional", "tool that may interrupt based on input", func(ctx *ToolContext, input conditionalToolInput) (string, error) { if input.Interrupt { return "", ctx.Interrupt(&InterruptOptions{ @@ -1020,7 +1020,7 @@ func TestToolInterruptsAndResume(t *testing.T) { }, ) - resumableTool := DefineTool(r, "resumable", "tool that can be resumed", + resumableTool := defineTool(r, "resumable", "tool that can be resumed", func(ctx *ToolContext, input resumableToolInput) (string, error) { if ctx.Resumed != nil { resumedData, ok := ctx.Resumed["data"].(string) @@ -1040,7 +1040,7 @@ func TestToolInterruptsAndResume(t *testing.T) { }, } - toolModel := DefineModel(r, "test/toolmodel", info, + toolModel := defineModel(r, "test/toolmodel", info, func(ctx context.Context, mr *ModelRequest, msc ModelStreamCallback) (*ModelResponse, error) { return &ModelResponse{ Request: mr, @@ -1279,14 +1279,14 @@ func TestResourceProcessing(t *testing.T) { r := registry.New() // Create test resources using DefineResource - DefineResource(r, "test-file", &ResourceOptions{ + defineResource(r, "test-file", &ResourceOptions{ URI: "file:///test.txt", Description: "Test file resource", }, func(ctx context.Context, input *ResourceInput) (*ResourceOutput, error) { return &ResourceOutput{Content: []*Part{NewTextPart("FILE CONTENT")}}, nil }) - DefineResource(r, "test-api", &ResourceOptions{ + defineResource(r, "test-api", &ResourceOptions{ URI: "api://data/123", Description: "Test API resource", }, func(ctx context.Context, input *ResourceInput) (*ResourceOutput, error) { @@ -1579,7 +1579,7 @@ func TestMultipartTools(t *testing.T) { t.Run("define multipart tool registers as tool.v2 only", func(t *testing.T) { r := registry.New() - DefineMultipartTool(r, "multipartTest", "a multipart tool", + defineMultipartTool(r, "multipartTest", "a multipart tool", func(ctx *ToolContext, input struct{ Query string }) (*MultipartToolResponse, error) { return &MultipartToolResponse{ Output: "main output", @@ -1607,7 +1607,7 @@ func TestMultipartTools(t *testing.T) { t.Run("regular tool registers as both tool and tool.v2", func(t *testing.T) { r := registry.New() - DefineTool(r, "regularTestTool", "a regular tool", + defineTool(r, "regularTestTool", "a regular tool", func(ctx *ToolContext, input struct{ Value int }) (int, error) { return input.Value * 2, nil }, @@ -1634,7 +1634,7 @@ func TestMultipartTools(t *testing.T) { ConfigureFormats(r) DefineGenerateAction(context.Background(), r) - multipartTool := DefineMultipartTool(r, "imageGenerator", "generates images", + multipartTool := defineMultipartTool(r, "imageGenerator", "generates images", func(ctx *ToolContext, input struct{ Prompt string }) (*MultipartToolResponse, error) { return &MultipartToolResponse{ Output: map[string]any{"description": "generated image"}, @@ -1647,7 +1647,7 @@ func TestMultipartTools(t *testing.T) { ) // Create a model that requests the tool - multipartToolModel := DefineModel(r, "test/multipartToolModel", &metadata, func(ctx context.Context, gr *ModelRequest, msc ModelStreamCallback) (*ModelResponse, error) { + multipartToolModel := defineModel(r, "test/multipartToolModel", &metadata, func(ctx context.Context, gr *ModelRequest, msc ModelStreamCallback) (*ModelResponse, error) { // Check if we already have a tool response for _, msg := range gr.Messages { if msg.Role == RoleTool { @@ -1700,7 +1700,7 @@ func TestMultipartTools(t *testing.T) { t.Run("RunRawMultipart returns MultipartToolResponse for regular tool", func(t *testing.T) { r := registry.New() - tool := DefineTool(r, "multipartWrapperTest", "test multipart wrapper", + tool := defineTool(r, "multipartWrapperTest", "test multipart wrapper", func(ctx *ToolContext, input struct{ Value int }) (int, error) { return input.Value * 3, nil }, @@ -1729,7 +1729,7 @@ func TestMultipartTools(t *testing.T) { t.Run("RunRawMultipart returns full response for multipart tool", func(t *testing.T) { r := registry.New() - tool := DefineMultipartTool(r, "multipartFullTest", "test multipart", + tool := defineMultipartTool(r, "multipartFullTest", "test multipart", func(ctx *ToolContext, input struct{ Query string }) (*MultipartToolResponse, error) { return &MultipartToolResponse{ Output: "result", @@ -1772,7 +1772,7 @@ func TestGenerateStream(t *testing.T) { chunkTexts := []string{"Hello", " ", "World"} chunkIndex := 0 - streamModel := DefineModel(r, "test/streamModel", &ModelOptions{ + streamModel := defineModel(r, "test/streamModel", &ModelOptions{ Supports: &ModelSupports{Multiturn: true}, }, func(ctx context.Context, req *ModelRequest, cb ModelStreamCallback) (*ModelResponse, error) { if cb != nil { @@ -1825,7 +1825,7 @@ func TestGenerateStream(t *testing.T) { }) t.Run("handles no streaming callback gracefully", func(t *testing.T) { - noStreamModel := DefineModel(r, "test/noStreamModel", &ModelOptions{ + noStreamModel := defineModel(r, "test/noStreamModel", &ModelOptions{ Supports: &ModelSupports{Multiturn: true}, }, func(ctx context.Context, req *ModelRequest, cb ModelStreamCallback) (*ModelResponse, error) { return &ModelResponse{ @@ -1865,7 +1865,7 @@ func TestGenerateStream(t *testing.T) { t.Run("propagates generation errors", func(t *testing.T) { expectedErr := errors.New("generation failed") - errorModel := DefineModel(r, "test/errorModel", &ModelOptions{ + errorModel := defineModel(r, "test/errorModel", &ModelOptions{ Supports: &ModelSupports{Multiturn: true}, }, func(ctx context.Context, req *ModelRequest, cb ModelStreamCallback) (*ModelResponse, error) { return nil, expectedErr @@ -1894,7 +1894,7 @@ func TestGenerateStream(t *testing.T) { ctx, cancel := context.WithCancel(context.Background()) defer cancel() - streamModel := DefineModel(r, "test/cancelModel", &ModelOptions{ + streamModel := defineModel(r, "test/cancelModel", &ModelOptions{ Supports: &ModelSupports{Multiturn: true}, }, func(ctx context.Context, req *ModelRequest, cb ModelStreamCallback) (*ModelResponse, error) { if cb != nil { @@ -1940,7 +1940,7 @@ func TestGenerateStream(t *testing.T) { }) t.Run("should not yield after stop", func(t *testing.T) { - streamModel := DefineModel(r, "test/breakStreamModel", &ModelOptions{ + streamModel := defineModel(r, "test/breakStreamModel", &ModelOptions{ Supports: &ModelSupports{ Multiturn: true, }, @@ -1971,7 +1971,7 @@ func TestGenerateDataStream(t *testing.T) { DefineGenerateAction(context.Background(), r) t.Run("yields typed chunks and final output", func(t *testing.T) { - streamModel := DefineModel(r, "test/typedStreamModel", &ModelOptions{ + streamModel := defineModel(r, "test/typedStreamModel", &ModelOptions{ Supports: &ModelSupports{ Multiturn: true, Constrained: ConstrainedSupportAll, @@ -2026,7 +2026,7 @@ func TestGenerateDataStream(t *testing.T) { }) t.Run("final output is correctly typed", func(t *testing.T) { - streamModel := DefineModel(r, "test/finalTypedModel", &ModelOptions{ + streamModel := defineModel(r, "test/finalTypedModel", &ModelOptions{ Supports: &ModelSupports{ Multiturn: true, Constrained: ConstrainedSupportAll, @@ -2068,7 +2068,7 @@ func TestGenerateDataStream(t *testing.T) { t.Run("automatically sets output type", func(t *testing.T) { var capturedRequest *ModelRequest - streamModel := DefineModel(r, "test/autoOutputModel", &ModelOptions{ + streamModel := defineModel(r, "test/autoOutputModel", &ModelOptions{ Supports: &ModelSupports{ Multiturn: true, Constrained: ConstrainedSupportAll, @@ -2099,7 +2099,7 @@ func TestGenerateDataStream(t *testing.T) { }) t.Run("handles tool interrupts", func(t *testing.T) { - interruptTool := DefineTool(r, "streamInterruptor", "always interrupts", + interruptTool := defineTool(r, "streamInterruptor", "always interrupts", func(ctx *ToolContext, input any) (any, error) { return nil, ctx.Interrupt(&InterruptOptions{ Metadata: map[string]any{ @@ -2109,7 +2109,7 @@ func TestGenerateDataStream(t *testing.T) { }, ) - streamModel := DefineModel(r, "test/streamInterruptModel", &ModelOptions{ + streamModel := defineModel(r, "test/streamInterruptModel", &ModelOptions{ Supports: &ModelSupports{ Multiturn: true, Tools: true, @@ -2167,13 +2167,13 @@ func TestGenerateDataStream(t *testing.T) { }) t.Run("handles returnToolRequests", func(t *testing.T) { - greetTool := DefineTool(r, "streamGreeter", "greets", + greetTool := defineTool(r, "streamGreeter", "greets", func(ctx *ToolContext, input any) (any, error) { return "hello", nil }, ) - streamModel := DefineModel(r, "test/streamReturnToolModel", &ModelOptions{ + streamModel := defineModel(r, "test/streamReturnToolModel", &ModelOptions{ Supports: &ModelSupports{ Multiturn: true, Tools: true, @@ -2224,7 +2224,7 @@ func TestGenerateDataStream(t *testing.T) { }) t.Run("propagates chunk parsing errors", func(t *testing.T) { - streamModel := DefineModel(r, "test/parseErrorModel", &ModelOptions{ + streamModel := defineModel(r, "test/parseErrorModel", &ModelOptions{ Supports: &ModelSupports{ Multiturn: true, Constrained: ConstrainedSupportAll, @@ -2258,7 +2258,7 @@ func TestGenerateDataStream(t *testing.T) { }) t.Run("should not yield after stop", func(t *testing.T) { - streamModel := DefineModel(r, "test/breakDataStreamModel", &ModelOptions{ + streamModel := defineModel(r, "test/breakDataStreamModel", &ModelOptions{ Supports: &ModelSupports{ Multiturn: true, Constrained: ConstrainedSupportAll, @@ -2290,7 +2290,7 @@ func TestGenerateDataStream(t *testing.T) { func TestGenerateText(t *testing.T) { r := newTestRegistry(t) - echoModel := DefineModel(r, "test/echoTextModel", nil, func(ctx context.Context, req *ModelRequest, cb ModelStreamCallback) (*ModelResponse, error) { + echoModel := defineModel(r, "test/echoTextModel", nil, func(ctx context.Context, req *ModelRequest, cb ModelStreamCallback) (*ModelResponse, error) { return &ModelResponse{ Request: req, Message: NewModelTextMessage("echo: " + req.Messages[0].Content[0].Text), @@ -2318,7 +2318,7 @@ func TestGenerateData(t *testing.T) { Value int `json:"value"` } - jsonModel := DefineModel(r, "test/jsonDataModel", &ModelOptions{ + jsonModel := defineModel(r, "test/jsonDataModel", &ModelOptions{ Supports: &ModelSupports{ Constrained: ConstrainedSupportAll, }, @@ -2479,7 +2479,7 @@ func TestGenerateWithMarkdownJSON(t *testing.T) { DefineGenerateAction(context.Background(), r) // A model that returns JSON wrapped in markdown - markdownModel := DefineModel(r, "test/markdownJson", &ModelOptions{ + markdownModel := defineModel(r, "test/markdownJson", &ModelOptions{ Supports: &ModelSupports{Constrained: ConstrainedSupportAll}, }, func(ctx context.Context, req *ModelRequest, cb ModelStreamCallback) (*ModelResponse, error) { jsonContent := "{\"name\": \"test\", \"value\": 123}" @@ -2490,7 +2490,7 @@ func TestGenerateWithMarkdownJSON(t *testing.T) { }) // A model that returns JSON wrapped in markdown with loose formatting (spaces) - looseMarkdownModel := DefineModel(r, "test/looseMarkdownJson", &ModelOptions{ + looseMarkdownModel := defineModel(r, "test/looseMarkdownJson", &ModelOptions{ Supports: &ModelSupports{Constrained: ConstrainedSupportAll}, }, func(ctx context.Context, req *ModelRequest, cb ModelStreamCallback) (*ModelResponse, error) { jsonContent := "{\"name\": \"test\", \"value\": 123}" @@ -2553,20 +2553,20 @@ func TestGenerateNoGoroutineLeak(t *testing.T) { done := make(chan struct{}) - slowTool := DefineTool(r, "slow", "slow", + slowTool := defineTool(r, "slow", "slow", func(*ToolContext, any) (any, error) { <-done return nil, nil }, ) - failTool := DefineTool(r, "fail", "fail", + failTool := defineTool(r, "fail", "fail", func(*ToolContext, any) (any, error) { return nil, errors.New("boom") }, ) - testModel := DefineModel(r, "test/testModel", &ModelOptions{ + testModel := defineModel(r, "test/testModel", &ModelOptions{ Supports: &ModelSupports{ Multiturn: true, Tools: true, diff --git a/go/ai/middleware.go b/go/ai/middleware.go index 96c0c8e38c..601790e8b6 100644 --- a/go/ai/middleware.go +++ b/go/ai/middleware.go @@ -152,13 +152,6 @@ func NewMiddleware[M Middleware](description string, prototype M) *MiddlewareDes } } -// DefineMiddleware creates and registers a middleware descriptor in one step. -func DefineMiddleware[M Middleware](r api.Registry, description string, prototype M) *MiddlewareDesc { - d := NewMiddleware(description, prototype) - d.Register(r) - return d -} - // MiddlewareFunc adapts a per-call factory closure to the [Middleware] // interface for ad-hoc inline use, without a registered descriptor or plugin // wiring. The adapted middleware does not appear in the Dev UI. diff --git a/go/ai/middleware_test.go b/go/ai/middleware_test.go index d77e7b9474..5f3235ef44 100644 --- a/go/ai/middleware_test.go +++ b/go/ai/middleware_test.go @@ -567,7 +567,7 @@ func TestWrapToolValidationErrorReturnedToModel(t *testing.T) { handler: modelHandler, }) - DefineTool(r, "validateMe", "A tool that requires a numeric value", + defineTool(r, "validateMe", "A tool that requires a numeric value", func(ctx *ToolContext, input any) (string, error) { m := input.(map[string]any) return fmt.Sprintf("success: %v", m["value"]), nil @@ -695,7 +695,7 @@ func TestMiddlewareHookOrderOnToolRestart(t *testing.T) { Interrupt bool `json:"interrupt"` } - tool := DefineTool(r, "restartable", "interrupts, then runs on resume", + tool := defineTool(r, "restartable", "interrupts, then runs on resume", func(ctx *ToolContext, in restartInput) (string, error) { if in.Interrupt { return "", ctx.Interrupt(&InterruptOptions{}) @@ -706,7 +706,7 @@ func TestMiddlewareHookOrderOnToolRestart(t *testing.T) { // Requests the tool on the first turn, returns a final text response once a // tool response is present in history. - model := DefineModel(r, "test/restartModel", &ModelOptions{ + model := defineModel(r, "test/restartModel", &ModelOptions{ Supports: &ModelSupports{Multiturn: true, Tools: true}, }, func(ctx context.Context, req *ModelRequest, _ ModelStreamCallback) (*ModelResponse, error) { for _, msg := range req.Messages { diff --git a/go/ai/prompt_test.go b/go/ai/prompt_test.go index 7db48122ef..5dfb320047 100644 --- a/go/ai/prompt_test.go +++ b/go/ai/prompt_test.go @@ -36,7 +36,7 @@ type InputOutput struct { } func testTool(reg api.Registry, name string) Tool { - return DefineTool(reg, name, "use when need to execute a test", + return defineTool(reg, name, "use when need to execute a test", func(ctx *ToolContext, input struct { Test string }, @@ -150,7 +150,7 @@ type HelloPromptInput struct { } func definePromptModel(reg api.Registry) Model { - return DefineModel(reg, "test/chat", + return defineModel(reg, "test/chat", &ModelOptions{Supports: &ModelSupports{ Tools: true, Multiturn: true, @@ -692,7 +692,7 @@ func TestOptionsPatternExecute(t *testing.T) { ConfigureFormats(reg) - testModel := DefineModel(reg, "options/test", nil, testGenerate) + testModel := defineModel(reg, "options/test", nil, testGenerate) t.Run("Streaming", func(t *testing.T) { p := DefinePrompt(reg, "TestExecute", WithInputType(InputOutput{}), WithPrompt("TestExecute")) @@ -744,7 +744,7 @@ func TestDefaultsOverride(t *testing.T) { // Set up default formats ConfigureFormats(reg) - testModel := DefineModel(reg, "defineoptions/test", nil, testGenerate) + testModel := defineModel(reg, "defineoptions/test", nil, testGenerate) model := definePromptModel(reg) tests := []struct { @@ -1575,7 +1575,7 @@ Generate a recipe for {{food}}. reg := registry.New() ConfigureFormats(reg) - DefineModel(reg, "test-model", &ModelOptions{ + defineModel(reg, "test-model", &ModelOptions{ Supports: &ModelSupports{Constrained: ConstrainedSupportAll}, }, func(ctx context.Context, req *ModelRequest, cb ModelStreamCallback) (*ModelResponse, error) { // Mock response that matches the expected schema structure @@ -1685,7 +1685,7 @@ Generate a recipe. reg := registry.New() ConfigureFormats(reg) - DefineModel(reg, "test-model", &ModelOptions{ + defineModel(reg, "test-model", &ModelOptions{ Supports: &ModelSupports{Constrained: ConstrainedSupportAll}, }, func(ctx context.Context, req *ModelRequest, cb ModelStreamCallback) (*ModelResponse, error) { return &ModelResponse{ @@ -1713,7 +1713,7 @@ func TestWithOutputSchemaName_DefinePrompt(t *testing.T) { reg := registry.New() ConfigureFormats(reg) - DefineModel(reg, "test-model", &ModelOptions{ + defineModel(reg, "test-model", &ModelOptions{ Supports: &ModelSupports{Constrained: ConstrainedSupportAll}, }, func(ctx context.Context, req *ModelRequest, cb ModelStreamCallback) (*ModelResponse, error) { return &ModelResponse{ @@ -1751,7 +1751,7 @@ func TestWithOutputSchemaName_DefinePrompt_Missing(t *testing.T) { reg := registry.New() ConfigureFormats(reg) - DefineModel(reg, "test-model", &ModelOptions{ + defineModel(reg, "test-model", &ModelOptions{ Supports: &ModelSupports{Constrained: ConstrainedSupportAll}, }, func(ctx context.Context, req *ModelRequest, cb ModelStreamCallback) (*ModelResponse, error) { return &ModelResponse{ @@ -1793,7 +1793,7 @@ func TestDataPromptExecute(t *testing.T) { t.Run("typed input and output", func(t *testing.T) { var capturedInput any - testModel := DefineModel(r, "test/dataPromptModel", &ModelOptions{ + testModel := defineModel(r, "test/dataPromptModel", &ModelOptions{ Supports: &ModelSupports{ Multiturn: true, Constrained: ConstrainedSupportAll, @@ -1835,7 +1835,7 @@ func TestDataPromptExecute(t *testing.T) { }) t.Run("string output type", func(t *testing.T) { - testModel := DefineModel(r, "test/stringDataPromptModel", &ModelOptions{ + testModel := defineModel(r, "test/stringDataPromptModel", &ModelOptions{ Supports: &ModelSupports{Multiturn: true}, }, func(ctx context.Context, req *ModelRequest, cb ModelStreamCallback) (*ModelResponse, error) { return &ModelResponse{ @@ -1874,7 +1874,7 @@ func TestDataPromptExecute(t *testing.T) { t.Run("additional options passed through", func(t *testing.T) { var capturedConfig any - testModel := DefineModel(r, "test/optionsDataPromptModel", &ModelOptions{ + testModel := defineModel(r, "test/optionsDataPromptModel", &ModelOptions{ Supports: &ModelSupports{ Multiturn: true, Constrained: ConstrainedSupportAll, @@ -1912,7 +1912,7 @@ func TestDataPromptExecute(t *testing.T) { }) t.Run("returns error for invalid output parsing", func(t *testing.T) { - testModel := DefineModel(r, "test/parseFailDataPromptModel", &ModelOptions{ + testModel := defineModel(r, "test/parseFailDataPromptModel", &ModelOptions{ Supports: &ModelSupports{ Multiturn: true, Constrained: ConstrainedSupportAll, @@ -1951,7 +1951,7 @@ func TestDataPromptExecuteStream(t *testing.T) { } t.Run("typed streaming with struct output", func(t *testing.T) { - testModel := DefineModel(r, "test/streamDataPromptModel", &ModelOptions{ + testModel := defineModel(r, "test/streamDataPromptModel", &ModelOptions{ Supports: &ModelSupports{ Multiturn: true, Constrained: ConstrainedSupportAll, @@ -2008,7 +2008,7 @@ func TestDataPromptExecuteStream(t *testing.T) { }) t.Run("string output streaming", func(t *testing.T) { - testModel := DefineModel(r, "test/stringStreamDataPromptModel", &ModelOptions{ + testModel := defineModel(r, "test/stringStreamDataPromptModel", &ModelOptions{ Supports: &ModelSupports{Multiturn: true}, }, func(ctx context.Context, req *ModelRequest, cb ModelStreamCallback) (*ModelResponse, error) { if cb != nil { @@ -2078,7 +2078,7 @@ func TestDataPromptExecuteStream(t *testing.T) { t.Run("handles options passed at execute time", func(t *testing.T) { var capturedConfig any - testModel := DefineModel(r, "test/optionsStreamModel", &ModelOptions{ + testModel := defineModel(r, "test/optionsStreamModel", &ModelOptions{ Supports: &ModelSupports{ Multiturn: true, Constrained: ConstrainedSupportAll, @@ -2121,7 +2121,7 @@ func TestDataPromptExecuteStream(t *testing.T) { t.Run("propagates errors", func(t *testing.T) { expectedErr := errors.New("stream failed") - testModel := DefineModel(r, "test/errorStreamDataPromptModel", &ModelOptions{ + testModel := defineModel(r, "test/errorStreamDataPromptModel", &ModelOptions{ Supports: &ModelSupports{Multiturn: true}, }, func(ctx context.Context, req *ModelRequest, cb ModelStreamCallback) (*ModelResponse, error) { return nil, expectedErr @@ -2149,7 +2149,7 @@ func TestDataPromptExecuteStream(t *testing.T) { }) t.Run("should not yield after stop", func(t *testing.T) { - testModel := DefineModel(r, "test/breakDataPromptStreamModel", &ModelOptions{ + testModel := defineModel(r, "test/breakDataPromptStreamModel", &ModelOptions{ Supports: &ModelSupports{ Multiturn: true, Constrained: ConstrainedSupportAll, @@ -2191,7 +2191,7 @@ func TestPromptExecuteStream(t *testing.T) { t.Run("yields chunks then final response", func(t *testing.T) { chunkTexts := []string{"A", "B", "C"} - testModel := DefineModel(r, "test/promptStreamModel", &ModelOptions{ + testModel := defineModel(r, "test/promptStreamModel", &ModelOptions{ Supports: &ModelSupports{Multiturn: true}, }, func(ctx context.Context, req *ModelRequest, cb ModelStreamCallback) (*ModelResponse, error) { if cb != nil { @@ -2262,7 +2262,7 @@ func TestPromptExecuteStream(t *testing.T) { t.Run("handles execution options", func(t *testing.T) { var capturedConfig any - testModel := DefineModel(r, "test/optionsPromptExecModel", &ModelOptions{ + testModel := defineModel(r, "test/optionsPromptExecModel", &ModelOptions{ Supports: &ModelSupports{Multiturn: true}, }, func(ctx context.Context, req *ModelRequest, cb ModelStreamCallback) (*ModelResponse, error) { capturedConfig = req.Config @@ -2297,7 +2297,7 @@ func TestPromptExecuteStream(t *testing.T) { }) t.Run("should not yield after stop", func(t *testing.T) { - testModel := DefineModel(r, "test/breakPromptStreamModel", &ModelOptions{ + testModel := defineModel(r, "test/breakPromptStreamModel", &ModelOptions{ Supports: &ModelSupports{ Multiturn: true, }, @@ -2339,7 +2339,7 @@ func TestSessionStateInjection(t *testing.T) { t.Run("state accessible in prompt template", func(t *testing.T) { var capturedPrompt string - testModel := DefineModel(r, "test/sessionStateModel", &ModelOptions{ + testModel := defineModel(r, "test/sessionStateModel", &ModelOptions{ Supports: &ModelSupports{Multiturn: true}, }, func(ctx context.Context, req *ModelRequest, cb ModelStreamCallback) (*ModelResponse, error) { capturedPrompt = req.Messages[0].Text() @@ -2378,7 +2378,7 @@ func TestSessionStateInjection(t *testing.T) { t.Run("prompt works without state in context", func(t *testing.T) { var capturedPrompt string - testModel := DefineModel(r, "test/noSessionModel", &ModelOptions{ + testModel := defineModel(r, "test/noSessionModel", &ModelOptions{ Supports: &ModelSupports{Multiturn: true}, }, func(ctx context.Context, req *ModelRequest, cb ModelStreamCallback) (*ModelResponse, error) { capturedPrompt = req.Messages[0].Text() @@ -2413,7 +2413,7 @@ func TestSessionStateInjection(t *testing.T) { t.Run("state and input variables can be used together", func(t *testing.T) { var capturedPrompt string - testModel := DefineModel(r, "test/mixedModel", &ModelOptions{ + testModel := defineModel(r, "test/mixedModel", &ModelOptions{ Supports: &ModelSupports{Multiturn: true}, }, func(ctx context.Context, req *ModelRequest, cb ModelStreamCallback) (*ModelResponse, error) { capturedPrompt = req.Messages[0].Text() diff --git a/go/ai/resource.go b/go/ai/resource.go index 8ca02332c7..8744005303 100644 --- a/go/ai/resource.go +++ b/go/ai/resource.go @@ -126,14 +126,6 @@ type Resource interface { Register(r api.Registry) } -// DefineResource creates a resource and registers it with the given Registry. -func DefineResource(r api.Registry, name string, opts *ResourceOptions, fn ResourceFunc) Resource { - metadata := resourceMetadata(name, opts) - a := core.NewActionOf(api.ActionTypeResource, name, &core.ActionOptions{Metadata: metadata}, fn) - a.Register(r) - return &resource{Action: *a} -} - // NewResource creates a resource but does not register it in the registry. // It can be registered later via the Register method. func NewResource(name string, opts *ResourceOptions, fn ResourceFunc) Resource { diff --git a/go/ai/resource_test.go b/go/ai/resource_test.go index 30f15b807c..3e92147872 100644 --- a/go/ai/resource_test.go +++ b/go/ai/resource_test.go @@ -26,7 +26,7 @@ func TestStaticResource(t *testing.T) { g := registry.New() // Define static resource - DefineResource(g, "test-doc", &ResourceOptions{ + defineResource(g, "test-doc", &ResourceOptions{ URI: "file:///test.txt", }, func(ctx context.Context, input *ResourceInput) (*ResourceOutput, error) { return &ResourceOutput{ @@ -107,7 +107,7 @@ func TestResourceInGeneration(t *testing.T) { ConfigureFormats(r) // Define mock model - DefineModel(r, "test", nil, func(ctx context.Context, req *ModelRequest, cb ModelStreamCallback) (*ModelResponse, error) { + defineModel(r, "test", nil, func(ctx context.Context, req *ModelRequest, cb ModelStreamCallback) (*ModelResponse, error) { // Extract resource parts from the prompt var responseText strings.Builder for _, msg := range req.Messages { @@ -128,7 +128,7 @@ func TestResourceInGeneration(t *testing.T) { }) // Define resource - DefineResource(r, "policy", &ResourceOptions{ + defineResource(r, "policy", &ResourceOptions{ URI: "file:///policy.txt", }, func(ctx context.Context, input *ResourceInput) (*ResourceOutput, error) { return &ResourceOutput{ @@ -163,7 +163,7 @@ func TestDynamicResourceInGeneration(t *testing.T) { ConfigureFormats(r) // Define mock model - DefineModel(r, "test", nil, func(ctx context.Context, req *ModelRequest, cb ModelStreamCallback) (*ModelResponse, error) { + defineModel(r, "test", nil, func(ctx context.Context, req *ModelRequest, cb ModelStreamCallback) (*ModelResponse, error) { var responseText strings.Builder for _, msg := range req.Messages { for _, part := range msg.Content { @@ -219,7 +219,7 @@ func TestMultipleDynamicResourcesInGeneration(t *testing.T) { ConfigureFormats(r) // Define mock model - DefineModel(r, "test", nil, func(ctx context.Context, req *ModelRequest, cb ModelStreamCallback) (*ModelResponse, error) { + defineModel(r, "test", nil, func(ctx context.Context, req *ModelRequest, cb ModelStreamCallback) (*ModelResponse, error) { return &ModelResponse{ Request: req, Message: &Message{ @@ -276,7 +276,7 @@ func contains(s, substr string) bool { func TestLookupResource(t *testing.T) { t.Run("finds registered resource", func(t *testing.T) { r := registry.New() - DefineResource(r, "test/lookup", &ResourceOptions{ + defineResource(r, "test/lookup", &ResourceOptions{ URI: "lookup://test", }, func(ctx context.Context, input *ResourceInput) (*ResourceOutput, error) { return &ResourceOutput{ @@ -304,7 +304,7 @@ func TestLookupResource(t *testing.T) { t.Run("resource can be executed after lookup", func(t *testing.T) { r := registry.New() - DefineResource(r, "test/executable", &ResourceOptions{ + defineResource(r, "test/executable", &ResourceOptions{ URI: "exec://test", }, func(ctx context.Context, input *ResourceInput) (*ResourceOutput, error) { return &ResourceOutput{ @@ -328,7 +328,7 @@ func TestLookupResource(t *testing.T) { t.Run("resource matches and extracts variables after lookup", func(t *testing.T) { r := registry.New() - DefineResource(r, "test/template", &ResourceOptions{ + defineResource(r, "test/template", &ResourceOptions{ Template: "template://item/{id}", }, func(ctx context.Context, input *ResourceInput) (*ResourceOutput, error) { return &ResourceOutput{ diff --git a/go/ai/testutil_test.go b/go/ai/testutil_test.go index 6c606a28a8..b9135673d4 100644 --- a/go/ai/testutil_test.go +++ b/go/ai/testutil_test.go @@ -70,7 +70,7 @@ func defineFakeModel(t *testing.T, r api.Registry, cfg fakeModelConfig) Model { }, nil } } - return DefineModel(r, cfg.name, &ModelOptions{Supports: cfg.supports}, cfg.handler) + return defineModel(r, cfg.name, &ModelOptions{Supports: cfg.supports}, cfg.handler) } // echoModelHandler creates a handler that echoes back information about the request. @@ -275,7 +275,7 @@ func assertNoPanic(t *testing.T, fn func()) { // defineFakeTool creates a simple tool for testing. func defineFakeTool(t *testing.T, r api.Registry, name, description string) Tool { t.Helper() - return DefineTool(r, name, description, + return defineTool(r, name, description, func(ctx *ToolContext, input struct { Value string `json:"value"` }) (string, error) { @@ -286,7 +286,7 @@ func defineFakeTool(t *testing.T, r api.Registry, name, description string) Tool // defineFakeEmbedder creates a simple embedder for testing. func defineFakeEmbedder(t *testing.T, r api.Registry, name string) Embedder { t.Helper() - return DefineEmbedder(r, name, nil, func(ctx context.Context, req *EmbedRequest) (*EmbedResponse, error) { + return defineEmbedder(r, name, nil, func(ctx context.Context, req *EmbedRequest) (*EmbedResponse, error) { embeddings := make([]*Embedding, len(req.Input)) for i := range req.Input { embeddings[i] = &Embedding{ diff --git a/go/ai/tools.go b/go/ai/tools.go index cda0f3db93..16e890176f 100644 --- a/go/ai/tools.go +++ b/go/ai/tools.go @@ -299,53 +299,6 @@ func applyStrictMetadata(metadata map[string]any, strict *bool) { toolMeta[toolStrictKey] = *strict } -// DefineTool creates a new [ToolDef] and registers it. -// Use [WithInputSchema] to provide a custom JSON schema instead of inferring from the type parameter. -func DefineTool[In, Out any]( - r api.Registry, - name, description string, - fn ToolFunc[In, Out], - opts ...ToolOption, -) *ToolDef[In, Out] { - toolOpts := &toolOptions{} - for _, opt := range opts { - if err := opt.applyTool(toolOpts); err != nil { - panic(fmt.Errorf("ai.DefineTool %q: %w", name, err)) - } - } - - // If the user provided a custom input schema, enforce that In is 'any' - if toolOpts.InputSchema != nil { - typ := reflect.TypeFor[*In]() - if typ != nil && typ.Elem().Kind() != reflect.Interface { - panic(fmt.Errorf("ai.DefineTool %q: WithInputSchema requires In to be of type 'any', but got %v", name, typ.Elem())) - } - } - - metadata, wrappedFn := wrapToolFunc(name, description, fn) - applyStrictMetadata(metadata, toolOpts.StrictSchema) - action := core.NewActionOf(api.ActionTypeToolV2, name, &core.ActionOptions{Metadata: metadata, InputSchema: toolOpts.InputSchema}, wrappedFn) - action.Register(r) - - // Also register under the "tool" action type for backward compatibility. - provider, id := api.ParseName(name) - r.RegisterAction(api.NewKey(api.ActionTypeTool, provider, id), action) - - return &ToolDef[In, Out]{action: action, multipart: false, registry: r} -} - -// DefineToolWithInputSchema creates a new [ToolDef] with a custom input schema and registers it. -// -// Deprecated: Use [DefineTool] with [WithInputSchema] instead. -func DefineToolWithInputSchema[Out any]( - r api.Registry, - name, description string, - inputSchema map[string]any, - fn ToolFunc[any, Out], -) *ToolDef[any, Out] { - return DefineTool(r, name, description, fn, WithInputSchema(inputSchema)) -} - // NewTool creates a new [ToolDef]. It can be passed directly to [Generate]. // Use [WithInputSchema] to provide a custom JSON schema instead of inferring from the type parameter. func NewTool[In, Out any](name, description string, fn ToolFunc[In, Out], opts ...ToolOption) *ToolDef[In, Out] { @@ -379,29 +332,6 @@ func NewToolWithInputSchema[Out any](name, description string, inputSchema map[s return NewTool(name, description, fn, WithInputSchema(inputSchema)) } -// DefineMultipartTool creates a new multipart [ToolDef] and registers it. -// Multipart tools can return both output data and additional content parts (like media). -// Use [WithInputSchema] to provide a custom JSON schema instead of inferring from the type parameter. -func DefineMultipartTool[In any]( - r api.Registry, - name, description string, - fn MultipartToolFunc[In], - opts ...ToolOption, -) *ToolDef[In, *MultipartToolResponse] { - toolOpts := &toolOptions{} - for _, opt := range opts { - if err := opt.applyTool(toolOpts); err != nil { - panic(fmt.Errorf("ai.DefineMultipartTool %q: %w", name, err)) - } - } - - metadata, wrappedFn := wrapMultipartToolFunc(name, description, fn) - applyStrictMetadata(metadata, toolOpts.StrictSchema) - action := core.NewActionOf(api.ActionTypeToolV2, name, &core.ActionOptions{Metadata: metadata, InputSchema: toolOpts.InputSchema}, wrappedFn) - action.Register(r) - return &ToolDef[In, *MultipartToolResponse]{action: action, multipart: true, registry: r} -} - // NewMultipartTool creates a new multipart [ToolDef]. It can be passed directly to [Generate]. // Multipart tools can return both output data and additional content parts (like media). // Use [WithInputSchema] to provide a custom JSON schema instead of inferring from the type parameter. diff --git a/go/ai/tools_test.go b/go/ai/tools_test.go index be5f491065..11f95b9290 100644 --- a/go/ai/tools_test.go +++ b/go/ai/tools_test.go @@ -137,7 +137,7 @@ func (e *wrappedInterruptError) Unwrap() error { func TestDefineTool(t *testing.T) { t.Run("creates and registers tool", func(t *testing.T) { r := newTestRegistry(t) - tl := DefineTool(r, "provider/addNumbers", "Adds two numbers", func(ctx *ToolContext, input struct { + tl := defineTool(r, "provider/addNumbers", "Adds two numbers", func(ctx *ToolContext, input struct { A int `json:"a"` B int `json:"b"` }) (int, error) { @@ -159,7 +159,7 @@ func TestDefineTool(t *testing.T) { t.Run("tool can be looked up after registration", func(t *testing.T) { r := newTestRegistry(t) - DefineTool(r, "provider/multiply", "Multiplies", func(ctx *ToolContext, input struct { + defineTool(r, "provider/multiply", "Multiplies", func(ctx *ToolContext, input struct { X int `json:"x"` Y int `json:"y"` }) (int, error) { @@ -174,7 +174,7 @@ func TestDefineTool(t *testing.T) { t.Run("tool executes correctly", func(t *testing.T) { r := newTestRegistry(t) - tl := DefineTool(r, "provider/concat", "Concatenates strings", func(ctx *ToolContext, input struct { + tl := defineTool(r, "provider/concat", "Concatenates strings", func(ctx *ToolContext, input struct { A string `json:"a"` B string `json:"b"` }) (string, error) { @@ -206,7 +206,7 @@ func TestDefineToolWithInputSchema(t *testing.T) { "required": []any{"query"}, } - tl := DefineToolWithInputSchema(r, "provider/search", "Searches", customSchema, + tl := defineToolWithInputSchema(r, "provider/search", "Searches", customSchema, func(ctx *ToolContext, input any) (string, error) { m := input.(map[string]any) return "results for: " + m["query"].(string), nil @@ -299,7 +299,7 @@ func TestNewToolWithInputSchema(t *testing.T) { func TestDefineMultipartTool(t *testing.T) { t.Run("creates multipart tool", func(t *testing.T) { r := newTestRegistry(t) - tl := DefineMultipartTool(r, "provider/imageGen", "Generates images", + tl := defineMultipartTool(r, "provider/imageGen", "Generates images", func(ctx *ToolContext, input struct { Prompt string `json:"prompt"` }) (*MultipartToolResponse, error) { @@ -327,7 +327,7 @@ func TestDefineMultipartTool(t *testing.T) { t.Run("multipart tool returns parts", func(t *testing.T) { r := newTestRegistry(t) - tl := DefineMultipartTool(r, "provider/multiOut", "Returns multiple parts", + tl := defineMultipartTool(r, "provider/multiOut", "Returns multiple parts", func(ctx *ToolContext, input struct{}) (*MultipartToolResponse, error) { return &MultipartToolResponse{ Output: map[string]any{"status": "ok"}, @@ -385,7 +385,7 @@ func TestNewMultipartTool(t *testing.T) { func TestToolDefinition(t *testing.T) { t.Run("includes all fields", func(t *testing.T) { r := newTestRegistry(t) - tl := DefineTool(r, "provider/complete", "A complete tool", func(ctx *ToolContext, input struct { + tl := defineTool(r, "provider/complete", "A complete tool", func(ctx *ToolContext, input struct { Query string `json:"query"` }) (struct { Result string `json:"result"` @@ -431,7 +431,7 @@ func TestLookupTool(t *testing.T) { t.Run("finds registered tool", func(t *testing.T) { r := newTestRegistry(t) - DefineTool(r, "test/findMe", "Find me", func(ctx *ToolContext, input struct{}) (bool, error) { + defineTool(r, "test/findMe", "Find me", func(ctx *ToolContext, input struct{}) (bool, error) { return true, nil }) @@ -450,7 +450,7 @@ func TestWithStrictSchema(t *testing.T) { t.Run("absent by default", func(t *testing.T) { r := newTestRegistry(t) - tl := DefineTool(r, "strict/default", "no strict opt", func(ctx *ToolContext, input struct{}) (string, error) { + tl := defineTool(r, "strict/default", "no strict opt", func(ctx *ToolContext, input struct{}) (string, error) { return "", nil }) def := tl.Definition() @@ -475,7 +475,7 @@ func TestWithStrictSchema(t *testing.T) { t.Run("DefineTool with WithStrictSchema(true) is surfaced on Definition", func(t *testing.T) { r := newTestRegistry(t) - tl := DefineTool(r, "strict/registered-true", "registered strict", + tl := defineTool(r, "strict/registered-true", "registered strict", func(ctx *ToolContext, input struct{}) (string, error) { return "", nil }, WithStrictSchema(true), ) @@ -484,7 +484,7 @@ func TestWithStrictSchema(t *testing.T) { t.Run("DefineTool with WithStrictSchema(false) is surfaced on Definition", func(t *testing.T) { r := newTestRegistry(t) - tl := DefineTool(r, "strict/registered-false", "registered loose", + tl := defineTool(r, "strict/registered-false", "registered loose", func(ctx *ToolContext, input struct{}) (string, error) { return "", nil }, WithStrictSchema(false), ) @@ -493,7 +493,7 @@ func TestWithStrictSchema(t *testing.T) { t.Run("LookupTool round-trips the strict flag for registered tools", func(t *testing.T) { r := newTestRegistry(t) - DefineTool(r, "strict/lookup-true", "registered strict", + defineTool(r, "strict/lookup-true", "registered strict", func(ctx *ToolContext, input struct{}) (string, error) { return "", nil }, WithStrictSchema(true), ) @@ -522,7 +522,7 @@ func TestWithStrictSchema(t *testing.T) { t.Run("DefineMultipartTool plumbs strict the same way", func(t *testing.T) { r := newTestRegistry(t) - tl := DefineMultipartTool(r, "strict/multipart", "multipart strict", + tl := defineMultipartTool(r, "strict/multipart", "multipart strict", func(ctx *ToolContext, input struct{}) (*MultipartToolResponse, error) { return &MultipartToolResponse{}, nil }, @@ -547,7 +547,7 @@ func TestWithStrictSchema(t *testing.T) { } }() r := newTestRegistry(t) - DefineTool(r, "strict/double-set", "double set", + defineTool(r, "strict/double-set", "double set", func(ctx *ToolContext, input struct{}) (string, error) { return "", nil }, WithStrictSchema(true), WithStrictSchema(false), @@ -558,7 +558,7 @@ func TestWithStrictSchema(t *testing.T) { func TestToolIsMultipart(t *testing.T) { t.Run("regular tool is not multipart", func(t *testing.T) { r := newTestRegistry(t) - tl := DefineTool(r, "provider/regular", "Regular tool", func(ctx *ToolContext, input struct{}) (string, error) { + tl := defineTool(r, "provider/regular", "Regular tool", func(ctx *ToolContext, input struct{}) (string, error) { return "ok", nil }) @@ -570,7 +570,7 @@ func TestToolIsMultipart(t *testing.T) { t.Run("multipart tool is multipart", func(t *testing.T) { r := newTestRegistry(t) - tl := DefineMultipartTool(r, "provider/multi", "Multi tool", + tl := defineMultipartTool(r, "provider/multi", "Multi tool", func(ctx *ToolContext, input struct{}) (*MultipartToolResponse, error) { return &MultipartToolResponse{}, nil }) @@ -585,7 +585,7 @@ func TestToolIsMultipart(t *testing.T) { func TestToolRunRaw(t *testing.T) { t.Run("returns output from regular tool", func(t *testing.T) { r := newTestRegistry(t) - tl := DefineTool(r, "provider/sum", "Sums numbers", func(ctx *ToolContext, input struct { + tl := defineTool(r, "provider/sum", "Sums numbers", func(ctx *ToolContext, input struct { Nums []int `json:"nums"` }) (int, error) { sum := 0 @@ -610,7 +610,7 @@ func TestToolRunRaw(t *testing.T) { t.Run("returns error from tool", func(t *testing.T) { r := newTestRegistry(t) - tl := DefineTool(r, "provider/fail", "Always fails", func(ctx *ToolContext, input struct{}) (string, error) { + tl := defineTool(r, "provider/fail", "Always fails", func(ctx *ToolContext, input struct{}) (string, error) { return "", errors.New("intentional failure") }) @@ -624,7 +624,7 @@ func TestToolRunRaw(t *testing.T) { func TestToolRunRawMultipart(t *testing.T) { t.Run("returns full response from multipart tool", func(t *testing.T) { r := newTestRegistry(t) - tl := DefineMultipartTool(r, "provider/fullResp", "Full response", + tl := defineMultipartTool(r, "provider/fullResp", "Full response", func(ctx *ToolContext, input struct{}) (*MultipartToolResponse, error) { return &MultipartToolResponse{ Output: "main output", @@ -650,7 +650,7 @@ func TestToolRunRawMultipart(t *testing.T) { func TestToolRespond(t *testing.T) { r := newTestRegistry(t) - tl := DefineTool(r, "provider/responder", "Test responder", func(ctx *ToolContext, input struct{}) (string, error) { + tl := defineTool(r, "provider/responder", "Test responder", func(ctx *ToolContext, input struct{}) (string, error) { return "ok", nil }) @@ -718,7 +718,7 @@ func TestToolRespond(t *testing.T) { func TestToolRestart(t *testing.T) { r := newTestRegistry(t) - tl := DefineTool(r, "provider/restarter", "Test restarter", func(ctx *ToolContext, input struct { + tl := defineTool(r, "provider/restarter", "Test restarter", func(ctx *ToolContext, input struct { Value int `json:"value"` }) (int, error) { return input.Value, nil @@ -816,7 +816,7 @@ func TestToolRestart(t *testing.T) { func TestToolInterrupt(t *testing.T) { t.Run("tool can interrupt execution", func(t *testing.T) { r := newTestRegistry(t) - tl := DefineTool(r, "provider/interrupter", "Can interrupt", + tl := defineTool(r, "provider/interrupter", "Can interrupt", func(ctx *ToolContext, input struct { ShouldInterrupt bool `json:"shouldInterrupt"` }) (string, error) { @@ -847,7 +847,7 @@ func TestToolInterrupt(t *testing.T) { t.Run("tool completes without interrupt", func(t *testing.T) { r := newTestRegistry(t) - tl := DefineTool(r, "provider/noInterrupt", "No interrupt", + tl := defineTool(r, "provider/noInterrupt", "No interrupt", func(ctx *ToolContext, input struct { ShouldInterrupt bool `json:"shouldInterrupt"` }) (string, error) { @@ -880,7 +880,7 @@ func TestToolWithInputSchemaOption(t *testing.T) { }, } - tl := DefineTool(r, "provider/customInput", "Custom input schema", + tl := defineTool(r, "provider/customInput", "Custom input schema", func(ctx *ToolContext, input any) (string, error) { m := input.(map[string]any) return m["customField"].(string), nil @@ -917,10 +917,10 @@ func TestToolWithInputSchemaOption(t *testing.T) { func TestResolveUniqueTools(t *testing.T) { t.Run("resolves tools from registry", func(t *testing.T) { r := newTestRegistry(t) - DefineTool(r, "provider/tool1", "Tool 1", func(ctx *ToolContext, input struct{}) (bool, error) { + defineTool(r, "provider/tool1", "Tool 1", func(ctx *ToolContext, input struct{}) (bool, error) { return true, nil }) - DefineTool(r, "provider/tool2", "Tool 2", func(ctx *ToolContext, input struct{}) (bool, error) { + defineTool(r, "provider/tool2", "Tool 2", func(ctx *ToolContext, input struct{}) (bool, error) { return true, nil }) @@ -981,7 +981,7 @@ func TestResolveUniqueTools(t *testing.T) { func TestIsMultipart(t *testing.T) { t.Run("returns false for standard tool", func(t *testing.T) { r := newTestRegistry(t) - tl := DefineTool(r, "provider/standard", "Standard tool", + tl := defineTool(r, "provider/standard", "Standard tool", func(ctx *ToolContext, input struct{}) (string, error) { return "result", nil }) @@ -1004,7 +1004,7 @@ func TestIsMultipart(t *testing.T) { t.Run("returns true for multipart tool", func(t *testing.T) { r := newTestRegistry(t) - tl := DefineMultipartTool(r, "provider/multipart", "Multipart tool", + tl := defineMultipartTool(r, "provider/multipart", "Multipart tool", func(ctx *ToolContext, input struct{}) (*MultipartToolResponse, error) { return &MultipartToolResponse{ Content: []*Part{NewTextPart("hello"), NewTextPart("world")}, diff --git a/go/genkit/exp/tools.go b/go/genkit/exp/tools.go index 1d0edbd703..a09974dac1 100644 --- a/go/genkit/exp/tools.go +++ b/go/genkit/exp/tools.go @@ -63,7 +63,9 @@ import ( // fmt.Println(resp.Text()) func DefineTool[In, Out any](g *genkit.Genkit, name, description string, fn aix.ToolFunc[In, Out], opts ...ai.ToolOption) *aix.Tool[In, Out] { requireExperimental(g, "DefineTool") - return aix.DefineTool(genkitbridge.RegistryOf(g), name, description, fn, opts...) + t := aix.NewTool(name, description, fn, opts...) + t.Register(genkitbridge.RegistryOf(g)) + return t } // DefineInterruptibleTool defines a tool that supports typed interrupt/resume, @@ -146,5 +148,7 @@ func DefineTool[In, Out any](g *genkit.Genkit, name, description string, fn aix. // } func DefineInterruptibleTool[In, Out, Resume any](g *genkit.Genkit, name, description string, fn aix.InterruptibleToolFunc[In, Out, Resume], opts ...ai.ToolOption) *aix.InterruptibleTool[In, Out, Resume] { requireExperimental(g, "DefineInterruptibleTool") - return aix.DefineInterruptibleTool(genkitbridge.RegistryOf(g), name, description, fn, opts...) + t := aix.NewInterruptibleTool(name, description, fn, opts...) + t.Register(genkitbridge.RegistryOf(g)) + return t } diff --git a/go/genkit/genkit.go b/go/genkit/genkit.go index 787992122a..a909da0fb2 100644 --- a/go/genkit/genkit.go +++ b/go/genkit/genkit.go @@ -571,7 +571,9 @@ func DefineTypedModel[Config any](g *Genkit, name string, opts *ai.ModelOptions, // Deprecated: Use [DefineTypedModel], which passes the request's config // to fn as a typed value instead of leaving it type-erased on the request. func DefineModel(g *Genkit, name string, opts *ai.ModelOptions, fn ai.ModelFunc) ai.Model { - return ai.DefineModel(g.reg, name, opts, fn) + m := ai.NewModel(name, opts, fn) + m.Register(g.reg) + return m } // DefineTypedBackgroundModel defines a background model, registers it as @@ -597,7 +599,9 @@ func DefineTypedBackgroundModel[Config any](g *Genkit, name string, opts *ai.Bac // request's config to startFn as a typed value instead of leaving it // type-erased on the request. func DefineBackgroundModel(g *Genkit, name string, opts *ai.BackgroundModelOptions, startFn ai.StartModelOpFunc, checkFn ai.CheckModelOpFunc) ai.BackgroundModel { - return ai.DefineBackgroundModel(g.reg, name, opts, startFn, checkFn) + m := ai.NewBackgroundModel(name, opts, startFn, checkFn) + m.Register(g.reg) + return m } // LookupModel retrieves a registered [ai.Model] by its provider and name. @@ -660,7 +664,9 @@ func LookupBackgroundModel(g *Genkit, name string) ai.BackgroundModel { // // fmt.Println(resp.Text()) // Might output something like "The weather in Paris is Sunny, 25°C." func DefineTool[In, Out any](g *Genkit, name, description string, fn ai.ToolFunc[In, Out], opts ...ai.ToolOption) *ai.ToolDef[In, Out] { - return ai.DefineTool(g.reg, name, description, fn, opts...) + t := ai.NewTool(name, description, fn, opts...) + t.Register(g.reg) + return t } // DefineToolWithInputSchema defines a tool with a custom input schema that can be used by models during generation, @@ -709,7 +715,9 @@ func DefineTool[In, Out any](g *Genkit, name, description string, fn ai.ToolFunc // ai.WithToolInputSchema(inputSchema), // ) func DefineToolWithInputSchema[Out any](g *Genkit, name, description string, inputSchema map[string]any, fn ai.ToolFunc[any, Out]) *ai.ToolDef[any, Out] { - return ai.DefineTool(g.reg, name, description, fn, ai.WithInputSchema(inputSchema)) + t := ai.NewTool(name, description, fn, ai.WithInputSchema(inputSchema)) + t.Register(g.reg) + return t } // DefineMultipartTool defines a multipart tool that can be used by models during generation, @@ -769,7 +777,9 @@ func DefineToolWithInputSchema[Out any](g *Genkit, name, description string, inp // // fmt.Println(resp.Text()) func DefineMultipartTool[In any](g *Genkit, name, description string, fn ai.MultipartToolFunc[In], opts ...ai.ToolOption) *ai.ToolDef[In, *ai.MultipartToolResponse] { - return ai.DefineMultipartTool(g.reg, name, description, fn, opts...) + t := ai.NewMultipartTool(name, description, fn, opts...) + t.Register(g.reg) + return t } // LookupTool retrieves a registered tool by its name. @@ -831,7 +841,9 @@ func LookupTool(g *Genkit, name string) ai.Tool { // ai.WithUse(Trace{Label: "debug"}), // ) func DefineMiddleware[M ai.Middleware](g *Genkit, description string, prototype M) *ai.MiddlewareDesc { - return ai.DefineMiddleware(g.reg, description, prototype) + d := ai.NewMiddleware(description, prototype) + d.Register(g.reg) + return d } // LookupMiddleware retrieves a registered middleware descriptor by its name. @@ -1398,7 +1410,9 @@ func DefineTypedEmbedder[Config any](g *Genkit, name string, opts *ai.EmbedderOp // options to fn as a typed value instead of leaving them type-erased on the // request. func DefineEmbedder(g *Genkit, name string, opts *ai.EmbedderOptions, fn ai.EmbedderFunc) ai.Embedder { - return ai.DefineEmbedder(g.reg, name, opts, fn) + e := ai.NewEmbedder(name, opts, fn) + e.Register(g.reg) + return e } // LookupEmbedder retrieves a registered [ai.Embedder] by its provider and name. @@ -1442,7 +1456,9 @@ func DefineTypedEvaluator[Config any](g *Genkit, name string, opts *ai.Evaluator // options to fn as a typed value instead of leaving them type-erased on the // request. func DefineEvaluator(g *Genkit, name string, opts *ai.EvaluatorOptions, fn ai.EvaluatorFunc) ai.Evaluator { - return ai.DefineEvaluator(g.reg, name, opts, fn) + e := ai.NewEvaluator(name, opts, fn) + e.Register(g.reg) + return e } // DefineTypedBatchEvaluator defines an evaluator that processes the @@ -1469,7 +1485,9 @@ func DefineTypedBatchEvaluator[Config any](g *Genkit, name string, opts *ai.Eval // request's options to fn as a typed value instead of leaving them // type-erased on the request. func DefineBatchEvaluator(g *Genkit, name string, opts *ai.EvaluatorOptions, fn ai.BatchEvaluatorFunc) ai.Evaluator { - return ai.DefineBatchEvaluator(g.reg, name, opts, fn) + e := ai.NewBatchEvaluator(name, opts, fn) + e.Register(g.reg) + return e } // LookupEvaluator retrieves a registered [ai.Evaluator] by its provider and name. @@ -1739,7 +1757,9 @@ func IsDefinedFormat(g *Genkit, name string) bool { // }, nil // }) func DefineResource(g *Genkit, name string, opts *ai.ResourceOptions, fn ai.ResourceFunc) ai.Resource { - return ai.DefineResource(g.reg, name, opts, fn) + res := ai.NewResource(name, opts, fn) + res.Register(g.reg) + return res } // FindMatchingResource finds a resource that matches the given URI. diff --git a/go/internal/fakeembedder/fakeembedder_test.go b/go/internal/fakeembedder/fakeembedder_test.go index 9b309906a1..237ec94d06 100644 --- a/go/internal/fakeembedder/fakeembedder_test.go +++ b/go/internal/fakeembedder/fakeembedder_test.go @@ -37,7 +37,8 @@ func TestFakeEmbedder(t *testing.T) { }, ConfigSchema: nil, } - emb := ai.DefineEmbedder(r, "fake/embed", emdOpts, embed.Embed) + emb := ai.NewEmbedder("fake/embed", emdOpts, embed.Embed) + emb.Register(r) d := ai.DocumentFromText("fakeembedder test", nil) vals := []float32{1, 2} diff --git a/go/plugins/middleware/define_test.go b/go/plugins/middleware/define_test.go new file mode 100644 index 0000000000..0790f9d810 --- /dev/null +++ b/go/plugins/middleware/define_test.go @@ -0,0 +1,43 @@ +// Copyright 2026 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 middleware + +import ( + "github.com/firebase/genkit/go/ai" + "github.com/firebase/genkit/go/core/api" +) + +// Test-local define helpers: New* + Register in one call, mirroring the +// removed registry-taking ai.Define* helpers so tests stay concise. + +func registerTestModel(r api.Registry, name string, opts *ai.ModelOptions, fn ai.ModelFunc) ai.Model { + m := ai.NewModel(name, opts, fn) + m.Register(r) + return m +} + +func registerTestTool[In, Out any](r api.Registry, name, description string, fn ai.ToolFunc[In, Out], opts ...ai.ToolOption) *ai.ToolDef[In, Out] { + t := ai.NewTool(name, description, fn, opts...) + t.Register(r) + return t +} + +func registerTestMiddleware[M ai.Middleware](r api.Registry, description string, prototype M) *ai.MiddlewareDesc { + d := ai.NewMiddleware(description, prototype) + d.Register(r) + return d +} diff --git a/go/plugins/middleware/filesystem_test.go b/go/plugins/middleware/filesystem_test.go index 181eb51bcd..2721d31351 100644 --- a/go/plugins/middleware/filesystem_test.go +++ b/go/plugins/middleware/filesystem_test.go @@ -72,7 +72,7 @@ func scriptedModel(t *testing.T, r *registry.Registry, name string, script []mod t.Helper() var seen []*ai.ModelRequest idx := 0 - m := ai.DefineModel(r, name, &ai.ModelOptions{ + m := registerTestModel(r, name, &ai.ModelOptions{ Supports: &ai.ModelSupports{Multiturn: true, SystemRole: true, Tools: true, Media: true}, }, func(ctx context.Context, req *ai.ModelRequest, cb ai.ModelStreamCallback) (*ai.ModelResponse, error) { seen = append(seen, req) @@ -106,7 +106,7 @@ func TestFilesystemListFilesNonRecursive(t *testing.T) { }) fs := &Filesystem{RootDir: root} - ai.DefineMiddleware(r, "filesystem", fs) + registerTestMiddleware(r, "filesystem", fs) resp, err := ai.Generate(ctx, r, ai.WithModel(m), ai.WithPrompt("go"), ai.WithUse(fs)) if err != nil { @@ -162,7 +162,7 @@ func TestFilesystemListFilesRecursive(t *testing.T) { }) fs := &Filesystem{RootDir: root} - ai.DefineMiddleware(r, "filesystem", fs) + registerTestMiddleware(r, "filesystem", fs) resp, err := ai.Generate(ctx, r, ai.WithModel(m), ai.WithPrompt("go"), ai.WithUse(fs)) if err != nil { @@ -203,7 +203,7 @@ func TestFilesystemReadFileInjectsUserMessage(t *testing.T) { }) fs := &Filesystem{RootDir: root} - ai.DefineMiddleware(r, "filesystem", fs) + registerTestMiddleware(r, "filesystem", fs) if _, err := ai.Generate(ctx, r, ai.WithModel(m), ai.WithPrompt("go"), ai.WithUse(fs)); err != nil { t.Fatal(err) @@ -258,7 +258,7 @@ func TestFilesystemReadFileInjectsImageAsMedia(t *testing.T) { }) fs := &Filesystem{RootDir: root} - ai.DefineMiddleware(r, "filesystem", fs) + registerTestMiddleware(r, "filesystem", fs) if _, err := ai.Generate(ctx, r, ai.WithModel(m), ai.WithPrompt("go"), ai.WithUse(fs)); err != nil { t.Fatal(err) @@ -294,7 +294,7 @@ func TestFilesystemStreamsInjectedFileContents(t *testing.T) { // Streaming model that emits the same payload it returns, so each turn // produces at least one chunk. turn := 0 - m := ai.DefineModel(r, "test/fs-stream", &ai.ModelOptions{ + m := registerTestModel(r, "test/fs-stream", &ai.ModelOptions{ Supports: &ai.ModelSupports{Multiturn: true, SystemRole: true, Tools: true, Media: true}, }, func(ctx context.Context, req *ai.ModelRequest, cb ai.ModelStreamCallback) (*ai.ModelResponse, error) { turn++ @@ -320,7 +320,7 @@ func TestFilesystemStreamsInjectedFileContents(t *testing.T) { }) fs := &Filesystem{RootDir: root} - ai.DefineMiddleware(r, "filesystem", fs) + registerTestMiddleware(r, "filesystem", fs) var chunks []*ai.ModelResponseChunk if _, err := ai.Generate(ctx, r, @@ -386,7 +386,7 @@ func TestFilesystemWriteFileRequiresAllowWriteAccess(t *testing.T) { // Only list_files/read_file should be registered; write_file must not be. fs := &Filesystem{RootDir: root} - ai.DefineMiddleware(r, "filesystem", fs) + registerTestMiddleware(r, "filesystem", fs) hooks, err := fs.New(ctx) if err != nil { @@ -420,7 +420,7 @@ func TestFilesystemWriteFileCreatesAndOverwrites(t *testing.T) { }) fs := &Filesystem{RootDir: root, AllowWriteAccess: true} - ai.DefineMiddleware(r, "filesystem", fs) + registerTestMiddleware(r, "filesystem", fs) if _, err := ai.Generate(ctx, r, ai.WithModel(m), ai.WithPrompt("go"), ai.WithUse(fs)); err != nil { t.Fatal(err) @@ -460,7 +460,7 @@ func TestFilesystemEditFile(t *testing.T) { }) fs := &Filesystem{RootDir: root, AllowWriteAccess: true} - ai.DefineMiddleware(r, "filesystem", fs) + registerTestMiddleware(r, "filesystem", fs) if _, err := ai.Generate(ctx, r, ai.WithModel(m), ai.WithPrompt("go"), ai.WithUse(fs)); err != nil { t.Fatal(err) @@ -503,7 +503,7 @@ func TestFilesystemEditFileNotFound(t *testing.T) { }) fs := &Filesystem{RootDir: root, AllowWriteAccess: true} - ai.DefineMiddleware(r, "filesystem", fs) + registerTestMiddleware(r, "filesystem", fs) resp, err := ai.Generate(ctx, r, ai.WithModel(m), ai.WithPrompt("go"), ai.WithUse(fs)) if err != nil { @@ -564,7 +564,7 @@ func TestFilesystemRejectsTraversal(t *testing.T) { }) fs := &Filesystem{RootDir: root} - ai.DefineMiddleware(r, "filesystem", fs) + registerTestMiddleware(r, "filesystem", fs) if _, err := ai.Generate(ctx, r, ai.WithModel(m), ai.WithPrompt("go"), ai.WithUse(fs)); err != nil { t.Fatal(err) @@ -630,7 +630,7 @@ func TestFilesystemRespectsConfigOverride(t *testing.T) { } // Register the middleware with the proto root. - ai.DefineMiddleware(r, "filesystem", &Filesystem{RootDir: protoRoot}) + registerTestMiddleware(r, "filesystem", &Filesystem{RootDir: protoRoot}) m, seen := scriptedModel(t, r, "test/fs-devui", []modelTurn{ {Tools: []*ai.ToolRequest{{Name: "read_file", Input: map[string]any{"filePath": "marker.txt"}}}}, @@ -679,7 +679,7 @@ func TestFilesystemMissingRootDirIsAnError(t *testing.T) { }) fs := &Filesystem{} - ai.DefineMiddleware(r, "filesystem", fs) + registerTestMiddleware(r, "filesystem", fs) _, err := ai.Generate(ctx, r, ai.WithModel(m), ai.WithPrompt("go"), ai.WithUse(fs)) if err == nil { @@ -705,7 +705,7 @@ func TestFilesystemListFilesIncludesSize(t *testing.T) { }) fs := &Filesystem{RootDir: root} - ai.DefineMiddleware(r, "filesystem", fs) + registerTestMiddleware(r, "filesystem", fs) resp, err := ai.Generate(ctx, r, ai.WithModel(m), ai.WithPrompt("go"), ai.WithUse(fs)) if err != nil { t.Fatal(err) @@ -766,7 +766,7 @@ func TestFilesystemReadFileIncludesLineMetadata(t *testing.T) { }) fs := &Filesystem{RootDir: root} - ai.DefineMiddleware(r, "filesystem", fs) + registerTestMiddleware(r, "filesystem", fs) if _, err := ai.Generate(ctx, r, ai.WithModel(m), ai.WithPrompt("go"), ai.WithUse(fs)); err != nil { t.Fatal(err) } @@ -825,7 +825,7 @@ func TestFilesystemEditFileRequiresPriorRead(t *testing.T) { }) fs := &Filesystem{RootDir: root, AllowWriteAccess: true} - ai.DefineMiddleware(r, "filesystem", fs) + registerTestMiddleware(r, "filesystem", fs) if _, err := ai.Generate(ctx, r, ai.WithModel(m), ai.WithPrompt("go"), ai.WithUse(fs)); err != nil { t.Fatal(err) } @@ -868,7 +868,7 @@ func TestFilesystemEditFileDetectsExternalModification(t *testing.T) { } turn := 0 - m := ai.DefineModel(r, "test/fs-edit-stale", &ai.ModelOptions{ + m := registerTestModel(r, "test/fs-edit-stale", &ai.ModelOptions{ Supports: &ai.ModelSupports{Multiturn: true, SystemRole: true, Tools: true}, }, func(ctx context.Context, req *ai.ModelRequest, _ ai.ModelStreamCallback) (*ai.ModelResponse, error) { // Between the read and the edit, simulate the user editing the file. @@ -892,7 +892,7 @@ func TestFilesystemEditFileDetectsExternalModification(t *testing.T) { }) fs := &Filesystem{RootDir: root, AllowWriteAccess: true} - ai.DefineMiddleware(r, "filesystem", fs) + registerTestMiddleware(r, "filesystem", fs) if _, err := ai.Generate(ctx, r, ai.WithModel(m), ai.WithPrompt("go"), ai.WithUse(fs)); err != nil { t.Fatal(err) } @@ -926,7 +926,7 @@ func TestFilesystemReadFileDedupsUnchanged(t *testing.T) { }) fs := &Filesystem{RootDir: root} - ai.DefineMiddleware(r, "filesystem", fs) + registerTestMiddleware(r, "filesystem", fs) if _, err := ai.Generate(ctx, r, ai.WithModel(m), ai.WithPrompt("go"), ai.WithUse(fs)); err != nil { t.Fatal(err) } @@ -986,7 +986,7 @@ func TestFilesystemWriteFileRequiresReadOnOverwrite(t *testing.T) { }) fs := &Filesystem{RootDir: root, AllowWriteAccess: true} - ai.DefineMiddleware(r, "filesystem", fs) + registerTestMiddleware(r, "filesystem", fs) if _, err := ai.Generate(ctx, r, ai.WithModel(m), ai.WithPrompt("go"), ai.WithUse(fs)); err != nil { t.Fatal(err) } @@ -1019,7 +1019,7 @@ func TestFilesystemReadFileOffsetLimit(t *testing.T) { }) fs := &Filesystem{RootDir: root} - ai.DefineMiddleware(r, "filesystem", fs) + registerTestMiddleware(r, "filesystem", fs) if _, err := ai.Generate(ctx, r, ai.WithModel(m), ai.WithPrompt("go"), ai.WithUse(fs)); err != nil { t.Fatal(err) } diff --git a/go/plugins/middleware/retry_test.go b/go/plugins/middleware/retry_test.go index 679a4ccec3..00d41058ea 100644 --- a/go/plugins/middleware/retry_test.go +++ b/go/plugins/middleware/retry_test.go @@ -36,7 +36,7 @@ func newTestRegistry(t *testing.T) *registry.Registry { func defineModel(t *testing.T, r *registry.Registry, name string, fn ai.ModelFunc) ai.Model { t.Helper() - return ai.DefineModel(r, name, &ai.ModelOptions{ + return registerTestModel(r, name, &ai.ModelOptions{ Supports: &ai.ModelSupports{Multiturn: true, SystemRole: true}, }, fn) } @@ -55,7 +55,7 @@ func TestRetrySucceedsOnFirstAttempt(t *testing.T) { }) retry := &Retry{} - ai.DefineMiddleware(r, "retry", retry) + registerTestMiddleware(r, "retry", retry) resp, err := ai.Generate(ctx, r, ai.WithModel(m), ai.WithPrompt("hello"), ai.WithUse(retry)) if err != nil { @@ -81,7 +81,7 @@ func TestRetryRecoversAfterTransientError(t *testing.T) { }) retry := &Retry{} - ai.DefineMiddleware(r, "retry", retry) + registerTestMiddleware(r, "retry", retry) resp, err := ai.Generate(ctx, r, ai.WithModel(m), ai.WithPrompt("hello"), ai.WithUse(retry)) if err != nil { @@ -104,7 +104,7 @@ func TestRetryExhaustsMaxRetries(t *testing.T) { }) retry := &Retry{MaxRetries: 2} - ai.DefineMiddleware(r, "retry", retry) + registerTestMiddleware(r, "retry", retry) _, err := ai.Generate(ctx, r, ai.WithModel(m), ai.WithPrompt("hello"), ai.WithUse(retry)) if err == nil { @@ -125,7 +125,7 @@ func TestRetryDoesNotRetryNonMatchingGenkitError(t *testing.T) { }) retry := &Retry{} - ai.DefineMiddleware(r, "retry", retry) + registerTestMiddleware(r, "retry", retry) _, err := ai.Generate(ctx, r, ai.WithModel(m), ai.WithPrompt("hello"), ai.WithUse(retry)) if err == nil { @@ -148,7 +148,7 @@ func TestRetryRetriesNonGenkitErrors(t *testing.T) { }) retry := &Retry{} - ai.DefineMiddleware(r, "retry", retry) + registerTestMiddleware(r, "retry", retry) resp, err := ai.Generate(ctx, r, ai.WithModel(m), ai.WithPrompt("hello"), ai.WithUse(retry)) if err != nil { @@ -174,7 +174,7 @@ func TestRetryCustomStatuses(t *testing.T) { Statuses: []core.StatusName{core.PERMISSION_DENIED}, MaxRetries: 1, } - ai.DefineMiddleware(r, "retry", retry) + registerTestMiddleware(r, "retry", retry) _, err := ai.Generate(ctx, r, ai.WithModel(m), ai.WithPrompt("hello"), ai.WithUse(retry)) if err == nil { @@ -203,7 +203,7 @@ func TestRetryBackoffDelays(t *testing.T) { BackoffFactor: 2, NoJitter: true, } - ai.DefineMiddleware(r, "retry", retry) + registerTestMiddleware(r, "retry", retry) _, _ = ai.Generate(ctx, r, ai.WithModel(m), ai.WithPrompt("hello"), ai.WithUse(retry)) @@ -237,7 +237,7 @@ func TestRetryMaxDelayClamp(t *testing.T) { BackoffFactor: 2, NoJitter: true, } - ai.DefineMiddleware(r, "retry", retry) + registerTestMiddleware(r, "retry", retry) _, _ = ai.Generate(ctx, r, ai.WithModel(m), ai.WithPrompt("hello"), ai.WithUse(retry)) @@ -272,7 +272,7 @@ func TestRetryStopsWhenContextCanceledDuringBackoff(t *testing.T) { defer func() { sleepFunc = origSleep }() retry := &Retry{MaxRetries: 5} - ai.DefineMiddleware(r, "retry", retry) + registerTestMiddleware(r, "retry", retry) _, err := ai.Generate(reqCtx, r, ai.WithModel(m), ai.WithPrompt("hello"), ai.WithUse(retry)) if err == nil { diff --git a/go/plugins/middleware/skills_test.go b/go/plugins/middleware/skills_test.go index de61288c3e..a1dc7ec94a 100644 --- a/go/plugins/middleware/skills_test.go +++ b/go/plugins/middleware/skills_test.go @@ -63,7 +63,7 @@ func setupSkillsDir(t *testing.T) string { func captureModel(t *testing.T, r *registry.Registry, name string) (ai.Model, *[]*ai.Message) { t.Helper() var captured []*ai.Message - m := ai.DefineModel(r, name, &ai.ModelOptions{ + m := registerTestModel(r, name, &ai.ModelOptions{ Supports: &ai.ModelSupports{Multiturn: true, SystemRole: true, Tools: true}, }, func(ctx context.Context, req *ai.ModelRequest, cb ai.ModelStreamCallback) (*ai.ModelResponse, error) { captured = req.Messages @@ -77,7 +77,7 @@ func captureModel(t *testing.T, r *registry.Registry, name string) (ai.Model, *[ // messages. func toolCallingModel(t *testing.T, r *registry.Registry, name, toolName string, input map[string]any) ai.Model { t.Helper() - return ai.DefineModel(r, name, &ai.ModelOptions{ + return registerTestModel(r, name, &ai.ModelOptions{ Supports: &ai.ModelSupports{Multiturn: true, SystemRole: true, Tools: true}, }, func(ctx context.Context, req *ai.ModelRequest, cb ai.ModelStreamCallback) (*ai.ModelResponse, error) { for _, msg := range req.Messages { @@ -106,7 +106,7 @@ func TestSkillsInjectsSystemPrompt(t *testing.T) { m, captured := captureModel(t, r, "test/capture") s := &Skills{SkillPaths: []string{skillsDir}} - ai.DefineMiddleware(r, "skills", s) + registerTestMiddleware(r, "skills", s) if _, err := ai.Generate(ctx, r, ai.WithModel(m), ai.WithPrompt("hello"), ai.WithUse(s)); err != nil { t.Fatal(err) @@ -139,7 +139,7 @@ func TestSkillsRegistersUseSkillTool(t *testing.T) { m := toolCallingModel(t, r, "test/toolcaller", useSkillToolName, map[string]any{"skillName": "python"}) s := &Skills{SkillPaths: []string{skillsDir}} - ai.DefineMiddleware(r, "skills", s) + registerTestMiddleware(r, "skills", s) resp, err := ai.Generate(ctx, r, ai.WithModel(m), ai.WithPrompt("use python"), ai.WithUse(s)) if err != nil { @@ -172,7 +172,7 @@ func TestSkillsUnknownSkillReturnsError(t *testing.T) { m := toolCallingModel(t, r, "test/unknown", useSkillToolName, map[string]any{"skillName": "nonexistent"}) s := &Skills{SkillPaths: []string{skillsDir}} - ai.DefineMiddleware(r, "skills", s) + registerTestMiddleware(r, "skills", s) _, err := ai.Generate(ctx, r, ai.WithModel(m), ai.WithPrompt("use skill"), ai.WithUse(s)) if err == nil { @@ -190,7 +190,7 @@ func TestSkillsPromptInjectionIsIdempotent(t *testing.T) { m, captured := captureModel(t, r, "test/idempotent") s := &Skills{SkillPaths: []string{skillsDir}} - ai.DefineMiddleware(r, "skills", s) + registerTestMiddleware(r, "skills", s) // First call produces a system message with a single block. resp, err := ai.Generate(ctx, r, ai.WithModel(m), ai.WithPrompt("hello"), ai.WithUse(s)) @@ -238,7 +238,7 @@ func TestSkillsNoopWhenNoSkillsFound(t *testing.T) { // Point at an empty directory — no skills, so the middleware should // leave the request untouched. s := &Skills{SkillPaths: []string{t.TempDir()}} - ai.DefineMiddleware(r, "skills", s) + registerTestMiddleware(r, "skills", s) if _, err := ai.Generate(ctx, r, ai.WithModel(m), ai.WithPrompt("hello"), ai.WithUse(s)); err != nil { t.Fatal(err) diff --git a/go/plugins/middleware/tool_approval_test.go b/go/plugins/middleware/tool_approval_test.go index 5796f07b86..ed3dac5a9e 100644 --- a/go/plugins/middleware/tool_approval_test.go +++ b/go/plugins/middleware/tool_approval_test.go @@ -27,14 +27,14 @@ import ( func defineToolModel(t *testing.T, r *registry.Registry, name string, fn ai.ModelFunc) ai.Model { t.Helper() - return ai.DefineModel(r, name, &ai.ModelOptions{ + return registerTestModel(r, name, &ai.ModelOptions{ Supports: &ai.ModelSupports{Multiturn: true, SystemRole: true, Tools: true}, }, fn) } func defineTool(t *testing.T, r api.Registry, name string) ai.Tool { t.Helper() - return ai.DefineTool(r, name, "test tool", + return registerTestTool(r, name, "test tool", func(ctx *ai.ToolContext, input struct { V string `json:"v"` }) (string, error) { @@ -77,7 +77,7 @@ func TestToolApprovalAllowsApprovedTools(t *testing.T) { alsoAllowed := defineTool(t, r, "alsoAllowed") ta := &ToolApproval{AllowedTools: []string{"allowed", "alsoAllowed"}} - ai.DefineMiddleware(r, "toolApproval", ta) + registerTestMiddleware(r, "toolApproval", ta) resp, err := ai.Generate(ctx, r, ai.WithModel(m), @@ -104,7 +104,7 @@ func TestToolApprovalInterruptsUnapprovedTools(t *testing.T) { dangerous := defineTool(t, r, "dangerous") ta := &ToolApproval{AllowedTools: []string{"safe"}} - ai.DefineMiddleware(r, "toolApproval", ta) + registerTestMiddleware(r, "toolApproval", ta) resp, err := ai.Generate(ctx, r, ai.WithModel(m), @@ -164,7 +164,7 @@ func TestToolApprovalEmptyListInterruptsAll(t *testing.T) { myTool := defineTool(t, r, "myTool") ta := &ToolApproval{} - ai.DefineMiddleware(r, "toolApproval", ta) + registerTestMiddleware(r, "toolApproval", ta) resp, err := ai.Generate(ctx, r, ai.WithModel(m), @@ -204,7 +204,7 @@ func TestToolApprovalResumedCallRuns(t *testing.T) { needsApproval := defineTool(t, r, "needsApproval") ta := &ToolApproval{} // deny all - ai.DefineMiddleware(r, "toolApproval", ta) + registerTestMiddleware(r, "toolApproval", ta) resp, err := ai.Generate(ctx, r, ai.WithModel(m), @@ -268,7 +268,7 @@ func TestToolApprovalResumedWithoutApprovalInterrupts(t *testing.T) { needsApproval := defineTool(t, r, "needsApproval") ta := &ToolApproval{} - ai.DefineMiddleware(r, "toolApproval", ta) + registerTestMiddleware(r, "toolApproval", ta) resp, err := ai.Generate(ctx, r, ai.WithModel(m), From 3a7fe0afa436cf14e2e10277df45f20f52c660ca Mon Sep 17 00:00:00 2001 From: Alex Pascal Date: Thu, 30 Jul 2026 10:31:22 -0700 Subject: [PATCH 05/10] feat(go/ai): return concrete *Action types from the typed constructors Exports the concrete primitive types as ModelAction, EmbedderAction, EvaluatorAction, and BackgroundModelAction, and returns them from the NewTyped* constructors and the genkit.DefineTyped* wrappers instead of the interfaces. A concrete return lets methods be added without the interface-growth breakage problem, satisfies api.Action structurally so plugin Init bodies can append models without type assertions, and matches the existing Tool/ToolDef pattern. The types embed core.Action through an unexported alias so the action methods promote without exporting the embedded field itself. The interfaces stay as the accepted-argument and lookup types, and the deprecated constructors keep their interface returns. --- go/ai/background_model.go | 27 +++++++++++++++++------- go/ai/config_test.go | 7 +++---- go/ai/embedder.go | 32 +++++++++++++++++----------- go/ai/embedder_test.go | 2 +- go/ai/evaluator.go | 42 ++++++++++++++++++++++--------------- go/ai/generate.go | 44 ++++++++++++++++++++++++++------------- go/genkit/genkit.go | 10 ++++----- 7 files changed, 102 insertions(+), 62 deletions(-) diff --git a/go/ai/background_model.go b/go/ai/background_model.go index c982e88d55..dad98626ff 100644 --- a/go/ai/background_model.go +++ b/go/ai/background_model.go @@ -41,11 +41,22 @@ type BackgroundModel interface { SupportsCancel() bool } -// backgroundModel is the concrete implementation of BackgroundModel interface. -type backgroundModel struct { - core.BackgroundAction[*ModelRequest, *ModelResponse] +// backgroundAction is an unexported alias of [core.BackgroundAction] used as +// the embedded field in [BackgroundModelAction]; see the action alias in +// generate.go for why. +type backgroundAction[In, Out any] = core.BackgroundAction[In, Out] + +// BackgroundModelAction is a background model backed by registry actions. It +// is the concrete type returned by [NewTypedBackgroundModel]; return it from +// a plugin's Init for the framework to register. +type BackgroundModelAction struct { + backgroundAction[*ModelRequest, *ModelResponse] } +// BackgroundModelAction can be passed anywhere a [BackgroundModel] is +// accepted. +var _ BackgroundModel = (*BackgroundModelAction)(nil) + // ModelOperation is a background operation for a model. type ModelOperation = core.Operation[*ModelResponse] @@ -79,16 +90,16 @@ func LookupBackgroundModel(r api.Registry, name string) BackgroundModel { if action == nil { return nil } - return &backgroundModel{*action} + return &BackgroundModelAction{*action} } -// NewTypedBackgroundModel creates a new [BackgroundModel]. Register it -// with [BackgroundModel.Register] to make it resolvable by name. +// NewTypedBackgroundModel creates a new [BackgroundModelAction]. Register it +// with [BackgroundModelAction.Register] to make it resolvable by name. // // Config is the model's typed configuration; it is usually inferred from // startFn's signature. See [NewTypedModel] for how the request's config // is deserialized. -func NewTypedBackgroundModel[Config any](name string, opts *BackgroundModelOptions, startFn TypedStartModelOpFunc[Config], checkFn CheckModelOpFunc) BackgroundModel { +func NewTypedBackgroundModel[Config any](name string, opts *BackgroundModelOptions, startFn TypedStartModelOpFunc[Config], checkFn CheckModelOpFunc) *BackgroundModelAction { if name == "" { panic("ai.NewTypedBackgroundModel: name is required") } @@ -161,7 +172,7 @@ func NewTypedBackgroundModel[Config any](name string, opts *BackgroundModelOptio return modelOpFromResponse(resp) } - return &backgroundModel{*core.NewBackgroundActionOf(api.ActionTypeBackgroundModel, name, &core.BackgroundActionOptions{ + return &BackgroundModelAction{*core.NewBackgroundActionOf(api.ActionTypeBackgroundModel, name, &core.BackgroundActionOptions{ Metadata: metadata, InputSchema: inputSchema, }, wrappedFn, checkFn, opts.Cancel)} diff --git a/go/ai/config_test.go b/go/ai/config_test.go index abf1ba05a3..c9c71c68e7 100644 --- a/go/ai/config_test.go +++ b/go/ai/config_test.go @@ -21,7 +21,6 @@ import ( "strings" "testing" - "github.com/firebase/genkit/go/core/api" "github.com/firebase/genkit/go/internal/registry" ) @@ -156,9 +155,9 @@ func TestModelConfigSchemaInference(t *testing.T) { } } - configSchemaOf := func(t *testing.T, m Model) any { + configSchemaOf := func(t *testing.T, m *ModelAction) any { t.Helper() - desc := m.(api.Action).Desc() + desc := m.Desc() modelMeta, ok := desc.Metadata["model"].(map[string]any) if !ok { t.Fatalf("missing model metadata: %+v", desc.Metadata) @@ -207,7 +206,7 @@ func TestModelConfigSchemaInference(t *testing.T) { m := NewModel("test/legacy-schema", nil, func(ctx context.Context, req *ModelRequest, cb ModelStreamCallback) (*ModelResponse, error) { return &ModelResponse{Message: NewModelTextMessage("ok"), Request: req}, nil }) - if s, _ := configSchemaOf(t, m).(map[string]any); len(s) != 0 { + if s, _ := configSchemaOf(t, m.(*ModelAction)).(map[string]any); len(s) != 0 { t.Errorf("customOptions = %v, want no inferred schema for legacy models", s) } }) diff --git a/go/ai/embedder.go b/go/ai/embedder.go index cc7c9281ec..5b79d63b67 100644 --- a/go/ai/embedder.go +++ b/go/ai/embedder.go @@ -88,13 +88,23 @@ type EmbedderOptions struct { Dimensions int `json:"dimensions,omitempty"` } -// embedder is an action with functions specific to converting documents to multidimensional vectors such as Embed(). -type embedder struct { - core.Action[*EmbedRequest, *EmbedResponse, struct{}] +// EmbedderAction is an embedder backed by a registry action. It is the +// concrete type returned by [NewTypedEmbedder]; pass it to [WithEmbedder] to +// use it for embedding, or return it from a plugin's Init for the framework +// to register. +type EmbedderAction struct { + action[*EmbedRequest, *EmbedResponse, struct{}] } -// NewTypedEmbedder creates a new [Embedder]. Register it with -// [Embedder.Register] to make it resolvable by name. +// EmbedderAction is a full registry action and can be passed anywhere an +// [api.Action] is accepted as well as anywhere an [Embedder] is accepted. +var ( + _ api.Action = (*EmbedderAction)(nil) + _ Embedder = (*EmbedderAction)(nil) +) + +// NewTypedEmbedder creates a new [EmbedderAction]. Register it with +// [EmbedderAction.Register] to make it resolvable by name. // // Config is the embedder's typed configuration; it is usually inferred from // fn's signature. The framework deserializes the request's raw options into @@ -104,7 +114,7 @@ type embedder struct { // normalized to the converted value, so it always matches the typed // parameter. The config's JSON schema is inferred from Config unless // [EmbedderOptions.ConfigSchema] overrides it. -func NewTypedEmbedder[Config any](name string, opts *EmbedderOptions, fn TypedEmbedderFunc[Config]) Embedder { +func NewTypedEmbedder[Config any](name string, opts *EmbedderOptions, fn TypedEmbedderFunc[Config]) *EmbedderAction { if name == "" { panic("ai.NewTypedEmbedder: name is required") } @@ -147,8 +157,8 @@ func NewTypedEmbedder[Config any](name string, opts *EmbedderOptions, fn TypedEm return fn(ctx, req, cfg) } - return &embedder{ - Action: *core.NewActionOf(api.ActionTypeEmbedder, name, &core.ActionOptions{ + return &EmbedderAction{ + action: *core.NewActionOf(api.ActionTypeEmbedder, name, &core.ActionOptions{ Metadata: metadata, InputSchema: inputSchema, }, rawFn), @@ -176,13 +186,11 @@ func LookupEmbedder(r api.Registry, name string) Embedder { if action == nil { return nil } - return &embedder{ - Action: *action, - } + return &EmbedderAction{*action} } // Embed runs the given [Embedder]. -func (e *embedder) Embed(ctx context.Context, req *EmbedRequest) (*EmbedResponse, error) { +func (e *EmbedderAction) Embed(ctx context.Context, req *EmbedRequest) (*EmbedResponse, error) { if e == nil { return nil, core.NewError(core.INVALID_ARGUMENT, "Embedder.Embed: embedder called on a nil embedder; check that all embedders are defined") } diff --git a/go/ai/embedder_test.go b/go/ai/embedder_test.go index c7a4e025a6..b78803fa2a 100644 --- a/go/ai/embedder_test.go +++ b/go/ai/embedder_test.go @@ -199,7 +199,7 @@ func TestEmbedderEmbed(t *testing.T) { }) t.Run("returns error on nil embedder", func(t *testing.T) { - var e *embedder + var e *EmbedderAction _, err := e.Embed(context.Background(), &EmbedRequest{}) if err == nil { t.Error("expected error for nil embedder") diff --git a/go/ai/evaluator.go b/go/ai/evaluator.go index 64340457e9..1e36ef6983 100644 --- a/go/ai/evaluator.go +++ b/go/ai/evaluator.go @@ -81,11 +81,21 @@ func (e EvaluatorRef) Config() any { return e.config } -// evaluator is an action with functions specific to evaluating a dataset. -type evaluator struct { - core.Action[*EvaluatorRequest, *EvaluatorResponse, struct{}] +// EvaluatorAction is an evaluator backed by a registry action. It is the +// concrete type returned by [NewTypedEvaluator] and [NewTypedBatchEvaluator]; +// pass it to [WithEvaluator], or return it from a plugin's Init for the +// framework to register. +type EvaluatorAction struct { + action[*EvaluatorRequest, *EvaluatorResponse, struct{}] } +// EvaluatorAction is a full registry action and can be passed anywhere an +// [api.Action] is accepted as well as anywhere an [Evaluator] is accepted. +var ( + _ api.Action = (*EvaluatorAction)(nil) + _ Evaluator = (*EvaluatorAction)(nil) +) + // Example is a single example that requires evaluation type Example struct { TestCaseId string `json:"testCaseId,omitempty"` @@ -185,8 +195,8 @@ func evaluatorMetadata(opts *EvaluatorOptions) map[string]any { } } -// NewTypedEvaluator creates a new [Evaluator]. Register it with -// [Evaluator.Register] to make it resolvable by name. +// NewTypedEvaluator creates a new [EvaluatorAction]. Register it with +// [EvaluatorAction.Register] to make it resolvable by name. // This method processes the input dataset one-by-one. // // Config is the evaluator's typed configuration; it is usually inferred from @@ -195,7 +205,7 @@ func evaluatorMetadata(opts *EvaluatorOptions) map[string]any { // map[string]any (from the Dev UI and other JSON callers) are accepted, and // mismatched types are rejected. The config's JSON schema is inferred from // Config unless [EvaluatorOptions.ConfigSchema] overrides it. -func NewTypedEvaluator[Config any](name string, opts *EvaluatorOptions, fn TypedEvaluatorFunc[Config]) Evaluator { +func NewTypedEvaluator[Config any](name string, opts *EvaluatorOptions, fn TypedEvaluatorFunc[Config]) *EvaluatorAction { if name == "" { panic("ai.NewTypedEvaluator: evaluator name is required") } @@ -206,8 +216,8 @@ func NewTypedEvaluator[Config any](name string, opts *EvaluatorOptions, fn Typed _, inputSchema := actionConfigSchemas[Config](opts.ConfigSchema, EvaluatorRequest{}, "options") - return &evaluator{ - Action: *core.NewActionOf(api.ActionTypeEvaluator, name, &core.ActionOptions{ + return &EvaluatorAction{ + action: *core.NewActionOf(api.ActionTypeEvaluator, name, &core.ActionOptions{ Metadata: evaluatorMetadata(opts), InputSchema: inputSchema, }, func(ctx context.Context, req *EvaluatorRequest) (output *EvaluatorResponse, err error) { @@ -273,15 +283,15 @@ func NewTypedEvaluator[Config any](name string, opts *EvaluatorOptions, fn Typed } } -// NewTypedBatchEvaluator creates a new [Evaluator]. Register it with -// [Evaluator.Register] to make it resolvable by name. +// NewTypedBatchEvaluator creates a new [EvaluatorAction]. Register it with +// [EvaluatorAction.Register] to make it resolvable by name. // This method provides the full [EvaluatorRequest] to the callback function, // giving more flexibility to the user for processing the data, such as batching or parallelization. // // Config is the evaluator's typed configuration; it is usually inferred from // fn's signature. See [NewTypedEvaluator] for how the request's options // are deserialized. -func NewTypedBatchEvaluator[Config any](name string, opts *EvaluatorOptions, fn TypedBatchEvaluatorFunc[Config]) Evaluator { +func NewTypedBatchEvaluator[Config any](name string, opts *EvaluatorOptions, fn TypedBatchEvaluatorFunc[Config]) *EvaluatorAction { if name == "" { panic("ai.NewTypedBatchEvaluator: batch evaluator name is required") } @@ -303,8 +313,8 @@ func NewTypedBatchEvaluator[Config any](name string, opts *EvaluatorOptions, fn return fn(ctx, req, cfg) } - return &evaluator{ - Action: *core.NewActionOf(api.ActionTypeEvaluator, name, &core.ActionOptions{ + return &EvaluatorAction{ + action: *core.NewActionOf(api.ActionTypeEvaluator, name, &core.ActionOptions{ Metadata: evaluatorMetadata(opts), InputSchema: inputSchema, }, rawFn), @@ -349,13 +359,11 @@ func LookupEvaluator(r api.Registry, name string) Evaluator { if action == nil { return nil } - return &evaluator{ - Action: *action, - } + return &EvaluatorAction{*action} } // Evaluate runs the given [Evaluator]. -func (e *evaluator) Evaluate(ctx context.Context, req *EvaluatorRequest) (*EvaluatorResponse, error) { +func (e *EvaluatorAction) Evaluate(ctx context.Context, req *EvaluatorRequest) (*EvaluatorResponse, error) { if e == nil { return nil, core.NewError(core.INVALID_ARGUMENT, "Evaluator.Evaluate: evaluator called on a nil evaluator; check that all evaluators are defined") } diff --git a/go/ai/generate.go b/go/ai/generate.go index dfdc0ccbf7..641a78a105 100644 --- a/go/ai/generate.go +++ b/go/ai/generate.go @@ -82,11 +82,27 @@ type ModelStreamCallback = func(context.Context, *ModelResponseChunk) error // Deprecated: Use [Middleware] interface with [WithUse] instead, which supports Generate, Model, and Tool hooks. type ModelMiddleware = core.Middleware[*ModelRequest, *ModelResponse, *ModelResponseChunk] -// model is an action with functions specific to model generation such as Generate(). -type model struct { - core.Action[*ModelRequest, *ModelResponse, *ModelResponseChunk] +// action is an unexported alias of [core.Action] used as the embedded field +// in the ai primitives (ModelAction, EmbedderAction, EvaluatorAction). +// Embedding via the alias promotes Action's methods without exporting the +// field itself, so the containment stays an internal detail of each primitive. +type action[In, Out, Stream any] = core.Action[In, Out, Stream] + +// ModelAction is a generative model backed by a registry action. It is the +// concrete type returned by [NewTypedModel]; pass it to [WithModel] to use it +// for generation, or return it from a plugin's Init for the framework to +// register. +type ModelAction struct { + action[*ModelRequest, *ModelResponse, *ModelResponseChunk] } +// ModelAction is a full registry action and can be passed anywhere an +// [api.Action] is accepted as well as anywhere a [Model] is accepted. +var ( + _ api.Action = (*ModelAction)(nil) + _ Model = (*ModelAction)(nil) +) + // generateAction is the type for a utility model generation action that takes in a GenerateActionOptions instead of a ModelRequest. type generateAction = core.Action[*GenerateActionOptions, *ModelResponse, *ModelResponseChunk] @@ -140,8 +156,8 @@ func DefineGenerateAction(ctx context.Context, r api.Registry) *generateAction { return (*generateAction)(a) } -// NewTypedModel creates a new [Model]. Register it with -// [Model.Register] to make it resolvable by name. +// NewTypedModel creates a new [ModelAction]. Register it with +// [ModelAction.Register] to make it resolvable by name. // // Config is the model's typed configuration; it is usually inferred from fn's // signature. The framework deserializes the request's raw config into Config @@ -157,7 +173,7 @@ func DefineGenerateAction(ctx context.Context, r api.Registry) *generateAction { // wrapper types like Opt[float64] that marshal to primitives but reflect as // objects), set [ModelOptions.ConfigSchema] explicitly or requests will be // rejected at the action boundary. -func NewTypedModel[Config any](name string, opts *ModelOptions, fn TypedModelFunc[Config]) Model { +func NewTypedModel[Config any](name string, opts *ModelOptions, fn TypedModelFunc[Config]) *ModelAction { if name == "" { panic("ai.NewTypedModel: name is required") } @@ -215,7 +231,7 @@ func NewTypedModel[Config any](name string, opts *ModelOptions, fn TypedModelFun addAutomaticTelemetry(), )(typedFn) - return &model{*core.NewStreamingActionOf(api.ActionTypeModel, name, &core.ActionOptions{ + return &ModelAction{*core.NewStreamingActionOf(api.ActionTypeModel, name, &core.ActionOptions{ Metadata: metadata, InputSchema: inputSchema, }, rawFn)} @@ -242,9 +258,7 @@ func LookupModel(r api.Registry, name string) Model { if action == nil { return nil } - return &model{ - Action: *action, - } + return &ModelAction{*action} } // GenerateWithRequest is the central generation implementation for ai.Generate(), prompt.Execute(), and the GenerateAction direct call. @@ -337,7 +351,7 @@ func GenerateWithRequest(ctx context.Context, r api.Registry, opts *GenerateActi // Native constrained output is enabled only when the user has // requested it, the model supports it, and there's a JSON schema. outputCfg.Constrained = opts.Output.JsonSchema != nil && - opts.Output.Constrained && outputCfg.Constrained && m != nil && m.(*model).supportsConstrained(len(toolDefs) > 0) + opts.Output.Constrained && outputCfg.Constrained && m != nil && m.(*ModelAction).supportsConstrained(len(toolDefs) > 0) // Add schema instructions to prompt when not using native constraints. // This is a no-op for unstructured output requests. @@ -913,21 +927,21 @@ func GenerateDataStream[Out any](ctx context.Context, r api.Registry, opts ...Ge } // Generate applies the [Action] to provided request. -func (m *model) Generate(ctx context.Context, req *ModelRequest, cb ModelStreamCallback) (*ModelResponse, error) { +func (m *ModelAction) Generate(ctx context.Context, req *ModelRequest, cb ModelStreamCallback) (*ModelResponse, error) { if m == nil { return nil, core.NewError(core.INVALID_ARGUMENT, "Model.Generate: generate called on a nil model; check that all models are defined") } - return m.Action.Run(ctx, req, cb) + return m.Run(ctx, req, cb) } // supportsConstrained returns whether the model supports constrained output. -func (m *model) supportsConstrained(hasTools bool) bool { +func (m *ModelAction) supportsConstrained(hasTools bool) bool { if m == nil { return false } - metadata := m.Action.Desc().Metadata + metadata := m.Desc().Metadata if metadata == nil { return false } diff --git a/go/genkit/genkit.go b/go/genkit/genkit.go index a909da0fb2..6df6da8a0f 100644 --- a/go/genkit/genkit.go +++ b/go/genkit/genkit.go @@ -559,7 +559,7 @@ func ListTools(g *Genkit) []ai.Tool { // return resp, nil // }, // ) -func DefineTypedModel[Config any](g *Genkit, name string, opts *ai.ModelOptions, fn ai.TypedModelFunc[Config]) ai.Model { +func DefineTypedModel[Config any](g *Genkit, name string, opts *ai.ModelOptions, fn ai.TypedModelFunc[Config]) *ai.ModelAction { m := ai.NewTypedModel(name, opts, fn) m.Register(g.reg) return m @@ -586,7 +586,7 @@ func DefineModel(g *Genkit, name string, opts *ai.ModelOptions, fn ai.ModelFunc) // Config is the model's typed configuration; it is usually inferred from // startFn's signature. See [ai.NewTypedModel] for how the request's // config is deserialized and validated. -func DefineTypedBackgroundModel[Config any](g *Genkit, name string, opts *ai.BackgroundModelOptions, startFn ai.TypedStartModelOpFunc[Config], checkFn ai.CheckModelOpFunc) ai.BackgroundModel { +func DefineTypedBackgroundModel[Config any](g *Genkit, name string, opts *ai.BackgroundModelOptions, startFn ai.TypedStartModelOpFunc[Config], checkFn ai.CheckModelOpFunc) *ai.BackgroundModelAction { m := ai.NewTypedBackgroundModel(name, opts, startFn, checkFn) m.Register(g.reg) return m @@ -1397,7 +1397,7 @@ func LookupRetriever(g *Genkit, name string) ai.Retriever { // // For embedders that don't need to be registered (e.g., for plugin development), // use [ai.NewTypedEmbedder] instead. -func DefineTypedEmbedder[Config any](g *Genkit, name string, opts *ai.EmbedderOptions, fn ai.TypedEmbedderFunc[Config]) ai.Embedder { +func DefineTypedEmbedder[Config any](g *Genkit, name string, opts *ai.EmbedderOptions, fn ai.TypedEmbedderFunc[Config]) *ai.EmbedderAction { e := ai.NewTypedEmbedder(name, opts, fn) e.Register(g.reg) return e @@ -1443,7 +1443,7 @@ func LookupPlugin(g *Genkit, name string) api.Plugin { // Config is the evaluator's typed configuration; it is usually inferred from // fn's signature. See [ai.NewTypedEvaluator] for how the request's // options are deserialized. -func DefineTypedEvaluator[Config any](g *Genkit, name string, opts *ai.EvaluatorOptions, fn ai.TypedEvaluatorFunc[Config]) ai.Evaluator { +func DefineTypedEvaluator[Config any](g *Genkit, name string, opts *ai.EvaluatorOptions, fn ai.TypedEvaluatorFunc[Config]) *ai.EvaluatorAction { e := ai.NewTypedEvaluator(name, opts, fn) e.Register(g.reg) return e @@ -1472,7 +1472,7 @@ func DefineEvaluator(g *Genkit, name string, opts *ai.EvaluatorOptions, fn ai.Ev // Config is the evaluator's typed configuration; it is usually inferred from // fn's signature. See [ai.NewTypedEvaluator] for how the request's // options are deserialized. -func DefineTypedBatchEvaluator[Config any](g *Genkit, name string, opts *ai.EvaluatorOptions, fn ai.TypedBatchEvaluatorFunc[Config]) ai.Evaluator { +func DefineTypedBatchEvaluator[Config any](g *Genkit, name string, opts *ai.EvaluatorOptions, fn ai.TypedBatchEvaluatorFunc[Config]) *ai.EvaluatorAction { e := ai.NewTypedBatchEvaluator(name, opts, fn) e.Register(g.reg) return e From f23f654ce7bf26b0bd1c6caf7c4d01a636fa1f34 Mon Sep 17 00:00:00 2001 From: Alex Pascal Date: Thu, 30 Jul 2026 10:37:46 -0700 Subject: [PATCH 06/10] refactor(go/ai): take the background cancel function positionally NewTypedBackgroundModel and genkit.DefineTypedBackgroundModel gain a positional nil-able cancelFn, matching core.NewBackgroundActionOf and the invariant that options structs hold descriptor data only. The shipped BackgroundModelOptions.Cancel field is deprecated and kept as a fallback when cancelFn is nil, so the deprecated constructors and existing option-based callers behave unchanged. Also asserts api.Action on BackgroundModelAction: the plugin path registers Init-returned actions via their Register method, so the overriding Register handles the check and cancel sub-actions correctly. --- go/ai/background_model.go | 33 +++++++++++++++++++++++++-------- go/ai/config_test.go | 3 ++- go/genkit/genkit.go | 10 ++++++---- 3 files changed, 33 insertions(+), 13 deletions(-) diff --git a/go/ai/background_model.go b/go/ai/background_model.go index dad98626ff..ff6c99300a 100644 --- a/go/ai/background_model.go +++ b/go/ai/background_model.go @@ -53,9 +53,14 @@ type BackgroundModelAction struct { backgroundAction[*ModelRequest, *ModelResponse] } -// BackgroundModelAction can be passed anywhere a [BackgroundModel] is -// accepted. -var _ BackgroundModel = (*BackgroundModelAction)(nil) +// BackgroundModelAction is a full registry action and can be passed anywhere +// an [api.Action] is accepted as well as anywhere a [BackgroundModel] is +// accepted. Its Register method registers the start, check, and cancel +// actions together. +var ( + _ api.Action = (*BackgroundModelAction)(nil) + _ BackgroundModel = (*BackgroundModelAction)(nil) +) // ModelOperation is a background operation for a model. type ModelOperation = core.Operation[*ModelResponse] @@ -78,8 +83,12 @@ type CancelModelOpFunc = func(ctx context.Context, op *ModelOperation) (*ModelOp // 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. + // Cancel is the function that cancels a background model operation. + // + // Deprecated: Pass cancelFn to [NewTypedBackgroundModel] instead; + // options hold descriptor data only. + Cancel CancelModelOpFunc + Metadata map[string]any // Additional metadata. } // LookupBackgroundModel looks up a registered [BackgroundModel] by name. @@ -96,10 +105,14 @@ func LookupBackgroundModel(r api.Registry, name string) BackgroundModel { // NewTypedBackgroundModel creates a new [BackgroundModelAction]. Register it // with [BackgroundModelAction.Register] to make it resolvable by name. // +// cancelFn is optional; nil means the model does not support canceling +// operations. When cancelFn is nil, the deprecated +// [BackgroundModelOptions.Cancel] is used as a fallback. +// // Config is the model's typed configuration; it is usually inferred from // startFn's signature. See [NewTypedModel] for how the request's config // is deserialized. -func NewTypedBackgroundModel[Config any](name string, opts *BackgroundModelOptions, startFn TypedStartModelOpFunc[Config], checkFn CheckModelOpFunc) *BackgroundModelAction { +func NewTypedBackgroundModel[Config any](name string, opts *BackgroundModelOptions, startFn TypedStartModelOpFunc[Config], checkFn CheckModelOpFunc, cancelFn CancelModelOpFunc) *BackgroundModelAction { if name == "" { panic("ai.NewTypedBackgroundModel: name is required") } @@ -172,10 +185,14 @@ func NewTypedBackgroundModel[Config any](name string, opts *BackgroundModelOptio return modelOpFromResponse(resp) } + if cancelFn == nil { + cancelFn = opts.Cancel + } + return &BackgroundModelAction{*core.NewBackgroundActionOf(api.ActionTypeBackgroundModel, name, &core.BackgroundActionOptions{ Metadata: metadata, InputSchema: inputSchema, - }, wrappedFn, checkFn, opts.Cancel)} + }, wrappedFn, checkFn, cancelFn)} } // NewBackgroundModel defines a new model that runs in the background. @@ -189,7 +206,7 @@ func NewBackgroundModel(name string, opts *BackgroundModelOptions, startFn Start } return NewTypedBackgroundModel(name, opts, func(ctx context.Context, req *ModelRequest, _ any) (*ModelOperation, error) { return startFn(ctx, req) - }, checkFn) + }, checkFn, nil) } // GenerateOperation generates a model response as a long-running operation based on the provided options. diff --git a/go/ai/config_test.go b/go/ai/config_test.go index c9c71c68e7..8f614ccb0e 100644 --- a/go/ai/config_test.go +++ b/go/ai/config_test.go @@ -125,7 +125,8 @@ func TestBackgroundModelConfigValidation(t *testing.T) { }, func(ctx context.Context, op *ModelOperation) (*ModelOperation, error) { return op, nil - }) + }, + nil) bm.Register(r) // A config that violates the inferred schema is rejected by input diff --git a/go/genkit/genkit.go b/go/genkit/genkit.go index 6df6da8a0f..954fc7e6bb 100644 --- a/go/genkit/genkit.go +++ b/go/genkit/genkit.go @@ -576,18 +576,20 @@ func DefineModel(g *Genkit, name string, opts *ai.ModelOptions, fn ai.ModelFunc) return m } -// DefineTypedBackgroundModel defines a background model, registers it as -// an [ai.BackgroundModel], and returns an [ai.BackgroundModel]. +// DefineTypedBackgroundModel defines a background model, registers it, and +// returns the concrete [ai.BackgroundModelAction]. // // The `name` is the identifier the model uses to request the background model. The `opts` // are the options for the background model. The `startFn` is the function that starts the background model. // The `checkFn` is the function that checks the status of the background model. +// The `cancelFn` is optional; nil means the model does not support canceling +// operations. // // Config is the model's typed configuration; it is usually inferred from // startFn's signature. See [ai.NewTypedModel] for how the request's // config is deserialized and validated. -func DefineTypedBackgroundModel[Config any](g *Genkit, name string, opts *ai.BackgroundModelOptions, startFn ai.TypedStartModelOpFunc[Config], checkFn ai.CheckModelOpFunc) *ai.BackgroundModelAction { - m := ai.NewTypedBackgroundModel(name, opts, startFn, checkFn) +func DefineTypedBackgroundModel[Config any](g *Genkit, name string, opts *ai.BackgroundModelOptions, startFn ai.TypedStartModelOpFunc[Config], checkFn ai.CheckModelOpFunc, cancelFn ai.CancelModelOpFunc) *ai.BackgroundModelAction { + m := ai.NewTypedBackgroundModel(name, opts, startFn, checkFn, cancelFn) m.Register(g.reg) return m } From 523faa4d5d8b0636f6f68c598e6def5ab856aa31 Mon Sep 17 00:00:00 2001 From: Alex Pascal Date: Thu, 30 Jul 2026 10:43:20 -0700 Subject: [PATCH 07/10] style(go): break the typed constructor signatures across lines One parameter per line on the NewTyped* and DefineTyped* family so the signatures render readably on pkg.go.dev, matching the formatting of core.NewBackgroundActionOf. --- go/ai/background_model.go | 8 +++++++- go/ai/embedder.go | 6 +++++- go/ai/evaluator.go | 12 ++++++++++-- go/ai/generate.go | 6 +++++- go/genkit/genkit.go | 37 ++++++++++++++++++++++++++++++++----- 5 files changed, 59 insertions(+), 10 deletions(-) diff --git a/go/ai/background_model.go b/go/ai/background_model.go index ff6c99300a..d02bf16349 100644 --- a/go/ai/background_model.go +++ b/go/ai/background_model.go @@ -112,7 +112,13 @@ func LookupBackgroundModel(r api.Registry, name string) BackgroundModel { // Config is the model's typed configuration; it is usually inferred from // startFn's signature. See [NewTypedModel] for how the request's config // is deserialized. -func NewTypedBackgroundModel[Config any](name string, opts *BackgroundModelOptions, startFn TypedStartModelOpFunc[Config], checkFn CheckModelOpFunc, cancelFn CancelModelOpFunc) *BackgroundModelAction { +func NewTypedBackgroundModel[Config any]( + name string, + opts *BackgroundModelOptions, + startFn TypedStartModelOpFunc[Config], + checkFn CheckModelOpFunc, + cancelFn CancelModelOpFunc, +) *BackgroundModelAction { if name == "" { panic("ai.NewTypedBackgroundModel: name is required") } diff --git a/go/ai/embedder.go b/go/ai/embedder.go index 5b79d63b67..53ecd8b8c6 100644 --- a/go/ai/embedder.go +++ b/go/ai/embedder.go @@ -114,7 +114,11 @@ var ( // normalized to the converted value, so it always matches the typed // parameter. The config's JSON schema is inferred from Config unless // [EmbedderOptions.ConfigSchema] overrides it. -func NewTypedEmbedder[Config any](name string, opts *EmbedderOptions, fn TypedEmbedderFunc[Config]) *EmbedderAction { +func NewTypedEmbedder[Config any]( + name string, + opts *EmbedderOptions, + fn TypedEmbedderFunc[Config], +) *EmbedderAction { if name == "" { panic("ai.NewTypedEmbedder: name is required") } diff --git a/go/ai/evaluator.go b/go/ai/evaluator.go index 1e36ef6983..84a998a0cd 100644 --- a/go/ai/evaluator.go +++ b/go/ai/evaluator.go @@ -205,7 +205,11 @@ func evaluatorMetadata(opts *EvaluatorOptions) map[string]any { // map[string]any (from the Dev UI and other JSON callers) are accepted, and // mismatched types are rejected. The config's JSON schema is inferred from // Config unless [EvaluatorOptions.ConfigSchema] overrides it. -func NewTypedEvaluator[Config any](name string, opts *EvaluatorOptions, fn TypedEvaluatorFunc[Config]) *EvaluatorAction { +func NewTypedEvaluator[Config any]( + name string, + opts *EvaluatorOptions, + fn TypedEvaluatorFunc[Config], +) *EvaluatorAction { if name == "" { panic("ai.NewTypedEvaluator: evaluator name is required") } @@ -291,7 +295,11 @@ func NewTypedEvaluator[Config any](name string, opts *EvaluatorOptions, fn Typed // Config is the evaluator's typed configuration; it is usually inferred from // fn's signature. See [NewTypedEvaluator] for how the request's options // are deserialized. -func NewTypedBatchEvaluator[Config any](name string, opts *EvaluatorOptions, fn TypedBatchEvaluatorFunc[Config]) *EvaluatorAction { +func NewTypedBatchEvaluator[Config any]( + name string, + opts *EvaluatorOptions, + fn TypedBatchEvaluatorFunc[Config], +) *EvaluatorAction { if name == "" { panic("ai.NewTypedBatchEvaluator: batch evaluator name is required") } diff --git a/go/ai/generate.go b/go/ai/generate.go index 641a78a105..2d995874f0 100644 --- a/go/ai/generate.go +++ b/go/ai/generate.go @@ -173,7 +173,11 @@ func DefineGenerateAction(ctx context.Context, r api.Registry) *generateAction { // wrapper types like Opt[float64] that marshal to primitives but reflect as // objects), set [ModelOptions.ConfigSchema] explicitly or requests will be // rejected at the action boundary. -func NewTypedModel[Config any](name string, opts *ModelOptions, fn TypedModelFunc[Config]) *ModelAction { +func NewTypedModel[Config any]( + name string, + opts *ModelOptions, + fn TypedModelFunc[Config], +) *ModelAction { if name == "" { panic("ai.NewTypedModel: name is required") } diff --git a/go/genkit/genkit.go b/go/genkit/genkit.go index 954fc7e6bb..557e5659de 100644 --- a/go/genkit/genkit.go +++ b/go/genkit/genkit.go @@ -559,7 +559,12 @@ func ListTools(g *Genkit) []ai.Tool { // return resp, nil // }, // ) -func DefineTypedModel[Config any](g *Genkit, name string, opts *ai.ModelOptions, fn ai.TypedModelFunc[Config]) *ai.ModelAction { +func DefineTypedModel[Config any]( + g *Genkit, + name string, + opts *ai.ModelOptions, + fn ai.TypedModelFunc[Config], +) *ai.ModelAction { m := ai.NewTypedModel(name, opts, fn) m.Register(g.reg) return m @@ -588,7 +593,14 @@ func DefineModel(g *Genkit, name string, opts *ai.ModelOptions, fn ai.ModelFunc) // Config is the model's typed configuration; it is usually inferred from // startFn's signature. See [ai.NewTypedModel] for how the request's // config is deserialized and validated. -func DefineTypedBackgroundModel[Config any](g *Genkit, name string, opts *ai.BackgroundModelOptions, startFn ai.TypedStartModelOpFunc[Config], checkFn ai.CheckModelOpFunc, cancelFn ai.CancelModelOpFunc) *ai.BackgroundModelAction { +func DefineTypedBackgroundModel[Config any]( + g *Genkit, + name string, + opts *ai.BackgroundModelOptions, + startFn ai.TypedStartModelOpFunc[Config], + checkFn ai.CheckModelOpFunc, + cancelFn ai.CancelModelOpFunc, +) *ai.BackgroundModelAction { m := ai.NewTypedBackgroundModel(name, opts, startFn, checkFn, cancelFn) m.Register(g.reg) return m @@ -1399,7 +1411,12 @@ func LookupRetriever(g *Genkit, name string) ai.Retriever { // // For embedders that don't need to be registered (e.g., for plugin development), // use [ai.NewTypedEmbedder] instead. -func DefineTypedEmbedder[Config any](g *Genkit, name string, opts *ai.EmbedderOptions, fn ai.TypedEmbedderFunc[Config]) *ai.EmbedderAction { +func DefineTypedEmbedder[Config any]( + g *Genkit, + name string, + opts *ai.EmbedderOptions, + fn ai.TypedEmbedderFunc[Config], +) *ai.EmbedderAction { e := ai.NewTypedEmbedder(name, opts, fn) e.Register(g.reg) return e @@ -1445,7 +1462,12 @@ func LookupPlugin(g *Genkit, name string) api.Plugin { // Config is the evaluator's typed configuration; it is usually inferred from // fn's signature. See [ai.NewTypedEvaluator] for how the request's // options are deserialized. -func DefineTypedEvaluator[Config any](g *Genkit, name string, opts *ai.EvaluatorOptions, fn ai.TypedEvaluatorFunc[Config]) *ai.EvaluatorAction { +func DefineTypedEvaluator[Config any]( + g *Genkit, + name string, + opts *ai.EvaluatorOptions, + fn ai.TypedEvaluatorFunc[Config], +) *ai.EvaluatorAction { e := ai.NewTypedEvaluator(name, opts, fn) e.Register(g.reg) return e @@ -1474,7 +1496,12 @@ func DefineEvaluator(g *Genkit, name string, opts *ai.EvaluatorOptions, fn ai.Ev // Config is the evaluator's typed configuration; it is usually inferred from // fn's signature. See [ai.NewTypedEvaluator] for how the request's // options are deserialized. -func DefineTypedBatchEvaluator[Config any](g *Genkit, name string, opts *ai.EvaluatorOptions, fn ai.TypedBatchEvaluatorFunc[Config]) *ai.EvaluatorAction { +func DefineTypedBatchEvaluator[Config any]( + g *Genkit, + name string, + opts *ai.EvaluatorOptions, + fn ai.TypedBatchEvaluatorFunc[Config], +) *ai.EvaluatorAction { e := ai.NewTypedBatchEvaluator(name, opts, fn) e.Register(g.reg) return e From ed2a43baa0cafbcb027e6a19e3408436b92468b3 Mon Sep 17 00:00:00 2001 From: Alex Pascal Date: Thu, 30 Jul 2026 10:51:24 -0700 Subject: [PATCH 08/10] feat(go/ai): add a flat TypedBackgroundModelOptions for the typed constructor NewTypedBackgroundModel and genkit.DefineTypedBackgroundModel take a dedicated options struct that inlines the ModelOptions fields instead of embedding them, plus a descriptor Metadata slot. The Metadata field is actually wired into the action descriptor (it was silently ignored before, both here and on main), with the reserved type and model keys winning over caller values. BackgroundModelOptions is deprecated and only feeds the deprecated NewBackgroundModel, which converts and forwards its Cancel field positionally, so the typed path no longer reads implementation functions from options. --- go/ai/background_model.go | 99 +++++++++++++++++++++++++-------------- go/ai/config_test.go | 29 ++++++++++++ go/genkit/genkit.go | 2 +- 3 files changed, 93 insertions(+), 37 deletions(-) diff --git a/go/ai/background_model.go b/go/ai/background_model.go index d02bf16349..1167e7ebeb 100644 --- a/go/ai/background_model.go +++ b/go/ai/background_model.go @@ -19,6 +19,7 @@ package ai import ( "context" "errors" + "maps" "github.com/firebase/genkit/go/core" "github.com/firebase/genkit/go/core/api" @@ -80,13 +81,25 @@ type CheckModelOpFunc = func(ctx context.Context, op *ModelOperation) (*ModelOpe // CancelModelOpFunc cancels a background model operation. type CancelModelOpFunc = func(ctx context.Context, op *ModelOperation) (*ModelOperation, error) -// BackgroundModelOptions holds configuration for defining a background model +// TypedBackgroundModelOptions configures a background model created with +// [NewTypedBackgroundModel]. It holds descriptor data only; the cancel +// function is a constructor argument. +type TypedBackgroundModelOptions struct { + ConfigSchema map[string]any // JSON schema for the model's config. + Label string // User-friendly name for the model. + Stage ModelStage // Indicates the maturity stage of the model. + Supports *ModelSupports // Capabilities of the model. + Versions []string // Available versions of the model. + Metadata map[string]any // Arbitrary key-value data attached to the action descriptor. +} + +// BackgroundModelOptions holds configuration for defining a background model. +// +// Deprecated: Use [TypedBackgroundModelOptions] with +// [NewTypedBackgroundModel]. type BackgroundModelOptions struct { ModelOptions // Cancel is the function that cancels a background model operation. - // - // Deprecated: Pass cancelFn to [NewTypedBackgroundModel] instead; - // options hold descriptor data only. Cancel CancelModelOpFunc Metadata map[string]any // Additional metadata. } @@ -106,15 +119,14 @@ func LookupBackgroundModel(r api.Registry, name string) BackgroundModel { // with [BackgroundModelAction.Register] to make it resolvable by name. // // cancelFn is optional; nil means the model does not support canceling -// operations. When cancelFn is nil, the deprecated -// [BackgroundModelOptions.Cancel] is used as a fallback. +// operations. // // Config is the model's typed configuration; it is usually inferred from // startFn's signature. See [NewTypedModel] for how the request's config // is deserialized. func NewTypedBackgroundModel[Config any]( name string, - opts *BackgroundModelOptions, + opts *TypedBackgroundModelOptions, startFn TypedStartModelOpFunc[Config], checkFn CheckModelOpFunc, cancelFn CancelModelOpFunc, @@ -130,7 +142,7 @@ func NewTypedBackgroundModel[Config any]( } if opts == nil { - opts = &BackgroundModelOptions{} + opts = &TypedBackgroundModelOptions{} } if opts.Label == "" { opts.Label = name @@ -141,26 +153,27 @@ func NewTypedBackgroundModel[Config any]( configSchema, inputSchema := actionConfigSchemas[Config](opts.ConfigSchema, ModelRequest{}, "config") - 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": configSchema, + // The reserved type and model keys win over caller-provided metadata. + metadata := make(map[string]any, len(opts.Metadata)+2) + maps.Copy(metadata, opts.Metadata) + metadata["type"] = api.ActionTypeBackgroundModel + metadata["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": configSchema, } typedStartFn := func(ctx context.Context, req *ModelRequest) (*ModelOperation, error) { @@ -173,13 +186,21 @@ func NewTypedBackgroundModel[Config any]( return startFn(ctx, req, cfg) } + mopts := &ModelOptions{ + ConfigSchema: opts.ConfigSchema, + Label: opts.Label, + Stage: opts.Stage, + Supports: opts.Supports, + Versions: opts.Versions, + } + // normalizeConfig runs outermost so that the built-in wrappers and the // start function all see the typed, converted config on the request. fn := core.ChainMiddleware( normalizeConfig[Config](name, opts.Versions), - simulateSystemPrompt(&opts.ModelOptions, nil), - augmentWithContext(&opts.ModelOptions, nil), - validateSupport(name, &opts.ModelOptions), + simulateSystemPrompt(mopts, nil), + augmentWithContext(mopts, nil), + validateSupport(name, mopts), )(backgroundModelToModelFn(typedStartFn)) wrappedFn := func(ctx context.Context, req *ModelRequest) (*ModelOperation, error) { @@ -191,10 +212,6 @@ func NewTypedBackgroundModel[Config any]( return modelOpFromResponse(resp) } - if cancelFn == nil { - cancelFn = opts.Cancel - } - return &BackgroundModelAction{*core.NewBackgroundActionOf(api.ActionTypeBackgroundModel, name, &core.BackgroundActionOptions{ Metadata: metadata, InputSchema: inputSchema, @@ -210,9 +227,19 @@ func NewBackgroundModel(name string, opts *BackgroundModelOptions, startFn Start if startFn == nil { panic("ai.NewBackgroundModel: startFn is required") } - return NewTypedBackgroundModel(name, opts, func(ctx context.Context, req *ModelRequest, _ any) (*ModelOperation, error) { + if opts == nil { + opts = &BackgroundModelOptions{} + } + return NewTypedBackgroundModel(name, &TypedBackgroundModelOptions{ + ConfigSchema: opts.ConfigSchema, + Label: opts.Label, + Stage: opts.Stage, + Supports: opts.Supports, + Versions: opts.Versions, + Metadata: opts.Metadata, + }, func(ctx context.Context, req *ModelRequest, _ any) (*ModelOperation, error) { return startFn(ctx, req) - }, checkFn, nil) + }, checkFn, opts.Cancel) } // GenerateOperation generates a model response as a long-running operation based on the provided options. diff --git a/go/ai/config_test.go b/go/ai/config_test.go index 8f614ccb0e..a77d4ee499 100644 --- a/go/ai/config_test.go +++ b/go/ai/config_test.go @@ -21,6 +21,7 @@ import ( "strings" "testing" + "github.com/firebase/genkit/go/core/api" "github.com/firebase/genkit/go/internal/registry" ) @@ -116,6 +117,34 @@ func TestModelConfigNormalizedBeforeBuiltins(t *testing.T) { } } +func TestTypedBackgroundModelMetadata(t *testing.T) { + bm := NewTypedBackgroundModel("test/bg-metadata", + &TypedBackgroundModelOptions{ + Metadata: map[string]any{ + "custom": "value", + "model": "caller values must not clobber the reserved keys", + }, + }, + func(ctx context.Context, req *ModelRequest, cfg testTypedConfig) (*ModelOperation, error) { + return &ModelOperation{ID: "op1"}, nil + }, + func(ctx context.Context, op *ModelOperation) (*ModelOperation, error) { + return op, nil + }, + nil) + + metadata := bm.Desc().Metadata + if got := metadata["custom"]; got != "value" { + t.Errorf(`Metadata["custom"] = %v, want "value"`, got) + } + if got := metadata["type"]; got != api.ActionTypeBackgroundModel { + t.Errorf(`Metadata["type"] = %v, want %v`, got, api.ActionTypeBackgroundModel) + } + if _, ok := metadata["model"].(map[string]any); !ok { + t.Errorf(`Metadata["model"] = %v, want the built model info map`, metadata["model"]) + } +} + func TestBackgroundModelConfigValidation(t *testing.T) { r := registry.New() diff --git a/go/genkit/genkit.go b/go/genkit/genkit.go index 557e5659de..a5d0284fa7 100644 --- a/go/genkit/genkit.go +++ b/go/genkit/genkit.go @@ -596,7 +596,7 @@ func DefineModel(g *Genkit, name string, opts *ai.ModelOptions, fn ai.ModelFunc) func DefineTypedBackgroundModel[Config any]( g *Genkit, name string, - opts *ai.BackgroundModelOptions, + opts *ai.TypedBackgroundModelOptions, startFn ai.TypedStartModelOpFunc[Config], checkFn ai.CheckModelOpFunc, cancelFn ai.CancelModelOpFunc, From 072cf035187843987d4c8e1ff64e7ccb769392d2 Mon Sep 17 00:00:00 2001 From: Alex Pascal Date: Thu, 30 Jul 2026 10:59:14 -0700 Subject: [PATCH 09/10] feat(go/ai): add a Metadata slot to ModelOptions Caller-provided metadata merges into the model action's descriptor with the reserved type and model keys winning, matching the background model options. Adding the field is additive; the deprecated constructors pick it up through delegation. --- go/ai/config_test.go | 24 ++++++++++++++++++++++++ go/ai/generate.go | 41 ++++++++++++++++++++++------------------- 2 files changed, 46 insertions(+), 19 deletions(-) diff --git a/go/ai/config_test.go b/go/ai/config_test.go index a77d4ee499..c7f98dc3e1 100644 --- a/go/ai/config_test.go +++ b/go/ai/config_test.go @@ -117,6 +117,30 @@ func TestModelConfigNormalizedBeforeBuiltins(t *testing.T) { } } +func TestTypedModelMetadata(t *testing.T) { + m := NewTypedModel("test/model-metadata", + &ModelOptions{ + Metadata: map[string]any{ + "custom": "value", + "model": "caller values must not clobber the reserved keys", + }, + }, + func(ctx context.Context, req *ModelRequest, cfg testTypedConfig, cb ModelStreamCallback) (*ModelResponse, error) { + return &ModelResponse{Message: NewModelTextMessage("ok"), Request: req}, nil + }) + + metadata := m.Desc().Metadata + if got := metadata["custom"]; got != "value" { + t.Errorf(`Metadata["custom"] = %v, want "value"`, got) + } + if got := metadata["type"]; got != api.ActionTypeModel { + t.Errorf(`Metadata["type"] = %v, want %v`, got, api.ActionTypeModel) + } + if _, ok := metadata["model"].(map[string]any); !ok { + t.Errorf(`Metadata["model"] = %v, want the built model info map`, metadata["model"]) + } +} + func TestTypedBackgroundModelMetadata(t *testing.T) { bm := NewTypedBackgroundModel("test/bg-metadata", &TypedBackgroundModelOptions{ diff --git a/go/ai/generate.go b/go/ai/generate.go index 2d995874f0..b0e99912aa 100644 --- a/go/ai/generate.go +++ b/go/ai/generate.go @@ -22,6 +22,7 @@ import ( "errors" "fmt" "iter" + "maps" "slices" "strings" "sync" @@ -134,6 +135,7 @@ type ModelOptions struct { Stage ModelStage // Indicates the maturity stage of the model. Supports *ModelSupports // Capabilities of the model. Versions []string // Available versions of the model. + Metadata map[string]any // Arbitrary key-value data attached to the action descriptor. } // DefineGenerateAction defines a utility generate action. @@ -193,26 +195,27 @@ func NewTypedModel[Config any]( configSchema, inputSchema := actionConfigSchemas[Config](opts.ConfigSchema, ModelRequest{}, "config") - metadata := map[string]any{ - "type": api.ActionTypeModel, - "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": configSchema, + // The reserved type and model keys win over caller-provided metadata. + metadata := make(map[string]any, len(opts.Metadata)+2) + maps.Copy(metadata, opts.Metadata) + metadata["type"] = api.ActionTypeModel + metadata["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": configSchema, } typedFn := func(ctx context.Context, req *ModelRequest, cb ModelStreamCallback) (*ModelResponse, error) { From d712b81b2467d4adb5acb3f88e102c3f22366cdb Mon Sep 17 00:00:00 2001 From: Alex Pascal Date: Thu, 30 Jul 2026 11:24:52 -0700 Subject: [PATCH 10/10] refactor(go/ai): move the optional cancel hook into TypedBackgroundModelOptions Provider-surface rule: required functions stay positional (they are compile-enforced and drive type inference), while optional lifecycle hooks live in the options struct so future additions cost a field instead of a signature break. Cancel is the first such hook. The core constructor keeps its positional cancel because holding CancelOpFunc[Out] would force BackgroundActionOptions to become generic. --- go/ai/background_model.go | 17 +++++++++-------- go/ai/config_test.go | 6 ++---- go/genkit/genkit.go | 5 +---- 3 files changed, 12 insertions(+), 16 deletions(-) diff --git a/go/ai/background_model.go b/go/ai/background_model.go index 1167e7ebeb..8c6ec24451 100644 --- a/go/ai/background_model.go +++ b/go/ai/background_model.go @@ -82,8 +82,8 @@ type CheckModelOpFunc = func(ctx context.Context, op *ModelOperation) (*ModelOpe type CancelModelOpFunc = func(ctx context.Context, op *ModelOperation) (*ModelOperation, error) // TypedBackgroundModelOptions configures a background model created with -// [NewTypedBackgroundModel]. It holds descriptor data only; the cancel -// function is a constructor argument. +// [NewTypedBackgroundModel]. It holds descriptor data plus optional lifecycle +// hooks; the required start and check functions are constructor arguments. type TypedBackgroundModelOptions struct { ConfigSchema map[string]any // JSON schema for the model's config. Label string // User-friendly name for the model. @@ -91,6 +91,10 @@ type TypedBackgroundModelOptions struct { Supports *ModelSupports // Capabilities of the model. Versions []string // Available versions of the model. Metadata map[string]any // Arbitrary key-value data attached to the action descriptor. + + // Cancel cancels a running operation. Optional: nil means the model does + // not support canceling operations. + Cancel CancelModelOpFunc } // BackgroundModelOptions holds configuration for defining a background model. @@ -118,9 +122,6 @@ func LookupBackgroundModel(r api.Registry, name string) BackgroundModel { // NewTypedBackgroundModel creates a new [BackgroundModelAction]. Register it // with [BackgroundModelAction.Register] to make it resolvable by name. // -// cancelFn is optional; nil means the model does not support canceling -// operations. -// // Config is the model's typed configuration; it is usually inferred from // startFn's signature. See [NewTypedModel] for how the request's config // is deserialized. @@ -129,7 +130,6 @@ func NewTypedBackgroundModel[Config any]( opts *TypedBackgroundModelOptions, startFn TypedStartModelOpFunc[Config], checkFn CheckModelOpFunc, - cancelFn CancelModelOpFunc, ) *BackgroundModelAction { if name == "" { panic("ai.NewTypedBackgroundModel: name is required") @@ -215,7 +215,7 @@ func NewTypedBackgroundModel[Config any]( return &BackgroundModelAction{*core.NewBackgroundActionOf(api.ActionTypeBackgroundModel, name, &core.BackgroundActionOptions{ Metadata: metadata, InputSchema: inputSchema, - }, wrappedFn, checkFn, cancelFn)} + }, wrappedFn, checkFn, opts.Cancel)} } // NewBackgroundModel defines a new model that runs in the background. @@ -237,9 +237,10 @@ func NewBackgroundModel(name string, opts *BackgroundModelOptions, startFn Start Supports: opts.Supports, Versions: opts.Versions, Metadata: opts.Metadata, + Cancel: opts.Cancel, }, func(ctx context.Context, req *ModelRequest, _ any) (*ModelOperation, error) { return startFn(ctx, req) - }, checkFn, opts.Cancel) + }, checkFn) } // GenerateOperation generates a model response as a long-running operation based on the provided options. diff --git a/go/ai/config_test.go b/go/ai/config_test.go index c7f98dc3e1..88a5b76f6f 100644 --- a/go/ai/config_test.go +++ b/go/ai/config_test.go @@ -154,8 +154,7 @@ func TestTypedBackgroundModelMetadata(t *testing.T) { }, func(ctx context.Context, op *ModelOperation) (*ModelOperation, error) { return op, nil - }, - nil) + }) metadata := bm.Desc().Metadata if got := metadata["custom"]; got != "value" { @@ -178,8 +177,7 @@ func TestBackgroundModelConfigValidation(t *testing.T) { }, func(ctx context.Context, op *ModelOperation) (*ModelOperation, error) { return op, nil - }, - nil) + }) bm.Register(r) // A config that violates the inferred schema is rejected by input diff --git a/go/genkit/genkit.go b/go/genkit/genkit.go index a5d0284fa7..cbb6fd48a6 100644 --- a/go/genkit/genkit.go +++ b/go/genkit/genkit.go @@ -587,8 +587,6 @@ func DefineModel(g *Genkit, name string, opts *ai.ModelOptions, fn ai.ModelFunc) // The `name` is the identifier the model uses to request the background model. The `opts` // are the options for the background model. The `startFn` is the function that starts the background model. // The `checkFn` is the function that checks the status of the background model. -// The `cancelFn` is optional; nil means the model does not support canceling -// operations. // // Config is the model's typed configuration; it is usually inferred from // startFn's signature. See [ai.NewTypedModel] for how the request's @@ -599,9 +597,8 @@ func DefineTypedBackgroundModel[Config any]( opts *ai.TypedBackgroundModelOptions, startFn ai.TypedStartModelOpFunc[Config], checkFn ai.CheckModelOpFunc, - cancelFn ai.CancelModelOpFunc, ) *ai.BackgroundModelAction { - m := ai.NewTypedBackgroundModel(name, opts, startFn, checkFn, cancelFn) + m := ai.NewTypedBackgroundModel(name, opts, startFn, checkFn) m.Register(g.reg) return m }