Improve vector health check

grafana · Oct 17, 2023 · a1b3bec · a1b3bec
1 parent bbca65b
commit a1b3bec
Show file tree

Hide file tree

Showing 6 changed files with 41 additions and 6 deletions.
diff --git 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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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{},