[Go] ai.Model is a separate type from its action

This continues the work begun in #402 of making the main ai types
distinct from their underlying actions, instead of aliases.

This allows the types to have methods, unstead of using top-level
functions.

It also clarifies documentation and other output, like panic stack
traces.

Author: jba

go/ai/generate.go

@@ -28,8 +28,8 @@ import (
 	"github.com/firebase/genkit/go/internal/atype"
 )
 
-// A ModelAction is used to generate content from an AI model.
-type ModelAction = core.Action[*GenerateRequest, *GenerateResponse, *GenerateResponseChunk]
+// A Model is used to generate content from an AI model.
+type Model core.Action[*GenerateRequest, *GenerateResponse, *GenerateResponseChunk]
 
 // ModelStreamingCallback is the type for the streaming callback of a model.
 type ModelStreamingCallback = func(context.Context, *GenerateResponseChunk) error
@@ -50,7 +50,7 @@ type ModelMetadata struct {
 
 // DefineModel registers the given generate function as an action, and returns a
 // [ModelAction] that runs it.
-func DefineModel(provider, name string, metadata *ModelMetadata, generate func(context.Context, *GenerateRequest, ModelStreamingCallback) (*GenerateResponse, error)) *ModelAction {
+func DefineModel(provider, name string, metadata *ModelMetadata, generate func(context.Context, *GenerateRequest, ModelStreamingCallback) (*GenerateResponse, error)) *Model {
 	metadataMap := map[string]any{}
 	if metadata != nil {
 		if metadata.Label != "" {
@@ -64,25 +64,26 @@ func DefineModel(provider, name string, metadata *ModelMetadata, generate func(c
 		}
 		metadataMap["supports"] = supports
 	}
-	return core.DefineStreamingAction(provider, name, atype.Model, map[string]any{
+	return (*Model)(core.DefineStreamingAction(provider, name, atype.Model, map[string]any{
 		"model": metadataMap,
-	}, generate)
+	}, generate))
 }
 
 // LookupModel looks up a [ModelAction] registered by [DefineModel].
 // It returns nil if the model was not defined.
-func LookupModel(provider, name string) *ModelAction {
-	return core.LookupActionFor[*GenerateRequest, *GenerateResponse, *GenerateResponseChunk](atype.Model, provider, name)
+func LookupModel(provider, name string) *Model {
+	return (*Model)(core.LookupActionFor[*GenerateRequest, *GenerateResponse, *GenerateResponseChunk](atype.Model, provider, name))
 }
 
-// Generate applies a [ModelAction] to some input, handling tool requests.
-func Generate(ctx context.Context, g *ModelAction, req *GenerateRequest, cb ModelStreamingCallback) (*GenerateResponse, error) {
+// Generate applies the [Model] to some input, handling tool requests.
+func (m *Model) Generate(ctx context.Context, req *GenerateRequest, cb ModelStreamingCallback) (*GenerateResponse, error) {
 	if err := conformOutput(req); err != nil {
 		return nil, err
 	}
 
+	a := (*core.Action[*GenerateRequest, *GenerateResponse, *GenerateResponseChunk])(m)
 	for {
-		resp, err := g.Run(ctx, req, cb)
+		resp, err := a.Run(ctx, req, cb)
 		if err != nil {
 			return nil, err
 		}

go/plugins/dotprompt/dotprompt.go

@@ -81,12 +81,12 @@ type Config struct {
 	// The prompt variant.
 	Variant string
 	// The name of the model for which the prompt is input.
-	// If this is non-empty, ModelAction should be nil.
-	Model string
+	// If this is non-empty, Model should be nil.
+	ModelName string
 
-	// The ModelAction to use.
+	// The Model to use.
 	// If this is non-nil, Model should be the empty string.
-	ModelAction *ai.ModelAction
+	Model *ai.Model
 
 	// TODO(iant): document
 	Tools []*ai.ToolDefinition
@@ -224,7 +224,7 @@ func parseFrontmatter(data []byte) (name string, c Config, rest []byte, err erro
 
 	ret := Config{
 		Variant:          fy.Variant,
-		Model:            fy.Model,
+		ModelName:        fy.Model,
 		Tools:            fy.Tools,
 		Candidates:       fy.Candidates,
 		GenerationConfig: fy.Config,
@@ -289,10 +289,10 @@ func Define(name, templateText string, cfg Config) (*Prompt, error) {
 // This may be used for testing or for direct calls not using the
 // genkit action and flow mechanisms.
 func New(name, templateText string, cfg Config) (*Prompt, error) {
-	if cfg.Model == "" && cfg.ModelAction == nil {
+	if cfg.ModelName == "" && cfg.Model == nil {
 		return nil, errors.New("dotprompt.New: config must specify either Model or ModelAction")
 	}
-	if cfg.Model != "" && cfg.ModelAction != nil {
+	if cfg.ModelName != "" && cfg.Model != nil {
 		return nil, errors.New("dotprompt.New: config must specify exactly one of Model and ModelAction")
 	}
 	hash := fmt.Sprintf("%02x", sha256.Sum256([]byte(templateText)))

go/plugins/dotprompt/dotprompt_test.go

@@ -110,8 +110,8 @@ func TestPrompts(t *testing.T) {
 				t.Fatal(err)
 			}
 
-			if prompt.Model != test.model {
-				t.Errorf("got model %q want %q", prompt.Model, test.model)
+			if prompt.ModelName != test.model {
+				t.Errorf("got model %q want %q", prompt.ModelName, test.model)
 			}
 			if diff := cmpSchema(t, prompt.InputSchema, test.input); diff != "" {
 				t.Errorf("input schema mismatch (-want, +got):\n%s", diff)

go/plugins/dotprompt/genkit.go

@@ -187,9 +187,9 @@ func (p *Prompt) Generate(ctx context.Context, pr *PromptRequest, cb func(contex
 		genReq.Context = pr.Context
 	}
 
-	model := p.ModelAction
+	model := p.Model
 	if model == nil {
-		modelName := p.Model
+		modelName := p.ModelName
 		if pr.Model != "" {
 			modelName = pr.Model
 		}
@@ -207,7 +207,7 @@ func (p *Prompt) Generate(ctx context.Context, pr *PromptRequest, cb func(contex
 		}
 	}
 
-	resp, err := ai.Generate(ctx, model, genReq, cb)
+	resp, err := model.Generate(ctx, genReq, cb)
 	if err != nil {
 		return nil, err
 	}

go/plugins/dotprompt/genkit_test.go

@@ -43,7 +43,7 @@ func testGenerate(ctx context.Context, req *ai.GenerateRequest, cb func(context.
 
 func TestExecute(t *testing.T) {
 	testModel := ai.DefineModel("test", "test", nil, testGenerate)
-	p, err := New("TestExecute", "TestExecute", Config{ModelAction: testModel})
+	p, err := New("TestExecute", "TestExecute", Config{Model: testModel})
 	if err != nil {
 		t.Fatal(err)
 	}

go/plugins/googleai/googleai.go

@@ -90,7 +90,7 @@ func IsKnownModel(name string) bool {
 // For known models, it can be nil, or if non-nil it will override the known value.
 // It must be supplied for unknown models.
 // Use [IsKnownModel] to determine if a model is known.
-func DefineModel(name string, caps *ai.ModelCapabilities) (*ai.ModelAction, error) {
+func DefineModel(name string, caps *ai.ModelCapabilities) (*ai.Model, error) {
 	state.mu.Lock()
 	defer state.mu.Unlock()
 	if !state.initted {
@@ -110,7 +110,7 @@ func DefineModel(name string, caps *ai.ModelCapabilities) (*ai.ModelAction, erro
 }
 
 // requires state.mu
-func defineModel(name string, caps ai.ModelCapabilities) *ai.ModelAction {
+func defineModel(name string, caps ai.ModelCapabilities) *ai.Model {
 	meta := &ai.ModelMetadata{
 		Label:    "Google AI - " + name,
 		Supports: caps,
@@ -147,7 +147,7 @@ func defineEmbedder(name string) *ai.EmbedderAction {
 
 // Model returns the [ai.ModelAction] with the given name.
 // It returns nil if the model was not configured.
-func Model(name string) *ai.ModelAction {
+func Model(name string) *ai.Model {
 	return ai.LookupModel(provider, name)
 }
 

go/plugins/googleai/googleai_test.go

@@ -116,7 +116,7 @@ func TestLive(t *testing.T) {
 			},
 		}
 
-		resp, err := ai.Generate(ctx, model, req, nil)
+		resp, err := model.Generate(ctx, req, nil)
 		if err != nil {
 			t.Fatal(err)
 		}
@@ -145,7 +145,7 @@ func TestLive(t *testing.T) {
 
 		out := ""
 		parts := 0
-		final, err := ai.Generate(ctx, model, req, func(ctx context.Context, c *ai.GenerateResponseChunk) error {
+		final, err := model.Generate(ctx, req, func(ctx context.Context, c *ai.GenerateResponseChunk) error {
 			parts++
 			out += c.Content[0].Text
 			return nil
@@ -184,7 +184,7 @@ func TestLive(t *testing.T) {
 			Tools: []*ai.ToolDefinition{toolDef},
 		}
 
-		resp, err := ai.Generate(ctx, model, req, nil)
+		resp, err := model.Generate(ctx, req, nil)
 		if err != nil {
 			t.Fatal(err)
 		}

go/plugins/ollama/ollama.go

@@ -49,9 +49,9 @@ func defineModel(model ModelDefinition, serverAddress string) {
 	ai.DefineModel(provider, model.Name, meta, g.generate)
 }
 
-// Model returns the [ai.ModelAction] with the given name.
+// Model returns the [ai.Model] with the given name.
 // It returns nil if the model was not configured.
-func Model(name string) *ai.ModelAction {
+func Model(name string) *ai.Model {
 	return ai.LookupModel(provider, name)
 }
 

go/plugins/vertexai/vertexai.go

@@ -72,7 +72,7 @@ func Init(ctx context.Context, projectID, location string) error {
 }
 
 // DefineModel defines a model with the given name.
-func DefineModel(name string) *ai.ModelAction {
+func DefineModel(name string) *ai.Model {
 	state.mu.Lock()
 	defer state.mu.Unlock()
 	if !state.initted {
@@ -101,9 +101,9 @@ func DefineEmbedder(name string) *ai.EmbedderAction {
 	})
 }
 
-// Model returns the [ai.ModelAction] with the given name.
+// Model returns the [ai.Model] with the given name.
 // It returns nil if the model was not configured.
-func Model(name string) *ai.ModelAction {
+func Model(name string) *ai.Model {
 	return ai.LookupModel(provider, name)
 }
 

go/plugins/vertexai/vertexai_test.go

@@ -91,7 +91,7 @@ func TestLive(t *testing.T) {
 			},
 		}
 
-		resp, err := ai.Generate(ctx, model, req, nil)
+		resp, err := model.Generate(ctx, req, nil)
 		if err != nil {
 			t.Fatal(err)
 		}
@@ -120,7 +120,7 @@ func TestLive(t *testing.T) {
 		out := ""
 		parts := 0
 		model := vertexai.Model(modelName)
-		final, err := ai.Generate(ctx, model, req, func(ctx context.Context, c *ai.GenerateResponseChunk) error {
+		final, err := model.Generate(ctx, req, func(ctx context.Context, c *ai.GenerateResponseChunk) error {
 			parts++
 			for _, p := range c.Content {
 				out += p.Text
@@ -162,7 +162,7 @@ func TestLive(t *testing.T) {
 			Tools: []*ai.ToolDefinition{toolDef},
 		}
 
-		resp, err := ai.Generate(ctx, model, req, nil)
+		resp, err := model.Generate(ctx, req, nil)
 		if err != nil {
 			t.Fatal(err)
 		}