Support openai embeddings endpoints in ai proxy

go/ai-proxy/api/bedrock/bedrock.go

@@ -1,5 +1 @@
 package bedrock
-
-type Endpoint string
-
-const EndpointChat = "/chat/bedrock"

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

@@ -15,6 +15,7 @@ type Endpoint string
 const (
 	EndpointChat            = "/openai/v1/chat/completions"
 	EndpointChatCompletions = "/v1/chat/completions"
+	EndpointEmbeddings      = "/openai/v1/embeddings"
 )
 
 type ChatCompletionRequest struct {

go/ai-proxy/api/provider.go

@@ -24,6 +24,8 @@ func ToProvider(s string) (Provider, error) {
 		return ProviderAnthropic, nil
 	case ProviderVertex.String():
 		return ProviderVertex, nil
+	case ProviderBedrock.String():
+		return ProviderBedrock, nil
 	}
 
 	return "", fmt.Errorf("invalid provider: %s", s)

go/ai-proxy/args/args.go

@@ -65,6 +65,10 @@ func Provider() api.Provider {
 }
 
 func ProviderHost() string {
+	if Provider() == api.ProviderBedrock {
+		return ""
+	}
+
 	if len(*argProviderHost) == 0 {
 		panic(fmt.Errorf("provider host is required"))
 	}
@@ -108,3 +112,7 @@ func Address() string {
 
 	return fmt.Sprintf("%s:%d", *argAddress, *argPort)
 }
+
+func OpenAICompatible() bool {
+	return Provider() == api.ProviderOpenAI || Provider() == api.ProviderBedrock
+}

go/ai-proxy/go.mod

@@ -53,6 +53,7 @@ require (
 	github.com/modern-go/reflect2 v1.0.2 // 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.37.0 // indirect
 	github.com/twitchyliquid64/golang-asm v0.15.1 // indirect
 	github.com/ugorji/go/codec v1.2.12 // indirect
 	golang.org/x/arch v0.11.0 // indirect

go/ai-proxy/go.sum

@@ -109,6 +109,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.37.0 h1:hQQowgYm4OXJ1Z/wTrE+XZaO20BYsL0R3uRPSpfNZkY=
+github.com/sashabaranov/go-openai v1.37.0/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=

go/ai-proxy/main.go

@@ -8,7 +8,6 @@ import (
 	"k8s.io/klog/v2"
 
 	"github.com/pluralsh/console/go/ai-proxy/api"
-	"github.com/pluralsh/console/go/ai-proxy/api/bedrock"
 	"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"
@@ -20,29 +19,30 @@ import (
 func main() {
 	klog.V(log.LogLevelMinimal).InfoS("Starting AI Proxy", "provider", args.Provider(), "version", environment.Version, "commit", environment.Commit)
 
+	router := mux.NewRouter()
 	p, err := proxy.NewOllamaTranslationProxy(args.Provider(), args.ProviderHost(), args.ProviderCredentials())
-	if err != nil {
+	if err == nil {
+		router.HandleFunc(ollama.EndpointChat, p.Proxy())
+	} else if args.Provider() != api.ProviderBedrock {
 		klog.ErrorS(err, "Could not create proxy")
 		os.Exit(1)
 	}
 
-	op, err := proxy.NewOpenAIProxy(api.ProviderOpenAI, args.ProviderHost(), args.ProviderCredentials())
-	if err != nil {
-		klog.ErrorS(err, "Could not create proxy")
-		os.Exit(1)
+	if args.OpenAICompatible() {
+		op, err := proxy.NewOpenAIProxy(args.Provider(), args.ProviderHost(), args.ProviderCredentials())
+		if err != nil {
+			klog.ErrorS(err, "Could not create proxy")
+			os.Exit(1)
+		}
+		ep, err := proxy.NewOpenAIEmbeddingsProxy(args.Provider(), args.ProviderHost(), args.ProviderCredentials())
+		if err != nil {
+			klog.ErrorS(err, "Could not create embedding proxy")
+			os.Exit(1)
+		}
+		router.HandleFunc(openai.EndpointChat, op.Proxy())
+		router.HandleFunc(openai.EndpointEmbeddings, ep.Proxy())
 	}
 
-	bp, err := proxy.NewBedrockProxy(api.ProviderBedrock, 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())
-	router.HandleFunc(bedrock.EndpointChat, bp.Proxy())
-
 	klog.V(log.LogLevelMinimal).InfoS("Listening and serving HTTP", "address", args.Address())
 	if err := http.ListenAndServe(args.Address(), router); err != nil {
 		klog.ErrorS(err, "Could not run the router")

go/ai-proxy/proxy/bedrock/embeddings.go

@@ -0,0 +1,137 @@
+package bedrock
+
+import (
+	"context"
+	"encoding/json"
+	"fmt"
+	"net/http"
+	"strings"
+
+	"github.com/aws/aws-sdk-go-v2/aws"
+	"github.com/aws/aws-sdk-go-v2/config"
+	"github.com/aws/aws-sdk-go-v2/service/bedrockruntime"
+	"github.com/ollama/ollama/openai"
+	"k8s.io/klog/v2"
+
+	"github.com/pluralsh/console/go/ai-proxy/api"
+)
+
+const (
+	Titan  = "titan"
+	Cohere = "cohere"
+)
+
+type titanEmbeddingResponse struct {
+	Embedding []float32 `json:"embedding"`
+}
+
+type cohereEmbeddingResponse struct {
+	Embedding []float32 `json:"embeddings"`
+}
+
+type BedrockEmbeddingsProxy struct {
+	bedrockClient *bedrockruntime.Client
+}
+
+func NewBedrockEmbeddingsProxy(region string) (api.OpenAIProxy, error) {
+	ctx := context.Background()
+
+	var loadOptions []func(options *config.LoadOptions) error
+	if region != "" {
+		loadOptions = append(loadOptions, config.WithRegion(region))
+	}
+
+	sdkConfig, err := config.LoadDefaultConfig(ctx, loadOptions...)
+	if err != nil {
+		klog.ErrorS(err, "Couldn't load default configuration. Have you set up your AWS account?")
+		return nil, err
+	}
+	bedrockClient := bedrockruntime.NewFromConfig(sdkConfig)
+	return &BedrockEmbeddingsProxy{
+		bedrockClient: bedrockClient,
+	}, nil
+}
+
+func (b *BedrockEmbeddingsProxy) Proxy() http.HandlerFunc {
+	return func(w http.ResponseWriter, r *http.Request) {
+		var openAIReq openai.EmbedRequest
+		if err := json.NewDecoder(r.Body).Decode(&openAIReq); err != nil {
+			http.Error(w, "failed to parse openai request", http.StatusBadRequest)
+			return
+		}
+		b.handleEmbeddingBedrock(w, &openAIReq)
+	}
+}
+
+func (b *BedrockEmbeddingsProxy) handleEmbeddingBedrock(
+	w http.ResponseWriter,
+	req *openai.EmbedRequest,
+) {
+	input := map[string]interface{}{}
+
+	switch {
+	case strings.Contains(strings.ToLower(req.Model), Titan):
+		input["inputText"] = req.Input
+	case strings.Contains(strings.ToLower(req.Model), Cohere):
+		input["texts"] = []string{req.Input.(string)}
+	default:
+		klog.Errorf("model doesn't support embedding at this time %s", req.Model)
+		return
+	}
+
+	payloadBytes, err := json.Marshal(input)
+	if err != nil {
+		klog.ErrorS(err, "failed to convert to bedrock request")
+		return
+	}
+
+	output, err := b.bedrockClient.InvokeModel(context.Background(), &bedrockruntime.InvokeModelInput{
+		ModelId: aws.String(req.Model),
+		Body:    payloadBytes,
+	})
+	if err != nil {
+		klog.ErrorS(err, "request to bedrock failed")
+		return
+	}
+
+	response, err := convertBedrockEmbeddingToOpenAI(output, req.Model)
+	if err != nil {
+		klog.ErrorS(err, "failed to convert bedrock response to openai format")
+		return
+	}
+
+	w.Header().Set("Content-Type", "application/json")
+	if err := json.NewEncoder(w).Encode(response); err != nil {
+		klog.Errorf("Error encoding response: %v", err)
+		return
+	}
+}
+
+func convertBedrockEmbeddingToOpenAI(output *bedrockruntime.InvokeModelOutput, model string) (*openai.EmbeddingList, error) {
+	switch {
+	case strings.Contains(model, Titan):
+		var embed titanEmbeddingResponse
+		if err := json.Unmarshal(output.Body, &embed); err != nil {
+			return nil, fmt.Errorf("failed to unmarshal Titan embedding response: %v", err)
+		}
+		var embedding openai.Embedding
+		embedding.Embedding = embed.Embedding
+		return &openai.EmbeddingList{
+			Model: model,
+			Data:  []openai.Embedding{embedding},
+		}, nil
+	case strings.Contains(model, Cohere):
+		var embed cohereEmbeddingResponse
+		if err := json.Unmarshal(output.Body, &embed); err != nil {
+			return nil, fmt.Errorf("failed to unmarshal Cohere embedding response: %v", err)
+		}
+		var embedding openai.Embedding
+		embedding.Embedding = embed.Embedding
+		return &openai.EmbeddingList{
+			Model: model,
+			Data:  []openai.Embedding{embedding},
+		}, nil
+	default:
+		return nil, fmt.Errorf("model doesn't support embedding at this time %s", model)
+	}
+}

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

@@ -0,0 +1,50 @@
+package openai
+
+import (
+	"fmt"
+	"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"
+)
+
+func NewOpenAIEmbeddingsProxy(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.EndpointEmbeddings)
+			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/proxy.go

@@ -24,18 +24,22 @@ 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) {
+func NewOpenAIProxy(p api.Provider, host, credentials string) (api.OpenAIProxy, error) {
 	switch p {
 	case api.ProviderOpenAI:
-		return openai.NewOpenAIProxy(host, token)
+		return openai.NewOpenAIProxy(host, credentials)
+	case api.ProviderBedrock:
+		return bedrock.NewBedrockProxy(credentials)
 	}
 	return nil, fmt.Errorf("invalid provider: %s", p)
 }
 
-func NewBedrockProxy(p api.Provider, region string) (api.OpenAIProxy, error) {
+func NewOpenAIEmbeddingsProxy(p api.Provider, host, credentials string) (api.OpenAIProxy, error) {
 	switch p {
+	case api.ProviderOpenAI:
+		return openai.NewOpenAIEmbeddingsProxy(host, credentials)
 	case api.ProviderBedrock:
-		return bedrock.NewBedrockProxy(region)
+		return bedrock.NewBedrockEmbeddingsProxy(credentials)
 	}
 	return nil, fmt.Errorf("invalid provider: %s", p)
 }