easyai-ai-gateway/apps/api/internal/store/candidates.go

package store

import (
	"context"
	"fmt"

	"github.com/easyai/easyai-ai-gateway/apps/api/internal/auth"
)

func (s *Store) ListModelCandidates(ctx context.Context, model string, modelType string, user *auth.User) ([]RuntimeModelCandidate, error) {
	rows, err := s.pool.Query(ctx, `
SELECT p.id::text, p.platform_key, p.name, p.provider,
       COALESCE(NULLIF(p.config->>'specType', ''), NULLIF(cp.provider_type, ''), NULLIF(p.config->>'sourceSpecType', ''), p.provider) AS spec_type,
       COALESCE(p.base_url, ''),
       p.auth_type, p.credentials, p.config, p.default_pricing_mode,
       p.default_discount_factor::float8, COALESCE(p.pricing_rule_set_id::text, ''),
       p.retry_policy, p.rate_limit_policy,
       COALESCE(p.dynamic_priority, p.priority) AS effective_priority,
       m.id::text, COALESCE(m.base_model_id::text, ''), COALESCE(b.canonical_model_key, ''),
       COALESCE(b.provider_model_name, ''), m.model_name, COALESCE(m.model_alias, ''),
       m.model_type, m.display_name, m.capabilities, m.capability_override,
       COALESCE(b.base_billing_config, '{}'::jsonb), m.billing_config, m.billing_config_override,
       m.pricing_mode, COALESCE(m.discount_factor, 0)::float8, COALESCE(m.pricing_rule_set_id::text, ''),
       m.permission_config, m.retry_policy, m.rate_limit_policy, COALESCE(m.runtime_policy_set_id::text, COALESCE(b.runtime_policy_set_id::text, '')),
       COALESCE(NULLIF(m.runtime_policy_override, '{}'::jsonb), b.runtime_policy_override, '{}'::jsonb)
FROM platform_models m
JOIN integration_platforms p ON p.id = m.platform_id
LEFT JOIN model_catalog_providers cp ON cp.provider_key = p.provider OR cp.provider_code = p.provider
LEFT JOIN base_model_catalog b ON b.id = m.base_model_id
LEFT JOIN runtime_client_states s
  ON s.client_id = p.platform_key || ':' || m.model_type || ':' || m.model_name
WHERE p.status = 'enabled'
  AND p.deleted_at IS NULL
  AND m.enabled = true
  AND (
    m.model_type = $2
    OR ($2 = 'text_generate' AND m.model_type IN ('chat', 'responses', 'text'))
    OR ($2 = 'image_generate' AND m.model_type IN ('image', 'images.generations'))
    OR ($2 = 'image_edit' AND m.model_type IN ('images.edits'))
    OR ($2 = 'video_generate' AND m.model_type IN ('video', 'videos.generations', 'video_generate', 'text_to_video', 'image_to_video', 'omni_video', 'video_edit', 'video_reference', 'video_first_last_frame'))
  )
  AND (p.cooldown_until IS NULL OR p.cooldown_until <= now())
  AND (
    m.model_name = $1
    OR m.model_alias = $1
    OR b.canonical_model_key = $1
    OR b.provider_model_name = $1
  )
ORDER BY effective_priority ASC,
         COALESCE(s.limiter_ratio, 0) ASC,
         COALESCE(s.running_count, 0) ASC,
         COALESCE(s.waiting_count, 0) ASC,
         COALESCE(s.last_assigned_at, to_timestamp(0)) ASC,
         m.created_at ASC`, model, modelType)
	if err != nil {
		return nil, err
	}
	defer rows.Close()

	items := make([]RuntimeModelCandidate, 0)
	for rows.Next() {
		var item RuntimeModelCandidate
		var credentials []byte
		var platformConfig []byte
		var platformRetryPolicy []byte
		var platformRateLimitPolicy []byte
		var capabilities []byte
		var capabilityOverride []byte
		var baseBilling []byte
		var billing []byte
		var billingOverride []byte
		var permissionConfig []byte
		var modelRetryPolicy []byte
		var modelRateLimitPolicy []byte
		var runtimePolicyOverride []byte
		if err := rows.Scan(
			&item.PlatformID,
			&item.PlatformKey,
			&item.PlatformName,
			&item.Provider,
			&item.SpecType,
			&item.BaseURL,
			&item.AuthType,
			&credentials,
			&platformConfig,
			&item.DefaultPricingMode,
			&item.DefaultDiscountFactor,
			&item.PlatformPricingRuleSetID,
			&platformRetryPolicy,
			&platformRateLimitPolicy,
			&item.PlatformPriority,
			&item.PlatformModelID,
			&item.BaseModelID,
			&item.CanonicalModelKey,
			&item.ProviderModelName,
			&item.ModelName,
			&item.ModelAlias,
			&item.ModelType,
			&item.DisplayName,
			&capabilities,
			&capabilityOverride,
			&baseBilling,
			&billing,
			&billingOverride,
			&item.PricingMode,
			&item.DiscountFactor,
			&item.ModelPricingRuleSetID,
			&permissionConfig,
			&modelRetryPolicy,
			&modelRateLimitPolicy,
			&item.RuntimePolicySetID,
			&runtimePolicyOverride,
		); err != nil {
			return nil, err
		}
		item.Credentials = decodeObject(credentials)
		item.PlatformConfig = decodeObject(platformConfig)
		item.PlatformRetryPolicy = decodeObject(platformRetryPolicy)
		item.PlatformRateLimitPolicy = decodeObject(platformRateLimitPolicy)
		item.Capabilities = decodeObject(capabilities)
		item.CapabilityOverride = decodeObject(capabilityOverride)
		item.BaseBillingConfig = decodeObject(baseBilling)
		item.BillingConfig = decodeObject(billing)
		item.BillingConfigOverride = decodeObject(billingOverride)
		item.PermissionConfig = decodeObject(permissionConfig)
		item.ModelRetryPolicy = decodeObject(modelRetryPolicy)
		item.ModelRateLimitPolicy = decodeObject(modelRateLimitPolicy)
		item.RuntimePolicyOverride = decodeObject(runtimePolicyOverride)
		item.ClientID = fmt.Sprintf("%s:%s:%s", item.PlatformKey, item.ModelType, item.ModelName)
		item.QueueKey = item.ClientID
		items = append(items, item)
	}
	if err := rows.Err(); err != nil {
		return nil, err
	}
	if len(items) == 0 {
		return nil, ErrNoModelCandidate
	}
	items, err = s.filterCandidatesByAccessRules(ctx, user, items)
	if err != nil {
		return nil, err
	}
	if len(items) == 0 {
		return nil, ErrNoModelCandidate
	}
	return items, nil
}