internal/gemini/gemini.go

// Copyright 2024 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// Package gemini implements access to Google's Gemini model.
//
// [Client] implements [llm.Embedder]. Use [NewClient] to connect.
package gemini

import (
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"log/slog"
	"maps"
	"net/http"
	"slices"
	"strings"

	"github.com/google/generative-ai-go/genai"
	"golang.org/x/oscar/internal/httprr"
	"golang.org/x/oscar/internal/llm"
	"golang.org/x/oscar/internal/secret"
	"google.golang.org/api/option"
)

// Scrub is a request scrubber for use with [rsc.io/httprr].
func Scrub(req *http.Request) error {
	delete(req.Header, "x-goog-api-key")    // genai does not canonicalize
	req.Header.Del("X-Goog-Api-Key")        // in case it starts
	delete(req.Header, "x-goog-api-client") // contains version numbers
	req.Header.Del("X-Goog-Api-Client")

	if ctype := req.Header.Get("Content-Type"); ctype == "application/json" || strings.HasPrefix(ctype, "application/json;") {
		// Canonicalize JSON body.
		// google.golang.org/protobuf/internal/encoding.json
		// goes out of its way to randomize the JSON encodings
		// of protobuf messages by adding or not adding spaces
		// after commas. Derandomize by compacting the JSON.
		b := req.Body.(*httprr.Body)
		var buf bytes.Buffer
		if err := json.Compact(&buf, b.Data); err == nil {
			b.Data = buf.Bytes()
		}
	}
	return nil
}

// A Client represents a connection to Gemini.
type Client struct {
	slog  *slog.Logger
	genai *genai.Client
}

// NewClient returns a connection to Gemini, using the given logger and HTTP client.
// It expects to find a secret of the form "AIza..." or "user:AIza..." in sdb
// under the name "ai.google.dev".
func NewClient(ctx context.Context, lg *slog.Logger, sdb secret.DB, hc *http.Client) (*Client, error) {
	key, ok := sdb.Get("ai.google.dev")
	if !ok {
		return nil, fmt.Errorf("missing api key for ai.google.dev")
	}
	// If key is from .netrc, ignore user name.
	if _, pass, ok := strings.Cut(key, ":"); ok {
		key = pass
	}

	// Ideally this would use use “option.WithAPIKey(key), option.WithHTTPClient(hc),”
	// but using option.WithHTTPClient bypasses the code that passes along the API key.
	// Instead we make our own derived http.Client that re-adds the key.
	// And then we still have to say option.WithAPIKey("ignored") because
	// otherwise NewClient complains that we haven't passed in a key.
	// (If we pass in the key, it ignores it, but if we don't pass it in,
	// it complains that we didn't give it a key.)
	ai, err := genai.NewClient(ctx,
		option.WithAPIKey("ignored"),
		option.WithHTTPClient(withKey(hc, key)))
	if err != nil {
		return nil, err
	}

	return &Client{slog: lg, genai: ai}, nil
}

// withKey returns a new http.Client that is the same as hc
// except that it adds "x-goog-api-key: key" to every request.
func withKey(hc *http.Client, key string) *http.Client {
	c := *hc
	t := c.Transport
	if t == nil {
		t = http.DefaultTransport
	}
	c.Transport = &transportWithKey{t, key}
	return &c
}

// transportWithKey is the same as rt
// except that it adds "x-goog-api-key: key" to every request.
type transportWithKey struct {
	rt  http.RoundTripper
	key string
}

func (t *transportWithKey) RoundTrip(req *http.Request) (resp *http.Response, err error) {
	r := *req
	r.Header = maps.Clone(req.Header)
	r.Header["x-goog-api-key"] = []string{t.key}
	return t.rt.RoundTrip(&r)
}

const maxBatch = 100 // empirical limit

// EmbedDocs returns the vector embeddings for the docs,
// implementing [llm.Embedder].
func (c *Client) EmbedDocs(ctx context.Context, docs []llm.EmbedDoc) ([]llm.Vector, error) {
	model := c.genai.EmbeddingModel("text-embedding-004")
	var vecs []llm.Vector
	for docs := range slices.Chunk(docs, maxBatch) {
		b := model.NewBatch()
		for _, d := range docs {
			b.AddContentWithTitle(d.Title, genai.Text(d.Text))
		}
		resp, err := model.BatchEmbedContents(ctx, b)
		if err != nil {
			return vecs, err
		}
		for _, e := range resp.Embeddings {
			vecs = append(vecs, e.Values)
		}
	}
	return vecs, nil
}