History

Merge pull request #237 from TUGOhost/feature/support_auto_model

Luis Pater committed 1 year ago

Commit 09ecba62 parents d6bd6f3 92f42783 files · +78 −2Browse files

Unified · Split

3 files
  1. internal/registry/model_registry.go +48 −0
  2. internal/util/provider.go +25 −0
  3. sdk/api/handlers/handlers.go +5 −2
internal/registry/model_registry.go+48−0
@@ -4,6 +4,7 @@ · 3 unmodified lines
4package registry4package registry
5 5
6import (6import (
7 "fmt"
7 "sort"8 "sort"
8 "strings"9 "strings"
9 "sync"10 "sync"
@@ -800,6 +801,9 @@ func (r *ModelRegistry) convertModelToMap(model *ModelInfo, handlerType string) · 790 unmodified lines
800 if model.Type != "" {801 if model.Type != "" {
801 result["type"] = model.Type802 result["type"] = model.Type
802 }803 }
804 if model.Created != 0 {
805 result["created"] = model.Created
806 }
803 return result807 return result
804 }808 }
805}809}
@@ -821,3 +825,47 @@ func (r *ModelRegistry) CleanupExpiredQuotas() { · 15 unmodified lines
821 }825 }
822 }826 }
823}827}
828
829
830// GetFirstAvailableModel returns the first available model for the given handler type.
831// It prioritizes models by their creation timestamp (newest first) and checks if they have
832// available clients that are not suspended or over quota.
833//
834// Parameters:
835// - handlerType: The API handler type (e.g., "openai", "claude", "gemini")
836//
837// Returns:
838// - string: The model ID of the first available model, or empty string if none available
839// - error: An error if no models are available
840func (r *ModelRegistry) GetFirstAvailableModel(handlerType string) (string, error) {
841 r.mutex.RLock()
842 defer r.mutex.RUnlock()
843
844 // Get all available models for this handler type
845 models := r.GetAvailableModels(handlerType)
846 if len(models) == 0 {
847 return "", fmt.Errorf("no models available for handler type: %s", handlerType)
848 }
849
850 // Sort models by creation timestamp (newest first)
851 sort.Slice(models, func(i, j int) bool {
852 // Extract created timestamps from map
853 createdI, okI := models[i]["created"].(int64)
854 createdJ, okJ := models[j]["created"].(int64)
855 if !okI || !okJ {
856 return false
857 }
858 return createdI > createdJ
859 })
860
861 // Find the first model with available clients
862 for _, model := range models {
863 if modelID, ok := model["id"].(string); ok {
864 if count := r.GetModelCount(modelID); count > 0 {
865 return modelID, nil
866 }
867 }
868 }
869
870 return "", fmt.Errorf("no available clients for any model in handler type: %s", handlerType)
871}
internal/util/provider.go+25−0
@@ -9,6 +9,7 @@ import ( · 8 unmodified lines
9 9
10 "github.com/router-for-me/CLIProxyAPI/v6/internal/config"10 "github.com/router-for-me/CLIProxyAPI/v6/internal/config"
11 "github.com/router-for-me/CLIProxyAPI/v6/internal/registry"11 "github.com/router-for-me/CLIProxyAPI/v6/internal/registry"
12 log "github.com/sirupsen/logrus"
12)13)
13 14
14// GetProviderName determines all AI service providers capable of serving a registered model.15// GetProviderName determines all AI service providers capable of serving a registered model.
@@ -59,6 +60,30 @@ func GetProviderName(modelName string) []string { · 44 unmodified lines
59 return providers60 return providers
60}61}
61 62
63// ResolveAutoModel resolves the "auto" model name to an actual available model.
64// It uses an empty handler type to get any available model from the registry.
65//
66// Parameters:
67// - modelName: The model name to check (should be "auto")
68//
69// Returns:
70// - string: The resolved model name, or the original if not "auto" or resolution fails
71func ResolveAutoModel(modelName string) string {
72 if modelName != "auto" {
73 return modelName
74 }
75
76 // Use empty string as handler type to get any available model
77 firstModel, err := registry.GetGlobalRegistry().GetFirstAvailableModel("")
78 if err != nil {
79 log.Warnf("Failed to resolve 'auto' model: %v, falling back to original model name", err)
80 return modelName
81 }
82
83 log.Infof("Resolved 'auto' model to: %s", firstModel)
84 return firstModel
85}
86
62// IsOpenAICompatibilityAlias checks if the given model name is an alias87// IsOpenAICompatibilityAlias checks if the given model name is an alias
63// configured for OpenAI compatibility routing.88// configured for OpenAI compatibility routing.
64//89//
sdk/api/handlers/handlers.go+5−2
@@ -295,11 +295,14 @@ func (h *BaseAPIHandler) ExecuteStreamWithAuthManager(ctx context.Context, handl · 294 unmodified lines
295}295}
296 296
297func (h *BaseAPIHandler) getRequestDetails(modelName string) (providers []string, normalizedModel string, metadata map[string]any, err *interfaces.ErrorMessage) {297func (h *BaseAPIHandler) getRequestDetails(modelName string) (providers []string, normalizedModel string, metadata map[string]any, err *interfaces.ErrorMessage) {
298 providerName, extractedModelName, isDynamic := h.parseDynamicModel(modelName)298 // Resolve "auto" model to an actual available model first
299 resolvedModelName := util.ResolveAutoModel(modelName)
300
301 providerName, extractedModelName, isDynamic := h.parseDynamicModel(resolvedModelName)
299 302
300 // First, normalize the model name to handle suffixes like "-thinking-128"303 // First, normalize the model name to handle suffixes like "-thinking-128"
301 // This needs to happen before determining the provider for non-dynamic models.304 // This needs to happen before determining the provider for non-dynamic models.
302 normalizedModel, metadata = normalizeModelMetadata(modelName)305 normalizedModel, metadata = normalizeModelMetadata(resolvedModelName)
303 306
304 if isDynamic {307 if isDynamic {
305 providers = []string{providerName}308 providers = []string{providerName}