// Filename: internal/service/resource_service.go package service import ( "errors" apperrors "gemini-balancer/internal/errors" "gemini-balancer/internal/models" "gemini-balancer/internal/repository" "gemini-balancer/internal/settings" "sort" "sync" "time" "github.com/sirupsen/logrus" ) var ( ErrNoResourceAvailable = errors.New("no available resource found for the request") ) type RequestResources struct { KeyGroup *models.KeyGroup APIKey *models.APIKey UpstreamEndpoint *models.UpstreamEndpoint ProxyConfig *models.ProxyConfig RequestConfig *models.RequestConfig } type ResourceService struct { settingsManager *settings.SettingsManager groupManager *GroupManager keyRepo repository.KeyRepository apiKeyService *APIKeyService logger *logrus.Entry initOnce sync.Once } func NewResourceService( sm *settings.SettingsManager, gm *GroupManager, kr repository.KeyRepository, aks *APIKeyService, logger *logrus.Logger, ) *ResourceService { logger.Debugf("[FORENSIC PROBE | INJECTION | ResourceService] Received 'keyRepo' param. Fingerprint: %p", kr) rs := &ResourceService{ settingsManager: sm, groupManager: gm, keyRepo: kr, apiKeyService: aks, logger: logger.WithField("component", "ResourceService📦️"), } rs.initOnce.Do(func() { go rs.preWarmCache(logger) }) return rs } // --- [模式一:智能聚合模式] --- func (s *ResourceService) GetResourceFromBasePool(authToken *models.AuthToken, modelName string) (*RequestResources, error) { log := s.logger.WithFields(logrus.Fields{"token_id": authToken.ID, "model_name": modelName, "mode": "BasePool"}) log.Debug("Entering BasePool resource acquisition.") // 1.筛选出所有符合条件的候选组,并按优先级排序 candidateGroups := s.filterAndSortCandidateGroups(modelName, authToken.AllowedGroups) if len(candidateGroups) == 0 { log.Warn("No candidate groups found for BasePool construction.") return nil, apperrors.ErrNoKeysAvailable } // 2.从 BasePool中,根据系统全局策略选择一个Key basePool := &repository.BasePool{ CandidateGroups: candidateGroups, PollingStrategy: s.settingsManager.GetSettings().PollingStrategy, } apiKey, selectedGroup, err := s.keyRepo.SelectOneActiveKeyFromBasePool(basePool) if err != nil { log.WithError(err).Warn("Failed to select a key from the BasePool.") return nil, apperrors.ErrNoKeysAvailable } // 3. 组装最终资源 // [关键] 在此模式下,RequestConfig 永远是空的,以保证透明性。 resources, err := s.assembleRequestResources(selectedGroup, apiKey) if err != nil { log.WithError(err).Error("Failed to assemble resources after selecting key from BasePool.") return nil, err } resources.RequestConfig = &models.RequestConfig{} // 强制为空 log.Infof("Successfully selected KeyID %d from GroupID %d for the BasePool.", apiKey.ID, selectedGroup.ID) return resources, nil } // --- [模式二:精确路由模式] --- func (s *ResourceService) GetResourceFromGroup(authToken *models.AuthToken, groupName string) (*RequestResources, error) { log := s.logger.WithFields(logrus.Fields{"token_id": authToken.ID, "group_name": groupName, "mode": "PreciseRoute"}) log.Debug("Entering PreciseRoute resource acquisition.") targetGroup, ok := s.groupManager.GetGroupByName(groupName) if !ok { return nil, apperrors.NewAPIError(apperrors.ErrGroupNotFound, "The specified group does not exist.") } if !s.isTokenAllowedForGroup(authToken, targetGroup.ID) { return nil, apperrors.NewAPIError(apperrors.ErrPermissionDenied, "Token does not have permission to access this group.") } apiKey, _, err := s.keyRepo.SelectOneActiveKey(targetGroup) if err != nil { log.WithError(err).Warn("Failed to select a key from the precisely targeted group.") return nil, apperrors.ErrNoKeysAvailable } resources, err := s.assembleRequestResources(targetGroup, apiKey) if err != nil { log.WithError(err).Error("Failed to assemble resources for precise route.") return nil, err } resources.RequestConfig = targetGroup.RequestConfig log.Infof("Successfully selected KeyID %d by precise routing to GroupID %d.", apiKey.ID, targetGroup.ID) return resources, nil } func (s *ResourceService) GetAllowedModelsForToken(authToken *models.AuthToken) []string { allGroups := s.groupManager.GetAllGroups() if len(allGroups) == 0 { return []string{} } allowedModelsSet := make(map[string]struct{}) if authToken.IsAdmin { for _, group := range allGroups { for _, modelMapping := range group.AllowedModels { allowedModelsSet[modelMapping.ModelName] = struct{}{} } } } else { allowedGroupIDs := make(map[uint]bool) for _, ag := range authToken.AllowedGroups { allowedGroupIDs[ag.ID] = true } for _, group := range allGroups { if _, ok := allowedGroupIDs[group.ID]; ok { for _, modelMapping := range group.AllowedModels { allowedModelsSet[modelMapping.ModelName] = struct{}{} } } } } result := make([]string, 0, len(allowedModelsSet)) for modelName := range allowedModelsSet { result = append(result, modelName) } sort.Strings(result) return result } func (s *ResourceService) assembleRequestResources(group *models.KeyGroup, apiKey *models.APIKey) (*RequestResources, error) { selectedUpstream := s.selectUpstreamForGroup(group) if selectedUpstream == nil { return nil, apperrors.NewAPIError(apperrors.ErrConfigurationError, "Selected group has no valid upstream and no global default is set.") } var proxyConfig *models.ProxyConfig // [注意] 代理逻辑需要一个 proxyModule 实例,我们暂时置空。后续需要重新注入依赖。 // if group.EnableProxy && s.proxyModule != nil { // var err error // proxyConfig, err = s.proxyModule.AssignProxyIfNeeded(apiKey) // if err != nil { // s.logger.WithError(err).Warnf("Failed to assign proxy for API key %d.", apiKey.ID) // } // } return &RequestResources{ KeyGroup: group, APIKey: apiKey, UpstreamEndpoint: selectedUpstream, ProxyConfig: proxyConfig, }, nil } func (s *ResourceService) selectUpstreamForGroup(group *models.KeyGroup) *models.UpstreamEndpoint { if len(group.AllowedUpstreams) > 0 { return group.AllowedUpstreams[0] } globalSettings := s.settingsManager.GetSettings() if globalSettings.DefaultUpstreamURL != "" { return &models.UpstreamEndpoint{URL: globalSettings.DefaultUpstreamURL, Status: "active"} } return nil } func (s *ResourceService) preWarmCache(logger *logrus.Logger) error { time.Sleep(2 * time.Second) s.logger.Info("Performing initial key cache pre-warming...") if err := s.keyRepo.LoadAllKeysToStore(); err != nil { logger.WithError(err).Error("Failed to perform initial key cache pre-warming.") return err } s.logger.Info("Initial key cache pre-warming completed successfully.") return nil } func (s *ResourceService) GetResourcesForRequest(modelName string, allowedGroups []*models.KeyGroup) (*RequestResources, error) { return nil, errors.New("GetResourcesForRequest is deprecated; use GetResourceFromBasePool or GetResourceFromGroup") } func (s *ResourceService) filterAndSortCandidateGroups(modelName string, allowedGroupsFromToken []*models.KeyGroup) []*models.KeyGroup { allGroupsFromCache := s.groupManager.GetAllGroups() var candidateGroups []*models.KeyGroup // 1. 确定权限范围 allowedGroupIDs := make(map[uint]bool) isTokenRestricted := len(allowedGroupsFromToken) > 0 if isTokenRestricted { for _, ag := range allowedGroupsFromToken { allowedGroupIDs[ag.ID] = true } } // 2. 筛选 for _, group := range allGroupsFromCache { // 检查Token权限 if isTokenRestricted && !allowedGroupIDs[group.ID] { continue } // 检查模型是否被允许 isModelAllowed := false if len(group.AllowedModels) == 0 { // 如果组不限制模型,则允许 isModelAllowed = true } else { for _, m := range group.AllowedModels { if m.ModelName == modelName { isModelAllowed = true break } } } if isModelAllowed { candidateGroups = append(candidateGroups, group) } } // 3.按 Order 字段升序排序 sort.SliceStable(candidateGroups, func(i, j int) bool { return candidateGroups[i].Order < candidateGroups[j].Order }) return candidateGroups } func (s *ResourceService) isTokenAllowedForGroup(authToken *models.AuthToken, groupID uint) bool { if authToken.IsAdmin { return true } for _, allowedGroup := range authToken.AllowedGroups { if allowedGroup.ID == groupID { return true } } return false } func (s *ResourceService) ReportRequestResult(resources *RequestResources, success bool, apiErr *apperrors.APIError) { if resources == nil || resources.KeyGroup == nil || resources.APIKey == nil { return } s.apiKeyService.HandleRequestResult(resources.KeyGroup, resources.APIKey, success, apiErr) }