geminiclient.go

package geminiclient

import (
	"context"
	"errors"
	"fmt"
	"net/http"

	"github.com/google/generative-ai-go/genai"
	"google.golang.org/api/iterator"
	"google.golang.org/api/option"
)

const (
	defaultBaseURL     = ""
	defaultModel       = "gemini-pro"
	defaultVisionModel = "gemini-pro-vision"
	defaultEmbedModel  = "embedding-001"
)

// ErrEmptyResponse is returned when the OpenAI API returns an empty response.
var ErrEmptyResponse = errors.New("empty response")
var ErrMissToken = errors.New("api key is not set")
var ErrRateLimitResponse = errors.New("rate limit reached for text-embedding-ada-002 in organization")

// Client is a client for the OpenAI API.
type Client struct {
	token      string
	Model      string
	baseURL    string
	embedModel string
}

// New returns a new Gemini client.
func New(token string, model string, baseURL string) (*Client, error) {
	c := &Client{
		token:      token,
		Model:      model,
		baseURL:    baseURL,
		embedModel: defaultEmbedModel,
	}
	if c.Model == "" {
		c.Model = defaultModel
	}

	if c.baseURL == "" {
		c.baseURL = defaultBaseURL
	}

	return c, nil
}

// Completion is a completion.
type Completion struct {
	Text string `json:"text"`
}

// CompletionRequest is a request to complete a completion.
type CompletionRequest struct {
	Model            string   `json:"model"`
	Prompt           string   `json:"prompt"`
	Temperature      float64  `json:"temperature,omitempty"`
	MaxTokens        int      `json:"max_tokens,omitempty"`
	N                int      `json:"n,omitempty"`
	FrequencyPenalty float64  `json:"frequency_penalty,omitempty"`
	PresencePenalty  float64  `json:"presence_penalty,omitempty"`
	TopP             float64  `json:"top_p,omitempty"`
	StopWords        []string `json:"stop,omitempty"`
	Images           [][]byte `json:"images"`

	// StreamingFunc is a function to be called for each chunk of a streaming response.
	// Return an error to stop streaming early.
	StreamingFunc func(ctx context.Context, chunk []byte) error `json:"-"`
}

// CreateCompletion creates a completion.
func (c *Client) CreateCompletion(ctx context.Context, r *CompletionRequest) (*Completion, error) {
	// Access your API key as an environment variable (see "Set up your API key" above)
	opts := make([]option.ClientOption, 0, 2)
	opts = append(opts, option.WithAPIKey(c.token))
	if c.baseURL != "" {
		opts = append(opts, option.WithEndpoint(c.baseURL))
	}

	client, err := genai.NewClient(ctx, opts...)
	if err != nil {
		return nil, err
	}
	defer client.Close()

	model := client.GenerativeModel(c.Model)

	if len(r.Images) > 0 {
		model = client.GenerativeModel(defaultVisionModel)
	}
	model.StopSequences = r.StopWords
	// model.SetTemperature(float32(r.Temperature))
	// model.SetTopP(float32(r.TopP))
	if r.MaxTokens > 0 {
		model.SetMaxOutputTokens(int32(r.MaxTokens))
	}

	blobs := make([]genai.Part, 0, len(r.Images)+1)
	for _, image := range r.Images {
		blobs = append(blobs, genai.ImageData("jpeg", image))
	}

	blobs = append(blobs, genai.Text(r.Prompt))

	iter := model.GenerateContentStream(ctx, blobs...)

	content := ""
	for {
		resp, err := iter.Next()
		if err == iterator.Done {
			break
		}
		if err != nil {
			return nil, err
		}
		stream := fmt.Sprintf("%s", resp.Candidates[0].Content.Parts[0])
		if r.StreamingFunc != nil {
			_ = r.StreamingFunc(ctx, []byte(stream))
		}
		content += stream
	}
	return &Completion{
		Text: content,
	}, nil
}

func (c *Client) setHeaders(req *http.Request) {
	req.Header.Set("Content-Type", "application/json")
}