diff --git a/internal/modeldata/fetcher_test.go b/internal/modeldata/fetcher_test.go index 00501eb8..da9c52db 100644 --- a/internal/modeldata/fetcher_test.go +++ b/internal/modeldata/fetcher_test.go @@ -133,14 +133,20 @@ func TestParse_BuildsReverseIndex(t *testing.T) { raw := []byte(`{ "version": 1, "updated_at": "2025-01-01T00:00:00Z", - "providers": {}, + "providers": { + "openai": {"display_name": "OpenAI"} + }, "models": { - "gpt-4o": {"display_name": "GPT-4o", "modes": ["chat"]} + "gpt-4o": { + "display_name": "GPT-4o", + "modes": ["chat"], + "aliases": ["gpt-4o-latest", "openai/gpt-4o-latest"] + } }, "provider_models": { "openai/gpt-4o": { "model_ref": "gpt-4o", - "custom_model_id": "gpt-4o-2024-08-06", + "provider_model_id": "gpt-4o-2024-08-06", "enabled": true } } @@ -159,6 +165,29 @@ func TestParse_BuildsReverseIndex(t *testing.T) { if compositeKey != "openai/gpt-4o" { t.Errorf("reverse index = %s, want openai/gpt-4o", compositeKey) } + targets := list.aliasTargetsByID["gpt-4o-latest"] + if len(targets) != 2 { + t.Fatalf("expected 2 alias targets for gpt-4o-latest, got %d", len(targets)) + } + var sawGeneric bool + var sawProviderSpecific bool + for _, target := range targets { + if target.ModelRef != "gpt-4o" { + t.Fatalf("alias target ModelRef = %q, want gpt-4o", target.ModelRef) + } + if target.ProviderType == "" { + sawGeneric = true + } + if target.ProviderType == "openai" { + sawProviderSpecific = true + } + } + if !sawGeneric { + t.Fatal("expected generic alias target for gpt-4o-latest") + } + if !sawProviderSpecific { + t.Fatal("expected provider-qualified alias target for gpt-4o-latest") + } } func TestParse_InvalidJSON(t *testing.T) { diff --git a/internal/modeldata/merger.go b/internal/modeldata/merger.go index b487bed9..2dc364f4 100644 --- a/internal/modeldata/merger.go +++ b/internal/modeldata/merger.go @@ -14,52 +14,141 @@ func Resolve(list *ModelList, providerType string, modelID string) *core.ModelMe return nil } - // Try provider_model lookup first: "providerType/modelID" - var pm *ProviderModelEntry - key := providerType + "/" + modelID - if entry, ok := list.ProviderModels[key]; ok { - pm = &entry + model, pm := resolveEntries(list, providerType, modelID) + if model == nil && pm == nil { + return nil } - // Determine the base model entry - var model *ModelEntry - if pm != nil { - // Use model_ref from provider_model to find the base model - if entry, ok := list.Models[pm.ModelRef]; ok { - model = &entry + return buildMetadata(model, pm) +} + +func resolveEntries(list *ModelList, providerType string, modelID string) (*ModelEntry, *ProviderModelEntry) { + if model, pm := resolveDirect(list, providerType, modelID); model != nil || pm != nil { + return model, pm + } + if model, pm := resolveReverseProviderModelID(list, providerType, modelID); model != nil || pm != nil { + return model, pm + } + return resolveAlias(list, providerType, modelID) +} + +func resolveDirect(list *ModelList, providerType string, modelID string) (*ModelEntry, *ProviderModelEntry) { + if providerType != "" { + if entry, ok := list.ProviderModels[providerType+"/"+modelID]; ok { + pm := entry + return resolveModelRef(list, providerType, pm.ModelRef, &pm) + } + } + if entry, ok := list.Models[modelID]; ok { + model := entry + return &model, nil + } + return nil, nil +} + +func resolveReverseProviderModelID(list *ModelList, providerType string, modelID string) (*ModelEntry, *ProviderModelEntry) { + if providerType == "" || list.providerModelByActualID == nil { + return nil, nil + } + compositeKey, ok := list.providerModelByActualID[providerType+"/"+modelID] + if !ok { + return nil, nil + } + pmEntry, ok := list.ProviderModels[compositeKey] + if !ok { + return nil, nil + } + pm := pmEntry + return resolveModelRef(list, providerType, pm.ModelRef, &pm) +} + +func resolveAlias(list *ModelList, providerType string, modelID string) (*ModelEntry, *ProviderModelEntry) { + if list.aliasTargetsByID == nil { + return nil, nil + } + targets := list.aliasTargetsByID[modelID] + if len(targets) == 0 { + return nil, nil + } + modelRef, ok := selectAliasModelRef(list, providerType, targets) + if !ok { + return nil, nil + } + return resolveModelRef(list, providerType, modelRef, nil) +} + +func selectAliasModelRef(list *ModelList, providerType string, targets []aliasTarget) (string, bool) { + if len(targets) == 0 { + return "", false + } + + bestScoreByModelRef := make(map[string]int, len(targets)) + for _, target := range targets { + score := aliasTargetScore(providerType, target) + if score == 0 { + continue } - } else { - // Fall back to direct model ID lookup - if entry, ok := list.Models[modelID]; ok { - model = &entry + if score > bestScoreByModelRef[target.ModelRef] { + bestScoreByModelRef[target.ModelRef] = score } } + if len(bestScoreByModelRef) == 0 { + return "", false + } - // No match at all — try reverse lookup via provider_model_id. - // Inline the lookup instead of recursing to avoid unbounded recursion - // if the reverse index contains cycles or stale entries. - if model == nil && pm == nil { - if list.providerModelByActualID != nil { - reverseKey := providerType + "/" + modelID - if compositeKey, ok := list.providerModelByActualID[reverseKey]; ok { - canonicalModelID := compositeKey[len(providerType)+1:] - // Look up the provider_model and base model directly - if entry, ok := list.ProviderModels[compositeKey]; ok { - pm = &entry - if baseEntry, ok := list.Models[entry.ModelRef]; ok { - model = &baseEntry - } - } else if baseEntry, ok := list.Models[canonicalModelID]; ok { - model = &baseEntry - } + bestScore := 0 + bestRefs := make([]string, 0, len(bestScoreByModelRef)) + for modelRef, score := range bestScoreByModelRef { + switch { + case score > bestScore: + bestScore = score + bestRefs = []string{modelRef} + case score == bestScore: + bestRefs = append(bestRefs, modelRef) + } + } + if len(bestRefs) == 1 { + return bestRefs[0], true + } + + if providerType != "" { + withProviderOverride := make([]string, 0, len(bestRefs)) + for _, modelRef := range bestRefs { + if _, ok := list.ProviderModels[providerType+"/"+modelRef]; ok { + withProviderOverride = append(withProviderOverride, modelRef) } } - if model == nil && pm == nil { - return nil + if len(withProviderOverride) == 1 { + return withProviderOverride[0], true } } - return buildMetadata(model, pm) + return "", false +} + +func aliasTargetScore(providerType string, target aliasTarget) int { + switch { + case target.ProviderType == "": + return 1 + case providerType != "" && target.ProviderType == providerType: + return 2 + default: + return 0 + } +} + +func resolveModelRef(list *ModelList, providerType, modelRef string, pm *ProviderModelEntry) (*ModelEntry, *ProviderModelEntry) { + if pm == nil && providerType != "" { + if entry, ok := list.ProviderModels[providerType+"/"+modelRef]; ok { + providerModel := entry + pm = &providerModel + } + } + if entry, ok := list.Models[modelRef]; ok { + model := entry + return &model, pm + } + return nil, pm } // buildMetadata merges base model fields with provider-model overrides into ModelMetadata. diff --git a/internal/modeldata/merger_test.go b/internal/modeldata/merger_test.go index a7b6b55f..b42d6480 100644 --- a/internal/modeldata/merger_test.go +++ b/internal/modeldata/merger_test.go @@ -421,3 +421,74 @@ func TestResolve_ReverseIndexWithProviderModelOverride(t *testing.T) { t.Errorf("OutputPerMtok = %f, want 12.00 (provider override)", *meta.Pricing.OutputPerMtok) } } + +func TestResolve_ModelAliasUsesProviderOverride(t *testing.T) { + list := &ModelList{ + Providers: map[string]ProviderEntry{ + "gemini": {DisplayName: "Gemini"}, + }, + Models: map[string]ModelEntry{ + "claude-4-opus": { + DisplayName: "Claude 4 Opus", + Modes: []string{"chat"}, + Aliases: []string{"claude-opus-4", "gemini/claude-opus-4"}, + Pricing: &core.ModelPricing{ + Currency: "USD", + InputPerMtok: new(15.00), + OutputPerMtok: new(75.00), + }, + }, + }, + ProviderModels: map[string]ProviderModelEntry{ + "gemini/claude-4-opus": { + ModelRef: "claude-4-opus", + Enabled: true, + ContextWindow: new(200000), + Pricing: &core.ModelPricing{ + Currency: "USD", + InputPerMtok: new(12.00), + OutputPerMtok: new(60.00), + }, + }, + }, + } + list.buildReverseIndex() + + meta := Resolve(list, "gemini", "claude-opus-4") + if meta == nil { + t.Fatal("expected non-nil metadata via model alias") + } + if meta.DisplayName != "Claude 4 Opus" { + t.Fatalf("DisplayName = %q, want Claude 4 Opus", meta.DisplayName) + } + if meta.ContextWindow == nil || *meta.ContextWindow != 200000 { + t.Fatalf("ContextWindow = %v, want 200000", meta.ContextWindow) + } + if meta.Pricing == nil || meta.Pricing.InputPerMtok == nil || *meta.Pricing.InputPerMtok != 12.00 { + t.Fatalf("InputPerMtok = %#v, want 12.00", meta.Pricing) + } +} + +func TestResolve_AmbiguousModelAliasReturnsNil(t *testing.T) { + list := &ModelList{ + Models: map[string]ModelEntry{ + "model-a": { + DisplayName: "Model A", + Modes: []string{"chat"}, + Aliases: []string{"shared-alias"}, + }, + "model-b": { + DisplayName: "Model B", + Modes: []string{"chat"}, + Aliases: []string{"shared-alias"}, + }, + }, + ProviderModels: map[string]ProviderModelEntry{}, + } + list.buildReverseIndex() + + meta := Resolve(list, "openai", "shared-alias") + if meta != nil { + t.Fatalf("expected nil for ambiguous alias, got %+v", meta) + } +} diff --git a/internal/modeldata/types.go b/internal/modeldata/types.go index 365f26b9..d5906649 100644 --- a/internal/modeldata/types.go +++ b/internal/modeldata/types.go @@ -8,28 +8,53 @@ import ( "gomodel/internal/core" ) - // ModelList represents the top-level structure of models.json. type ModelList struct { - Version int `json:"version"` - UpdatedAt string `json:"updated_at"` - Providers map[string]ProviderEntry `json:"providers"` - Models map[string]ModelEntry `json:"models"` + Version int `json:"version"` + UpdatedAt string `json:"updated_at"` + Providers map[string]ProviderEntry `json:"providers"` + Models map[string]ModelEntry `json:"models"` ProviderModels map[string]ProviderModelEntry `json:"provider_models"` - // providerModelByActualID maps "providerType/actualModelID" → composite key + // providerModelByActualID maps "providerType/actualModelID" -> composite key // in ProviderModels, enabling reverse lookup from a provider's response model // (e.g., "openai/gpt-4o-2024-08-06") to the canonical registry key // (e.g., "openai/gpt-4o"). Built by buildReverseIndex(). providerModelByActualID map[string]string + + // aliasTargetsByID maps a non-canonical model ID to one or more canonical + // model refs. ProviderType is set when the alias is provider-qualified in the + // external model list, e.g. "gemini/claude-opus-4" -> alias key + // "claude-opus-4" with ProviderType "gemini". + aliasTargetsByID map[string][]aliasTarget +} + +type aliasTarget struct { + ModelRef string + ProviderType string } // buildReverseIndex populates providerModelByActualID from ProviderModels entries -// that have a non-nil custom_model_id differing from the key's model portion. +// and aliasTargetsByID from model aliases. func (l *ModelList) buildReverseIndex() { l.providerModelByActualID = make(map[string]string) + l.aliasTargetsByID = make(map[string][]aliasTarget) + + for modelRef, model := range l.Models { + for _, alias := range model.Aliases { + l.addAliasTarget(alias, aliasTarget{ModelRef: modelRef}) + if providerType, aliasID, ok := l.splitProviderQualifiedAlias(alias); ok { + l.addAliasTarget(aliasID, aliasTarget{ + ModelRef: modelRef, + ProviderType: providerType, + }) + } + } + } + for compositeKey, pm := range l.ProviderModels { - if pm.CustomModelID == nil { + actualID := pm.actualModelID() + if actualID == "" { continue } // compositeKey is "providerType/modelID" @@ -37,7 +62,6 @@ func (l *ModelList) buildReverseIndex() { if !ok { continue } - actualID := *pm.CustomModelID reverseKey := providerType + "/" + actualID // Only add if the actual ID differs from the key's model portion if reverseKey != compositeKey { @@ -46,6 +70,30 @@ func (l *ModelList) buildReverseIndex() { } } +func (l *ModelList) addAliasTarget(alias string, target aliasTarget) { + if alias == "" { + return + } + existing := l.aliasTargetsByID[alias] + for _, candidate := range existing { + if candidate == target { + return + } + } + l.aliasTargetsByID[alias] = append(existing, target) +} + +func (l *ModelList) splitProviderQualifiedAlias(alias string) (providerType, modelID string, ok bool) { + providerType, modelID, ok = strings.Cut(alias, "/") + if !ok || providerType == "" || modelID == "" { + return "", "", false + } + if _, exists := l.Providers[providerType]; !exists { + return "", "", false + } + return providerType, modelID, true +} + // ProviderEntry represents a provider in the registry. type ProviderEntry struct { DisplayName string `json:"display_name"` @@ -63,43 +111,56 @@ type ProviderEntry struct { // ModelEntry represents a model in the registry. type ModelEntry struct { - DisplayName string `json:"display_name"` - Description *string `json:"description"` - OwnedBy *string `json:"owned_by"` - Family *string `json:"family"` - ReleaseDate *string `json:"release_date"` - DeprecationDate *string `json:"deprecation_date"` - Tags []string `json:"tags"` - Modes []string `json:"modes"` - SourceURL *string `json:"source_url"` - Modalities *Modalities `json:"modalities"` - Capabilities map[string]bool `json:"capabilities"` - ContextWindow *int `json:"context_window"` - MaxOutputTokens *int `json:"max_output_tokens"` - MaxImagesPerRequest *int `json:"max_images_per_request"` - MaxVideosPerRequest *int `json:"max_videos_per_request"` - MaxAudioPerRequest *int `json:"max_audio_per_request"` - MaxAudioLengthSeconds *int `json:"max_audio_length_seconds"` - MaxVideoLengthSeconds *int `json:"max_video_length_seconds"` - MaxPDFSizeMB *int `json:"max_pdf_size_mb"` - OutputVectorSize *int `json:"output_vector_size"` - Parameters map[string]ParameterSpec `json:"parameters"` - Rankings map[string]RankingEntry `json:"rankings"` - Pricing *core.ModelPricing `json:"pricing"` + DisplayName string `json:"display_name"` + Description *string `json:"description"` + OwnedBy *string `json:"owned_by"` + Family *string `json:"family"` + ReleaseDate *string `json:"release_date"` + DeprecationDate *string `json:"deprecation_date"` + Tags []string `json:"tags"` + Modes []string `json:"modes"` + SourceURL *string `json:"source_url"` + Modalities *Modalities `json:"modalities"` + Capabilities map[string]bool `json:"capabilities"` + ContextWindow *int `json:"context_window"` + MaxOutputTokens *int `json:"max_output_tokens"` + MaxImagesPerRequest *int `json:"max_images_per_request"` + MaxVideosPerRequest *int `json:"max_videos_per_request"` + MaxAudioPerRequest *int `json:"max_audio_per_request"` + MaxAudioLengthSeconds *int `json:"max_audio_length_seconds"` + MaxVideoLengthSeconds *int `json:"max_video_length_seconds"` + MaxPDFSizeMB *int `json:"max_pdf_size_mb"` + OutputVectorSize *int `json:"output_vector_size"` + Parameters map[string]ParameterSpec `json:"parameters"` + Rankings map[string]RankingEntry `json:"rankings"` + Pricing *core.ModelPricing `json:"pricing"` + Aliases []string `json:"aliases"` } // ProviderModelEntry represents a provider-specific model override. type ProviderModelEntry struct { - ModelRef string `json:"model_ref"` + ModelRef string `json:"model_ref"` + ProviderModelID *string `json:"provider_model_id"` + // CustomModelID is retained for compatibility with older cached payloads. CustomModelID *string `json:"custom_model_id"` Enabled bool `json:"enabled"` Pricing *core.ModelPricing `json:"pricing"` - ContextWindow *int `json:"context_window"` - MaxOutputTokens *int `json:"max_output_tokens"` - Capabilities map[string]bool `json:"capabilities"` - RateLimits *RateLimits `json:"rate_limits"` - Endpoints []string `json:"endpoints"` - Regions []string `json:"regions"` + ContextWindow *int `json:"context_window"` + MaxOutputTokens *int `json:"max_output_tokens"` + Capabilities map[string]bool `json:"capabilities"` + RateLimits *RateLimits `json:"rate_limits"` + Endpoints []string `json:"endpoints"` + Regions []string `json:"regions"` +} + +func (e ProviderModelEntry) actualModelID() string { + if e.ProviderModelID != nil && *e.ProviderModelID != "" { + return *e.ProviderModelID + } + if e.CustomModelID != nil && *e.CustomModelID != "" { + return *e.CustomModelID + } + return "" } // ParameterSpec describes a model parameter's constraints. diff --git a/internal/providers/registry_test.go b/internal/providers/registry_test.go index a90037f1..8fa0b9e0 100644 --- a/internal/providers/registry_test.go +++ b/internal/providers/registry_test.go @@ -303,6 +303,84 @@ func TestModelRegistry(t *testing.T) { } }) + t.Run("EnrichModelsUsesAliasesWithoutAddingSyntheticModels", func(t *testing.T) { + registry := NewModelRegistry() + + mock := ®istryMockProvider{ + name: "gemini-provider", + modelsResponse: &core.ModelsResponse{ + Object: "list", + Data: []core.Model{ + { + ID: "claude-opus-4", + Object: "model", + OwnedBy: "gemini", + }, + }, + }, + } + registry.RegisterProviderWithType(mock, "gemini") + if err := registry.Initialize(context.Background()); err != nil { + t.Fatalf("Initialize() error = %v", err) + } + + raw := []byte(`{ + "version": 1, + "updated_at": "2025-01-01T00:00:00Z", + "providers": { + "gemini": { + "display_name": "Gemini", + "api_type": "openai", + "supported_modes": ["chat"] + } + }, + "models": { + "claude-4-opus": { + "display_name": "Claude 4 Opus", + "modes": ["chat"], + "aliases": ["claude-opus-4", "gemini/claude-opus-4"] + } + }, + "provider_models": { + "gemini/claude-4-opus": { + "model_ref": "claude-4-opus", + "enabled": true, + "context_window": 200000 + } + } + }`) + list, err := modeldata.Parse(raw) + if err != nil { + t.Fatalf("Parse() error = %v", err) + } + registry.SetModelList(list, raw) + registry.EnrichModels() + + if registry.ModelCount() != 1 { + t.Fatalf("ModelCount() = %d, want 1", registry.ModelCount()) + } + if synthetic := registry.GetModel("claude-4-opus"); synthetic != nil { + t.Fatalf("expected canonical alias target to NOT be materialized, got %+v", synthetic) + } + + info := registry.GetModel("claude-opus-4") + if info == nil { + t.Fatal("expected upstream model ID to remain registered") + } + if info.Model.ID != "claude-opus-4" { + t.Fatalf("Model.ID = %q, want claude-opus-4", info.Model.ID) + } + if info.Model.Metadata == nil { + t.Fatal("expected metadata to be enriched via alias") + } + if info.Model.Metadata.DisplayName != "Claude 4 Opus" { + t.Fatalf("DisplayName = %q, want Claude 4 Opus", info.Model.Metadata.DisplayName) + } + if info.Model.Metadata.ContextWindow == nil || *info.Model.Metadata.ContextWindow != 200000 { + t.Fatalf("ContextWindow = %v, want 200000", info.Model.Metadata.ContextWindow) + } + }) + t.Run("DuplicateModels", func(t *testing.T) { registry := NewModelRegistry() mock1 := ®istryMockProvider{