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

go/ai/generate.go

@@ -389,7 +389,18 @@ func GenerateWithRequest(ctx context.Context, r api.Registry, opts *GenerateActi
 				return resp, nil
 			}
 
-			return generate(ctx, newReq, currentTurn+1, currentIndex+1)
+			finalResp, err := generate(ctx, newReq, currentTurn+1, currentIndex+1)
+			if err != nil {
+				return nil, err
+			}
+
+			if finalResp.Message != nil && resp.Message != nil {
+				var reasoningParts []*Part
+				reasoningParts = append(reasoningParts, resp.Message.Content...)
+				finalResp.Message.Content = append(reasoningParts, finalResp.Message.Content...)
+			}
+			return finalResp, nil
+			// return generate(ctx, newReq, currentTurn+1, currentIndex+1)
 		})
 	}
 

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,17 @@ func toGeminiPart(p *ai.Part) (*genai.Part, error) {
 				"content": toolResp.Output,
 			}
 		}
-		fr := genai.NewPartFromFunctionResponse(toolResp.Name, output)
-		return fr, nil
+		if multiPart, ok := p.Metadata["multipart"].(bool); ok {
+			if multiPart {
+				toolRespParts, err := toGeminiFunctionResponsePart(toolResp.Content)
+				if err != nil {
+					return nil, err
+				}
+				return genai.NewPartFromFunctionResponseWithParts(toolResp.Name, output, toolRespParts), nil
+			}
+		}
+		fmt.Printf("tool response: %#v\n", toolResp.Content)
+		return genai.NewPartFromFunctionResponse(toolResp.Name, output), nil
 	case p.IsToolRequest():
 		toolReq := p.ToolRequest
 		var input map[string]any
@@ -941,6 +986,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

@@ -717,3 +717,64 @@ func genToolName(length int, chars string) string {
 	}
 	return string(r)
 }
+
+func TestToGeminiParts_MultipartToolResponse(t *testing.T) {
+	// Create a tool response with both output and additional content (text)
+	toolResp := &ai.ToolResponse{
+		Name:   "generateImage",
+		Output: map[string]any{"status": "success"},
+		Content: []*ai.Part{
+			ai.NewTextPart("Generated image description"),
+		},
+	}
+
+	// Create an ai.Part wrapping the tool response
+	part := ai.NewToolResponsePart(toolResp)
+
+	// Convert to Gemini parts
+	geminiParts, err := toGeminiParts([]*ai.Part{part})
+	if err != nil {
+		t.Fatalf("toGeminiParts failed: %v", err)
+	}
+
+	// Expecting 2 parts: 1 for function response, 1 for text content
+	if len(geminiParts) != 2 {
+		t.Fatalf("expected 2 Gemini parts, got %d", len(geminiParts))
+	}
+
+	// Check first part (FunctionResponse)
+	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)
+	}
+
+	// Check second part (Text)
+	if geminiParts[1].Text != "Generated image description" {
+		t.Errorf("expected second part text 'Generated image description', got %q", geminiParts[1].Text)
+	}
+}
+
+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")
+	}
+}

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,26 +33,44 @@ 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 and returns content (description)
+	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.NewTextPart(fmt.Sprintf("I created an invitation card for %s for their %s. It features a beautiful rectangle.", input.Name, input.Occasion)),
+					ai.NewMediaPart("image/png", rectangle),
+				},
+			}, nil
+		},
+	)
 
-		genkit.DefineSchemaFor[Joke](g)
+	type InvitationCard struct {
+		Ocassion 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. Describe what you made.", input.Ocassion)),
+		)
 		if err != nil {
 			return "", err
 		}