139 lines
4.5 KiB
Go
139 lines
4.5 KiB
Go
package relay
|
|
|
|
import (
|
|
"errors"
|
|
"fmt"
|
|
"io"
|
|
"net/http"
|
|
|
|
"github.com/QuantumNous/new-api/common"
|
|
"github.com/QuantumNous/new-api/constant"
|
|
"github.com/QuantumNous/new-api/logger"
|
|
relaycommon "github.com/QuantumNous/new-api/relay/common"
|
|
"github.com/QuantumNous/new-api/relay/helper"
|
|
"github.com/QuantumNous/new-api/relaykit/dto"
|
|
"github.com/QuantumNous/new-api/relaykit/types"
|
|
"github.com/QuantumNous/new-api/service"
|
|
|
|
"github.com/gin-gonic/gin"
|
|
)
|
|
|
|
func AlphaSearchHelper(c *gin.Context, info *relaycommon.RelayInfo) (newAPIError *types.NewAPIError) {
|
|
info.InitChannelMeta(c)
|
|
|
|
switch info.ChannelType {
|
|
case constant.ChannelTypeSub2API,
|
|
constant.ChannelTypeNewAPI,
|
|
constant.ChannelTypeCodex,
|
|
constant.ChannelTypeAdvancedCustom:
|
|
default:
|
|
// Allow retry onto another channel that may support this endpoint.
|
|
return types.NewError(
|
|
errors.New("channel does not support /v1/alpha/search"),
|
|
types.ErrorCodeInvalidRequest,
|
|
)
|
|
}
|
|
|
|
request, ok := info.Request.(*dto.AlphaSearchRequest)
|
|
if !ok {
|
|
return types.NewErrorWithStatusCode(
|
|
fmt.Errorf("invalid request type, expected *dto.AlphaSearchRequest, got %T", info.Request),
|
|
types.ErrorCodeInvalidRequest,
|
|
http.StatusBadRequest,
|
|
types.ErrOptionWithSkipRetry(),
|
|
)
|
|
}
|
|
|
|
err := helper.ModelMappedHelper(c, info, request)
|
|
if err != nil {
|
|
return types.NewError(err, types.ErrorCodeChannelModelMappedError, types.ErrOptionWithSkipRetry())
|
|
}
|
|
|
|
jsonData, err := buildAlphaSearchRequestBody(request.RawBody, info.OriginModelName, info.UpstreamModelName)
|
|
if err != nil {
|
|
return types.NewError(err, types.ErrorCodeConvertRequestFailed, types.ErrOptionWithSkipRetry())
|
|
}
|
|
|
|
if len(info.ParamOverride) > 0 {
|
|
jsonData, err = relaycommon.ApplyParamOverrideWithRelayInfo(jsonData, info)
|
|
if err != nil {
|
|
return newAPIErrorFromParamOverride(err)
|
|
}
|
|
}
|
|
|
|
logger.LogDebug(c, "requestBody: %s", jsonData)
|
|
body, closer, err := relaycommon.NewOutboundJSONBody(jsonData)
|
|
if err != nil {
|
|
return types.NewError(err, types.ErrorCodeConvertRequestFailed, types.ErrOptionWithSkipRetry())
|
|
}
|
|
defer closer.Close()
|
|
|
|
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)
|
|
|
|
resp, err := adaptor.DoRequest(c, info, body)
|
|
if err != nil {
|
|
return types.NewOpenAIError(err, types.ErrorCodeDoRequestFailed, http.StatusInternalServerError)
|
|
}
|
|
|
|
statusCodeMappingStr := c.GetString("status_code_mapping")
|
|
httpResp, ok := resp.(*http.Response)
|
|
if !ok || httpResp == nil {
|
|
return types.NewOpenAIError(errors.New("invalid http response"), types.ErrorCodeDoRequestFailed, http.StatusInternalServerError)
|
|
}
|
|
defer httpResp.Body.Close()
|
|
|
|
if httpResp.StatusCode < 200 || httpResp.StatusCode >= 300 {
|
|
newAPIError = service.RelayErrorHandler(c.Request.Context(), httpResp, false)
|
|
service.ResetStatusCode(newAPIError, statusCodeMappingStr)
|
|
return newAPIError
|
|
}
|
|
|
|
if contentType := httpResp.Header.Get("Content-Type"); contentType != "" {
|
|
c.Writer.Header().Set("Content-Type", contentType)
|
|
}
|
|
c.Writer.WriteHeader(httpResp.StatusCode)
|
|
if _, err := io.Copy(c.Writer, httpResp.Body); err != nil {
|
|
return types.NewError(err, types.ErrorCodeDoRequestFailed, types.ErrOptionWithSkipRetry())
|
|
}
|
|
|
|
// Upstream alpha search returns no usage; bill one web_search_preview call.
|
|
if info.ResponsesUsageInfo == nil {
|
|
info.ResponsesUsageInfo = &relaycommon.ResponsesUsageInfo{
|
|
BuiltInTools: make(map[string]*relaycommon.BuildInToolInfo),
|
|
}
|
|
}
|
|
if info.ResponsesUsageInfo.BuiltInTools == nil {
|
|
info.ResponsesUsageInfo.BuiltInTools = make(map[string]*relaycommon.BuildInToolInfo)
|
|
}
|
|
info.ResponsesUsageInfo.BuiltInTools[dto.BuildInToolWebSearchPreview] = &relaycommon.BuildInToolInfo{
|
|
ToolName: dto.BuildInToolWebSearchPreview,
|
|
CallCount: 1,
|
|
}
|
|
|
|
usage := &dto.Usage{}
|
|
service.PostTextConsumeQuota(c, info, usage, nil)
|
|
return nil
|
|
}
|
|
|
|
// buildAlphaSearchRequestBody returns RawBody unchanged unless the model was
|
|
// mapped, in which case only the "model" field is rewritten so unknown fields
|
|
// are preserved.
|
|
func buildAlphaSearchRequestBody(rawBody []byte, originModel, upstreamModel string) ([]byte, error) {
|
|
if len(rawBody) == 0 {
|
|
return nil, errors.New("empty alpha search request body")
|
|
}
|
|
if upstreamModel == "" || upstreamModel == originModel {
|
|
return rawBody, nil
|
|
}
|
|
var body map[string]any
|
|
if err := common.Unmarshal(rawBody, &body); err != nil {
|
|
return nil, err
|
|
}
|
|
body["model"] = upstreamModel
|
|
return common.Marshal(body)
|
|
}
|