add vertexai API client

cmd/vertexai/main.go

@@ -0,0 +1,65 @@
+package main
+
+import (
+	"context"
+	"flag"
+	"fmt"
+	"log"
+
+	"github.com/milosgajdos/go-embeddings/vertexai"
+	"golang.org/x/oauth2/google"
+)
+
+var (
+	input    string
+	model    string
+	truncate bool
+	taskType string
+	title    string
+)
+
+func init() {
+	flag.StringVar(&input, "input", "what is life", "input data")
+	flag.StringVar(&model, "model", string(vertexai.EmbedGeckoV2), "model name")
+	flag.BoolVar(&truncate, "truncate", false, "truncate type")
+	flag.StringVar(&taskType, "task-type", string(vertexai.RetrQueryTask), "task type")
+	flag.StringVar(&title, "title", "", "title: only relevant for retrival document tasks")
+}
+
+func main() {
+	flag.Parse()
+
+	ctx := context.Background()
+
+	ts, err := google.DefaultTokenSource(ctx, vertexai.Scopes)
+	if err != nil {
+		log.Fatalf("token source: %v", err)
+	}
+
+	c, err := vertexai.NewClient()
+	if err != nil {
+		log.Fatal(err)
+	}
+	c.WithTokenSrc(ts)
+	c.WithModelID(model)
+
+	embReq := &vertexai.EmbeddingRequest{
+		Instances: []vertexai.Instance{
+			{
+				Content:  input,
+				TaskType: vertexai.TaskType(taskType),
+				Title:    title,
+			},
+		},
+		Params: vertexai.Params{
+			AutoTruncate: truncate,
+		},
+	}
+
+	embs, err := c.Embeddings(context.Background(), embReq)
+	if err != nil {
+		log.Fatal(err)
+	}
+
+	fmt.Printf("got %d embeddings", len(embs))
+}

go.mod

@@ -1,3 +1,14 @@
 module github.com/milosgajdos/go-embeddings
 
 go 1.20
+
+require golang.org/x/oauth2 v0.15.0
+
+require (
+	cloud.google.com/go/compute v1.20.1 // indirect
+	cloud.google.com/go/compute/metadata v0.2.3 // indirect
+	github.com/golang/protobuf v1.5.3 // indirect
+	golang.org/x/net v0.19.0 // indirect
+	google.golang.org/appengine v1.6.7 // indirect
+	google.golang.org/protobuf v1.31.0 // indirect
+)

go.sum

