Add functionality to AI Proxy for full Open AI protocol

go/ai-proxy/api/openai_standard/openai_standard.go

@@ -0,0 +1,7 @@
+package openai_standard
+
+type Endpoint string
+
+const (
+	EndpointChat = "/openai/chat/completions"
+)

go/ai-proxy/api/provider.go

@@ -20,6 +20,8 @@ func ToProvider(s string) (Provider, error) {
 		return ProviderOllama, nil
 	case ProviderOpenAI.String():
 		return ProviderOpenAI, nil
+	case ProviderOpenAIStandard.String():
+		return ProviderOpenAIStandard, nil
 	case ProviderAnthropic.String():
 		return ProviderAnthropic, nil
 	case ProviderVertex.String():
@@ -30,10 +32,11 @@ func ToProvider(s string) (Provider, error) {
 }
 
 const (
-	ProviderOpenAI    Provider = "openai"
-	ProviderAnthropic Provider = "anthropic"
-	ProviderOllama    Provider = "ollama"
-	ProviderVertex    Provider = "vertex"
+	ProviderOpenAI         Provider = "openai"
+	ProviderOpenAIStandard Provider = "openai_standard"
+	ProviderAnthropic      Provider = "anthropic"
+	ProviderOllama         Provider = "ollama"
+	ProviderVertex         Provider = "vertex"
 )
 
 type OllamaAPI string

go/ai-proxy/api/proxy.go

@@ -10,3 +10,7 @@ type TranslationProxy interface {
 	ModifyRequest(*httputil.ProxyRequest)
 	ModifyResponse(*http.Response) error
 }
+
+type OpenAIProxy interface {
+	Proxy() http.HandlerFunc
+}

go/ai-proxy/args/args.go

@@ -75,6 +75,10 @@ func ProviderCredentials() string {
 		return *argProviderToken
 	}
 
+	if len(*argProviderToken) > 0 && Provider() == api.ProviderOpenAIStandard {
+		return *argProviderToken
+	}
+
 	if len(*argProviderServiceAccount) > 0 && Provider() == api.ProviderVertex {
 		return *argProviderServiceAccount
 	}

go/ai-proxy/go.mod

@@ -47,8 +47,14 @@ require (
 	github.com/mattn/go-isatty v0.0.20 // indirect
 	github.com/modern-go/concurrent v0.0.0-20180306012644-bacd9c7ef1dd // indirect
 	github.com/modern-go/reflect2 v1.0.2 // indirect
+	github.com/openai/openai-go v0.1.0-alpha.45 // indirect
 	github.com/pelletier/go-toml/v2 v2.2.3 // indirect
 	github.com/rogpeppe/go-internal v1.8.0 // indirect
+	github.com/sashabaranov/go-openai v1.36.1 // indirect
+	github.com/tidwall/gjson v1.14.4 // indirect
+	github.com/tidwall/match v1.1.1 // indirect
+	github.com/tidwall/pretty v1.2.1 // indirect
+	github.com/tidwall/sjson v1.2.5 // indirect
 	github.com/twitchyliquid64/golang-asm v0.15.1 // indirect
 	github.com/ugorji/go/codec v1.2.12 // indirect
 	go.opencensus.io v0.24.0 // indirect

go/ai-proxy/go.sum

@@ -120,6 +120,8 @@ github.com/modern-go/reflect2 v1.0.2 h1:xBagoLtFs94CBntxluKeaWgTMpvLxC4ur3nMaC9G
 github.com/modern-go/reflect2 v1.0.2/go.mod h1:yWuevngMOJpCy52FWWMvUC8ws7m/LJsjYzDa0/r8luk=
 github.com/ollama/ollama v0.3.14 h1:e94+Fb1PDqmD3O90g5cqUSkSxfNm9U3fHMIyaKQ8aSc=
 github.com/ollama/ollama v0.3.14/go.mod h1:YrWoNkFnPOYsnDvsf/Ztb1wxU9/IXrNsQHqcxbY2r94=
+github.com/openai/openai-go v0.1.0-alpha.45 h1:PAj4Rj+ofOIh9ziT56FaTqb0as6PoUfbKPIvlUAOy6M=
+github.com/openai/openai-go v0.1.0-alpha.45/go.mod h1:3SdE6BffOX9HPEQv8IL/fi3LYZ5TUpRYaqGQZbyk11A=
 github.com/pelletier/go-toml/v2 v2.2.3 h1:YmeHyLY8mFWbdkNWwpr+qIL2bEqT0o95WSdkNHvL12M=
 github.com/pelletier/go-toml/v2 v2.2.3/go.mod h1:MfCQTFTvCcUyyvvwm1+G6H/jORL20Xlb6rzQu9GuUkc=
 github.com/pkg/diff v0.0.0-20210226163009-20ebb0f2a09e/go.mod h1:pJLUxLENpZxwdsKMEsNbx1VGcRFpLqf3715MtcvvzbA=
@@ -133,6 +135,8 @@ github.com/rogpeppe/go-internal v1.8.0 h1:FCbCCtXNOY3UtUuHUYaghJg4y7Fd14rXifAYUA
 github.com/rogpeppe/go-internal v1.8.0/go.mod h1:WmiCO8CzOY8rg0OYDC4/i/2WRWAB6poM+XZ2dLUbcbE=
 github.com/samber/lo v1.47.0 h1:z7RynLwP5nbyRscyvcD043DWYoOcYRv3mV8lBeqOCLc=
 github.com/samber/lo v1.47.0/go.mod h1:RmDH9Ct32Qy3gduHQuKJ3gW1fMHAnE/fAzQuf6He5cU=
+github.com/sashabaranov/go-openai v1.36.1 h1:EVfRXwIlW2rUzpx6vR+aeIKCK/xylSrVYAx1TMTSX3g=
+github.com/sashabaranov/go-openai v1.36.1/go.mod h1:lj5b/K+zjTSFxVLijLSTDZuP7adOgerWeFyZLUhAKRg=
 github.com/spf13/pflag v1.0.5 h1:iy+VFUOCP1a+8yFto/drg2CJ5u0yRoB7fZw3DKv/JXA=
 github.com/spf13/pflag v1.0.5/go.mod h1:McXfInJRrz4CZXVZOBLb0bTZqETkiAhM9Iw0y3An2Bg=
 github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME=
@@ -145,6 +149,16 @@ github.com/stretchr/testify v1.8.0/go.mod h1:yNjHg4UonilssWZ8iaSj1OCr/vHnekPRkoO
 github.com/stretchr/testify v1.8.1/go.mod h1:w2LPCIKwWwSfY2zedu0+kehJoqGctiVI29o6fzry7u4=
 github.com/stretchr/testify v1.9.0 h1:HtqpIVDClZ4nwg75+f6Lvsy/wHu+3BoSGCbBAcpTsTg=
 github.com/stretchr/testify v1.9.0/go.mod h1:r2ic/lqez/lEtzL7wO/rwa5dbSLXVDPFyf8C91i36aY=
+github.com/tidwall/gjson v1.14.2/go.mod h1:/wbyibRr2FHMks5tjHJ5F8dMZh3AcwJEMf5vlfC0lxk=
+github.com/tidwall/gjson v1.14.4 h1:uo0p8EbA09J7RQaflQ1aBRffTR7xedD2bcIVSYxLnkM=
+github.com/tidwall/gjson v1.14.4/go.mod h1:/wbyibRr2FHMks5tjHJ5F8dMZh3AcwJEMf5vlfC0lxk=
+github.com/tidwall/match v1.1.1 h1:+Ho715JplO36QYgwN9PGYNhgZvoUSc9X2c80KVTi+GA=
+github.com/tidwall/match v1.1.1/go.mod h1:eRSPERbgtNPcGhD8UCthc6PmLEQXEWd3PRB5JTxsfmM=
+github.com/tidwall/pretty v1.2.0/go.mod h1:ITEVvHYasfjBbM0u2Pg8T2nJnzm8xPwvNhhsoaGGjNU=
+github.com/tidwall/pretty v1.2.1 h1:qjsOFOWWQl+N3RsoF5/ssm1pHmJJwhjlSbZ51I6wMl4=
+github.com/tidwall/pretty v1.2.1/go.mod h1:ITEVvHYasfjBbM0u2Pg8T2nJnzm8xPwvNhhsoaGGjNU=
+github.com/tidwall/sjson v1.2.5 h1:kLy8mja+1c9jlljvWTlSazM7cKDRfJuR/bOJhcY5NcY=
+github.com/tidwall/sjson v1.2.5/go.mod h1:Fvgq9kS/6ociJEDnK0Fk1cpYF4FIW6ZF7LAe+6jwd28=
 github.com/twitchyliquid64/golang-asm v0.15.1 h1:SU5vSMR7hnwNxj24w34ZyCi/FmDZTkS4MhqMhdFk5YI=
 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=

go/ai-proxy/main.go

@@ -4,28 +4,58 @@ import (
 	"net/http"
 	"os"
 
+	"github.com/gorilla/mux"
 	"k8s.io/klog/v2"
 
+	"github.com/pluralsh/console/go/ai-proxy/api"
+	"github.com/pluralsh/console/go/ai-proxy/api/ollama"
+	"github.com/pluralsh/console/go/ai-proxy/api/openai_standard"
 	"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/proxy"
-	"github.com/pluralsh/console/go/ai-proxy/router"
 )
 
 func main() {
-	klog.V(log.LogLevelMinimal).InfoS("Starting AI Proxy", "provider", args.Provider(), "version", environment.Version, "commit", environment.Commit)
+	provider := args.Provider()
+	host := args.ProviderHost()
+	creds := args.ProviderCredentials()
 
-	p, err := proxy.NewOllamaTranslationProxy(args.Provider(), args.ProviderHost(), args.ProviderCredentials())
-	if err != nil {
-		klog.ErrorS(err, "Could not create proxy")
-		os.Exit(1)
+	var translationProxy api.TranslationProxy
+	if provider != api.ProviderOpenAIStandard {
+		tp, err := proxy.NewTranslationProxy(provider, host, creds)
+		if err != nil {
+			klog.ErrorS(err, "Could not create translation proxy")
+			os.Exit(1)
+		}
+		translationProxy = tp
+	} else {
+		translationProxy = nil
+	}
+
+	var openaiProxy api.OpenAIProxy
+	if provider == api.ProviderOpenAIStandard {
+		op, err := proxy.NewOpenAIProxy(provider, host, creds)
+		if err != nil {
+			klog.ErrorS(err, "Could not create openai proxy")
+			os.Exit(1)
+		}
+		openaiProxy = op
+	} else {
+		openaiProxy = nil
 	}
 
-	handler := router.NewRouter(p)
+	router := mux.NewRouter()
+
+	if translationProxy != nil {
+		router.HandleFunc(ollama.EndpointChat, translationProxy.Proxy())
+	}
+
+	if openaiProxy != nil {
+		router.HandleFunc(openai_standard.EndpointChat, openaiProxy.Proxy())
+	}
 
 	klog.V(log.LogLevelMinimal).InfoS("Listening and serving HTTP", "address", args.Address())
-	if err := http.ListenAndServe(args.Address(), handler); err != nil {
+	if err := http.ListenAndServe(args.Address(), router); err != nil {
 		klog.ErrorS(err, "Could not run the router")
 		os.Exit(1)
 	}

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

@@ -0,0 +1,102 @@
+package openai
+
+import (
+	"bytes"
+	"compress/gzip"
+	"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/api/openai"
+	"github.com/pluralsh/console/go/ai-proxy/internal/log"
+)
+
+const headerContentEncoding = "Content-Encoding"
+
+type OpenAIProxy struct {
+	proxy *httputil.ReverseProxy
+	token string
+}
+
+func (o *OpenAIProxy) Proxy() http.HandlerFunc {
+	return func(w http.ResponseWriter, r *http.Request) {
+		o.proxy.ServeHTTP(w, r)
+	}
+}
+
+func NewOpenAIStandardProxy(host, token string) (api.OpenAIProxy, error) {
+	parsedURL, err := url.Parse(host)
+	if err != nil {
+		return nil, err
+	}
+
+	reverse := &httputil.ReverseProxy{
+		Rewrite: func(r *httputil.ProxyRequest) {
+			r.Out.Header.Set("Authorization", "Bearer "+token)
+
+			r.SetXForwarded()
+
+			targetURL, err := url.Parse(openai.EndpointChat)
+			if err != nil {
+				klog.ErrorS(err, "failed to parse target url")
+				return
+			}
+
+			r.Out.URL.Scheme = parsedURL.Scheme
+			r.Out.URL.Host = parsedURL.Host
+			r.Out.Host = parsedURL.Host
+			r.Out.URL.Path = targetURL.Path
+
+			klog.V(log.LogLevelDebug).InfoS(
+				"proxying request",
+				"from", fmt.Sprintf("%s %s", r.In.Method, r.In.URL.Path),
+				"to", r.Out.URL.String(),
+			)
+		},
+
+		ModifyResponse: func(resp *http.Response) error {
+			contentEncoding := resp.Header.Get(headerContentEncoding)
+			if contentEncoding == "" {
+				return nil
+			}
+
+			var reader io.Reader
+			switch strings.TrimSpace(contentEncoding) {
+			case "br":
+				resp.Header.Del(headerContentEncoding)
+				reader = brotli.NewReader(resp.Body)
+			case "gzip":
+				resp.Header.Del(headerContentEncoding)
+				gzr, err := gzip.NewReader(resp.Body)
+				if err != nil {
+					return err
+				}
+				reader = gzr
+			default:
+				return nil
+			}
+
+			decompressed, err := io.ReadAll(reader)
+			if err != nil {
+				return err
+			}
+
+			resp.Body = io.NopCloser(bytes.NewReader(decompressed))
+			resp.ContentLength = int64(len(decompressed))
+
+			return nil
+		},
+	}
+
+	return &OpenAIProxy{
+		proxy: reverse,
+		token: token,
+	}, nil
+}

go/ai-proxy/proxy/proxy.go

@@ -4,10 +4,11 @@ import (
 	"fmt"
 
 	"github.com/pluralsh/console/go/ai-proxy/api"
+	"github.com/pluralsh/console/go/ai-proxy/proxy/openai"
 	"github.com/pluralsh/console/go/ai-proxy/proxy/provider"
 )
 
-func NewOllamaTranslationProxy(p api.Provider, host string, credentials string) (api.TranslationProxy, error) {
+func NewTranslationProxy(p api.Provider, host string, credentials string) (api.TranslationProxy, error) {
 	switch p {
 	case api.ProviderOllama:
 		return provider.NewOllamaProxy(host)
@@ -21,3 +22,11 @@ func NewOllamaTranslationProxy(p api.Provider, host string, credentials string)
 
 	return nil, fmt.Errorf("invalid provider: %s", p)
 }
+
+func NewOpenAIProxy(p api.Provider, host, token string) (api.OpenAIProxy, error) {
+	switch p {
+	case api.ProviderOpenAIStandard:
+		return openai.NewOpenAIStandardProxy(host, token)
+	}
+	return nil, fmt.Errorf("invalid provider: %s", p)
+}