new-api-custom/controller/detection_worker.go

package controller

import (
	"context"
	"errors"
	"fmt"
	"sync"
	"time"

	"github.com/QuantumNous/new-api/common"
	"github.com/QuantumNous/new-api/model"
	benchmark "github.com/QuantumNous/new-api/pkg/model_benchmark"
	fingerprint "github.com/QuantumNous/new-api/pkg/model_fingerprint"
	"github.com/QuantumNous/new-api/service"
	"github.com/tidwall/gjson"
)

type detectionHandler struct{}

func (detectionHandler) Type() string { return "group_detection" }
func (detectionHandler) Enabled() bool {
	return common.GetEnvOrDefaultBool("MODEL_DETECTION_ENABLED", true)
}
func (detectionHandler) Interval() time.Duration { return 15 * time.Second }
func (detectionHandler) NewPayload() any         { return nil }

func (detectionHandler) Run(ctx context.Context, task *model.SystemTask, runner string) {
	_ = model.DeliverDetectionLogs()
	// Keep accounting facts, but expire bulky generated artwork after the public window.
	_ = model.DB.Model(&model.DetectionRun{}).Where("settled = ? AND created_at < ? AND html <> ?", true, time.Now().Add(-6*time.Hour).Unix(), "").Update("html", "").Error
	scheduleErr := scheduleBenchmarkDetections(ctx)
	if candyErr := scheduleAutomaticDetections(ctx, "candy"); scheduleErr == nil {
		scheduleErr = candyErr
	}
	var err error
	{
		var runs []model.DetectionRun
		err = model.DB.Where("settled = ?", false).Order("created_at asc").Find(&runs).Error
		if err == nil {
			jobs := make(chan model.DetectionRun)
			var workers sync.WaitGroup
			for range 2 {
				workers.Go(func() {
					for run := range jobs {
						processDetection(ctx, run)
					}
				})
			}
			for _, run := range runs {
				select {
				case jobs <- run:
				case <-ctx.Done():
				}
			}
			close(jobs)
			workers.Wait()
		}
	}
	if err == nil {
		err = scheduleErr
	}
	status := model.SystemTaskStatusSucceeded
	if err != nil {
		status = model.SystemTaskStatusFailed
	}
	finishSystemTaskHandler(task, runner, status, nil, err)
}

func scheduleBenchmarkDetections(ctx context.Context) error {
	return scheduleAutomaticDetections(ctx, "benchmark")
}

func scheduleAutomaticDetections(ctx context.Context, kind string) error {
	channels, schedules, err := service.AutomaticDetectionChannelsDue(ctx, kind, time.Now())
	if err != nil {
		return err
	}
	now := time.Now().Unix()
	for _, channel := range channels {
		groups := channel.GetGroups()
		if len(groups) == 0 {
			continue
		}
		price, reserve, priceErr := service.DetectionPricing("", groups[0], kind)
		if priceErr != nil {
			return priceErr
		}
		created := time.Unix(now, 0)
		price.Request.At = &created
		encoded, _ := common.Marshal(price)
		id, err := common.GenerateRandomCharsKey(32)
		if err != nil {
			return err
		}
		run := &model.DetectionRun{ID: id, RequestKey: fmt.Sprintf("%s:%d:%s", kind, channel.Id, schedules[channel.Id].LastResultID), ChannelID: channel.Id, GroupName: groups[0], Kind: kind, Status: "pending", CreatedAt: now, Price: string(encoded), Reserved: reserve}
		if err = model.CreateDetection(run); err != nil {
			return err
		}
	}
	return nil
}

func processDetection(ctx context.Context, run model.DetectionRun) {
	if ctx.Err() != nil {
		return
	}
	if run.Deadline == 0 {
		run.Deadline = time.Now().Unix() + 600
		if model.DB.Model(&model.DetectionRun{}).Where("id = ? AND deadline = 0", run.ID).Updates(map[string]any{"deadline": run.Deadline, "status": "running"}).Error != nil {
			return
		}
	}
	if time.Now().Unix() >= run.Deadline {
		var calls []model.DetectionCall
		if model.DB.Where("run_id = ?", run.ID).Find(&calls).Error != nil {
			return
		}
		for _, call := range calls {
			if call.State != "running" {
				continue
			}
			if cancel, ok := detectionCalls.Load(call.ID); ok {
				cancel.(context.CancelFunc)()
				return
			}
			// All proxy contexts expire before the task deadline. A vanished process
			// cannot resume or recharge this consumed credential.
			call.Error = "interrupted_usage_unknown"
			call.Quota = 0
			_ = model.FinishDetectionCall(&call)
		}
		finishDetection(run, "failed", "detection_timeout")
		return
	}
	if run.Status == "succeeded" || run.Status == "failed" {
		settleDetectionAndSchedule(run)
		return
	}
	if run.Kind != "fingerprint" && run.Kind != "benchmark" && run.Kind != "candy" {
		finishDetection(run, "failed", "detection_retired")
		return
	}
	processPaidFingerprint(ctx, run)
}

func finishDetection(run model.DetectionRun, status, reason string) {
	updates := map[string]any{"status": status, "error": reason}
	if status == "failed" && run.Kind == "benchmark" {
		updates["result"] = ""
	}
	if model.DB.Model(&model.DetectionRun{}).Where("id = ?", run.ID).Updates(updates).Error != nil {
		return
	}
	settleDetectionAndSchedule(run)
}

