334 lines
11 KiB
Go
334 lines
11 KiB
Go
package service
|
|
|
|
import (
|
|
"context"
|
|
"sort"
|
|
"strings"
|
|
"time"
|
|
|
|
"github.com/QuantumNous/new-api/common"
|
|
"github.com/QuantumNous/new-api/model"
|
|
fingerprint "github.com/QuantumNous/new-api/pkg/model_fingerprint"
|
|
)
|
|
|
|
const AstraSamplingTaskType = "modeltrace_sampling"
|
|
|
|
const (
|
|
AstraSamplingWindow = 6 * time.Hour
|
|
AstraSamplingDefaultInterval = 30 * time.Minute
|
|
AstraSamplingUnstableInterval = 30 * time.Minute
|
|
AstraAvailabilityInterval = 15 * time.Minute
|
|
)
|
|
|
|
type AstraSamplingRound struct {
|
|
ID string `json:"id"`
|
|
ChannelID int `json:"channel_id"`
|
|
StartedAt int64 `json:"started_at"`
|
|
FinishedAt int64 `json:"finished_at,omitempty"`
|
|
Sent int `json:"sent"`
|
|
Version string `json:"version"`
|
|
Result fingerprint.Result `json:"result"`
|
|
Reason string `json:"reason,omitempty"`
|
|
Usage []fingerprint.ProbeResponse `json:"usage,omitempty"`
|
|
}
|
|
type AstraSamplingState struct {
|
|
Rounds []AstraSamplingRound `json:"rounds"`
|
|
}
|
|
|
|
type ChannelAvailabilityRound struct {
|
|
ChannelID int `json:"channel_id"`
|
|
At int64 `json:"at"`
|
|
Success bool `json:"success"`
|
|
}
|
|
|
|
type ChannelAvailabilityState struct {
|
|
Rounds []ChannelAvailabilityRound `json:"rounds"`
|
|
}
|
|
|
|
type AstraSamplingHistory struct {
|
|
UserInitiated bool `json:"user_initiated,omitempty"`
|
|
ID string `json:"id"`
|
|
At int64 `json:"at"`
|
|
Status string `json:"status"`
|
|
ValidSamples int `json:"valid_samples"`
|
|
Closest string `json:"closest,omitempty"`
|
|
Gap float64 `json:"score_gap,omitempty"`
|
|
}
|
|
type AstraSamplingSummary struct {
|
|
LatestValid *AstraSamplingHistory `json:"-"`
|
|
Enabled bool `json:"enabled"`
|
|
Available bool `json:"available"`
|
|
Channels int `json:"channels"`
|
|
SampledChannels int `json:"sampled_channels"`
|
|
Matches int `json:"matches"`
|
|
Mismatches int `json:"mismatches"`
|
|
Inconclusive int `json:"inconclusive"`
|
|
Errors int `json:"errors"`
|
|
Running int `json:"running"`
|
|
Rate *float64 `json:"rate"`
|
|
WindowHours int `json:"window_hours"`
|
|
IntervalMinutes int `json:"interval_minutes"`
|
|
IntervalMaxMinutes int `json:"interval_max_minutes"`
|
|
AvailabilityOnly bool `json:"availability_only,omitempty"`
|
|
AvailabilityIntervalMinutes int `json:"availability_interval_minutes,omitempty"`
|
|
AvailabilityChecks int `json:"availability_checks,omitempty"`
|
|
AvailabilitySuccesses int `json:"availability_successes,omitempty"`
|
|
AvailabilityRate *float64 `json:"availability_rate,omitempty"`
|
|
History []AstraSamplingHistory `json:"history"`
|
|
}
|
|
|
|
func AstraSamplingEnabled() bool {
|
|
return common.GetEnvOrDefaultBool("ASTRA_FINGERPRINT_ENABLED", false)
|
|
}
|
|
|
|
func AstraAvailabilityEnabled() bool {
|
|
return common.GetEnvOrDefaultBool("CHANNEL_AVAILABILITY_PROBE_ENABLED", true)
|
|
}
|
|
|
|
func ChannelSupportsAstra(channel *model.Channel) bool {
|
|
if channel == nil {
|
|
return false
|
|
}
|
|
for _, name := range channel.GetModels() {
|
|
if strings.TrimSpace(name) == fingerprint.Model {
|
|
return true
|
|
}
|
|
}
|
|
return false
|
|
}
|
|
|
|
func AstraSamplingIntervalForGroup(group string) time.Duration {
|
|
return AstraSamplingDefaultInterval
|
|
}
|
|
|
|
func AstraSamplingIntervalForChannel(channel model.Channel) time.Duration {
|
|
interval := AstraSamplingDefaultInterval
|
|
for _, group := range channel.GetGroups() {
|
|
interval = min(interval, AstraSamplingIntervalForGroup(group))
|
|
}
|
|
return interval
|
|
}
|
|
|
|
func AstraSamplingSchedulerInterval() time.Duration {
|
|
return 15 * time.Second
|
|
}
|
|
|
|
func AstraAvailabilityIntervalMinutes() int {
|
|
minutes := common.GetEnvOrDefault("CHANNEL_AVAILABILITY_INTERVAL_MINUTES", int(AstraAvailabilityInterval/time.Minute))
|
|
if minutes <= 0 {
|
|
return int(AstraAvailabilityInterval / time.Minute)
|
|
}
|
|
return minutes
|
|
}
|
|
|
|
// Public statistics classify completed rounds by their closest valid fingerprint,
|
|
// including partial rounds. Raw probe results retain the scorer's confidence.
|
|
func AstraSamplingResultStatus(result fingerprint.Result) string {
|
|
if result.Engine != "" {
|
|
return result.Status
|
|
}
|
|
closest := strings.TrimSpace(result.Closest)
|
|
if result.Status != "running" && result.ValidSamples > 0 && closest != "" {
|
|
if strings.EqualFold(closest, fingerprint.Model) {
|
|
return "match"
|
|
}
|
|
return "mismatch"
|
|
}
|
|
if result.Status == "match" || result.Status == "mismatch" {
|
|
return "inconclusive"
|
|
}
|
|
return result.Status
|
|
}
|
|
|
|
// No credentials are fetched by the public status query. Each eligible channel
|
|
// is sampled once even when it belongs to multiple groups.
|
|
func AstraSamplingChannels(ctx context.Context) ([]model.Channel, error) {
|
|
return samplingChannels(ctx, true)
|
|
}
|
|
|
|
func samplingChannels(ctx context.Context, requireAstra bool) ([]model.Channel, error) {
|
|
var rows []model.Channel
|
|
err := model.DB.WithContext(ctx).Select("id", "type", "group", "models", "status").Where("status = ?", common.ChannelStatusEnabled).Find(&rows).Error
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
channels := make([]model.Channel, 0, len(rows))
|
|
for _, channel := range rows {
|
|
if !requireAstra {
|
|
channels = append(channels, channel)
|
|
continue
|
|
}
|
|
if ChannelSupportsAstra(&channel) {
|
|
channels = append(channels, channel)
|
|
}
|
|
}
|
|
return channels, nil
|
|
}
|
|
|
|
func AstraSamplingChannelsDue(ctx context.Context, now time.Time) ([]model.Channel, error) {
|
|
due, _, err := AutomaticDetectionChannelsDue(ctx, "modeltrace", now)
|
|
return due, err
|
|
}
|
|
|
|
func AstraSamplingSummaries(ctx context.Context, now time.Time, settings map[string]model.StatusPageGroup) (map[string]*AstraSamplingSummary, error) {
|
|
channels, err := samplingChannels(ctx, false)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
var tasks []*model.SystemTask
|
|
err = model.DB.WithContext(ctx).Select("status", "state").Where("type = ? AND created_at >= ?", AstraSamplingTaskType, now.Add(-AstraSamplingWindow).Unix()).Order("id asc").Find(&tasks).Error
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
var availabilityTasks []*model.SystemTask
|
|
err = model.DB.WithContext(ctx).Select("status", "state").Where("type = ? AND created_at >= ?", model.SystemTaskTypeChannelAvailability, now.Add(-AstraSamplingWindow).Unix()).Order("id asc").Find(&availabilityTasks).Error
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
summaries := aggregateAstraSampling(channels, tasks, now, settings)
|
|
AggregateChannelAvailability(summaries, channels, availabilityTasks, now)
|
|
ranges, err := DetectionScheduleRanges(ctx, channels, "modeltrace")
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
for group, summary := range summaries {
|
|
if r, ok := ranges[group]; ok {
|
|
summary.IntervalMinutes, summary.IntervalMaxMinutes = r[0], r[1]
|
|
} else {
|
|
summary.IntervalMaxMinutes = summary.IntervalMinutes
|
|
}
|
|
}
|
|
return summaries, nil
|
|
}
|
|
|
|
func AggregateAstraSampling(channels []model.Channel, tasks []*model.SystemTask, now time.Time) map[string]*AstraSamplingSummary {
|
|
return aggregateAstraSampling(channels, tasks, now, nil)
|
|
}
|
|
|
|
func aggregateAstraSampling(channels []model.Channel, tasks []*model.SystemTask, now time.Time, settings map[string]model.StatusPageGroup) map[string]*AstraSamplingSummary {
|
|
groups := map[string]*AstraSamplingSummary{}
|
|
channelGroups := map[int][]string{}
|
|
seen := map[string]map[int]bool{}
|
|
for _, channel := range channels {
|
|
channelGroups[channel.Id] = channel.GetGroups()
|
|
astra := ChannelSupportsAstra(&channel)
|
|
for _, group := range channel.GetGroups() {
|
|
if groups[group] == nil {
|
|
groups[group] = &AstraSamplingSummary{Enabled: AstraSamplingEnabled(), Available: true, WindowHours: int(AstraSamplingWindow / time.Hour), IntervalMinutes: int(AstraSamplingIntervalForGroup(group) / time.Minute), History: []AstraSamplingHistory{}}
|
|
seen[group] = map[int]bool{}
|
|
}
|
|
if astra {
|
|
groups[group].Channels++
|
|
}
|
|
}
|
|
}
|
|
for _, task := range tasks {
|
|
var state AstraSamplingState
|
|
if task.DecodeState(&state) != nil {
|
|
continue
|
|
}
|
|
for _, round := range state.Rounds {
|
|
if round.Result.Engine != fingerprint.Version && round.Result.Status != "running" {
|
|
continue
|
|
}
|
|
if round.StartedAt < now.Add(-AstraSamplingWindow).Unix() {
|
|
continue
|
|
}
|
|
if !containsAstraChannel(channels, round.ChannelID) {
|
|
continue
|
|
}
|
|
status := AstraSamplingResultStatus(round.Result)
|
|
if status == "running" && (task.Status == model.SystemTaskStatusFailed || task.Status == model.SystemTaskStatusSucceeded) {
|
|
status = "error"
|
|
}
|
|
for _, group := range channelGroups[round.ChannelID] {
|
|
if cutoff := settings[group].ResetBefore; cutoff > 0 && round.StartedAt <= cutoff {
|
|
continue
|
|
}
|
|
s := groups[group]
|
|
switch status {
|
|
case "match":
|
|
s.Matches++
|
|
case "mismatch":
|
|
s.Mismatches++
|
|
case "inconclusive":
|
|
s.Inconclusive++
|
|
case "error":
|
|
s.Errors++
|
|
case "running":
|
|
s.Running++
|
|
default:
|
|
continue
|
|
}
|
|
if round.FinishedAt > 0 {
|
|
seen[group][round.ChannelID] = true
|
|
}
|
|
s.History = append(s.History, AstraSamplingHistory{ID: round.ID, At: round.StartedAt, Status: status, ValidSamples: round.Result.ValidSamples, Closest: round.Result.Closest, Gap: round.Result.Gap})
|
|
}
|
|
}
|
|
}
|
|
for group, s := range groups {
|
|
if s.Channels == 0 {
|
|
s.AvailabilityOnly = true
|
|
s.IntervalMinutes = AstraAvailabilityIntervalMinutes()
|
|
s.AvailabilityIntervalMinutes = AstraAvailabilityIntervalMinutes()
|
|
}
|
|
s.SampledChannels = len(seen[group])
|
|
if n := s.Matches + s.Mismatches; n > 0 {
|
|
rate := float64(s.Mismatches) / float64(n)
|
|
s.Rate = &rate
|
|
}
|
|
sort.SliceStable(s.History, func(i, j int) bool { return s.History[i].At > s.History[j].At })
|
|
for i := range s.History {
|
|
round := &s.History[i]
|
|
if round.ValidSamples > 0 && round.Closest != "" && (round.Status == "match" || round.Status == "mismatch") {
|
|
s.LatestValid = round
|
|
break
|
|
}
|
|
}
|
|
}
|
|
return groups
|
|
}
|
|
|
|
func containsAstraChannel(channels []model.Channel, id int) bool {
|
|
for i := range channels {
|
|
if channels[i].Id == id {
|
|
return ChannelSupportsAstra(&channels[i])
|
|
}
|
|
}
|
|
return false
|
|
}
|
|
|
|
func AggregateChannelAvailability(summaries map[string]*AstraSamplingSummary, channels []model.Channel, tasks []*model.SystemTask, now time.Time) {
|
|
channelGroups := map[int][]string{}
|
|
for i := range channels {
|
|
if !ChannelSupportsAstra(&channels[i]) {
|
|
channelGroups[channels[i].Id] = channels[i].GetGroups()
|
|
}
|
|
}
|
|
for _, task := range tasks {
|
|
var state ChannelAvailabilityState
|
|
if task.DecodeState(&state) != nil {
|
|
continue
|
|
}
|
|
for _, round := range state.Rounds {
|
|
if round.At < now.Add(-AstraSamplingWindow).Unix() {
|
|
continue
|
|
}
|
|
for _, group := range channelGroups[round.ChannelID] {
|
|
s := summaries[group]
|
|
s.AvailabilityChecks++
|
|
if round.Success {
|
|
s.AvailabilitySuccesses++
|
|
}
|
|
}
|
|
}
|
|
}
|
|
for _, s := range summaries {
|
|
if s.AvailabilityChecks > 0 {
|
|
rate := float64(s.AvailabilitySuccesses) / float64(s.AvailabilityChecks)
|
|
s.AvailabilityRate = &rate
|
|
}
|
|
}
|
|
}
|