new-api-custom/service/group_quality.go

package service

import (
	"context"
	"errors"
	"time"

	"github.com/QuantumNous/new-api/common"
	"github.com/QuantumNous/new-api/model"
	fingerprint "github.com/QuantumNous/new-api/pkg/model_fingerprint"
)

// Routing evidence has no statistical window. Each detector recovers separately.
func readGroupQuality(ctx context.Context, now time.Time) (map[string]AstraSamplingHistory, map[string]DetectionHistory, map[string]DetectionHistory, error) {
	settings, err := model.GetStatusPageGroups(ctx)
	if err != nil {
		return nil, nil, nil, err
	}
	channels, err := samplingChannels(ctx, false)
	if err != nil {
		return nil, nil, nil, err
	}
	channelGroups := map[int][]string{}
	for _, channel := range channels {
		channelGroups[channel.Id] = channel.GetGroups()
	}
	latest, visions := map[string]AstraSamplingHistory{}, map[string]DetectionHistory{}
	ciallos := map[string]DetectionHistory{}
	observe := func(channelID int, h AstraSamplingHistory) {
		if h.At <= 0 || h.At > now.Unix() || h.ValidSamples <= 0 || h.Closest == "" || (h.Status != "match" && h.Status != "mismatch") {
			return
		}
		for _, group := range channelGroups[channelID] {
			if h.At > settings[group].ResetBefore && h.At > latest[group].At {
				latest[group] = h
			}
		}
	}
	var tasks []model.SystemTask
	if err = model.DB.WithContext(ctx).Select("state").Where("type = ?", AstraSamplingTaskType).Find(&tasks).Error; err != nil {
		return nil, nil, nil, err
	}
	for _, task := range tasks {
		var state AstraSamplingState
		if task.DecodeState(&state) != nil {
			continue
		}
		for _, round := range state.Rounds {
			if round.FinishedAt <= 0 || round.Result.Engine != fingerprint.Version {
				continue
			}
			observe(round.ChannelID, AstraSamplingHistory{ID: round.ID, At: round.StartedAt, Status: AstraSamplingResultStatus(round.Result), ValidSamples: round.Result.ValidSamples, Closest: round.Result.Closest})
		}
	}
	var runs []model.DetectionRun
	if err = model.DB.WithContext(ctx).Select("id", "channel_id", "kind", "created_at", "quality", "result", "status", "error").Where("(status = ? OR (kind IN ? AND status = ?)) AND settled = ? AND created_at <= ?", "succeeded", []string{"vision", "ciallo"}, "failed", true, now.Unix()).Find(&runs).Error; err != nil {
		return nil, nil, nil, err
	}
	for _, run := range runs {
		run = NormalizeVisionRun(run)
		if run.Status != "succeeded" {
			continue
		}
		if run.Kind == "ciallo" {
			var result fingerprint.Result
			if common.UnmarshalJsonStr(run.Result, &result) == nil {
				run.Quality = CialloQuality(result.Closest)
			}
		}
		if run.Kind == "fingerprint" {
			var result fingerprint.Result
			if common.UnmarshalJsonStr(run.Result, &result) != nil || result.Engine != fingerprint.Version {
				continue
			}
			observe(run.ChannelID, AstraSamplingHistory{ID: run.ID, At: run.CreatedAt, Status: AstraSamplingResultStatus(result), ValidSamples: result.ValidSamples, Closest: result.Closest})
		}

		if (run.Kind == "vision" || run.Kind == "ciallo") && (run.Quality == "normal" || run.Quality == "degraded" || (run.Kind == "vision" && run.Quality == "incomplete")) {
			target := visions
			if run.Kind == "ciallo" {
				target = ciallos
			}
			for _, group := range channelGroups[run.ChannelID] {
				if run.CreatedAt > settings[group].ResetBefore && run.CreatedAt > target[group].At {
					target[group] = DetectionHistory{ID: run.ID, At: run.CreatedAt, Status: run.Status, Quality: run.Quality}
				}
			}
		}
	}
	return latest, visions, ciallos, nil
}

func cachedGroupQuality(ctx context.Context) (map[string]AstraSamplingHistory, map[string]DetectionHistory, map[string]DetectionHistory, error) {
	groupDegradationSnapshot.Lock()
	defer groupDegradationSnapshot.Unlock()
	now := time.Now()
	if now.Sub(groupDegradationSnapshot.updated) >= 15*time.Second {
		ctx, cancel := context.WithTimeout(ctx, 2*time.Second)
		latest, visions, ciallos, err := readGroupQuality(ctx, now)
		cancel()
		if err == nil {
			groupDegradationSnapshot.latest, groupDegradationSnapshot.vision = latest, visions
			groupDegradationSnapshot.ciallo = ciallos
		} else if groupDegradationSnapshot.latest == nil {
			return nil, nil, nil, errors.New("group quality temporarily unavailable")
		}
		groupDegradationSnapshot.updated = now
	}
	return groupDegradationSnapshot.latest, groupDegradationSnapshot.vision, groupDegradationSnapshot.ciallo, nil
}

func RecentDegradedGroups(ctx context.Context) (map[string]bool, error) {
	latest, visions, ciallos, err := cachedGroupQuality(ctx)
	if err != nil {
		return nil, err
	}
	result := map[string]bool{}
	cutoff := time.Now().Add(-AstraSamplingWindow).Unix()
	for group, h := range latest {
		if h.At >= cutoff && h.Closest != fingerprint.Model {
			result[group] = true
		}
	}
	for _, detector := range []map[string]DetectionHistory{visions, ciallos} {
		for group, h := range detector {
			if h.At >= cutoff && h.Quality == "degraded" {
				result[group] = true
			}
		}
	}
	return result, nil
}

func VisionUnavailableGroups(ctx context.Context) (map[string]bool, error) {
	_, visions, _, err := cachedGroupQuality(ctx)
	if err != nil {
		return nil, err
	}
	result := map[string]bool{}
	for group, h := range visions {
		if h.Quality == "incomplete" {
			result[group] = true
		}
	}
	return result, nil
}

func InvalidateGroupQualityCache() {
	groupDegradationSnapshot.Lock()
	defer groupDegradationSnapshot.Unlock()
	groupDegradationSnapshot.updated = time.Time{}
}