persona-community-2/pkg/gemini/client.go

// Package gemini provides a Go client for Google's Generative AI APIs.
//
// This package wraps the official google.golang.org/genai SDK to provide
// a simplified interface for image generation (Imagen) and video generation (Veo).
//
// Basic usage:
//
//	client, err := gemini.NewClient(gemini.Config{
//	    APIKey: os.Getenv("GEMINI_API_KEY"),
//	})
//	if err != nil {
//	    log.Fatal(err)
//	}
//	defer client.Close()
//
//	// Generate an image
//	resp, err := client.GenerateImage(ctx, gemini.ImageRequest{
//	    Prompt: "A serene Japanese garden",
//	})
//
// The client handles authentication and provides methods for both image and
// video generation with configurable models and parameters.
package gemini

import (
	"context"
	"fmt"
	"log/slog"
	"time"

	"google.golang.org/genai"
)

const (
	defaultTimeout        = 120 * time.Second
	defaultVideoPollDelay = 10 * time.Second
	defaultVideoMaxWait   = 5 * time.Minute

	// Retry configuration defaults
	defaultMaxRetries   = 3
	defaultInitialDelay = 100 * time.Millisecond
	defaultMaxDelay     = 2 * time.Second
)

// Config holds configuration options for the Gemini client
type Config struct {
	APIKey         string        // Required: API key for authentication
	Timeout        time.Duration // Optional: defaults to 120s
	VideoPollDelay time.Duration // Optional: delay between video status polls, defaults to 10s
	VideoMaxWait   time.Duration // Optional: max time to wait for video generation, defaults to 5m
	Logger         *slog.Logger  // Optional: defaults to slog.Default()

	// Retry configuration for transient errors (5xx, timeouts)
	MaxRetries   int           // Optional: max retry attempts, defaults to 3 (0 disables retry)
	InitialDelay time.Duration // Optional: initial delay between retries, defaults to 100ms
	MaxDelay     time.Duration // Optional: max delay between retries, defaults to 2s
}

// Client is the Gemini API client.
// Supports automatic retry with exponential backoff for transient errors.
type Client struct {
	genaiClient  *genai.Client
	config       *Config
	logger       *slog.Logger
	maxRetries   int
	initialDelay time.Duration
	maxDelay     time.Duration
}

// NewClient creates a new Gemini API client with automatic retry for transient errors.
func NewClient(ctx context.Context, config Config) (*Client, error) {
	if config.APIKey == "" {
		return nil, fmt.Errorf("%w: API key is required", ErrInvalidConfig)
	}

	if config.Timeout == 0 {
		config.Timeout = defaultTimeout
	}

	if config.VideoPollDelay == 0 {
		config.VideoPollDelay = defaultVideoPollDelay
	}

	if config.VideoMaxWait == 0 {
		config.VideoMaxWait = defaultVideoMaxWait
	}

	if config.Logger == nil {
		config.Logger = slog.Default()
	}

	// Set retry defaults
	maxRetries := config.MaxRetries
	if maxRetries == 0 {
		maxRetries = defaultMaxRetries
	}

	initialDelay := config.InitialDelay
	if initialDelay == 0 {
		initialDelay = defaultInitialDelay
	}

	maxDelay := config.MaxDelay
	if maxDelay == 0 {
		maxDelay = defaultMaxDelay
	}

	// Create genai client with API key
	genaiClient, err := genai.NewClient(ctx, &genai.ClientConfig{
		APIKey:  config.APIKey,
		Backend: genai.BackendGeminiAPI,
	})
	if err != nil {
		return nil, fmt.Errorf("create genai client: %w", err)
	}

	return &Client{
		genaiClient:  genaiClient,
		config:       &config,
		logger:       config.Logger,
		maxRetries:   maxRetries,
		initialDelay: initialDelay,
		maxDelay:     maxDelay,
	}, nil
}

// Close closes the underlying genai client
func (c *Client) Close() error {
	// The genai client doesn't have a Close method, but we keep this for future compatibility
	return nil
}

// retryWithBackoff executes fn with exponential backoff for transient errors.
// Returns the result of fn if successful, or the last error if all retries fail.
func (c *Client) retryWithBackoff(ctx context.Context, operation string, fn func() error) error {
	var lastErr error

	for attempt := 0; attempt <= c.maxRetries; attempt++ {
		if err := ctx.Err(); err != nil {
			return err
		}

		lastErr = fn()
		if lastErr == nil {
			return nil
		}

		// Only retry retryable errors
		if !IsRetryableError(lastErr) {
			return lastErr
		}

		// Don't retry on last attempt
		if attempt == c.maxRetries {
			break
		}

		// Calculate backoff delay
		delay := c.initialDelay * time.Duration(1<<attempt)
		if delay > c.maxDelay {
			delay = c.maxDelay
		}

		c.logger.Warn("retrying after transient error",
			"operation", operation,
			"attempt", attempt+1,
			"max_retries", c.maxRetries,
			"delay", delay,
			"error", lastErr,
		)

		// Wait for backoff or context cancellation
		select {
		case <-ctx.Done():
			return ctx.Err()
		case <-time.After(delay):
		}
	}

	return lastErr
}

// Health checks if the client can communicate with the Gemini API
// by listing available models
func (c *Client) Health(ctx context.Context) error {
	iter, err := c.genaiClient.Models.List(ctx, nil)
	if err != nil {
		return fmt.Errorf("health check failed: %w", err)
	}

	// Try to get at least one model to verify connectivity
	_, err = iter.Next(ctx)
	if err != nil {
		return fmt.Errorf("health check failed: %w", err)
	}

	return nil
}