main.go

Overview

The main.go file demonstrates the creation and execution of a custom agent designed to generate images using Vertex AI's Imagen model. The agent is capable of generating images based on user prompts, storing those images as artifacts, and saving the images to the local filesystem upon request. This file orchestrates the integration of several components such as the Gemini model interface, function tools for generating and saving images, artifact management, and the agent lifecycle launcher.

This implementation serves as an example of building an AI agent that leverages llmagent capabilities (LLM Integration and Agents) with function tools (Function Tools) and artifact management (Artifact Management) to perform image generation and storage workflows.

Package and Imports

• Package main: The executable entry point for the application.

• Imports: Includes context management, logging, OS and file path utilities, and various Google ADK packages for AI models, agents, artifact services, and command launchers.

Main Components

Main Function

func main()

• Purpose: Initializes the agent and launches it.

• Workflow:

  1. Creates a Gemini model instance (gemini.NewModel) representing the Vertex AI Gemini-2.0 Flash model.

  2. Creates two function tools:

     • generateImageTool: Generates images using Vertex AI Imagen.

     • saveImageTool: Saves images from artifact storage to the local filesystem.

  3. Constructs an llmagent.Agent that uses the model and tools, with explicit instructions describing its role.

  4. Sets up an in-memory artifact service for temporary artifact storage.

  5. Configures and starts a full command launcher (full.NewLauncher) to run the agent lifecycle and process command-line arguments.

• Error Handling: Fatal logs on any initialization failures to halt execution.

• Context: Uses context.Background() as the root context propagated through calls.

Agent Configuration and Tools

• Agent Name: "image_generator"

• Description: Generates images, describes them, and saves locally if requested.

• Instruction: Guides the agent to generate images, save them as artifacts with provided filenames, and save locally on demand.

• Tools Registered:

  • loadartifactstool.New(): Loads artifacts into the session (Artifact Loading Tool).

  • generateImageTool: Custom function tool to generate images.

  • saveImageTool: Custom function tool to save images locally.

Function Tools

generateImage

func generateImage(ctx tool.Context, input generateImageInput) (generateImageResult, error)

• Purpose: Generates an image using the Vertex AI Imagen model and saves the image as an artifact with a specified filename.

• Parameters:

  • ctx: Tool invocation context providing access to artifacts and cancellation.

  • input: Struct containing:

    • Prompt (string): Text prompt for image generation.

    • Filename (string): Name to save the generated image artifact.

• Returns:

  • generateImageResult: Contains the filename and status "success" on successful generation.

  • error: Non-nil if any operation fails.

• Implementation Details:

  1. Creates a GenAI client configured with Google Cloud project and location environment variables.

  2. Calls GenerateImages on the Imagen model backend with the prompt, requesting a single image.

  3. Saves the generated image bytes as an artifact using the filename provided.

• Usage Example:

  {
    "prompt": "A futuristic city skyline at sunset",
    "filename": "city_sunset.png"
  }
  

• Errors: Propagates errors from client creation, image generation, or artifact saving.

saveImage

func saveImage(ctx tool.Context, input saveImageInput) (saveImageResult, error)

• Purpose: Loads an image artifact by filename and saves it to the local filesystem under an output directory.

• Parameters:

  • ctx: Tool invocation context with artifact access.

  • input: Struct containing:

    • Filename (string): The artifact filename to load and save locally.

• Returns:

  • saveImageResult: Contains a status "success" on successful save.

  • error: Non-nil if loading or saving fails.

• Implementation Details:

  1. Loads the artifact from the artifact service.

  2. Checks for valid inline data in the artifact.

  3. Ensures the filename has a .png extension for the local file.

  4. Creates an output directory if missing.

  5. Writes the image bytes to a file in the output directory with appropriate permissions.

• Error Handling: Logs errors related to missing artifact data, directory creation failures, or file write failures.

• Usage Example:

  {
    "filename": "city_sunset.png"
  }
  

Data Structures

generateImageInput

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

• Input schema for the generateImage function tool.

• Contains the image prompt and the artifact filename.

generateImageResult

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

• Result schema for generateImage.

• Returns the filename and operation status.

saveImageInput

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

• Input schema for the saveImage function tool.

• Specifies the artifact filename to be saved locally.

saveImageResult

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

• Result schema for saveImage.

• Indicates success status.

Artifact Service Interaction

• Uses an in-memory artifact service (artifact.InMemoryService()) for storing and retrieving image blobs during runtime (In-Memory Artifact Service).

• Images generated are saved as artifacts via ctx.Artifacts().Save.

• Images to be saved locally are loaded from artifacts via ctx.Artifacts().Load.

• This approach decouples image generation from persistent storage and local saving.

Agent Execution and Launcher

• The agent lifecycle is managed by the full.NewLauncher() which supports command-line interface execution and integrates the artifact service and agent loader (Agent Lifecycle and Callbacks, REST API and Web Launchers).

• The launcher executes with the provided context and command line arguments, handling the agent's run loop.

Interactions and Workflow Summary

The file ties together several components to implement this workflow:

flowchart TD
A[main] --> B[Create Gemini Model]
B --> C[Create generateImage Tool]
C --> D[Create saveImage Tool]
D --> E[Create LLM Agent with Model and Tools]
E --> F[Setup In-Memory Artifact Service]
F --> G[Launch Agent Lifecycle]
subgraph generateImage Tool
G1[Create GenAI Client] --> G2[Call GenerateImages API]
G2 --> G3[Save Generated Image Artifact]
end
subgraph saveImage Tool
S1[Load Image Artifact] --> S2[Check Inline Data]
S2 --> S3[Ensure .png Extension]
S3 --> S4[Create Output Directory]
S4 --> S5[Write Image to Local File]
end
G --> generateImage Tool
G --> saveImage Tool

• The launcher invokes the agent which can call generateImage to create and save images as artifacts.

• Upon user request, the agent can call saveImage to load the artifact and save the image locally.

Important Implementation Notes

• Environment Variables: The Google Cloud project and location must be set as environment variables (GOOGLE_CLOUD_PROJECT and GOOGLE_CLOUD_LOCATION) for Vertex AI client initialization.

• Filename Handling: When saving locally, the code appends .png if the filename lacks that extension, ensuring correct file format.

• Error Propagation: Each step logs and returns errors to avoid silent failures.

• Output Directory: Images saved locally are placed under an output directory created relative to the current working directory.

• Tool Configuration: The function tools are wrapped with descriptive metadata (name, description) and integrated into the agent's toolset.

System Interaction References

• This file heavily relies on the LLM Integration and Agents framework ([80562]) for agent creation and configuration.

• It uses the Function Tools system ([80570]) to wrap the image generation and saving logic as callable tools.

• Artifact storage and retrieval are handled using the Artifact Management interfaces ([80557]), specifically the in-memory artifact service ([80580]).

• The agent lifecycle and execution are managed through Agent Lifecycle and Callbacks ([80573]) and REST API and Web Launchers ([80564]) for command-line execution.

• The Vertex AI model interaction is done via the Gemini model abstraction ([80562]).

This detailed documentation covers the purpose, components, and workflow of main.go in generating, storing, and saving images using a custom AI agent integrated with Google Cloud Vertex AI and the ADK agent framework.