panic-cli/internal/ai/gemini.go

package ai

import (
	"context"
	"fmt"
	"io"

	"google.golang.org/genai"
)

type GeminiModel string

const (
	GeminiModel1p5Flash     GeminiModel = "gemini-1.5-flash"
	GeminiModel20FlashExp   GeminiModel = "gemini-2.0-flash-exp"
	GeminiModel2p5Flash     GeminiModel = "gemini-2.5-flash"
	GeminiModel2p5FlashLite GeminiModel = "gemini-2.5-flash-lite"
)

type GeminiProvider struct {
	client *genai.Client
	model  GeminiModel
	config *genai.GenerateContentConfig
}

func NewGeminiProvider(ctx context.Context, apiKey string, modelName GeminiModel, config *genai.GenerateContentConfig) (*GeminiProvider, error) {
	cc := genai.ClientConfig{
		APIKey: apiKey,
	}
	client, err := genai.NewClient(ctx, &cc)
	if err != nil {
		return nil, err
	}

	return &GeminiProvider{client: client, model: modelName, config: config}, nil
}

func (p *GeminiProvider) Summarize(ctx context.Context, reader io.Reader, writer io.Writer) error {

	inputBytes, err := io.ReadAll(reader)
	if err != nil {
		return fmt.Errorf("failed to read input: %w", err)
	}

	prompt := fmt.Sprintf("Summarize the following:\n%s", string(inputBytes))
	stream := p.client.Models.GenerateContentStream(
		ctx,
		string(p.model),
		genai.Text(prompt),
		p.config,
	)

	for chunk, err := range stream {
		if err != nil {
			return fmt.Errorf("failed to generate content: %w", err)
		}
		part := chunk.Candidates[0].Content.Parts[0]
		_, err = writer.Write([]byte(part.Text))
		if err != nil {
			return fmt.Errorf("failed to write output: %w", err)
		}
	}
	return nil
}