Merge pull request #46 from grafana/vectorservice-query-fix

vectorservice query fix
grafana · Sep 14, 2023 · 67d7394 · 67d7394
2 parents 2936f84 + 8dd1622
commit 67d7394
Show file tree

Hide file tree

Showing 5 changed files with 16 additions and 13 deletions.
diff --git a/pkg/plugin/resources.go b/pkg/plugin/resources.go
@@ -64,9 +64,9 @@ func newOpenAIProxy() http.Handler {
 }
 
 type vectorSearchRequest struct {
-	Text       string `json:"text"`
+	Query      string `json:"query"`
 	Collection string `json:"collection"`
-	Limit      uint64 `json:"limit"`
+	TopK       uint64 `json:"topK"`
 }
 
 type vectorSearchResponse struct {
@@ -87,10 +87,10 @@ func (app *App) handleVectorSearch(w http.ResponseWriter, req *http.Request) {
 		http.Error(w, err.Error(), http.StatusBadRequest)
 		return
 	}
-	if body.Limit == 0 {
-		body.Limit = 10
+	if body.TopK == 0 {
+		body.TopK = 10
 	}
-	results, err := app.vectorService.Search(req.Context(), body.Collection, body.Text, body.Limit)
+	results, err := app.vectorService.Search(req.Context(), body.Collection, body.Query, body.TopK)
 	if err != nil {
 		http.Error(w, err.Error(), http.StatusInternalServerError)
 		return

diff --git a/pkg/plugin/vector/service.go b/pkg/plugin/vector/service.go
@@ -12,7 +12,7 @@ import (
 )
 
 type Service interface {
-	Search(ctx context.Context, collection string, query string, limit uint64) ([]store.SearchResult, error)
+	Search(ctx context.Context, collection string, query string, topK uint64) ([]store.SearchResult, error)
 	Cancel()
 }
 
@@ -57,7 +57,10 @@ func NewService(s VectorSettings, secrets map[string]string) (Service, error) {
 	}, nil
 }
 
-func (v *vectorService) Search(ctx context.Context, collection string, query string, limit uint64) ([]store.SearchResult, error) {
+func (v *vectorService) Search(ctx context.Context, collection string, query string, topK uint64) ([]store.SearchResult, error) {
+	if query == "" {
+		return nil, fmt.Errorf("query cannot be empty")
+	}
 	exists, err := v.store.CollectionExists(ctx, collection)
 	if err != nil {
 		return nil, fmt.Errorf("vector store collections: %w", err)
@@ -75,7 +78,7 @@ func (v *vectorService) Search(ctx context.Context, collection string, query str
 
 	log.DefaultLogger.Info("Searching", "collection", collection, "query", query)
 	// Search the vector store for similar vectors.
-	results, err := v.store.Search(ctx, collection, e, limit)
+	results, err := v.store.Search(ctx, collection, e, topK)
 	if err != nil {
 		return nil, fmt.Errorf("vector store search: %w", err)
 	}

diff --git a/pkg/plugin/vector/store/qdrant.go b/pkg/plugin/vector/store/qdrant.go
@@ -82,14 +82,14 @@ func (q *qdrantStore) CollectionExists(ctx context.Context, collection string) (
 	return true, nil
 }
 
-func (q *qdrantStore) Search(ctx context.Context, collection string, vector []float32, limit uint64) ([]SearchResult, error) {
+func (q *qdrantStore) Search(ctx context.Context, collection string, vector []float32, topK uint64) ([]SearchResult, error) {
 	if q.md != nil {
 		ctx = metadata.NewOutgoingContext(ctx, *q.md)
 	}
 	result, err := q.pointsClient.Search(ctx, &qdrant.SearchPoints{
 		CollectionName: collection,
 		Vector:         vector,
-		Limit:          limit,
+		Limit:          topK,
 		// Include all payloads in the search result
 		WithVectors: &qdrant.WithVectorsSelector{SelectorOptions: &qdrant.WithVectorsSelector_Enable{Enable: false}},
 		WithPayload: &qdrant.WithPayloadSelector{SelectorOptions: &qdrant.WithPayloadSelector_Enable{Enable: true}},

diff --git a/pkg/plugin/vector/store/store.go b/pkg/plugin/vector/store/store.go
@@ -20,7 +20,7 @@ type SearchResult struct {
 
 type ReadVectorStore interface {
 	CollectionExists(ctx context.Context, collection string) (bool, error)
-	Search(ctx context.Context, collection string, vector []float32, limit uint64) ([]SearchResult, error)
+	Search(ctx context.Context, collection string, vector []float32, topK uint64) ([]SearchResult, error)
 }
 
 type WriteVectorStore interface {

diff --git a/pkg/plugin/vector/store/vectorapi.go b/pkg/plugin/vector/store/vectorapi.go
@@ -31,14 +31,14 @@ func (g *grafanaVectorAPI) CollectionExists(ctx context.Context, collection stri
 	return true, nil
 }
 
-func (g *grafanaVectorAPI) Search(ctx context.Context, collection string, vector []float32, limit uint64) ([]SearchResult, error) {
+func (g *grafanaVectorAPI) Search(ctx context.Context, collection string, vector []float32, topK uint64) ([]SearchResult, error) {
 	type queryPointsRequest struct {
 		Query []float32 `json:"query"`
 		TopK  uint64    `json:"top_k"`
 	}
 	reqBody := queryPointsRequest{
 		Query: vector,
-		TopK:  limit,
+		TopK:  topK,
 	}
 	reqJSON, err := json.Marshal(reqBody)
 	if err != nil {