- revert suggestion on h.Write calls in prefix plugin

- refactoring based on review

Signed-off-by: Maroon Ayoub <[email protected]>

Author: vMaroon

pkg/epp/scheduling/framework/plugins/multi/prefix/plugin.go

@@ -123,8 +123,10 @@ func (s *SchedulingContextState) Clone() plugins.StateData {
 }
 
 // compile-time type assertion
-var _ framework.Scorer = &Plugin{}
-var _ requestcontrol.PreRequest = &Plugin{}
+var (
+	_ framework.Scorer          = &Plugin{}
+	_ requestcontrol.PreRequest = &Plugin{}
+)
 
 // PrefixCachePluginFactory defines the factory function for Prefix plugin.
 func PrefixCachePluginFactory(name string, rawParameters json.RawMessage, handle plugins.Handle) (plugins.Plugin, error) {
@@ -238,7 +240,6 @@ func (p *Plugin) matchLongestPrefix(ctx context.Context, hashes []BlockHash) map
 			for server := range cachedServers {
 				// Update servers with their longest prefix match.
 				res[server]++
-
 			}
 		}
 	}
@@ -269,17 +270,17 @@ func hashPrompt(ctx context.Context, request *types.LLMRequest, cacheBlockSize i
 		loggerDebug.Info("Truncating input", "size", len(userInput), "max prefix blocks", maxPrefixBlocks, "block size", cacheBlockSize)
 		userInput = userInput[:maxPrefixBlocks*cacheBlockSize]
 	}
-	// Split the body into blocks of size cacheBlockSize. 
+	// Split the body into blocks of size cacheBlockSize.
 	// If the last block is smaller than cacheBlockSize, it will be ignored.
 	res := make([]BlockHash, 0, len(userInput)/cacheBlockSize)
 	// Add the model to the first block hash so that different models have different hashes even with the same body.
 	h := xxhash.New()
