new-api-custom/controller/astra_sampling.go

package controller

import (
	"bytes"
	"context"
	"crypto/rand"
	"errors"
	"fmt"
	"math/big"
	"net/http"
	"net/http/httptest"
	"sync"
	"time"

	"github.com/QuantumNous/new-api/common"
	"github.com/QuantumNous/new-api/middleware"
	"github.com/QuantumNous/new-api/model"
	fingerprint "github.com/QuantumNous/new-api/pkg/model_fingerprint"
	"github.com/QuantumNous/new-api/relay"
	relaycommon "github.com/QuantumNous/new-api/relay/common"
	"github.com/QuantumNous/new-api/relaykit/dto"
	"github.com/QuantumNous/new-api/relaykit/types"
	"github.com/QuantumNous/new-api/service"
	"github.com/gin-gonic/gin"
)

type astraSamplingHandler struct{}

type astraSamplingTaskPayload struct {
	Force bool `json:"force,omitempty"`
}

func (astraSamplingHandler) Type() string               { return service.AstraSamplingTaskType }
func (astraSamplingHandler) Enabled() bool              { return service.AstraSamplingEnabled() }
func (astraSamplingHandler) Interval() time.Duration    { return service.AstraSamplingSchedulerInterval() }
func (astraSamplingHandler) ScheduleFromCreation() bool { return true }
func (astraSamplingHandler) NewPayload() any            { return nil }

func (astraSamplingHandler) ScheduleDue(ctx context.Context, now time.Time) (bool, error) {
	channels, err := service.AstraSamplingChannelsDue(ctx, now)
	return len(channels) > 0, err
}

func (astraSamplingHandler) Run(parent context.Context, task *model.SystemTask, runnerID string) {
	ctx, cancel := context.WithTimeout(parent, 25*time.Minute)
	defer cancel()
	payload := astraSamplingTaskPayload{}
	if err := task.DecodePayload(&payload); err != nil {
		finishSystemTaskHandler(task, runnerID, model.SystemTaskStatusFailed, nil, err)
		return
	}
	channels, err := service.AstraSamplingChannelsDue(ctx, time.Now())
	if payload.Force {
		channels, err = service.AstraSamplingChannels(ctx)
	}
	if err != nil {
		finishSystemTaskHandler(task, runnerID, model.SystemTaskStatusFailed, nil, errors.New("channel query failed"))
		return
	}
	state := service.AstraSamplingState{Rounds: []service.AstraSamplingRound{}}
	var mu sync.Mutex
	var persistErr error
	// Persist before sending, and after every probe. A failed lease or restart
	// never resumes a paid probe; unfinished records remain visible as errors.
	persist := func() bool {
		if persistErr != nil {
			return false
		}
		persistErr = model.UpdateSystemTaskState(task.TaskID, runnerID, state)
		if persistErr != nil {
			cancel()
		}
		return persistErr == nil
	}
	jobs := make(chan int)
	var workers sync.WaitGroup
	for worker := 0; worker < 2; worker++ {
		workers.Add(1)
		go func() {
			defer workers.Done()
			for id := range jobs {
				if ctx.Err() != nil {
					continue
				}
				mu.Lock()
				index := len(state.Rounds)
				state.Rounds = append(state.Rounds, service.AstraSamplingRound{ID: common.GetRandomString(16), ChannelID: id, StartedAt: time.Now().Unix(), Version: fingerprint.Version, Result: fingerprint.Result{Status: "running"}})
				ok := persist()
				mu.Unlock()
				if !ok {
					continue
				}
				channel, loadErr := model.GetChannelById(id, true)
				samples := []fingerprint.Sample{}
				reason := ""
				usage := []fingerprint.ProbeResponse{}
				if loadErr != nil || channel.Status != common.ChannelStatusEnabled {
					reason = "channel_unavailable"
				} else {
					for probe := 0; probe < 3 && ctx.Err() == nil; probe++ {
						mu.Lock()
						state.Rounds[index].Sent++
						ok = persist()
						mu.Unlock()
						if !ok {
							break
						}
						answer, count, probeErr := runAstraFingerprintProbe(ctx, channel)
						usage = append(usage, answer)
						if probeErr != nil {
							if errors.Is(probeErr, fingerprint.ErrInvalidSample) {
								reason = "invalid_sample"
							} else {
								reason = "transport_error"
								break
							}
						} else {
							samples = append(samples, fingerprint.Sample{Text: answer.Text, Count: count})
						}
						mu.Lock()
						state.Rounds[index].Usage = append([]fingerprint.ProbeResponse(nil), usage...)
						ok = persist()
						mu.Unlock()
						if !ok {
							break
						}
					}
				}
				result, scoreErr := fingerprint.Score(samples)
				if scoreErr != nil {
					reason = "scoring_error"
				}
				if ctx.Err() != nil {
					reason = "interrupted"
				}
				if reason != "" && reason != "invalid_sample" {
					result.Status = "error"
				}
				mu.Lock()
				state.Rounds[index].Result = result
				state.Rounds[index].Reason = reason
				state.Rounds[index].Usage = usage
				state.Rounds[index].FinishedAt = time.Now().Unix()
				if persist() && !payload.Force {
					if err := model.ApplyDetectionScheduleResult(ctx, id, "fingerprint", state.Rounds[index].ID, state.Rounds[index].FinishedAt, service.FingerprintScheduleOutcome(service.AstraSamplingResultStatus(result))); err != nil {
						persistErr = err
						cancel()
					}
				}
				mu.Unlock()
			}
		}()
	}
	for _, channel := range channels {
		select {
		case jobs <- channel.Id:
		case <-ctx.Done():
		}
	}
	close(jobs)
	workers.Wait()
	if persistErr != nil || ctx.Err() != nil {
		finishSystemTaskHandler(task, runnerID, model.SystemTaskStatusFailed, nil, errors.New("sampling interrupted"))
		return
	}
	finishSystemTaskHandler(task, runnerID, model.SystemTaskStatusSucceeded, map[string]any{"channels": len(state.Rounds), "model": fingerprint.Model}, nil)
}

