feat(go): add custom configs for all primitives (#2883)

Author: sahdev77

go/ai/embedder.go

@@ -21,6 +21,7 @@ import (
 	"fmt"
 
 	"github.com/firebase/genkit/go/core"
+	"github.com/firebase/genkit/go/internal/base"
 	"github.com/firebase/genkit/go/internal/registry"
 )
 
@@ -32,6 +33,32 @@ type Embedder interface {
 	Embed(ctx context.Context, req *EmbedRequest) (*EmbedResponse, error)
 }
 
+// EmbedderInfo represents the structure of the embedder information object.
+type EmbedderInfo struct {
+	// Label is a user-friendly name for the embedder model (e.g., "Google AI - Gemini Pro").
+	Label string `json:"label,omitempty"`
+	// Supports defines the capabilities of the embedder, such as input types and multilingual support.
+	Supports *EmbedderSupports `json:"supports,omitempty"`
+	// Dimensions specifies the number of dimensions in the embedding vector.
+	Dimensions int `json:"dimensions,omitempty"`
+}
+
+// EmbedderSupports represents the supported capabilities of the embedder model.
+type EmbedderSupports struct {
+	// Input lists the types of data the model can process (e.g., "text", "image", "video").
+	Input []string `json:"input,omitempty"`
+	// Multilingual indicates whether the model supports multiple languages.
+	Multilingual bool `json:"multilingual,omitempty"`
+}
+
+// EmbedderOptions represents the configuration options for an embedder.
+type EmbedderOptions struct {
+	// ConfigSchema defines the schema for the embedder's configuration options.
+	ConfigSchema any `json:"configSchema,omitempty"`
+	// Info contains metadata about the embedder, such as its label and capabilities.
+	Info *EmbedderInfo `json:"info,omitempty"`
+}
+
 // An embedder is used to convert a document to a multidimensional vector.
 type embedder core.ActionDef[*EmbedRequest, *EmbedResponse, struct{}]
 
@@ -40,9 +67,22 @@ type embedder core.ActionDef[*EmbedRequest, *EmbedResponse, struct{}]
 func DefineEmbedder(
 	r *registry.Registry,
 	provider, name string,
+	opts *EmbedderOptions,
 	embed func(context.Context, *EmbedRequest) (*EmbedResponse, error),
 ) Embedder {
-	return (*embedder)(core.DefineAction(r, provider, name, core.ActionTypeEmbedder, nil, embed))
+	metadata := map[string]any{}
+	metadata["type"] = "embedder"
+	metadata["info"] = opts.Info
+	if opts.ConfigSchema != nil {
+		metadata["embedder"] = map[string]any{"customOptions": base.ToSchemaMap(opts.ConfigSchema)}
+	}
+	inputSchema := base.InferJSONSchema(EmbedRequest{})
+	if inputSchema.Properties != nil && opts.ConfigSchema != nil {
+		if _, ok := inputSchema.Properties.Get("options"); ok {
+			inputSchema.Properties.Set("options", base.InferJSONSchema(opts.ConfigSchema))
+		}
+	}
+	return (*embedder)(core.DefineActionWithInputSchema(r, provider, name, core.ActionTypeEmbedder, metadata, inputSchema, embed))
 }
 
 // LookupEmbedder looks up an [Embedder] registered by [DefineEmbedder].

go/ai/retriever.go

@@ -22,6 +22,7 @@ import (
 	"fmt"
 
 	"github.com/firebase/genkit/go/core"
+	"github.com/firebase/genkit/go/internal/base"
 	"github.com/firebase/genkit/go/internal/registry"
 )
 
@@ -35,12 +36,39 @@ type Retriever interface {
 	Retrieve(ctx context.Context, req *RetrieverRequest) (*RetrieverResponse, error)
 }
 
+// RetrieverInfo contains metadata about the retriever, such as its label and capabilities.
+type RetrieverInfo struct {
+	// Label is a user-friendly name for the retriever.
+	Label string `json:"label,omitempty"`
+	// Supports defines the capabilities of the retriever, such as media support.
+	Supports *RetrieverSupports `json:"supports,omitempty"`
+}
+
+// RetrieverSupports defines the supported capabilities of the retriever.
+type RetrieverSupports struct {
+	// Media indicates whether the retriever supports media content.
+	Media bool `json:"media,omitempty"`
+}
+
+// RetrieverOptions represents the configuration options for a retriever.
+type RetrieverOptions struct {
+	// ConfigSchema holds the configuration schema for the retriever.
+	ConfigSchema any
+	// Info contains metadata about the retriever, such as its label and capabilities.
+	Info *RetrieverInfo
+}
 type retriever core.ActionDef[*RetrieverRequest, *RetrieverResponse, struct{}]
 
 // DefineRetriever registers the given retrieve function as an action, and returns a
 // [Retriever] that runs it.
-func DefineRetriever(r *registry.Registry, provider, name string, fn RetrieverFunc) Retriever {
-	return (*retriever)(core.DefineAction(r, provider, name, core.ActionTypeRetriever, nil, fn))
+func DefineRetriever(r *registry.Registry, provider, name string, opts *RetrieverOptions, fn RetrieverFunc) Retriever {
+	metadata := map[string]any{}
+	metadata["type"] = "retriever"
+	metadata["info"] = opts.Info
+	if opts.ConfigSchema != nil {
+		metadata["retriever"] = map[string]any{"customOptions": base.InferJSONSchema(opts.ConfigSchema)}
+	}
+	return (*retriever)(core.DefineAction(r, provider, name, core.ActionTypeRetriever, metadata, fn))
 }
 
 // LookupRetriever looks up a [Retriever] registered by [DefineRetriever].

go/core/action.go

@@ -97,7 +97,8 @@ type noStream = func(context.Context, struct{}) error
 // DefineAction creates a new non-streaming Action and registers it.
 func DefineAction[In, Out any](
 	r *registry.Registry,
-	provider, name string,
+	provider,
+	name string,
 	atype ActionType,
 	metadata map[string]any,
 	fn Func[In, Out],
@@ -140,24 +141,25 @@ func DefineStreamingAction[In, Out, Stream any](
 // This differs from DefineAction in that the input schema is
 // defined dynamically; the static input type is "any".
 // This is used for prompts and tools that need custom input validation.
-func DefineActionWithInputSchema[Out any](
+func DefineActionWithInputSchema[In, Out any](
 	r *registry.Registry,
 	provider, name string,
 	atype ActionType,
 	metadata map[string]any,
 	inputSchema *jsonschema.Schema,
-	fn Func[any, Out],
-) *ActionDef[any, Out, struct{}] {
+	fn Func[In, Out],
+) *ActionDef[In, Out, struct{}] {
 	return defineAction(r, provider, name, atype, metadata, inputSchema,
-		func(ctx context.Context, in any, _ noStream) (Out, error) {
+		func(ctx context.Context, in In, _ noStream) (Out, error) {
 			return fn(ctx, in)
 		})
 }
 
 // defineAction creates an action and registers it with the given Registry.
 func defineAction[In, Out, Stream any](
 	r *registry.Registry,
-	provider, name string,
+	provider,
+	name string,
 	atype ActionType,
 	metadata map[string]any,
 	inputSchema *jsonschema.Schema,

go/genkit/genkit.go

@@ -731,8 +731,8 @@ func GenerateData[Out any](ctx context.Context, g *Genkit, opts ...ai.GenerateOp
 // The `provider` and `name` form the unique identifier. The `ret` function
 // contains the logic to process an [ai.RetrieverRequest] (containing the query)
 // and return an [ai.RetrieverResponse] (containing the relevant documents).
-func DefineRetriever(g *Genkit, provider, name string, ret func(context.Context, *ai.RetrieverRequest) (*ai.RetrieverResponse, error)) ai.Retriever {
-	return ai.DefineRetriever(g.reg, provider, name, ret)
+func DefineRetriever(g *Genkit, provider, name string, opts *ai.RetrieverOptions, ret func(context.Context, *ai.RetrieverRequest) (*ai.RetrieverResponse, error)) ai.Retriever {
+	return ai.DefineRetriever(g.reg, provider, name, opts, ret)
 }
 
 // LookupRetriever retrieves a registered [ai.Retriever] by its provider and name.
@@ -749,8 +749,8 @@ func LookupRetriever(g *Genkit, provider, name string) ai.Retriever {
 // The `provider` and `name` form the unique identifier. The `embed` function
 // contains the logic to process an [ai.EmbedRequest] (containing documents or a query)
 // and return an [ai.EmbedResponse] (containing the corresponding embeddings).
-func DefineEmbedder(g *Genkit, provider, name string, embed func(context.Context, *ai.EmbedRequest) (*ai.EmbedResponse, error)) ai.Embedder {
-	return ai.DefineEmbedder(g.reg, provider, name, embed)
+func DefineEmbedder(g *Genkit, provider string, name string, opts *ai.EmbedderOptions, embed func(context.Context, *ai.EmbedRequest) (*ai.EmbedResponse, error)) ai.Embedder {
+	return ai.DefineEmbedder(g.reg, provider, name, opts, embed)
 }
 
 // LookupEmbedder retrieves a registered [ai.Embedder] by its provider and name.

go/internal/base/json.go

@@ -160,3 +160,9 @@ func GetJsonObjectLines(text string) []string {
 	// Return the slice containing the filtered and trimmed lines.
 	return result
 }
+
+func ToSchemaMap(config any) map[string]any {
+	schema := InferJSONSchema(config)
+	result := SchemaAsMap(schema)
+	return result
+}

go/internal/doc-snippets/pinecone.go

@@ -48,10 +48,19 @@ func pineconeEx(ctx context.Context) error {
 	var docChunks []*ai.Document
 
 	// [START defineretriever]
+	retOpts := &ai.RetrieverOptions{
+		ConfigSchema: pinecone.PineconeRetrieverOptions{},
+		Info: &ai.RetrieverInfo{
+			Label: "Pinecone",
+			Supports: &ai.RetrieverSupports{
+				Media: false,
+			},
+		},
+	}
 	ds, menuRetriever, err := pinecone.DefineRetriever(ctx, g, pinecone.Config{
 		IndexID:  "menu_data",                                           // Your Pinecone index
 		Embedder: googlegenai.GoogleAIEmbedder(g, "text-embedding-004"), // Embedding model of your choice
-	})
+	}, retOpts)
 	if err != nil {
 		return err
 	}

go/internal/doc-snippets/rag/main.go

@@ -49,13 +49,23 @@ func main() {
 	if err != nil {
 		log.Fatal(err)
 	}
+	retOpts := &ai.RetrieverOptions{
+		ConfigSchema: localvec.RetrieverOptions{},
+		Info: &ai.RetrieverInfo{
+			Label: "menuQA",
+			Supports: &ai.RetrieverSupports{
+				Media: false,
+			},
+		},
+	}
 
 	docStore, _, err := localvec.DefineRetriever(
 		g,
 		"menuQA",
 		localvec.Config{
 			Embedder: googlegenai.VertexAIEmbedder(g, "text-embedding-004"),
 		},
+		retOpts,
 	)
 	if err != nil {
 		log.Fatal(err)
@@ -155,12 +165,23 @@ func menuQA() {
 
 	model := googlegenai.VertexAIModel(g, "gemini-1.5-flash")
 
+	retOpts := &ai.RetrieverOptions{
+		ConfigSchema: localvec.RetrieverOptions{},
+		Info: &ai.RetrieverInfo{
+			Label: "menuQA",
+			Supports: &ai.RetrieverSupports{
+				Media: false,
+			},
+		},
+	}
+
 	_, menuPdfRetriever, err := localvec.DefineRetriever(
 		g,
 		"menuQA",
 		localvec.Config{
 			Embedder: googlegenai.VertexAIEmbedder(g, "text-embedding-004"),
 		},
+		retOpts,
 	)
 	if err != nil {
 		log.Fatal(err)
@@ -207,23 +228,44 @@ func customret() {
 		log.Fatal(err)
 	}
 
+	retOpts := &ai.RetrieverOptions{
+		ConfigSchema: localvec.RetrieverOptions{},
+		Info: &ai.RetrieverInfo{
+			Label: "menuQA",
+			Supports: &ai.RetrieverSupports{
+				Media: false,
+			},
+		},
+	}
+
 	_, menuPDFRetriever, _ := localvec.DefineRetriever(
 		g,
 		"menuQA",
 		localvec.Config{
 			Embedder: googlegenai.VertexAIEmbedder(g, "text-embedding-004"),
 		},
+		retOpts,
 	)
 
 	// [START customret]
 	type CustomMenuRetrieverOptions struct {
 		K          int
 		PreRerankK int
 	}
+	genRetOpts := &ai.RetrieverOptions{
+		ConfigSchema: CustomMenuRetrieverOptions{},
+		Info: &ai.RetrieverInfo{
+			Label: "advancedMenuRetriever",
+			Supports: &ai.RetrieverSupports{
+				Media: false,
+			},
+		},
+	}
 	advancedMenuRetriever := genkit.DefineRetriever(
 		g,
 		"custom",
 		"advancedMenuRetriever",
+		genRetOpts,
 		func(ctx context.Context, req *ai.RetrieverRequest) (*ai.RetrieverResponse, error) {
 			// Handle options passed using our custom type.
 			opts, _ := req.Options.(CustomMenuRetrieverOptions)

go/internal/fakeembedder/fakeembedder_test.go

@@ -32,7 +32,17 @@ func TestFakeEmbedder(t *testing.T) {
 	}
 
 	embed := New()
-	emb := ai.DefineEmbedder(r, "fake", "embed", embed.Embed)
+	emdOpts := &ai.EmbedderOptions{
+		Info: &ai.EmbedderInfo{
+			Dimensions: 32,
+			Label:      "embed",
+			Supports: &ai.EmbedderSupports{
+				Input: []string{"text"},
+			},
+		},
+		ConfigSchema: nil,
+	}
+	emb := ai.DefineEmbedder(r, "fake", "embed", emdOpts, embed.Embed)
 	d := ai.DocumentFromText("fakeembedder test", nil)
 
 	vals := []float32{1, 2}

go/plugins/alloydb/genkit.go

@@ -88,7 +88,7 @@ func DefineRetriever(ctx context.Context, g *genkit.Genkit, p *Postgres, cfg *Co
 		return nil, nil, err
 	}
 
-	return ds, genkit.DefineRetriever(g, provider, ds.config.TableName, ds.Retrieve), nil
+	return ds, genkit.DefineRetriever(g, provider, ds.config.TableName, nil, ds.Retrieve), nil
 }
 
 // Retriever returns the retriever with the given index name.

go/plugins/alloydb/genkit_test.go

@@ -170,7 +170,7 @@ func TestPostgres(t *testing.T) {
 		IDColumn:              CustomIdColumn,
 		MetadataJSONColumn:    CustomMetadataColumn,
 		IgnoreMetadataColumns: []string{"created_at", "updated_at"},
-		Embedder:              genkit.DefineEmbedder(g, "fake", "embedder3", embedder.Embed),
+		Embedder:              genkit.DefineEmbedder(g, "fake", "embedder3", nil, embedder.Embed),
 		EmbedderOptions:       nil,
 	}