-
Notifications
You must be signed in to change notification settings - Fork 255
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Make stream include usage as optional
Signed-off-by: Varun Gupta <[email protected]>
- Loading branch information
1 parent
31e3cf9
commit 9223769
Showing
4 changed files
with
151 additions
and
105 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,122 @@ | ||
| package gateway | ||
|
|
||
| import ( | ||
| "encoding/json" | ||
| "slices" | ||
| "strings" | ||
|
|
||
| configPb "github.com/envoyproxy/go-control-plane/envoy/config/core/v3" | ||
| extProcPb "github.com/envoyproxy/go-control-plane/envoy/service/ext_proc/v3" | ||
| envoyTypePb "github.com/envoyproxy/go-control-plane/envoy/type/v3" | ||
| "github.com/vllm-project/aibrix/pkg/utils" | ||
| "k8s.io/klog/v2" | ||
| ) | ||
|
|
||
| // validateStreamOptions validates whether stream options to include usage is set for user request | ||
| func validateStreamOptions(requestID string, user utils.User, jsonMap map[string]interface{}) *extProcPb.ProcessingResponse { | ||
| if user.Tpm != 0 { | ||
| streamOptions, ok := jsonMap["stream_options"].(map[string]interface{}) | ||
| if !ok { | ||
| klog.ErrorS(nil, "no stream option available", "requestID", requestID, "jsonMap", jsonMap) | ||
| return generateErrorResponse(envoyTypePb.StatusCode_InternalServerError, | ||
| []*configPb.HeaderValueOption{{Header: &configPb.HeaderValue{ | ||
| Key: HeaderErrorNoStreamOptions, RawValue: []byte("stream options not set")}}}, | ||
| "no stream option available") | ||
| } | ||
| includeUsage, ok := streamOptions["include_usage"].(bool) | ||
| if !includeUsage || !ok { | ||
| klog.ErrorS(nil, "no stream with usage option available", "requestID", requestID, "jsonMap", jsonMap) | ||
| return generateErrorResponse(envoyTypePb.StatusCode_InternalServerError, | ||
| []*configPb.HeaderValueOption{{Header: &configPb.HeaderValue{ | ||
| Key: HeaderErrorStreamOptionsIncludeUsage, RawValue: []byte("include usage for stream options not set")}}}, | ||
| "no stream with usage option available") | ||
| } | ||
| } | ||
| return nil | ||
| } | ||
|
|
||
| // getRoutingStrategy retrieves the routing strategy from the headers or environment variable | ||
| // It returns the routing strategy value and whether custom routing strategy is enabled. | ||
| func getRoutingStrategy(headers []*configPb.HeaderValue) (string, bool) { | ||
| var routingStrategy string | ||
| routingStrategyEnabled := false | ||
|
|
||
| // Check headers for routing strategy | ||
| for _, header := range headers { | ||
| if strings.ToLower(header.Key) == HeaderRoutingStrategy { | ||
| routingStrategy = string(header.RawValue) | ||
| routingStrategyEnabled = true | ||
| break // Prioritize header value over environment variable | ||
| } | ||
| } | ||
|
|
||
| // If header not set, check environment variable | ||
| if !routingStrategyEnabled { | ||
| if value, exists := utils.CheckEnvExists(EnvRoutingAlgorithm); exists { | ||
| routingStrategy = value | ||
| routingStrategyEnabled = true | ||
| } | ||
| } | ||
|
|
||
| return routingStrategy, routingStrategyEnabled | ||
| } | ||
|
|
||
| // validateRoutingStrategy validates if user provided routing strategy is supported by gateway | ||
| func validateRoutingStrategy(routingStrategy string) bool { | ||
| routingStrategy = strings.TrimSpace(routingStrategy) | ||
| return slices.Contains(routingStrategies, routingStrategy) | ||
| } | ||
|
|
||
| // getRequestMessage returns input request message field which has user prompt | ||
| func getRequestMessage(jsonMap map[string]interface{}) (string, *extProcPb.ProcessingResponse) { | ||
| messages, ok := jsonMap["messages"] | ||
| if !ok { | ||
| return "", generateErrorResponse(envoyTypePb.StatusCode_InternalServerError, | ||
| []*configPb.HeaderValueOption{{Header: &configPb.HeaderValue{Key: HeaderErrorRequestBodyProcessing, RawValue: []byte("true")}}}, | ||
| "no messages in the request body") | ||
| } | ||
| messagesJSON, err := json.Marshal(messages) | ||
| if err != nil { | ||
| return "", generateErrorResponse(envoyTypePb.StatusCode_InternalServerError, | ||
| []*configPb.HeaderValueOption{{Header: &configPb.HeaderValue{Key: HeaderErrorRequestBodyProcessing, RawValue: []byte("true")}}}, | ||
| "unable to marshal messages from request body") | ||
| } | ||
| return string(messagesJSON), nil | ||
| } | ||
|
|
||
| // generateErrorResponse construct envoy proxy error response | ||
| func generateErrorResponse(statusCode envoyTypePb.StatusCode, headers []*configPb.HeaderValueOption, body string) *extProcPb.ProcessingResponse { | ||
| // Set the Content-Type header to application/json | ||
| headers = append(headers, &configPb.HeaderValueOption{ | ||
| Header: &configPb.HeaderValue{ | ||
| Key: "Content-Type", | ||
| Value: "application/json", | ||
| }, | ||
| }) | ||
|
|
||
| return &extProcPb.ProcessingResponse{ | ||
| Response: &extProcPb.ProcessingResponse_ImmediateResponse{ | ||
| ImmediateResponse: &extProcPb.ImmediateResponse{ | ||
| Status: &envoyTypePb.HttpStatus{ | ||
| Code: statusCode, | ||
| }, | ||
| Headers: &extProcPb.HeaderMutation{ | ||
| SetHeaders: headers, | ||
| }, | ||
| Body: generateErrorMessage(body, int(statusCode)), | ||
| }, | ||
| }, | ||
| } | ||
| } | ||
|
|
||
| // generateErrorMessage constructs a JSON error message | ||
| func generateErrorMessage(message string, code int) string { | ||
| errorStruct := map[string]interface{}{ | ||
| "error": map[string]interface{}{ | ||
| "message": message, | ||
| "code": code, | ||
| }, | ||
| } | ||
| jsonData, _ := json.Marshal(errorStruct) | ||
| return string(jsonData) | ||
| } |