vecna/pkg/embedclient/google.go

package embedclient

import (
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"net/http"
)

type googleClient struct {
	baseURL    string
	apiKey     string
	model      string
	httpClient *http.Client
}

// NewGoogle returns a Client that speaks the Google Gemini batchEmbedContents API.
func NewGoogle(baseURL, apiKey, model string, httpClient *http.Client) Client {
	if httpClient == nil {
		httpClient = http.DefaultClient
	}
	return &googleClient{baseURL: baseURL, apiKey: apiKey, model: model, httpClient: httpClient}
}

type googleBatchRequest struct {
	Requests []googleEmbedRequest `json:"requests"`
}

type googleEmbedRequest struct {
	Model   string        `json:"model"`
	Content googleContent `json:"content"`
}

type googleContent struct {
	Parts []googlePart `json:"parts"`
}

type googlePart struct {
	Text string `json:"text"`
}

type googleBatchResponse struct {
	Embeddings []struct {
		Values []float32 `json:"values"`
	} `json:"embeddings"`
}

func (c *googleClient) Embed(ctx context.Context, req Request) (Response, error) {
	requests := make([]googleEmbedRequest, len(req.Texts))
	for i, text := range req.Texts {
		requests[i] = googleEmbedRequest{
			Model:   "models/" + c.model,
			Content: googleContent{Parts: []googlePart{{Text: text}}},
		}
	}

	body, err := json.Marshal(googleBatchRequest{Requests: requests})
	if err != nil {
		return Response{}, fmt.Errorf("google embed marshal: %w", err)
	}

	url := fmt.Sprintf("%s/v1/models/%s:batchEmbedContents", c.baseURL, c.model)
	if c.apiKey != "" {
		url += "?key=" + c.apiKey
	}

	httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost, url, bytes.NewReader(body))
	if err != nil {
		return Response{}, fmt.Errorf("google embed request: %w", err)
	}
	httpReq.Header.Set("Content-Type", "application/json")

	resp, err := c.httpClient.Do(httpReq)
	if err != nil {
		return Response{}, fmt.Errorf("google embed do: %w", err)
	}
	defer func() { _ = resp.Body.Close() }()

	if resp.StatusCode != http.StatusOK {
		return Response{}, fmt.Errorf("google embed: unexpected status %d", resp.StatusCode)
	}

	var gResp googleBatchResponse
	if err := json.NewDecoder(resp.Body).Decode(&gResp); err != nil {
		return Response{}, fmt.Errorf("google embed decode: %w", err)
	}

	embeddings := make([][]float32, len(gResp.Embeddings))
	for i, e := range gResp.Embeddings {
		embeddings[i] = e.Values
	}

	return Response{
		Embeddings: embeddings,
		Model:      c.model,
	}, nil
}