@@ -0,0 +1,27 @@
+cloud.google.com/go/compute v1.20.1 h1:6aKEtlUiwEpJzM001l0yFkpXmUVXaN8W+fbkb2AZNbg=
+cloud.google.com/go/compute v1.20.1/go.mod h1:4tCnrn48xsqlwSAiLf1HXMQk8CONslYbdiEZc9FEIbM=
+cloud.google.com/go/compute/metadata v0.2.3 h1:mg4jlk7mCAj6xXp9UJ4fjI9VUI5rubuGBW5aJ7UnBMY=
+cloud.google.com/go/compute/metadata v0.2.3/go.mod h1:VAV5nSsACxMJvgaAuX6Pk2AawlZn8kiOGuCv6gTkwuA=
+github.com/golang/protobuf v1.3.1/go.mod h1:6lQm79b+lXiMfvg/cZm0SGofjICqVBUtrP5yJMmIC1U=
+github.com/golang/protobuf v1.5.0/go.mod h1:FsONVRAS9T7sI+LIUmWTfcYkHO4aIWwzhcaSAoJOfIk=
+github.com/golang/protobuf v1.5.3 h1:KhyjKVUg7Usr/dYsdSqoFveMYd5ko72D+zANwlG1mmg=
+github.com/golang/protobuf v1.5.3/go.mod h1:XVQd3VNwM+JqD3oG2Ue2ip4fOMUkwXdXDdiuN0vRsmY=
+github.com/google/go-cmp v0.5.5/go.mod h1:v8dTdLbMG2kIc/vJvl+f65V22dbkXbowE6jgT/gNBxE=
+github.com/google/go-cmp v0.5.9 h1:O2Tfq5qg4qc4AmwVlvv0oLiVAGB7enBSJ2x2DqQFi38=
+golang.org/x/crypto v0.0.0-20190308221718-c2843e01d9a2/go.mod h1:djNgcEr1/C05ACkg1iLfiJU5Ep61QUkGW8qpdssI0+w=
+golang.org/x/net v0.0.0-20190603091049-60506f45cf65/go.mod h1:HSz+uSET+XFnRR8LxR5pz3Of3rY3CfYBVs4xY44aLks=
+golang.org/x/net v0.19.0 h1:zTwKpTd2XuCqf8huc7Fo2iSy+4RHPd10s4KzeTnVr1c=
+golang.org/x/net v0.19.0/go.mod h1:CfAk/cbD4CthTvqiEl8NpboMuiuOYsAr/7NOjZJtv1U=
+golang.org/x/oauth2 v0.15.0 h1:s8pnnxNVzjWyrvYdFUQq5llS1PX2zhPXmccZv99h7uQ=
+golang.org/x/oauth2 v0.15.0/go.mod h1:q48ptWNTY5XWf+JNten23lcvHpLJ0ZSxF5ttTHKVCAM=
+golang.org/x/sys v0.0.0-20190215142949-d0b11bdaac8a/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY=
+golang.org/x/text v0.3.0/go.mod h1:NqM8EUOU14njkJ3fqMW+pc6Ldnwhi/IjpwHt7yyuwOQ=
+golang.org/x/text v0.3.2/go.mod h1:bEr9sfX3Q8Zfm5fL9x+3itogRgK3+ptLWKqgva+5dAk=
+golang.org/x/tools v0.0.0-20180917221912-90fa682c2a6e/go.mod h1:n7NCudcB/nEzxVGmLbDWY5pfWTLqBcC2KZ6jyYvM4mQ=
+golang.org/x/xerrors v0.0.0-20191204190536-9bdfabe68543/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0=
+google.golang.org/appengine v1.6.7 h1:FZR1q0exgwxzPzp/aF+VccGrSfxfPpkBqjIIEq3ru6c=
+google.golang.org/appengine v1.6.7/go.mod h1:8WjMMxjGQR8xUklV/ARdw2HLXBOI7O7uCIDZVag1xfc=
+google.golang.org/protobuf v1.26.0-rc.1/go.mod h1:jlhhOSvTdKEhbULTjvd4ARK9grFBp09yW+WbY/TyQbw=
+google.golang.org/protobuf v1.26.0/go.mod h1:9q0QmTI4eRPtz6boOQmLYwt+qCgq0jsYwAQnmE0givc=
+google.golang.org/protobuf v1.31.0 h1:g0LDEJHgrBl9N9r17Ru3sqWhkIx2NB67okBHPwC7hs8=
+google.golang.org/protobuf v1.31.0/go.mod h1:HV8QOd/L58Z+nl8r43ehVNZIU/HEI6OcFqwMG9pJV4I=

vertexai/client.go

