Atlas/internal/provider/models.go

package provider

import (
	"context"
	"encoding/json"
	"fmt"
	"net/url"
	"strings"

	"github.com/rose_cat707/Atlas/internal/store"
)

type ModelEntry struct {
	ID   string `json:"id"`
	Name string `json:"name"`
}

type ModelsService struct {
	store  *store.Store
	router *Router
}

func NewModelsService(st *store.Store, router *Router) *ModelsService {
	return &ModelsService{store: st, router: router}
}

func DefaultModel(models []ModelEntry) string {
	if len(models) > 0 {
		return models[0].ID
	}
	return ""
}

// ListEffective returns upstream models first, then manual models not already present.
func (m *ModelsService) ListEffective(ctx context.Context, cfg *Config) ([]ModelEntry, error) {
	upstream, _ := m.fetchUpstream(ctx, cfg)
	manual, err := m.listManual(ctx, cfg.ID)
	if err != nil {
		return nil, err
	}
	return mergeModels(upstream, manual), nil
}

func (m *ModelsService) ListForServiceType(ctx context.Context, serviceType string) ([]ModelEntry, error) {
	providers, err := m.store.ListProviders(ctx, serviceType, true)
	if err != nil {
		return nil, err
	}
	seen := map[string]struct{}{}
	var out []ModelEntry
	for _, p := range providers {
		cfg, err := m.router.providerToConfig(&p)
		if err != nil {
			continue
		}
		models, err := m.ListEffective(ctx, cfg)
		if err != nil {
			continue
		}
		for _, model := range models {
			if _, ok := seen[model.ID]; ok {
				continue
			}
			seen[model.ID] = struct{}{}
			out = append(out, model)
		}
	}
	return out, nil
}

type ModelsBreakdown struct {
	Upstream  []ModelEntry `json:"upstream"`
	Manual    []ModelEntry `json:"manual"`
	Effective []ModelEntry `json:"effective"`
}

func (m *ModelsService) Breakdown(ctx context.Context, providerID int64) (*ModelsBreakdown, error) {
	p, err := m.store.GetProvider(ctx, providerID)
	if err != nil {
		return nil, err
	}
	cfg, err := m.router.providerToConfig(p)
	if err != nil {
		return nil, err
	}
	upstream, _ := m.fetchUpstream(ctx, cfg)
	manual, err := m.listManual(ctx, providerID)
	if err != nil {
		return nil, err
	}
	return &ModelsBreakdown{
		Upstream:  upstream,
		Manual:    manual,
		Effective: mergeModels(upstream, manual),
	}, nil
}

func (m *ModelsService) AddManual(ctx context.Context, providerID int64, modelID, name string) error {
	modelID = strings.TrimSpace(modelID)
	name = strings.TrimSpace(name)
	if modelID == "" || name == "" {
		return fmt.Errorf("model_id and name required")
	}
	_, err := m.store.CreateProviderModel(ctx, &store.ProviderModel{
		ProviderID: providerID,
		ModelID:    modelID,
		Name:       name,
	})
	if err != nil {
		return fmt.Errorf("model already exists or invalid")
	}
	m.invalidateCache(providerID)
	return nil
}

func (m *ModelsService) RemoveManual(ctx context.Context, providerID int64, modelID string) error {
	if err := m.store.DeleteProviderModel(ctx, providerID, modelID); err != nil {
		return err
	}
	m.invalidateCache(providerID)
	return nil
}

func (m *ModelsService) RefreshUpstream(ctx context.Context, providerID int64) ([]ModelEntry, error) {
	p, err := m.store.GetProvider(ctx, providerID)
	if err != nil {
		return nil, err
	}
	cfg, err := m.router.providerToConfig(p)
	if err != nil {
		return nil, err
	}
	m.invalidateCache(providerID)
	return m.fetchUpstream(ctx, cfg)
}

// FetchUpstream loads models from upstream using the provider protocol.
func (m *ModelsService) FetchUpstream(ctx context.Context, cfg *Config) ([]ModelEntry, error) {
	return m.fetchUpstream(ctx, cfg)
}

