Add vector service concept, and resource endpoint to handle vector se…

…arch

This cannibalizes some of #23 but focusses on the read path: searching a
configured vector store using the configured embedding engine.

docker-compose.yaml

@@ -6,7 +6,10 @@ services:
     build:
       context: ./.config
       args:
-        grafana_version: ${GRAFANA_VERSION:-9.5.2}
+        grafana_version: ${GRAFANA_VERSION:-10.1.0}
+    environment:
+      GF_FEATURE_TOGGLES_ENABLE: 'externalServiceAuth'
+      OPENAI_API_KEY: $OPENAI_API_KEY
     ports:
       - 3000:3000/tcp
     volumes:

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,12 +2,13 @@ package plugin
 
 import (
 	"context"
-	"encoding/json"
 	"net/http"
 
 	"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"
+	"github.com/grafana/llm/pkg/plugin/vector"
 )
 
 // Make sure App implements required interfaces. This is important to do
@@ -21,32 +22,16 @@ var (
 	_ backend.StreamHandler         = (*App)(nil)
 )
 
-const openAIKey = "openAIKey"
-
-type Settings struct {
-	OpenAIURL            string `json:"openAIUrl"`
-	OpenAIOrganizationID string `json:"openAIOrganizationId"`
-
-	openAIKey string
-}
-
-func loadSettings(appSettings backend.AppInstanceSettings) Settings {
-	settings := Settings{
-		OpenAIURL: "https://api.openai.com",
-	}
-	_ = json.Unmarshal(appSettings.JSONData, &settings)
-
-	settings.openAIKey = 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.Info("Creating new app instance")
 	var app App
 
 	// Use a httpadapter (provided by the SDK) for resource calls. This allows us
@@ -56,17 +41,23 @@ func NewApp(appSettings backend.AppInstanceSettings) (instancemgmt.Instance, err
 	app.registerRoutes(mux)
 	app.CallResourceHandler = httpadapter.New(mux)
 
+	settings := loadSettings(appSettings)
+	var err error
+	app.vectorService, err = vector.NewService(settings.EmbeddingSettings, settings.VectorStoreSettings)
+	if err != nil {
+		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
-}
+func (a *App) Dispose() {}
 
 // 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

@@ -9,6 +9,7 @@ import (
 
 	"github.com/grafana/grafana-plugin-sdk-go/backend/log"
 	"github.com/grafana/grafana-plugin-sdk-go/backend/resource/httpadapter"
+	"github.com/grafana/llm/pkg/plugin/vector/store"
 )
 
 // /api/plugins/app-with-backend/resources/ping
@@ -62,9 +63,37 @@ func newOpenAIProxy() http.Handler {
 	}
 }
 
+type vectorSearchRequest struct {
+	Text       string `json:"text"`
+	Collection string `json:"collection"`
+}
+
+type vectorSearchResponse struct {
+	Results []store.SearchResult `json:"results"`
+}
+
+func (app *App) handleVectorSearch(w http.ResponseWriter, req *http.Request) {
+	body := vectorSearchRequest{}
+	if err := json.NewDecoder(req.Body).Decode(&body); err != nil {
+		http.Error(w, err.Error(), http.StatusBadRequest)
+		return
+	}
+	results, err := app.vectorService.Search(req.Context(), body.Collection, body.Text)
+	if err != nil {
+		http.Error(w, err.Error(), http.StatusInternalServerError)
+	}
+	resp := vectorSearchResponse{Results: results}
+	bodyJSON, err := json.Marshal(resp)
+	if err != nil {
+		http.Error(w, err.Error(), http.StatusInternalServerError)
+	}
+	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,32 @@
+package plugin
+
+import (
+	"encoding/json"
+
+	"github.com/grafana/grafana-plugin-sdk-go/backend"
+	"github.com/grafana/llm/pkg/plugin/vector/embed"
+	"github.com/grafana/llm/pkg/plugin/vector/store"
+)
+
+const openAIKey = "openAIKey"
+
+type Settings struct {
+	OpenAIURL            string `json:"openAIUrl"`
+	OpenAIOrganizationID string `json:"openAIOrganizationId"`
+
+	openAIKey string
+
+	EmbeddingSettings   embed.Settings `json:"embeddings"`
+	VectorStoreSettings store.Settings `json:"vectorStore"`
+}
+
+func loadSettings(appSettings backend.AppInstanceSettings) Settings {
+	settings := Settings{
+		OpenAIURL: "https://api.openai.com",
+	}
+	_ = json.Unmarshal(appSettings.JSONData, &settings)
+
+	settings.openAIKey = appSettings.DecryptedSecureJSONData[openAIKey]
+	settings.EmbeddingSettings.OpenAI.APIKey = settings.openAIKey
+	return settings
+}

pkg/plugin/vector/embed/embedder.go

@@ -0,0 +1,28 @@
+package embed
+
+import "context"
+
+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) (Embedder, error) {
+	switch EmbedderType(s.Type) {
+	case EmbedderOpenAI:
+		return newOpenAIEmbedder(s.OpenAI), nil
+	}
+	return nil, nil
+}

pkg/plugin/vector/embed/openai.go

@@ -0,0 +1,93 @@
+package embed
+
+import (
+	"bytes"
+	"context"
+	"encoding/json"
+	"fmt"
+	"io"
+	"net/http"
+	"strings"
+
+	"github.com/grafana/grafana-plugin-sdk-go/backend/log"
+)
+
+type openAISettings struct {
+	URL    string `json:"url"`
+	APIKey string `json:"apiKey"`
+}
+
+type openAILLMClient struct {
+	client *http.Client
+	url    string
+	apiKey string
+}
+
+type openAIEmbeddingsRequest struct {
+	Model string `json:"model"`
+	Input string `json:"input"`
+}
+
+type openAIEmbeddingsResponse struct {
+	Data []openAIEmbeddingData `json:"data"`
+}
+
+type openAIEmbeddingData struct {
+	Embedding []float32 `json:"embedding"`
+}
+
+func (o *openAILLMClient) Embed(ctx context.Context, model string, payload string) ([]float32, error) {
+	// TODO: ensure payload is under 8191 tokens, somehow.
+	url := o.url
+	if url == "" {
+		url = "https://api.openai.com"
+	}
+	url = strings.TrimSuffix(url, "/")
+	url = url + "/v1/embeddings"
+	reqBody := openAIEmbeddingsRequest{
+		Model: model,
+		Input: payload,
+	}
+	bodyJSON, err := json.Marshal(reqBody)
+	if err != nil {
+		return nil, fmt.Errorf("marshal request: %w", err)
+	}
+	req, err := http.NewRequestWithContext(ctx, http.MethodPost, url, bytes.NewReader(bodyJSON))
+	if err != nil {
+		return nil, fmt.Errorf("create request: %w", err)
+	}
+	req.Header.Set("Content-Type", "application/json")
+	req.Header.Set("Authorization", "Bearer "+o.apiKey)
+	resp, err := o.client.Do(req)
+	if err != nil {
+		return nil, fmt.Errorf("make request: %w", err)
+	}
+	defer func() {
+		if err := resp.Body.Close(); err != nil {
+			log.DefaultLogger.Warn("failed to close response body", "err", err)
+		}
+	}()
+	if resp.StatusCode/100 != 2 {
+		return nil, fmt.Errorf("got non-2xx status from OpenAI: %s", resp.Status)
+	}
+	respBody, err := io.ReadAll(io.LimitReader(resp.Body, 1024*1024*2))
+	if err != nil {
+		return nil, fmt.Errorf("read response body: %w", err)
+	}
+	var body openAIEmbeddingsResponse
+	err = json.Unmarshal(respBody, &body)
+	if err != nil {
+		return nil, fmt.Errorf("unmarshal response body: %w", err)
+	}
+	return body.Data[0].Embedding, nil
+}
+
+// newOpenAIEmbedder creates a new Embedder using OpenAI's API.
+func newOpenAIEmbedder(settings openAISettings) Embedder {
+	impl := openAILLMClient{
+		client: &http.Client{},
+		url:    settings.URL,
+		apiKey: settings.APIKey,
+	}
+	return &impl
+}

pkg/plugin/vector/service.go

@@ -0,0 +1,70 @@
+// package vector provides a service for searching vector embeddings.
+// It combines the embedding engine and the vector store.
+package vector
+
+import (
+	"context"
+	"fmt"
+
+	"github.com/grafana/llm/pkg/plugin/vector/embed"
+	"github.com/grafana/llm/pkg/plugin/vector/store"
+)
+
+type Collection struct {
+	Name      string `json:"name"`
+	Dimension int    `json:"dimension"`
+	Model     string `json:"model"`
+}
+
+type Service interface {
+	Search(ctx context.Context, collection string, query string) ([]store.SearchResult, error)
+}
+
+type vectorService struct {
+	embedder         embed.Embedder
+	store            store.ReadVectorStore
+	collectionConfig map[string]Collection
+}
+
+func NewService(embedSettings embed.Settings, storeSettings store.Settings) (Service, error) {
+	em, err := embed.NewEmbedder(embedSettings)
+	if err != nil {
+		return nil, fmt.Errorf("new embedder: %w", err)
+	}
+	if em == nil {
+		return nil, nil
+	}
+	st, err := store.NewReadVectorStore(storeSettings)
+	if err != nil {
+		return nil, fmt.Errorf("new vector store: %w", err)
+	}
+	if st == nil {
+		return nil, nil
+	}
+	return &vectorService{
+		embedder: em,
+		store:    st,
+	}, nil
+}
+
+func (g vectorService) Search(ctx context.Context, collection string, query string) ([]store.SearchResult, error) {
+	// Determine which model was used to embed this collection.
+	c := g.collectionConfig[collection]
+	if c.Name == "" {
+		return nil, fmt.Errorf("unknown collection %s", collection)
+	}
+
+	// Get the embedding for the search query.
+	e, err := g.embedder.Embed(ctx, c.Model, query)
+	if err != nil {
+		return nil, fmt.Errorf("embed query: %w", err)
+	}
+
+	// Search the vector store for similar vectors.
+	results, err := g.store.Search(ctx, collection, e, 10)
+	if err != nil {
+		return nil, fmt.Errorf("vector store search: %w", err)
+	}
+
+	return results, nil
+}