841 lines
29 KiB
Go
841 lines
29 KiB
Go
package store
|
|
|
|
import (
|
|
"context"
|
|
"fmt"
|
|
"math"
|
|
"sort"
|
|
"strings"
|
|
"unicode"
|
|
|
|
"github.com/easyai/easyai-ai-gateway/apps/api/internal/auth"
|
|
)
|
|
|
|
type ListModelCandidatesOptions struct {
|
|
CacheAffinityKey string
|
|
CacheAffinityPolicy map[string]any
|
|
}
|
|
|
|
func (s *Store) ListModelCandidates(ctx context.Context, model string, modelType string, user *auth.User, options ...ListModelCandidatesOptions) ([]RuntimeModelCandidate, error) {
|
|
exactModel := strings.TrimSpace(model)
|
|
modelMatchKey := normalizeModelMatchKey(exactModel)
|
|
listOptions := normalizeListModelCandidatesOptions(modelType, options...)
|
|
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(NULLIF(m.provider_model_name, ''), m.model_name), m.model_name, COALESCE(m.model_alias, ''),
|
|
$2::text AS requested_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, ''),
|
|
COALESCE(b.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),
|
|
COALESCE(rp.retry_policy, '{}'::jsonb), COALESCE(rp.rate_limit_policy, '{}'::jsonb),
|
|
COALESCE(rp.auto_disable_policy, '{}'::jsonb), COALESCE(rp.degrade_policy, '{}'::jsonb),
|
|
COALESCE(con.active, 0)::float8,
|
|
COALESCE(queued.waiting, 0)::float8,
|
|
COALESCE(rpm.used_value, 0)::float8, COALESCE(rpm.reserved_value, 0)::float8,
|
|
COALESCE(tpm.used_value, 0)::float8, COALESCE(tpm.reserved_value, 0)::float8,
|
|
COALESCE(s.running_count, 0)::float8,
|
|
COALESCE(s.waiting_count, 0)::float8,
|
|
COALESCE(s.limiter_ratio, 0)::float8,
|
|
COALESCE(EXTRACT(EPOCH FROM s.last_assigned_at), 0)::float8,
|
|
COALESCE(ca.request_count, 0)::float8,
|
|
COALESCE(ca.input_tokens, 0)::float8,
|
|
COALESCE(ca.cached_input_tokens, 0)::float8,
|
|
COALESCE(ca.ema_hit_ratio, 0)::float8,
|
|
COALESCE(ca.last_hit_ratio, 0)::float8,
|
|
COALESCE(EXTRACT(EPOCH FROM ca.last_observed_at), 0)::float8
|
|
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 model_runtime_policy_sets rp ON rp.id = COALESCE(m.runtime_policy_set_id, b.runtime_policy_set_id)
|
|
LEFT JOIN runtime_client_states s
|
|
ON s.client_id = p.platform_key || ':' || $2::text || ':' || COALESCE(NULLIF(m.provider_model_name, ''), m.model_name)
|
|
LEFT JOIN gateway_cache_affinity_stats ca
|
|
ON ca.client_id = p.platform_key || ':' || $2::text || ':' || COALESCE(NULLIF(m.provider_model_name, ''), m.model_name)
|
|
AND ca.cache_affinity_key = NULLIF($4::text, '')
|
|
AND ($5::int <= 0 OR ca.last_observed_at >= now() - ($5::int * interval '1 second'))
|
|
LEFT JOIN (
|
|
SELECT scope_key, SUM(lease_value) AS active
|
|
FROM gateway_concurrency_leases
|
|
WHERE scope_type = 'platform_model'
|
|
AND released_at IS NULL
|
|
AND expires_at > now()
|
|
GROUP BY scope_key
|
|
) con ON con.scope_key = m.id::text
|
|
LEFT JOIN (
|
|
SELECT queued_sources.platform_model_id, COUNT(DISTINCT queued_sources.task_id) AS waiting
|
|
FROM (
|
|
SELECT t.id::text AS task_id, qm.id::text AS platform_model_id
|
|
FROM gateway_tasks t
|
|
JOIN integration_platforms qp ON TRUE
|
|
JOIN platform_models qm ON qm.platform_id = qp.id
|
|
WHERE t.async_mode = true
|
|
AND t.status = 'queued'
|
|
AND NULLIF(t.model_type, '') IS NOT NULL
|
|
AND qm.model_type @> jsonb_build_array(t.model_type)
|
|
AND t.queue_key = qp.platform_key || ':' || t.model_type || ':' || COALESCE(NULLIF(qm.provider_model_name, ''), qm.model_name)
|
|
AND NOT EXISTS (SELECT 1 FROM gateway_task_attempts existing_attempt WHERE existing_attempt.task_id = t.id)
|
|
UNION ALL
|
|
SELECT latest.task_id::text AS task_id, latest.platform_model_id
|
|
FROM (
|
|
SELECT DISTINCT ON (a.task_id) a.task_id, a.platform_model_id::text AS platform_model_id
|
|
FROM gateway_tasks t
|
|
JOIN gateway_task_attempts a ON a.task_id = t.id
|
|
WHERE t.async_mode = true
|
|
AND t.status = 'queued'
|
|
AND a.platform_model_id IS NOT NULL
|
|
ORDER BY a.task_id, a.attempt_no DESC, a.started_at DESC
|
|
) latest
|
|
) queued_sources
|
|
GROUP BY queued_sources.platform_model_id
|
|
) queued ON queued.platform_model_id = m.id::text
|
|
LEFT JOIN (
|
|
SELECT DISTINCT ON (scope_key) scope_key, used_value, reserved_value
|
|
FROM gateway_rate_limit_counters
|
|
WHERE scope_type = 'platform_model'
|
|
AND metric = 'rpm'
|
|
AND reset_at > now()
|
|
ORDER BY scope_key, window_start DESC
|
|
) rpm ON rpm.scope_key = m.id::text
|
|
LEFT JOIN (
|
|
SELECT scope_key, SUM(used_value) AS used_value, SUM(reserved_value) AS reserved_value
|
|
FROM gateway_rate_limit_counters
|
|
WHERE scope_type = 'platform_model'
|
|
AND metric LIKE 'tpm%'
|
|
AND reset_at > now()
|
|
GROUP BY scope_key
|
|
) tpm ON tpm.scope_key = m.id::text
|
|
WHERE p.status = 'enabled'
|
|
AND p.deleted_at IS NULL
|
|
AND m.enabled = true
|
|
AND m.model_type @> jsonb_build_array($2::text)
|
|
AND (p.cooldown_until IS NULL OR p.cooldown_until <= now())
|
|
AND (m.cooldown_until IS NULL OR m.cooldown_until <= now())
|
|
AND (
|
|
(
|
|
$2::text IN ('audio_generate', 'text_to_speech', 'voice_clone')
|
|
AND (
|
|
m.model_alias = $1::text
|
|
OR m.model_name = $1::text
|
|
OR b.canonical_model_key = $1::text
|
|
OR b.provider_model_name = $1::text
|
|
OR (
|
|
NULLIF($3::text, '') IS NOT NULL
|
|
AND (
|
|
regexp_replace(COALESCE(m.model_alias, ''), '[[:space:]]+', '', 'g') = $3::text
|
|
OR regexp_replace(COALESCE(m.model_name, ''), '[[:space:]]+', '', 'g') = $3::text
|
|
OR regexp_replace(COALESCE(b.canonical_model_key, ''), '[[:space:]]+', '', 'g') = $3::text
|
|
OR regexp_replace(COALESCE(b.provider_model_name, ''), '[[:space:]]+', '', 'g') = $3::text
|
|
)
|
|
)
|
|
)
|
|
)
|
|
OR (
|
|
$2::text NOT IN ('audio_generate', 'text_to_speech', 'voice_clone')
|
|
AND (
|
|
(
|
|
COALESCE(m.model_alias, '') <> ''
|
|
AND (
|
|
m.model_alias = $1::text
|
|
OR (
|
|
NULLIF($3::text, '') IS NOT NULL
|
|
AND regexp_replace(COALESCE(m.model_alias, ''), '[[:space:]]+', '', 'g') = $3::text
|
|
)
|
|
)
|
|
)
|
|
OR (
|
|
COALESCE(m.model_alias, '') = ''
|
|
AND (
|
|
m.model_name = $1::text
|
|
OR b.canonical_model_key = $1::text
|
|
OR b.provider_model_name = $1::text
|
|
OR (
|
|
NULLIF($3::text, '') IS NOT NULL
|
|
AND (
|
|
regexp_replace(COALESCE(m.model_name, ''), '[[:space:]]+', '', 'g') = $3::text
|
|
OR regexp_replace(COALESCE(b.canonical_model_key, ''), '[[:space:]]+', '', 'g') = $3::text
|
|
OR regexp_replace(COALESCE(b.provider_model_name, ''), '[[:space:]]+', '', 'g') = $3::text
|
|
)
|
|
)
|
|
)
|
|
)
|
|
)
|
|
)
|
|
)
|
|
ORDER BY effective_priority 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`, exactModel, modelType, modelMatchKey, listOptions.CacheAffinityKey, cacheAffinityStaleAfterSeconds(listOptions.CacheAffinityPolicy))
|
|
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
|
|
var runtimeRetryPolicy []byte
|
|
var runtimeRateLimitPolicy []byte
|
|
var autoDisablePolicy []byte
|
|
var degradePolicy []byte
|
|
var concurrentActive float64
|
|
var queuedWaiting float64
|
|
var rpmUsed float64
|
|
var rpmReserved float64
|
|
var tpmUsed float64
|
|
var tpmReserved float64
|
|
var stateRunningCount float64
|
|
var stateWaitingCount float64
|
|
var stateLimiterRatio float64
|
|
var lastAssignedUnix float64
|
|
var cacheRequestCount float64
|
|
var cacheInputTokens float64
|
|
var cacheCachedInputTokens float64
|
|
var cacheEMAHitRatio float64
|
|
var cacheLastHitRatio float64
|
|
var cacheLastObservedUnix float64
|
|
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,
|
|
&item.BasePricingRuleSetID,
|
|
&permissionConfig,
|
|
&modelRetryPolicy,
|
|
&modelRateLimitPolicy,
|
|
&item.RuntimePolicySetID,
|
|
&runtimePolicyOverride,
|
|
&runtimeRetryPolicy,
|
|
&runtimeRateLimitPolicy,
|
|
&autoDisablePolicy,
|
|
°radePolicy,
|
|
&concurrentActive,
|
|
&queuedWaiting,
|
|
&rpmUsed,
|
|
&rpmReserved,
|
|
&tpmUsed,
|
|
&tpmReserved,
|
|
&stateRunningCount,
|
|
&stateWaitingCount,
|
|
&stateLimiterRatio,
|
|
&lastAssignedUnix,
|
|
&cacheRequestCount,
|
|
&cacheInputTokens,
|
|
&cacheCachedInputTokens,
|
|
&cacheEMAHitRatio,
|
|
&cacheLastHitRatio,
|
|
&cacheLastObservedUnix,
|
|
); 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.RuntimeRetryPolicy = decodeObject(runtimeRetryPolicy)
|
|
item.RuntimeRateLimitPolicy = decodeObject(runtimeRateLimitPolicy)
|
|
item.AutoDisablePolicy = decodeObject(autoDisablePolicy)
|
|
item.DegradePolicy = decodeObject(degradePolicy)
|
|
upstreamModelName := firstNonEmpty(item.ProviderModelName, item.ModelName)
|
|
item.ClientID = fmt.Sprintf("%s:%s:%s", item.PlatformKey, item.ModelType, upstreamModelName)
|
|
item.QueueKey = item.ClientID
|
|
item.RunningCount = stateRunningCount
|
|
item.WaitingCount = maxFloat(queuedWaiting, stateWaitingCount)
|
|
item.LastAssignedUnix = lastAssignedUnix
|
|
applyRuntimeCandidateCacheAffinity(&item, listOptions, runtimeCandidateCacheAffinityInput{
|
|
RequestCount: int(cacheRequestCount),
|
|
InputTokens: int(cacheInputTokens),
|
|
CachedInputTokens: int(cacheCachedInputTokens),
|
|
EMAHitRatio: cacheEMAHitRatio,
|
|
LastHitRatio: cacheLastHitRatio,
|
|
LastObservedUnix: cacheLastObservedUnix,
|
|
})
|
|
applyRuntimeCandidateLoad(&item, runtimeCandidateLoadInput{
|
|
Policy: effectiveModelRateLimitPolicy(item.PlatformRateLimitPolicy, item.RuntimeRateLimitPolicy, item.RuntimePolicySetID, item.RuntimePolicyOverride, item.ModelRateLimitPolicy),
|
|
ConcurrentActive: concurrentActive,
|
|
QueuedWaiting: queuedWaiting,
|
|
RPMUsed: rpmUsed,
|
|
RPMReserved: rpmReserved,
|
|
TPMUsed: tpmUsed,
|
|
TPMReserved: tpmReserved,
|
|
StateRunningCount: stateRunningCount,
|
|
StateWaitingCount: stateWaitingCount,
|
|
StateLimiterRatio: stateLimiterRatio,
|
|
})
|
|
items = append(items, item)
|
|
}
|
|
if err := rows.Err(); err != nil {
|
|
return nil, err
|
|
}
|
|
if len(items) == 0 {
|
|
if unavailableErr, err := s.modelCandidateCooldownError(ctx, model, modelType); err != nil {
|
|
return nil, err
|
|
} else if unavailableErr != nil {
|
|
return nil, unavailableErr
|
|
}
|
|
return nil, ErrNoModelCandidate
|
|
}
|
|
items, err = s.filterCandidatesByAccessRules(ctx, user, items)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
if len(items) == 0 {
|
|
return nil, ErrNoModelCandidate
|
|
}
|
|
sortRuntimeModelCandidates(items)
|
|
return items, nil
|
|
}
|
|
|
|
func (s *Store) GetRuntimeModelCandidateForVoiceCloneDeletion(ctx context.Context, platformModelID string, platformID string) (RuntimeModelCandidate, bool, error) {
|
|
platformModelID = strings.TrimSpace(platformModelID)
|
|
platformID = strings.TrimSpace(platformID)
|
|
if platformModelID == "" && platformID == "" {
|
|
return RuntimeModelCandidate{}, false, nil
|
|
}
|
|
var item RuntimeModelCandidate
|
|
var credentials []byte
|
|
var platformConfig []byte
|
|
err := s.pool.QueryRow(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,
|
|
COALESCE(m.id::text, ''), COALESCE(NULLIF(m.provider_model_name, ''), m.model_name, 'voice_clone'),
|
|
COALESCE(m.model_name, 'voice_clone'), COALESCE(m.model_alias, '')
|
|
FROM integration_platforms p
|
|
LEFT JOIN platform_models m
|
|
ON m.platform_id = p.id
|
|
AND (
|
|
(NULLIF($1, '')::uuid IS NOT NULL AND m.id = NULLIF($1, '')::uuid)
|
|
OR (
|
|
NULLIF($1, '') IS NULL
|
|
AND m.model_type @> jsonb_build_array('voice_clone'::text)
|
|
)
|
|
)
|
|
LEFT JOIN model_catalog_providers cp ON cp.provider_key = p.provider OR cp.provider_code = p.provider
|
|
WHERE p.deleted_at IS NULL
|
|
AND (
|
|
(NULLIF($1, '')::uuid IS NOT NULL AND m.id = NULLIF($1, '')::uuid)
|
|
OR (NULLIF($2, '')::uuid IS NOT NULL AND p.id = NULLIF($2, '')::uuid)
|
|
)
|
|
ORDER BY CASE WHEN NULLIF($1, '')::uuid IS NOT NULL AND m.id = NULLIF($1, '')::uuid THEN 0 ELSE 1 END,
|
|
m.created_at DESC
|
|
LIMIT 1`, platformModelID, platformID).Scan(
|
|
&item.PlatformID,
|
|
&item.PlatformKey,
|
|
&item.PlatformName,
|
|
&item.Provider,
|
|
&item.SpecType,
|
|
&item.BaseURL,
|
|
&item.AuthType,
|
|
&credentials,
|
|
&platformConfig,
|
|
&item.PlatformModelID,
|
|
&item.ProviderModelName,
|
|
&item.ModelName,
|
|
&item.ModelAlias,
|
|
)
|
|
if err != nil {
|
|
if IsNotFound(err) {
|
|
return RuntimeModelCandidate{}, false, nil
|
|
}
|
|
return RuntimeModelCandidate{}, false, err
|
|
}
|
|
item.Credentials = decodeObject(credentials)
|
|
item.PlatformConfig = decodeObject(platformConfig)
|
|
item.ModelType = "voice_clone"
|
|
item.ClientID = item.PlatformKey + ":voice_clone:" + firstNonEmpty(item.ProviderModelName, item.ModelName)
|
|
item.QueueKey = item.ClientID
|
|
return item, true, nil
|
|
}
|
|
|
|
type runtimeCandidateLoadInput struct {
|
|
Policy map[string]any
|
|
ConcurrentActive float64
|
|
QueuedWaiting float64
|
|
RPMUsed float64
|
|
RPMReserved float64
|
|
TPMUsed float64
|
|
TPMReserved float64
|
|
StateRunningCount float64
|
|
StateWaitingCount float64
|
|
StateLimiterRatio float64
|
|
}
|
|
|
|
type runtimeCandidateCacheAffinityInput struct {
|
|
RequestCount int
|
|
InputTokens int
|
|
CachedInputTokens int
|
|
EMAHitRatio float64
|
|
LastHitRatio float64
|
|
LastObservedUnix float64
|
|
}
|
|
|
|
func normalizeListModelCandidatesOptions(modelType string, options ...ListModelCandidatesOptions) ListModelCandidatesOptions {
|
|
if len(options) == 0 {
|
|
return ListModelCandidatesOptions{}
|
|
}
|
|
out := options[0]
|
|
out.CacheAffinityKey = strings.TrimSpace(out.CacheAffinityKey)
|
|
if out.CacheAffinityKey == "" || !cacheAffinityPolicyEnabled(out.CacheAffinityPolicy, modelType) {
|
|
out.CacheAffinityKey = ""
|
|
}
|
|
return out
|
|
}
|
|
|
|
func applyRuntimeCandidateCacheAffinity(candidate *RuntimeModelCandidate, options ListModelCandidatesOptions, input runtimeCandidateCacheAffinityInput) {
|
|
affinity := RuntimeCandidateCacheAffinity{
|
|
Key: options.CacheAffinityKey,
|
|
RequestCount: input.RequestCount,
|
|
InputTokens: input.InputTokens,
|
|
CachedInputTokens: input.CachedInputTokens,
|
|
EMAHitRatio: boundedRatio(input.EMAHitRatio),
|
|
LastHitRatio: boundedRatio(input.LastHitRatio),
|
|
LastObservedUnix: input.LastObservedUnix,
|
|
AdjustedPriority: float64(candidate.PlatformPriority),
|
|
}
|
|
minSamples := cacheAffinityMinSamples(options.CacheAffinityPolicy)
|
|
hasObservedCachedHit := affinity.CachedInputTokens > 0 || affinity.LastHitRatio > 0
|
|
hasSampledAffinity := input.RequestCount >= minSamples && affinity.EMAHitRatio > 0
|
|
if options.CacheAffinityKey == "" || input.RequestCount <= 0 || (!hasObservedCachedHit && !hasSampledAffinity) {
|
|
candidate.CacheAffinity = affinity
|
|
return
|
|
}
|
|
affinity.Confidence = float64(input.RequestCount) / float64(minSamples)
|
|
if affinity.Confidence > 1 {
|
|
affinity.Confidence = 1
|
|
}
|
|
if hasObservedCachedHit && affinity.Confidence < 1 {
|
|
affinity.Confidence = 1
|
|
}
|
|
affinity.Score = runtimeCandidateCacheAffinityScore(affinity)
|
|
maxBoost := cacheAffinityMaxPriorityBoost(options.CacheAffinityPolicy)
|
|
affinity.Boost = affinity.Score * maxBoost
|
|
if affinity.Boost > maxBoost {
|
|
affinity.Boost = maxBoost
|
|
}
|
|
if hasObservedCachedHit || affinity.Score > 0 {
|
|
affinity.Applied = true
|
|
}
|
|
if affinity.Boost > 0 {
|
|
affinity.AdjustedPriority = float64(candidate.PlatformPriority) - affinity.Boost
|
|
}
|
|
candidate.CacheAffinity = affinity
|
|
}
|
|
|
|
func cacheAffinityPolicyEnabled(policy map[string]any, modelType string) bool {
|
|
if enabled, ok := policy["enabled"].(bool); ok && !enabled {
|
|
return false
|
|
}
|
|
modelTypes := stringListFromAny(policy["modelTypes"])
|
|
if len(modelTypes) == 0 {
|
|
modelTypes = []string{"chat", "text_generate", "responses"}
|
|
}
|
|
modelType = strings.TrimSpace(modelType)
|
|
for _, candidate := range modelTypes {
|
|
if strings.EqualFold(strings.TrimSpace(candidate), modelType) {
|
|
return true
|
|
}
|
|
}
|
|
return false
|
|
}
|
|
|
|
func cacheAffinityMinSamples(policy map[string]any) int {
|
|
value := intFromStorePolicy(policy, "minSamples", 3)
|
|
if value <= 0 {
|
|
return 1
|
|
}
|
|
return value
|
|
}
|
|
|
|
func cacheAffinityMinInputTokens(policy map[string]any) int {
|
|
value := intFromStorePolicy(policy, "minInputTokens", 512)
|
|
if value < 0 {
|
|
return 0
|
|
}
|
|
return value
|
|
}
|
|
|
|
func cacheAffinityStaleAfterSeconds(policy map[string]any) int {
|
|
return intFromStorePolicy(policy, "staleAfterSeconds", 86400)
|
|
}
|
|
|
|
func cacheAffinityMaxPriorityBoost(policy map[string]any) float64 {
|
|
value := floatFromStorePolicy(policy, "maxPriorityBoost", 20)
|
|
if value < 0 {
|
|
return 0
|
|
}
|
|
return value
|
|
}
|
|
|
|
func cacheAffinityEMAAlpha(policy map[string]any) float64 {
|
|
value := floatFromStorePolicy(policy, "emaAlpha", 0.35)
|
|
if value <= 0 || value > 1 {
|
|
return 0.35
|
|
}
|
|
return value
|
|
}
|
|
|
|
func runtimeCandidateCacheAffinityScore(affinity RuntimeCandidateCacheAffinity) float64 {
|
|
hitRatio := affinity.EMAHitRatio
|
|
if affinity.LastHitRatio > hitRatio {
|
|
hitRatio = affinity.LastHitRatio
|
|
}
|
|
hitRatio = boundedRatio(hitRatio)
|
|
confidence := boundedRatio(affinity.Confidence)
|
|
if hitRatio <= 0 || confidence <= 0 {
|
|
return 0
|
|
}
|
|
tokenFactor := 0.0
|
|
if affinity.CachedInputTokens > 0 {
|
|
tokenFactor = float64(affinity.CachedInputTokens) / 8192
|
|
if tokenFactor > 1 {
|
|
tokenFactor = 1
|
|
}
|
|
}
|
|
return math.Round(hitRatio*confidence*(1+tokenFactor)*1_000_000) / 1_000_000
|
|
}
|
|
|
|
func intFromStorePolicy(policy map[string]any, key string, fallback int) int {
|
|
if policy == nil {
|
|
return fallback
|
|
}
|
|
switch value := policy[key].(type) {
|
|
case int:
|
|
return value
|
|
case int64:
|
|
return int(value)
|
|
case float64:
|
|
return int(value)
|
|
default:
|
|
return fallback
|
|
}
|
|
}
|
|
|
|
func floatFromStorePolicy(policy map[string]any, key string, fallback float64) float64 {
|
|
if policy == nil {
|
|
return fallback
|
|
}
|
|
switch value := policy[key].(type) {
|
|
case int:
|
|
return float64(value)
|
|
case int64:
|
|
return float64(value)
|
|
case float64:
|
|
return value
|
|
default:
|
|
return fallback
|
|
}
|
|
}
|
|
|
|
func boundedRatio(value float64) float64 {
|
|
if value < 0 {
|
|
return 0
|
|
}
|
|
if value > 1 {
|
|
return 1
|
|
}
|
|
return value
|
|
}
|
|
|
|
func applyRuntimeCandidateLoad(candidate *RuntimeModelCandidate, input runtimeCandidateLoadInput) {
|
|
rpmLimit := rateLimitForMetric(input.Policy, "rpm")
|
|
tpmLimitValue := tpmLimit(input.Policy)
|
|
concurrentLimit := rateLimitForMetric(input.Policy, "concurrent")
|
|
rpmCurrent := input.RPMUsed + input.RPMReserved
|
|
tpmCurrent := input.TPMUsed + input.TPMReserved
|
|
concurrentCurrent := input.ConcurrentActive + input.QueuedWaiting
|
|
metrics := RuntimeCandidateLoadMetrics{
|
|
RPMCurrent: rpmCurrent,
|
|
RPMLimit: rpmLimit,
|
|
RPMRatio: ratioIfLimited(rpmCurrent, rpmLimit),
|
|
TPMCurrent: tpmCurrent,
|
|
TPMLimit: tpmLimitValue,
|
|
TPMRatio: ratioIfLimited(tpmCurrent, tpmLimitValue),
|
|
ConcurrentCurrent: concurrentCurrent,
|
|
ConcurrentLimit: concurrentLimit,
|
|
ConcurrentRatio: ratioIfLimited(concurrentCurrent, concurrentLimit),
|
|
QueuedCount: input.QueuedWaiting,
|
|
StateRunningCount: input.StateRunningCount,
|
|
StateWaitingCount: input.StateWaitingCount,
|
|
StateLimiterRatio: input.StateLimiterRatio,
|
|
}
|
|
candidate.LoadMetrics = metrics
|
|
candidate.LoadLimited = rpmLimit > 0 || tpmLimitValue > 0 || concurrentLimit > 0
|
|
candidate.LoadRatio = maxFloat(metrics.RPMRatio, metrics.TPMRatio, metrics.ConcurrentRatio)
|
|
}
|
|
|
|
func ratioIfLimited(current float64, limit float64) float64 {
|
|
if limit <= 0 {
|
|
return 0
|
|
}
|
|
return current / limit
|
|
}
|
|
|
|
func sortRuntimeModelCandidates(items []RuntimeModelCandidate) {
|
|
hasFull := false
|
|
hasNonFull := false
|
|
for index := range items {
|
|
items[index].LoadAvoided = false
|
|
if !items[index].CacheAffinity.Applied && items[index].CacheAffinity.AdjustedPriority == 0 {
|
|
items[index].CacheAffinity.AdjustedPriority = float64(items[index].PlatformPriority)
|
|
}
|
|
if runtimeCandidateFull(items[index]) {
|
|
hasFull = true
|
|
} else {
|
|
hasNonFull = true
|
|
}
|
|
}
|
|
if hasFull && hasNonFull {
|
|
for index := range items {
|
|
items[index].LoadAvoided = runtimeCandidateFull(items[index])
|
|
}
|
|
}
|
|
sort.SliceStable(items, func(i, j int) bool {
|
|
aFull := runtimeCandidateFull(items[i])
|
|
bFull := runtimeCandidateFull(items[j])
|
|
if aFull != bFull {
|
|
return !aFull
|
|
}
|
|
if items[i].CacheAffinity.Applied != items[j].CacheAffinity.Applied {
|
|
return items[i].CacheAffinity.Applied
|
|
}
|
|
if items[i].CacheAffinity.Applied && items[j].CacheAffinity.Applied {
|
|
if items[i].CacheAffinity.Score != items[j].CacheAffinity.Score {
|
|
return items[i].CacheAffinity.Score > items[j].CacheAffinity.Score
|
|
}
|
|
if items[i].CacheAffinity.EMAHitRatio != items[j].CacheAffinity.EMAHitRatio {
|
|
return items[i].CacheAffinity.EMAHitRatio > items[j].CacheAffinity.EMAHitRatio
|
|
}
|
|
if items[i].CacheAffinity.LastHitRatio != items[j].CacheAffinity.LastHitRatio {
|
|
return items[i].CacheAffinity.LastHitRatio > items[j].CacheAffinity.LastHitRatio
|
|
}
|
|
if items[i].CacheAffinity.CachedInputTokens != items[j].CacheAffinity.CachedInputTokens {
|
|
return items[i].CacheAffinity.CachedInputTokens > items[j].CacheAffinity.CachedInputTokens
|
|
}
|
|
}
|
|
if items[i].CacheAffinity.AdjustedPriority != items[j].CacheAffinity.AdjustedPriority {
|
|
return items[i].CacheAffinity.AdjustedPriority < items[j].CacheAffinity.AdjustedPriority
|
|
}
|
|
if items[i].LoadRatio != items[j].LoadRatio {
|
|
return items[i].LoadRatio < items[j].LoadRatio
|
|
}
|
|
if items[i].RunningCount != items[j].RunningCount {
|
|
return items[i].RunningCount < items[j].RunningCount
|
|
}
|
|
if items[i].WaitingCount != items[j].WaitingCount {
|
|
return items[i].WaitingCount < items[j].WaitingCount
|
|
}
|
|
if items[i].LastAssignedUnix != items[j].LastAssignedUnix {
|
|
return items[i].LastAssignedUnix < items[j].LastAssignedUnix
|
|
}
|
|
return false
|
|
})
|
|
}
|
|
|
|
func runtimeCandidateFull(candidate RuntimeModelCandidate) bool {
|
|
return candidate.LoadLimited && candidate.LoadRatio >= 1
|
|
}
|
|
|
|
func (s *Store) modelCandidateCooldownError(ctx context.Context, model string, modelType string) (error, error) {
|
|
exactModel := strings.TrimSpace(model)
|
|
modelMatchKey := normalizeModelMatchKey(exactModel)
|
|
rows, err := s.pool.Query(ctx, `
|
|
SELECT p.name,
|
|
COALESCE(NULLIF(m.display_name, ''), NULLIF(m.model_alias, ''), m.model_name),
|
|
COALESCE(to_char(p.cooldown_until AT TIME ZONE 'UTC', 'YYYY-MM-DD"T"HH24:MI:SS.MS"Z"'), ''),
|
|
GREATEST(COALESCE(EXTRACT(EPOCH FROM p.cooldown_until - now()), 0), 0)::float8,
|
|
COALESCE(to_char(m.cooldown_until AT TIME ZONE 'UTC', 'YYYY-MM-DD"T"HH24:MI:SS.MS"Z"'), ''),
|
|
GREATEST(COALESCE(EXTRACT(EPOCH FROM m.cooldown_until - now()), 0), 0)::float8
|
|
FROM platform_models m
|
|
JOIN integration_platforms p ON p.id = m.platform_id
|
|
LEFT JOIN base_model_catalog b ON b.id = m.base_model_id
|
|
WHERE p.status = 'enabled'
|
|
AND p.deleted_at IS NULL
|
|
AND m.enabled = true
|
|
AND m.model_type @> jsonb_build_array($2::text)
|
|
AND (
|
|
(
|
|
$2::text IN ('audio_generate', 'text_to_speech', 'voice_clone')
|
|
AND (
|
|
m.model_alias = $1::text
|
|
OR m.model_name = $1::text
|
|
OR b.canonical_model_key = $1::text
|
|
OR b.provider_model_name = $1::text
|
|
OR (
|
|
NULLIF($3::text, '') IS NOT NULL
|
|
AND (
|
|
regexp_replace(COALESCE(m.model_alias, ''), '[[:space:]]+', '', 'g') = $3::text
|
|
OR regexp_replace(COALESCE(m.model_name, ''), '[[:space:]]+', '', 'g') = $3::text
|
|
OR regexp_replace(COALESCE(b.canonical_model_key, ''), '[[:space:]]+', '', 'g') = $3::text
|
|
OR regexp_replace(COALESCE(b.provider_model_name, ''), '[[:space:]]+', '', 'g') = $3::text
|
|
)
|
|
)
|
|
)
|
|
)
|
|
OR (
|
|
$2::text NOT IN ('audio_generate', 'text_to_speech', 'voice_clone')
|
|
AND (
|
|
(
|
|
COALESCE(m.model_alias, '') <> ''
|
|
AND (
|
|
m.model_alias = $1::text
|
|
OR (
|
|
NULLIF($3::text, '') IS NOT NULL
|
|
AND regexp_replace(COALESCE(m.model_alias, ''), '[[:space:]]+', '', 'g') = $3::text
|
|
)
|
|
)
|
|
)
|
|
OR (
|
|
COALESCE(m.model_alias, '') = ''
|
|
AND (
|
|
m.model_name = $1::text
|
|
OR b.canonical_model_key = $1::text
|
|
OR b.provider_model_name = $1::text
|
|
OR (
|
|
NULLIF($3::text, '') IS NOT NULL
|
|
AND (
|
|
regexp_replace(COALESCE(m.model_name, ''), '[[:space:]]+', '', 'g') = $3::text
|
|
OR regexp_replace(COALESCE(b.canonical_model_key, ''), '[[:space:]]+', '', 'g') = $3::text
|
|
OR regexp_replace(COALESCE(b.provider_model_name, ''), '[[:space:]]+', '', 'g') = $3::text
|
|
)
|
|
)
|
|
)
|
|
)
|
|
)
|
|
)
|
|
)
|
|
ORDER BY GREATEST(COALESCE(p.cooldown_until, to_timestamp(0)), COALESCE(m.cooldown_until, to_timestamp(0))) DESC,
|
|
p.priority ASC,
|
|
m.created_at ASC`, exactModel, modelType, modelMatchKey)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
defer rows.Close()
|
|
|
|
for rows.Next() {
|
|
var platformName string
|
|
var displayName string
|
|
var platformCooldownUntil string
|
|
var platformRemainingSeconds float64
|
|
var modelCooldownUntil string
|
|
var modelRemainingSeconds float64
|
|
if err := rows.Scan(
|
|
&platformName,
|
|
&displayName,
|
|
&platformCooldownUntil,
|
|
&platformRemainingSeconds,
|
|
&modelCooldownUntil,
|
|
&modelRemainingSeconds,
|
|
); err != nil {
|
|
return nil, err
|
|
}
|
|
if modelRemainingSeconds > 0 {
|
|
return &ModelCandidateUnavailableError{
|
|
Code: "model_cooling_down",
|
|
Message: cooldownErrorMessage("模型", displayName, modelRemainingSeconds, modelCooldownUntil),
|
|
}, nil
|
|
}
|
|
if platformRemainingSeconds > 0 {
|
|
return &ModelCandidateUnavailableError{
|
|
Code: "platform_cooling_down",
|
|
Message: cooldownErrorMessage("平台", platformName, platformRemainingSeconds, platformCooldownUntil),
|
|
}, nil
|
|
}
|
|
}
|
|
if err := rows.Err(); err != nil {
|
|
return nil, err
|
|
}
|
|
return nil, nil
|
|
}
|
|
|
|
func cooldownErrorMessage(scope string, name string, remainingSeconds float64, cooldownUntil string) string {
|
|
name = strings.TrimSpace(name)
|
|
if name == "" {
|
|
name = "候选"
|
|
}
|
|
remainingMinutes := remainingSeconds / 60
|
|
if remainingMinutes < 0.1 {
|
|
remainingMinutes = 0.1
|
|
}
|
|
message := fmt.Sprintf("%s %s 冷却中，剩余 %.1f 分钟", scope, name, remainingMinutes)
|
|
if strings.TrimSpace(cooldownUntil) != "" {
|
|
message += "，预计恢复时间 " + cooldownUntil
|
|
}
|
|
return message
|
|
}
|
|
|
|
func normalizeModelMatchKey(value string) string {
|
|
value = strings.TrimSpace(value)
|
|
if value == "" {
|
|
return ""
|
|
}
|
|
var builder strings.Builder
|
|
builder.Grow(len(value))
|
|
for _, char := range value {
|
|
if unicode.IsSpace(char) {
|
|
continue
|
|
}
|
|
builder.WriteRune(char)
|
|
}
|
|
return builder.String()
|
|
}
|