Repository: google/adk-go
File: examples/vertexai/imagegenerator/main.go (Go)

Code before:

type generateImageInput struct {
	Prompt   string `json:"prompt"`
	Filename string `json:"filename"`
}

type generateImageResult struct {
	Filename string `json:"filename"`
	Status   string `json:"Status"`
}

func generateImage(ctx tool.Context, input generateImageInput) generateImageResult {
	client, err := genai.NewClient(ctx, &genai.ClientConfig{
		Project:  os.Getenv("GOOGLE_CLOUD_PROJECT"),
		Location: os.Getenv("GOOGLE_CLOUD_LOCATION"),
		Backend:  genai.BackendVertexAI,
	})
	if err != nil {
		log.Fatalf("Failed to create client: %v", err)
		return generateImageResult{
			Status: "fail",
		}
	}

	response, err := client.Models.GenerateImages(
		context.Background(),
		"imagen-3.0-generate-002",
		input.Prompt,
		&genai.GenerateImagesConfig{NumberOfImages: 1})
	if err != nil {
		log.Fatalf("Failed to generate an image: %v", err)
		return generateImageResult{
			Status: "fail",
		}
	}

	if err := ctx.Artifacts().Save(input.Filename, *genai.NewPartFromBytes(response.GeneratedImages[0].Image.ImageBytes, "image/png")); err != nil {
		log.Fatalf("Failed to save artifact: %v", err)
		return generateImageResult{
			Status: "fail",
		}
	}
	return generateImageResult{
		Status:   "success",
		Filename: input.Filename,
	}
}

type saveImageInput struct {
	Filename string `json:"filename"`
}

Review comment (suggestion):
Can't we pass `ctx` here instead of `context.Background()`? (I find it's not very obvious that `tool.Context` can be used as a `context.Context`....)

Code after:
	"google.golang.org/adk/tool"
	"google.golang.org/genai"
)

func main() {
	ctx := context.Background()

	model, err := gemini.NewModel(ctx, "gemini-2.0-flash-001", nil)
	if err != nil {
		log.Fatalf("Failed to create model: %v", err)
	}

	generateImageTool, err := tool.NewFunctionTool(
		tool.FunctionToolConfig{
			Name:        "generate_image",
			Description: "Generates image and saves in artifact service.",
		},
		generateImage)
	if err != nil {
		log.Fatalf("Failed to create generate image tool: %v", err)
	}

	saveImageTool, err := tool.NewFunctionTool(
		tool.FunctionToolConfig{
			Name:        "save_image_locally",
			Description: "Saves images locally based on the filename.",
		},
		saveImage)
	if err != nil {
		log.Fatalf("Failed to create generate image tool: %v", err)
	}

	agent, err := llmagent.New(llmagent.Config{
		Name:        "image_generator",
		Model:       model,
		Description: "Agent to generate pictures, answers questions about it and saves it locally if asked.",
		Instruction: "You are an agent whose job is to generate an image, describe it and save it locally if asked." +
			" Also user will provide the filename and you should save it in the artifacts with that filename." +
			" When user ask to save image locally you can call save_image_locally to do it.",
		Tools: []tool.Tool{
			tool.NewLoadArtifactsTool(), generateImageTool, saveImageTool,
		},
	})
	if err != nil {
		log.Fatalf("Failed to create agent: %v", err)
	}

	examples.Run(ctx, agent, &examples.RunConfig{ArtifactService: artifactservice.Mem()})
}

// This is a function tool to generate images using Vertex AI's Imagen model.