add support for decoding response with brotli

go/ai-proxy/go.mod

@@ -12,6 +12,7 @@ require (
 )
 
 require (
+	github.com/andybalholm/brotli v1.1.1 // indirect
 	github.com/bytedance/sonic v1.12.3 // indirect
 	github.com/bytedance/sonic/loader v0.2.1 // indirect
 	github.com/cenkalti/backoff v2.2.1+incompatible // indirect

go/ai-proxy/go.sum

@@ -1,3 +1,5 @@
+github.com/andybalholm/brotli v1.1.1 h1:PR2pgnyFznKEugtsUo0xLdDop5SKXd5Qf5ysW+7XdTA=
+github.com/andybalholm/brotli v1.1.1/go.mod h1:05ib4cKhjx3OQYUY22hTVd34Bc8upXjOLL2rKwwZBoA=
 github.com/bytedance/sonic v1.12.3 h1:W2MGa7RCU1QTeYRTPE3+88mVC0yXmsRQRChiyVocVjU=
 github.com/bytedance/sonic v1.12.3/go.mod h1:B8Gt/XvtZ3Fqj+iSKMypzymZxw/FVwgIGKzMzT9r/rk=
 github.com/bytedance/sonic/loader v0.1.1/go.mod h1:ncP89zfokxS5LZrJxl5z0UJcsk4M4yY2JpfqGeCtNLU=
@@ -89,6 +91,7 @@ github.com/twitchyliquid64/golang-asm v0.15.1 h1:SU5vSMR7hnwNxj24w34ZyCi/FmDZTkS
 github.com/twitchyliquid64/golang-asm v0.15.1/go.mod h1:a1lVb/DtPvCB8fslRZhAngC2+aY1QWCk3Cedj/Gdt08=
 github.com/ugorji/go/codec v1.2.12 h1:9LC83zGrHhuUA9l16C9AHXAqEV/2wBQ4nkvumAE65EE=
 github.com/ugorji/go/codec v1.2.12/go.mod h1:UNopzCgEMSXjBc6AOMqYvWC1ktqTAfzJZUZgYf6w6lg=
+github.com/xyproto/randomstring v1.0.5/go.mod h1:rgmS5DeNXLivK7YprL0pY+lTuhNQW3iGxZ18UQApw/E=
 golang.org/x/arch v0.11.0 h1:KXV8WWKCXm6tRpLirl2szsO5j/oOODwZf4hATmGVNs4=
 golang.org/x/arch v0.11.0/go.mod h1:FEVrYAQjsQXMVJ1nsMoVVXPZg6p2JE2mx8psSWTDQys=
 golang.org/x/crypto v0.28.0 h1:GBDwsMXVQi34v5CCYUm2jkJvu4cbtru2U4TN2PSyQnw=

go/ai-proxy/main.go

@@ -4,29 +4,20 @@ import (
 	"net/http"
 	"os"
 
-	"github.com/gorilla/mux"
 	"k8s.io/klog/v2"
 
 	"github.com/pluralsh/console/go/ai-proxy/args"
 	"github.com/pluralsh/console/go/ai-proxy/environment"
 	"github.com/pluralsh/console/go/ai-proxy/internal/log"
+	"github.com/pluralsh/console/go/ai-proxy/router"
 )
 
-var (
-	router *mux.Router
-)
-
-func init() {
-	router = mux.NewRouter()
-
-	// Registers all routes under /api group
-	register(args.Provider(), args.ProviderHost(), args.ProviderToken())
-}
-
 func main() {
+	handler := router.NewRouter()
+
 	klog.V(log.LogLevelMinimal).InfoS("Starting AI Proxy", "provider", args.Provider(), "version", environment.Version, "commit", environment.Commit)
 	klog.V(log.LogLevelMinimal).InfoS("Listening and serving HTTP", "address", args.Address())
-	if err := http.ListenAndServe(args.Address(), router); err != nil {
+	if err := http.ListenAndServe(args.Address(), handler); err != nil {
 		klog.ErrorS(err, "Could not run the router")
 		os.Exit(1)
 	}

go/ai-proxy/proxy/provider/base.go

@@ -1,15 +1,23 @@
 package provider
 
 import (
+	"bytes"
 	"fmt"
+	"io"
 	"net/http"
 	"net/http/httputil"
 	"net/url"
 	"strings"
 
+	"github.com/andybalholm/brotli"
+	"k8s.io/klog/v2"
+
 	"github.com/pluralsh/console/go/ai-proxy/api"
 	"github.com/pluralsh/console/go/ai-proxy/internal/log"
-	"k8s.io/klog/v2"
+)
+
+const (
+	headerContentEncoding = "Content-Encoding"
 )
 
 type baseTranslationProxy struct {
@@ -41,7 +49,26 @@ func (in *baseTranslationProxy) ModifyRequest(r *httputil.ProxyRequest) {
 		"to", targetURL)
 }
 
-func (in *baseTranslationProxy) ModifyResponse(*http.Response) error {
+func (in *baseTranslationProxy) ModifyResponse(r *http.Response) error {
+	var reader io.Reader
+	contentEncoding := r.Header.Get(headerContentEncoding)
+	if len(contentEncoding) == 0 {
+		return nil
+	}
+
+	switch strings.TrimSpace(contentEncoding) {
+	case "br": // brotli
+		r.Header.Del(headerContentEncoding)
+		reader = brotli.NewReader(r.Body)
+	}
+
+	respBody, err := io.ReadAll(reader)
+	if err != nil {
+		return err
+	}
+
+	r.ContentLength = int64(len(respBody))
+	r.Body = io.NopCloser(bytes.NewReader(respBody))
 	return nil
 }
 

go/ai-proxy/proxy/provider/openai.go

@@ -30,6 +30,10 @@ func (in *OpenAIProxy) ModifyRequest(r *httputil.ProxyRequest) {
 }
 
 func (in *OpenAIProxy) ModifyResponse(r *http.Response) error {
+	if err := in.baseTranslationProxy.ModifyResponse(r); err != nil {
+		return err
+	}
+
 	err := in.modifyResponseBody(r)
 	if err != nil {
 		klog.ErrorS(err, "failed to map request body")

go/ai-proxy/register.go → go/ai-proxy/router/register.go

@@ -1,17 +1,21 @@
-package main
+package router
 
 import (
+	"github.com/gorilla/mux"
+
 	"github.com/pluralsh/console/go/ai-proxy/api"
 	"github.com/pluralsh/console/go/ai-proxy/api/ollama"
 	"github.com/pluralsh/console/go/ai-proxy/proxy"
 )
 
-func register(provider api.Provider, host string, token string) {
+func register(router *mux.Router, provider api.Provider, host string, token string) {
 	p, err := proxy.NewOllamaTranslationProxy(provider, host, token)
 	if err != nil {
 		panic(err)
 	}
 
+	router.Use()
+
 	// Register all Ollama API routes that should be proxied.
 	router.HandleFunc(ollama.EndpointChat, p.Proxy())
 }

go/ai-proxy/router/router.go

@@ -0,0 +1,18 @@
+package router
+
+import (
+	"net/http"
+
+	"github.com/gorilla/mux"
+
+	"github.com/pluralsh/console/go/ai-proxy/args"
+)
+
+func NewRouter() http.Handler {
+	router := mux.NewRouter()
+
+	// Registers all routes under /api group
+	register(router, args.Provider(), args.ProviderHost(), args.ProviderToken())
+
+	return router
+}

go/ai-proxy/test/main_test.go

@@ -0,0 +1,31 @@
+package router_test
+
+import (
+	"testing"
+)
+
+func TestOpenAIProxy(t *testing.T) {
+	//applicationHandler := router.NewRouter()
+	//
+	//ts := httptest.NewServer(applicationHandler)
+	//defer ts.Close()
+	//
+	//res, err := http.Get(ts.URL)
+	//
+	//if err != nil {
+	//	t.Fatalf("%s", err)
+	//}
+	//
+	//defer res.Body.Close()
+	//greeting, err := ioutil.ReadAll(res.Body)
+	//
+	//if err != nil {
+	//	log.Fatalf("%s", err)
+	//}
+	//
+	//want := []byte("Hello World!")
+	//
+	//if !bytes.Equal(want, greeting) {
+	//	t.Errorf("Expected greeting %s; got %s", want, greeting)
+	//}
+}