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

@@ -5,6 +5,7 @@ import (
 
 	"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/api/openai_standard"
 	"github.com/pluralsh/console/go/ai-proxy/api/vertex"
 )
 
@@ -20,6 +21,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 +33,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
@@ -47,6 +51,9 @@ var (
 	ollamaToVertex ProviderAPIMapping = map[string]string{
 		ollama.EndpointChat: vertex.EndpointChat,
 	}
+	openAIToOpenAI ProviderAPIMapping = map[string]string{
+		openai_standard.EndpointChat: openai.EndpointChat,
+	}
 )
 
 func ToProviderAPIPath(target Provider, path string) string {
@@ -59,6 +66,13 @@ func ToProviderAPIPath(target Provider, path string) string {
 			panic(fmt.Sprintf("path %s not registered for provider %s", path, target))
 		}
 
+		return targetPath
+	case ProviderOpenAIStandard:
+		targetPath, exists := openAIToOpenAI[path]
+		if !exists {
+			panic(fmt.Sprintf("path %s not registered for provider %s", path, target))
+		}
+
 		return targetPath
 	case ProviderVertex:
 		targetPath, exists := ollamaToVertex[path]

go/ai-proxy/args/args.go

@@ -79,7 +79,7 @@ func ProviderCredentials() string {
 		return *argProviderServiceAccount
 	}
 
-	if Provider() == defaultProvider {
+	if Provider() == defaultProvider || Provider() == api.ProviderOpenAIStandard {
 		return ""
 	}
 

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

@@ -16,7 +16,7 @@ import (
 func main() {
 	klog.V(log.LogLevelMinimal).InfoS("Starting AI Proxy", "provider", args.Provider(), "version", environment.Version, "commit", environment.Commit)
 
-	p, err := proxy.NewOllamaTranslationProxy(args.Provider(), args.ProviderHost(), args.ProviderCredentials())
+	p, err := proxy.NewTranslationProxy(args.Provider(), args.ProviderHost(), args.ProviderCredentials())
 	if err != nil {
 		klog.ErrorS(err, "Could not create proxy")
 		os.Exit(1)

go/ai-proxy/proxy/provider/openai_standard.go

@@ -0,0 +1,20 @@
+package provider
+
+import (
+	"github.com/pluralsh/console/go/ai-proxy/api"
+)
+
+type OpenAIStandardProxy struct {
+	api.TranslationProxy
+}
+
+func NewOpenAIStandardProxy(host string) (api.TranslationProxy, error) {
+	base, err := newBaseTranslationProxy(host, api.ProviderOpenAIStandard, nil, nil, nil)
+	if err != nil {
+		return nil, err
+	}
+
+	return &OpenAIStandardProxy{
+		TranslationProxy: base,
+	}, nil
+}

go/ai-proxy/proxy/proxy.go

@@ -7,12 +7,14 @@ import (
 	"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)
 	case api.ProviderOpenAI:
 		return provider.NewOpenAIProxy(host, credentials)
+	case api.ProviderOpenAIStandard:
+		return provider.NewOpenAIStandardProxy(host)
 	case api.ProviderVertex:
 		return provider.NewVertexProxy(host, credentials)
 	case api.ProviderAnthropic:

go/ai-proxy/router/router.go

@@ -7,6 +7,7 @@ import (
 
 	"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"
 )
 
 func NewRouter(proxy api.TranslationProxy) http.Handler {
@@ -15,5 +16,8 @@ func NewRouter(proxy api.TranslationProxy) http.Handler {
 	// Register all Ollama API routes that should be proxied.
 	router.HandleFunc(ollama.EndpointChat, proxy.Proxy())
 
+	// Register OpenAI API routes  (ollama/openai --> openai
+	router.HandleFunc(openai_standard.EndpointChat, proxy.Proxy())
+
 	return router
 }

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

@@ -16,7 +16,7 @@ import (
 )
 
 func SetupServer() (*httptest.Server, error) {
-	p, err := proxy.NewOllamaTranslationProxy(args.Provider(), args.ProviderHost(), args.ProviderCredentials())
+	p, err := proxy.NewTranslationProxy(args.Provider(), args.ProviderHost(), args.ProviderCredentials())
 	if err != nil {
 		return nil, err
 	}
@@ -69,6 +69,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 {

go/ai-proxy/test/openai_standard/openai_standard_test.go

@@ -0,0 +1,149 @@
+package openai_standard
+
+import (
+	"bytes"
+	"encoding/json"
+	"errors"
+	"net/http"
+	"net/http/httptest"
+	"os"
+	"testing"
+
+	"github.com/ollama/ollama/openai"
+	"github.com/spf13/pflag"
+	"k8s.io/klog/v2"
+
+	"github.com/pluralsh/console/go/ai-proxy/test/helpers"
+)
+
+func init() {
+	_ = pflag.Set("provider", "ollama")
+	_ = pflag.Set("provider-host", "localhost:8081")
+	_ = pflag.Set("provider-token", "test")
+}
+
+var (
+	server         *httptest.Server
+	providerServer *httptest.Server
+	// TODO: for parallel runs concurrent map should be used
+	handlers = make(map[string]http.HandlerFunc)
+)
+
+const endpoint = "/openai/chat/completions"
+
+func TestMain(m *testing.M) {
+	var err error
+	server, err = helpers.SetupServer()
+	if err != nil {
+		klog.Fatal(err)
+	}
+
+	providerServer, err = helpers.SetupProviderServer(handlers)
+	if err != nil {
+		klog.Fatal(err)
+	}
+
+	os.Exit(m.Run())
+}
+
+func TestOpenAIStandardProxy(t *testing.T) {
+	cases := []helpers.TestStruct[any, any]{
+		{
+			Name:     "chat request should return correct openai response",
+			Method:   "POST",
+			Endpoint: endpoint,
+			Request: openai.ChatCompletionRequest{
+				Model: "testmodel",
+				Messages: []openai.Message{{
+					Role:    "user",
+					Content: "test prompt",
+				}},
+			},
+			WantData: openai.ChatCompletion{
+				Model: "testmodel",
+				Choices: []openai.Choice{{
+					Message: openai.Message{Role: "assistant", Content: "test response"},
+				}},
+			},
+			WantErr:    nil,
+			WantStatus: http.StatusOK,
+		},
+	}
+
+	for _, tc := range cases {
+		t.Run(tc.Name, func(t *testing.T) {
+			wantDataBytes, err := json.Marshal(tc.WantData)
+			if err != nil {
+				t.Fatal(err)
+			}
+
+			mockResponseFunc := helpers.MockResponse(tc.Endpoint, wantDataBytes, tc.WantErr, tc.WantStatus)
+			err = mockResponseFunc(handlers)
+			if err != nil {
+				t.Fatal(err)
+			}
+
+			requestFunc := helpers.CreateRequest(tc.Method, tc.Endpoint, tc.Request)
+			res, err := requestFunc(server, providerServer)
+			if !errors.Is(err, tc.WantErr) {
+				t.Fatalf("\nwant:\n%v\ngot:\n%v", tc.WantErr, err)
+			}
+
+			if !bytes.Equal(wantDataBytes, res) {
+				t.Errorf("\nwant:\n%s\ngot:\n%s", tc.WantData, res)
+			}
+		})
+	}
+}
+
+func TestOpenAIStandardProxy_Streaming(t *testing.T) {
+	wantSSE := "data: test streaming response\n\n"
+
+	streamTest := helpers.TestStruct[openai.ChatCompletionRequest, any]{
+		Name:     "chat request with streaming should return SSE headers",
+		Method:   "POST",
+		Endpoint: endpoint,
+		Request: openai.ChatCompletionRequest{
+			Model:  "testmodel",
+			Stream: true,
+			Messages: []openai.Message{
+				{Role: "user", Content: "test streaming prompt"},
+			},
+		},
+		WantStatus: http.StatusOK,
+	}
+
+	t.Run(streamTest.Name, func(t *testing.T) {
+		mockResponse := func(h map[string]http.HandlerFunc) error {
+			h[streamTest.Endpoint] = func(w http.ResponseWriter, r *http.Request) {
+				w.Header().Set("Content-Type", "text/event-stream; charset=utf-8")
+				w.WriteHeader(http.StatusOK)
+				_, _ = w.Write([]byte(wantSSE))
+			}
+			return nil
+		}
+
+		if err := mockResponse(handlers); err != nil {
+			t.Fatal(err)
+		}
+
+		requestFunc := helpers.CreateRequestWithResponse(streamTest.Method, streamTest.Endpoint, streamTest.Request)
+		resBody, resp, err := requestFunc(server, providerServer)
+		if err != nil {
+			t.Fatalf("unexpected error: %v", err)
+		}
+
+		if resp.StatusCode != streamTest.WantStatus {
+			t.Fatalf("want status: %d, got: %d", streamTest.WantStatus, resp.StatusCode)
+		}
+
+		contentType := resp.Header.Get("Content-Type")
+		if contentType != "text/event-stream; charset=utf-8" {
+			t.Errorf("expected Content-Type = text/event-stream; charset=utf-8, got: %s", contentType)
+		}
+
+		if string(resBody) != wantSSE {
+			t.Errorf("expected SSE body %q, got %q", wantSSE, string(resBody))
+		}
+	})
+}