chore: add an index method to insert a vector (#129)

henomis · Sep 13, 2023 · 15a40a1 · 15a40a1
1 parent f97cc8b
commit 15a40a1
Show file tree

Hide file tree

Showing 3 changed files with 79 additions and 0 deletions.
diff --git a/index/pinecone/pinecone.go b/index/pinecone/pinecone.go
@@ -126,6 +126,31 @@ func (p *Index) IsEmpty(ctx context.Context) (bool, error) {
 
 }
 
+func (p *Index) Add(ctx context.Context, item *index.Data) error {
+	err := p.createIndexIfRequired(ctx)
+	if err != nil {
+		return fmt.Errorf("%s: %w", index.ErrInternal, err)
+	}
+
+	if item.ID == "" {
+		id, err := uuid.NewUUID()
+		if err != nil {
+			return err
+		}
+		item.ID = id.String()
+	}
+
+	return p.vectorUpsert(ctx,
+		[]pineconerequest.Vector{
+			{
+				ID:       item.ID,
+				Values:   item.Values,
+				Metadata: item.Metadata,
+			},
+		},
+	)
+}
+
 func (p *Index) Search(ctx context.Context, values []float64, opts ...option.Option) (index.SearchResults, error) {
 
 	pineconeOptions := &option.Options{

diff --git a/index/qdrant/qdrant.go b/index/qdrant/qdrant.go
@@ -115,6 +115,31 @@ func (p *Index) IsEmpty(ctx context.Context) (bool, error) {
 
 }
 
+func (q *Index) Add(ctx context.Context, item *index.Data) error {
+	err := q.createCollectionIfRequired(ctx)
+	if err != nil {
+		return fmt.Errorf("%s: %w", index.ErrInternal, err)
+	}
+
+	if item.ID == "" {
+		id, err := uuid.NewUUID()
+		if err != nil {
+			return err
+		}
+		item.ID = id.String()
+	}
+
+	return q.pointUpsert(ctx,
+		[]qdrantrequest.Point{
+			{
+				ID:      item.ID,
+				Vector:  item.Values,
+				Payload: item.Metadata,
+			},
+		},
+	)
+}
+
 func (q *Index) Search(ctx context.Context, values []float64, opts ...option.Option) (index.SearchResults, error) {
 	qdrantOptions := &option.Options{
 		TopK: defaultTopK,

diff --git a/index/simpleVectorIndex/simpleVectorIndex.go b/index/simpleVectorIndex/simpleVectorIndex.go
@@ -6,6 +6,7 @@ import (
 	"fmt"
 	"math"
 	"os"
+	"strconv"
 	"strings"
 
 	"github.com/henomis/lingoose/document"
@@ -140,6 +141,34 @@ func (s *Index) IsEmpty() (bool, error) {
 	return len(s.data) == 0, nil
 }
 
+func (s *Index) Add(ctx context.Context, item *index.Data) error {
+	err := s.load()
+	if err != nil {
+		return fmt.Errorf("%s: %w", index.ErrInternal, err)
+	}
+
+	if item.ID == "" {
+		lastID := s.data[len(s.data)-1].ID
+		lastIDAsInt, err := strconv.Atoi(lastID)
+		if err != nil {
+			return fmt.Errorf("%s: %w", index.ErrInternal, err)
+		}
+
+		item.ID = fmt.Sprintf("%d", lastIDAsInt+1)
+	}
+
+	s.data = append(
+		s.data,
+		data{
+			ID:       item.ID,
+			Values:   item.Values,
+			Metadata: item.Metadata,
+		},
+	)
+
+	return s.save()
+}
+
 func (s *Index) Search(ctx context.Context, values []float64, opts ...option.Option) (index.SearchResults, error) {
 	sviOptions := &option.Options{
 		TopK: defaultTopK,