Cleanup

* updated README
* added more linters
* NewClient no longer returns error
* updated godoc comments

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

Author: milosgajdos

.golangci.yml

@@ -12,6 +12,15 @@ linters:
     - errcheck
     - goimports
     - prealloc
+    - gosec
+    - goconst
+    - gosimple
+    - errname
+    - gocritic
+    - tparallel
+    - tenv
+    - unconvert
+    - whitespace
 
 run:
   deadline: 2m

README.md

@@ -5,3 +5,10 @@
 [![License: Apache-2.0](https://img.shields.io/badge/License-Apache--2.0-blue.svg)](https://opensource.org/licenses/Apache-2.0)
 
 This project provides an implementation for fetching embeddings from various LLMs.
+
+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)
+
+There are also simple command line tools provided by this project that let you query the APIs for text embeddings passed in via cli flags.

cmd/cohere/main.go

@@ -26,10 +26,7 @@ func init() {
 func main() {
 	flag.Parse()
 
-	c, err := cohere.NewClient()
-	if err != nil {
-		log.Fatal(err)
-	}
+	c := cohere.NewClient()
 
 	embReq := &cohere.EmbeddingRequest{
 		Texts:     []string{input},

cmd/openai/main.go

@@ -24,10 +24,7 @@ func init() {
 func main() {
 	flag.Parse()
 
-	c, err := openai.NewClient()
-	if err != nil {
-		log.Fatal(err)
-	}
+	c := openai.NewClient()
 
 	embReq := &openai.EmbeddingRequest{
 		Input:          input,

cmd/vertexai/main.go

@@ -36,12 +36,9 @@ func main() {
 		log.Fatalf("token source: %v", err)
 	}
 
-	c, err := vertexai.NewClient()
-	if err != nil {
-		log.Fatal(err)
-	}
-	c.WithTokenSrc(ts)
-	c.WithModelID(model)
+	c := vertexai.NewClient().
+		WithTokenSrc(ts).
+		WithModelID(model)
 
 	embReq := &vertexai.EmbeddingRequest{
 		Instances: []vertexai.Instance{

cohere/client.go

@@ -25,18 +25,17 @@ type Client struct {
 	hc      *http.Client
 }
 
-// NewClient creates a new HTTP client and returns it.
-// It reads the Cohere API key from COHERE_API_KEY env var
-// and uses the default Go http.Client.
-// You can override the default options by using the
-// client methods.
-func NewClient() (*Client, error) {
+// NewClient creates a new HTTP API client and returns it.
+// By default it reads the Cohere API key from COHERE_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() *Client {
 	return &Client{
 		apiKey:  os.Getenv("COHERE_API_KEY"),
 		baseURL: BaseURL,
 		version: EmbedAPIVersion,
 		hc:      &http.Client{},
-	}, nil
+	}
 }
 
 // WithAPIKey sets the API key.

cohere/cohere.go

@@ -1,6 +1,6 @@
 package cohere
 
-// Model is embedding model.
+// Model is an embedding model.
 type Model string
 
 const (
@@ -13,7 +13,7 @@ const (
 	MultiLingV2      Model = "embed-multilingual-v2.0"
 )
 
-// InputType is embedding input type.
+// InputType is an embedding input type.
 type InputType string
 
 const (
@@ -24,6 +24,8 @@ const (
 )
 
 // Truncate controls input truncating.
+// It controls how the API handles inputs
+// longer than the maximum token length (recommended: <512)
 type Truncate string
 
 const (

cohere/embedding.go

@@ -9,7 +9,7 @@ import (
 	"net/url"
 )
 
-// Embedding is cohere API vector embedding.
+// Embedding is vector embedding.
 type Embedding struct {
 	Vector []float64 `json:"vector"`
 }
@@ -22,13 +22,13 @@ type EmbeddingRequest struct {
 	Truncate  Truncate  `json:"truncate,omitempty"`
 }
 
-// EmbedddingResponse received from API endpoint.
+// EmbedddingResponse received from API.
 type EmbedddingResponse struct {
 	Embeddings [][]float64 `json:"embeddings"`
 	Meta       *Meta       `json:"meta,omitempty"`
 }
 
-// Meta stores API response metadata
+// Meta stores API response metadata.
 type Meta struct {
 	APIVersion *APIVersion `json:"api_version,omitempty"`
 }
@@ -38,7 +38,9 @@ type APIVersion struct {
 	Version string `json:"version"`
 }
 
-func ToEmbeddings(r io.Reader) ([]*Embedding, error) {
+// toEmbeddings decodes the raw API response,
+// parses it into a slice of embeddings and returns it.
+func toEmbeddings(r io.Reader) ([]*Embedding, error) {
 	var resp EmbedddingResponse
 	if err := json.NewDecoder(r).Decode(&resp); err != nil {
 		return nil, err
@@ -80,5 +82,5 @@ func (c *Client) Embeddings(ctx context.Context, embReq *EmbeddingRequest) ([]*E
 	}
 	defer resp.Body.Close()
 
-	return ToEmbeddings(resp.Body)
+	return toEmbeddings(resp.Body)
 }

cohere/error.go

@@ -2,10 +2,12 @@ package cohere
 
 import "encoding/json"
 
+// APIError is Cohere API error.
 type APIError struct {
 	Message string `json:"message"`
 }
 
+// Error implements errors interface.
 func (e APIError) Error() string {
 	b, err := json.Marshal(e)
 	if err != nil {

openai/client.go

@@ -13,13 +13,13 @@ import (
 const (
 	// BaseURL is OpenAI HTTP API base URL.
 	BaseURL = "https://api.openai.com"
-	// EmbedAPIVersion is the latest stable embedding API version.
+	// EmbedAPIVersion is the latest stable embeddings API version.
 	EmbedAPIVersion = "v1"
-	// Org header
+	// OrgHeader is an Organization header
 	OrgHeader = "OpenAI-Organization"
 )
 
-// Client is OpenAI HTTP API client.
+// Client is an OpenAI HTTP API client.
 type Client struct {
 	apiKey  string
 	baseURL string
@@ -28,19 +28,18 @@ type Client struct {
 	hc      *http.Client
 }
 
-// NewClient creates a new HTTP client and returns it.
-// It reads the OpenAI API key from OPENAI_API_KEY env var
-// and uses the default Go http.Client.
-// You can override the default options by using the
-// client methods.
-func NewClient() (*Client, error) {
+// NewClient creates a new HTTP API client and returns it.
+// By default it reads the OpenAI API key from OPENAI_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() *Client {
 	return &Client{
 		apiKey:  os.Getenv("OPENAI_API_KEY"),
 		baseURL: BaseURL,
 		version: EmbedAPIVersion,
 		orgID:   "",
 		hc:      &http.Client{},
-	}, nil
+	}
 }
 
 // WithAPIKey sets the API key.