more

Signed-off-by: Takeshi Yoneda <[email protected]>

Author: mathetake

internal/dynamic_module/router_filter.go

@@ -31,8 +31,10 @@ type (
 	routerFilter struct {
 		fc                     *routerFilterConfig
 		endpoint               endpoint
+		originalHeaders        map[string]string
 		originalRequestBody    interface{}
 		originalRequestBodyRaw []byte
+		attemptCount           int
 	}
 )
 
@@ -118,17 +120,17 @@ func (f *routerFilter) RequestBody(e sdk.EnvoyHTTPFilter, endOfStream bool) sdk.
 	var modelName string
 	switch f.endpoint {
 	case chatCompletionsEndpoint:
-		parsed, modelName, err = modelBodyParser(raw, func(req *openai.ChatCompletionRequest) string { return req.Model })
+		parsed, modelName, err = parseBodyWithModel(raw, func(req *openai.ChatCompletionRequest) string { return req.Model })
 	case completionsEndpoint:
-		parsed, modelName, err = modelBodyParser(raw, func(req *openai.CompletionRequest) string { return req.Model })
+		parsed, modelName, err = parseBodyWithModel(raw, func(req *openai.CompletionRequest) string { return req.Model })
 	case embeddingsEndpoint:
-		parsed, modelName, err = modelBodyParser(raw, func(req *openai.EmbeddingRequest) string { return req.Model })
+		parsed, modelName, err = parseBodyWithModel(raw, func(req *openai.EmbeddingRequest) string { return req.Model })
 	case imagesGenerationsEndpoint:
-		parsed, modelName, err = modelBodyParser(raw, func(req *openaisdk.ImageGenerateParams) string { return req.Model })
+		parsed, modelName, err = parseBodyWithModel(raw, func(req *openaisdk.ImageGenerateParams) string { return req.Model })
 	case rerankEndpoint:
-		parsed, modelName, err = modelBodyParser(raw, func(req *cohereschema.RerankV2Request) string { return req.Model })
+		parsed, modelName, err = parseBodyWithModel(raw, func(req *cohereschema.RerankV2Request) string { return req.Model })
 	case messagesEndpoint:
-		parsed, modelName, err = modelBodyParser(raw, func(req *anthropic.MessagesRequest) string { return req.GetModel() })
+		parsed, modelName, err = parseBodyWithModel(raw, func(req *anthropic.MessagesRequest) string { return req.GetModel() })
 	default:
 		e.SendLocalReply(500, nil, []byte("BUG: unsupported endpoint at body parsing: "+fmt.Sprintf("%d", f.endpoint)))
 	}
@@ -144,6 +146,8 @@ func (f *routerFilter) RequestBody(e sdk.EnvoyHTTPFilter, endOfStream bool) sdk.
 	// Store the pointer to the filter in dynamic metadata for later retrieval in the upstream filter.
 	e.SetDynamicMetadataString(internalapi.AIGatewayFilterMetadataNamespace, routerFilterPointerDynamicMetadataKey,
 		fmt.Sprintf("%d", uintptr(unsafe.Pointer(f))))
+
+	f.originalHeaders = multiValueHeadersToSingleValue(e.GetRequestHeaders())
 	return sdk.RequestBodyStatusContinue
 }
 
@@ -168,10 +172,21 @@ func (f *routerFilter) handleModelsEndpoint(e sdk.EnvoyHTTPFilter) sdk.RequestHe
 	return sdk.RequestHeadersStatusStopIteration
 }
 
