Skip to content

Commit

Permalink
qdrant: add wait option when adding documents
Browse files Browse the repository at this point in the history
  • Loading branch information
n3sser96 committed Aug 12, 2024
1 parent 1975058 commit 07163e5
Show file tree
Hide file tree
Showing 3 changed files with 20 additions and 5 deletions.
14 changes: 12 additions & 2 deletions vectorstores/qdrant/options.go
Original file line number Diff line number Diff line change
Expand Up @@ -9,7 +9,8 @@ import (
)

const (
defaultContentKey = "content"
defaultContentKey = "content"
defaultWaitForInsert = true
)

// ErrInvalidOptions is returned when the options given are invalid.
Expand Down Expand Up @@ -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 {
Expand Down
1 change: 1 addition & 0 deletions vectorstores/qdrant/qdrant.go
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,7 @@ type Store struct {
qdrantURL url.URL
apiKey string
contentKey string
waitForInsert bool
}

var _ vectorstores.VectorStore = Store{}
Expand Down
10 changes: 7 additions & 3 deletions vectorstores/qdrant/rest.go
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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
}
Expand Down

0 comments on commit 07163e5

Please sign in to comment.