fix(go/ai): fixed custom options schema and dynamic embedder input schema (#3538)

Author: apascal07

go/ai/embedder.go

@@ -121,13 +121,13 @@ func NewEmbedder(name string, opts *EmbedderOptions, fn EmbedderFunc) Embedder {
 
 	inputSchema := core.InferSchemaMap(EmbedRequest{})
 	if inputSchema != nil && opts.ConfigSchema != nil {
-		if _, ok := inputSchema["options"]; ok {
-			inputSchema["options"] = opts.ConfigSchema
+		if props, ok := inputSchema["properties"].(map[string]any); ok {
+			props["options"] = opts.ConfigSchema
 		}
 	}
 
 	return &embedder{
-		ActionDef: *core.NewAction(name, api.ActionTypeEmbedder, metadata, nil, fn),
+		ActionDef: *core.NewAction(name, api.ActionTypeEmbedder, metadata, inputSchema, fn),
 	}
 }
 

go/ai/evaluator.go

@@ -184,8 +184,8 @@ func NewEvaluator(name string, opts *EvaluatorOptions, fn EvaluatorFunc) Evaluat
 
 	inputSchema := core.InferSchemaMap(EvaluatorRequest{})
 	if inputSchema != nil && opts.ConfigSchema != nil {
-		if _, ok := inputSchema["options"]; ok {
-			inputSchema["options"] = opts.ConfigSchema
+		if props, ok := inputSchema["properties"].(map[string]any); ok {
+			props["options"] = opts.ConfigSchema
 		}
 	}
 

go/ai/generate.go

@@ -173,8 +173,8 @@ func NewModel(name string, opts *ModelOptions, fn ModelFunc) Model {
 
 	inputSchema := core.InferSchemaMap(ModelRequest{})
 	if inputSchema != nil && opts.ConfigSchema != nil {
-		if _, ok := inputSchema["config"]; ok {
-			inputSchema["config"] = opts.ConfigSchema
+		if props, ok := inputSchema["properties"].(map[string]any); ok {
+			props["config"] = opts.ConfigSchema
 		}
 	}
 

go/ai/retriever.go

@@ -115,8 +115,8 @@ func NewRetriever(name string, opts *RetrieverOptions, fn RetrieverFunc) Retriev
 
 	inputSchema := core.InferSchemaMap(RetrieverRequest{})
 	if inputSchema != nil && opts.ConfigSchema != nil {
-		if _, ok := inputSchema["options"]; ok {
-			inputSchema["options"] = opts.ConfigSchema
+		if props, ok := inputSchema["properties"].(map[string]any); ok {
+			props["options"] = opts.ConfigSchema
 		}
 	}
 

go/internal/base/json.go

@@ -94,8 +94,8 @@ func InferJSONSchema(x any) (s *jsonschema.Schema) {
 		},
 	}
 	s = r.Reflect(x)
-	// TODO: Unwind this change once Monaco Editor supports newer than JSON schema draft-07.
 	s.Version = ""
+	s.ID = ""
 	return s
 }
 

go/internal/base/json_test.go

@@ -78,7 +78,6 @@ func TestSchemaAsMap(t *testing.T) {
 	}
 
 	want := map[string]any{
-		"$id":                  string("https://github.com/firebase/genkit/go/internal/base/foo"),
 		"additionalProperties": bool(false),
 		"properties": map[string]any{
 			"BarField": map[string]any{

go/plugins/googlegenai/gemini.go

@@ -40,9 +40,6 @@ import (
 )
 
 const (
-	// Thinking budget limit
-	thinkingBudgetMax = 24576
-
 	// Tool name regex
 	toolNameRegex = "^[a-zA-Z_][a-zA-Z0-9_.-]{0,63}$"
 )
@@ -203,8 +200,8 @@ func newEmbedder(client *genai.Client, name string, embedOpts *ai.EmbedderOption
 		var content []*genai.Content
 		var embedConfig *genai.EmbedContentConfig
 
-		if options, _ := req.Options.(*genai.EmbedContentConfig); options != nil {
-			embedConfig = options
+		if config, ok := req.Options.(*genai.EmbedContentConfig); ok {
+			embedConfig = config
 		}
 
 		for _, doc := range req.Input {