diff --git a/pkg/plugin/resources.go b/pkg/plugin/resources.go index e2d6c7bc..2f81e408 100644 --- 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 index 879a0805..b1a6b4bb 100644 --- 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 index 28aed551..c7ba1b1b 100644 --- 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 index f96465a3..91bcee38 100644 --- 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 index 0ebe9bfe..230d4e71 100644 --- 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 {