qdrant: add wait option when adding documents

tmc · Aug 12, 2024 · ada9083 · ada9083
1 parent 1975058
commit ada9083
Show file tree

Hide file tree

Showing 4 changed files with 26 additions and 5 deletions.
diff --git 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
@@ -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
@@ -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
@@ -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
 	}