GopherGate/internal/models/registry.go

package models

import "strings"

// CanonicalProviders lists the original model creators in priority order.
// When a model name exists in multiple providers (e.g. deepseek-v4-pro in
// deepseek, ollama-cloud, openrouter, etc.), these providers take precedence
// so the proxy uses authoritative metadata (pricing, limits) rather than a
// reseller's values.
var CanonicalProviders = []string{
	"openai",
	"google",
	"deepseek",
	"xai",
	"moonshotai",
	"moonshotai-cn",
	"anthropic",
	"mistral",
	"cohere",
	"minimax",
	"xiaomi",
}

type ModelRegistry struct {
	Providers map[string]ProviderInfo `json:"-"`
}

type ProviderInfo struct {
	ID     string                   `json:"id"`
	Name   string                   `json:"name"`
	Models map[string]ModelMetadata `json:"models"`
}

type ModelMetadata struct {
	ID         string           `json:"id"`
	Name       string           `json:"name"`
	Cost       *ModelCost       `json:"cost,omitempty"`
	Limit      *ModelLimit      `json:"limit,omitempty"`
	Modalities *ModelModalities `json:"modalities,omitempty"`
	ToolCall   *bool            `json:"tool_call,omitempty"`
	Reasoning  *bool            `json:"reasoning,omitempty"`
}

type ModelCost struct {
	Input      float64  `json:"input"`
	Output     float64  `json:"output"`
	CacheRead  *float64 `json:"cache_read,omitempty"`
	CacheWrite *float64 `json:"cache_write,omitempty"`
}

type ModelLimit struct {
	Context uint32 `json:"context"`
	Output  uint32 `json:"output"`
}

type ModelModalities struct {
	Input  []string `json:"input"`
	Output []string `json:"output"`
}

// findInCanonical searches the canonical providers in order for an exact model
// key match. Returns the metadata and true if found.
func (r *ModelRegistry) findInCanonical(modelID string) (*ModelMetadata, bool) {
	for _, key := range CanonicalProviders {
		if p, ok := r.Providers[key]; ok {
			if m, ok := p.Models[modelID]; ok {
				return &m, true
			}
		}
	}
	return nil, false
}

// findInAll searches all providers (map iteration, random order) for an exact
// model key match. Used as fallback when canonical search fails.
func (r *ModelRegistry) findInAll(modelID string) (*ModelMetadata, bool) {
	for _, p := range r.Providers {
		if m, ok := p.Models[modelID]; ok {
			return &m, true
		}
	}
	return nil, false
}

// findInCanonicalByID searches canonical providers for a model whose metadata
// ID field matches modelID.
func (r *ModelRegistry) findInCanonicalByID(modelID string) (*ModelMetadata, bool) {
	for _, key := range CanonicalProviders {
		if p, ok := r.Providers[key]; ok {
			for _, m := range p.Models {
				if m.ID == modelID {
					return &m, true
				}
			}
		}
	}
	return nil, false
}

// findInAllByID searches all providers for a model whose metadata ID field
// matches modelID.
func (r *ModelRegistry) findInAllByID(modelID string) (*ModelMetadata, bool) {
	for _, p := range r.Providers {
		for _, m := range p.Models {
			if m.ID == modelID {
				return &m, true
			}
		}
	}
	return nil, false
}

// findCanonicalReverseFuzzy searches canonical providers for any model whose
// key starts with modelID.
func (r *ModelRegistry) findCanonicalReverseFuzzy(modelID string) (*ModelMetadata, bool) {
	for _, key := range CanonicalProviders {
		if p, ok := r.Providers[key]; ok {
			for id, m := range p.Models {
				if strings.HasPrefix(id, modelID) {
					return &m, true
				}
			}
		}
	}
	return nil, false
}

// findAllReverseFuzzy searches all providers for any model whose key starts
// with modelID.
func (r *ModelRegistry) findAllReverseFuzzy(modelID string) (*ModelMetadata, bool) {
	for _, p := range r.Providers {
		for id, m := range p.Models {
			if strings.HasPrefix(id, modelID) {
				return &m, true
			}
		}
	}
	return nil, false
}

// findCanonicalForwardFuzzy searches canonical providers for any model whose
// key is a prefix of modelID.
func (r *ModelRegistry) findCanonicalForwardFuzzy(modelID string) (*ModelMetadata, bool) {
	for _, key := range CanonicalProviders {
		if p, ok := r.Providers[key]; ok {
			for id, m := range p.Models {
				if strings.HasPrefix(modelID, id) {
					return &m, true
				}
			}
		}
	}
	return nil, false
}

// findAllForwardFuzzy searches all providers for any model whose key is a
// prefix of modelID.
func (r *ModelRegistry) findAllForwardFuzzy(modelID string) (*ModelMetadata, bool) {
	for _, p := range r.Providers {
		for id, m := range p.Models {
			if strings.HasPrefix(modelID, id) {
				return &m, true
			}
		}
	}
	return nil, false
}

// FindModel looks up model metadata by ID. It searches canonical providers
// first at each strategy level (exact key, metadata ID, reverse fuzzy,
// forward fuzzy) and falls back to all providers only when canonical search
// yields no result. This prevents reseller entries (ollama-cloud, openrouter,
// etc.) from overriding the original provider's authoritative pricing and
// limits.
func (r *ModelRegistry) FindModel(modelID string) *ModelMetadata {
	// 1. Exact key match — canonical first, then all
	if m, ok := r.findInCanonical(modelID); ok {
		return m
	}
	if m, ok := r.findInAll(modelID); ok {
		return m
	}

	// 2. Match by metadata ID field — canonical first, then all
	if m, ok := r.findInCanonicalByID(modelID); ok {
		return m
	}
	if m, ok := r.findInAllByID(modelID); ok {
		return m
	}

	// 3. Reverse fuzzy: model key starts with modelID
	//    e.g. 'gpt-5.4-mini' matching 'gpt-5.4-mini-2026-04-01'
	if m, ok := r.findCanonicalReverseFuzzy(modelID); ok {
		return m
	}
	if m, ok := r.findAllReverseFuzzy(modelID); ok {
		return m
	}

	// 4. Forward fuzzy: modelID starts with model key
	//    e.g. 'gpt-4o-2024-05-13' matching 'gpt-4o'
	if m, ok := r.findCanonicalForwardFuzzy(modelID); ok {
		return m
	}
	if m, ok := r.findAllForwardFuzzy(modelID); ok {
		return m
	}

	return nil
}