WiP fixes

Author: davidbreitgand

pkg/bbr/README.md

@@ -1,6 +1,6 @@
 # Body-Based Routing
 This package provides an extension that can be deployed to write the `model`
-HTTP body parameter as a header (X-Gateway-Model-Name) so as to enable routing capabilities on the
+HTTP body parameter as a header (`X-Gateway-Model-Name`) so as to enable routing capabilities on the
 model name.
 
 As per OpenAI spec, it is standard for the model name to be included in the

pkg/bbr/framework/interfaces.go

@@ -17,7 +17,6 @@ limitations under the License.
 package framework
 
 import (
-	"github.com/openai/openai-go/v3"
 	bbrplugins "sigs.k8s.io/gateway-api-inference-extension/pkg/bbr/plugins"
 )
 
@@ -48,12 +47,8 @@ type PluginsChain interface {
 	AddPluginAtInd(typeKey string, i int, r PluginRegistry) error               //only affects the instance of the plugin chain
 	GetPlugin(index int, registry PluginRegistry) (bbrplugins.BBRPlugin, error) //retrieves i-th plugin as defined in the chain from the registry
 	Length() int
-	ParseChatCompletion(data []byte) (openai.ChatCompletionNewParams, error)
-	ParseCompletion(data []byte) (openai.CompletionNewParams, error)
-	GetSharedChatCompletion() openai.ChatCompletionNewParams
-	GetSharedCompletion() openai.CompletionNewParams
-	GetSharedMemory(which string) interface{}
-	Run(bodyBytes []byte, metaDataKeys []string, registry PluginRegistry) ([]byte, map[string]string, error) //return potentially mutated body and all headers map safely merged
+	GetPlugins() []string
+	Run(bodyBytes []byte, metaDataKeys []string, registry PluginRegistry) (map[string]string, []byte, error) //return potentially mutated body and all headers map safely merged
 	String() string
 }
 

pkg/bbr/framework/registry.go

@@ -20,24 +20,13 @@ import (
 	"fmt"
 	"slices"
 
-	"github.com/openai/openai-go/v3" //imported as openai
 	bbrplugins "sigs.k8s.io/gateway-api-inference-extension/pkg/bbr/plugins"
 )
 
 // -------------------- INTERFACES -----------------------------------------------------------------------
 // Interfaces are defined in "sigs.k8s.io/gateway-api-inference-extension/pkg/bbr/framework/interfaces.go"
 
 // --------------------- PluginRegistry implementation ---------------------------------------------------
