feat: add ollama embeddings API (#23)

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

README.md

@@ -8,9 +8,10 @@ This project provides an implementation of API clients for fetching embeddings f
 
 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] [Cohere](https://docs.cohere.com/reference/embed)
+* [x] [Google Vertex](https://cloud.google.com/vertex-ai/docs/generative-ai/embeddings/get-text-embeddings)
 * [x] [VoyageAI](https://docs.voyageai.com/reference/embeddings-api)
+* [x] [Ollama](https://ollama.com/)
 
 You can find sample programs that demonstrate how to use the client packages to fetch the embeddings in `cmd` directory of this project.
 
@@ -19,7 +20,10 @@ It's essentially a Go rewrite of character and recursive character text splitter
 
 ## Environment variables
 
-Each client package lets you initialize a default API client for a specific embeddings provider by reading the API keys from environment variables.
+> [!NOTE]
+> Each client package lets you initialize a default API client for a specific embeddings provider by reading the API keys from environment variables
+
+Here's a list of the env vars for each supported client
 
 ### OpenAI
 
@@ -36,6 +40,10 @@ Each client package lets you initialize a default API client for a specific embe
 * `GOOGLE_PROJECT_ID`: Google Project ID
 * `VOYAGE_API_KEY`: VoyageAI API key
 
+### Voyage
+
+* `VOYAGE_API_KEY`: Voyage AI API key
+
 ## nix
 
 The project provides a simple `nix` flake tha leverages [gomod2nix](https://github.com/nix-community/gomod2nix) for consistent Go environments and builds.

cmd/ollama/main.go

@@ -0,0 +1,42 @@
+package main
+
+import (
+	"context"
+	"flag"
+	"fmt"
+	"log"
+
+	"github.com/milosgajdos/go-embeddings/ollama"
+)
+
+var (
+	prompt string
+	model  string
+)
+
+func init() {
+	flag.StringVar(&prompt, "prompt", "what is life", "input prompt")
+	flag.StringVar(&model, "model", "", "model name")
+}
+
+func main() {
+	flag.Parse()
+
+	if model == "" {
+		log.Fatal("missing ollama model")
+	}
+
+	c := ollama.NewClient()
+
+	embReq := &ollama.EmbeddingRequest{
+		Prompt: prompt,
+		Model:  model,
+	}
+
+	embs, err := c.Embed(context.Background(), embReq)
+	if err != nil {
+		log.Fatal(err)
+	}
+
+	fmt.Printf("got %d embeddings", len(embs))
+}

ollama/client.go

@@ -0,0 +1,61 @@
+package ollama
+
+import (
+	"github.com/milosgajdos/go-embeddings"
+	"github.com/milosgajdos/go-embeddings/client"
+)
+
+const (
+	// BaseURL is Ollama HTTP API embeddings base URL.
+	BaseURL = "http://localhost:11434/api"
+)
+
+// Client is an OpenAI HTTP API client.
+type Client struct {
+	opts Options
+}
+
+type Options struct {
+	BaseURL    string
+	HTTPClient *client.HTTP
+}
+
+// Option is functional graph option.
+type Option func(*Options)
+
+// NewClient creates a new Ollama HTTP API client and returns it.
+// You can override the default options via the client methods.
+func NewClient(opts ...Option) *Client {
+	options := Options{
+		BaseURL:    BaseURL,
+		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...)
+}
+
+// WithBaseURL sets the API base URL.
+func WithBaseURL(baseURL string) Option {
+	return func(o *Options) {
+		o.BaseURL = baseURL
+	}
+}
+
+// WithVersion sets the API version.
+// WithHTTPClient sets the HTTP client.
+func WithHTTPClient(httpClient *client.HTTP) Option {
+	return func(o *Options) {
+		o.HTTPClient = httpClient
+	}
+}

ollama/client_test.go

@@ -0,0 +1,28 @@
+package ollama
+
+import (
+	"testing"
+
+	"github.com/milosgajdos/go-embeddings/client"
+	"github.com/stretchr/testify/assert"
+)
+
+func TestClient(t *testing.T) {
+	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("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)
+	})
+}

ollama/embedding.go

@@ -0,0 +1,67 @@
+package ollama
+
+import (
+	"bytes"
+	"context"
+	"encoding/json"
+	"net/http"
+	"net/url"
+
+	"github.com/milosgajdos/go-embeddings"
+	"github.com/milosgajdos/go-embeddings/request"
+)
+
+// EmbeddingRequest is serialized and sent to the API server.
+type EmbeddingRequest struct {
+	Prompt any    `json:"prompt"`
+	Model  string `json:"model"`
+}
+
+// EmbedddingResponse received from API.
+type EmbedddingResponse struct {
+	Embedding []float64 `json:"embedding"`
+}
+
+// ToEmbeddings converts the API response,
+// into a slice of embeddings and returns it.
+func (e *EmbedddingResponse) ToEmbeddings() ([]*embeddings.Embedding, error) {
+	floats := make([]float64, len(e.Embedding))
+	copy(floats, e.Embedding)
+	return []*embeddings.Embedding{
+		{Vector: floats},
+	}, nil
+}
+
+// Embed returns embeddings for every object in EmbeddingRequest.
+func (c *Client) Embed(ctx context.Context, embReq *EmbeddingRequest) ([]*embeddings.Embedding, error) {
+	u, err := url.Parse(c.opts.BaseURL + "/embeddings")
+	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
+	}
+
+	options := []request.Option{}
+	req, err := request.NewHTTP(ctx, http.MethodPost, u.String(), body, options...)
+	if err != nil {
+		return nil, err
+	}
+
+	resp, err := request.Do[APIError](c.opts.HTTPClient, req)
+	if err != nil {
+		return nil, err
+	}
+	defer resp.Body.Close()
+
+	e := new(EmbedddingResponse)
+	if err := json.NewDecoder(resp.Body).Decode(e); err != nil {
+		return nil, err
+	}
+
+	return e.ToEmbeddings()
+}

ollama/error.go

@@ -0,0 +1,17 @@
+package ollama
+
+import "encoding/json"
+
+// APIError is Ollama API error.
+type APIError struct {
+	ErrorMessage string `json:"error"`
+}
+
+// Error implements errors interface.
+func (e APIError) Error() string {
+	b, err := json.Marshal(e)
+	if err != nil {
+		return "unknown error"
+	}
+	return string(b)
+}

openai/client.go

@@ -32,7 +32,7 @@ type Options struct {
 // Option is functional graph option.
 type Option func(*Options)
 
-// NewClient creates a new HTTP API client and returns it.
+// NewClient creates a new OpenAI 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.