From 1bba04f5288c0c46c763340895160cc0ef4ee3c4 Mon Sep 17 00:00:00 2001 From: Erik Dubbelboer Date: Sun, 21 Apr 2024 13:39:42 +0200 Subject: [PATCH 1/3] Add NewEmbeddingFuncOllamaWithURL Allows the user to specify the base URL of Ollama instead of the default localhost one. --- embed_ollama.go | 18 +++++++++++++++--- embed_ollama_test.go | 5 +---- 2 files changed, 16 insertions(+), 7 deletions(-) diff --git a/embed_ollama.go b/embed_ollama.go index d2231f7..3a5ac4c 100644 --- a/embed_ollama.go +++ b/embed_ollama.go @@ -11,9 +11,7 @@ import ( "sync" ) -// TODO: Turn into const and use as default, but allow user to pass custom URL -// as well as custom API key, in case Ollama runs on a remote (secured) server. -var baseURLOllama = "http://localhost:11434/api" +const defaultBaseURLOllama = "http://localhost:11434/api" type ollamaResponse struct { Embedding []float32 `json:"embedding"` @@ -24,6 +22,20 @@ type ollamaResponse struct { // that supports embeddings. A good one as of 2024-03-02 is "nomic-embed-text". // See https://ollama.com/library/nomic-embed-text func NewEmbeddingFuncOllama(model string) EmbeddingFunc { + return NewEmbeddingFuncOllamaWithURL(model, defaultBaseURLOllama) +} + +// NewEmbeddingFuncOllamaWithURL returns a function that creates embeddings for a text +// using Ollama's embedding API. You can pass any model that Ollama supports and +// that supports embeddings. A good one as of 2024-03-02 is "nomic-embed-text". +// See https://ollama.com/library/nomic-embed-text +// baseURLOllama is the base URL of the Ollama API. If it's empty, +// "http://localhost:11434/api" is used. +func NewEmbeddingFuncOllamaWithURL(model string, baseURLOllama string) EmbeddingFunc { + if baseURLOllama == "" { + baseURLOllama = defaultBaseURLOllama + } + // We don't set a default timeout here, although it's usually a good idea. // In our case though, the library user can set the timeout on the context, // and it might have to be a long timeout, depending on the text length. diff --git a/embed_ollama_test.go b/embed_ollama_test.go index a3af70a..5d96ff9 100644 --- a/embed_ollama_test.go +++ b/embed_ollama_test.go @@ -64,11 +64,8 @@ func TestNewEmbeddingFuncOllama(t *testing.T) { if err != nil { t.Fatal("unexpected error:", err) } - // TODO: It's bad to overwrite a global var for testing. Follow-up with a change - // to allow passing custom URLs to the function. - baseURLOllama = strings.Replace(baseURLOllama, "11434", u.Port(), 1) - f := NewEmbeddingFuncOllama(model) + f := NewEmbeddingFuncOllamaWithURL(model, strings.Replace(defaultBaseURLOllama, "11434", u.Port(), 1)) res, err := f(context.Background(), prompt) if err != nil { t.Fatal("expected nil, got", err) From 929853584a1acbe3e6a34f377c9c5708670bcf9f Mon Sep 17 00:00:00 2001 From: Erik Dubbelboer Date: Sun, 21 Apr 2024 19:29:45 +0200 Subject: [PATCH 2/3] Make it a breaking change --- embed_ollama.go | 10 +--------- embed_ollama_test.go | 2 +- 2 files changed, 2 insertions(+), 10 deletions(-) diff --git a/embed_ollama.go b/embed_ollama.go index 3a5ac4c..019c669 100644 --- a/embed_ollama.go +++ b/embed_ollama.go @@ -21,17 +21,9 @@ type ollamaResponse struct { // using Ollama's embedding API. You can pass any model that Ollama supports and // that supports embeddings. A good one as of 2024-03-02 is "nomic-embed-text". // See https://ollama.com/library/nomic-embed-text -func NewEmbeddingFuncOllama(model string) EmbeddingFunc { - return NewEmbeddingFuncOllamaWithURL(model, defaultBaseURLOllama) -} - -// NewEmbeddingFuncOllamaWithURL returns a function that creates embeddings for a text -// using Ollama's embedding API. You can pass any model that Ollama supports and -// that supports embeddings. A good one as of 2024-03-02 is "nomic-embed-text". -// See https://ollama.com/library/nomic-embed-text // baseURLOllama is the base URL of the Ollama API. If it's empty, // "http://localhost:11434/api" is used. -func NewEmbeddingFuncOllamaWithURL(model string, baseURLOllama string) EmbeddingFunc { +func NewEmbeddingFuncOllama(model string, baseURLOllama string) EmbeddingFunc { if baseURLOllama == "" { baseURLOllama = defaultBaseURLOllama } diff --git a/embed_ollama_test.go b/embed_ollama_test.go index 5d96ff9..4cd5a63 100644 --- a/embed_ollama_test.go +++ b/embed_ollama_test.go @@ -65,7 +65,7 @@ func TestNewEmbeddingFuncOllama(t *testing.T) { t.Fatal("unexpected error:", err) } - f := NewEmbeddingFuncOllamaWithURL(model, strings.Replace(defaultBaseURLOllama, "11434", u.Port(), 1)) + f := NewEmbeddingFuncOllama(model, strings.Replace(defaultBaseURLOllama, "11434", u.Port(), 1)) res, err := f(context.Background(), prompt) if err != nil { t.Fatal("expected nil, got", err) From f51b5920d78bad99770d966bb1a6075da263e70c Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Philipp=20Gill=C3=A9?= Date: Sun, 21 Apr 2024 20:50:49 +0200 Subject: [PATCH 3/3] Update example to use changed constructor function --- examples/rag-wikipedia-ollama/main.go | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/examples/rag-wikipedia-ollama/main.go b/examples/rag-wikipedia-ollama/main.go index 3b865da..f754eca 100644 --- a/examples/rag-wikipedia-ollama/main.go +++ b/examples/rag-wikipedia-ollama/main.go @@ -49,7 +49,7 @@ func main() { // variable to be set. // For this example we choose to use a locally running embedding model though. // It requires Ollama to serve its API at "http://localhost:11434/api". - collection, err := db.GetOrCreateCollection("Wikipedia", nil, chromem.NewEmbeddingFuncOllama(embeddingModel)) + collection, err := db.GetOrCreateCollection("Wikipedia", nil, chromem.NewEmbeddingFuncOllama(embeddingModel, "")) if err != nil { panic(err) }