@@ -0,0 +1,162 @@
+package vertexai
+
+import (
+	"bytes"
+	"context"
+	"encoding/json"
+	"fmt"
+	"io"
+	"net/http"
+	"os"
+
+	"golang.org/x/oauth2"
+)
+
+const (
+	// BaseURL is Vertex AI HTTP API base URL
+	BaseURL = "https://us-central1-aiplatform.googleapis.com/v1/projects"
+	// ModelURI is Vertex AI HTTP API model URI.
+	ModelURI = "locations/us-central1/publishers/google/models"
+	// EmbedAction is embedding API action.
+	EmbedAction = ":predict"
+)
+
+// Client is vertex AI HTTP API client.
+type Client struct {
+	token     string
+	tokenSrc  oauth2.TokenSource
+	projectID string
+	modelID   string
+	baseURL   string
+	hc        *http.Client
+}
+
+// NewClient creates a new HTTP client and returns it.
+// It reads the Google API token from VERTEXAI_TOKEN env var
+// just like the project ID is read from GOOGLE_PROJECT_ID env var
+// and uses the default Go http.Client.
+// You can override the default options by using the
+// client methods.
+func NewClient() (*Client, error) {
+	return &Client{
+		token:     os.Getenv("VERTEXAI_TOKEN"),
+		modelID:   os.Getenv("VERTEXAI_MODEL_ID"),
+		projectID: os.Getenv("GOOGLE_PROJECT_ID"),
+		baseURL:   BaseURL,
+		hc:        &http.Client{},
+	}, nil
+}
+
+// WithToken sets the API key.
+func (c *Client) WithToken(token string) *Client {
+	c.token = token
+	return c
+}
+
+// WithTokenSrc sets the API token source.
+func (c *Client) WithTokenSrc(ts oauth2.TokenSource) *Client {
+	c.tokenSrc = ts
+	return c
+}
+
+// WithProjectID sets the project ID.
+func (c *Client) WithProjectID(id string) *Client {
+	c.projectID = id
+	return c
+}
+
+// WithModelID sets the model ID.
+func (c *Client) WithModelID(id string) *Client {
+	c.modelID = id
+	return c
+}
+
+// WithBaseURL sets the API base URL.
+func (c *Client) WithBaseURL(baseURL string) *Client {
+	c.baseURL = baseURL
+	return c
+}
+
+// WithHTTPClient sets the HTTP client.
+func (c *Client) WithHTTPClient(httpClient *http.Client) *Client {
+	c.hc = httpClient
+	return c
+}
+
+// ReqOption is http requestion functional option.
+type ReqOption func(*http.Request)
+
+// WithSetHeader sets the header key to value val.
+func WithSetHeader(key, val string) ReqOption {
+	return func(req *http.Request) {
+		if req.Header == nil {
+			req.Header = make(http.Header)
+		}
+		req.Header.Set(key, val)
+	}
+}
+
+// WithAddHeader adds the val to key header.
+func WithAddHeader(key, val string) ReqOption {
+	return func(req *http.Request) {
+		if req.Header == nil {
+			req.Header = make(http.Header)
+		}
+		req.Header.Add(key, val)
+	}
+}
+
+func (c *Client) newRequest(ctx context.Context, method, url string, body io.Reader, opts ...ReqOption) (*http.Request, error) {
+	if ctx == nil {
+		ctx = context.Background()
+	}
+	if body == nil {
+		body = &bytes.Reader{}
+	}
+
+	req, err := http.NewRequestWithContext(ctx, method, url, body)
+	if err != nil {
+		return nil, err
+	}
+
+	for _, setOption := range opts {
+		setOption(req)
+	}
+
+	if c.token == "" {
+		var err error
+		c.token, err = GetToken(c.tokenSrc)
+		if err != nil {
+			return nil, err
+		}
+	}
+
+	req.Header.Set("Authorization", fmt.Sprintf("Bearer %s", c.token))
+	req.Header.Set("Accept", "application/json; charset=utf-8")
+	if body != nil {
+		// if no content-type is specified we default to json
+		if ct := req.Header.Get("Content-Type"); len(ct) == 0 {
+			req.Header.Set("Content-Type", "application/json; charset=utf-8")
+		}
+	}
+
+	return req, nil
+}
+
+func (c *Client) doRequest(req *http.Request) (*http.Response, error) {
+	resp, err := c.hc.Do(req)
+	if err != nil {
+		return nil, err
+	}
+	if resp.StatusCode >= http.StatusOK && resp.StatusCode < http.StatusBadRequest {
+		return resp, nil
+	}
+	defer resp.Body.Close()
+
+	var apiErr APIError
+	if err := json.NewDecoder(resp.Body).Decode(&apiErr); err != nil {
+		return nil, err
+	}
+
+	return nil, apiErr
+}

vertexai/embedding.go

