From ada90835b5cf7f52d92134091b74db0d059ec10a Mon Sep 17 00:00:00 2001
From: nasser <nasser.subhi@rihal.om>
Date: Mon, 12 Aug 2024 10:54:41 +0400
Subject: [PATCH] qdrant: add wait option when adding documents

---
 vectorstores/qdrant/options.go     | 14 ++++++++++++--
 vectorstores/qdrant/qdrant.go      |  1 +
 vectorstores/qdrant/qdrant_test.go |  6 ++++++
 vectorstores/qdrant/rest.go        | 10 +++++++---
 4 files changed, 26 insertions(+), 5 deletions(-)

diff --git a/vectorstores/qdrant/options.go b/vectorstores/qdrant/options.go
index 60c407591..ba0d5b21d 100644
--- a/vectorstores/qdrant/options.go
+++ b/vectorstores/qdrant/options.go
@@ -9,7 +9,8 @@ import (
 )
 
 const (
-	defaultContentKey = "content"
+	defaultContentKey    = "content"
+	defaultWaitForInsert = true
 )
 
 // ErrInvalidOptions is returned when the options given are invalid.
@@ -56,9 +57,18 @@ func WithContentKey(contentKey string) Option {
 	}
 }
 
+// WithWaitForInsert returns an Option for setting whether to wait for the insert to complete.
+// Optional. Defaults to true.
+func WithWaitForInsert(wait bool) Option {
+	return func(p *Store) {
+		p.waitForInsert = wait
+	}
+}
+
 func applyClientOptions(opts ...Option) (Store, error) {
 	o := &Store{
-		contentKey: defaultContentKey,
+		contentKey:    defaultContentKey,
+		waitForInsert: defaultWaitForInsert,
 	}
 
 	for _, opt := range opts {
diff --git a/vectorstores/qdrant/qdrant.go b/vectorstores/qdrant/qdrant.go
index 6896c424f..e0ef186e0 100644
--- a/vectorstores/qdrant/qdrant.go
+++ b/vectorstores/qdrant/qdrant.go
@@ -16,6 +16,7 @@ type Store struct {
 	qdrantURL      url.URL
 	apiKey         string
 	contentKey     string
+	waitForInsert  bool
 }
 
 var _ vectorstores.VectorStore = Store{}
diff --git a/vectorstores/qdrant/qdrant_test.go b/vectorstores/qdrant/qdrant_test.go
index 788430258..36fd90b26 100644
--- a/vectorstores/qdrant/qdrant_test.go
+++ b/vectorstores/qdrant/qdrant_test.go
@@ -43,6 +43,7 @@ func TestQdrantStore(t *testing.T) {
 		qdrant.WithAPIKey(apiKey),
 		qdrant.WithCollectionName(collectionName),
 		qdrant.WithEmbedder(e),
+		qdrant.WithWaitForInsert(true),
 	)
 	require.NoError(t, err)
 
@@ -81,6 +82,7 @@ func TestQdrantStoreWithScoreThreshold(t *testing.T) {
 		qdrant.WithAPIKey(apiKey),
 		qdrant.WithCollectionName(collectionName),
 		qdrant.WithEmbedder(e),
+		qdrant.WithWaitForInsert(true),
 	)
 	require.NoError(t, err)
 
@@ -136,6 +138,7 @@ func TestSimilaritySearchWithInvalidScoreThreshold(t *testing.T) {
 		qdrant.WithAPIKey(apiKey),
 		qdrant.WithCollectionName(collectionName),
 		qdrant.WithEmbedder(e),
+		qdrant.WithWaitForInsert(true),
 	)
 	require.NoError(t, err)
 
@@ -187,6 +190,7 @@ func TestQdrantAsRetriever(t *testing.T) {
 		qdrant.WithAPIKey(apiKey),
 		qdrant.WithCollectionName(collectionName),
 		qdrant.WithEmbedder(e),
+		qdrant.WithWaitForInsert(true),
 	)
 	require.NoError(t, err)
 
@@ -235,6 +239,7 @@ func TestQdrantRetrieverScoreThreshold(t *testing.T) {
 		qdrant.WithAPIKey(apiKey),
 		qdrant.WithCollectionName(collectionName),
 		qdrant.WithEmbedder(e),
+		qdrant.WithWaitForInsert(true),
 	)
 	require.NoError(t, err)
 
@@ -287,6 +292,7 @@ func TestQdrantRetrieverFilter(t *testing.T) {
 		qdrant.WithAPIKey(apiKey),
 		qdrant.WithCollectionName(collectionName),
 		qdrant.WithEmbedder(e),
+		qdrant.WithWaitForInsert(true),
 	)
 	require.NoError(t, err)
 
diff --git a/vectorstores/qdrant/rest.go b/vectorstores/qdrant/rest.go
index e4fee5d2b..63eabde03 100644
--- a/vectorstores/qdrant/rest.go
+++ b/vectorstores/qdrant/rest.go
@@ -33,11 +33,15 @@ func (s Store) upsertPoints(
 		},
 	}
 
-	url := baseURL.JoinPath("collections", s.collectionName, "points")
+	upsertUrl := baseURL.JoinPath("collections", s.collectionName, "points")
+	queryParams := url.Values{}
+	queryParams.Add("wait", fmt.Sprintf("%t", s.waitForInsert))
+	upsertUrl.RawQuery = queryParams.Encode()
+
 	body,
 		status,
 		err := DoRequest(
-		ctx, *url,
+		ctx, *upsertUrl,
 		s.apiKey,
 		http.MethodPut,
 		payload,
@@ -134,7 +138,7 @@ func DoRequest(ctx context.Context,
 	}
 	body := bytes.NewReader(payloadBytes)
 
-	req, err := http.NewRequestWithContext(ctx, method, url.String()+"?wait=true", body)
+	req, err := http.NewRequestWithContext(ctx, method, url.String(), body)
 	if err != nil {
 		return nil, 0, err
 	}