diff --git a/controller/relay.go b/controller/relay.go index 10dfd502fbd0..a91c0deadf66 100644 --- a/controller/relay.go +++ b/controller/relay.go @@ -1,11 +1,13 @@ package controller import ( + "encoding/json" "errors" "fmt" "io" "log" "net/http" + "strconv" "strings" "time" @@ -22,6 +24,7 @@ import ( "github.com/QuantumNous/new-api/service" "github.com/QuantumNous/new-api/setting" "github.com/QuantumNous/new-api/setting/operation_setting" + "github.com/QuantumNous/new-api/setting/ratio_setting" "github.com/QuantumNous/new-api/types" "github.com/bytedance/gopkg/util/gopool" @@ -69,13 +72,15 @@ func Relay(c *gin.Context, relayFormat types.RelayFormat) { requestId := c.GetString(common.RequestIdKey) //group := common.GetContextKeyString(c, constant.ContextKeyUsingGroup) //originalModel := common.GetContextKeyString(c, constant.ContextKeyOriginalModel) + responsesWs := relayFormat == types.RelayFormatOpenAIResponses && websocket.IsWebSocketUpgrade(c.Request) var ( newAPIError *types.NewAPIError + request dto.Request ws *websocket.Conn ) - if relayFormat == types.RelayFormatOpenAIRealtime { + if relayFormat == types.RelayFormatOpenAIRealtime || responsesWs { var err error ws, err = upgrader.Upgrade(c.Writer, c.Request, nil) if err != nil { @@ -83,6 +88,22 @@ func Relay(c *gin.Context, relayFormat types.RelayFormat) { return } defer ws.Close() + + if responsesWs { + request, err = readFirstResponsesWSRequest(ws) + if err != nil { + helper.WssError(c, ws, types.NewError(err, types.ErrorCodeInvalidRequest, types.ErrOptionWithSkipRetry()).ToOpenAIError()) + return + } + if responsesReq, ok := request.(*dto.OpenAIResponsesRequest); ok { + c.Set("first_wss_request", responsesReq) + newAPIError = setupResponsesWSChannel(c, responsesReq.Model) + if newAPIError != nil { + helper.WssError(c, ws, newAPIError.ToOpenAIError()) + return + } + } + } } defer func() { @@ -92,6 +113,14 @@ func Relay(c *gin.Context, relayFormat types.RelayFormat) { switch relayFormat { case types.RelayFormatOpenAIRealtime: helper.WssError(c, ws, newAPIError.ToOpenAIError()) + case types.RelayFormatOpenAIResponses: + if responsesWs { + helper.WssError(c, ws, newAPIError.ToOpenAIError()) + return + } + c.JSON(newAPIError.StatusCode, gin.H{ + "error": newAPIError.ToOpenAIError(), + }) case types.RelayFormatClaude: c.JSON(newAPIError.StatusCode, gin.H{ "type": "error", @@ -105,15 +134,18 @@ func Relay(c *gin.Context, relayFormat types.RelayFormat) { } }() - request, err := helper.GetAndValidateRequest(c, relayFormat) - if err != nil { - // Map "request body too large" to 413 so clients can handle it correctly - if common.IsRequestBodyTooLargeError(err) || errors.Is(err, common.ErrRequestBodyTooLarge) { - newAPIError = types.NewErrorWithStatusCode(err, types.ErrorCodeReadRequestBodyFailed, http.StatusRequestEntityTooLarge, types.ErrOptionWithSkipRetry()) - } else { - newAPIError = types.NewError(err, types.ErrorCodeInvalidRequest) + if request == nil { + var err error + request, err = helper.GetAndValidateRequest(c, relayFormat) + if err != nil { + // Map "request body too large" to 413 so clients can handle it correctly + if common.IsRequestBodyTooLargeError(err) || errors.Is(err, common.ErrRequestBodyTooLarge) { + newAPIError = types.NewErrorWithStatusCode(err, types.ErrorCodeReadRequestBodyFailed, http.StatusRequestEntityTooLarge, types.ErrOptionWithSkipRetry()) + } else { + newAPIError = types.NewError(err, types.ErrorCodeInvalidRequest) + } + return } - return } relayInfo, err := relaycommon.GenRelayInfo(c, relayFormat, request, ws) @@ -121,6 +153,9 @@ func Relay(c *gin.Context, relayFormat types.RelayFormat) { newAPIError = types.NewError(err, types.ErrorCodeGenRelayInfoFailed) return } + if responsesWs { + relayInfo.ClientWs = ws + } needSensitiveCheck := setting.ShouldCheckPromptSensitive() needCountToken := constant.CountToken @@ -211,6 +246,12 @@ func Relay(c *gin.Context, relayFormat types.RelayFormat) { switch relayFormat { case types.RelayFormatOpenAIRealtime: newAPIError = relay.WssHelper(c, relayInfo) + case types.RelayFormatOpenAIResponses: + if responsesWs { + newAPIError = relay.WssResponsesHelper(c, relayInfo) + break + } + newAPIError = relayHandler(c, relayInfo) case types.RelayFormatClaude: newAPIError = relay.ClaudeHelper(c, relayInfo) case types.RelayFormatGemini: @@ -241,6 +282,92 @@ func Relay(c *gin.Context, relayFormat types.RelayFormat) { } } +func readFirstResponsesWSRequest(ws *websocket.Conn) (dto.Request, error) { + if ws == nil { + return nil, errors.New("websocket connection is nil") + } + _, payload, err := ws.ReadMessage() + if err != nil { + return nil, fmt.Errorf("read first websocket message failed: %w", err) + } + request := &dto.OpenAIResponsesRequest{} + if err := json.Unmarshal(payload, request); err != nil { + return nil, fmt.Errorf("parse first websocket message failed: %w", err) + } + if request.Model == "" { + return nil, errors.New("model is required") + } + if request.Input == nil { + return nil, errors.New("input is required") + } + return request, nil +} + +func setupResponsesWSChannel(c *gin.Context, modelName string) *types.NewAPIError { + if modelName == "" { + return types.NewError(errors.New("model is required"), types.ErrorCodeInvalidRequest, types.ErrOptionWithSkipRetry()) + } + if common.GetContextKeyInt(c, constant.ContextKeyChannelId) != 0 && + common.GetContextKeyString(c, constant.ContextKeyOriginalModel) != "" { + return nil + } + + modelLimitEnable := common.GetContextKeyBool(c, constant.ContextKeyTokenModelLimitEnabled) + if modelLimitEnable { + s, ok := common.GetContextKey(c, constant.ContextKeyTokenModelLimit) + if !ok { + return types.NewErrorWithStatusCode(errors.New("token has no model access"), types.ErrorCodeInvalidRequest, http.StatusForbidden, types.ErrOptionWithSkipRetry()) + } + tokenModelLimit, ok := s.(map[string]bool) + if !ok { + tokenModelLimit = map[string]bool{} + } + matchName := ratio_setting.FormatMatchingModelName(modelName) + if _, ok := tokenModelLimit[matchName]; !ok { + return types.NewErrorWithStatusCode(fmt.Errorf("model %s is forbidden for this token", modelName), types.ErrorCodeInvalidRequest, http.StatusForbidden, types.ErrOptionWithSkipRetry()) + } + } + + var selectedChannel *model.Channel + if channelId, ok := common.GetContextKey(c, constant.ContextKeyTokenSpecificChannelId); ok { + id, err := strconv.Atoi(channelId.(string)) + if err != nil { + return types.NewErrorWithStatusCode(err, types.ErrorCodeGetChannelFailed, http.StatusBadRequest, types.ErrOptionWithSkipRetry()) + } + channel, err := model.GetChannelById(id, true) + if err != nil { + return types.NewErrorWithStatusCode(err, types.ErrorCodeGetChannelFailed, http.StatusBadRequest, types.ErrOptionWithSkipRetry()) + } + if channel.Status != common.ChannelStatusEnabled { + return types.NewErrorWithStatusCode(errors.New("channel is disabled"), types.ErrorCodeGetChannelFailed, http.StatusForbidden, types.ErrOptionWithSkipRetry()) + } + selectedChannel = channel + } else { + usingGroup := common.GetContextKeyString(c, constant.ContextKeyUsingGroup) + if usingGroup == "" { + usingGroup = common.GetContextKeyString(c, constant.ContextKeyUserGroup) + } + channel, selectGroup, err := service.CacheGetRandomSatisfiedChannel(&service.RetryParam{ + Ctx: c, + ModelName: modelName, + TokenGroup: usingGroup, + Retry: common.GetPointer(0), + }) + if err != nil { + return types.NewErrorWithStatusCode(err, types.ErrorCodeGetChannelFailed, http.StatusServiceUnavailable, types.ErrOptionWithSkipRetry()) + } + if channel == nil { + return types.NewErrorWithStatusCode(fmt.Errorf("no available channel for model %s", modelName), types.ErrorCodeGetChannelFailed, http.StatusServiceUnavailable, types.ErrOptionWithSkipRetry()) + } + if usingGroup == "auto" && selectGroup != "" { + common.SetContextKey(c, constant.ContextKeyAutoGroup, selectGroup) + } + selectedChannel = channel + } + + return middleware.SetupContextForSelectedChannel(c, selectedChannel, modelName) +} + var upgrader = websocket.Upgrader{ Subprotocols: []string{"realtime"}, // WS 握手支持的协议,如果有使用 Sec-WebSocket-Protocol,则必须在此声明对应的 Protocol TODO add other protocol CheckOrigin: func(r *http.Request) bool { diff --git a/dto/openai_request.go b/dto/openai_request.go index a6fc3f66bb22..0e9b035232fa 100644 --- a/dto/openai_request.go +++ b/dto/openai_request.go @@ -812,6 +812,7 @@ type WebSearchOptions struct { // https://platform.openai.com/docs/api-reference/responses/create type OpenAIResponsesRequest struct { + Type string `json:"type,omitempty"` Model string `json:"model"` Input json.RawMessage `json:"input,omitempty"` Include json.RawMessage `json:"include,omitempty"` diff --git a/middleware/distributor.go b/middleware/distributor.go index db57998caaf4..add7911829e0 100644 --- a/middleware/distributor.go +++ b/middleware/distributor.go @@ -20,6 +20,7 @@ import ( "github.com/QuantumNous/new-api/types" "github.com/gin-gonic/gin" + "github.com/gorilla/websocket" ) type ModelRequest struct { @@ -177,6 +178,15 @@ func getModelRequest(c *gin.Context) (*ModelRequest, bool, error) { var modelRequest ModelRequest shouldSelectChannel := true var err error + if c.Request.Method == http.MethodGet && + websocket.IsWebSocketUpgrade(c.Request) && + strings.HasPrefix(c.Request.URL.Path, "/v1/responses") { + modelRequest.Model = c.Query("model") + if strings.HasPrefix(c.Request.URL.Path, "/v1/responses/compact") && modelRequest.Model != "" { + modelRequest.Model = ratio_setting.WithCompactModelSuffix(modelRequest.Model) + } + return &modelRequest, false, nil + } if strings.Contains(c.Request.URL.Path, "/mj/") { relayMode := relayconstant.Path2RelayModeMidjourney(c.Request.URL.Path) if relayMode == relayconstant.RelayModeMidjourneyTaskFetch || diff --git a/relay/channel/api_request.go b/relay/channel/api_request.go index 8dfb61d40093..411341680861 100644 --- a/relay/channel/api_request.go +++ b/relay/channel/api_request.go @@ -356,6 +356,11 @@ func DoWssRequest(a Adaptor, c *gin.Context, info *common.RelayInfo, requestBody if err != nil { return nil, fmt.Errorf("get request url failed: %w", err) } + if strings.HasPrefix(fullRequestURL, "https://") { + fullRequestURL = "wss://" + strings.TrimPrefix(fullRequestURL, "https://") + } else if strings.HasPrefix(fullRequestURL, "http://") { + fullRequestURL = "ws://" + strings.TrimPrefix(fullRequestURL, "http://") + } targetHeader := http.Header{} err = a.SetupRequestHeader(c, &targetHeader, info) if err != nil { diff --git a/relay/websocket.go b/relay/websocket.go index 57a51895b006..4bdbfd78b16e 100644 --- a/relay/websocket.go +++ b/relay/websocket.go @@ -1,13 +1,19 @@ package relay import ( + "errors" "fmt" + "time" "github.com/QuantumNous/new-api/dto" + "github.com/QuantumNous/new-api/logger" + "github.com/QuantumNous/new-api/relay/channel" relaycommon "github.com/QuantumNous/new-api/relay/common" + "github.com/QuantumNous/new-api/relay/helper" "github.com/QuantumNous/new-api/service" "github.com/QuantumNous/new-api/types" + "github.com/bytedance/gopkg/util/gopool" "github.com/gin-gonic/gin" "github.com/gorilla/websocket" ) @@ -44,3 +50,103 @@ func WssHelper(c *gin.Context, info *relaycommon.RelayInfo) (newAPIError *types. service.PostWssConsumeQuota(c, info, info.UpstreamModelName, usage.(*dto.RealtimeUsage), "") return nil } + +func WssResponsesHelper(c *gin.Context, info *relaycommon.RelayInfo) (newAPIError *types.NewAPIError) { + info.InitChannelMeta(c) + if info.ClientWs == nil { + return types.NewError(errors.New("client websocket connection is nil"), types.ErrorCodeBadResponse, types.ErrOptionWithSkipRetry()) + } + + adaptor := GetAdaptor(info.ApiType) + if adaptor == nil { + return types.NewError(fmt.Errorf("invalid api type: %d", info.ApiType), types.ErrorCodeInvalidApiType, types.ErrOptionWithSkipRetry()) + } + adaptor.Init(info) + + targetWs, err := channel.DoWssRequest(adaptor, c, info, nil) + if err != nil { + return types.NewError(err, types.ErrorCodeDoRequestFailed) + } + info.TargetWs = targetWs + defer info.TargetWs.Close() + + if err := sendInitialResponsesWSRequest(c, adaptor, info); err != nil { + return types.NewError(err, types.ErrorCodeDoRequestFailed, types.ErrOptionWithSkipRetry()) + } + if err := proxyResponsesWS(c, info.ClientWs, info.TargetWs); err != nil { + return types.NewError(err, types.ErrorCodeBadResponse, types.ErrOptionWithSkipRetry()) + } + if err := service.SettleBilling(c, info, info.FinalPreConsumedQuota); err != nil { + logger.LogError(c, "responses websocket settle billing failed: "+err.Error()) + } + return nil +} + +func sendInitialResponsesWSRequest(c *gin.Context, adaptor channel.Adaptor, info *relaycommon.RelayInfo) error { + request, ok := info.Request.(*dto.OpenAIResponsesRequest) + if !ok || request == nil { + return errors.New("invalid responses websocket request") + } + initialRequest := *request + if initialRequest.Type == "" { + initialRequest.Type = "response.create" + } + converted, err := adaptor.ConvertOpenAIResponsesRequest(c, info, initialRequest) + if err != nil { + return err + } + switch v := converted.(type) { + case string: + return helper.WssString(c, info.TargetWs, v) + case []byte: + return info.TargetWs.WriteMessage(websocket.TextMessage, v) + default: + return helper.WssObject(c, info.TargetWs, v) + } +} + +func proxyResponsesWS(c *gin.Context, clientConn, targetConn *websocket.Conn) error { + errChan := make(chan error, 2) + forward := func(src, dst *websocket.Conn, direction string) { + defer func() { + if r := recover(); r != nil { + errChan <- fmt.Errorf("%s panic: %v", direction, r) + } + }() + for { + messageType, payload, err := src.ReadMessage() + if err != nil { + if websocket.IsCloseError(err, websocket.CloseNormalClosure, websocket.CloseGoingAway) { + errChan <- nil + return + } + errChan <- fmt.Errorf("%s read failed: %w", direction, err) + return + } + if err := dst.WriteMessage(messageType, payload); err != nil { + errChan <- fmt.Errorf("%s write failed: %w", direction, err) + return + } + } + } + + gopool.Go(func() { + forward(clientConn, targetConn, "client->target") + }) + gopool.Go(func() { + forward(targetConn, clientConn, "target->client") + }) + + select { + case <-c.Request.Context().Done(): + return nil + case err := <-errChan: + if err == nil { + deadline := time.Now().Add(500 * time.Millisecond) + _ = clientConn.WriteControl(websocket.CloseMessage, websocket.FormatCloseMessage(websocket.CloseNormalClosure, ""), deadline) + _ = targetConn.WriteControl(websocket.CloseMessage, websocket.FormatCloseMessage(websocket.CloseNormalClosure, ""), deadline) + return nil + } + return err + } +} diff --git a/router/relay-router.go b/router/relay-router.go index 17a13cad7fd6..dc17bf5a1546 100644 --- a/router/relay-router.go +++ b/router/relay-router.go @@ -78,6 +78,9 @@ func SetRelayRouter(router *gin.Engine) { wsRouter.GET("/realtime", func(c *gin.Context) { controller.Relay(c, types.RelayFormatOpenAIRealtime) }) + wsRouter.GET("/responses", func(c *gin.Context) { + controller.Relay(c, types.RelayFormatOpenAIResponses) + }) } { //http router