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 21e3817b95..8c6ec24451 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" @@ -41,31 +42,73 @@ 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 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] // StartModelOpFunc starts a background model operation. type StartModelOpFunc = func(ctx context.Context, req *ModelRequest) (*ModelOperation, error) +// 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 +// [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) // 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 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. + 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. + + // 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. +// +// Deprecated: Use [TypedBackgroundModelOptions] with +// [NewTypedBackgroundModel]. 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. + Cancel CancelModelOpFunc + 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) @@ -73,23 +116,33 @@ func LookupBackgroundModel(r api.Registry, name string) BackgroundModel { if action == nil { return nil } - return &backgroundModel{*action} + return &BackgroundModelAction{*action} } -// NewBackgroundModel defines a new model that runs in the background. -func NewBackgroundModel(name string, opts *BackgroundModelOptions, startFn StartModelOpFunc, checkFn CheckModelOpFunc) BackgroundModel { +// 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 *TypedBackgroundModelOptions, + startFn TypedStartModelOpFunc[Config], + checkFn CheckModelOpFunc, +) *BackgroundModelAction { if name == "" { - panic("ai.NewBackgroundModel: name is required") + panic("ai.NewTypedBackgroundModel: name is required") } if startFn == nil { - panic("ai.NewBackgroundModel: startFn is required") + panic("ai.NewTypedBackgroundModel: startFn is required") } if checkFn == nil { - panic("ai.NewBackgroundModel: checkFn is required") + panic("ai.NewTypedBackgroundModel: checkFn is required") } if opts == nil { - opts = &BackgroundModelOptions{} + opts = &TypedBackgroundModelOptions{} } if opts.Label == "" { opts.Label = name @@ -98,41 +151,57 @@ func NewBackgroundModel(name string, opts *BackgroundModelOptions, startFn Start opts.Supports = &ModelSupports{} } - metadata := map[string]any{ - "type": api.ActionTypeBackgroundModel, - "model": map[string]any{ - "label": opts.Label, - "supports": map[string]any{ - "media": opts.Supports.Media, - "context": opts.Supports.Context, - "multiturn": opts.Supports.Multiturn, - "systemRole": opts.Supports.SystemRole, - "tools": opts.Supports.Tools, - "toolChoice": opts.Supports.ToolChoice, - "constrained": opts.Supports.Constrained, - "output": opts.Supports.Output, - "contentType": opts.Supports.ContentType, - "longRunning": opts.Supports.LongRunning, - }, - "versions": opts.Versions, - "stage": opts.Stage, - "customOptions": opts.ConfigSchema, + configSchema, inputSchema := actionConfigSchemas[Config](opts.ConfigSchema, ModelRequest{}, "config") + + // 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, } - 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{ - simulateSystemPrompt(&opts.ModelOptions, nil), - augmentWithContext(&opts.ModelOptions, nil), - validateSupport(name, &opts.ModelOptions), + mopts := &ModelOptions{ + ConfigSchema: opts.ConfigSchema, + Label: opts.Label, + Stage: opts.Stage, + Supports: opts.Supports, + Versions: opts.Versions, } - fn := core.ChainMiddleware(mws...)(backgroundModelToModelFn(startFn)) + + // 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(mopts, nil), + augmentWithContext(mopts, nil), + validateSupport(name, mopts), + )(backgroundModelToModelFn(typedStartFn)) wrappedFn := func(ctx context.Context, req *ModelRequest) (*ModelOperation, error) { resp, err := fn(ctx, req, nil) @@ -143,14 +212,35 @@ func NewBackgroundModel(name string, opts *BackgroundModelOptions, startFn Start return modelOpFromResponse(resp) } - return &backgroundModel{*core.NewBackgroundAction(name, api.ActionTypeBackgroundModel, metadata, wrappedFn, checkFn, opts.Cancel)} + return &BackgroundModelAction{*core.NewBackgroundActionOf(api.ActionTypeBackgroundModel, name, &core.BackgroundActionOptions{ + Metadata: metadata, + InputSchema: inputSchema, + }, wrappedFn, checkFn, opts.Cancel)} } -// DefineBackgroundModel defines and registers a new model that runs in the background. -func DefineBackgroundModel(r *registry.Registry, name string, opts *BackgroundModelOptions, fn StartModelOpFunc, checkFn CheckModelOpFunc) BackgroundModel { - m := NewBackgroundModel(name, opts, fn, checkFn) - m.Register(r) - return m +// NewBackgroundModel defines a new model that runs in the background. +// +// 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") + } + 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, + Cancel: opts.Cancel, + }, func(ctx context.Context, req *ModelRequest, _ any) (*ModelOperation, error) { + return startFn(ctx, req) + }, checkFn) } // GenerateOperation generates a model response as a long-running operation based on the provided options. diff --git a/go/ai/config.go b/go/ai/config.go new file mode 100644 index 0000000000..aeccb05603 --- /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 +// 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. + +// 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..88a5b76f6f --- /dev/null +++ b/go/ai/config_test.go @@ -0,0 +1,363 @@ +// 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 NewTyped* constructors wrap 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 := 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 + }) + 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 := 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 + }) + 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 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{ + 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 + }) + + 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() + + bm := NewTypedBackgroundModel("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() TypedModelFunc[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 *ModelAction) any { + t.Helper() + desc := m.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 := 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)) + } + 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 := NewTypedModel("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 := 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 { + 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.(*ModelAction)).(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 := NewTypedEmbedder("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 := NewTypedEvaluator("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 := NewTypedBatchEvaluator("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/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 5801da9592..53ecd8b8c6 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) +// 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 [NewTypedEmbedder]). +type TypedEmbedderFunc[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. @@ -83,15 +88,39 @@ 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{}] } -// NewEmbedder creates a new [Embedder]. -func NewEmbedder(name string, opts *EmbedderOptions, fn EmbedderFunc) Embedder { +// 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 +// 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 NewTypedEmbedder[Config any]( + name string, + opts *EmbedderOptions, + fn TypedEmbedderFunc[Config], +) *EmbedderAction { if name == "" { - panic("ai.NewEmbedder: name is required") + panic("ai.NewTypedEmbedder: name is required") } if opts == nil { @@ -103,6 +132,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,31 +146,43 @@ 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), + return &EmbedderAction{ + action: *core.NewActionOf(api.ActionTypeEmbedder, name, &core.ActionOptions{ + Metadata: metadata, + InputSchema: inputSchema, + }, rawFn), } } -// DefineEmbedder registers the given embed function as an action, and returns an -// [Embedder] that runs it. -func DefineEmbedder(r api.Registry, name string, opts *EmbedderOptions, fn EmbedderFunc) Embedder { - e := NewEmbedder(name, opts, fn) - e.Register(r) - return e +// NewEmbedder creates a new [Embedder]. +// +// 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 NewTypedEmbedder(name, opts, func(ctx context.Context, req *EmbedRequest, _ any) (*EmbedResponse, error) { + return fn(ctx, req) + }) } -// 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 { @@ -147,13 +190,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 43404479f1..b78803fa2a 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 { @@ -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") @@ -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 1ab7335932..84a998a0cd 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) +// 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 [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) +// 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 +// [NewTypedBatchEvaluator]). +type TypedBatchEvaluatorFunc[Config any] = func(context.Context, *EvaluatorRequest, Config) (*EvaluatorResponse, error) + // Evaluator represents a evaluator action. type Evaluator interface { // Name returns the name of the evaluator. @@ -70,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"` @@ -161,19 +182,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 +193,46 @@ 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 - } +// 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 +// 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 NewTypedEvaluator[Config any]( + name string, + opts *EvaluatorOptions, + fn TypedEvaluatorFunc[Config], +) *EvaluatorAction { + if name == "" { + panic("ai.NewTypedEvaluator: evaluator name is required") } - return &evaluator{ - Action: *core.NewAction(name, api.ActionTypeEvaluator, metadata, inputSchema, func(ctx context.Context, req *EvaluatorRequest) (output *EvaluatorResponse, err error) { + if opts == nil { + opts = &EvaluatorOptions{} + } + + _, inputSchema := actionConfigSchemas[Config](opts.ConfigSchema, EvaluatorRequest{}, "options") + + 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) { + 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 +250,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,63 +287,91 @@ 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. -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]. +// 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. -func NewBatchEvaluator(name string, opts *EvaluatorOptions, fn BatchEvaluatorFunc) Evaluator { +// +// 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 { if name == "" { - panic("ai.NewBatchEvaluator: batch evaluator name is required") + panic("ai.NewTypedBatchEvaluator: 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, - }, + _, 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.NewAction(name, api.ActionTypeEvaluator, metadata, nil, fn), + return &EvaluatorAction{ + action: *core.NewActionOf(api.ActionTypeEvaluator, name, &core.ActionOptions{ + Metadata: evaluatorMetadata(opts), + InputSchema: inputSchema, + }, rawFn), } } -// DefineBatchEvaluator creates a new [Evaluator] and registers it. +// NewEvaluator creates a new [Evaluator]. +// This method processes the input dataset one-by-one. +// +// 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 NewTypedEvaluator(name, opts, func(ctx context.Context, req *EvaluatorCallbackRequest, _ any) (*EvaluatorCallbackResponse, error) { + return fn(ctx, req) + }) +} + +// 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. -func DefineBatchEvaluator(r api.Registry, name string, opts *EvaluatorOptions, fn BatchEvaluatorFunc) Evaluator { - e := NewBatchEvaluator(name, opts, fn) - e.(*evaluator).Register(r) - return e +// +// 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 NewTypedBatchEvaluator(name, opts, func(ctx context.Context, req *EvaluatorRequest, _ any) (*EvaluatorResponse, error) { + return fn(ctx, req) + }) } -// 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) 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/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 a535903de5..b0e99912aa 100644 --- a/go/ai/generate.go +++ b/go/ai/generate.go @@ -22,6 +22,7 @@ import ( "errors" "fmt" "iter" + "maps" "slices" "strings" "sync" @@ -69,6 +70,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] +// 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 [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 @@ -77,11 +83,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] @@ -113,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. @@ -135,10 +158,30 @@ 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 { +// 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 +// 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 NewTypedModel[Config any]( + name string, + opts *ModelOptions, + fn TypedModelFunc[Config], +) *ModelAction { if name == "" { - panic("ai.NewModel: name is required") + panic("ai.NewTypedModel: name is required") } if opts == nil { @@ -150,54 +193,71 @@ func NewModel(name string, opts *ModelOptions, fn ModelFunc) Model { opts.Supports = &ModelSupports{} } - 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": opts.ConfigSchema, + configSchema, inputSchema := actionConfigSchemas[Config](opts.ConfigSchema, ModelRequest{}, "config") + + // 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, } - 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 &ModelAction{*core.NewStreamingActionOf(api.ActionTypeModel, name, &core.ActionOptions{ + Metadata: metadata, + InputSchema: inputSchema, + }, rawFn)} } -// DefineModel creates a new [Model] and registers it. -func DefineModel(r api.Registry, name string, opts *ModelOptions, fn ModelFunc) Model { - m := NewModel(name, opts, fn) - m.Register(r) - return m +// NewModel creates a new [Model]. +// +// 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 NewTypedModel(name, opts, func(ctx context.Context, req *ModelRequest, _ any, cb ModelStreamCallback) (*ModelResponse, error) { + return fn(ctx, req, cb) + }) } -// 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 { @@ -205,9 +265,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. @@ -300,7 +358,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. @@ -876,21 +934,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/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/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 { 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/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/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 3f9b3c14da..cbb6fd48a6 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. +// 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.ModelFunc]) implements the actual generation logic, handling -// input requests ([ai.ModelRequest]) and producing responses ([ai.ModelResponse]), +// 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.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.NewModel] instead. +// use [ai.NewTypedModel] instead. // // Example: // -// echoModel := genkit.DefineModel(g, "custom/echo", +// echoModel := genkit.DefineTypedModel(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,18 +559,60 @@ 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 { + m := ai.NewTypedModel(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 [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 } -// DefineBackgroundModel defines a background model, registers it as a [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. +// +// 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.TypedBackgroundModelOptions, + startFn ai.TypedStartModelOpFunc[Config], + checkFn ai.CheckModelOpFunc, +) *ai.BackgroundModelAction { + m := ai.NewTypedBackgroundModel(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 [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 { - 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. @@ -629,7 +675,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, @@ -678,7 +726,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, @@ -738,7 +788,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. @@ -800,7 +852,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. @@ -1339,18 +1393,42 @@ 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). +// 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). // // 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.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.NewEmbedder] instead. +// use [ai.NewTypedEmbedder] instead. +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 +} + +// DefineEmbedder defines a custom text embedding implementation, registers it as a +// [core.Action] of type Embedder, and returns an [ai.Embedder]. +// +// 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 { - 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. @@ -1370,35 +1448,72 @@ 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. +// 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. // -// 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]. -func DefineEvaluator(g *Genkit, name string, opts *ai.EvaluatorOptions, fn ai.EvaluatorFunc) ai.Evaluator { - return ai.DefineEvaluator(g.reg, name, opts, fn) +// 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 { + e := ai.NewTypedEvaluator(name, opts, fn) + e.Register(g.reg) + return e } -// DefineBatchEvaluator defines an evaluator that processes the entire dataset at once, +// 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 [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 { + e := ai.NewEvaluator(name, opts, fn) + e.Register(g.reg) + return e +} + +// 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]. +// // 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.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 { + e := ai.NewTypedBatchEvaluator(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 [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 { - 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. @@ -1668,7 +1783,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),