| |
| |
| |
| package registry |
|
|
| import ( |
| "context" |
| "fmt" |
| "sort" |
| "strings" |
| "sync" |
| "time" |
|
|
| misc "github.com/router-for-me/CLIProxyAPI/v6/internal/misc" |
| log "github.com/sirupsen/logrus" |
| ) |
|
|
| |
| type ModelInfo struct { |
| |
| ID string `json:"id"` |
| |
| Object string `json:"object"` |
| |
| Created int64 `json:"created"` |
| |
| OwnedBy string `json:"owned_by"` |
| |
| Type string `json:"type"` |
| |
| DisplayName string `json:"display_name,omitempty"` |
| |
| Name string `json:"name,omitempty"` |
| |
| Version string `json:"version,omitempty"` |
| |
| Description string `json:"description,omitempty"` |
| |
| ContextLength int `json:"context_length,omitempty"` |
| |
| MaxCompletionTokens int `json:"max_completion_tokens,omitempty"` |
| |
| SupportedGenerationMethods []string `json:"supportedGenerationMethods,omitempty"` |
| |
| SupportedParameters []string `json:"supported_parameters,omitempty"` |
|
|
| |
| |
| Thinking *ThinkingSupport `json:"thinking,omitempty"` |
|
|
| |
| |
| |
| UserDefined bool `json:"-"` |
| } |
|
|
| |
| |
| type ThinkingSupport struct { |
| |
| Min int `json:"min,omitempty"` |
| |
| Max int `json:"max,omitempty"` |
| |
| ZeroAllowed bool `json:"zero_allowed,omitempty"` |
| |
| DynamicAllowed bool `json:"dynamic_allowed,omitempty"` |
| |
| |
| Levels []string `json:"levels,omitempty"` |
| } |
|
|
| |
| type ModelRegistration struct { |
| |
| Info *ModelInfo |
| |
| InfoByProvider map[string]*ModelInfo |
| |
| Count int |
| |
| LastUpdated time.Time |
| |
| QuotaExceededClients map[string]*time.Time |
| |
| Providers map[string]int |
| |
| SuspendedClients map[string]string |
| } |
|
|
| |
| |
| type ModelRegistryHook interface { |
| OnModelsRegistered(ctx context.Context, provider, clientID string, models []*ModelInfo) |
| OnModelsUnregistered(ctx context.Context, provider, clientID string) |
| } |
|
|
| |
| type ModelRegistry struct { |
| |
| models map[string]*ModelRegistration |
| |
| clientModels map[string][]string |
| |
| |
| clientModelInfos map[string]map[string]*ModelInfo |
| |
| clientProviders map[string]string |
| |
| mutex *sync.RWMutex |
| |
| hook ModelRegistryHook |
| } |
|
|
| |
| var globalRegistry *ModelRegistry |
| var registryOnce sync.Once |
|
|
| |
| func GetGlobalRegistry() *ModelRegistry { |
| registryOnce.Do(func() { |
| globalRegistry = &ModelRegistry{ |
| models: make(map[string]*ModelRegistration), |
| clientModels: make(map[string][]string), |
| clientModelInfos: make(map[string]map[string]*ModelInfo), |
| clientProviders: make(map[string]string), |
| mutex: &sync.RWMutex{}, |
| } |
| }) |
| return globalRegistry |
| } |
|
|
| |
| func LookupModelInfo(modelID string, provider ...string) *ModelInfo { |
| modelID = strings.TrimSpace(modelID) |
| if modelID == "" { |
| return nil |
| } |
|
|
| p := "" |
| if len(provider) > 0 { |
| p = strings.ToLower(strings.TrimSpace(provider[0])) |
| } |
|
|
| if info := GetGlobalRegistry().GetModelInfo(modelID, p); info != nil { |
| return info |
| } |
| return LookupStaticModelInfo(modelID) |
| } |
|
|
| |
| func (r *ModelRegistry) SetHook(hook ModelRegistryHook) { |
| if r == nil { |
| return |
| } |
| r.mutex.Lock() |
| defer r.mutex.Unlock() |
| r.hook = hook |
| } |
|
|
| const defaultModelRegistryHookTimeout = 5 * time.Second |
|
|
| func (r *ModelRegistry) triggerModelsRegistered(provider, clientID string, models []*ModelInfo) { |
| hook := r.hook |
| if hook == nil { |
| return |
| } |
| modelsCopy := cloneModelInfosUnique(models) |
| go func() { |
| defer func() { |
| if recovered := recover(); recovered != nil { |
| log.Errorf("model registry hook OnModelsRegistered panic: %v", recovered) |
| } |
| }() |
| ctx, cancel := context.WithTimeout(context.Background(), defaultModelRegistryHookTimeout) |
| defer cancel() |
| hook.OnModelsRegistered(ctx, provider, clientID, modelsCopy) |
| }() |
| } |
|
|
| func (r *ModelRegistry) triggerModelsUnregistered(provider, clientID string) { |
| hook := r.hook |
| if hook == nil { |
| return |
| } |
| go func() { |
| defer func() { |
| if recovered := recover(); recovered != nil { |
| log.Errorf("model registry hook OnModelsUnregistered panic: %v", recovered) |
| } |
| }() |
| ctx, cancel := context.WithTimeout(context.Background(), defaultModelRegistryHookTimeout) |
| defer cancel() |
| hook.OnModelsUnregistered(ctx, provider, clientID) |
| }() |
| } |
|
|
| |
| |
| |
| |
| |
| func (r *ModelRegistry) RegisterClient(clientID, clientProvider string, models []*ModelInfo) { |
| r.mutex.Lock() |
| defer r.mutex.Unlock() |
|
|
| provider := strings.ToLower(clientProvider) |
| uniqueModelIDs := make([]string, 0, len(models)) |
| rawModelIDs := make([]string, 0, len(models)) |
| newModels := make(map[string]*ModelInfo, len(models)) |
| newCounts := make(map[string]int, len(models)) |
| for _, model := range models { |
| if model == nil || model.ID == "" { |
| continue |
| } |
| rawModelIDs = append(rawModelIDs, model.ID) |
| newCounts[model.ID]++ |
| if _, exists := newModels[model.ID]; exists { |
| continue |
| } |
| newModels[model.ID] = model |
| uniqueModelIDs = append(uniqueModelIDs, model.ID) |
| } |
|
|
| if len(uniqueModelIDs) == 0 { |
| |
| r.unregisterClientInternal(clientID) |
| delete(r.clientModels, clientID) |
| delete(r.clientModelInfos, clientID) |
| delete(r.clientProviders, clientID) |
| misc.LogCredentialSeparator() |
| return |
| } |
|
|
| now := time.Now() |
|
|
| oldModels, hadExisting := r.clientModels[clientID] |
| oldProvider := r.clientProviders[clientID] |
| providerChanged := oldProvider != provider |
| if !hadExisting { |
| |
| for _, modelID := range rawModelIDs { |
| model := newModels[modelID] |
| r.addModelRegistration(modelID, provider, model, now) |
| } |
| r.clientModels[clientID] = append([]string(nil), rawModelIDs...) |
| |
| clientInfos := make(map[string]*ModelInfo, len(newModels)) |
| for id, m := range newModels { |
| clientInfos[id] = cloneModelInfo(m) |
| } |
| r.clientModelInfos[clientID] = clientInfos |
| if provider != "" { |
| r.clientProviders[clientID] = provider |
| } else { |
| delete(r.clientProviders, clientID) |
| } |
| r.triggerModelsRegistered(provider, clientID, models) |
| log.Debugf("Registered client %s from provider %s with %d models", clientID, clientProvider, len(rawModelIDs)) |
| misc.LogCredentialSeparator() |
| return |
| } |
|
|
| oldCounts := make(map[string]int, len(oldModels)) |
| for _, id := range oldModels { |
| oldCounts[id]++ |
| } |
|
|
| added := make([]string, 0) |
| for _, id := range uniqueModelIDs { |
| if oldCounts[id] == 0 { |
| added = append(added, id) |
| } |
| } |
|
|
| removed := make([]string, 0) |
| for id := range oldCounts { |
| if newCounts[id] == 0 { |
| removed = append(removed, id) |
| } |
| } |
|
|
| |
| if providerChanged && oldProvider != "" { |
| for id, newCount := range newCounts { |
| if newCount == 0 { |
| continue |
| } |
| oldCount := oldCounts[id] |
| if oldCount == 0 { |
| continue |
| } |
| toRemove := newCount |
| if oldCount < toRemove { |
| toRemove = oldCount |
| } |
| if reg, ok := r.models[id]; ok && reg.Providers != nil { |
| if count, okProv := reg.Providers[oldProvider]; okProv { |
| if count <= toRemove { |
| delete(reg.Providers, oldProvider) |
| if reg.InfoByProvider != nil { |
| delete(reg.InfoByProvider, oldProvider) |
| } |
| } else { |
| reg.Providers[oldProvider] = count - toRemove |
| } |
| } |
| } |
| } |
| } |
|
|
| |
| for _, id := range removed { |
| oldCount := oldCounts[id] |
| for i := 0; i < oldCount; i++ { |
| r.removeModelRegistration(clientID, id, oldProvider, now) |
| } |
| } |
|
|
| for id, oldCount := range oldCounts { |
| newCount := newCounts[id] |
| if newCount == 0 || oldCount <= newCount { |
| continue |
| } |
| overage := oldCount - newCount |
| for i := 0; i < overage; i++ { |
| r.removeModelRegistration(clientID, id, oldProvider, now) |
| } |
| } |
|
|
| |
| for id, newCount := range newCounts { |
| oldCount := oldCounts[id] |
| if newCount <= oldCount { |
| continue |
| } |
| model := newModels[id] |
| diff := newCount - oldCount |
| for i := 0; i < diff; i++ { |
| r.addModelRegistration(id, provider, model, now) |
| } |
| } |
|
|
| |
| addedSet := make(map[string]struct{}, len(added)) |
| for _, id := range added { |
| addedSet[id] = struct{}{} |
| } |
| for _, id := range uniqueModelIDs { |
| model := newModels[id] |
| if reg, ok := r.models[id]; ok { |
| reg.Info = cloneModelInfo(model) |
| if provider != "" { |
| if reg.InfoByProvider == nil { |
| reg.InfoByProvider = make(map[string]*ModelInfo) |
| } |
| reg.InfoByProvider[provider] = cloneModelInfo(model) |
| } |
| reg.LastUpdated = now |
| if reg.QuotaExceededClients != nil { |
| delete(reg.QuotaExceededClients, clientID) |
| } |
| if reg.SuspendedClients != nil { |
| delete(reg.SuspendedClients, clientID) |
| } |
| if providerChanged && provider != "" { |
| if _, newlyAdded := addedSet[id]; newlyAdded { |
| continue |
| } |
| overlapCount := newCounts[id] |
| if oldCount := oldCounts[id]; oldCount < overlapCount { |
| overlapCount = oldCount |
| } |
| if overlapCount <= 0 { |
| continue |
| } |
| if reg.Providers == nil { |
| reg.Providers = make(map[string]int) |
| } |
| reg.Providers[provider] += overlapCount |
| } |
| } |
| } |
|
|
| |
| if len(rawModelIDs) > 0 { |
| r.clientModels[clientID] = append([]string(nil), rawModelIDs...) |
| } |
| |
| clientInfos := make(map[string]*ModelInfo, len(newModels)) |
| for id, m := range newModels { |
| clientInfos[id] = cloneModelInfo(m) |
| } |
| r.clientModelInfos[clientID] = clientInfos |
| if provider != "" { |
| r.clientProviders[clientID] = provider |
| } else { |
| delete(r.clientProviders, clientID) |
| } |
|
|
| r.triggerModelsRegistered(provider, clientID, models) |
| if len(added) == 0 && len(removed) == 0 && !providerChanged { |
| |
| return |
| } |
|
|
| log.Debugf("Reconciled client %s (provider %s) models: +%d, -%d", clientID, provider, len(added), len(removed)) |
| misc.LogCredentialSeparator() |
| } |
|
|
| func (r *ModelRegistry) addModelRegistration(modelID, provider string, model *ModelInfo, now time.Time) { |
| if model == nil || modelID == "" { |
| return |
| } |
| if existing, exists := r.models[modelID]; exists { |
| existing.Count++ |
| existing.LastUpdated = now |
| existing.Info = cloneModelInfo(model) |
| if existing.SuspendedClients == nil { |
| existing.SuspendedClients = make(map[string]string) |
| } |
| if existing.InfoByProvider == nil { |
| existing.InfoByProvider = make(map[string]*ModelInfo) |
| } |
| if provider != "" { |
| if existing.Providers == nil { |
| existing.Providers = make(map[string]int) |
| } |
| existing.Providers[provider]++ |
| existing.InfoByProvider[provider] = cloneModelInfo(model) |
| } |
| log.Debugf("Incremented count for model %s, now %d clients", modelID, existing.Count) |
| return |
| } |
|
|
| registration := &ModelRegistration{ |
| Info: cloneModelInfo(model), |
| InfoByProvider: make(map[string]*ModelInfo), |
| Count: 1, |
| LastUpdated: now, |
| QuotaExceededClients: make(map[string]*time.Time), |
| SuspendedClients: make(map[string]string), |
| } |
| if provider != "" { |
| registration.Providers = map[string]int{provider: 1} |
| registration.InfoByProvider[provider] = cloneModelInfo(model) |
| } |
| r.models[modelID] = registration |
| log.Debugf("Registered new model %s from provider %s", modelID, provider) |
| } |
|
|
| func (r *ModelRegistry) removeModelRegistration(clientID, modelID, provider string, now time.Time) { |
| registration, exists := r.models[modelID] |
| if !exists { |
| return |
| } |
| registration.Count-- |
| registration.LastUpdated = now |
| if registration.QuotaExceededClients != nil { |
| delete(registration.QuotaExceededClients, clientID) |
| } |
| if registration.SuspendedClients != nil { |
| delete(registration.SuspendedClients, clientID) |
| } |
| if registration.Count < 0 { |
| registration.Count = 0 |
| } |
| if provider != "" && registration.Providers != nil { |
| if count, ok := registration.Providers[provider]; ok { |
| if count <= 1 { |
| delete(registration.Providers, provider) |
| if registration.InfoByProvider != nil { |
| delete(registration.InfoByProvider, provider) |
| } |
| } else { |
| registration.Providers[provider] = count - 1 |
| } |
| } |
| } |
| log.Debugf("Decremented count for model %s, now %d clients", modelID, registration.Count) |
| if registration.Count <= 0 { |
| delete(r.models, modelID) |
| log.Debugf("Removed model %s as no clients remain", modelID) |
| } |
| } |
|
|
| func cloneModelInfo(model *ModelInfo) *ModelInfo { |
| if model == nil { |
| return nil |
| } |
| copyModel := *model |
| if len(model.SupportedGenerationMethods) > 0 { |
| copyModel.SupportedGenerationMethods = append([]string(nil), model.SupportedGenerationMethods...) |
| } |
| if len(model.SupportedParameters) > 0 { |
| copyModel.SupportedParameters = append([]string(nil), model.SupportedParameters...) |
| } |
| return ©Model |
| } |
|
|
| func cloneModelInfosUnique(models []*ModelInfo) []*ModelInfo { |
| if len(models) == 0 { |
| return nil |
| } |
| cloned := make([]*ModelInfo, 0, len(models)) |
| seen := make(map[string]struct{}, len(models)) |
| for _, model := range models { |
| if model == nil || model.ID == "" { |
| continue |
| } |
| if _, exists := seen[model.ID]; exists { |
| continue |
| } |
| seen[model.ID] = struct{}{} |
| cloned = append(cloned, cloneModelInfo(model)) |
| } |
| return cloned |
| } |
|
|
| |
| |
| |
| func (r *ModelRegistry) UnregisterClient(clientID string) { |
| r.mutex.Lock() |
| defer r.mutex.Unlock() |
| r.unregisterClientInternal(clientID) |
| } |
|
|
| |
| func (r *ModelRegistry) unregisterClientInternal(clientID string) { |
| models, exists := r.clientModels[clientID] |
| provider, hasProvider := r.clientProviders[clientID] |
| if !exists { |
| if hasProvider { |
| delete(r.clientProviders, clientID) |
| } |
| return |
| } |
|
|
| now := time.Now() |
| for _, modelID := range models { |
| if registration, isExists := r.models[modelID]; isExists { |
| registration.Count-- |
| registration.LastUpdated = now |
|
|
| |
| delete(registration.QuotaExceededClients, clientID) |
| if registration.SuspendedClients != nil { |
| delete(registration.SuspendedClients, clientID) |
| } |
|
|
| if hasProvider && registration.Providers != nil { |
| if count, ok := registration.Providers[provider]; ok { |
| if count <= 1 { |
| delete(registration.Providers, provider) |
| if registration.InfoByProvider != nil { |
| delete(registration.InfoByProvider, provider) |
| } |
| } else { |
| registration.Providers[provider] = count - 1 |
| } |
| } |
| } |
|
|
| log.Debugf("Decremented count for model %s, now %d clients", modelID, registration.Count) |
|
|
| |
| if registration.Count <= 0 { |
| delete(r.models, modelID) |
| log.Debugf("Removed model %s as no clients remain", modelID) |
| } |
| } |
| } |
|
|
| delete(r.clientModels, clientID) |
| delete(r.clientModelInfos, clientID) |
| if hasProvider { |
| delete(r.clientProviders, clientID) |
| } |
| log.Debugf("Unregistered client %s", clientID) |
| |
| misc.LogCredentialSeparator() |
| r.triggerModelsUnregistered(provider, clientID) |
| } |
|
|
| |
| |
| |
| |
| func (r *ModelRegistry) SetModelQuotaExceeded(clientID, modelID string) { |
| r.mutex.Lock() |
| defer r.mutex.Unlock() |
|
|
| if registration, exists := r.models[modelID]; exists { |
| now := time.Now() |
| registration.QuotaExceededClients[clientID] = &now |
| log.Debugf("Marked model %s as quota exceeded for client %s", modelID, clientID) |
| } |
| } |
|
|
| |
| |
| |
| |
| func (r *ModelRegistry) ClearModelQuotaExceeded(clientID, modelID string) { |
| r.mutex.Lock() |
| defer r.mutex.Unlock() |
|
|
| if registration, exists := r.models[modelID]; exists { |
| delete(registration.QuotaExceededClients, clientID) |
| |
| } |
| } |
|
|
| |
| |
| |
| |
| |
| func (r *ModelRegistry) SuspendClientModel(clientID, modelID, reason string) { |
| if clientID == "" || modelID == "" { |
| return |
| } |
| r.mutex.Lock() |
| defer r.mutex.Unlock() |
|
|
| registration, exists := r.models[modelID] |
| if !exists || registration == nil { |
| return |
| } |
| if registration.SuspendedClients == nil { |
| registration.SuspendedClients = make(map[string]string) |
| } |
| if _, already := registration.SuspendedClients[clientID]; already { |
| return |
| } |
| registration.SuspendedClients[clientID] = reason |
| registration.LastUpdated = time.Now() |
| if reason != "" { |
| log.Debugf("Suspended client %s for model %s: %s", clientID, modelID, reason) |
| } else { |
| log.Debugf("Suspended client %s for model %s", clientID, modelID) |
| } |
| } |
|
|
| |
| |
| |
| |
| func (r *ModelRegistry) ResumeClientModel(clientID, modelID string) { |
| if clientID == "" || modelID == "" { |
| return |
| } |
| r.mutex.Lock() |
| defer r.mutex.Unlock() |
|
|
| registration, exists := r.models[modelID] |
| if !exists || registration == nil || registration.SuspendedClients == nil { |
| return |
| } |
| if _, ok := registration.SuspendedClients[clientID]; !ok { |
| return |
| } |
| delete(registration.SuspendedClients, clientID) |
| registration.LastUpdated = time.Now() |
| log.Debugf("Resumed client %s for model %s", clientID, modelID) |
| } |
|
|
| |
| func (r *ModelRegistry) ClientSupportsModel(clientID, modelID string) bool { |
| clientID = strings.TrimSpace(clientID) |
| modelID = strings.TrimSpace(modelID) |
| if clientID == "" || modelID == "" { |
| return false |
| } |
|
|
| r.mutex.RLock() |
| defer r.mutex.RUnlock() |
|
|
| models, exists := r.clientModels[clientID] |
| if !exists || len(models) == 0 { |
| return false |
| } |
|
|
| for _, id := range models { |
| if strings.EqualFold(strings.TrimSpace(id), modelID) { |
| return true |
| } |
| } |
|
|
| return false |
| } |
|
|
| |
| |
| |
| |
| |
| |
| func (r *ModelRegistry) GetAvailableModels(handlerType string) []map[string]any { |
| r.mutex.RLock() |
| defer r.mutex.RUnlock() |
|
|
| models := make([]map[string]any, 0) |
| quotaExpiredDuration := 5 * time.Minute |
|
|
| for _, registration := range r.models { |
| |
| availableClients := registration.Count |
| now := time.Now() |
|
|
| |
| expiredClients := 0 |
| for _, quotaTime := range registration.QuotaExceededClients { |
| if quotaTime != nil && now.Sub(*quotaTime) < quotaExpiredDuration { |
| expiredClients++ |
| } |
| } |
|
|
| cooldownSuspended := 0 |
| otherSuspended := 0 |
| if registration.SuspendedClients != nil { |
| for _, reason := range registration.SuspendedClients { |
| if strings.EqualFold(reason, "quota") { |
| cooldownSuspended++ |
| continue |
| } |
| otherSuspended++ |
| } |
| } |
|
|
| effectiveClients := availableClients - expiredClients - otherSuspended |
| if effectiveClients < 0 { |
| effectiveClients = 0 |
| } |
|
|
| |
| if effectiveClients > 0 || (availableClients > 0 && (expiredClients > 0 || cooldownSuspended > 0) && otherSuspended == 0) { |
| model := r.convertModelToMap(registration.Info, handlerType) |
| if model != nil { |
| models = append(models, model) |
| } |
| } |
| } |
|
|
| return models |
| } |
|
|
| |
| |
| |
| |
| |
| |
| func (r *ModelRegistry) GetAvailableModelsByProvider(provider string) []*ModelInfo { |
| provider = strings.ToLower(strings.TrimSpace(provider)) |
| if provider == "" { |
| return nil |
| } |
|
|
| r.mutex.RLock() |
| defer r.mutex.RUnlock() |
|
|
| type providerModel struct { |
| count int |
| info *ModelInfo |
| } |
|
|
| providerModels := make(map[string]*providerModel) |
|
|
| for clientID, clientProvider := range r.clientProviders { |
| if clientProvider != provider { |
| continue |
| } |
| modelIDs := r.clientModels[clientID] |
| if len(modelIDs) == 0 { |
| continue |
| } |
| clientInfos := r.clientModelInfos[clientID] |
| for _, modelID := range modelIDs { |
| modelID = strings.TrimSpace(modelID) |
| if modelID == "" { |
| continue |
| } |
| entry := providerModels[modelID] |
| if entry == nil { |
| entry = &providerModel{} |
| providerModels[modelID] = entry |
| } |
| entry.count++ |
| if entry.info == nil { |
| if clientInfos != nil { |
| if info := clientInfos[modelID]; info != nil { |
| entry.info = info |
| } |
| } |
| if entry.info == nil { |
| if reg, ok := r.models[modelID]; ok && reg != nil && reg.Info != nil { |
| entry.info = reg.Info |
| } |
| } |
| } |
| } |
| } |
|
|
| if len(providerModels) == 0 { |
| return nil |
| } |
|
|
| quotaExpiredDuration := 5 * time.Minute |
| now := time.Now() |
| result := make([]*ModelInfo, 0, len(providerModels)) |
|
|
| for modelID, entry := range providerModels { |
| if entry == nil || entry.count <= 0 { |
| continue |
| } |
| registration, ok := r.models[modelID] |
|
|
| expiredClients := 0 |
| cooldownSuspended := 0 |
| otherSuspended := 0 |
| if ok && registration != nil { |
| if registration.QuotaExceededClients != nil { |
| for clientID, quotaTime := range registration.QuotaExceededClients { |
| if clientID == "" { |
| continue |
| } |
| if p, okProvider := r.clientProviders[clientID]; !okProvider || p != provider { |
| continue |
| } |
| if quotaTime != nil && now.Sub(*quotaTime) < quotaExpiredDuration { |
| expiredClients++ |
| } |
| } |
| } |
| if registration.SuspendedClients != nil { |
| for clientID, reason := range registration.SuspendedClients { |
| if clientID == "" { |
| continue |
| } |
| if p, okProvider := r.clientProviders[clientID]; !okProvider || p != provider { |
| continue |
| } |
| if strings.EqualFold(reason, "quota") { |
| cooldownSuspended++ |
| continue |
| } |
| otherSuspended++ |
| } |
| } |
| } |
|
|
| availableClients := entry.count |
| effectiveClients := availableClients - expiredClients - otherSuspended |
| if effectiveClients < 0 { |
| effectiveClients = 0 |
| } |
|
|
| if effectiveClients > 0 || (availableClients > 0 && (expiredClients > 0 || cooldownSuspended > 0) && otherSuspended == 0) { |
| if entry.info != nil { |
| result = append(result, entry.info) |
| continue |
| } |
| if ok && registration != nil && registration.Info != nil { |
| result = append(result, registration.Info) |
| } |
| } |
| } |
|
|
| return result |
| } |
|
|
| |
| |
| |
| |
| |
| |
| func (r *ModelRegistry) GetModelCount(modelID string) int { |
| r.mutex.RLock() |
| defer r.mutex.RUnlock() |
|
|
| if registration, exists := r.models[modelID]; exists { |
| now := time.Now() |
| quotaExpiredDuration := 5 * time.Minute |
|
|
| |
| expiredClients := 0 |
| for _, quotaTime := range registration.QuotaExceededClients { |
| if quotaTime != nil && now.Sub(*quotaTime) < quotaExpiredDuration { |
| expiredClients++ |
| } |
| } |
| suspendedClients := 0 |
| if registration.SuspendedClients != nil { |
| suspendedClients = len(registration.SuspendedClients) |
| } |
| result := registration.Count - expiredClients - suspendedClients |
| if result < 0 { |
| return 0 |
| } |
| return result |
| } |
| return 0 |
| } |
|
|
| |
| |
| |
| |
| |
| |
| func (r *ModelRegistry) GetModelProviders(modelID string) []string { |
| r.mutex.RLock() |
| defer r.mutex.RUnlock() |
|
|
| registration, exists := r.models[modelID] |
| if !exists || registration == nil || len(registration.Providers) == 0 { |
| return nil |
| } |
|
|
| type providerCount struct { |
| name string |
| count int |
| } |
| providers := make([]providerCount, 0, len(registration.Providers)) |
| |
| |
| |
| |
| |
| |
| |
| |
| for name, count := range registration.Providers { |
| if count <= 0 { |
| continue |
| } |
| |
| |
| |
| |
| |
| providers = append(providers, providerCount{name: name, count: count}) |
| } |
| if len(providers) == 0 { |
| return nil |
| } |
|
|
| sort.Slice(providers, func(i, j int) bool { |
| if providers[i].count == providers[j].count { |
| return providers[i].name < providers[j].name |
| } |
| return providers[i].count > providers[j].count |
| }) |
|
|
| result := make([]string, 0, len(providers)) |
| for _, item := range providers { |
| result = append(result, item.name) |
| } |
| return result |
| } |
|
|
| |
| func (r *ModelRegistry) GetModelInfo(modelID, provider string) *ModelInfo { |
| r.mutex.RLock() |
| defer r.mutex.RUnlock() |
| if reg, ok := r.models[modelID]; ok && reg != nil { |
| |
| if provider != "" && reg.InfoByProvider != nil { |
| if reg.Providers != nil { |
| if count, ok := reg.Providers[provider]; ok && count > 0 { |
| if info, ok := reg.InfoByProvider[provider]; ok && info != nil { |
| return info |
| } |
| } |
| } |
| } |
| |
| return reg.Info |
| } |
| return nil |
| } |
|
|
| |
| func (r *ModelRegistry) convertModelToMap(model *ModelInfo, handlerType string) map[string]any { |
| if model == nil { |
| return nil |
| } |
|
|
| switch handlerType { |
| case "openai": |
| result := map[string]any{ |
| "id": model.ID, |
| "object": "model", |
| "owned_by": model.OwnedBy, |
| } |
| if model.Created > 0 { |
| result["created"] = model.Created |
| } |
| if model.Type != "" { |
| result["type"] = model.Type |
| } |
| if model.DisplayName != "" { |
| result["display_name"] = model.DisplayName |
| } |
| if model.Version != "" { |
| result["version"] = model.Version |
| } |
| if model.Description != "" { |
| result["description"] = model.Description |
| } |
| if model.ContextLength > 0 { |
| result["context_length"] = model.ContextLength |
| |
| result["context_window"] = model.ContextLength |
| result["max_context_length"] = model.ContextLength |
| } |
| if model.MaxCompletionTokens > 0 { |
| result["max_completion_tokens"] = model.MaxCompletionTokens |
| } |
| if len(model.SupportedParameters) > 0 { |
| result["supported_parameters"] = model.SupportedParameters |
| } |
| return result |
|
|
| case "claude": |
| result := map[string]any{ |
| "id": model.ID, |
| "object": "model", |
| "owned_by": model.OwnedBy, |
| } |
| if model.Created > 0 { |
| result["created_at"] = model.Created |
| } |
| if model.Type != "" { |
| result["type"] = "model" |
| } |
| if model.DisplayName != "" { |
| result["display_name"] = model.DisplayName |
| } |
| return result |
|
|
| case "gemini": |
| result := map[string]any{} |
| if model.Name != "" { |
| result["name"] = model.Name |
| } else { |
| result["name"] = model.ID |
| } |
| if model.Version != "" { |
| result["version"] = model.Version |
| } |
| if model.DisplayName != "" { |
| result["displayName"] = model.DisplayName |
| } |
| if model.Description != "" { |
| result["description"] = model.Description |
| } |
| if model.ContextLength > 0 { |
| result["inputTokenLimit"] = model.ContextLength |
| } |
| if model.MaxCompletionTokens > 0 { |
| result["outputTokenLimit"] = model.MaxCompletionTokens |
| } |
| if len(model.SupportedGenerationMethods) > 0 { |
| result["supportedGenerationMethods"] = model.SupportedGenerationMethods |
| } |
| return result |
|
|
| default: |
| |
| result := map[string]any{ |
| "id": model.ID, |
| "object": "model", |
| } |
| if model.OwnedBy != "" { |
| result["owned_by"] = model.OwnedBy |
| } |
| if model.Type != "" { |
| result["type"] = model.Type |
| } |
| if model.Created != 0 { |
| result["created"] = model.Created |
| } |
| return result |
| } |
| } |
|
|
| |
| func (r *ModelRegistry) CleanupExpiredQuotas() { |
| r.mutex.Lock() |
| defer r.mutex.Unlock() |
|
|
| now := time.Now() |
| quotaExpiredDuration := 5 * time.Minute |
|
|
| for modelID, registration := range r.models { |
| for clientID, quotaTime := range registration.QuotaExceededClients { |
| if quotaTime != nil && now.Sub(*quotaTime) >= quotaExpiredDuration { |
| delete(registration.QuotaExceededClients, clientID) |
| log.Debugf("Cleaned up expired quota tracking for model %s, client %s", modelID, clientID) |
| } |
| } |
| } |
| } |
|
|
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| func (r *ModelRegistry) GetFirstAvailableModel(handlerType string) (string, error) { |
| r.mutex.RLock() |
| defer r.mutex.RUnlock() |
|
|
| |
| models := r.GetAvailableModels(handlerType) |
| if len(models) == 0 { |
| return "", fmt.Errorf("no models available for handler type: %s", handlerType) |
| } |
|
|
| |
| sort.Slice(models, func(i, j int) bool { |
| |
| createdI, okI := models[i]["created"].(int64) |
| createdJ, okJ := models[j]["created"].(int64) |
| if !okI || !okJ { |
| return false |
| } |
| return createdI > createdJ |
| }) |
|
|
| |
| for _, model := range models { |
| if modelID, ok := model["id"].(string); ok { |
| if count := r.GetModelCount(modelID); count > 0 { |
| return modelID, nil |
| } |
| } |
| } |
|
|
| return "", fmt.Errorf("no available clients for any model in handler type: %s", handlerType) |
| } |
|
|
| |
| |
| |
| |
| |
| |
| func (r *ModelRegistry) GetModelsForClient(clientID string) []*ModelInfo { |
| r.mutex.RLock() |
| defer r.mutex.RUnlock() |
|
|
| modelIDs, exists := r.clientModels[clientID] |
| if !exists || len(modelIDs) == 0 { |
| return nil |
| } |
|
|
| |
| clientInfos := r.clientModelInfos[clientID] |
|
|
| seen := make(map[string]struct{}) |
| result := make([]*ModelInfo, 0, len(modelIDs)) |
| for _, modelID := range modelIDs { |
| if _, dup := seen[modelID]; dup { |
| continue |
| } |
| seen[modelID] = struct{}{} |
|
|
| |
| if clientInfos != nil { |
| if info, ok := clientInfos[modelID]; ok && info != nil { |
| result = append(result, info) |
| continue |
| } |
| } |
| |
| if reg, ok := r.models[modelID]; ok && reg.Info != nil { |
| result = append(result, reg.Info) |
| } |
| } |
| return result |
| } |
|
|