fix(go/plugins/googlegenai): ignore response mime type if tools are present (#2794)

Author: hugoaguirre

go/plugins/googlegenai/gemini.go

@@ -340,7 +340,7 @@ func generate(
 		return nil, err
 	}
 
-	gcc, err := convertRequest(input, cache)
+	gcc, err := toGeminiRequest(input, cache)
 	if err != nil {
 		return nil, err
 	}
@@ -448,9 +448,9 @@ func generate(
 	return r, nil
 }
 
-// convertRequest translates from [*ai.ModelRequest] to
+// toGeminiRequest translates from [*ai.ModelRequest] to
 // *genai.GenerateContentParameters
-func convertRequest(input *ai.ModelRequest, cache *genai.CachedContent) (*genai.GenerateContentConfig, error) {
+func toGeminiRequest(input *ai.ModelRequest, cache *genai.CachedContent) (*genai.GenerateContentConfig, error) {
 	gcc := genai.GenerateContentConfig{
 		CandidateCount: 1,
 	}
@@ -483,8 +483,9 @@ func convertRequest(input *ai.ModelRequest, cache *genai.CachedContent) (*genai.
 	hasOutput := input.Output != nil
 	isJsonFormat := hasOutput && input.Output.Format == "json"
 	isJsonContentType := hasOutput && input.Output.ContentType == "application/json"
-	jsonMode := isJsonFormat || (isJsonContentType && len(input.Tools) == 0)
-	if jsonMode {
+	jsonMode := isJsonFormat || isJsonContentType
+	// this setting is not compatible with tools forcing controlled output generation
+	if jsonMode && len(input.Tools) == 0 {
 		gcc.ResponseMIMEType = "application/json"
 	}
 
@@ -507,7 +508,7 @@ func convertRequest(input *ai.ModelRequest, cache *genai.CachedContent) (*genai.
 		gcc.Tools = tools
 
 		// Then set up the tool configuration based on ToolChoice
-		tc, err := convertToolChoice(input.ToolChoice, input.Tools)
+		tc, err := toGeminiToolChoice(input.ToolChoice, input.Tools)
 		if err != nil {
 			return nil, err
 		}
@@ -708,7 +709,7 @@ func castToStringArray(i []any) []string {
 	return r
 }
 
-func convertToolChoice(toolChoice ai.ToolChoice, tools []*ai.ToolDefinition) (*genai.ToolConfig, error) {
+func toGeminiToolChoice(toolChoice ai.ToolChoice, tools []*ai.ToolDefinition) (*genai.ToolConfig, error) {
 	var mode genai.FunctionCallingConfigMode
 	switch toolChoice {
 	case "":

go/plugins/googlegenai/gemini_test.go

@@ -22,7 +22,8 @@ import (
 	"github.com/firebase/genkit/go/ai"
 )
 
-func TestGeminiTools(t *testing.T) {
+func TestConvertRequest(t *testing.T) {
+	text := "hello"
 	tool := &ai.ToolDefinition{
 		Description: "this is a dummy tool",
 		Name:        "myTool",
@@ -35,6 +36,92 @@ func TestGeminiTools(t *testing.T) {
 		OutputSchema: map[string]any{"type": string("string")},
 	}
 
+	req := &ai.ModelRequest{
+		Config: GeminiConfig{
+			MaxOutputTokens: 10,
+			StopSequences:   []string{"stop"},
+			Temperature:     0.4,
+			TopK:            1.0,
+			TopP:            1.0,
+			Version:         text,
+		},
+		Tools:      []*ai.ToolDefinition{tool},
+		ToolChoice: ai.ToolChoiceAuto,
+		Output: &ai.ModelOutputConfig{
+			Constrained: true,
+			Schema:      map[string]any{"type": string("string")},
+		},
+		Messages: []*ai.Message{
+			{
+				Role: ai.RoleUser,
+				Content: []*ai.Part{
+					{Text: text},
+				},
+			},
+			{
+				Role: ai.RoleSystem,
+				Content: []*ai.Part{
+					{Text: text},
+				},
+			},
+			{
+				Role: ai.RoleUser,
+				Content: []*ai.Part{
+					{Text: text},
+				},
+			},
+			{
+				Role: ai.RoleSystem,
+				Content: []*ai.Part{
+					{Text: text},
+				},
+			},
+		},
+	}
+	t.Run("convert request", func(t *testing.T) {
+		gcc, err := toGeminiRequest(req, nil)
+		if err != nil {
+			t.Fatal(err)
+		}
+		if gcc.SystemInstruction == nil {
+			t.Error("expecting system instructions to be populated")
+		}
+		if len(gcc.SystemInstruction.Parts) != 2 {
+			t.Errorf("got: %d, want: 2", len(gcc.SystemInstruction.Parts))
+		}
+		if gcc.SystemInstruction.Role != string(ai.RoleSystem) {
+			t.Errorf(" system instruction role: got: %q, want: %q", gcc.SystemInstruction.Role, string(ai.RoleSystem))
+		}
+		// this is explicitly set to 1 in source
+		if gcc.CandidateCount == nil {
+			t.Error("candidate count: got: nil, want: 1")
+		}
+		ogCfg, ok := req.Config.(GeminiConfig)
+		if !ok {
+			t.Fatalf("request config should have been of type: GeminiConfig, got: %T", req.Config)
+		}
+		if gcc.MaxOutputTokens == nil {
+			t.Errorf("max output tokens: got: nil, want %d", ogCfg.MaxOutputTokens)
+		}
+		if len(gcc.StopSequences) == 0 {
+			t.Errorf("stop sequences: got: 0, want: %d", len(ogCfg.StopSequences))
+		}
+		if gcc.Temperature == nil {
+			t.Errorf("temperature: got: nil, want %f", ogCfg.Temperature)
+		}
+		if gcc.TopP == nil {
+			t.Errorf("topP: got: nil, want %f", ogCfg.TopP)
+		}
+		if gcc.TopK == nil {
+			t.Errorf("topK: got: nil, want %d", ogCfg.TopK)
+		}
+		if gcc.ResponseMIMEType != "" {
+			t.Errorf("ResponseMIMEType should been empty if tools are present")
+		}
+		if gcc.ResponseSchema == nil {
+			t.Errorf("ResponseSchema should not be empty")
+		}
+	})
 	t.Run("convert tools with valid tool", func(t *testing.T) {
 		tools := []*ai.ToolDefinition{tool}
 		gt, err := toGeminiTools(tools)

go/plugins/googlegenai/googleai_live_test.go

@@ -20,6 +20,7 @@ import (
 	"context"
 	"encoding/base64"
 	"flag"
+	"fmt"
 	"io"
 	"log"
 	"math"
@@ -67,7 +68,7 @@ func TestGoogleAILive(t *testing.T) {
 		t.Fatal(err)
 	}
 
-	gablorkenTool := genkit.DefineTool(g, "gablorken", "use when need to calculate a gablorken",
+	gablorkenTool := genkit.DefineTool(g, "gablorken", "use this tool when the user asks to calculate a gablorken",
 		func(ctx *ai.ToolContext, input struct {
 			Value int
 			Over  float64
@@ -165,7 +166,39 @@ func TestGoogleAILive(t *testing.T) {
 			t.Errorf("got %q, expecting it to contain %q", out, want)
 		}
 	})
+	t.Run("tool with json output", func(t *testing.T) {
+		type weatherQuery struct {
+			Location string `json:"location"`
+		}
+
+		type weather struct {
+			Report string `json:"report"`
+		}
+
+		weatherTool := genkit.DefineTool(g, "weatherTool",
+			"Use this tool to get the weather report for a specific location",
+			func(ctx *ai.ToolContext, input weatherQuery) (string, error) {
+				report := fmt.Sprintf("The weather in %s is sunny and 70 degrees today.", input.Location)
+				return report, nil
+			},
+		)
 
+		resp, err := genkit.Generate(ctx, g,
+			ai.WithTools(weatherTool),
+			ai.WithPrompt("what's the weather in San Francisco?"),
+			ai.WithOutputType(weather{}),
+		)
+		if err != nil {
+			t.Fatal(err)
+		}
+		var w weather
+		if err = resp.Output(&w); err != nil {
+			t.Fatal(err)
+		}
+		if w.Report == "" {
+			t.Fatal("empty weather report, tool should have provided an output")
+		}
+	})
 	t.Run("avoid tool", func(t *testing.T) {
 		resp, err := genkit.Generate(ctx, g,
 			ai.WithPrompt("what is a gablorken of 2 over 3.5?"),