func runAstraFingerprintProbe(parent context.Context, channel *model.Channel, paid ...bool) (fingerprint.ProbeResponse, int, error) {
	ctx, cancel := context.WithTimeout(parent, 180*time.Second)
	defer cancel()
	seed, err := rand.Int(rand.Reader, big.NewInt(82))
	if err != nil {
		return fingerprint.ProbeResponse{}, 0, err
	}
	count := 292 + int(seed.Int64()/2)
	prompt := fmt.Sprintf("Directly choose %d integers from 1 through 355, inclusive. Allow repeats. Do not sort, balance frequencies, repair duplicates, or deliberately create an arithmetic pattern.\nReply directly with one JSON integer array and no explanation. Do not call tools, read files, execute code or ask another model. Do not continue the preceding task.", count)
	if seed.Int64()%2 == 0 {
		prompt = fmt.Sprintf("直接选择 %d 个 1 到 355（含边界）的整数。允许重复；不要排序、平衡频数、修复重复，也不要刻意构造等差规律。\n直接回答一个 JSON 整数数组，不要解释。不要调用工具、读文件、运行代码或让其他模型代答；不要继续之前的任务。", count)
	}
	input, _ := common.Marshal([]map[string]any{{"role": "user", "content": []map[string]string{{"type": "input_text", "text": prompt}}}})
	stream := true
	limit := uint(4096)
	request := &dto.OpenAIResponsesRequest{Model: fingerprint.Model, Input: input, Instructions: []byte(`""`), Stream: &stream, MaxOutputTokens: &limit, Store: []byte("false"), Tools: []byte("[]"), ToolChoice: []byte(`"none"`)}
	encoded, err := common.Marshal(request)
	if err != nil {
		return fingerprint.ProbeResponse{}, count, err
	}
	c, _ := gin.CreateTestContext(httptest.NewRecorder())
	c.Request = httptest.NewRequestWithContext(ctx, http.MethodPost, "/v1/responses", bytes.NewReader(encoded))
	c.Request.Header.Set("Content-Type", "application/json")
	if apiErr := middleware.SetupContextForSelectedChannel(c, channel, fingerprint.Model); apiErr != nil {
		return fingerprint.ProbeResponse{}, count, errors.New("channel setup failed")
	}
	info, err := relaycommon.GenRelayInfo(c, types.RelayFormatOpenAIResponses, request, nil)
	if err != nil {
		return fingerprint.ProbeResponse{}, count, err
	}
	info.IsChannelTest = true
	info.DisablePing = true
	info.AstraDetectionDisabled = true
	info.AstraAutoSwitch = false
	adaptor, body, closer, apiErr := relay.PrepareResponsesRequest(c, info, request)
	if closer != nil {
		defer closer.Close()
	}
	if apiErr != nil {
		return fingerprint.ProbeResponse{}, count, errors.New("probe preparation failed")
	}
	if len(paid) > 0 && paid[0] {
		bounded, boundedCloser, boundErr := boundDetectionBody(body, request)
		if boundErr != nil {
			return fingerprint.ProbeResponse{}, count, boundErr
		}
		defer boundedCloser.Close()
		body = bounded
	}
	response, err := adaptor.DoRequest(c, info, body)
	if err != nil {
		return fingerprint.ProbeResponse{}, count, errors.New("probe transport failed")
	}
	resp, ok := response.(*http.Response)
	if !ok || resp == nil {
		return fingerprint.ProbeResponse{}, count, errors.New("unsupported probe transport")
	}
	defer resp.Body.Close()
	if resp.StatusCode != http.StatusOK {
		return fingerprint.ProbeResponse{}, count, errors.New("probe HTTP failure")
	}
	// Deliberately bypass DoResponse: this is a platform-funded probe, with no
	// customer token, wallet charge, session binding or generated response relay.
	answer, err := fingerprint.ReadResponse(resp.Body, resp.Header.Get("Content-Type"))
	return answer, count, err
}

func TestAstraFingerprint(c *gin.Context) {
	if !service.AstraSamplingEnabled() {
		c.JSON(http.StatusConflict, gin.H{"success": false, "message": "Astra sampling is disabled"})
		return
	}
	latest, err := model.GetLatestSystemTasks([]string{service.AstraSamplingTaskType})
	if err != nil {
		common.ApiError(c, err)
		return
	}
	if last := latest[service.AstraSamplingTaskType]; last != nil && time.Now().Unix()-last.CreatedAt < 300 {
		c.JSON(http.StatusTooManyRequests, gin.H{"success": false, "message": "Please wait five minutes between sampling rounds"})
		return
	}
	task, created, err := service.EnqueueSystemTask(service.AstraSamplingTaskType, astraSamplingTaskPayload{Force: true})
	if err != nil {
		common.ApiError(c, err)
		return
	}
	c.JSON(http.StatusOK, gin.H{"success": true, "data": gin.H{"task_id": task.TaskID, "created": created}})
}