vectorstores: add Milvus (#352)

* adding milvus vector store

* adding hugging face text embedings

* adding milvus vector store example

embeddings/tei/doc.go

@@ -0,0 +1,8 @@
+/*
+Huggingface Text Embeddings Inference
+https://github.com/huggingface/text-embeddings-inference
+
+package is a wrapper for the Huggingface text embeddings inference project
+that can be run locally for creating vector embeddings.
+*/
+package tei

embeddings/tei/options.go

@@ -0,0 +1,97 @@
+package tei
+
+import (
+	"errors"
+	"runtime"
+	"time"
+
+	client "github.com/gage-technologies/tei-go"
+)
+
+const (
+	_defaultBatchSize       = 512
+	_defaultStripNewLines   = true
+	_defaultTimeNanoSeconds = 60 * 1000000000
+)
+
+var ErrMissingAPIBaseURL = errors.New("missing the API Base URL") //nolint:lll
+
+type Option func(emb *TextEmbeddingsInference)
+
+// WithStripNewLines is an option for specifying the should it strip new lines.
+func WithStripNewLines(stripNewLines bool) Option {
+	return func(p *TextEmbeddingsInference) {
+		p.StripNewLines = stripNewLines
+	}
+}
+
+// WithPoolSize is an option for specifying the number of goroutines.
+func WithPoolSize(poolSize int) Option {
+	return func(p *TextEmbeddingsInference) {
+		p.poolSize = poolSize
+	}
+}
+
+// WithBatchSize is an option for specifying the batch size.
+func WithBatchSize(batchSize int) Option {
+	return func(p *TextEmbeddingsInference) {
+		p.BatchSize = batchSize
+	}
+}
+
+// WithAPIBaseURL adds base url for api.
+func WithAPIBaseURL(url string) Option {
+	return func(emb *TextEmbeddingsInference) {
+		emb.baseURL = url
+	}
+}
+
+// WithHeaders add request headers.
+func WithHeaders(headers map[string]string) Option {
+	return func(emb *TextEmbeddingsInference) {
+		if emb.headers == nil {
+			emb.headers = make(map[string]string, len(headers))
+		}
+		for k, v := range headers {
+			emb.headers[k] = v
+		}
+	}
+}
+
+// WithCookies add request cookies.
+func WithCookies(cookies map[string]string) Option {
+	return func(emb *TextEmbeddingsInference) {
+		if emb.cookies == nil {
+			emb.cookies = make(map[string]string, len(cookies))
+		}
+		for k, v := range cookies {
+			emb.cookies[k] = v
+		}
+	}
+}
+
+// WithTimeout set the request timeout.
+func WithTimeout(dur time.Duration) Option {
+	return func(emb *TextEmbeddingsInference) {
+		emb.timeout = dur
+	}
+}
+
+func applyClientOptions(opts ...Option) (TextEmbeddingsInference, error) {
+	emb := TextEmbeddingsInference{
+		StripNewLines: _defaultStripNewLines,
+		BatchSize:     _defaultBatchSize,
+		timeout:       time.Duration(_defaultTimeNanoSeconds),
+		poolSize:      runtime.GOMAXPROCS(0),
+	}
+	for _, opt := range opts {
+		opt(&emb)
+	}
+	if emb.baseURL == "" {
+		return emb, ErrMissingAPIBaseURL
+	}
+	if emb.client == nil {
+		emb.client = client.NewClient(emb.baseURL, emb.headers, emb.cookies, emb.timeout)
+	}
+	return emb, nil
+}

embeddings/tei/text_embeddings_inference..go

@@ -0,0 +1,84 @@
+package tei
+
+import (
+	"context"
+	"strings"
+	"time"
+
+	client "github.com/gage-technologies/tei-go"
+	"github.com/sourcegraph/conc/pool"
+	"github.com/tmc/langchaingo/embeddings"
+)
+
+type TextEmbeddingsInference struct {
+	client        *client.Client
+	StripNewLines bool
+	BatchSize     int
+	baseURL       string
+	headers       map[string]string
+	cookies       map[string]string
+	timeout       time.Duration
+	poolSize      int
+}
+
+var _ embeddings.Embedder = TextEmbeddingsInference{}
+
+func New(opts ...Option) (TextEmbeddingsInference, error) {
+	emb, err := applyClientOptions(opts...)
+	if err != nil {
+		return emb, err
+	}
+	emb.client = client.NewClient(emb.baseURL, emb.headers, emb.cookies, emb.timeout)
+
+	return emb, nil
+}
+
+// EmbedDocuments creates one vector embedding for each of the texts.
+func (e TextEmbeddingsInference) EmbedDocuments(_ context.Context, texts []string) ([][]float32, error) {
+	batchedTexts := embeddings.BatchTexts(
+		embeddings.MaybeRemoveNewLines(texts, e.StripNewLines),
+		e.BatchSize,
+	)
+
+	emb := make([][]float32, 0, len(texts))
+
+	p := pool.New().WithMaxGoroutines(e.poolSize).WithErrors()
+
+	for _, txt := range batchedTexts {
+		p.Go(func() error {
+			curTextEmbeddings, err := e.client.Embed(strings.Join(txt, " "), false)
+			if err != nil {
+				return err
+			}
+
+			textLengths := make([]int, 0, len(txt))
+			for _, text := range txt {
+				textLengths = append(textLengths, len(text))
+			}
+
+			combined, err := embeddings.CombineVectors(curTextEmbeddings, textLengths)
+			if err != nil {
+				return err
+			}
+
+			emb = append(emb, combined)
+
+			return nil
+		})
+	}
+	return emb, p.Wait()
+}
+
+// EmbedQuery embeds a single text.
+func (e TextEmbeddingsInference) EmbedQuery(_ context.Context, text string) ([]float32, error) {
+	if e.StripNewLines {
+		text = strings.ReplaceAll(text, "\n", " ")
+	}
+
+	emb, err := e.client.Embed(text, false)
+	if err != nil {
+		return nil, err
+	}
+
+	return emb[0], nil
+}

examples/milvus-vectorstore-example/.gitignore

@@ -0,0 +1 @@
+volumes

examples/milvus-vectorstore-example/README.md

@@ -0,0 +1,17 @@
+# Milvus vector store with local embeddings 
+
+Dependencies:
+- [Text Embeddings Inference](https://github.com/huggingface/text-embeddings-inference)
+- [Ollama](https://ollama.ai/) 
+
+```shell
+# start milvus
+docker-compose up -d
+
+# start mistral on ollama
+ollama run mistral
+
+#start embedding server
+text-embeddings-router --model-id thenlper/gte-large --port 5500
+
+```

examples/milvus-vectorstore-example/Taskfile.yml

@@ -0,0 +1,13 @@
+# https://taskfile.dev
+
+version: '3'
+
+silent: true
+
+tasks:
+  default:
+    cmds:
+      - task --list-all
+  run:text-embeddings-inference:
+    cmds:
+      - text-embeddings-router --model-id thenlper/gte-large --port 5500

examples/milvus-vectorstore-example/docker-compose.yml

@@ -0,0 +1,47 @@
+version: "3.5"
+
+services:
+  etcd:
+    container_name: milvus-etcd
+    image: quay.io/coreos/etcd:v3.5.0
+    environment:
+      - ETCD_AUTO_COMPACTION_MODE=revision
+      - ETCD_AUTO_COMPACTION_RETENTION=1000
+      - ETCD_QUOTA_BACKEND_BYTES=4294967296
+    volumes:
+      - ${DOCKER_VOLUME_DIRECTORY:-.}/volumes/etcd:/etcd
+    command: etcd -advertise-client-urls=http://127.0.0.1:2379 -listen-client-urls http://0.0.0.0:2379 --data-dir /etcd
+
+  minio:
+    container_name: milvus-minio
+    image: minio/minio:RELEASE.2020-12-03T00-03-10Z
+    environment:
+      MINIO_ACCESS_KEY: minioadmin
+      MINIO_SECRET_KEY: minioadmin
+    volumes:
+      - ${DOCKER_VOLUME_DIRECTORY:-.}/volumes/minio:/minio_data
+    command: minio server /minio_data
+    healthcheck:
+      test: ["CMD", "curl", "-f", "http://localhost:9000/minio/health/live"]
+      interval: 30s
+      timeout: 20s
+      retries: 3
+
+  standalone:
+    container_name: milvus-standalone
+    image: milvusdb/milvus
+    command: ["milvus", "run", "standalone"]
+    environment:
+      ETCD_ENDPOINTS: etcd:2379
+      MINIO_ADDRESS: minio:9000
+    volumes:
+      - ${DOCKER_VOLUME_DIRECTORY:-.}/volumes/milvus:/var/lib/milvus
+    ports:
+      - "19530:19530"
+    depends_on:
+      - "etcd"
+      - "minio"
+
+networks:
+  default:
+    name: milvus

examples/milvus-vectorstore-example/go.mod

@@ -0,0 +1,42 @@
+module github.com/tmc/langchaingo/examples/ollama-milvus-vectorstore-example
+
+go 1.20
+
+// NOTE: remove the following line to use the official (rather than local development) version
+replace github.com/tmc/langchaingo => ../..
+
+require (
+	github.com/milvus-io/milvus-sdk-go/v2 v2.3.2
+	github.com/tmc/langchaingo v0.0.0-00010101000000-000000000000
+)
+
+require (
+	github.com/cockroachdb/errors v1.9.1 // indirect
+	github.com/cockroachdb/logtags v0.0.0-20211118104740-dabe8e521a4f // indirect
+	github.com/cockroachdb/redact v1.1.3 // indirect
+	github.com/dlclark/regexp2 v1.8.1 // indirect
+	github.com/gage-technologies/tei-go v0.2.0 // indirect
+	github.com/getsentry/sentry-go v0.12.0 // indirect
+	github.com/gogo/protobuf v1.3.2 // indirect
+	github.com/golang/protobuf v1.5.3 // indirect
+	github.com/google/uuid v1.3.0 // indirect
+	github.com/grpc-ecosystem/go-grpc-middleware v1.3.0 // indirect
+	github.com/kr/pretty v0.3.0 // indirect
+	github.com/kr/text v0.2.0 // indirect
+	github.com/milvus-io/milvus-proto/go-api/v2 v2.3.2 // indirect
+	github.com/pkg/errors v0.9.1 // indirect
+	github.com/pkoukk/tiktoken-go v0.1.2 // indirect
+	github.com/rogpeppe/go-internal v1.11.0 // indirect
+	github.com/sourcegraph/conc v0.3.0 // indirect
+	github.com/tidwall/gjson v1.14.4 // indirect
+	github.com/tidwall/match v1.1.1 // indirect
+	github.com/tidwall/pretty v1.2.0 // indirect
+	go.uber.org/atomic v1.7.0 // indirect
+	go.uber.org/multierr v1.9.0 // indirect
+	golang.org/x/net v0.17.0 // indirect
+	golang.org/x/sys v0.13.0 // indirect
+	golang.org/x/text v0.13.0 // indirect
+	google.golang.org/genproto/googleapis/rpc v0.0.0-20230530153820-e85fd2cbaebc // indirect
+	google.golang.org/grpc v1.57.1 // indirect
+	google.golang.org/protobuf v1.31.0 // indirect
+)