Add streaming response process.

Author: zetxqx

pkg/epp/handlers/response.go

@@ -17,16 +17,17 @@ limitations under the License.
 package handlers
 
 import (
+	"bytes"
 	"context"
 	"encoding/json"
-	"fmt"
 	"strings"
 
 	configPb "github.com/envoyproxy/go-control-plane/envoy/config/core/v3"
 	extProcPb "github.com/envoyproxy/go-control-plane/envoy/service/ext_proc/v3"
 	"sigs.k8s.io/controller-runtime/pkg/log"
 
 	"sigs.k8s.io/gateway-api-inference-extension/pkg/epp/metrics"
+	"sigs.k8s.io/gateway-api-inference-extension/pkg/epp/scheduling/types"
 	logutil "sigs.k8s.io/gateway-api-inference-extension/pkg/epp/util/logging"
 )
 
@@ -36,49 +37,50 @@ const (
 )
 
 // HandleResponseBody always returns the requestContext even in the error case, as the request context is used in error handling.
-func (s *StreamingServer) HandleResponseBody(ctx context.Context, reqCtx *RequestContext, response map[string]any) (*RequestContext, error) {
+func (s *StreamingServer) HandleResponseBody(ctx context.Context, reqCtx *RequestContext, body []byte) (*RequestContext, error) {
 	logger := log.FromContext(ctx)
-	responseBytes, err := json.Marshal(response)
+	llmResponse, err := types.NewLLMResponseFromBytes(body)
 	if err != nil {
-		return reqCtx, fmt.Errorf("error marshalling responseBody - %w", err)
+		logger.Error(err, "failed to create LLMResponse from bytes")
+		return reqCtx, err
 	}
-	if response["usage"] != nil {
-		usg := response["usage"].(map[string]any)
-		usage := Usage{
-			PromptTokens:     int(usg["prompt_tokens"].(float64)),
-			CompletionTokens: int(usg["completion_tokens"].(float64)),
-			TotalTokens:      int(usg["total_tokens"].(float64)),
-		}
+	reqCtx.SchedulingResponse = llmResponse
+	if usage := reqCtx.SchedulingResponse.Usage(); usage != nil {
 		reqCtx.Usage = usage
-		logger.V(logutil.VERBOSE).Info("Response generated", "usage", reqCtx.Usage)
+		logger.V(logutil.VERBOSE).Info("Response generated", "usage", usage)
 	}
-	reqCtx.ResponseSize = len(responseBytes)
+	reqCtx.ResponseSize = len(body)
 	// ResponseComplete is to indicate the response is complete. In non-streaming
 	// case, it will be set to be true once the response is processed; in
 	// streaming case, it will be set to be true once the last chunk is processed.
 	// TODO(https://github.com/kubernetes-sigs/gateway-api-inference-extension/issues/178)
 	// will add the processing for streaming case.
 	reqCtx.ResponseComplete = true
 
-	reqCtx.respBodyResp = generateResponseBodyResponses(responseBytes, true)
+	reqCtx.respBodyResp = generateResponseBodyResponses(body, true)
 
 	return s.director.HandleResponseBodyComplete(ctx, reqCtx)
 }
 
 // The function is to handle streaming response if the modelServer is streaming.
-func (s *StreamingServer) HandleResponseBodyModelStreaming(ctx context.Context, reqCtx *RequestContext, responseText string) {
+func (s *StreamingServer) HandleResponseBodyModelStreaming(ctx context.Context, reqCtx *RequestContext, streamBody []byte) {
 	logger := log.FromContext(ctx)
 	_, err := s.director.HandleResponseBodyStreaming(ctx, reqCtx, logger)
 	if err != nil {
 		logger.Error(err, "error in HandleResponseBodyStreaming")
 	}
-	if strings.Contains(responseText, streamingEndMsg) {
+	if bytes.Contains(streamBody, []byte(streamingEndMsg)) {
 		reqCtx.ResponseComplete = true
-		resp := parseRespForUsage(ctx, responseText)
-		reqCtx.Usage = resp.Usage
-		metrics.RecordInputTokens(reqCtx.IncomingModelName, reqCtx.TargetModelName, resp.Usage.PromptTokens)
-		metrics.RecordOutputTokens(reqCtx.IncomingModelName, reqCtx.TargetModelName, resp.Usage.CompletionTokens)
-		_, err := s.director.HandleResponseBodyComplete(ctx, reqCtx)
+		resp, err := types.NewLLMResponseFromStream(streamBody)
+		if err != nil {
+			logger.Error(err, "error in converting stream response to LLMResponse.")
+		}
+		if usage := resp.Usage(); usage != nil {
+			reqCtx.Usage = usage
+			metrics.RecordInputTokens(reqCtx.IncomingModelName, reqCtx.TargetModelName, usage.PromptTokens)
+			metrics.RecordOutputTokens(reqCtx.IncomingModelName, reqCtx.TargetModelName, usage.CompletionTokens)
+		}
+		_, err = s.director.HandleResponseBodyComplete(ctx, reqCtx)
 		if err != nil {
 			logger.Error(err, "error in HandleResponseBodyComplete")
 		}

pkg/epp/handlers/response_test.go

@@ -18,13 +18,13 @@ package handlers
 
 import (
 	"context"
-	"encoding/json"
 	"testing"
 
 	"github.com/go-logr/logr"
 	"github.com/google/go-cmp/cmp"
 
 	"sigs.k8s.io/gateway-api-inference-extension/pkg/epp/backend"
+	"sigs.k8s.io/gateway-api-inference-extension/pkg/epp/scheduling/types"
 	logutil "sigs.k8s.io/gateway-api-inference-extension/pkg/epp/util/logging"
 )
 
@@ -53,12 +53,33 @@ const (
 	}
 	`
 
-	streamingBodyWithoutUsage = `data: {"id":"cmpl-41764c93-f9d2-4f31-be08-3ba04fa25394","object":"text_completion","created":1740002445,"model":"food-review-0","choices":[],"usage":null}
-	`
+	streamingBodyWithoutUsage = `
+		    data: {"id":"chatcmpl-1","object":"chat.completion.chunk","choices":[{"index":0,"delta":{"role":"assistant"}}]} 
 
-	streamingBodyWithUsage = `data: {"id":"cmpl-41764c93-f9d2-4f31-be08-3ba04fa25394","object":"text_completion","created":1740002445,"model":"food-review-0","choices":[],"usage":{"prompt_tokens":7,"total_tokens":17,"completion_tokens":10}}
-data: [DONE]
-	`
+			data: {"id":"chatcmpl-1","object":"chat.completion.chunk","choices":[{"index":0,"delta":{"content":"Hello"}}]} 
+
+			data: {"id":"chatcmpl-1","object":"chat.completion.chunk","choices":[{"index":0,"delta":{"content":" world"}}]} 
+
+			data: {"id":"chatcmpl-1","object":"chat.completion.chunk","choices":[{"index":0,"delta":{},"finish_reason":"stop"}]} 
+
+			data: {"id":"chatcmpl-1","object":"chat.completion.chunk","choices":[],"usage":null} 
+
+			data: [DONE]
+	  		`
+
+	streamingBodyWithUsage = `
+			data: {"id":"chatcmpl-1","object":"chat.completion.chunk","choices":[{"index":0,"delta":{"role":"assistant"}}]} 
+
+			data: {"id":"chatcmpl-1","object":"chat.completion.chunk","choices":[{"index":0,"delta":{"content":"Hello"}}]} 
+
+			data: {"id":"chatcmpl-1","object":"chat.completion.chunk","choices":[{"index":0,"delta":{"content":" world"}}]} 
+
+			data: {"id":"chatcmpl-1","object":"chat.completion.chunk","choices":[{"index":0,"delta":{},"finish_reason":"stop"}]} 
+
+			data: {"id":"chatcmpl-1","object":"chat.completion.chunk","choices":[],"usage":{"prompt_tokens":5,"completion_tokens":7,"total_tokens":12}} 
+
+			data: [DONE]
+			`
 )
 
 type mockDirector struct{}
@@ -89,13 +110,13 @@ func TestHandleResponseBody(t *testing.T) {
 		name    string
 		body    []byte
 		reqCtx  *RequestContext
-		want    Usage
+		want    *types.Usage
 		wantErr bool
 	}{
 		{
 			name: "success",
 			body: []byte(body),
-			want: Usage{
+			want: &types.Usage{
 				PromptTokens:     11,
 				TotalTokens:      111,
 				CompletionTokens: 100,
@@ -111,12 +132,7 @@ func TestHandleResponseBody(t *testing.T) {
 			if reqCtx == nil {
 				reqCtx = &RequestContext{}
 			}
-			var responseMap map[string]any
-			marshalErr := json.Unmarshal(test.body, &responseMap)
-			if marshalErr != nil {
-				t.Error(marshalErr, "Error unmarshaling request body")
-			}
-			_, err := server.HandleResponseBody(ctx, reqCtx, responseMap)
+			_, err := server.HandleResponseBody(ctx, reqCtx, test.body)
 			if err != nil {
 				if !test.wantErr {
 					t.Fatalf("HandleResponseBody returned unexpected error: %v, want %v", err, test.wantErr)
@@ -137,7 +153,7 @@ func TestHandleStreamedResponseBody(t *testing.T) {
 		name    string
 		body    string
 		reqCtx  *RequestContext
-		want    Usage
+		want    *types.Usage
 		wantErr bool
 	}{
 		{
@@ -156,10 +172,10 @@ func TestHandleStreamedResponseBody(t *testing.T) {
 				modelServerStreaming: true,
 			},
 			wantErr: false,
-			want: Usage{
-				PromptTokens:     7,
-				TotalTokens:      17,
-				CompletionTokens: 10,
+			want: &types.Usage{
+				PromptTokens:     5,
+				TotalTokens:      12,
+				CompletionTokens: 7,
 			},
 		},
 	}
@@ -172,7 +188,7 @@ func TestHandleStreamedResponseBody(t *testing.T) {
 			if reqCtx == nil {
 				reqCtx = &RequestContext{}
 			}
-			server.HandleResponseBodyModelStreaming(ctx, reqCtx, test.body)
+			server.HandleResponseBodyModelStreaming(ctx, reqCtx, []byte(test.body))
 
 			if diff := cmp.Diff(test.want, reqCtx.Usage); diff != "" {
 				t.Errorf("HandleResponseBody returned unexpected response, diff(-want, +got): %v", diff)

pkg/epp/handlers/server.go

@@ -85,14 +85,15 @@ type RequestContext struct {
 	RequestReceivedTimestamp  time.Time
 	ResponseCompleteTimestamp time.Time
 	RequestSize               int
-	Usage                     Usage
+	Usage                     *schedulingtypes.Usage
 	ResponseSize              int
 	ResponseComplete          bool
 	ResponseStatusCode        string
 	RequestRunning            bool
 	Request                   *Request
 
-	SchedulingRequest *schedulingtypes.LLMRequest
+	SchedulingRequest  *schedulingtypes.LLMRequest
+	SchedulingResponse *schedulingtypes.LLMResponse
 
 	RequestState         StreamRequestState
 	modelServerStreaming bool
@@ -115,7 +116,6 @@ type Request struct {
 }
 type Response struct {
 	Headers map[string]string
-	Body    []byte
 }
 type StreamRequestState int
 
@@ -268,11 +268,10 @@ func (s *StreamingServer) Process(srv extProcPb.ExternalProcessor_ProcessServer)
 			reqCtx.respHeaderResp = s.generateResponseHeaderResponse(reqCtx)
 
 		case *extProcPb.ProcessingRequest_ResponseBody:
+			body = append(body, v.ResponseBody.Body...)
 			if reqCtx.modelServerStreaming {
 				// Currently we punt on response parsing if the modelServer is streaming, and we just passthrough.
-
-				responseText := string(v.ResponseBody.Body)
-				s.HandleResponseBodyModelStreaming(ctx, reqCtx, responseText)
+				s.HandleResponseBodyModelStreaming(ctx, reqCtx, body)
 				if v.ResponseBody.EndOfStream {
 					loggerTrace.Info("stream completed")
 
@@ -283,8 +282,6 @@ func (s *StreamingServer) Process(srv extProcPb.ExternalProcessor_ProcessServer)
 
 				reqCtx.respBodyResp = generateResponseBodyResponses(v.ResponseBody.Body, v.ResponseBody.EndOfStream)
 			} else {
-				body = append(body, v.ResponseBody.Body...)
-
 				// Message is buffered, we can read and decode.
 				if v.ResponseBody.EndOfStream {
 					loggerTrace.Info("stream completed")
@@ -303,8 +300,7 @@ func (s *StreamingServer) Process(srv extProcPb.ExternalProcessor_ProcessServer)
 						break
 					}
 
-					reqCtx.Response.Body = body
-					reqCtx, responseErr = s.HandleResponseBody(ctx, reqCtx, responseBody)
+					reqCtx, responseErr = s.HandleResponseBody(ctx, reqCtx, body)
 					if responseErr != nil {
 						if logger.V(logutil.DEBUG).Enabled() {
 							logger.V(logutil.DEBUG).Error(responseErr, "Failed to process response body", "request", req)

pkg/epp/requestcontrol/director.go

@@ -309,12 +309,11 @@ func (d *Director) HandleResponseBodyComplete(ctx context.Context, reqCtx *handl
 	requestID := reqCtx.Request.Headers[requtil.RequestIdHeaderKey]
 	logger := log.FromContext(ctx).WithValues("stage", "bodyChunk", requtil.RequestIdHeaderKey, requestID)
 	logger.V(logutil.DEBUG).Info("Entering HandleResponseBodyComplete")
-	llmResponse, err := schedulingtypes.NewLLMResponseFromBytes(reqCtx.Response.Body)
-	if err != nil {
-		logger.Error(err, "HandleResponseBodyComplete: failed to convert the response to LLMResponse.")
+	if reqCtx.SchedulingResponse == nil {
+		err := fmt.Errorf("nil scheduling reponse from reqCtx")
 		return reqCtx, err
 	}
-	d.runResponseCompletePlugins(ctx, reqCtx.SchedulingRequest, llmResponse, reqCtx.TargetPod)
+	d.runResponseCompletePlugins(ctx, reqCtx.SchedulingRequest, reqCtx.SchedulingResponse, reqCtx.TargetPod)
 
 	logger.V(logutil.DEBUG).Info("Exiting HandleResponseBodyComplete")
 	return reqCtx, nil

pkg/epp/requestcontrol/director_test.go

@@ -704,9 +704,9 @@ func TestDirector_HandleResponseComplete(t *testing.T) {
 		},
 		Response: &handlers.Response{
 			Headers: map[string]string{"X-Test-Complete-Header": "CompleteValue"},
-			Body:    []byte(chatCompletionJSON),
 		},
-		TargetPod: &backend.Pod{NamespacedName: types.NamespacedName{Namespace: "namespace1", Name: "test-pod-name"}},
+		SchedulingResponse: wantLLMResponse,
+		TargetPod:          &backend.Pod{NamespacedName: types.NamespacedName{Namespace: "namespace1", Name: "test-pod-name"}},
 	}
 
 	_, err = director.HandleResponseBodyComplete(ctx, reqCtx)