From a1b3bec9bf174159a1c18785f72ae98f57971271 Mon Sep 17 00:00:00 2001
From: Yasir Ekinci <yas.ekinci@grafana.com>
Date: Tue, 17 Oct 2023 11:20:57 +0200
Subject: [PATCH] Improve vector health check

---
 pkg/plugin/health.go                 | 15 +++++++++------
 pkg/plugin/health_test.go            |  4 ++++
 pkg/plugin/vector/service.go         |  5 +++++
 pkg/plugin/vector/store/qdrant.go    | 11 +++++++++++
 pkg/plugin/vector/store/store.go     |  1 +
 pkg/plugin/vector/store/vectorapi.go | 11 +++++++++++
 6 files changed, 41 insertions(+), 6 deletions(-)

diff --git a/pkg/plugin/health.go b/pkg/plugin/health.go
index 64733bd9..f8611159 100644
--- a/pkg/plugin/health.go
+++ b/pkg/plugin/health.go
@@ -13,7 +13,6 @@ import (
 )
 
 var openAIModels = []string{"gpt-3.5-turbo", "gpt-4"}
-var vectorCollections = []string{"grafana.core.dashboards"}
 
 type healthCheckClient interface {
 	Do(req *http.Request) (*http.Response, error)
@@ -118,12 +117,12 @@ func (a *App) openAIHealth(ctx context.Context) (openAIHealthDetails, error) {
 	return d, nil
 }
 
-func (a *App) testVectorService(ctx context.Context) error {
+func (a *App) testVectorService(ctx context.Context) (bool, error) {
 	if a.vectorService == nil {
-		return fmt.Errorf("vector service not configured")
+		return false, fmt.Errorf("vector service not configured")
 	}
-	_, err := a.vectorService.Search(ctx, vectorCollections[0], "test", 1, nil)
-	return err
+	result, err := a.vectorService.Health(ctx)
+	return result, err
 }
 
 func (a *App) vectorHealth(ctx context.Context) vectorHealthDetails {
@@ -135,11 +134,15 @@ func (a *App) vectorHealth(ctx context.Context) vectorHealthDetails {
 		d.OK = false
 		return d
 	}
-	err := a.testVectorService(ctx)
+	result, err := a.testVectorService(ctx)
 	if err != nil {
 		d.OK = false
 		d.Error = err.Error()
 	}
+	if !result {
+		d.OK = false
+		d.Error = "Vector service health check failed"
+	}
 	return d
 }
 
diff --git a/pkg/plugin/health_test.go b/pkg/plugin/health_test.go
index e24b4c9f..c5e13f9e 100644
--- a/pkg/plugin/health_test.go
+++ b/pkg/plugin/health_test.go
@@ -27,6 +27,10 @@ func (m *mockVectorService) Search(ctx context.Context, collection string, query
 	return []store.SearchResult{{Payload: map[string]any{"a": "b"}, Score: 1.0}}, nil
 }
 
+func (m *mockVectorService) Health(ctx context.Context) (bool, error) {
+	return true, nil
+}
+
 func (m *mockVectorService) Cancel() {}
 
 // TestCheckHealth tests CheckHealth calls, using backend.CheckHealthRequest and backend.CheckHealthResponse.
diff --git a/pkg/plugin/vector/service.go b/pkg/plugin/vector/service.go
index 6120c7e8..6a55ad07 100644
--- a/pkg/plugin/vector/service.go
+++ b/pkg/plugin/vector/service.go
@@ -13,6 +13,7 @@ import (
 
 type Service interface {
 	Search(ctx context.Context, collection string, query string, topK uint64, filter map[string]interface{}) ([]store.SearchResult, error)
+	Health(ctx context.Context) (bool, error)
 	Cancel()
 }
 
@@ -86,6 +87,10 @@ func (v *vectorService) Search(ctx context.Context, collection string, query str
 	return results, nil
 }
 
+func (v *vectorService) Health(ctx context.Context) (bool, error) {
+	return v.store.Health(ctx)
+}
+
 func (v vectorService) Cancel() {
 	if v.cancel != nil {
 		v.cancel()
diff --git a/pkg/plugin/vector/store/qdrant.go b/pkg/plugin/vector/store/qdrant.go
index 573bba0e..4123627d 100644
--- a/pkg/plugin/vector/store/qdrant.go
+++ b/pkg/plugin/vector/store/qdrant.go
@@ -62,6 +62,17 @@ func newQdrantStore(s qdrantSettings, secrets map[string]string) (ReadVectorStor
 	}, cancel, nil
 }
 
+func (q *qdrantStore) Health(ctx context.Context) (bool, error) {
+	if q.md != nil {
+		ctx = metadata.NewOutgoingContext(ctx, *q.md)
+	}
+	_, err := q.collectionsClient.List(ctx, &qdrant.ListCollectionsRequest{}, grpc.WaitForReady(true))
+	if err != nil {
+		return false, err
+	}
+	return true, nil
+}
+
 func (q *qdrantStore) CollectionExists(ctx context.Context, collection string) (bool, error) {
 	if q.md != nil {
 		ctx = metadata.NewOutgoingContext(ctx, *q.md)
diff --git a/pkg/plugin/vector/store/store.go b/pkg/plugin/vector/store/store.go
index e7f31a6c..a120bc1d 100644
--- a/pkg/plugin/vector/store/store.go
+++ b/pkg/plugin/vector/store/store.go
@@ -21,6 +21,7 @@ type SearchResult struct {
 type ReadVectorStore interface {
 	CollectionExists(ctx context.Context, collection string) (bool, error)
 	Search(ctx context.Context, collection string, vector []float32, topK uint64, filter map[string]interface{}) ([]SearchResult, error)
+	Health(ctx context.Context) (bool, error)
 }
 
 type WriteVectorStore interface {
diff --git a/pkg/plugin/vector/store/vectorapi.go b/pkg/plugin/vector/store/vectorapi.go
index 84cdddc7..8cd25676 100644
--- a/pkg/plugin/vector/store/vectorapi.go
+++ b/pkg/plugin/vector/store/vectorapi.go
@@ -86,6 +86,17 @@ func (g *grafanaVectorAPI) Search(ctx context.Context, collection string, vector
 	return results, nil
 }
 
+func (g *grafanaVectorAPI) Health(ctx context.Context) (bool, error) {
+	resp, err := g.client.Get(g.url + "/healthz")
+	if err != nil {
+		return false, fmt.Errorf("get health: %w", err)
+	}
+	if resp.StatusCode != http.StatusOK {
+		return false, fmt.Errorf("get health: %s", resp.Status)
+	}
+	return true, nil
+}
+
 func newGrafanaVectorAPI(s grafanaVectorAPISettings, secrets map[string]string) ReadVectorStore {
 	return &grafanaVectorAPI{
 		client: &http.Client{},