-// --------------------- Constructors lookup map ---------------------------------------------------------
-// This lookup map should be extended when a new concrete plugin implementation is added
-/* var pluginConstructors = map[string]PluginFactoryFunc{
-	"simple-model-extractor":  NewSimpleModelExtractor,
-	"semantic-model-selector": NewSemanticModelSelector,
-	"bad-words-blocker":       NewBadWordsBlocker,
-	"pid-disclosure-blocker":  NewPidDisclosureBlocker,
-	"extract-metadata-bykeys": NewExtractMetadataByKeys,
-} */
-// Registration: of constructors for implementations is done in the main's initPlugins based on the ConfigMap
 
 // pluginRegistry implements PluginRegistry
 type pluginRegistry struct {
@@ -123,15 +112,6 @@ func (r *pluginRegistry) UnregisterFactory(typeKey string) error {
 	return fmt.Errorf("plugin (%s) not found", typeKey)
 }
 
-// Removes  a plugin instance by type key
-func (r *pluginRegistry) UnregisterPlugin(typeKey string) error {
-	if _, ok := r.plugins[typeKey]; ok {
-		delete(r.plugins, typeKey)
-		return nil
-	}
-	return fmt.Errorf("plugin (%s) not found", typeKey)
-}
-
 // ListPlugins lists all registered plugins
 func (r *pluginRegistry) ListPlugins() []string {
 	typeKeys := make([]string, 0, len(r.plugins))
@@ -160,12 +140,6 @@ func (r *pluginRegistry) GetPlugins() map[string]bbrplugins.BBRPlugin {
 	return r.plugins
 }
 
-// Clear removes all registered factories and plugins
-func (r *pluginRegistry) Clear() {
-	r.pluginsFactory = make(map[string]PluginFactoryFunc)
-	r.plugins = make(map[string]bbrplugins.BBRPlugin)
-}
-
 // Checks for presense of a factory in this registry
 func (r *pluginRegistry) ContainsFactory(typeKey string) bool {
 	_, exists := r.pluginsFactory[typeKey]
@@ -186,17 +160,13 @@ func (r *pluginRegistry) String() string {
 
 // PluginsChain is a sequence of plugins to be executed in order inside the ext_proc server
 type pluginsChain struct {
-	plugins              []string
-	sharedChatCompletion openai.ChatCompletionNewParams //will be nil if an instance of pluginsChain does not contain a plugin that requires full parsing
-	sharedCompletion     openai.CompletionNewParams     //likewise
+	plugins []string
 }
 
 // NewPluginsChain creates a new PluginsChain instance
 func NewPluginsChain() PluginsChain {
 	return &pluginsChain{
-		plugins:              []string{},
-		sharedChatCompletion: openai.ChatCompletionNewParams{},
-		sharedCompletion:     openai.CompletionNewParams{},
+		plugins: []string{},
 	}
 }
 
@@ -212,17 +182,6 @@ func (pc *pluginsChain) AddPlugin(typeKey string, r PluginRegistry) error {
 	return nil
 }
 
-// DeletePlugin deletes a plugin from the chain
-func (pc *pluginsChain) DeletePlugin(p string) error {
-	for i := range len(pc.plugins) {
-		if pc.plugins[i] == p {
-			pc.plugins = append(pc.plugins[:i], pc.plugins[i+1:]...)
-			return nil
-		}
-	}
-	return fmt.Errorf("plugin %s not found in chain", p)
-}
-
 // GetPlugin retrieves the next plugin in the chain by index
 func (pc *pluginsChain) GetPlugin(index int, r PluginRegistry) (bbrplugins.BBRPlugin, error) {
 	if index < 0 || index >= len(pc.plugins) {
@@ -255,36 +214,8 @@ func (pc *pluginsChain) AddPluginAtInd(typeKey string, i int, r PluginRegistry)
 	return nil
 }
 
-func (pc *pluginsChain) ParseChatCompletion(data []byte) (openai.ChatCompletionNewParams, error) {
-	if err := pc.sharedChatCompletion.UnmarshalJSON(data); err != nil {
-		return pc.sharedChatCompletion, err
-	}
-	return pc.sharedChatCompletion, nil
-}
-
-func (pc *pluginsChain) ParseCompletion(data []byte) (openai.CompletionNewParams, error) {
-	if err := pc.sharedCompletion.UnmarshalJSON(data); err != nil {
-		return pc.sharedCompletion, err
-	}
-	return pc.sharedCompletion, nil
-}
-
-func (pc *pluginsChain) GetSharedChatCompletion() openai.ChatCompletionNewParams {
-	return pc.sharedChatCompletion
-}
-
-func (pc *pluginsChain) GetSharedCompletion() openai.CompletionNewParams {
-	return pc.sharedCompletion
-}
-
-func (pc *pluginsChain) GetSharedMemory(which string) interface{} {
-	if which == "/v1/completions" {
-		return pc.sharedCompletion
-	}
-	if which == "/v1/chat/completions" {
-		return pc.sharedChatCompletion
-	}
-	return nil
+func (pc *pluginsChain) GetPlugins() []string {
+	return pc.plugins
 }
 
 // MergeMaps copies all key/value pairs from src into dst and returns dst.
@@ -315,7 +246,7 @@ func (pc *pluginsChain) Run(
 	bodyBytes []byte,
 	metaDataKeys []string,
 	r PluginRegistry,
-) (mutateBodyBytes []byte, headers map[string]string, err error) {
+) (headers map[string]string, mutateBodyBytes []byte, err error) {
 
 	allHeaders := make(map[string]string)
 	mutatedBodyBytes := bodyBytes
@@ -327,20 +258,20 @@ func (pc *pluginsChain) Run(
 		metExtPlugin, err := r.GetPlugin(pluginType)
 
 		if err != nil {
-			return bodyBytes, allHeaders, err
+			return allHeaders, bodyBytes, err
 		}
 
 		// The plugin i in the chain receives the (potentially mutated) body from plugin i-1 in the chain
 		headers, mutatedBodyBytes, err := metExtPlugin.Execute(mutatedBodyBytes, metaDataKeys)
 
 		if err != nil {
-			return mutatedBodyBytes, headers, err
+			return headers, mutatedBodyBytes, err
 		}
 
 		//note that the existing overlapping keys are NOT over-written by merge
 		MergeMaps(allHeaders, headers)
 	}
-	return
+	return allHeaders, mutatedBodyBytes, nil
 }
 
 func (pc *pluginsChain) String() string {

pkg/bbr/handlers/request.go

@@ -18,6 +18,7 @@ package handlers
 
 import (
 	"context"
+	"strings"
 
 	basepb "github.com/envoyproxy/go-control-plane/envoy/config/core/v3"
 	eppb "github.com/envoyproxy/go-control-plane/envoy/service/ext_proc/v3"
@@ -31,45 +32,34 @@ import (
 // HandleRequestBody handles request bodies.
 func (s *Server) HandleRequestBody(ctx context.Context, requestBodyBytes []byte) ([]*eppb.ProcessingResponse, error) {
 	logger := log.FromContext(ctx)
+	var ret []*eppb.ProcessingResponse
 
-	allHeaders, mutatedBodyBytes, _ := helpers.RunPluginsChain( //TODO
-		ctx,
-		logger,
-		requestBodyBytes,
-		s.requestChain,
-		s.registry,
-		s.metaDataKeys)
-
-	//At this point, we have all the headers and a mutated body (note that actually, the body might not be mutated, but we do not care)
+	allHeaders, mutatedBodyBytes, err := s.requestChain.Run(requestBodyBytes, s.metaDataKeys, s.registry)
 
-	var ret []*eppb.ProcessingResponse
+	if err != nil {
+		//TODO: add metric in metrics.go to count "other errors"
+		logger.V(logutil.DEFAULT).Info("error processing body", "error", err)
+		ret, _ := buildEmptyResponsesForMissingModel(s.streaming, requestBodyBytes)
+		return ret, nil
+	}
 
-	// process headers
-	Model := allHeaders[bbrplugins.ModelHeader] //it is required that the ModelHeader is always set (i.e., that there always exist requestPluginsChain with at least one plugin that sets the model header)
+	model, ok := allHeaders[bbrplugins.ModelHeader]
 
-	logger.V(logutil.DEFAULT).Info("model extracted from request body", "model", Model)
+	if !ok {
+		//TODO: add metric in metrics.go to count "other errors"
+		logger.V(logutil.DEFAULT).Info("manadatory header X-Gateway-Model-Name value is undetermined")
+		ret, _ := buildEmptyResponsesForMissingModel(s.streaming, requestBodyBytes)
+		return ret, nil
+	}
 
-	if Model == "" {
+	if strings.TrimSpace(model) == "" {
 		metrics.RecordModelNotInBodyCounter()
-
-		if s.streaming {
-			ret = append(ret, &eppb.ProcessingResponse{
-				Response: &eppb.ProcessingResponse_RequestHeaders{
-					RequestHeaders: &eppb.HeadersResponse{},
-				},
-			})
-			ret = addStreamedBodyResponse(ret, requestBodyBytes)
-			return ret, nil
-		} else {
-			ret = append(ret, &eppb.ProcessingResponse{
-				Response: &eppb.ProcessingResponse_RequestBody{
-					RequestBody: &eppb.BodyResponse{},
-				},
-			})
-		}
+		ret, _ := buildEmptyResponsesForMissingModel(s.streaming, requestBodyBytes)
 		return ret, nil
 	}
 
+	logger.V(logutil.DEFAULT).Info("model extracted from request body", "model", model)
+
 	metrics.RecordSuccessCounter()
 
 	if s.streaming {
@@ -83,7 +73,7 @@ func (s *Server) HandleRequestBody(ctx context.Context, requestBodyBytes []byte)
 								{
 									Header: &basepb.HeaderValue{
 										Key:      bbrplugins.ModelHeader,
-										RawValue: []byte(Model),
+										RawValue: []byte(model),
 									},
 								},
 							},
@@ -111,7 +101,7 @@ func (s *Server) HandleRequestBody(ctx context.Context, requestBodyBytes []byte)
 								{
 									Header: &basepb.HeaderValue{
 										Key:      bbrplugins.ModelHeader,
-										RawValue: []byte(Model),
+										RawValue: []byte(model),
 									},
 								},
 							},
@@ -168,3 +158,31 @@ func (s *Server) HandleRequestTrailers(trailers *eppb.HttpTrailers) ([]*eppb.Pro
 		},
 	}, nil
 }
+
+// buildEmptyResponsesForMissingModel is a local helper that returns the appropriate empty responses
+// for the "model not found" branch depending on streaming mode.
+// It is also used to create empty responses in case of other errors related to running plugins on the body
+// This is not very clean and MUST be segregated in the future.
+// Corresponding metrics should be defined to make different errors observable
+func buildEmptyResponsesForMissingModel(streaming bool, requestBodyBytes []byte) ([]*eppb.ProcessingResponse, error) {
+	var ret []*eppb.ProcessingResponse
+
+	if streaming {
+		// Emit empty headers response, then stream body unchanged.
+		ret = append(ret, &eppb.ProcessingResponse{
+			Response: &eppb.ProcessingResponse_RequestHeaders{
+				RequestHeaders: &eppb.HeadersResponse{},
+			},
+		})
+		ret = addStreamedBodyResponse(ret, requestBodyBytes)
+		return ret, nil
+	}
+
+	// Non-streaming: emit empty body response.
+	ret = append(ret, &eppb.ProcessingResponse{
+		Response: &eppb.ProcessingResponse_RequestBody{
+			RequestBody: &eppb.BodyResponse{},
+		},
+	})
+	return ret, nil
+}

pkg/bbr/handlers/server_test.go

@@ -128,6 +128,7 @@ func TestProcessRequestBody(t *testing.T) {
 	}
 
 	//Initialize PluginRegistry and request/response PluginsChain instances
+	//Change to hermetic
 	registry, requestChain, responseChain, metaDataKeys, _ := bbrutils.InitPlugins()
 
 	for _, tc := range cases {

pkg/bbr/plugins/simple_model_extractor.go

@@ -32,7 +32,7 @@ import (
 
 type simpleModelExtractor struct { //implements the MetadataExtractor interface
 	typedName           plugins.TypedName
-	requiresFullParsing bool
+	requiresFullParsing bool //this field will be used to determine whether shared struct should be created in this chain
 }
 
 // NewSimpleModelExtractor is a factory that constructs SimpleModelExtractor plugin