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

Code before:
		"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"`
}

type saveImageResult struct {
	Status string `json:"Status"`
}

func saveImage(ctx tool.Context, input saveImageInput) saveImageResult {
	filename := input.Filename
	part, err := ctx.Artifacts().Load(filename)
	if err != nil {
		log.Printf("Failed to load artifact '%s': %v", filename, err)
		return saveImageResult{Status: "fail"}
	}

	if part.InlineData == nil || len(part.InlineData.Data) == 0 {
		log.Printf("Artifact '%s' has no inline data", filename)
		return saveImageResult{Status: "fail"}
	}

	// Ensure the filename has a .png extension for the local file.
	localFilename := filename
	if filepath.Ext(localFilename) != ".png" {
		localFilename += ".png"
	}

	// Create an "output" directory in the current working directory if it doesn't exist.
	outputDir := "output"

Review comment (suggestion):
Add a comment that this is the function tool, and place `generateImageInput/generateImageResult` after this function.

Code after:
		},
		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.
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 {
		return generateImageResult{
			Status: "fail",
		}
	}

	response, err := client.Models.GenerateImages(
		ctx,
		"imagen-3.0-generate-002",
		input.Prompt,
		&genai.GenerateImagesConfig{NumberOfImages: 1})
	if err != nil {
		return generateImageResult{
			Status: "fail",
		}
	}

	if err := ctx.Artifacts().Save(input.Filename, *genai.NewPartFromBytes(response.GeneratedImages[0].Image.ImageBytes, "image/png")); err != nil {
		return generateImageResult{
			Status: "fail",