Add functionality to AI Proxy for full Open AI protocol

go/ai-proxy/Dockerfile

@@ -19,7 +19,6 @@ COPY args ./args
 COPY environment ./environment
 COPY internal ./internal
 COPY proxy ./proxy
-COPY router ./router
 
 # Build
 # the GOARCH has not a default value to allow the binary be built according to the host where the command

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

@@ -13,7 +13,8 @@ import (
 type Endpoint string
 
 const (
-	EndpointChat = "/v1/chat/completions"
+	EndpointChat            = "/openai/chat/completions"
+	EndpointChatCompletions = "/v1/chat/completions"
 )
 
 type ChatCompletionRequest struct {

go/ai-proxy/api/provider.go

@@ -42,7 +42,7 @@ type ProviderAPIMapping map[string]string
 
 var (
 	ollamaToOpenAI ProviderAPIMapping = map[string]string{
-		ollama.EndpointChat: openai.EndpointChat,
+		ollama.EndpointChat: openai.EndpointChatCompletions,
 	}
 	ollamaToVertex ProviderAPIMapping = map[string]string{
 		ollama.EndpointChat: vertex.EndpointChat,

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/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,13 +4,16 @@ 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"
 	"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() {
@@ -22,10 +25,18 @@ func main() {
 		os.Exit(1)
 	}
 
-	handler := router.NewRouter(p)
+	op, err := proxy.NewOpenAIProxy(api.ProviderOpenAI, args.ProviderHost(), args.ProviderCredentials())
+	if err != nil {
+		klog.ErrorS(err, "Could not create proxy")
+		os.Exit(1)
+	}
+
+	router := mux.NewRouter()
+	router.HandleFunc(ollama.EndpointChat, p.Proxy())
+	router.HandleFunc(openai.EndpointChat, op.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,62 @@
+package openai
+
+import (
+	"fmt"
+	"net/http"
+	"net/http/httputil"
+	"net/url"
+
+	"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"
+)
+
+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 NewOpenAIProxy(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.EndpointChatCompletions)
+			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(),
+			)
+		},
+	}
+
+	return &OpenAIProxy{
+		proxy: reverse,
+		token: token,
+	}, nil
+}

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

@@ -46,7 +46,7 @@ func (in *OpenAIProxy) ModifyResponse(r *http.Response) error {
 func (in *OpenAIProxy) modifyRequestBody(r *httputil.ProxyRequest) error {
 	endpoint := r.Out.URL.Path
 	switch endpoint {
-	case openai.EndpointChat:
+	case openai.EndpointChatCompletions:
 		return replaceRequestBody(r, openai.ToChatCompletionRequest)
 	}
 
@@ -60,7 +60,7 @@ func (in *OpenAIProxy) modifyResponseBody(r *http.Response) error {
 
 	endpoint := r.Request.URL.Path
 	switch endpoint {
-	case openai.EndpointChat:
+	case openai.EndpointChatCompletions:
 		return replaceResponseBody(r, openai.FromChatCompletionResponse)
 	}
 

go/ai-proxy/proxy/proxy.go

@@ -4,6 +4,7 @@ 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"
 )
 
@@ -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.ProviderOpenAI:
+		return openai.NewOpenAIProxy(host, token)
+	}
+	return nil, fmt.Errorf("invalid provider: %s", p)
+}

go/ai-proxy/test/helpers/common.go

@@ -10,9 +10,13 @@ import (
 	"net/http/httptest"
 	"strings"
 
+	"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/api/openai"
 	"github.com/pluralsh/console/go/ai-proxy/args"
 	"github.com/pluralsh/console/go/ai-proxy/proxy"
-	"github.com/pluralsh/console/go/ai-proxy/router"
 )
 
 func SetupServer() (*httptest.Server, error) {
@@ -21,9 +25,16 @@ func SetupServer() (*httptest.Server, error) {
 		return nil, err
 	}
 
-	return httptest.NewServer(router.NewRouter(p)), nil
-}
+	op, err := proxy.NewOpenAIProxy(api.ProviderOpenAI, args.ProviderHost(), args.ProviderCredentials())
+	if err != nil {
+		return nil, err
+	}
+	router := mux.NewRouter()
+	router.HandleFunc(ollama.EndpointChat, p.Proxy())
+	router.HandleFunc(openai.EndpointChat, op.Proxy())
 
+	return httptest.NewServer(router), nil
+}
 func SetupProviderServer(handlers map[string]http.HandlerFunc) (*httptest.Server, error) {
 	server := httptest.NewUnstartedServer(http.HandlerFunc(func(response http.ResponseWriter, request *http.Request) {
 		if handler, exists := handlers[request.URL.Path]; exists {
@@ -69,6 +80,35 @@ func CreateRequest[T any](method string, endpoint string, body T) func(requestSe
 	}
 }
 
+func CreateRequestWithResponse[T any](method string, endpoint string, body T) func(requestServer *httptest.Server, externalServer *httptest.Server) ([]byte, *http.Response, error) {
+	return func(requestServer *httptest.Server, externalServer *httptest.Server) ([]byte, *http.Response, error) {
+		bodyBytes, err := json.Marshal(body)
+		if err != nil {
+			return nil, nil, err
+		}
+
+		url := fmt.Sprintf("%s/%s", externalServer.URL, strings.TrimLeft(endpoint, "/"))
+		req, err := http.NewRequest(method, url, bytes.NewReader(bodyBytes))
+		if err != nil {
+			return nil, nil, err
+		}
+		req.Header.Set("Content-Type", "application/json")
+
+		res, err := requestServer.Client().Do(req)
+		if err != nil {
+			return nil, nil, err
+		}
+		defer res.Body.Close()
+
+		responseBytes, err := io.ReadAll(res.Body)
+		if err != nil {
+			return nil, nil, err
+		}
+
+		return responseBytes, res, nil
+	}
+}
+
 func MockResponse(endpoint string, response []byte, err error, status int) func(handlers map[string]http.HandlerFunc) error {
 	return func(handlers map[string]http.HandlerFunc) error {
 		if _, exists := handlers[endpoint]; exists {