feature: Add support for VoyageAI embeddings (#22)

See: https://docs.voyageai.com/reference/embeddings-api

Signed-off-by: Milos Gajdos <[email protected]>

.gitignore

@@ -13,3 +13,4 @@
 
 # Dependency directories (remove the comment below to include it)
 # vendor/
+.env

README.md

@@ -10,11 +10,12 @@ Currently supported APIs:
 * [x] [OpenAI](https://platform.openai.com/docs/api-reference/embeddings)
 * [x] [Cohere AI](https://docs.cohere.com/reference/embed)
 * [x] [Google Vertex AI](https://cloud.google.com/vertex-ai/docs/generative-ai/embeddings/get-text-embeddings)
+* [x] [VoyageAI](https://docs.voyageai.com/reference/embeddings-api)
 
 You can find sample programs that demonstrate how to use the client packages to fetch the embeddings in `cmd` directory of this project.
 
 Finally, the `document` package provides an implementation of simple document text splitters, heavily inspired by the popular [Langchain framework](https://github.com/langchain-ai/langchain).
-It's essentially a Go rewrite of character and recursive character text splitters.
+It's essentially a Go rewrite of character and recursive character text splitters from the Langchain framework with minor modifications, but more or less identical results.
 
 ## Environment variables
 
@@ -33,6 +34,7 @@ Each client package lets you initialize a default API client for a specific embe
 * `VERTEXAI_TOKEN`: Google Vertex AI API token (can be fetch by `gcloud auth print-access-token` once you've authenticated)
 * `VERTEXAI_MODEL_ID`: Embeddings model (at the moment only `textembedding-gecko@00` or `multimodalembedding@001` are available)
 * `GOOGLE_PROJECT_ID`: Google Project ID
+* `VOYAGE_API_KEY`: VoyageAI API key
 
 ## nix
 

cmd/voyage/main.go

@@ -0,0 +1,44 @@
+package main
+
+import (
+ "context"
+ "flag"
+ "fmt"
+ "log"
+
+ "github.com/milosgajdos/go-embeddings/voyage"
+)
+
+var (
+ input string
+ model string
+ truncation bool
+ inputType string
+)
+
+func init() {
+ flag.StringVar(&input, "input", "what is life", "input data")
+ flag.StringVar(&model, "model", voyage.VoyageV2.String(), "model name")
+ flag.StringVar(&inputType, "input-type", voyage.DocInput.String(), "input type")
+ flag.BoolVar(&truncation, "truncate", false, "truncate type")
+}
+
+func main() {
+ flag.Parse()
+
+ c := voyage.NewClient()
+
+ embReq := &voyage.EmbeddingRequest{
+ Input: []string{input},
+ Model: voyage.Model(model),
+ InputType: voyage.InputType(inputType),
+ Truncation: truncation,
+ }
+
+ embs, err := c.Embed(context.Background(), embReq)
+ if err != nil {
+ log.Fatal(err)
+ }
+
+ fmt.Printf("got %d embeddings", len(embs))
+}

embeddings.go

@@ -1,6 +1,12 @@
 package embeddings
 
-import "context"
+import (
+ "context"
+ "encoding/base64"
+ "encoding/binary"
+ "fmt"
+ "math"
+)
 
 // Embedder fetches embeddings.
 type Embedder[T any] interface {
@@ -21,3 +27,27 @@ func (e Embedding) ToFloat32() []float32 {
  }
  return floats
 }
+
+// Base64 is base64 encoded embedding string.
+type Base64 string
+
+// Decode decodes base64 encoded string into a slice of floats.
+func (s Base64) Decode() ([]float64, error) {
+ decoded, err := base64.StdEncoding.DecodeString(string(s))
+ if err != nil {
+ return nil, err
+ }
+
+ if len(decoded)%8 != 0 {
+ return nil, fmt.Errorf("invalid base64 encoded string length")
+ }
+
+ floats := make([]float64, len(decoded)/8)
+
+ for i := 0; i < len(floats); i++ {
+ bits := binary.LittleEndian.Uint64(decoded[i*8 : (i+1)*8])
+ floats[i] = math.Float64frombits(bits)
+ }
+
+ return floats, nil
+}

openai/embedding.go

@@ -3,43 +3,15 @@ package openai
 import (
  "bytes"
  "context"
- "encoding/base64"
- "encoding/binary"
  "encoding/json"
- "fmt"
  "io"
- "math"
  "net/http"
  "net/url"
 
  "github.com/milosgajdos/go-embeddings"
  "github.com/milosgajdos/go-embeddings/request"
 )
 
-// EmbeddingString is base64 encoded embedding.
-type EmbeddingString string
-
-// Decode decodes base64 encoded string into a slice of floats.
-func (s EmbeddingString) Decode() ([]float64, error) {
- decoded, err := base64.StdEncoding.DecodeString(string(s))
- if err != nil {
- return nil, err
- }
-
- if len(decoded)%8 != 0 {
- return nil, fmt.Errorf("invalid base64 encoded string length")
- }
-
- floats := make([]float64, len(decoded)/8)
-
- for i := 0; i < len(floats); i++ {
- bits := binary.LittleEndian.Uint64(decoded[i*8 : (i+1)*8])
- floats[i] = math.Float64frombits(bits)
- }
-
- return floats, nil
-}
-
 // Usage tracks API token usage.
 type Usage struct {
  PromptTokens int `json:"prompt_tokens"`
@@ -110,7 +82,7 @@ func toEmbeddingResp[T any](resp io.Reader) (*EmbeddingResponse, error) {
  }
 
  switch e := any(data).(type) {
- case *EmbeddingResponseGen[EmbeddingString]:
+ case *EmbeddingResponseGen[embeddings.Base64]:
  embData := make([]Data, 0, len(e.Data))
  for _, d := range e.Data {
  floats, err := d.Embedding.Decode()
@@ -181,7 +153,7 @@ func (c *Client) Embed(ctx context.Context, embReq *EmbeddingRequest) ([]*embedd
 
  switch embReq.EncodingFormat {
  case EncodingBase64:
- embs, err = toEmbeddingResp[EmbeddingResponseGen[EmbeddingString]](resp.Body)
+ embs, err = toEmbeddingResp[EmbeddingResponseGen[embeddings.Base64]](resp.Body)
  case EncodingFloat:
  embs, err = toEmbeddingResp[EmbeddingResponseGen[[]float64]](resp.Body)
  default:

voyage/client.go

@@ -0,0 +1,85 @@
+package voyage
+
+import (
+ "os"
+
+ "github.com/milosgajdos/go-embeddings"
+ "github.com/milosgajdos/go-embeddings/client"
+)
+
+const (
+ // BaseURL is VoyageAI HTTP API base URL.
+ BaseURL = "https://api.voyageai.com"
+ // EmbedAPIVersion is the latest stable embedding API version.
+ EmbedAPIVersion = "v1"
+)
+
+// Client is Voyage HTTP API client.
+type Client struct {
+ opts Options
+}
+
+// Options are client options
+type Options struct {
+ APIKey string
+ BaseURL string
+ Version string
+ HTTPClient *client.HTTP
+}
+
+// Option is functional graph option.
+type Option func(*Options)
+
+// NewClient creates a new HTTP API client and returns it.
+// By default it reads the Voyage API key from VOYAGE_API_KEY
+// env var and uses the default Go http.Client for making API requests.
+// You can override the default options via the client methods.
+func NewClient(opts ...Option) *Client {
+ options := Options{
+ APIKey: os.Getenv("VOYAGE_API_KEY"),
+ BaseURL: BaseURL,
+ Version: EmbedAPIVersion,
+ HTTPClient: client.NewHTTP(),
+ }
+
+ for _, apply := range opts {
+ apply(&options)
+ }
+
+ return &Client{
+ opts: options,
+ }
+}
+
+// NewEmbedder creates a client that implements embeddings.Embedder
+func NewEmbedder(opts ...Option) embeddings.Embedder[*EmbeddingRequest] {
+ return NewClient(opts...)
+}
+
+// WithAPIKey sets the API key.
+func WithAPIKey(apiKey string) Option {
+ return func(o *Options) {
+ o.APIKey = apiKey
+ }
+}
+
+// WithBaseURL sets the API base URL.
+func WithBaseURL(baseURL string) Option {
+ return func(o *Options) {
+ o.BaseURL = baseURL
+ }
+}
+
+// WithVersion sets the API version.
+func WithVersion(version string) Option {
+ return func(o *Options) {
+ o.Version = version
+ }
+}
+
+// WithHTTPClient sets the HTTP client.
+func WithHTTPClient(httpClient *client.HTTP) Option {
+ return func(o *Options) {
+ o.HTTPClient = httpClient
+ }
+}

voyage/client_test.go

@@ -0,0 +1,52 @@
+package voyage
+
+import (
+ "testing"
+
+ "github.com/milosgajdos/go-embeddings/client"
+ "github.com/stretchr/testify/assert"
+)
+
+const (
+ voyageAPIKey = "somekey"
+)
+
+func TestClient(t *testing.T) {
+ t.Setenv("VOYAGE_API_KEY", voyageAPIKey)
+
+ t.Run("API key", func(t *testing.T) {
+ c := NewClient()
+ assert.Equal(t, c.opts.APIKey, voyageAPIKey)
+
+ testVal := "foo"
+ c = NewClient(WithAPIKey(testVal))
+ assert.Equal(t, c.opts.APIKey, testVal)
+ })
+
+ t.Run("BaseURL", func(t *testing.T) {
+ c := NewClient()
+ assert.Equal(t, c.opts.BaseURL, BaseURL)
+
+ testVal := "http://foo"
+ c = NewClient(WithBaseURL(testVal))
+ assert.Equal(t, c.opts.BaseURL, testVal)
+ })
+
+ t.Run("version", func(t *testing.T) {
+ c := NewClient()
+ assert.Equal(t, c.opts.Version, EmbedAPIVersion)
+
+ testVal := "v3"
+ c = NewClient(WithVersion(testVal))
+ assert.Equal(t, c.opts.Version, testVal)
+ })
+
+ t.Run("http client", func(t *testing.T) {
+ c := NewClient()
+ assert.NotNil(t, c.opts.HTTPClient)
+
+ testVal := client.NewHTTP()
+ c = NewClient(WithHTTPClient(testVal))
+ assert.NotNil(t, c.opts.HTTPClient)
+ })
+}