more

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

Author: mathetake

internal/dynamic_module/router_filter.go

@@ -11,6 +11,7 @@ import (
 	cohereschema "github.com/envoyproxy/ai-gateway/internal/apischema/cohere"
 	"github.com/envoyproxy/ai-gateway/internal/apischema/openai"
 	"github.com/envoyproxy/ai-gateway/internal/dynamic_module/sdk"
+	"github.com/envoyproxy/ai-gateway/internal/filterapi"
 	"github.com/envoyproxy/ai-gateway/internal/internalapi"
 	openaisdk "github.com/openai/openai-go/v2"
 )
@@ -33,8 +34,6 @@ type (
 		originalRequestBody    interface{}
 		originalRequestBodyRaw []byte
 	}
-
-	requestBodyParserFn func(body []byte) (parsed interface{}, modelName string, err error)
 )
 
 func newRouterFilterConfig(fcr *filterConfigReceiverImpl) *routerFilterConfig {
@@ -59,6 +58,11 @@ func (f *routerFilterConfig) NewFilter() sdk.HTTPFilter {
 	return &routerFilter{fc: f}
 }
 
+// filterConfig is a helper that returns the filter config for the current stream.
+func (f *routerFilter) filterConfig() *filterapi.RuntimeConfig {
+	return f.fc.fcr.fc
+}
+
 // RequestHeaders implements [sdk.HTTPFilter].
 func (f *routerFilter) RequestHeaders(e sdk.EnvoyHTTPFilter, _ bool) sdk.RequestHeadersStatus {
 	p, _ := e.GetRequestHeader(":path") // The :path pseudo header is always present.
@@ -110,24 +114,24 @@ func (f *routerFilter) RequestBody(e sdk.EnvoyHTTPFilter, endOfStream bool) sdk.
 		return sdk.RequestBodyStatusStopIterationAndBuffer
 	}
 	f.originalRequestBodyRaw = raw
-	var parserFn requestBodyParserFn
+	var parsed any
+	var modelName string
 	switch f.endpoint {
 	case chatCompletionsEndpoint:
-		parserFn = chatCompletionsBodyParser
+		parsed, modelName, err = modelBodyParser(raw, func(req *openai.ChatCompletionRequest) string { return req.Model })
 	case completionsEndpoint:
-		parserFn = completionsBodyParser
+		parsed, modelName, err = modelBodyParser(raw, func(req *openai.CompletionRequest) string { return req.Model })
 	case embeddingsEndpoint:
-		parserFn = embeddingsBodyParser
+		parsed, modelName, err = modelBodyParser(raw, func(req *openai.EmbeddingRequest) string { return req.Model })
 	case imagesGenerationsEndpoint:
-		parserFn = imagesGenerationsBodyParser
+		parsed, modelName, err = modelBodyParser(raw, func(req *openaisdk.ImageGenerateParams) string { return req.Model })
 	case rerankEndpoint:
-		parserFn = rerankBodyParser
+		parsed, modelName, err = modelBodyParser(raw, func(req *cohereschema.RerankV2Request) string { return req.Model })
 	case messagesEndpoint:
-		parserFn = messagesBodyParser
+		parsed, modelName, err = modelBodyParser(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)))
 	}
-	parsed, modelName, err := parserFn(raw)
 	if err != nil {
 		e.SendLocalReply(400, nil, []byte("failed to parse request body: "+err.Error()))
 		return sdk.RequestBodyStatusStopIterationAndBuffer
@@ -164,50 +168,10 @@ func (f *routerFilter) handleModelsEndpoint(e sdk.EnvoyHTTPFilter) sdk.RequestHe
 	return sdk.RequestHeadersStatusStopIteration
 }
 
