Merge pull request #33 from grafana/vector-service

Add vector service concept, and resource endpoint to handle vector search

cspell.config.json

@@ -27,14 +27,20 @@
     "AIURL",
     "datasource",
     "datasources",
+    "embedder",
+    "errcheck",
     "eventsource",
     "grafana",
     "httpadapter",
     "instancemgmt",
     "llms",
+    "nolint",
     "openai",
     "proxying",
+    "qdrant",
     "testid",
-    "unmarshalling"
+    "unmarshalling",
+    "Upsert",
+    "vectorapi"
   ]
 }

docker-compose.yaml

@@ -6,9 +6,29 @@ services:
     build:
       context: ./.config
       args:
-        grafana_version: ${GRAFANA_VERSION:-9.5.2}
+        grafana_version: ${GRAFANA_VERSION:-10.1.0}
+    environment:
+      OPENAI_API_KEY: $OPENAI_API_KEY
     ports:
       - 3000:3000/tcp
     volumes:
       - ./dist:/var/lib/grafana/plugins/grafana-llm-app
       - ./provisioning:/etc/grafana/provisioning
+
+  qdrant:
+    image: qdrant/qdrant
+    volumes:
+      - qdrant-storage:/qdrant/storage
+
+  # vectorapi:
+  #   image: vectorapi:latest
+  #   environment:
+  #     PORT: 8889
+  #   ports:
+  #     - 8889:8889
+  #   volumes:
+  #     - sentence-transformers:/app/.sentence_transfomers
+
+volumes:
+  # sentence-transformers:
+  qdrant-storage:

go.mod