// ModelsPath returns the upstream models API path for a provider configuration.
func ModelsPath(cfg *Config) string {
	q := url.Values{}
	if cfg.Protocol != "" {
		q.Set("protocol", cfg.Protocol)
	}
	if cfg.ServiceType != "" {
		q.Set("service_type", cfg.ServiceType)
	}
	if len(q) == 0 {
		return "/v1/models"
	}
	return "/v1/models?" + q.Encode()
}

func (m *ModelsService) listManual(ctx context.Context, providerID int64) ([]ModelEntry, error) {
	rows, err := m.store.ListProviderModels(ctx, providerID)
	if err != nil {
		return nil, err
	}
	out := make([]ModelEntry, 0, len(rows))
	for _, r := range rows {
		out = append(out, ModelEntry{ID: r.ModelID, Name: r.Name})
	}
	return out, nil
}

// ListEffectiveDetails returns upstream + manual models with input schema.
func (m *ModelsService) ListEffectiveDetails(ctx context.Context, cfg *Config) ([]ModelDetail, error) {
	upstream, _ := m.fetchUpstreamDetails(ctx, cfg)
	manual, err := m.listManual(ctx, cfg.ID)
	if err != nil {
		return nil, err
	}
	return mergeModelDetails(upstream, manual), nil
}

// FindModelDetail looks up a model by id for a provider configuration.
func (m *ModelsService) FindModelDetail(ctx context.Context, cfg *Config, modelID string) (*ModelDetail, error) {
	models, err := m.ListEffectiveDetails(ctx, cfg)
	if err != nil {
		return nil, err
	}
	for _, md := range models {
		if md.ID == modelID {
			copy := md
			return &copy, nil
		}
	}
	return nil, fmt.Errorf("model not found: %s", modelID)
}

// ListImageModelsForMeta aggregates image models from all image-related providers.
func (m *ModelsService) ListImageModelsForMeta(ctx context.Context) ([]ModelDetail, error) {
	seen := map[string]struct{}{}
	var out []ModelDetail
	for _, serviceType := range []string{ModelKindTextToImage, ModelKindImageToImage} {
		providers, err := m.store.ListProviders(ctx, serviceType, true)
		if err != nil {
			return nil, err
		}
		for _, p := range providers {
			cfg, err := m.router.providerToConfig(&p)
			if err != nil {
				continue
			}
			models, err := m.ListEffectiveDetails(ctx, cfg)
			if err != nil {
				continue
			}
			for _, model := range models {
				if _, ok := seen[model.ID]; ok {
					continue
				}
				seen[model.ID] = struct{}{}
				out = append(out, model)
			}
		}
	}
	return out, nil
}

func (m *ModelsService) fetchUpstream(ctx context.Context, cfg *Config) ([]ModelEntry, error) {
	details, err := m.fetchUpstreamDetails(ctx, cfg)
	if err != nil || len(details) == 0 {
		return nil, err
	}
	out := make([]ModelEntry, len(details))
	for i, d := range details {
		out[i] = ModelEntry{ID: d.ID, Name: d.Name}
	}
	return out, nil
}

func (m *ModelsService) fetchUpstreamDetails(ctx context.Context, cfg *Config) ([]ModelDetail, error) {
	key := detailsCacheKey(cfg.ID)
	if v, ok := m.store.Cache.Get(key); ok {
		if models, ok := v.([]ModelDetail); ok {
			return models, nil
		}
	}
	path := ModelsPath(cfg)
	meta := CallMeta{Method: "GET", Path: path}
	res, err := m.router.ProxyWithMeta(ctx, cfg, "GET", path, nil, nil, meta)
	if err != nil || res.StatusCode >= 400 {
		return nil, err
	}
	models := parseDetailsForProtocol(cfg.Protocol, res.Body)
	if len(models) > 0 {
		m.store.Cache.Set(key, models)
	}
	return models, nil
}

func parseModelsForProtocol(protocol string, body []byte) []ModelEntry {
	if protocol == "replicate" {
		if models := parseReplicateModels(body); len(models) > 0 {
			return models
		}
	}
	if models := parseOpenAIModels(body); len(models) > 0 {
		return models
	}
	if protocol == "replicate" {
		return parseReplicateModels(body)
	}
	return nil
}

