Initial Pulze Provider

pkg/plugin/health.go

@@ -11,7 +11,11 @@ import (
 	"github.com/grafana/grafana-plugin-sdk-go/build"
 )
 
-var openAIModels = []string{"gpt-3.5-turbo", "gpt-4"}
+// Define models for each provider to be included in the health check.
+var providerModels = map[string][]string{
+	"openai": {"gpt-3.5-turbo", "gpt-4"},
+	"pulze":  {"pulze", "openai/gpt-4"},
+}
 
 type healthCheckClient interface {
 	Do(req *http.Request) (*http.Response, error)
@@ -79,20 +83,24 @@ func (a *App) testOpenAIModel(ctx context.Context, model string) error {
 	return nil
 }
 
-// openAIHealth checks the health of the OpenAI configuration and caches the
+// openAIHealth performs a health check for the selected provider and caches the
 // result if successful. The caller must lock a.healthCheckMutex.
-func (a *App) openAIHealth(ctx context.Context, req *backend.CheckHealthRequest) (openAIHealthDetails, error) {
+func (a *App) openAIHealth(ctx context.Context, req *backend.CheckHealthRequest) openAIHealthDetails {
 	if a.healthOpenAI != nil {
-		return *a.healthOpenAI, nil
+		return *a.healthOpenAI
 	}
 
 	d := openAIHealthDetails{
 		OK:         true,
 		Configured: a.settings.OpenAI.apiKey != "" || a.settings.OpenAI.Provider == openAIProviderGrafana,
 		Models:     map[string]openAIModelHealth{},
 	}
+	models := providerModels["openai"]
+	if a.settings.OpenAI.Provider == openAIProviderPulze {
+		models = providerModels["pulze"]
+	}
 
-	for _, model := range openAIModels {
+	for _, model := range models {
 		health := openAIModelHealth{OK: false, Error: "OpenAI not configured"}
 		if d.Configured {
 			health.OK = true
@@ -121,7 +129,7 @@ func (a *App) openAIHealth(ctx context.Context, req *backend.CheckHealthRequest)
 	if d.OK {
 		a.healthOpenAI = &d
 	}
-	return d, nil
+	return d
 }
 
 // testVectorService checks the health of VectorAPI and caches the result if
@@ -137,6 +145,8 @@ func (a *App) testVectorService(ctx context.Context) error {
 	return nil
 }
 
+// vectorHealth performs a health check for the Vector service and caches the
+// result if successful. The caller must lock a.healthCheckMutex.
 func (a *App) vectorHealth(ctx context.Context) vectorHealthDetails {
 	if a.healthVector != nil {
 		return *a.healthVector
@@ -169,10 +179,9 @@ func (a *App) CheckHealth(ctx context.Context, req *backend.CheckHealthRequest)
 	a.healthCheckMutex.Lock()
 	defer a.healthCheckMutex.Unlock()
 
-	openAI, err := a.openAIHealth(ctx, req)
-	if err != nil {
-		openAI.OK = false
-		openAI.Error = err.Error()
+	openAI := a.openAIHealth(ctx, req)
+	if openAI.Error == "" {
+		a.healthOpenAI = &openAI
 	}
 
 	vector := a.vectorHealth(ctx)

pkg/plugin/openai.go

@@ -22,6 +22,8 @@ func (a *App) newAuthenticatedOpenAIRequest(ctx context.Context, method string,
 		req.Header.Set("OpenAI-Organization", a.settings.OpenAI.OrganizationID)
 	case openAIProviderAzure:
 		req.Header.Set("api-key", a.settings.OpenAI.apiKey)
+	case openAIProviderPulze:
+		req.Header.Set("Authorization", "Bearer "+a.settings.OpenAI.apiKey)
 	case openAIProviderGrafana:
 		req.SetBasicAuth(a.settings.Tenant, a.settings.GrafanaComAPIKey)
 		req.Header.Add("X-Scope-OrgID", a.settings.Tenant)
@@ -34,6 +36,8 @@ func (a *App) newOpenAIChatCompletionsRequest(ctx context.Context, body map[stri
 	var err error
 
 	switch a.settings.OpenAI.Provider {
+	case openAIProviderPulze:
+		fallthrough
 	case openAIProviderOpenAI:
 		url, err = url.Parse(a.settings.OpenAI.URL)
 		if err != nil {

pkg/plugin/resources.go

@@ -94,6 +94,83 @@ func newOpenAIProxy(settings Settings) http.Handler {
 	}
 }
 
+// pulzeOpenAIProxy is a reverse proxy for Pulze API calls.
+// It modifies the request to point to the configured OpenAI API, returning
+// a 400 error if the URL in settings cannot be parsed, then proxies the request
+// using the configured API key and setting the default pulze model if none given.
+type pulzeOpenAIProxy struct {
+	settings Settings
+	// rp is a reverse proxy handling the modified request. Use this rather than
+	// our own client, since it handles things like buffering.
+	rp *httputil.ReverseProxy
+}
+
+func newPulzeOpenAIProxy(settings Settings) http.Handler {
+	// We make all of the actual modifications in ServeHTTP, since they can fail
+	// and we want to early-return from HTTP requests in that case.
+	director := func(req *http.Request) {}
+	return &pulzeOpenAIProxy{
+		settings: settings,
+		rp:       &httputil.ReverseProxy{Director: director},
+	}
+}
+
+func (p *pulzeOpenAIProxy) modifyRequest(req *http.Request) error {
+	err := modifyURL(p.settings.OpenAI.URL, req)
+	if err != nil {
+		return fmt.Errorf("modify url: %w", err)
+	}
+
+	req.URL.Path = strings.TrimPrefix(req.URL.Path, "/pulze")
+	req.Header.Add("Authorization", "Bearer "+p.settings.OpenAI.apiKey)
+
+	// Read the body so we can determine if we need to add the configured pulze model
+	bodyBytes, _ := io.ReadAll(req.Body)
+	var requestBody map[string]interface{}
+	err = json.Unmarshal(bodyBytes, &requestBody)
+	if err != nil {
+		return fmt.Errorf("unmarshal request body: %w", err)
+	}
+
+	// check if model is empty or not present
+	if val, ok := requestBody["model"].(string); !ok || val == "" {
+		requestBody["model"] = p.settings.OpenAI.PulzeModel
+	}
+
+	newBodyBytes, err := json.Marshal(requestBody)
+	if err != nil {
+		return fmt.Errorf("unmarshal request body: %w", err)
+	}
+
+	req.Body = io.NopCloser(bytes.NewBuffer(newBodyBytes))
+	req.ContentLength = int64(len(newBodyBytes))
+	return nil
+}
+
+func (p *pulzeOpenAIProxy) ServeHTTP(w http.ResponseWriter, req *http.Request) {
+	err := p.modifyRequest(req)
+	if err != nil {
+		// Attempt to write the error as JSON.
+		jd, err := json.Marshal(map[string]string{"error": err.Error()})
+		if err != nil {
+			// We can't write JSON, so just write the error string.
+			w.WriteHeader(http.StatusInternalServerError)
+			_, err = w.Write([]byte(err.Error()))
+			if err != nil {
+				log.DefaultLogger.Error("Unable to write error response", "err", err)
+			}
+			return
+		}
+		w.WriteHeader(http.StatusBadRequest)
+		_, err = w.Write(jd)
+		if err != nil {
+			log.DefaultLogger.Error("Unable to write error response", "err", err)
+		}
+		return
+	}
+	p.rp.ServeHTTP(w, req)
+}
+
 // azureOpenAIProxy is a reverse proxy for Azure OpenAI API calls.
 // It modifies the request to point to the configured Azure OpenAI API, returning
 // a 400 error if the URL in settings cannot be parsed or if the request refers
@@ -417,6 +494,8 @@ func (a *App) registerRoutes(mux *http.ServeMux, settings Settings) {
 		mux.Handle("/openai/", newAzureOpenAIProxy(settings))
 	case openAIProviderGrafana:
 		mux.Handle("/openai/", newGrafanaOpenAIProxy(settings))
+	case openAIProviderPulze:
+		mux.Handle("/pulze/", newPulzeOpenAIProxy(settings))
 	default:
 		log.DefaultLogger.Warn("Unknown OpenAI provider configured", "provider", settings.OpenAI.Provider)
 	}

pkg/plugin/resources_test.go

@@ -274,6 +274,52 @@ func TestCallOpenAIProxy(t *testing.T) {
 			expReqPath: "/openai/v1/chat/completions",
 			expReqBody: []byte(`{"model": "gpt-3.5-turbo", "messages": ["some stuff"]}`),
 
+			expStatus: http.StatusOK,
+		},
+		{
+			name: "pulze with specific model",
+
+			settings: Settings{
+				OpenAI: OpenAISettings{
+					Provider:   openAIProviderPulze,
+					PulzeModel: "pulze",
+				},
+			},
+			apiKey: "abcd1234",
+
+			method: http.MethodPost,
+			path:   "/pulze/v1/chat/completions",
+			body:   []byte(`{"model": "openai/gpt-3.5-turbo", "messages": ["some stuff"]}`),
+
+			expReqHeaders: http.Header{
+				"Authorization": {"Bearer abcd1234"},
+			},
+			expReqPath: "/v1/chat/completions",
+			expReqBody: []byte(`{"model": "openai/gpt-3.5.turbo", "messages": ["some stuff"]}`),
+
+			expStatus: http.StatusOK,
+		},
+		{
+			name: "pulze without model",
+
+			settings: Settings{
+				OpenAI: OpenAISettings{
+					Provider:   openAIProviderPulze,
+					PulzeModel: "pulze",
+				},
+			},
+			apiKey: "abcd1234",
+
+			method: http.MethodPost,
+			path:   "/pulze/v1/chat/completions",
+			body:   []byte(`{"model": "", "messages": ["some stuff"]}`),
+
+			expReqHeaders: http.Header{
+				"Authorization": {"Bearer abcd1234"},
+			},
+			expReqPath: "/v1/chat/completions",
+			expReqBody: []byte(`{"model": "pulze", "messages": ["some stuff"]}`),
+
 			expStatus: http.StatusOK,
 		},
 	} {

pkg/plugin/settings.go

@@ -13,15 +13,17 @@ import (
 	"github.com/grafana/grafana-plugin-sdk-go/backend/log"
 )
 
-const openAIKey = "openAIKey"
-const encodedTenantAndTokenKey = "base64EncodedAccessToken"
-
 type openAIProvider string
 
 const (
 	openAIProviderOpenAI  openAIProvider = "openai"
 	openAIProviderAzure   openAIProvider = "azure"
 	openAIProviderGrafana openAIProvider = "grafana" // via llm-gateway
+	openAIProviderPulze   openAIProvider = "pulze"
+
+	openAIKey                = "openAIKey"
+	llmGatewayKey            = "llmGatewayKey"
+	encodedTenantAndTokenKey = "base64EncodedAccessToken"
 )
 
 // OpenAISettings contains the user-specified OpenAI connection details
@@ -38,6 +40,9 @@ type OpenAISettings struct {
 	// Model mappings required for Azure's OpenAI
 	AzureMapping [][]string `json:"azureModelMapping"`
 
+	// The pulze model to use
+	PulzeModel string `json:"pulzeModel"`
+
 	// apiKey is the user-specified  api key needed to authenticate requests to the OpenAI
 	// provider (excluding the LLMGateway). Stored securely.
 	apiKey string
@@ -121,6 +126,10 @@ func loadSettings(appSettings backend.AppInstanceSettings) (*Settings, error) {
 			log.DefaultLogger.Warn("Cannot use LLM Gateway as no URL specified, disabling it")
 			settings.OpenAI.Provider = ""
 		}
+	case openAIProviderPulze:
+		if settings.OpenAI.URL == "" {
+			settings.OpenAI.URL = "https://api.pulze.ai/v1"
+		}
 	default:
 		// Default to disabled LLM support if an unknown provider was specified.
 		log.DefaultLogger.Warn("Unknown OpenAI provider", "provider", settings.OpenAI.Provider)

src/components/AppConfig/AppConfig.test.tsx

@@ -1,7 +1,7 @@
+import React from 'react';
 import { PluginType } from '@grafana/data';
 import { render, screen } from '@testing-library/react';
 import { testIds } from 'components/testIds';
-import React from 'react';
 import { AppConfig, AppConfigProps } from './AppConfig';
 
 describe('Components/AppConfig', () => {
@@ -43,6 +43,7 @@ describe('Components/AppConfig', () => {
     // expect(screen.queryByTestId(testIds.appConfig.openAIUrl)).toBeInTheDocument();
     // expect(screen.queryByTestId(testIds.appConfig.model)).toBeInTheDocument();
     expect(screen.queryByRole('group', { name: /vector settings/i })).toBeInTheDocument();
+    expect(screen.queryByTestId(testIds.appConfig.model)).toBeInTheDocument();
     expect(screen.queryByTestId(testIds.appConfig.qdrantSecure)).toBeInTheDocument();
     expect(screen.queryByTestId(testIds.appConfig.qdrantAddress)).toBeInTheDocument();
     // Don't expect to see the Grafana vector API field when type is qdrant

src/components/AppConfig/OpenAI.tsx

@@ -7,7 +7,8 @@ import { testIds } from 'components/testIds';
 import { getStyles, Secrets, SecretsSet } from './AppConfig';
 import { AzureModelDeploymentConfig, AzureModelDeployments } from './AzureConfig';
 
-export type OpenAIProvider = 'openai' | 'azure' | 'grafana';
+export type OpenAIProvider = 'openai' | 'azure' | 'grafana' | 'pulze';
+export type PulzeModel = 'pulze' | 'pulze-v0'
 
 export interface OpenAISettings {
   // The URL to reach OpenAI.
@@ -18,6 +19,8 @@ export interface OpenAISettings {
   provider?: OpenAIProvider;
   // A mapping of OpenAI models to Azure deployment names.
   azureModelMapping?: AzureModelDeployments;
+  // Default pulze model to use if no model is specified.
+  pulzeModel?: PulzeModel;
 }
 
 export function OpenAIConfig({
@@ -51,6 +54,7 @@ export function OpenAIConfig({
             [
               { label: 'OpenAI', value: 'openai' },
               { label: 'Azure OpenAI', value: 'azure' },
+              { label: 'Pulze', value: 'pulze' },
             ] as Array<SelectableValue<OpenAIProvider>>
           }
           value={settings.provider ?? 'openai'}
@@ -90,7 +94,7 @@ export function OpenAIConfig({
         />
       </Field>
 
-      {settings.provider !== 'azure' && (
+      {settings.provider !== 'azure' && settings.provider !== 'pulze' && (
         <Field label="OpenAI API Organization ID" description="Your OpenAI API Organization ID">
           <Input
             width={60}
@@ -120,6 +124,26 @@ export function OpenAIConfig({
           />
         </Field>
       )}
+
+      {settings.provider === 'pulze' && (
+        <Field
+          label="Default Pulze Model"
+          description="The default pulze model to use"
+          data-testid={testIds.appConfig.pulzeModel}
+        >
+          <Select
+            options={
+              [
+                { label: 'pulze', value: 'pulze' },
+                { label: 'pulze-v0', value: 'pulze-v0' },
+              ] as Array<SelectableValue<PulzeModel>>
+            }
+            value={settings.pulzeModel ?? 'pulze'}
+            onChange={(e) => onChange({ ...settings, pulzeModel: e.value })}
+            width={60}
+          />
+        </Field>
+      )}
     </FieldSet>
   );
 }

src/components/testIds.ts

@@ -5,6 +5,7 @@ export const testIds = {
     openAIKey: 'data-testid ac-openai-api-key',
     openAIOrganizationID: 'data-testid ac-openai-api-organization-id',
     openAIUrl: 'data-testid ac-openai-api-url',
+    pulzeModel: 'data-testid ac-pulze-model',
     vectorEnabled: 'data-testid ac-vector-enabled',
     qdrantAddress: 'data-testid ac-qdrant-address',
     qdrantSecure: 'data-testid ac-qdrant-secure',