feat: web search adaptor plugin (labring#216)

* feat: web search adaptor plugin

* fix: ci lint

* fix: remove web_search_options field on web search plugin

* docs: add web search docs

Author: zijiren233

core/common/gin.go

@@ -40,6 +40,12 @@ func (l *LimitedReader) Read(p []byte) (n int, err error) {
 	return
 }
 
+func SetRequestBody(req *http.Request, body []byte) {
+	ctx := req.Context()
+	bufCtx := context.WithValue(ctx, RequestBodyKey{}, body)
+	*req = *req.WithContext(bufCtx)
+}
+
 func GetRequestBody(req *http.Request) ([]byte, error) {
 	contentType := req.Header.Get("Content-Type")
 	if contentType == "application/x-www-form-urlencoded" ||
@@ -78,9 +84,7 @@ func GetRequestBody(req *http.Request) ([]byte, error) {
 	if err != nil {
 		return nil, fmt.Errorf("request body read failed: %w", err)
 	}
-	ctx := req.Context()
-	bufCtx := context.WithValue(ctx, RequestBodyKey{}, buf)
-	*req = *req.WithContext(bufCtx)
+	SetRequestBody(req, buf)
 	return buf, nil
 }
 

core/controller/channel-test.go

@@ -33,11 +33,11 @@ import (
 const channelTestRequestID = "channel-test"
 
 var (
-	modelConfigCache     map[string]*model.ModelConfig = make(map[string]*model.ModelConfig)
+	modelConfigCache     map[string]model.ModelConfig = make(map[string]model.ModelConfig)
 	modelConfigCacheOnce sync.Once
 )
 
-func guessModelConfig(model string) *model.ModelConfig {
+func guessModelConfig(modelName string) model.ModelConfig {
 	modelConfigCacheOnce.Do(func() {
 		for _, c := range adaptors.ChannelAdaptor {
 			for _, m := range c.GetModelList() {
@@ -48,10 +48,10 @@ func guessModelConfig(model string) *model.ModelConfig {
 		}
 	})
 
-	if cachedConfig, ok := modelConfigCache[model]; ok {
+	if cachedConfig, ok := modelConfigCache[modelName]; ok {
 		return cachedConfig
 	}
-	return nil
+	return model.ModelConfig{}
 }
 
 // testSingleModel tests a single model in the channel
@@ -62,7 +62,7 @@ func testSingleModel(mc *model.ModelCaches, channel *model.Channel, modelName st
 	}
 	if modelConfig.Type == mode.Unknown {
 		newModelConfig := guessModelConfig(modelName)
-		if newModelConfig != nil {
+		if newModelConfig.Type != mode.Unknown {
 			modelConfig = newModelConfig
 		}
 	}

core/controller/dashboard.go

@@ -270,7 +270,7 @@ func GetGroupDashboardModels(c *gin.Context) {
 			}) {
 				continue
 			}
-			newEnabledModelConfigs = append(newEnabledModelConfigs, middleware.GetGroupAdjustedModelConfig(groupCache, *mc))
+			newEnabledModelConfigs = append(newEnabledModelConfigs, middleware.GetGroupAdjustedModelConfig(groupCache, mc))
 		}
 	}
 	middleware.SuccessResponse(c, newEnabledModelConfigs)

core/controller/model.go

@@ -55,14 +55,14 @@ func (c *BuiltinModelConfig) MarshalJSON() ([]byte, error) {
 	})
 }
 
-func SortBuiltinModelConfigsFunc(i, j *BuiltinModelConfig) int {
-	return model.SortModelConfigsFunc((*model.ModelConfig)(i), (*model.ModelConfig)(j))
+func SortBuiltinModelConfigsFunc(i, j BuiltinModelConfig) int {
+	return model.SortModelConfigsFunc((model.ModelConfig)(i), (model.ModelConfig)(j))
 }
 
 var (
-	builtinModels             []*BuiltinModelConfig
+	builtinModels             []BuiltinModelConfig
 	builtinModelsMap          map[string]*OpenAIModels
-	builtinChannelType2Models map[model.ChannelType][]*BuiltinModelConfig
+	builtinChannelType2Models map[model.ChannelType][]BuiltinModelConfig
 )
 
 var permission = []OpenAIModelPermission{
@@ -83,12 +83,12 @@ var permission = []OpenAIModelPermission{
 }
 
 func init() {
-	builtinChannelType2Models = make(map[model.ChannelType][]*BuiltinModelConfig)
+	builtinChannelType2Models = make(map[model.ChannelType][]BuiltinModelConfig)
 	builtinModelsMap = make(map[string]*OpenAIModels)
 	// https://platform.openai.com/docs/models/model-endpoint-compatibility
 	for i, adaptor := range adaptors.ChannelAdaptor {
 		modelNames := adaptor.GetModelList()
-		builtinChannelType2Models[i] = make([]*BuiltinModelConfig, len(modelNames))
+		builtinChannelType2Models[i] = make([]BuiltinModelConfig, len(modelNames))
 		for idx, _model := range modelNames {
 			if _model.Owner == "" {
 				_model.Owner = model.ModelOwner(i.String())
@@ -103,11 +103,11 @@ func init() {
 					Root:       _model.Model,
 					Parent:     nil,
 				}
-				builtinModels = append(builtinModels, (*BuiltinModelConfig)(_model))
+				builtinModels = append(builtinModels, (BuiltinModelConfig)(_model))
 			} else if v.OwnedBy != string(_model.Owner) {
 				log.Fatalf("model %s owner mismatch, expect %s, actual %s", _model.Model, string(_model.Owner), v.OwnedBy)
 			}
-			builtinChannelType2Models[i][idx] = (*BuiltinModelConfig)(_model)
+			builtinChannelType2Models[i][idx] = (BuiltinModelConfig)(_model)
 		}
 	}
 	for _, models := range builtinChannelType2Models {

core/controller/modelconfig.go

@@ -112,7 +112,7 @@ func SearchModelConfigs(c *gin.Context) {
 type SaveModelConfigsRequest struct {
 	CreatedAt int64 `json:"created_at"`
 	UpdatedAt int64 `json:"updated_at"`
-	*model.ModelConfig
+	model.ModelConfig
 }
 
 // SaveModelConfigs godoc
@@ -126,12 +126,12 @@ type SaveModelConfigsRequest struct {
 //	@Success		200		{object}	middleware.APIResponse
 //	@Router			/api/model_configs/ [post]
 func SaveModelConfigs(c *gin.Context) {
-	var configs []*SaveModelConfigsRequest
+	var configs []SaveModelConfigsRequest
 	if err := c.ShouldBindJSON(&configs); err != nil {
 		middleware.ErrorResponse(c, http.StatusBadRequest, err.Error())
 		return
 	}
-	modelConfigs := make([]*model.ModelConfig, len(configs))
+	modelConfigs := make([]model.ModelConfig, len(configs))
 	for i, config := range configs {
 		modelConfigs[i] = config.ModelConfig
 	}

core/controller/relay-controller.go

@@ -31,15 +31,17 @@ import (
 	"github.com/labring/aiproxy/core/relay/meta"
 	"github.com/labring/aiproxy/core/relay/mode"
 	relaymodel "github.com/labring/aiproxy/core/relay/model"
+	"github.com/labring/aiproxy/core/relay/plugin"
+	websearch "github.com/labring/aiproxy/core/relay/plugin/web-search"
 	log "github.com/sirupsen/logrus"
 )
 
 // https://platform.openai.com/docs/api-reference/chat
 
 type (
 	RelayHandler    func(*gin.Context, *meta.Meta) *controller.HandleResult
-	GetRequestUsage func(*gin.Context, *model.ModelConfig) (model.Usage, error)
-	GetRequestPrice func(*gin.Context, *model.ModelConfig) (model.Price, error)
+	GetRequestUsage func(*gin.Context, model.ModelConfig) (model.Usage, error)
+	GetRequestPrice func(*gin.Context, model.ModelConfig) (model.Price, error)
 )
 
 type RelayController struct {
@@ -48,6 +50,7 @@ type RelayController struct {
 	Handler         RelayHandler
 }
 
+// TODO: convert to plugin
 type wrapAdaptor struct {
 	adaptor.Adaptor
 }
@@ -163,7 +166,13 @@ func relayHandler(c *gin.Context, meta *meta.Meta) *controller.HandleResult {
 		}
 	}
 
-	return controller.Handle(&wrapAdaptor{adaptor}, c, meta)
+	a := plugin.WrapperAdaptor(&wrapAdaptor{adaptor},
+		websearch.NewWebSearchPlugin(func(modelName string) (*model.Channel, error) {
+			return getWebSearchChannel(c, modelName)
+		}),
+	)
+
+	return controller.Handle(a, c, meta)
 }
 
 func relayController(m mode.Mode) RelayController {
@@ -318,9 +327,17 @@ var (
 
 func GetRandomChannel(mc *model.ModelCaches, availableSet []string, modelName string, errorRates map[int64]float64, ignoreChannel ...int64) (*model.Channel, []*model.Channel, error) {
 	channelMap := make(map[int]*model.Channel)
-	for _, set := range availableSet {
-		for _, channel := range mc.EnabledModel2ChannelsBySet[set][modelName] {
-			channelMap[channel.ID] = channel
+	if len(availableSet) != 0 {
+		for _, set := range availableSet {
+			for _, channel := range mc.EnabledModel2ChannelsBySet[set][modelName] {
+				channelMap[channel.ID] = channel
+			}
+		}
+	} else {
+		for _, sets := range mc.EnabledModel2ChannelsBySet {
+			for _, channel := range sets[modelName] {
+				channelMap[channel.ID] = channel
+			}
 		}
 	}
 	migratedChannels := make([]*model.Channel, 0, len(channelMap))
@@ -403,12 +420,11 @@ func NewMetaByContext(c *gin.Context, channel *model.Channel, mode mode.Mode, op
 }
 
 func relay(c *gin.Context, mode mode.Mode, relayController RelayController) {
-	log := middleware.GetLogger(c)
 	requestModel := middleware.GetRequestModel(c)
 	mc := middleware.GetModelConfig(c)
 
 	// Get initial channel
-	initialChannel, err := getInitialChannel(c, requestModel, log)
+	initialChannel, err := getInitialChannel(c, requestModel)
 	if err != nil || initialChannel == nil || initialChannel.channel == nil {
 		middleware.AbortLogWithMessageWithMode(mode, c,
 			http.StatusServiceUnavailable,
@@ -486,7 +502,7 @@ func relay(c *gin.Context, mode mode.Mode, relayController RelayController) {
 	)
 
 	// Retry loop
-	retryLoop(c, mode, retryState, relayController.Handler, log)
+	retryLoop(c, mode, retryState, relayController.Handler)
 }
 
 // recordResult records the consumption for the final result
@@ -572,7 +588,8 @@ type initialChannel struct {
 	migratedChannels  []*model.Channel
 }
 
-func getInitialChannel(c *gin.Context, modelName string, log *log.Entry) (*initialChannel, error) {
+func getInitialChannel(c *gin.Context, modelName string) (*initialChannel, error) {
+	log := middleware.GetLogger(c)
 	if channel := middleware.GetChannel(c); channel != nil {
 		log.Data["designated_channel"] = "true"
 		return &initialChannel{channel: channel, designatedChannel: true}, nil
@@ -607,6 +624,29 @@ func getInitialChannel(c *gin.Context, modelName string, log *log.Entry) (*initi
 	}, nil
 }
 
+func getWebSearchChannel(c *gin.Context, modelName string) (*model.Channel, error) {
+	log := middleware.GetLogger(c)
+	mc := middleware.GetModelCaches(c)
+
+	ids, err := monitor.GetBannedChannelsWithModel(c.Request.Context(), modelName)
+	if err != nil {
+		log.Errorf("get %s auto banned channels failed: %+v", modelName, err)
+	}
+	log.Debugf("%s model banned channels: %+v", modelName, ids)
+
+	errorRates, err := monitor.GetModelChannelErrorRate(c.Request.Context(), modelName)
+	if err != nil {
+		log.Errorf("get channel model error rates failed: %+v", err)
+	}
+
+	channel, _, err := getChannelWithFallback(mc, nil, modelName, errorRates, ids...)
+	if err != nil {
+		return nil, err
+	}
+
+	return channel, nil
+}
+
 func handleRelayResult(c *gin.Context, bizErr adaptor.Error, retry bool, retryTimes int) (done bool) {
 	if bizErr == nil {
 		return true
@@ -645,7 +685,9 @@ func initRetryState(retryTimes int, channel *initialChannel, meta *meta.Meta, re
 	return state
 }
 
-func retryLoop(c *gin.Context, mode mode.Mode, state *retryState, relayController RelayHandler, log *log.Entry) {
+func retryLoop(c *gin.Context, mode mode.Mode, state *retryState, relayController RelayHandler) {
+	log := middleware.GetLogger(c)
+
 	// do not use for i := range state.retryTimes, because the retryTimes is constant
 	i := 0
 

core/middleware/distributor.go

@@ -132,10 +132,10 @@ func UpdateGroupModelTokennameTokensRequest(c *gin.Context, tpm, tps int64) {
 	// log.Data["tps"] = strconv.FormatInt(tps, 10)
 }
 
-func checkGroupModelRPMAndTPM(c *gin.Context, group *model.GroupCache, mc *model.ModelConfig, tokenName string) error {
+func checkGroupModelRPMAndTPM(c *gin.Context, group *model.GroupCache, mc model.ModelConfig, tokenName string) error {
 	log := GetLogger(c)
 
-	adjustedModelConfig := GetGroupAdjustedModelConfig(group, *mc)
+	adjustedModelConfig := GetGroupAdjustedModelConfig(group, mc)
 
 	groupModelCount, groupModelOverLimitCount, groupModelSecondCount := reqlimit.PushGroupModelRequest(c.Request.Context(), group.ID, mc.Model, adjustedModelConfig.RPM)
 	UpdateGroupModelRequest(c, group, groupModelCount+groupModelOverLimitCount, groupModelSecondCount)
@@ -461,8 +461,8 @@ func GetRequestMetadata(c *gin.Context) map[string]string {
 	return c.GetStringMapString(RequestMetadata)
 }
 
-func GetModelConfig(c *gin.Context) *model.ModelConfig {
-	return c.MustGet(ModelConfig).(*model.ModelConfig)
+func GetModelConfig(c *gin.Context) model.ModelConfig {
+	return c.MustGet(ModelConfig).(model.ModelConfig)
 }
 
 func NewMetaByContext(c *gin.Context,

core/model/cache.go

@@ -672,7 +672,7 @@ func CacheGetPublicMCPReusingParam(mcpID, groupID string) (*PublicMCPReusingPara
 
 //nolint:revive
 type ModelConfigCache interface {
-	GetModelConfig(model string) (*ModelConfig, bool)
+	GetModelConfig(model string) (ModelConfig, bool)
 }
 
 // read-only cache
@@ -684,9 +684,9 @@ type ModelCaches struct {
 	// map[set][]model
 	EnabledModelsBySet map[string][]string
 	// map[set][]modelconfig
-	EnabledModelConfigsBySet map[string][]*ModelConfig
+	EnabledModelConfigsBySet map[string][]ModelConfig
 	// map[model]modelconfig
-	EnabledModelConfigsMap map[string]*ModelConfig
+	EnabledModelConfigsMap map[string]ModelConfig
 
 	// map[set]map[model][]channel
 	EnabledModel2ChannelsBySet map[string]map[string][]*Channel
@@ -811,10 +811,10 @@ func LoadChannelByID(id int) (*Channel, error) {
 var _ ModelConfigCache = (*modelConfigMapCache)(nil)
 
 type modelConfigMapCache struct {
-	modelConfigMap map[string]*ModelConfig
+	modelConfigMap map[string]ModelConfig
 }
 
-func (m *modelConfigMapCache) GetModelConfig(model string) (*ModelConfig, bool) {
+func (m *modelConfigMapCache) GetModelConfig(model string) (ModelConfig, bool) {
 	config, ok := m.modelConfigMap[model]
 	return config, ok
 }
@@ -825,7 +825,7 @@ type disabledModelConfigCache struct {
 	modelConfigs ModelConfigCache
 }
 
-func (d *disabledModelConfigCache) GetModelConfig(model string) (*ModelConfig, bool) {
+func (d *disabledModelConfigCache) GetModelConfig(model string) (ModelConfig, bool) {
 	if config, ok := d.modelConfigs.GetModelConfig(model); ok {
 		return config, true
 	}
@@ -837,7 +837,7 @@ func initializeModelConfigCache() (ModelConfigCache, error) {
 	if err != nil {
 		return nil, err
 	}
-	newModelConfigMap := make(map[string]*ModelConfig)
+	newModelConfigMap := make(map[string]ModelConfig)
 	for _, modelConfig := range modelConfigs {
 		newModelConfigMap[modelConfig.Model] = modelConfig
 	}
@@ -905,16 +905,16 @@ func sortChannelsByPriorityBySet(modelMapBySet map[string]map[string][]*Channel)
 
 func buildEnabledModelsBySet(modelMapBySet map[string]map[string][]*Channel, modelConfigCache ModelConfigCache) (
 	map[string][]string,
-	map[string][]*ModelConfig,
-	map[string]*ModelConfig,
+	map[string][]ModelConfig,
+	map[string]ModelConfig,
 ) {
 	modelsBySet := make(map[string][]string)
-	modelConfigsBySet := make(map[string][]*ModelConfig)
-	modelConfigsMap := make(map[string]*ModelConfig)
+	modelConfigsBySet := make(map[string][]ModelConfig)
+	modelConfigsMap := make(map[string]ModelConfig)
 
 	for set, modelMap := range modelMapBySet {
 		models := make([]string, 0)
-		configs := make([]*ModelConfig, 0)
+		configs := make([]ModelConfig, 0)
 		appended := make(map[string]struct{})
 
 		for model := range modelMap {
@@ -940,7 +940,7 @@ func buildEnabledModelsBySet(modelMapBySet map[string]map[string][]*Channel, mod
 	return modelsBySet, modelConfigsBySet, modelConfigsMap
 }
 
-func SortModelConfigsFunc(i, j *ModelConfig) int {
+func SortModelConfigsFunc(i, j ModelConfig) int {
 	if i.Owner != j.Owner {
 		if natural.Less(string(i.Owner), string(j.Owner)) {
 			return -1