85 lines
3.1 KiB
Go
85 lines
3.1 KiB
Go
package store
|
|
|
|
import (
|
|
"context"
|
|
"strings"
|
|
)
|
|
|
|
type CacheAffinityObservationInput struct {
|
|
CacheAffinityKey string
|
|
CacheAffinityKeys []string
|
|
CacheAffinityPolicy map[string]any
|
|
PlatformID string
|
|
PlatformModelID string
|
|
ClientID string
|
|
ModelType string
|
|
InputTokens int
|
|
CachedInputTokens int
|
|
CachedInputTokensKnown bool
|
|
}
|
|
|
|
func (s *Store) RecordCacheAffinityObservation(ctx context.Context, input CacheAffinityObservationInput) error {
|
|
input.CacheAffinityKey = strings.TrimSpace(input.CacheAffinityKey)
|
|
input.CacheAffinityKeys = normalizedCacheAffinityKeys(input.CacheAffinityKey, input.CacheAffinityKeys)
|
|
input.ClientID = strings.TrimSpace(input.ClientID)
|
|
input.ModelType = strings.TrimSpace(input.ModelType)
|
|
if len(input.CacheAffinityKeys) == 0 || input.ClientID == "" || input.ModelType == "" {
|
|
return nil
|
|
}
|
|
if !input.CachedInputTokensKnown || !cacheAffinityPolicyEnabled(input.CacheAffinityPolicy, input.ModelType) {
|
|
return nil
|
|
}
|
|
minInputTokens := cacheAffinityMinInputTokens(input.CacheAffinityPolicy)
|
|
if input.InputTokens < minInputTokens || input.InputTokens <= 0 {
|
|
return nil
|
|
}
|
|
if input.CachedInputTokens < 0 {
|
|
input.CachedInputTokens = 0
|
|
}
|
|
if input.CachedInputTokens > input.InputTokens {
|
|
input.CachedInputTokens = input.InputTokens
|
|
}
|
|
alpha := cacheAffinityEMAAlpha(input.CacheAffinityPolicy)
|
|
hitRatio := float64(input.CachedInputTokens) / float64(input.InputTokens)
|
|
for _, key := range input.CacheAffinityKeys {
|
|
if err := s.recordCacheAffinityObservationKey(ctx, input, key, hitRatio, alpha); err != nil {
|
|
return err
|
|
}
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func (s *Store) recordCacheAffinityObservationKey(ctx context.Context, input CacheAffinityObservationInput, key string, hitRatio float64, alpha float64) error {
|
|
_, err := s.pool.Exec(ctx, `
|
|
INSERT INTO gateway_cache_affinity_stats (
|
|
client_id, cache_affinity_key, platform_id, platform_model_id, model_type,
|
|
request_count, input_tokens, cached_input_tokens, ema_hit_ratio, last_hit_ratio, last_observed_at
|
|
)
|
|
VALUES (
|
|
$1, $2, NULLIF($3::text, '')::uuid, NULLIF($4::text, '')::uuid, $5,
|
|
1, $6, $7, $8, $8, now()
|
|
)
|
|
ON CONFLICT (client_id, cache_affinity_key) DO UPDATE
|
|
SET platform_id = COALESCE(EXCLUDED.platform_id, gateway_cache_affinity_stats.platform_id),
|
|
platform_model_id = COALESCE(EXCLUDED.platform_model_id, gateway_cache_affinity_stats.platform_model_id),
|
|
model_type = EXCLUDED.model_type,
|
|
request_count = gateway_cache_affinity_stats.request_count + 1,
|
|
input_tokens = gateway_cache_affinity_stats.input_tokens + EXCLUDED.input_tokens,
|
|
cached_input_tokens = gateway_cache_affinity_stats.cached_input_tokens + EXCLUDED.cached_input_tokens,
|
|
ema_hit_ratio = ($9::float8 * EXCLUDED.last_hit_ratio) + ((1 - $9::float8) * gateway_cache_affinity_stats.ema_hit_ratio),
|
|
last_hit_ratio = EXCLUDED.last_hit_ratio,
|
|
last_observed_at = now(),
|
|
updated_at = now()`,
|
|
input.ClientID,
|
|
key,
|
|
input.PlatformID,
|
|
input.PlatformModelID,
|
|
input.ModelType,
|
|
input.InputTokens,
|
|
input.CachedInputTokens,
|
|
hitRatio,
|
|
alpha,
|
|
)
|
|
return err
|
|
}
|