@@ -5,6 +5,8 @@ go 1.19
 require (
 	github.com/grafana/grafana-plugin-sdk-go v0.175.0
 	github.com/launchdarkly/eventsource v1.7.1
+	github.com/qdrant/go-client v1.5.0
+	google.golang.org/grpc v1.57.0
 )
 
 require (
@@ -80,7 +82,6 @@ require (
 	golang.org/x/xerrors v0.0.0-20200804184101-5ec99f83aff1 // indirect
 	google.golang.org/genproto/googleapis/api v0.0.0-20230731193218-e0aa005b6bdf // indirect
 	google.golang.org/genproto/googleapis/rpc v0.0.0-20230731193218-e0aa005b6bdf // indirect
-	google.golang.org/grpc v1.57.0 // indirect
 	google.golang.org/protobuf v1.31.0 // indirect
 	gopkg.in/fsnotify/fsnotify.v1 v1.4.7 // indirect
 	gopkg.in/yaml.v2 v2.4.0 // indirect

go.sum

@@ -219,6 +219,8 @@ github.com/prometheus/common v0.42.0 h1:EKsfXEYo4JpWMHH5cg+KOUWeuJSov1Id8zGR8eeI
 github.com/prometheus/common v0.42.0/go.mod h1:xBwqVerjNdUDjgODMpudtOMwlOwf2SaTr1yjz4b7Zbc=
 github.com/prometheus/procfs v0.8.0 h1:ODq8ZFEaYeCaZOJlZZdJA2AbQR98dSHSM1KW/You5mo=
 github.com/prometheus/procfs v0.8.0/go.mod h1:z7EfXMXOkbkqb9IINtpCn86r/to3BnA0uaxHdg830/4=
+github.com/qdrant/go-client v1.5.0 h1:JxIp4oTFqZqsX7K5LZbfr2fCgfiT/mRsFaxJ1qCpdVM=
+github.com/qdrant/go-client v1.5.0/go.mod h1:680gkxNAsVtre0Z8hAQmtPzJtz1xFAyCu2TUxULtnoE=
 github.com/rogpeppe/fastuuid v1.2.0/go.mod h1:jVj6XXZzXRy/MSR5jhDC/2q6DgLz+nrA6LYCDYWNEvQ=
 github.com/rogpeppe/go-charset v0.0.0-20180617210344-2471d30d28b4/go.mod h1:qgYeAmZ5ZIpBWTGllZSQnw97Dj+woV0toclVaRGI8pc=
 github.com/rogpeppe/go-internal v1.11.0 h1:cWPaGQEPrBb5/AsnsZesgZZ9yb1OQ+GOISoDNXVBh4M=

pkg/main.go

@@ -9,6 +9,7 @@ import (
 )
 
 func main() {
+	log.DefaultLogger.Info("Starting plugin process")
 	// Start listening to requests sent from Grafana. This call is blocking so
 	// it won't finish until Grafana shuts down the process or the plugin choose
 	// to exit by itself using os.Exit. Manage automatically manages life cycle

pkg/plugin/app.go

@@ -2,11 +2,12 @@ package plugin
 
 import (
 	"context"
-	"encoding/json"
 	"net/http"
 
+	"github.com/grafana/grafana-llm-app/pkg/plugin/vector"
 	"github.com/grafana/grafana-plugin-sdk-go/backend"
 	"github.com/grafana/grafana-plugin-sdk-go/backend/instancemgmt"
+	"github.com/grafana/grafana-plugin-sdk-go/backend/log"
 	"github.com/grafana/grafana-plugin-sdk-go/backend/resource/httpadapter"
 )
 
@@ -21,37 +22,16 @@ var (
 	_ backend.StreamHandler         = (*App)(nil)
 )
 
-const openAIKey = "openAIKey"
-
-type OpenAISettings struct {
-	URL            string `json:"url"`
-	OrganizationID string `json:"organizationId"`
-	apiKey         string
-}
-
-type Settings struct {
-	OpenAI OpenAISettings `json:"openAI"`
-}
-
-func loadSettings(appSettings backend.AppInstanceSettings) Settings {
-	settings := Settings{
-		OpenAI: OpenAISettings{
-			URL: "https://api.openai.com",
-		},
-	}
-	_ = json.Unmarshal(appSettings.JSONData, &settings)
-
-	settings.OpenAI.apiKey = appSettings.DecryptedSecureJSONData[openAIKey]
-	return settings
-}
-
 // App is an example app backend plugin which can respond to data queries.
 type App struct {
 	backend.CallResourceHandler
+
+	vectorService vector.Service
 }
 
 // NewApp creates a new example *App instance.
 func NewApp(appSettings backend.AppInstanceSettings) (instancemgmt.Instance, error) {
+	log.DefaultLogger.Debug("Creating new app instance")
 	var app App
 
 	// Use a httpadapter (provided by the SDK) for resource calls. This allows us
@@ -61,17 +41,36 @@ func NewApp(appSettings backend.AppInstanceSettings) (instancemgmt.Instance, err
 	app.registerRoutes(mux)
 	app.CallResourceHandler = httpadapter.New(mux)
 
+	log.DefaultLogger.Debug("Loading settings")
+	settings := loadSettings(appSettings)
+	var err error
+
+	if settings.Vector.Enabled {
+		log.DefaultLogger.Debug("Creating vector service")
+		app.vectorService, err = vector.NewService(
+			settings.Vector,
+			appSettings.DecryptedSecureJSONData,
+		)
+		if err != nil {
+			log.DefaultLogger.Error("Error creating vector service", "err", err)
+			return nil, err
+		}
+	}
+
 	return &app, nil
 }
 
 // Dispose here tells plugin SDK that plugin wants to clean up resources when a new instance
 // created.
 func (a *App) Dispose() {
-	// cleanup
+	if a.vectorService != nil {
+		a.vectorService.Cancel()
+	}
 }
 
 // CheckHealth handles health checks sent from Grafana to the plugin.
 func (a *App) CheckHealth(_ context.Context, _ *backend.CheckHealthRequest) (*backend.CheckHealthResult, error) {
+	log.DefaultLogger.Info("check health")
 	return &backend.CheckHealthResult{
 		Status:  backend.HealthStatusOk,
 		Message: "ok",

pkg/plugin/resources.go

@@ -7,6 +7,7 @@ import (
 	"net/url"
 	"strings"
 
+	"github.com/grafana/grafana-llm-app/pkg/plugin/vector/store"
 	"github.com/grafana/grafana-plugin-sdk-go/backend/log"
 	"github.com/grafana/grafana-plugin-sdk-go/backend/resource/httpadapter"
 )
@@ -62,9 +63,52 @@ func newOpenAIProxy() http.Handler {
 	}
 }
 
+type vectorSearchRequest struct {
+	Text       string `json:"text"`
+	Collection string `json:"collection"`
+	Limit      uint64 `json:"limit"`
+}
+
+type vectorSearchResponse struct {
+	Results []store.SearchResult `json:"results"`
+}
+
+func (app *App) handleVectorSearch(w http.ResponseWriter, req *http.Request) {
+	if app.vectorService == nil {
+		w.WriteHeader(http.StatusServiceUnavailable)
+		return
+	}
+	if req.Method != http.MethodPost {
+		w.WriteHeader(http.StatusMethodNotAllowed)
+		return
+	}
+	body := vectorSearchRequest{}
+	if err := json.NewDecoder(req.Body).Decode(&body); err != nil {
+		http.Error(w, err.Error(), http.StatusBadRequest)
+		return
+	}
+	if body.Limit == 0 {
+		body.Limit = 10
+	}
+	results, err := app.vectorService.Search(req.Context(), body.Collection, body.Text, body.Limit)
+	if err != nil {
+		http.Error(w, err.Error(), http.StatusInternalServerError)
+		return
+	}
+	resp := vectorSearchResponse{Results: results}
+	bodyJSON, err := json.Marshal(resp)
+	if err != nil {
+		http.Error(w, err.Error(), http.StatusInternalServerError)
+		return
+	}
+	//nolint:errcheck // Just do our best to write.
+	w.Write(bodyJSON)
+}
+
 // registerRoutes takes a *http.ServeMux and registers some HTTP handlers.
 func (a *App) registerRoutes(mux *http.ServeMux) {
 	mux.HandleFunc("/ping", a.handlePing)
 	mux.HandleFunc("/echo", a.handleEcho)
 	mux.Handle("/openai/", newOpenAIProxy())
+	mux.HandleFunc("/vector/search", a.handleVectorSearch)
 }

pkg/plugin/settings.go

@@ -0,0 +1,36 @@
+package plugin
+
+import (
+	"encoding/json"
+
+	"github.com/grafana/grafana-llm-app/pkg/plugin/vector"
+	"github.com/grafana/grafana-plugin-sdk-go/backend"
+)
+
+const openAIKey = "openAIKey"
+
+type OpenAISettings struct {
+	URL            string `json:"url"`
+	OrganizationID string `json:"organizationId"`
+	apiKey         string
+}
+
+type Settings struct {
+	OpenAI OpenAISettings `json:"openAI"`
+
+	openAIKey string
+
+	Vector vector.VectorSettings `json:"vector"`
+}
+
+func loadSettings(appSettings backend.AppInstanceSettings) Settings {
+	settings := Settings{
+		OpenAI: OpenAISettings{
+			URL: "https://api.openai.com",
+		},
+	}
+	_ = json.Unmarshal(appSettings.JSONData, &settings)
+
+	settings.openAIKey = appSettings.DecryptedSecureJSONData[openAIKey]
+	return settings
+}

pkg/plugin/vector/embed/embedder.go

@@ -0,0 +1,33 @@
+package embed
+
+import (
+	"context"
+
+	"github.com/grafana/grafana-plugin-sdk-go/backend/log"
+)
+
+type EmbedderType string
+
+const (
+	EmbedderOpenAI EmbedderType = "openai"
+)
+
+type Embedder interface {
+	Embed(ctx context.Context, model string, text string) ([]float32, error)
+}
+
+type Settings struct {
+	Type string `json:"type"`
+
+	OpenAI openAISettings `json:"openai"`
+}
+
+// NewEmbedder creates a new embedder.
+func NewEmbedder(s Settings, secrets map[string]string) (Embedder, error) {
+	switch EmbedderType(s.Type) {
+	case EmbedderOpenAI:
+		log.DefaultLogger.Debug("Creating OpenAI embedder")
+		return newOpenAIEmbedder(s.OpenAI, secrets), nil
+	}
+	return nil, nil
+}