-func modelBodyParser[T any](body []byte, modelExtractFn func(req *T) string) (interface{}, string, error) {
+func parseBodyWithModel[T any](body []byte, modelExtractFn func(req *T) string) (interface{}, string, error) {
 	var req T
 	if err := json.Unmarshal(body, &req); err != nil {
 		return nil, "", fmt.Errorf("failed to unmarshal body: %w", err)
 	}
 	return req, modelExtractFn(&req), nil
 }
+
+// multiValueHeadersToSingleValue converts a map of headers with multiple values to a map of headers with single values by taking the first value for each header.
+//
+// TODO: this is purely for feature parity with the old filter where we ignore the case of multiple header values.
+func multiValueHeadersToSingleValue(headers map[string][]string) map[string]string {
+	singleValueHeaders := make(map[string]string, len(headers))
+	for k, v := range headers {
+		singleValueHeaders[k] = v[0]
+	}
+	return singleValueHeaders
+}

internal/dynamic_module/upstream_filter.go

@@ -1,11 +1,14 @@
 package main
 
 import (
+	"context"
 	"fmt"
 	"strconv"
 	"unsafe"
 
 	"github.com/envoyproxy/ai-gateway/internal/dynamic_module/sdk"
+	"github.com/envoyproxy/ai-gateway/internal/filterapi"
+	"github.com/envoyproxy/ai-gateway/internal/headermutator"
 	"github.com/envoyproxy/ai-gateway/internal/internalapi"
 )
 
@@ -14,7 +17,10 @@ type (
 	upstreamFilterConfig struct{}
 	// upstreamFilter implements [sdk.HTTPFilter].
 	upstreamFilter struct {
-		rf *routerFilter
+		rf         *routerFilter
+		backend    *filterapi.RuntimeBackend
+		reqHeaders map[string]string
+		onRetry    bool
 	}
 )
 
@@ -38,6 +44,8 @@ func (f *upstreamFilter) RequestHeaders(e sdk.EnvoyHTTPFilter, _ bool) sdk.Reque
 		return sdk.RequestHeadersStatusStopIteration
 	}
 	f.rf = (*routerFilter)(unsafe.Pointer(uintptr(rfPtr)))
+	f.rf.attemptCount++
+	f.onRetry = f.rf.attemptCount > 1
 
 	backend, ok := e.GetUpstreamHostMetadataString(internalapi.AIGatewayFilterMetadataNamespace, internalapi.InternalMetadataBackendNameKey)
 	if !ok {
@@ -49,11 +57,61 @@ func (f *upstreamFilter) RequestHeaders(e sdk.EnvoyHTTPFilter, _ bool) sdk.Reque
 		e.SendLocalReply(500, nil, []byte(fmt.Sprintf("backend %s not found in filter config", backend)))
 		return sdk.RequestHeadersStatusStopIteration
 	}
+
+	f.backend = b
+	f.reqHeaders = multiValueHeadersToSingleValue(e.GetRequestHeaders())
+
+	// Now mutate the headers based on the backend configuration.
+	if hm := b.Backend.HeaderMutation; hm != nil {
+		sets, removes := headermutator.NewHeaderMutator(b.Backend.HeaderMutation, f.rf.originalHeaders).Mutate(f.reqHeaders, f.onRetry)
+		for _, h := range sets {
+			if !e.SetRequestHeader(h.Key(), []byte(h.Value())) {
+				e.SendLocalReply(500, nil, []byte(fmt.Sprintf("failed to set header %s", h.Key())))
+				return sdk.RequestHeadersStatusStopIteration
+			}
+			f.reqHeaders[h.Key()] = h.Value()
+		}
+		for _, key := range removes {
+			if !e.SetRequestHeader(key, nil) {
+				e.SendLocalReply(500, nil, []byte(fmt.Sprintf("failed to remove header %s", key)))
+				return sdk.RequestHeadersStatusStopIteration
+			}
+			delete(f.reqHeaders, key)
+		}
+	}
 	return sdk.RequestHeadersStatusContinue
 }
 
 // RequestBody implements [sdk.HTTPFilter].
-func (f *upstreamFilter) RequestBody(sdk.EnvoyHTTPFilter, bool) sdk.RequestBodyStatus {
+func (f *upstreamFilter) RequestBody(e sdk.EnvoyHTTPFilter, endOfStream bool) sdk.RequestBodyStatus {
+	if !endOfStream {
+		// TODO: ideally, we should not buffer the entire body for the passthrough case.
+		return sdk.RequestBodyStatusStopIterationAndBuffer
+	}
+
+	b := f.backend
+
+	// TODO: endpoint specific logic such as translation.
+
+	if bm := b.Backend.BodyMutation; bm != nil {
+		// TODO: body mutation if needed.
+		_ = bm
+	}
+
+	// Next is to do the upstream auth if needed.
+	if b.Handler != nil {
+		authHeaders, err := b.Handler.Do(context.Background(), f.reqHeaders, f.rf.originalRequestBodyRaw)
+		if err != nil {
+			e.SendLocalReply(500, nil, []byte(fmt.Sprintf("failed to do backend auth: %v", err)))
+			return sdk.RequestBodyStatusStopIterationAndBuffer
+		}
+		for _, h := range authHeaders {
+			if !e.SetRequestHeader(h.Key(), []byte(h.Value())) {
+				e.SendLocalReply(500, nil, []byte(fmt.Sprintf("failed to set auth header %s", h.Key())))
+				return sdk.RequestBodyStatusStopIterationAndBuffer
+			}
+		}
+	}
 	return sdk.RequestBodyStatusContinue
 }