func settleDetectionAndSchedule(run model.DetectionRun) {
	if model.SettleDetection(run.ID) != nil || run.UserID != 0 || (run.Kind != "benchmark" && run.Kind != "candy") {
		return
	}
	// Reconciliation also repairs the clock if a process exits after settlement.
	if err := model.DB.Select("id", "channel_id", "finished_at", "status", "quality").First(&run, "id = ?", run.ID).Error; err != nil {
		return
	}
	outcome := "unknown"
	if run.Status == "succeeded" {
		outcome = run.Quality
	}
	_ = model.ApplyDetectionScheduleResult(context.Background(), run.ChannelID, run.Kind, run.ID, run.FinishedAt, outcome)
}

func processPaidFingerprint(ctx context.Context, run model.DetectionRun) {
	channel, err := model.GetChannelById(run.ChannelID, true)
	if err != nil || channel.Status != common.ChannelStatusEnabled {
		finishDetection(run, "failed", "channel_unavailable")
		return
	}
	output, calls := fingerprint.OutputTokens, fingerprint.ProbeCount
	if run.Kind == "benchmark" {
		output, calls = benchmark.OutputTokens, benchmark.ProbeCount
	}
	if run.Kind == "candy" {
		output, calls = service.CandyOutputTokens, 1
	}
	var price service.DetectionPrice
	if common.UnmarshalJsonStr(run.Price, &price) != nil || gjson.GetBytes(price.Request.Body, "max_output_tokens").Int() != int64(output) {
		finishDetection(run, "failed", "benchmark_changed")
		return
	}
	samples := []fingerprint.Sample{}
	for index := 0; index < calls; index++ {
		id := fmt.Sprintf("%s_%d", run.ID, index)
		var previous []model.DetectionCall
		if model.DB.Where("id = ?", id).Find(&previous).Error != nil {
			return
		}
		if len(previous) > 0 {
			if previous[0].State == "running" {
				return
			}
			var sample fingerprint.Sample
			if previous[0].Sample != "" && common.UnmarshalJsonStr(previous[0].Sample, &sample) == nil {
				samples = append(samples, sample)
			}
			continue
		}
		call := model.DetectionCall{ID: id, RunID: run.ID, SecretHash: model.DetectionSecretHash(id), State: "running", StartedAt: time.Now().Unix()}
		if model.DB.Create(&call).Error != nil {
			return
		}
		probeCtx, cancel := context.WithDeadline(ctx, time.Unix(run.Deadline-10, 0))
		detectionCalls.Store(call.ID, cancel)
		answer, count, probeErr := runModelDetectionProbe(probeCtx, channel, index, run.Kind)
		interrupted := probeCtx.Err() != nil
		cancel()
		available := probeErr == nil || answer.HasUsage
		call.Available = &available
		if !answer.HasUsage && (errors.Is(probeErr, fingerprint.ErrInvalidSample) || errors.Is(probeErr, fingerprint.ErrIncomplete) || interrupted) {
			call.Available = nil
		}
		if probeErr != nil {
			call.Error = "channel_call_failed"
		}
		if answer.HasUsage {
			usage, _ := common.Marshal(answer)
			call.Usage = string(usage)
			quota, quotaErr := service.DetectionUsageQuota(run.Price, answer.InputTokens, answer.OutputTokens, answer.CachedTokens, answer.CacheWriteTokens)
			if quotaErr != nil || quota > run.Reserved/calls {
				call.Error = "settlement_unavailable"
			} else {
				call.Quota = quota
			}
		} else {
			call.Error = "usage_unavailable"
		}
		if probeErr == nil && ((run.Kind == "candy" && answer.Text != "") || (run.Kind == "benchmark" && benchmark.ValidSample(answer.Text)) || (run.Kind == "fingerprint" && fingerprint.ValidArray(answer.Text, count))) {
			sample := fingerprint.Sample{Text: answer.Text, Count: count}
			samples = append(samples, sample)
			encoded, _ := common.Marshal(sample)
			call.Sample = string(encoded)
		}
		if model.FinishDetectionCall(&call) != nil {
			detectionCalls.Delete(call.ID)
			return
		}
		detectionCalls.Delete(call.ID)
		if ctx.Err() != nil {
			return
		}
	}
	var result fingerprint.Result
	if run.Kind == "benchmark" {
		secondary := make([]benchmark.Sample, len(samples))
		for i, sample := range samples {
			secondary[i] = benchmark.Sample{Text: sample.Text, Count: sample.Count}
		}
		scored, scoreErr := benchmark.Score(secondary)
		err = scoreErr
		result = fingerprint.Result{Status: scored.Status, Engine: scored.Engine, ValidSamples: scored.ValidSamples, Closest: scored.Closest, Weight: scored.Weight, Gap: scored.Gap, Matches: scored.Matches}
	} else if run.Kind == "candy" {
		if len(samples) != 1 {
			finishDetection(run, "failed", "answer_unavailable")
			return
		}
		result = fingerprint.Result{Status: service.CandyQuality(samples[0].Text), Engine: "candy-v1", ValidSamples: 1, Closest: samples[0].Text}
	} else {
		result, err = fingerprint.Score(samples)
	}
	if err != nil {
		finishDetection(run, "failed", "scoring_failed")
		return
	}
	quality := result.Status
	if run.Kind == "benchmark" {
		quality = service.FingerprintScheduleOutcome(result.Status)
	}
	encoded, _ := common.Marshal(result)
	if model.DB.Model(&model.DetectionRun{}).Where("id = ?", run.ID).Updates(map[string]any{"quality": quality, "result": string(encoded)}).Error != nil {
		return
	}
	finishDetection(run, "succeeded", "")
}