@@ -0,0 +1,103 @@
+package vertexai
+
+import (
+	"bytes"
+	"context"
+	"encoding/json"
+	"io"
+	"net/http"
+	"net/url"
+)
+
+// Embedding is cohere API vector embedding.
+type Embedding struct {
+	Vector []float64 `json:"vector"`
+}
+
+// EmbeddingRequest sent to API endpoint.
+// https://cloud.google.com/vertex-ai/docs/generative-ai/embeddings/get-text-embeddings#generative-ai-get-text-embedding-drest
+type EmbeddingRequest struct {
+	Instances []Instance `json:"instances"`
+	Params    Params     `json:"parameters"`
+}
+
+// NOTE: Title is only valid with TaskType set to RetrDocTask
+// https://cloud.google.com/vertex-ai/docs/generative-ai/embeddings/get-text-embeddings#api_changes_to_models_released_on_or_after_august_2023
+type Instance struct {
+	TaskType TaskType `json:"task_type"`
+	Title    string   `json:"title,omitempty"`
+	Content  string   `json:"content"`
+}
+
+// Params are additional API request parameters passed via body.
+type Params struct {
+	// If set to false, text that exceeds the token limit (3.072)
+	// causes the request to fail. The default value is true
+	AutoTruncate bool `json:"autoTruncate"`
+}
+
+// EmbedddingResponse received from API endpoint.
+// https://cloud.google.com/vertex-ai/docs/generative-ai/model-reference/text-embeddings#response_body
+type EmbedddingResponse struct {
+	Predictions []Predictions  `json:"predictions"`
+	Metadata    map[string]any `json:"metadata"`
+}
+
+// Predictions is the generated response
+type Predictions struct {
+	Embeddings struct {
+		Values     []float64  `json:"values"`
+		Statistics Statistics `json:"statistics"`
+	} `json:"embeddings"`
+}
+
+// Statistics define the statistics for a text embedding
+type Statistics struct {
+	TokenCount int  `json:"token_count"`
+	Truncated  bool `json:"truncated"`
+}
+
+func ToEmbeddings(r io.Reader) ([]*Embedding, error) {
+	var resp EmbedddingResponse
+	if err := json.NewDecoder(r).Decode(&resp); err != nil {
+		return nil, err
+	}
+	embs := make([]*Embedding, 0, len(resp.Predictions))
+	for _, p := range resp.Predictions {
+		floats := make([]float64, len(p.Embeddings.Values))
+		copy(floats, p.Embeddings.Values)
+		emb := &Embedding{
+			Vector: floats,
+		}
+		embs = append(embs, emb)
+	}
+	return embs, nil
+}
+
+// Embeddings returns embeddings for every object in EmbeddingRequest.
+func (c *Client) Embeddings(ctx context.Context, embReq *EmbeddingRequest) ([]*Embedding, error) {
+	u, err := url.Parse(c.baseURL + "/" + c.projectID + "/" + ModelURI + "/" + c.modelID + EmbedAction)
+	if err != nil {
+		return nil, err
+	}
+
+	var body = &bytes.Buffer{}
+	enc := json.NewEncoder(body)
+	enc.SetEscapeHTML(false)
+	if err := enc.Encode(embReq); err != nil {
+		return nil, err
+	}
+
+	req, err := c.newRequest(ctx, http.MethodPost, u.String(), body)
+	if err != nil {
+		return nil, err
+	}
+
+	resp, err := c.doRequest(req)
+	if err != nil {
+		return nil, err
+	}
+	defer resp.Body.Close()
+
+	return ToEmbeddings(resp.Body)
+}

vertexai/error.go

@@ -0,0 +1,19 @@
+package vertexai
+
+import "encoding/json"
+
+type APIError struct {
+	RespError struct {
+		Code    int    `json:"code"`
+		Message string `json:"message"`
+		Status  string `json:"status"`
+	} `json:"error"`
+}
+
+func (e APIError) Error() string {
+	b, err := json.Marshal(e)
+	if err != nil {
+		return "unknown error"
+	}
+	return string(b)
+}