Repository: vllm-project/aibrix
File: pkg/plugins/gateway/gateway.go (Go)

Code before:
			if len(evt.Choices) == 0 {
				// Do not overwrite model, res can be empty.
				usage = evt.Usage
			}
		}
		if err := streaming.Err(); err != nil {
			klog.ErrorS(err, "error to unmarshal response", "requestID", requestID, "responseBody", string(b.ResponseBody.GetBody()))
			complete = true
			return generateErrorResponse(
				envoyTypePb.StatusCode_InternalServerError,
				[]*configPb.HeaderValueOption{{Header: &configPb.HeaderValue{
					Key: "x-streaming-error", RawValue: []byte("true"),
				}}},
				err.Error()), complete
		}
	} else {
		if err := json.Unmarshal(b.ResponseBody.Body, &res); err != nil {
			klog.ErrorS(err, "error to unmarshal response", "requestID", requestID, "responseBody", string(b.ResponseBody.GetBody()))
			complete = true
			return generateErrorResponse(
				envoyTypePb.StatusCode_InternalServerError,
				[]*configPb.HeaderValueOption{{Header: &configPb.HeaderValue{
					Key: "x-error-response-unmarshal", RawValue: []byte("true"),
				}}},
				err.Error()), complete
		} else if res.Model != model {
			err = ErrorUnknownResponse
			klog.ErrorS(err, "unexpected response", "requestID", requestID, "responseBody", string(b.ResponseBody.GetBody()))
			complete = true
			return generateErrorResponse(
				envoyTypePb.StatusCode_InternalServerError,
				[]*configPb.HeaderValueOption{{Header: &configPb.HeaderValue{
					Key: "x-error-response-unknown", RawValue: []byte("true"),
				}}},
				err.Error()), complete
		}
		// Do not overwrite model, res can be empty.
		usage = res.Usage
	}

	var requestEnd string
	if usage.TotalTokens != 0 {
		complete = true
		// Update promptTokens and completeTokens
		promptTokens = usage.PromptTokens
		completionTokens = usage.CompletionTokens
		// Count token per user.
		if user.Name != "" {
			tpm, err := s.ratelimiter.Incr(ctx, fmt.Sprintf("%v_TPM_CURRENT", user), res.Usage.TotalTokens)
			if err != nil {
				return generateErrorResponse(

Review comment (suggestion):
This check is used to handle error messages from the vLLM model. I assumed "model" from request header should equal to "res.Model" returned from the response. If  that is not the case, replacing with len(res.Model) == 0 is ok.

Code after:
			if len(evt.Choices) == 0 {
				// Do not overwrite model, res can be empty.
				usage = evt.Usage
			}
		}
		if err := streaming.Err(); err != nil {
			klog.ErrorS(err, "error to unmarshal response", "requestID", requestID, "responseBody", string(b.ResponseBody.GetBody()))
			complete = true
			return generateErrorResponse(
				envoyTypePb.StatusCode_InternalServerError,
				[]*configPb.HeaderValueOption{{Header: &configPb.HeaderValue{
					Key: "x-streaming-error", RawValue: []byte("true"),
				}}},
				err.Error()), complete
		}
	} else {
		if err := json.Unmarshal(b.ResponseBody.Body, &res); err != nil {
			klog.ErrorS(err, "error to unmarshal response", "requestID", requestID, "responseBody", string(b.ResponseBody.GetBody()))
			complete = true
			return generateErrorResponse(
				envoyTypePb.StatusCode_InternalServerError,
				[]*configPb.HeaderValueOption{{Header: &configPb.HeaderValue{
					Key: "x-error-response-unmarshal", RawValue: []byte("true"),
				}}},
				err.Error()), complete
		}
		// } else if res.Model != model {
		// 	err = ErrorUnknownResponse
		// 	klog.ErrorS(err, "unexpected response", "requestID", requestID, "responseBody", string(b.ResponseBody.GetBody()))
		// 	complete = true
		// 	return generateErrorResponse(
		// 		envoyTypePb.StatusCode_InternalServerError,
		// 		[]*configPb.HeaderValueOption{{Header: &configPb.HeaderValue{
		// 			Key: "x-error-response-unknown", RawValue: []byte("true"),
		// 		}}},
		// 		err.Error()), complete
		// }
		// Do not overwrite model, res can be empty.
		usage = res.Usage
	}

	var requestEnd string
	if usage.TotalTokens != 0 {
		complete = true
		// Update promptTokens and completeTokens
		promptTokens = usage.PromptTokens
		completionTokens = usage.CompletionTokens
		// Count token per user.
		if user.Name != "" {
			tpm, err := s.ratelimiter.Incr(ctx, fmt.Sprintf("%v_TPM_CURRENT", user), res.Usage.TotalTokens)
			if err != nil {