func parseOpenAIModels(body []byte) []ModelEntry {
	var openai struct {
		Data []struct {
			ID   string `json:"id"`
			Name string `json:"name"`
		} `json:"data"`
	}
	if err := json.Unmarshal(body, &openai); err == nil && len(openai.Data) > 0 {
		out := make([]ModelEntry, 0, len(openai.Data))
		for _, item := range openai.Data {
			name := item.Name
			if name == "" {
				name = item.ID
			}
			if item.ID != "" {
				out = append(out, ModelEntry{ID: item.ID, Name: name})
			}
		}
		return out
	}
	var arr []struct {
		ID   string `json:"id"`
		Name string `json:"name"`
	}
	if err := json.Unmarshal(body, &arr); err == nil && len(arr) > 0 {
		out := make([]ModelEntry, 0, len(arr))
		for _, item := range arr {
			name := item.Name
			if name == "" {
				name = item.ID
			}
			if item.ID != "" {
				out = append(out, ModelEntry{ID: item.ID, Name: name})
			}
		}
		return out
	}
	return nil
}

func parseReplicateModels(body []byte) []ModelEntry {
	var resp struct {
		Results []struct {
			Owner       string `json:"owner"`
			Name        string `json:"name"`
			Description string `json:"description"`
		} `json:"results"`
	}
	if err := json.Unmarshal(body, &resp); err != nil || len(resp.Results) == 0 {
		return nil
	}
	out := make([]ModelEntry, 0, len(resp.Results))
	for _, item := range resp.Results {
		if item.Owner == "" || item.Name == "" {
			continue
		}
		id := item.Owner + "/" + item.Name
		name := item.Name
		if item.Description != "" {
			name = item.Description
		} else {
			name = id
		}
		out = append(out, ModelEntry{ID: id, Name: name})
	}
	return out
}

func parseDetailsForProtocol(protocol string, body []byte) []ModelDetail {
	if protocol == "replicate" {
		if models := parseReplicateModelDetails(body); len(models) > 0 {
			return models
		}
	}
	if entries := parseOpenAIModels(body); len(entries) > 0 {
		out := make([]ModelDetail, len(entries))
		for i, e := range entries {
			out[i] = manualModelDetail(e.ID, e.Name)
		}
		return out
	}
	if protocol == "replicate" {
		if entries := parseReplicateModels(body); len(entries) > 0 {
			out := make([]ModelDetail, len(entries))
			for i, e := range entries {
				out[i] = manualModelDetail(e.ID, e.Name)
			}
			return out
		}
	}
	return nil
}

func mergeModelDetails(upstream []ModelDetail, manual []ModelEntry) []ModelDetail {
	seen := map[string]struct{}{}
	out := make([]ModelDetail, 0, len(upstream)+len(manual))
	for _, m := range upstream {
		if _, ok := seen[m.ID]; ok {
			continue
		}
		seen[m.ID] = struct{}{}
		out = append(out, m)
	}
	for _, m := range manual {
		if _, ok := seen[m.ID]; ok {
			continue
		}
		seen[m.ID] = struct{}{}
		out = append(out, manualModelDetail(m.ID, m.Name))
	}
	return out
}

func mergeModels(upstream, manual []ModelEntry) []ModelEntry {
	seen := map[string]struct{}{}
	out := make([]ModelEntry, 0, len(upstream)+len(manual))
	for _, m := range upstream {
		if _, ok := seen[m.ID]; ok {
			continue
		}
		seen[m.ID] = struct{}{}
		out = append(out, m)
	}
	for _, m := range manual {
		if _, ok := seen[m.ID]; ok {
			continue
		}
		seen[m.ID] = struct{}{}
		out = append(out, m)
	}
	return out
}

func cacheKey(providerID int64) string {
	return fmt.Sprintf("upstream_models:%d", providerID)
}

func detailsCacheKey(providerID int64) string {
	return fmt.Sprintf("upstream_model_details:%d", providerID)
}

func (m *ModelsService) invalidateCache(providerID int64) {
	m.store.Cache.Delete(cacheKey(providerID))
	m.store.Cache.Delete(detailsCacheKey(providerID))
}

func FilterModelsByKind(models []ModelDetail, kind string) []ModelDetail {
	out := make([]ModelDetail, 0, len(models))
	for _, m := range models {
		if m.Kind == kind {
			out = append(out, m)
		}
	}
	return out
}

func DefaultModelDetail(models []ModelDetail) string {
	if len(models) > 0 {
		return models[0].ID
	}
	return ""
}