thread-master/apps/backend/internal/module/ai/registry.go

package ai

import (
	"context"
	"crypto/sha256"
	"encoding/hex"
	"fmt"
	"strings"
	"sync"
	"time"
)

const (
	ProviderXAI        = "xai"
	ProviderOpenCodeGo = "opencode-go"

	// ModelsCacheTTL — list models cache window
	ModelsCacheTTL = 5 * time.Minute
)

// Registry holds provider clients + models cache.
type Registry struct {
	clients map[string]Client
	mu      sync.Mutex
	cache   map[string]modelsCacheEntry
}

type modelsCacheEntry struct {
	Models    []string
	ExpiresAt time.Time
}

func NewRegistry() *Registry {
	return &Registry{
		clients: map[string]Client{
			ProviderXAI:        NewOpenAICompatible(ProviderXAI, "https://api.x.ai/v1"),
			ProviderOpenCodeGo: NewOpenAICompatible(ProviderOpenCodeGo, "https://opencode.ai/zen/go/v1"),
		},
		cache: map[string]modelsCacheEntry{},
	}
}

func NormalizeProvider(p string) string {
	switch strings.ToLower(strings.TrimSpace(p)) {
	case "opencode-go", "opencode", "opencode_go":
		return ProviderOpenCodeGo
	case "xai", "grok", "":
		return ProviderXAI
	default:
		return strings.ToLower(strings.TrimSpace(p))
	}
}

func (r *Registry) Client(provider string) (Client, error) {
	id := NormalizeProvider(provider)
	c, ok := r.clients[id]
	if !ok {
		return nil, fmt.Errorf("unsupported provider: %s (use xai or opencode-go)", provider)
	}
	return c, nil
}

// ListModelsCached fetches models via provider API; caches 5 minutes per provider+key.
// fromCache indicates whether result was served from cache.
func (r *Registry) ListModelsCached(ctx context.Context, provider, apiKey string) (models []string, fromCache bool, err error) {
	id := NormalizeProvider(provider)
	key := cacheKey(id, apiKey)

	r.mu.Lock()
	if ent, ok := r.cache[key]; ok && time.Now().Before(ent.ExpiresAt) {
		out := append([]string(nil), ent.Models...)
		r.mu.Unlock()
		return out, true, nil
	}
	r.mu.Unlock()

	c, err := r.Client(id)
	if err != nil {
		return nil, false, err
	}
	list, err := c.ListModels(ctx, apiKey)
	if err != nil {
		return nil, false, err
	}

	r.mu.Lock()
	r.cache[key] = modelsCacheEntry{
		Models:    append([]string(nil), list...),
		ExpiresAt: time.Now().Add(ModelsCacheTTL),
	}
	r.mu.Unlock()
	return list, false, nil
}

func cacheKey(provider, apiKey string) string {
	sum := sha256.Sum256([]byte(provider + "\n" + apiKey))
	return hex.EncodeToString(sum[:16])
}

// FallbackModels when no key / provider error (never empty for UI).
// OpenCode Go 模型 id 見 https://opencode.ai/docs/go/（不是 OpenAI gpt-*）
func FallbackModels(provider string) []string {
	switch NormalizeProvider(provider) {
	case ProviderOpenCodeGo:
		return []string{
			"kimi-k2.6",
			"glm-5.1",
			"glm-5.2",
			"deepseek-v4-flash",
			"deepseek-v4-pro",
			"qwen3.7-plus",
			"minimax-m2.7",
			"mimo-v2.5",
		}
	default:
		return []string{"grok-3", "grok-3-mini", "grok-2"}
	}
}

// DefaultModel for provider when member settings model empty / invalid.
func DefaultModel(provider string) string {
	list := FallbackModels(provider)
	if len(list) == 0 {
		return "grok-3"
	}
	return list[0]
}

// SanitizeModel remaps invalid model ids for a provider (e.g. gpt-4o-mini on opencode-go).
func SanitizeModel(provider, model string) string {
	provider = NormalizeProvider(provider)
	model = strings.TrimSpace(model)
	if model == "" {
		return DefaultModel(provider)
	}
	if provider != ProviderOpenCodeGo {
		// xAI：允許 grok-*；若誤存 opencode 模型名，仍交由 API 驗證
		return model
	}
	// OpenCode Go：常見錯誤 — 從 xAI / OpenAI 帶過來的模型名
	lower := strings.ToLower(model)
	if strings.HasPrefix(lower, "gpt-") ||
		strings.HasPrefix(lower, "o1") ||
		strings.HasPrefix(lower, "o3") ||
		strings.HasPrefix(lower, "grok") ||
		strings.HasPrefix(lower, "claude") ||
		strings.HasPrefix(lower, "gemini") {
		return DefaultModel(provider)
	}
	return model
}