Added the azure openai proxy resource endpoint

CHANGELOG.md

@@ -3,6 +3,7 @@
 ## Unreleased
 
 * Add Go package providing an OpenAI client to use the LLM app from backend Go code
+* Add support for Azure OpenAI. The plugin must be configured to use OpenAI and provide a link between OpenAI model names and Azure deployment names
 
 ## 0.2.1
 

cspell.config.json

@@ -37,6 +37,7 @@
     "llms",
     "nolint",
     "openai",
+    "proxied",
     "proxying",
     "qdrant",
     "testid",

pkg/plugin/app.go

@@ -34,15 +34,16 @@ func NewApp(ctx context.Context, appSettings backend.AppInstanceSettings) (insta
 	log.DefaultLogger.Debug("Creating new app instance")
 	var app App
 
+	log.DefaultLogger.Debug("Loading settings")
+	settings := loadSettings(appSettings)
+
 	// Use a httpadapter (provided by the SDK) for resource calls. This allows us
 	// to use a *http.ServeMux for resource calls, so we can map multiple routes
 	// to CallResource without having to implement extra logic.
 	mux := http.NewServeMux()
-	app.registerRoutes(mux)
+	app.registerRoutes(mux, settings)
 	app.CallResourceHandler = httpadapter.New(mux)
 
-	log.DefaultLogger.Debug("Loading settings")
-	settings := loadSettings(appSettings)
 	var err error
 
 	if settings.Vector.Enabled {

pkg/plugin/resources.go

@@ -1,29 +1,168 @@
 package plugin
 
 import (
+	"bytes"
 	"encoding/json"
+	"fmt"
+	"io"
 	"net/http"
 	"net/http/httputil"
 	"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"
 )
 
-func newOpenAIProxy() http.Handler {
-	return &httputil.ReverseProxy{
-		Rewrite: func(r *httputil.ProxyRequest) {
-			config := httpadapter.PluginConfigFromContext(r.In.Context())
-			settings := loadSettings(*config.AppInstanceSettings)
-			u, _ := url.Parse(settings.OpenAI.URL)
-			r.SetURL(u)
-			r.Out.Header.Set("Authorization", "Bearer "+settings.OpenAI.apiKey)
-			organizationID := settings.OpenAI.OrganizationID
-			r.Out.Header.Set("OpenAI-Organization", organizationID)
-			r.Out.URL.Path = strings.TrimPrefix(r.In.URL.Path, "/openai")
-			log.DefaultLogger.Info("proxying to url", "url", r.Out.URL.String())
+// modifyURL modifies the request URL to point to the configured OpenAI API.
+func modifyURL(openAI OpenAISettings, req *http.Request) error {
+	u, err := url.Parse(openAI.URL)
+	if err != nil {
+		log.DefaultLogger.Error("Unable to parse OpenAI URL", "err", err)
+		return fmt.Errorf("parse OpenAI URL: %w", err)
+	}
+	req.URL.Scheme = u.Scheme
+	req.URL.Host = u.Host
+	return nil
+}
+
+// openAIProxy is a reverse proxy for OpenAI 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 OpenAI organization.
+type openAIProxy 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 (a *openAIProxy) ServeHTTP(w http.ResponseWriter, req *http.Request) {
+	err := modifyURL(a.settings.OpenAI, 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)
+		}
+	}
+	a.rp.ServeHTTP(w, req)
+}
+
+func newOpenAIProxy(settings Settings) http.Handler {
+	director := func(req *http.Request) {
+		req.URL.Path = strings.TrimPrefix(req.URL.Path, "/openai")
+		req.Header.Add("Authorization", "Bearer "+settings.OpenAI.apiKey)
+		req.Header.Add("OpenAI-Organization", settings.OpenAI.OrganizationID)
+	}
+	return &openAIProxy{
+		settings: settings,
+		rp:       &httputil.ReverseProxy{Director: director},
+	}
+}
+
+// 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
+// to a model without a corresponding deployment in settings. It then proxies the request
+// using the configured API key and deployment.
+type azureOpenAIProxy 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 (a *azureOpenAIProxy) modifyRequest(req *http.Request) error {
+	err := modifyURL(a.settings.OpenAI, req)
+	if err != nil {
+		return fmt.Errorf("modify url: %w", err)
+	}
+
+	// Read the body so we can determine the deployment to use
+	// by mapping the model in the request to a deployment in settings.
+	// Azure OpenAI API requires this deployment name in the URL.
+	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)
+	}
+
+	// Find the deployment for the model.
+	// Models are mapped to deployments in settings.OpenAI.AzureMapping.
+	var deployment string = ""
+	for _, v := range a.settings.OpenAI.AzureMapping {
+		if val, ok := requestBody["model"].(string); ok && val == v[0] {
+			deployment = v[1]
+			break
+		}
+	}
+
+	if deployment == "" {
+		return fmt.Errorf("no deployment found for model: %s", requestBody["model"])
+	}
+
+	// We've got a deployment, so finish modifying the request.
+	req.URL.Path = fmt.Sprintf("/openai/deployments/%s/%s", deployment, strings.TrimPrefix(req.URL.Path, "/openai/v1/"))
+	req.Header.Add("api-key", a.settings.OpenAI.apiKey)
+	req.URL.RawQuery = "api-version=2023-03-15-preview"
+
+	// Remove extra fields
+	delete(requestBody, "model")
+
+	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 (a *azureOpenAIProxy) ServeHTTP(w http.ResponseWriter, req *http.Request) {
+	err := a.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
+	}
+	a.rp.ServeHTTP(w, req)
+}
+
+func newAzureOpenAIProxy(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 &azureOpenAIProxy{
+		settings: settings,
+		rp: &httputil.ReverseProxy{
+			Director: director,
 		},
 	}
 }
@@ -71,7 +210,11 @@ func (app *App) handleVectorSearch(w http.ResponseWriter, req *http.Request) {
 }
 
 // registerRoutes takes a *http.ServeMux and registers some HTTP handlers.
-func (a *App) registerRoutes(mux *http.ServeMux) {
-	mux.Handle("/openai/", newOpenAIProxy())
+func (a *App) registerRoutes(mux *http.ServeMux, settings Settings) {
+	if settings.OpenAI.UseAzure {
+		mux.Handle("/openai/", newAzureOpenAIProxy(settings))
+	} else {
+		mux.Handle("/openai/", newOpenAIProxy(settings))
+	}
 	mux.HandleFunc("/vector/search", a.handleVectorSearch)
 }