-func chatCompletionsBodyParser(body []byte) (interface{}, string, error) {
-	var req openai.ChatCompletionRequest
-	if err := json.Unmarshal(body, &req); err != nil {
-		return nil, "", fmt.Errorf("failed to unmarshal body: %w", err)
-	}
-	return req, req.Model, nil
-}
-
-func completionsBodyParser(body []byte) (interface{}, string, error) {
-	var req openai.CompletionRequest
+func modelBodyParser[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, req.Model, nil
-}
-
-func embeddingsBodyParser(body []byte) (interface{}, string, error) {
-	var req openai.EmbeddingRequest
-	if err := json.Unmarshal(body, &req); err != nil {
-		return nil, "", fmt.Errorf("failed to unmarshal body: %w", err)
-	}
-	return req, req.Model, nil
-}
-
-func imagesGenerationsBodyParser(body []byte) (interface{}, string, error) {
-	var req openaisdk.ImageGenerateParams
-	if err := json.Unmarshal(body, &req); err != nil {
-		return nil, "", fmt.Errorf("failed to unmarshal body: %w", err)
-	}
-	return req, req.Model, nil
-}
-
-func rerankBodyParser(body []byte) (interface{}, string, error) {
-	var req cohereschema.RerankV2Request
-	if err := json.Unmarshal(body, &req); err != nil {
-		return nil, "", fmt.Errorf("failed to unmarshal body: %w", err)
-	}
-	return req, req.Model, nil
-}
-
-func messagesBodyParser(body []byte) (interface{}, string, error) {
-	var anthropicReq anthropic.MessagesRequest
-	if err := json.Unmarshal(body, &anthropicReq); err != nil {
-		return nil, "", fmt.Errorf("failed to unmarshal body: %w", err)
-	}
-	return anthropicReq, anthropicReq.GetModel(), nil
+	return req, modelExtractFn(&req), nil
 }

internal/dynamic_module/sdk/abi_envoy_v1.36.go

@@ -551,17 +551,33 @@ func (e envoyFilter) GetResponseBody() (io.Reader, bool) {
 	return &bodyReader{chunks: chunks}, true
 }
 
+// https://github.com/envoyproxy/envoy/blob/dc2d3098ae5641555f15c71d5bb5ce0060a8015c/source/extensions/dynamic_modules/abi.h#L271-L282
+type metadataSource uint32
+
+const (
+	metadataSourceDynamic      metadataSource = 0
+	metadataSourceUpstreamHost metadataSource = 3
+)
+
 // GetDynamicMetadataString implements [EnvoyHTTPFilter].
 func (e envoyFilter) GetDynamicMetadataString(namespace string, key string) (string, bool) {
+	return e.getMetadataString(namespace, key, metadataSourceDynamic)
+}
+
+// GetUpstreamHostMetadataString implements [EnvoyHTTPFilter].
+func (e envoyFilter) GetUpstreamHostMetadataString(namespace string, key string) (string, bool) {
+	return e.getMetadataString(namespace, key, metadataSourceUpstreamHost)
+}
+
+func (e envoyFilter) getMetadataString(namespace string, key string, source metadataSource) (string, bool) {
 	namespacePtr := uintptr(unsafe.Pointer(unsafe.StringData(namespace)))
 	keyPtr := uintptr(unsafe.Pointer(unsafe.StringData(key)))
 	var resultBufferPtr *byte
 	var resultBufferLengthPtr C.size_t
 
 	ret := C.envoy_dynamic_module_callback_http_get_metadata_string(
 		C.uintptr_t(e.raw),
-		// https://github.com/envoyproxy/envoy/blob/dc2d3098ae5641555f15c71d5bb5ce0060a8015c/source/extensions/dynamic_modules/abi.h#L273
-		C.uint32_t(0), // source
+		C.uint32_t(source),
 		C.uintptr_t(namespacePtr),
 		C.size_t(len(namespace)),
 		C.uintptr_t(keyPtr),

internal/dynamic_module/sdk/gosdk.go

@@ -54,6 +54,8 @@ type EnvoyHTTPFilter interface {
 	SendLocalReply(statusCode uint32, headers [][2]string, body []byte)
 	// GetDynamicMetadataString gets the dynamic metadata value for the given namespace and key. Returns the value and true if the value is found.
 	GetDynamicMetadataString(namespace string, key string) (string, bool)
+	// GetUpstreamHostMetadataString gets the upstream host metadata value for the given namespace and key. Returns the value and true if the value is found.
+	GetUpstreamHostMetadataString(namespace string, key string) (string, bool)
 	// SetDynamicMetadataString sets the dynamic metadata value for the given namespace and key. Returns true if the value is set successfully.
 	SetDynamicMetadataString(namespace string, key string, value string) bool
 }

internal/dynamic_module/upstream_filter.go

@@ -38,6 +38,17 @@ func (f *upstreamFilter) RequestHeaders(e sdk.EnvoyHTTPFilter, _ bool) sdk.Reque
 		return sdk.RequestHeadersStatusStopIteration
 	}
 	f.rf = (*routerFilter)(unsafe.Pointer(uintptr(rfPtr)))
+
+	backend, ok := e.GetUpstreamHostMetadataString(internalapi.AIGatewayFilterMetadataNamespace, internalapi.InternalMetadataBackendNameKey)
+	if !ok {
+		e.SendLocalReply(500, nil, []byte("backend name not found in upstream host metadata"))
+		return sdk.RequestHeadersStatusStopIteration
+	}
+	b, ok := f.rf.filterConfig().Backends[backend]
+	if !ok {
+		e.SendLocalReply(500, nil, []byte(fmt.Sprintf("backend %s not found in filter config", backend)))
+		return sdk.RequestHeadersStatusStopIteration
+	}
 	return sdk.RequestHeadersStatusContinue
 }