feat(go/plugins/googlegenai): add multi-part tools support

go/go.mod

@@ -41,7 +41,7 @@ require (
 	golang.org/x/exp v0.0.0-20250408133849-7e4ce0ab07d0
 	golang.org/x/tools v0.34.0
 	google.golang.org/api v0.236.0
-	google.golang.org/genai v1.36.0
+	google.golang.org/genai v1.40.0
 )
 
 require (

go/go.sum

@@ -537,8 +537,8 @@ google.golang.org/appengine v1.1.0/go.mod h1:EbEs0AVv82hx2wNQdGPgUI5lhzA/G0D9Ywl
 google.golang.org/appengine v1.4.0/go.mod h1:xpcJRLb0r/rnEns0DIKYYv+WjYCduHsrkT7/EB5XEv4=
 google.golang.org/appengine/v2 v2.0.6 h1:LvPZLGuchSBslPBp+LAhihBeGSiRh1myRoYK4NtuBIw=
 google.golang.org/appengine/v2 v2.0.6/go.mod h1:WoEXGoXNfa0mLvaH5sV3ZSGXwVmy8yf7Z1JKf3J3wLI=
-google.golang.org/genai v1.36.0 h1:sJCIjqTAmwrtAIaemtTiKkg2TO1RxnYEusTmEQ3nGxM=
-google.golang.org/genai v1.36.0/go.mod h1:A3kkl0nyBjyFlNjgxIwKq70julKbIxpSxqKO5gw/gmk=
+google.golang.org/genai v1.40.0 h1:kYxyQSH+vsib8dvsgyLJzsVEIv5k3ZmHJyVqdvGncmc=
+google.golang.org/genai v1.40.0/go.mod h1:A3kkl0nyBjyFlNjgxIwKq70julKbIxpSxqKO5gw/gmk=
 google.golang.org/genproto v0.0.0-20180817151627-c66870c02cf8/go.mod h1:JiN7NxoALGmiZfu7CAH4rXhgtRTLTxftemlI0sWmxmc=
 google.golang.org/genproto v0.0.0-20190819201941-24fa4b261c55/go.mod h1:DMBHOl98Agz4BDEuKkezgsaosCRResVns1a3J2ZsMNc=
 google.golang.org/genproto v0.0.0-20200526211855-cb27e3aa2013/go.mod h1:NbSheEEYHJ7i3ixzK3sjbqSGDJWnxyFXZblF3eUsNvo=

go/plugins/googlegenai/gemini.go

@@ -478,6 +478,34 @@ func toGeminiTools(inTools []*ai.ToolDefinition) ([]*genai.Tool, error) {
 	return outTools, nil
 }
 
+// toGeminiFunctionResponsePart translates a slice of [ai.Part] to a slice of [genai.FunctionResponsePart]
+func toGeminiFunctionResponsePart(parts []*ai.Part) ([]*genai.FunctionResponsePart, error) {
+	frp := []*genai.FunctionResponsePart{}
+	for _, p := range parts {
+		switch {
+		case p.IsData():
+			contentType, data, err := uri.Data(p)
+			if err != nil {
+				return nil, err
+			}
+			frp = append(frp, genai.NewFunctionResponsePartFromBytes(data, contentType))
+		case p.IsMedia():
+			if strings.HasPrefix(p.Text, "data:") {
+				contentType, data, err := uri.Data(p)
+				if err != nil {
+					return nil, err
+				}
+				frp = append(frp, genai.NewFunctionResponsePartFromBytes(data, contentType))
+				continue
+			}
+			frp = append(frp, genai.NewFunctionResponsePartFromURI(p.Text, p.ContentType))
+		default:
+			return nil, fmt.Errorf("unsupported function response part type: %d", p.Kind)
+		}
+	}
+	return frp, nil
+}
+
 // mergeTools consolidates all FunctionDeclarations into a single Tool
 // while preserving non-function tools (Retrieval, GoogleSearch, CodeExecution, etc.)
 func mergeTools(ts []*genai.Tool) []*genai.Tool {
@@ -808,6 +836,14 @@ func translateCandidate(cand *genai.Candidate) (*ai.ModelResponse, error) {
 				Name:  part.FunctionCall.Name,
 				Input: part.FunctionCall.Args,
 			})
+			// FunctionCall parts may contain a ThoughtSignature that must be preserved
+			// and returned in subsequent requests for the tool call to be valid.
+			if len(part.ThoughtSignature) > 0 {
+				if p.Metadata == nil {
+					p.Metadata = make(map[string]any)
+				}
+				p.Metadata["signature"] = part.ThoughtSignature
+			}
 		}
 		if part.CodeExecutionResult != nil {
 			partFound++
@@ -888,7 +924,7 @@ func toGeminiParts(parts []*ai.Part) ([]*genai.Part, error) {
 func toGeminiPart(p *ai.Part) (*genai.Part, error) {
 	switch {
 	case p.IsReasoning():
-		// TODO: go-genai does not support genai.NewPartFromThought()
+		// NOTE: go-genai does not support genai.NewPartFromThought()
 		signature := []byte{}
 		if p.Metadata != nil {
 			if sig, ok := p.Metadata["signature"].([]byte); ok {
@@ -928,8 +964,22 @@ func toGeminiPart(p *ai.Part) (*genai.Part, error) {
 				"content": toolResp.Output,
 			}
 		}
-		fr := genai.NewPartFromFunctionResponse(toolResp.Name, output)
-		return fr, nil
+
+		var isMultipart bool
+		if multiPart, ok := p.Metadata["multipart"].(bool); ok {
+			isMultipart = multiPart
+		}
+		if len(toolResp.Content) > 0 {
+			isMultipart = true
+		}
+		if isMultipart {
+			toolRespParts, err := toGeminiFunctionResponsePart(toolResp.Content)
+			if err != nil {
+				return nil, err
+			}
+			return genai.NewPartFromFunctionResponseWithParts(toolResp.Name, output, toolRespParts), nil
+		}
+		return genai.NewPartFromFunctionResponse(toolResp.Name, output), nil
 	case p.IsToolRequest():
 		toolReq := p.ToolRequest
 		var input map[string]any
@@ -941,6 +991,12 @@ func toGeminiPart(p *ai.Part) (*genai.Part, error) {
 			}
 		}
 		fc := genai.NewPartFromFunctionCall(toolReq.Name, input)
+		// Restore ThoughtSignature if present in metadata
+		if p.Metadata != nil {
+			if sig, ok := p.Metadata["signature"].([]byte); ok {
+				fc.ThoughtSignature = sig
+			}
+		}
 		return fc, nil
 	default:
 		panic("unknown part type in a request")

go/plugins/googlegenai/gemini_test.go

@@ -707,6 +707,82 @@ func TestValidToolName(t *testing.T) {
 	}
 }
 
+func TestToGeminiParts_MultipartToolResponse(t *testing.T) {
+	t.Run("ValidPartType", func(t *testing.T) {
+		// Create a tool response with both output and additional content (media)
+		toolResp := &ai.ToolResponse{
+			Name:   "generateImage",
+			Output: map[string]any{"status": "success"},
+			Content: []*ai.Part{
+				ai.NewMediaPart("image/png", "data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mP8z8BQDwAEhQGAhKmMIQAAAABJRU5ErkJggg=="),
+			},
+		}
+
+		// create a mock ToolResponsePart, setting "multipart" to true is required
+		part := ai.NewToolResponsePart(toolResp)
+		part.Metadata = map[string]any{"multipart": true}
+
+		geminiParts, err := toGeminiParts([]*ai.Part{part})
+		if err != nil {
+			t.Fatalf("toGeminiParts failed: %v", err)
+		}
+
+		// Expecting 1 part which contains the function response with internal parts
+		if len(geminiParts) != 1 {
+			t.Fatalf("expected 1 Gemini part, got %d", len(geminiParts))
+		}
+
+		if geminiParts[0].FunctionResponse == nil {
+			t.Error("expected first part to be FunctionResponse")
+		}
+		if geminiParts[0].FunctionResponse.Name != "generateImage" {
+			t.Errorf("expected function name 'generateImage', got %q", geminiParts[0].FunctionResponse.Name)
+		}
+	})
+
+	t.Run("UnsupportedPartType", func(t *testing.T) {
+		// Create a tool response with text content (unsupported for multipart)
+		toolResp := &ai.ToolResponse{
+			Name:   "generateText",
+			Output: map[string]any{"status": "success"},
+			Content: []*ai.Part{
+				ai.NewTextPart("Generated text"),
+			},
+		}
+
+		part := ai.NewToolResponsePart(toolResp)
+		part.Metadata = map[string]any{"multipart": true}
+
+		_, err := toGeminiParts([]*ai.Part{part})
+		if err == nil {
+			t.Fatal("expected error for unsupported text part in multipart response, got nil")
+		}
+	})
+}
+
+func TestToGeminiParts_SimpleToolResponse(t *testing.T) {
+	// Create a simple tool response (no content)
+	toolResp := &ai.ToolResponse{
+		Name:   "search",
+		Output: map[string]any{"result": "foo"},
+	}
+
+	part := ai.NewToolResponsePart(toolResp)
+
+	geminiParts, err := toGeminiParts([]*ai.Part{part})
+	if err != nil {
+		t.Fatalf("toGeminiParts failed: %v", err)
+	}
+
+	if len(geminiParts) != 1 {
+		t.Fatalf("expected 1 Gemini part, got %d", len(geminiParts))
+	}
+
+	if geminiParts[0].FunctionResponse == nil {
+		t.Error("expected part to be FunctionResponse")
+	}
+}
+
 // genToolName generates a string of a specified length using only
 // the valid characters for a Gemini Tool name
 func genToolName(length int, chars string) string {

go/plugins/googlegenai/googleai_live_test.go

@@ -170,7 +170,7 @@ func TestGoogleAILive(t *testing.T) {
 			t.Fatal(err)
 		}
 
-		out := resp.Message.Content[0].Text
+		out := resp.Text()
 		const want = "11.31"
 		if !strings.Contains(out, want) {
 			t.Errorf("got %q, expecting it to contain %q", out, want)
@@ -219,7 +219,7 @@ func TestGoogleAILive(t *testing.T) {
 			t.Fatal(err)
 		}
 
-		out := resp.Message.Content[0].Text
+		out := resp.Text()
 		const want = "11.31"
 		if !strings.Contains(out, want) {
 			t.Errorf("got %q, expecting it to contain %q", out, want)
@@ -307,7 +307,7 @@ func TestGoogleAILive(t *testing.T) {
 			t.Fatal(err)
 		}
 
-		out := resp.Message.Content[0].Text
+		out := resp.Text()
 		const doNotWant = "11.31"
 		if strings.Contains(out, doNotWant) {
 			t.Errorf("got %q, expecting it NOT to contain %q", out, doNotWant)
@@ -582,6 +582,37 @@ func TestGoogleAILive(t *testing.T) {
 			t.Fatal("thoughts tokens should be zero")
 		}
 	})
+	t.Run("multipart tool", func(t *testing.T) {
+		m := googlegenai.GoogleAIModel(g, "gemini-3-pro-preview")
+		img64, err := fetchImgAsBase64()
+		if err != nil {
+			t.Fatal(err)
+		}
+
+		tool := genkit.DefineMultipartTool(g, "getImage", "returns a misterious image",
+			func(ctx *ai.ToolContext, input any) (*ai.MultipartToolResponse, error) {
+				return &ai.MultipartToolResponse{
+					Output: map[string]any{"status": "success"},
+					Content: []*ai.Part{
+						ai.NewMediaPart("image/jpeg", "data:image/jpeg;base64,"+img64),
+					},
+				}, nil
+			},
+		)
+
+		resp, err := genkit.Generate(ctx, g,
+			ai.WithModel(m),
+			ai.WithTools(tool),
+			ai.WithPrompt("get an image and tell me what is in it"),
+		)
+		if err != nil {
+			t.Fatal(err)
+		}
+
+		if !strings.Contains(strings.ToLower(resp.Text()), "cat") {
+			t.Errorf("expected response to contain 'cat', got: %s", resp.Text())
+		}
+	})
 }
 
 func TestCacheHelper(t *testing.T) {

go/samples/basic-gemini/main.go

@@ -16,6 +16,7 @@ package main
 
 import (
 	"context"
+	"fmt"
 
 	"github.com/firebase/genkit/go/ai"
 	"github.com/firebase/genkit/go/genkit"
@@ -32,31 +33,62 @@ func main() {
 	// practice.
 	g := genkit.Init(ctx, genkit.WithPlugins(&googlegenai.GoogleAI{}))
 
-	// Define a simple flow that generates jokes about a given topic
-	genkit.DefineStreamingFlow(g, "jokesFlow", func(ctx context.Context, input string, cb ai.ModelStreamCallback) (string, error) {
-		type Joke struct {
-			Joke     string `json:"joke"`
-			Category string `json:"jokeCategory" description:"What is the joke about"`
-		}
+	// Define a multipart tool.
+	// This simulates a tool that "generates" an invitation card
+	invitationTool := genkit.DefineMultipartTool(g, "createInvitationCard", "Creates a greeting card",
+		func(ctx *ai.ToolContext, input struct {
+			Name     string `json:"name"`
+			Occasion string `json:"occasion"`
+		},
+		) (*ai.MultipartToolResponse, error) {
+			rectangle := "data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAHIAAABUAQMAAABk5vEVAAAABlBMVEX///8AAABVwtN+" +
+				"AAAAI0lEQVR4nGNgGHaA/z8UHIDwOWASDqP8Uf7w56On/1FAQwAAVM0exw1hqwkAAAAASUVORK5CYII="
+			return &ai.MultipartToolResponse{
+				Output: map[string]any{"success": true},
+				Content: []*ai.Part{
+					ai.NewMediaPart("image/png", rectangle),
+				},
+			}, nil
+		},
+	)
 
-		genkit.DefineSchemaFor[Joke](g)
+	type InvitationCard struct {
+		Occasion string `json:"occasion"`
+	}
 
+	// Define a simple flow that uses the multipart tool
+	genkit.DefineStreamingFlow(g, "cardFlow", func(ctx context.Context, input InvitationCard, cb ai.ModelStreamCallback) ([]string, error) {
 		resp, err := genkit.Generate(ctx, g,
-			ai.WithModelName("googleai/gemini-2.5-flash"),
+			ai.WithModelName("googleai/gemini-3-pro-preview"),
 			ai.WithConfig(&genai.GenerateContentConfig{
 				Temperature: genai.Ptr[float32](1.0),
 				ThinkingConfig: &genai.ThinkingConfig{
-					ThinkingBudget: genai.Ptr[int32](0),
+					ThinkingLevel: genai.ThinkingLevelHigh,
 				},
 			}),
+			ai.WithTools(invitationTool),
 			ai.WithStreaming(cb),
-			ai.WithOutputSchemaName("Joke"),
-			ai.WithPrompt(`Tell short jokes about %s`, input))
+			ai.WithPrompt(fmt.Sprintf("Create an invitation card for the following ocassion: %s. Create one for Alex and another one for Pavel", input.Occasion)),
+		)
 		if err != nil {
-			return "", err
+			return nil, err
 		}
 
-		return resp.Text(), nil
+		invitations := []string{}
+		for _, m := range resp.History() {
+			if m.Role == ai.RoleTool {
+				for _, p := range m.Content {
+					if p.IsToolResponse() {
+						for _, contentPart := range p.ToolResponse.Content {
+							if contentPart.IsMedia() {
+								invitations = append(invitations, contentPart.Text)
+							}
+						}
+					}
+				}
+			}
+		}
+		return invitations, nil
 	})
 
 	<-ctx.Done()