-	h.Write([]byte(request.TargetModel))
+	_, _ = h.Write([]byte(request.TargetModel))
 	prevBlockHash := BlockHash(h.Sum64())
 	for i := 0; i+cacheBlockSize <= len(userInput); i += cacheBlockSize {
 		h.Reset()
-		h.Write(userInput[i : i+cacheBlockSize])
-		h.Write(toBytes(prevBlockHash))
+		_, _ = h.Write(userInput[i : i+cacheBlockSize])
+		_, _ = h.Write(toBytes(prevBlockHash))
 		res = append(res, BlockHash(h.Sum64()))
 
 		prevBlockHash = res[len(res)-1]

pkg/epp/scheduling/types/types.go

@@ -23,6 +23,8 @@ import (
 	backendmetrics "sigs.k8s.io/gateway-api-inference-extension/pkg/epp/backend/metrics"
 )
 
+const nilString = "<nil>"
+
 // LLMRequest is a structured representation of the fields we parse out of the LLMRequest body.
 type LLMRequest struct {
 	// RequestId is the Envoy generated Id for the request being processed
@@ -36,6 +38,10 @@ type LLMRequest struct {
 }
 
 func (r *LLMRequest) String() string {
+	if r == nil {
+		return nilString
+	}
+
 	return fmt.Sprintf("RequestID: %s, TargetModel: %s, RequestData: %s, Headers: %v",
 		r.RequestId, r.TargetModel, r.Data, r.Headers)
 }
@@ -51,12 +57,19 @@ type LLMRequestData struct {
 }
 
 func (r *LLMRequestData) String() string {
+	if r == nil {
+		return nilString
+	}
+
 	if r.Completions != nil {
-		return "Completions: " + r.Completions.String()
+		return r.Completions.String()
+	}
+
+	if r.ChatCompletions != nil {
+		return r.ChatCompletions.String()
 	}
 
-	// Must be a ChatCompletionsRequest
-	return "ChatCompletions: " + r.ChatCompletions.String()
+	return ""
 }
 
 // CompletionsRequest is a structured representation of the fields we parse out of the
@@ -65,10 +78,14 @@ func (r *LLMRequestData) String() string {
 // API spec.
 type CompletionsRequest struct {
 	// Prompt is the prompt that was sent in the request body.
-	Prompt string
+	Prompt string `json:"prompt,omitempty"`
 }
 
 func (r *CompletionsRequest) String() string {
+	if r == nil {
+		return nilString
+	}
+
 	return fmt.Sprintf("{PromptLength: %d}", len(r.Prompt))
 }
 
@@ -78,7 +95,7 @@ func (r *CompletionsRequest) String() string {
 // API spec.
 type ChatCompletionsRequest struct {
 	/* parameters from the official OpenAI chat-completions API */
-	Messages []Message
+	Messages []Message     `json:"messages,omitempty"`
 	Tools    []interface{} `json:"tools,omitempty"`
 	/* parameters from the HuggingFace transformers chat-templates API */
 	Documents                 []interface{}          `json:"documents,omitempty"`
@@ -90,6 +107,10 @@ type ChatCompletionsRequest struct {
 }
 
 func (r *ChatCompletionsRequest) String() string {
+	if r == nil {
+		return nilString
+	}
+
 	messagesLen := 0
 	for _, msg := range r.Messages {
 		messagesLen += len(msg.Content)
@@ -117,8 +138,9 @@ type ScoredPod struct {
 
 func (pm *PodMetrics) String() string {
 	if pm == nil {
-		return ""
+		return nilString
 	}
+
 	return fmt.Sprintf("%+v", *pm)
 }
 

pkg/epp/util/request/body.go

@@ -55,14 +55,5 @@ func validateChatCompletionsMessages(messages []types.Message) error {
 		return errutil.Error{Code: errutil.BadRequest, Msg: "chat-completions request must have at least one message"}
 	}
 
-	for i, msg := range messages {
-		if msg.Role == "" {
-			return errutil.Error{Code: errutil.BadRequest, Msg: "message at index " + string(rune(i)) + " is missing role"}
-		}
-		if msg.Content == "" {
-			return errutil.Error{Code: errutil.BadRequest, Msg: "message at index " + string(rune(i)) + " is missing content"}
-		}
-	}
-
 	return nil
 }

pkg/epp/util/request/body_test.go

@@ -19,6 +19,7 @@ package request
 import (
 	"testing"
 
+	"github.com/google/go-cmp/cmp"
 	"sigs.k8s.io/gateway-api-inference-extension/pkg/epp/scheduling/types"
 )
 
@@ -127,26 +128,6 @@ func TestExtractRequestData(t *testing.T) {
 			},
 			wantErr: true,
 		},
-		{
-			name: "message missing role",
-			body: map[string]any{
-				"model": "test",
-				"messages": []any{
-					map[string]any{"content": "hello"},
-				},
-			},
-			wantErr: true,
-		},
-		{
-			name: "message missing content",
-			body: map[string]any{
-				"model": "test",
-				"messages": []any{
-					map[string]any{"role": "user"},
-				},
-			},
-			wantErr: true,
-		},
 		{
 			name: "message with non-string role",
 			body: map[string]any{
@@ -257,111 +238,13 @@ func TestExtractRequestData(t *testing.T) {
 				return
 			}
 
-			// Compare the results
-			if !compareResults(got, tt.want, t) {
-				t.Errorf("ExtractRequestData() result mismatch")
+			if diff := cmp.Diff(tt.want, got); diff != "" {
+				t.Errorf("ExtractRequestData() mismatch (-want +got):\n%s", diff)
 			}
 		})
 	}
 }
 
-func compareResults(got, want *types.LLMRequestData, t *testing.T) bool {
-	switch {
-	case got.Completions != nil && want.Completions != nil:
-		return compareCompletionsRequest(got.Completions, want.Completions, t)
-	case got.ChatCompletions != nil && want.ChatCompletions != nil:
-		return compareChatCompletionsRequest(got.ChatCompletions, want.ChatCompletions, t)
-	case got.Completions == nil && want.Completions == nil && got.ChatCompletions == nil && want.ChatCompletions == nil:
-		return true
-	default:
-		t.Errorf("Result type mismatch: got completions=%v, chatCompletions=%v; want completions=%v, chatCompletions=%v",
-			got.Completions != nil, got.ChatCompletions != nil, want.Completions != nil, want.ChatCompletions != nil)
-		return false
-	}
-}
-
-func compareCompletionsRequest(got, want *types.CompletionsRequest, t *testing.T) bool {
-	if got.Prompt != want.Prompt {
-		t.Errorf("CompletionsRequest.Prompt = %v, want %v", got.Prompt, want.Prompt)
-		return false
-	}
-	return true
-}
-
-func compareChatCompletionsRequest(got, want *types.ChatCompletionsRequest, t *testing.T) bool {
-	// Compare messages
-	if len(got.Messages) != len(want.Messages) {
-		t.Errorf("Messages length = %v, want %v", len(got.Messages), len(want.Messages))
-		return false
-	}
-	for i, msg := range got.Messages {
-		wantMsg := want.Messages[i]
-		if msg.Role != wantMsg.Role || msg.Content != wantMsg.Content {
-			t.Errorf("Message[%d] = %v, want %v", i, msg, wantMsg)
-			return false
-		}
-	}
-
-	// Compare optional fields
-	if got.ChatTemplate != want.ChatTemplate {
-		t.Errorf("ChatTemplate = %v, want %v", got.ChatTemplate, want.ChatTemplate)
-		return false
-	}
-	if got.ReturnAssistantTokensMask != want.ReturnAssistantTokensMask {
-		t.Errorf("ReturnAssistantTokensMask = %v, want %v", got.ReturnAssistantTokensMask, want.ReturnAssistantTokensMask)
-		return false
-	}
-	if got.ContinueFinalMessage != want.ContinueFinalMessage {
-		t.Errorf("ContinueFinalMessage = %v, want %v", got.ContinueFinalMessage, want.ContinueFinalMessage)
-		return false
-	}
-	if got.AddGenerationPrompt != want.AddGenerationPrompt {
-		t.Errorf("AddGenerationPrompt = %v, want %v", got.AddGenerationPrompt, want.AddGenerationPrompt)
-		return false
-	}
-
-	// Compare tools (shallow comparison for test purposes)
-	if !compareSliceAny(got.Tools, want.Tools) {
-		t.Errorf("Tools mismatch")
-		return false
-	}
-
-	// Compare documents (shallow comparison for test purposes)
-	if !compareSliceAny(got.Documents, want.Documents) {
-		t.Errorf("Documents mismatch")
-		return false
-	}
-
-	// Compare chat template kwargs (shallow comparison for test purposes)
-	if !compareMapAny(got.ChatTemplateKWArgs, want.ChatTemplateKWArgs) {
-		t.Errorf("ChatTemplateKWArgs mismatch")
-		return false
-	}
-
-	return true
-}
-
-func compareSliceAny(got, want []any) bool {
-	if len(got) != len(want) {
-		return false
-	}
-	// For test purposes, we'll do a simple length check and type check
-	// In practice, you might want deeper comparison depending on your needs
-	return true
-}
-
-func compareMapAny(got, want map[string]any) bool {
-	if len(got) != len(want) {
-		return false
-	}
-	for k, v := range want {
-		if gotV, exists := got[k]; !exists || gotV != v {
-			return false
-		}
-	}
-	return true
-}
-
 // Benchmark tests for performance comparison
 func BenchmarkExtractRequestData_Completions(b *testing.B) {
 	body := map[string]any{