From b41e1c4eb9812051816bfeea44d465b31987d9c1 Mon Sep 17 00:00:00 2001 From: Saika Date: Fri, 17 Jul 2026 17:10:35 +0800 Subject: [PATCH] fix: support codex alpha search relay --- controller/relay.go | 2 + dto/codex_alpha_search.go | 45 ++++++ middleware/distributor.go | 4 + relay/alpha_search_handler.go | 154 +++++++++++++++++++++ relay/channel/codex/adaptor.go | 47 ++++++- relay/channel/openai/adaptor.go | 18 +++ relay/channel/openai/relay_alpha_search.go | 26 ++++ relay/common/relay_info.go | 12 ++ relay/common/request_conversion.go | 2 + relay/constant/relay_mode.go | 6 + relay/helper/valid_request.go | 13 ++ router/relay-router.go | 17 +++ router/web-router.go | 6 +- service/tool_billing.go | 123 +++++++++++++++- setting/operation_setting/tools.go | 1 + types/relay_format.go | 1 + 16 files changed, 463 insertions(+), 14 deletions(-) create mode 100644 dto/codex_alpha_search.go create mode 100644 relay/alpha_search_handler.go create mode 100644 relay/channel/openai/relay_alpha_search.go diff --git a/controller/relay.go b/controller/relay.go index 6e91ccb6050..fea4133c7df 100644 --- a/controller/relay.go +++ b/controller/relay.go @@ -49,6 +49,8 @@ func relayHandler(c *gin.Context, info *relaycommon.RelayInfo) *types.NewAPIErro err = relay.EmbeddingHelper(c, info) case relayconstant.RelayModeResponses, relayconstant.RelayModeResponsesCompact: err = relay.ResponsesHelper(c, info) + case relayconstant.RelayModeCodexAlphaSearch: + err = relay.AlphaSearchHelper(c, info) default: err = relay.TextHelper(c, info) } diff --git a/dto/codex_alpha_search.go b/dto/codex_alpha_search.go new file mode 100644 index 00000000000..59f3f79b53e --- /dev/null +++ b/dto/codex_alpha_search.go @@ -0,0 +1,45 @@ +package dto + +import ( + "encoding/json" + "strings" + + "github.com/QuantumNous/new-api/types" + "github.com/gin-gonic/gin" +) + +// CodexAlphaSearchRequest is the standalone search request used by Codex CLI. +// It is not an OpenAI Responses request; unknown fields must be preserved by +// the relay path and forwarded as-is. +type CodexAlphaSearchRequest struct { + Model string `json:"model,omitempty"` + ID string `json:"id,omitempty"` + Query string `json:"query,omitempty"` + Commands json.RawMessage `json:"commands,omitempty"` + Input json.RawMessage `json:"input,omitempty"` + Settings json.RawMessage `json:"settings,omitempty"` +} + +func (r *CodexAlphaSearchRequest) GetTokenCountMeta() *types.TokenCountMeta { + parts := make([]string, 0, 4) + if strings.TrimSpace(r.Query) != "" { + parts = append(parts, r.Query) + } + for _, raw := range []json.RawMessage{r.Commands, r.Input, r.Settings} { + if len(raw) > 0 { + parts = append(parts, string(raw)) + } + } + return &types.TokenCountMeta{ + CombineText: strings.Join(parts, "\n"), + TokenType: types.TokenTypeTokenizer, + } +} + +func (r *CodexAlphaSearchRequest) IsStream(c *gin.Context) bool { + return false +} + +func (r *CodexAlphaSearchRequest) SetModelName(modelName string) { + r.Model = modelName +} diff --git a/middleware/distributor.go b/middleware/distributor.go index 4234011c9f7..03bb19d9cb8 100644 --- a/middleware/distributor.go +++ b/middleware/distributor.go @@ -409,6 +409,10 @@ func getModelRequest(c *gin.Context) (*ModelRequest, bool, error) { common.SetContextKey(c, constant.ContextKeyTokenGroup, modelRequest.Group) } + if relayconstant.Path2RelayMode(c.Request.URL.Path) == relayconstant.RelayModeCodexAlphaSearch { + c.Set("relay_mode", relayconstant.RelayModeCodexAlphaSearch) + } + if strings.HasPrefix(c.Request.URL.Path, "/v1/responses/compact") && modelRequest.Model != "" { modelRequest.Model = ratio_setting.WithCompactModelSuffix(modelRequest.Model) } diff --git a/relay/alpha_search_handler.go b/relay/alpha_search_handler.go new file mode 100644 index 00000000000..f5c498e7c7c --- /dev/null +++ b/relay/alpha_search_handler.go @@ -0,0 +1,154 @@ +package relay + +import ( + "encoding/json" + "fmt" + "net/http" + "strings" + + "github.com/QuantumNous/new-api/common" + "github.com/QuantumNous/new-api/constant" + "github.com/QuantumNous/new-api/dto" + 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/gin-gonic/gin" +) + +var codexAlphaSearchUnsupportedBodyFields = [...]string{ + // Codex alpha/search is a standalone SearchRequest protocol, not a + // /v1/responses subrequest. Codex clients and proxy chains may carry these + // Responses-only fields; ChatGPT alpha/search rejects them as unknown. + "prompt_cache_key", + "prompt_cache_retention", +} + +func buildCodexAlphaSearchBody(body []byte, model string) ([]byte, error) { + if len(body) == 0 { + return body, nil + } + var obj map[string]json.RawMessage + if err := common.Unmarshal(body, &obj); err != nil || obj == nil { + return body, err + } + + changed := false + for _, field := range codexAlphaSearchUnsupportedBodyFields { + if _, ok := obj[field]; ok { + delete(obj, field) + changed = true + } + } + + model = strings.TrimSpace(model) + if _, hasModel := obj["model"]; hasModel && model != "" { + marshaledModel, err := common.Marshal(model) + if err != nil { + return nil, err + } + if string(obj["model"]) != string(marshaledModel) { + obj["model"] = marshaledModel + changed = true + } + } + + if !changed { + return body, nil + } + out, err := common.Marshal(obj) + if err != nil { + return nil, err + } + return out, nil +} + +func alphaSearchRelaySupported(apiType int) bool { + switch apiType { + case constant.APITypeOpenAI, + constant.APITypeOpenRouter, + constant.APITypeXinference, + constant.APITypeCodex, + constant.APITypeAdvancedCustom: + return true + default: + return false + } +} + +func AlphaSearchHelper(c *gin.Context, info *relaycommon.RelayInfo) (newAPIError *types.NewAPIError) { + info.InitChannelMeta(c) + if !alphaSearchRelaySupported(info.ApiType) { + return types.NewErrorWithStatusCode( + fmt.Errorf("unsupported endpoint %q for api type %d", "/v1/alpha/search", info.ApiType), + types.ErrorCodeInvalidRequest, + http.StatusBadRequest, + types.ErrOptionWithSkipRetry(), + ) + } + + request, ok := info.Request.(*dto.CodexAlphaSearchRequest) + if !ok { + return types.NewErrorWithStatusCode( + fmt.Errorf("invalid request type, expected dto.CodexAlphaSearchRequest, got %T", info.Request), + types.ErrorCodeInvalidRequest, + http.StatusBadRequest, + types.ErrOptionWithSkipRetry(), + ) + } + requestCopy := *request + if err := helper.ModelMappedHelper(c, info, &requestCopy); err != nil { + return types.NewError(err, types.ErrorCodeChannelModelMappedError, 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) + + storage, err := common.GetBodyStorage(c) + if err != nil { + return types.NewErrorWithStatusCode(err, types.ErrorCodeReadRequestBodyFailed, http.StatusBadRequest, types.ErrOptionWithSkipRetry()) + } + bodyBytes, err := storage.Bytes() + if err != nil { + return types.NewErrorWithStatusCode(err, types.ErrorCodeReadRequestBodyFailed, http.StatusBadRequest, types.ErrOptionWithSkipRetry()) + } + outboundBytes, err := buildCodexAlphaSearchBody(bodyBytes, requestCopy.Model) + if err != nil { + return types.NewErrorWithStatusCode(err, types.ErrorCodeInvalidRequest, http.StatusBadRequest, types.ErrOptionWithSkipRetry()) + } + + requestBody, size, closer, err := relaycommon.NewOutboundJSONBody(outboundBytes) + if err != nil { + return types.NewError(err, types.ErrorCodeConvertRequestFailed, types.ErrOptionWithSkipRetry()) + } + defer closer.Close() + info.UpstreamRequestBodySize = size + + resp, err := adaptor.DoRequest(c, info, requestBody) + if err != nil { + return types.NewOpenAIError(err, types.ErrorCodeDoRequestFailed, http.StatusInternalServerError) + } + httpResp, ok := resp.(*http.Response) + if !ok || httpResp == nil { + return types.NewOpenAIError(fmt.Errorf("invalid alpha search response type: %T", resp), types.ErrorCodeBadResponse, http.StatusInternalServerError) + } + + statusCodeMappingStr := c.GetString("status_code_mapping") + if httpResp.StatusCode < http.StatusOK || httpResp.StatusCode >= http.StatusMultipleChoices { + newAPIError = service.RelayErrorHandler(c.Request.Context(), httpResp, false) + service.ResetStatusCode(newAPIError, statusCodeMappingStr) + return newAPIError + } + + if _, newAPIError = adaptor.DoResponse(c, httpResp, info); newAPIError != nil { + service.ResetStatusCode(newAPIError, statusCodeMappingStr) + return newAPIError + } + + service.PostCodexAlphaSearchConsumeQuota(c, info) + return nil +} diff --git a/relay/channel/codex/adaptor.go b/relay/channel/codex/adaptor.go index ef4d4fa0412..2a92ae64a75 100644 --- a/relay/channel/codex/adaptor.go +++ b/relay/channel/codex/adaptor.go @@ -112,10 +112,16 @@ func (a *Adaptor) DoRequest(c *gin.Context, info *relaycommon.RelayInfo, request } func (a *Adaptor) DoResponse(c *gin.Context, resp *http.Response, info *relaycommon.RelayInfo) (usage any, err *types.NewAPIError) { - if info.RelayMode != relayconstant.RelayModeResponses && info.RelayMode != relayconstant.RelayModeResponsesCompact { + if info.RelayMode != relayconstant.RelayModeResponses && + info.RelayMode != relayconstant.RelayModeResponsesCompact && + info.RelayMode != relayconstant.RelayModeCodexAlphaSearch { return nil, types.NewError(errors.New("codex channel: endpoint not supported"), types.ErrorCodeInvalidRequest) } + if info.RelayMode == relayconstant.RelayModeCodexAlphaSearch { + return openai.OaiAlphaSearchHandler(c, resp) + } + if info.RelayMode == relayconstant.RelayModeResponsesCompact { return openai.OaiResponsesCompactionHandler(c, resp) } @@ -135,11 +141,18 @@ func (a *Adaptor) GetChannelName() string { } func (a *Adaptor) GetRequestURL(info *relaycommon.RelayInfo) (string, error) { - if info.RelayMode != relayconstant.RelayModeResponses && info.RelayMode != relayconstant.RelayModeResponsesCompact { - return "", errors.New("codex channel: only /v1/responses and /v1/responses/compact are supported") + if info.RelayMode != relayconstant.RelayModeResponses && + info.RelayMode != relayconstant.RelayModeResponsesCompact && + info.RelayMode != relayconstant.RelayModeCodexAlphaSearch { + return "", errors.New("codex channel: only /v1/responses, /v1/responses/compact and /v1/alpha/search are supported") } path := "/backend-api/codex/responses" - if info.RelayMode == relayconstant.RelayModeResponsesCompact { + if info.RelayMode == relayconstant.RelayModeCodexAlphaSearch { + path = "/backend-api/codex/alpha/search" + if idx := strings.Index(info.RequestURLPath, "?"); idx >= 0 { + path += info.RequestURLPath[idx:] + } + } else if info.RelayMode == relayconstant.RelayModeResponsesCompact { path = "/backend-api/codex/responses/compact" } return relaycommon.GetFullRequestURL(info.ChannelBaseUrl, path, info.ChannelType), nil @@ -171,7 +184,27 @@ func (a *Adaptor) SetupRequestHeader(c *gin.Context, req *http.Header, info *rel req.Set("Authorization", "Bearer "+accessToken) req.Set("chatgpt-account-id", accountID) - if req.Get("OpenAI-Beta") == "" { + if info.RelayMode == relayconstant.RelayModeCodexAlphaSearch { + req.Del("OpenAI-Beta") + req.Del("Session_ID") + req.Del("Conversation_ID") + req.Del("X-Codex-Beta-Features") + req.Del("X-Codex-Turn-State") + + if c != nil && c.Request != nil { + for _, name := range []string{ + "Version", + "User-Agent", + "Session_id", + "X-Client-Request-Id", + "X-Codex-Turn-Metadata", + } { + if value := strings.TrimSpace(c.GetHeader(name)); value != "" { + req.Set(name, value) + } + } + } + } else if req.Get("OpenAI-Beta") == "" { req.Set("OpenAI-Beta", "responses=experimental") } if req.Get("originator") == "" { @@ -182,7 +215,9 @@ func (a *Adaptor) SetupRequestHeader(c *gin.Context, req *http.Header, info *rel // Clients may omit it or include parameters like `application/json; charset=utf-8`, // which can be rejected by the upstream. Force the exact media type. req.Set("Content-Type", "application/json") - if info.IsStream { + if info.RelayMode == relayconstant.RelayModeCodexAlphaSearch { + req.Set("Accept", "application/json") + } else if info.IsStream { req.Set("Accept", "text/event-stream") } else if req.Get("Accept") == "" { req.Set("Accept", "application/json") diff --git a/relay/channel/openai/adaptor.go b/relay/channel/openai/adaptor.go index e118252352d..b07b4f8aba9 100644 --- a/relay/channel/openai/adaptor.go +++ b/relay/channel/openai/adaptor.go @@ -114,8 +114,18 @@ func (a *Adaptor) GetRequestURL(info *relaycommon.RelayInfo) (string, error) { info.ChannelBaseUrl = baseUrl } } + alphaSearchPath := func() string { + path := "/v1/alpha/search" + if idx := strings.Index(info.RequestURLPath, "?"); idx >= 0 { + path += info.RequestURLPath[idx:] + } + return path + } switch info.ChannelType { case constant.ChannelTypeAzure: + if info.RelayMode == relayconstant.RelayModeCodexAlphaSearch { + return "", fmt.Errorf("azure channel does not support /v1/alpha/search") + } apiVersion := info.ApiVersion if apiVersion == "" { apiVersion = constant.AzureDefaultAPIVersion @@ -167,10 +177,16 @@ func (a *Adaptor) GetRequestURL(info *relaycommon.RelayInfo) (string, error) { //case constant.ChannelTypeMiniMax: // return minimax.GetRequestURL(info) case constant.ChannelTypeCustom: + if info.RelayMode == relayconstant.RelayModeCodexAlphaSearch { + return relaycommon.GetFullRequestURL(info.ChannelBaseUrl, alphaSearchPath(), info.ChannelType), nil + } url := info.ChannelBaseUrl url = strings.Replace(url, "{model}", info.UpstreamModelName, -1) return url, nil default: + if info.RelayMode == relayconstant.RelayModeCodexAlphaSearch { + return relaycommon.GetFullRequestURL(info.ChannelBaseUrl, alphaSearchPath(), info.ChannelType), nil + } if (info.RelayFormat == types.RelayFormatClaude || info.RelayFormat == types.RelayFormatGemini) && info.RelayMode != relayconstant.RelayModeResponses && info.RelayMode != relayconstant.RelayModeResponsesCompact { @@ -651,6 +667,8 @@ func (a *Adaptor) DoResponse(c *gin.Context, resp *http.Response, info *relaycom } case relayconstant.RelayModeResponsesCompact: usage, err = OaiResponsesCompactionHandler(c, resp) + case relayconstant.RelayModeCodexAlphaSearch: + usage, err = OaiAlphaSearchHandler(c, resp) default: if info.IsStream { usage, err = OaiStreamHandler(c, info, resp) diff --git a/relay/channel/openai/relay_alpha_search.go b/relay/channel/openai/relay_alpha_search.go new file mode 100644 index 00000000000..f5eedeeceb6 --- /dev/null +++ b/relay/channel/openai/relay_alpha_search.go @@ -0,0 +1,26 @@ +package openai + +import ( + "fmt" + "io" + "net/http" + + "github.com/QuantumNous/new-api/dto" + "github.com/QuantumNous/new-api/service" + "github.com/QuantumNous/new-api/types" + "github.com/gin-gonic/gin" +) + +func OaiAlphaSearchHandler(c *gin.Context, resp *http.Response) (*dto.Usage, *types.NewAPIError) { + if resp == nil || resp.Body == nil { + return nil, types.NewOpenAIError(fmt.Errorf("invalid response"), types.ErrorCodeBadResponse, http.StatusInternalServerError) + } + defer service.CloseResponseBodyGracefully(resp) + + responseBody, err := io.ReadAll(resp.Body) + if err != nil { + return nil, types.NewOpenAIError(err, types.ErrorCodeReadResponseBodyFailed, http.StatusInternalServerError) + } + service.IOCopyBytesGracefully(c, resp, responseBody) + return &dto.Usage{}, nil +} diff --git a/relay/common/relay_info.go b/relay/common/relay_info.go index 9f460ce5c6a..fc0e3db1bea 100644 --- a/relay/common/relay_info.go +++ b/relay/common/relay_info.go @@ -416,6 +416,13 @@ func GenRelayInfoResponses(c *gin.Context, request *dto.OpenAIResponsesRequest) return info } +func GenRelayInfoCodexAlphaSearch(c *gin.Context, request *dto.CodexAlphaSearchRequest) *RelayInfo { + info := genBaseRelayInfo(c, request) + info.RelayMode = relayconstant.RelayModeCodexAlphaSearch + info.RelayFormat = types.RelayFormatCodexAlphaSearch + return info +} + func GenRelayInfoGemini(c *gin.Context, request dto.Request) *RelayInfo { info := genBaseRelayInfo(c, request) info.RelayFormat = types.RelayFormatGemini @@ -575,6 +582,11 @@ func GenRelayInfo(c *gin.Context, relayFormat types.RelayFormat, request dto.Req return GenRelayInfoResponsesCompaction(c, request), nil } return nil, errors.New("request is not a OpenAIResponsesCompactionRequest") + case types.RelayFormatCodexAlphaSearch: + if request, ok := request.(*dto.CodexAlphaSearchRequest); ok { + return GenRelayInfoCodexAlphaSearch(c, request), nil + } + return nil, errors.New("request is not a CodexAlphaSearchRequest") case types.RelayFormatTask: info = genBaseRelayInfo(c, nil) info.TaskRelayInfo = &TaskRelayInfo{} diff --git a/relay/common/request_conversion.go b/relay/common/request_conversion.go index 96b728d217d..e9ce55d3620 100644 --- a/relay/common/request_conversion.go +++ b/relay/common/request_conversion.go @@ -11,6 +11,8 @@ func GuessRelayFormatFromRequest(req any) (types.RelayFormat, bool) { return types.RelayFormatOpenAI, true case *dto.OpenAIResponsesRequest, dto.OpenAIResponsesRequest: return types.RelayFormatOpenAIResponses, true + case *dto.CodexAlphaSearchRequest, dto.CodexAlphaSearchRequest: + return types.RelayFormatCodexAlphaSearch, true case *dto.ClaudeRequest, dto.ClaudeRequest: return types.RelayFormatClaude, true case *dto.GeminiChatRequest, dto.GeminiChatRequest: diff --git a/relay/constant/relay_mode.go b/relay/constant/relay_mode.go index 25671567921..e5e22292405 100644 --- a/relay/constant/relay_mode.go +++ b/relay/constant/relay_mode.go @@ -52,6 +52,8 @@ const ( RelayModeGemini RelayModeResponsesCompact + + RelayModeCodexAlphaSearch ) func Path2RelayMode(path string) int { @@ -72,6 +74,10 @@ func Path2RelayMode(path string) int { relayMode = RelayModeImagesEdits } else if strings.HasPrefix(path, "/v1/edits") { relayMode = RelayModeEdits + } else if strings.HasPrefix(path, "/v1/alpha/search") || + strings.HasPrefix(path, "/alpha/search") || + strings.HasPrefix(path, "/backend-api/codex/alpha/search") { + relayMode = RelayModeCodexAlphaSearch } else if strings.HasPrefix(path, "/v1/responses/compact") { relayMode = RelayModeResponsesCompact } else if strings.HasPrefix(path, "/v1/responses") { diff --git a/relay/helper/valid_request.go b/relay/helper/valid_request.go index afbd59ff968..3f575f2574d 100644 --- a/relay/helper/valid_request.go +++ b/relay/helper/valid_request.go @@ -38,6 +38,8 @@ func GetAndValidateRequest(c *gin.Context, format types.RelayFormat) (request dt request, err = GetAndValidateResponsesRequest(c) case types.RelayFormatOpenAIResponsesCompaction: request, err = GetAndValidateResponsesCompactionRequest(c) + case types.RelayFormatCodexAlphaSearch: + request, err = GetAndValidateCodexAlphaSearchRequest(c) case types.RelayFormatOpenAIImage: request, err = GetAndValidOpenAIImageRequest(c, relayMode) @@ -55,6 +57,17 @@ func GetAndValidateRequest(c *gin.Context, format types.RelayFormat) (request dt return request, err } +func GetAndValidateCodexAlphaSearchRequest(c *gin.Context) (*dto.CodexAlphaSearchRequest, error) { + request := &dto.CodexAlphaSearchRequest{} + if err := common.UnmarshalBodyReusable(c, request); err != nil { + return nil, err + } + if strings.TrimSpace(request.Model) == "" { + return nil, errors.New("model is required") + } + return request, nil +} + func GetAndValidAudioRequest(c *gin.Context, relayMode int) (*dto.AudioRequest, error) { audioRequest := &dto.AudioRequest{} err := common.UnmarshalBodyReusable(c, audioRequest) diff --git a/router/relay-router.go b/router/relay-router.go index 17a13cad7fd..6cc5141bb14 100644 --- a/router/relay-router.go +++ b/router/relay-router.go @@ -104,6 +104,9 @@ func SetRelayRouter(router *gin.Engine) { httpRouter.POST("/responses/compact", func(c *gin.Context) { controller.Relay(c, types.RelayFormatOpenAIResponsesCompaction) }) + httpRouter.POST("/alpha/search", func(c *gin.Context) { + controller.Relay(c, types.RelayFormatCodexAlphaSearch) + }) // image related routes httpRouter.POST("/edits", func(c *gin.Context) { @@ -165,6 +168,20 @@ func SetRelayRouter(router *gin.Engine) { httpRouter.DELETE("/models/:model", controller.RelayNotImplemented) } + alphaSearchHandler := func(c *gin.Context) { + controller.Relay(c, types.RelayFormatCodexAlphaSearch) + } + alphaSearchRouter := router.Group("") + alphaSearchRouter.Use(middleware.RouteTag("relay")) + alphaSearchRouter.Use(middleware.SystemPerformanceCheck()) + alphaSearchRouter.Use(middleware.TokenAuth()) + alphaSearchRouter.Use(middleware.ModelRequestRateLimit()) + alphaSearchRouter.Use(middleware.Distribute()) + { + alphaSearchRouter.POST("/alpha/search", alphaSearchHandler) + alphaSearchRouter.POST("/backend-api/codex/alpha/search", alphaSearchHandler) + } + relayMjRouter := router.Group("/mj") relayMjRouter.Use(middleware.RouteTag("relay")) relayMjRouter.Use(middleware.SystemPerformanceCheck()) diff --git a/router/web-router.go b/router/web-router.go index 0d475e90d54..5bfbd92c314 100644 --- a/router/web-router.go +++ b/router/web-router.go @@ -32,7 +32,11 @@ func SetWebRouter(router *gin.Engine, assets ThemeAssets) { router.Use(static.Serve("/", themeFS)) router.NoRoute(func(c *gin.Context) { c.Set(middleware.RouteTagKey, "web") - if strings.HasPrefix(c.Request.RequestURI, "/v1") || strings.HasPrefix(c.Request.RequestURI, "/api") || strings.HasPrefix(c.Request.RequestURI, "/assets") { + if strings.HasPrefix(c.Request.RequestURI, "/v1") || + strings.HasPrefix(c.Request.RequestURI, "/api") || + strings.HasPrefix(c.Request.RequestURI, "/assets") || + strings.HasPrefix(c.Request.RequestURI, "/backend-api/") || + c.Request.URL.Path == "/alpha/search" { controller.RelayNotFound(c) return } diff --git a/service/tool_billing.go b/service/tool_billing.go index cadc750cda9..ab188807e94 100644 --- a/service/tool_billing.go +++ b/service/tool_billing.go @@ -1,10 +1,20 @@ package service import ( + "fmt" + "math" + "time" + "github.com/QuantumNous/new-api/common" + "github.com/QuantumNous/new-api/logger" + "github.com/QuantumNous/new-api/model" + relaycommon "github.com/QuantumNous/new-api/relay/common" "github.com/QuantumNous/new-api/setting/operation_setting" + "github.com/gin-gonic/gin" ) +const ToolNameCodexAlphaSearch = "codex_alpha_search" + // ToolCallUsage captures all tool call counts from a single request. type ToolCallUsage struct { ModelName string @@ -27,8 +37,10 @@ type ToolCallItem struct { // ToolCallResult holds the aggregated tool call billing for a request. type ToolCallResult struct { - TotalQuota int `json:"total_quota"` - Items []ToolCallItem `json:"items,omitempty"` + TotalQuota int `json:"total_quota"` + Items []ToolCallItem `json:"items,omitempty"` + Clamp *common.QuotaClamp `json:"-"` + Err error `json:"-"` } // ComputeToolCallQuota calculates the total quota for all tool calls in a @@ -37,17 +49,36 @@ type ToolCallResult struct { func ComputeToolCallQuota(usage ToolCallUsage, groupRatio float64) ToolCallResult { var items []ToolCallItem totalQuota := 0 + var quotaClamp *common.QuotaClamp + var err error + + noteClamp := func(clamp *common.QuotaClamp) { + if clamp != nil && quotaClamp == nil { + quotaClamp = clamp + } + } + addQuota := func(quota int) { + nextTotal, clamp := common.QuotaFromFloatChecked(float64(totalQuota) + float64(quota)) + noteClamp(clamp) + totalQuota = nextTotal + } + + if groupRatio <= 0 || math.IsNaN(groupRatio) || math.IsInf(groupRatio, 0) { + err = fmt.Errorf("invalid group ratio for tool call billing: %g", groupRatio) + return ToolCallResult{Err: err} + } addItem := func(toolName string, count int) { if count <= 0 { return } pricePer1K := operation_setting.GetToolPriceForModel(toolName, usage.ModelName) - if pricePer1K <= 0 { + if pricePer1K <= 0 || math.IsNaN(pricePer1K) || math.IsInf(pricePer1K, 0) { return } totalPrice := pricePer1K * float64(count) / 1000 - quota := common.QuotaRound(totalPrice * common.QuotaPerUnit * groupRatio) + quota, clamp := common.QuotaRoundChecked(totalPrice * common.QuotaPerUnit * groupRatio) + noteClamp(clamp) items = append(items, ToolCallItem{ Name: toolName, CallCount: count, @@ -55,7 +86,7 @@ func ComputeToolCallQuota(usage ToolCallUsage, groupRatio float64) ToolCallResul TotalPrice: totalPrice, Quota: quota, }) - totalQuota += quota + addQuota(quota) } if usage.WebSearchCalls > 0 && usage.WebSearchToolName != "" { @@ -68,7 +99,16 @@ func ComputeToolCallQuota(usage ToolCallUsage, groupRatio float64) ToolCallResul if usage.ImageGenerationCall { price := operation_setting.GetGPTImage1PriceOnceCall(usage.ImageGenerationQuality, usage.ImageGenerationSize) - quota := common.QuotaRound(price * common.QuotaPerUnit * groupRatio) + if price <= 0 || math.IsNaN(price) || math.IsInf(price, 0) { + return ToolCallResult{ + TotalQuota: totalQuota, + Items: items, + Clamp: quotaClamp, + Err: err, + } + } + quota, clamp := common.QuotaRoundChecked(price * common.QuotaPerUnit * groupRatio) + noteClamp(clamp) items = append(items, ToolCallItem{ Name: "image_generation", CallCount: 1, @@ -76,11 +116,80 @@ func ComputeToolCallQuota(usage ToolCallUsage, groupRatio float64) ToolCallResul TotalPrice: price, Quota: quota, }) - totalQuota += quota + addQuota(quota) } return ToolCallResult{ TotalQuota: totalQuota, Items: items, + Clamp: quotaClamp, + Err: err, + } +} + +func PostCodexAlphaSearchConsumeQuota(ctx *gin.Context, relayInfo *relaycommon.RelayInfo) { + if relayInfo == nil { + return } + result := ComputeToolCallQuota(ToolCallUsage{ + ModelName: relayInfo.OriginModelName, + WebSearchCalls: 1, + WebSearchToolName: ToolNameCodexAlphaSearch, + }, relayInfo.PriceData.GroupRatioInfo.GroupRatio) + if result.Err != nil { + logger.LogError(ctx, "error computing codex alpha search billing: "+result.Err.Error()) + return + } + if result.Clamp != nil && relayInfo.QuotaClamp == nil { + relayInfo.QuotaClamp = result.Clamp + } + + quota := result.TotalQuota + model.UpdateUserUsedQuotaAndRequestCount(relayInfo.UserId, quota) + if quota > 0 { + model.UpdateChannelUsedQuota(relayInfo.ChannelId, quota) + } + + if err := SettleBilling(ctx, relayInfo, quota); err != nil { + logger.LogError(ctx, "error settling codex alpha search billing: "+err.Error()) + } + + pricePer1K := operation_setting.GetToolPriceForModel(ToolNameCodexAlphaSearch, relayInfo.OriginModelName) + content := fmt.Sprintf("Codex Alpha Search called 1 time, quota cost %d", quota) + if pricePer1K > 0 { + content = fmt.Sprintf("Codex Alpha Search called 1 time, tool price %.4f USD/1K calls, quota cost %d", pricePer1K, quota) + } + + other := GenerateTextOtherInfo(ctx, relayInfo, + relayInfo.PriceData.ModelRatio, + relayInfo.PriceData.GroupRatioInfo.GroupRatio, + relayInfo.PriceData.CompletionRatio, + 0, + relayInfo.PriceData.CacheRatio, + relayInfo.PriceData.ModelPrice, + relayInfo.PriceData.GroupRatioInfo.GroupSpecialRatio, + ) + other["web_search"] = true + other["web_search_call_count"] = 1 + other["web_search_price"] = pricePer1K + other["codex_alpha_search"] = true + if len(result.Items) > 0 { + other["tool_call_items"] = result.Items + } + attachQuotaSaturation(ctx, relayInfo, other) + + model.RecordConsumeLog(ctx, relayInfo.UserId, model.RecordConsumeLogParams{ + ChannelId: relayInfo.ChannelId, + PromptTokens: 0, + CompletionTokens: 0, + ModelName: relayInfo.OriginModelName, + TokenName: ctx.GetString("token_name"), + Quota: quota, + Content: content, + TokenId: relayInfo.TokenId, + UseTimeSeconds: int(time.Since(relayInfo.StartTime).Seconds()), + IsStream: relayInfo.IsStream, + Group: relayInfo.UsingGroup, + Other: other, + }) } diff --git a/setting/operation_setting/tools.go b/setting/operation_setting/tools.go index 0eb2da0e506..03178ebfe52 100644 --- a/setting/operation_setting/tools.go +++ b/setting/operation_setting/tools.go @@ -22,6 +22,7 @@ import ( var defaultToolPrices = map[string]float64{ "web_search": 10.0, // OpenAI web search (all models) / Claude web search "web_search_preview": 10.0, // OpenAI web search preview (default: reasoning models) + "codex_alpha_search": 10.0, // Protective proxy billing for Codex standalone alpha/search. "file_search": 2.5, // OpenAI file search (Responses API) "google_search": 14.0, // Gemini Grounding with Google Search } diff --git a/types/relay_format.go b/types/relay_format.go index 9b4c86f2493..76b941396ec 100644 --- a/types/relay_format.go +++ b/types/relay_format.go @@ -8,6 +8,7 @@ const ( RelayFormatGemini = "gemini" RelayFormatOpenAIResponses = "openai_responses" RelayFormatOpenAIResponsesCompaction = "openai_responses_compaction" + RelayFormatCodexAlphaSearch = "codex_alpha_search" RelayFormatOpenAIAudio = "openai_audio" RelayFormatOpenAIImage = "openai_image" RelayFormatOpenAIRealtime = "openai_realtime"