feat: sync deployed v1.0.0-rc.38-custom.67

controller/astra_sampling.go

@@ -3,10 +3,7 @@ package controller
 import (
 	"bytes"
 	"context"
-	"crypto/rand"
 	"errors"
-	"fmt"
-	"math/big"
 	"net/http"
 	"net/http/httptest"
 	"sync"
@@ -15,6 +12,7 @@ import (
 	"github.com/QuantumNous/new-api/common"
 	"github.com/QuantumNous/new-api/middleware"
 	"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/relay"
 	relaycommon "github.com/QuantumNous/new-api/relay/common"
@@ -97,7 +95,7 @@ func (astraSamplingHandler) Run(parent context.Context, task *model.SystemTask,
 				if loadErr != nil || channel.Status != common.ChannelStatusEnabled {
 					reason = "channel_unavailable"
 				} else {
-					for probe := 0; probe < 3 && ctx.Err() == nil; probe++ {
+					for probe := 0; probe < fingerprint.ProbeCount && ctx.Err() == nil; probe++ {
 						mu.Lock()
 						state.Rounds[index].Sent++
 						ok = persist()
@@ -105,7 +103,7 @@ func (astraSamplingHandler) Run(parent context.Context, task *model.SystemTask,
 						if !ok {
 							break
 						}
-						answer, count, probeErr := runAstraFingerprintProbe(ctx, channel)
+						answer, count, probeErr := runAstraFingerprintProbe(ctx, channel, probe)
 						usage = append(usage, answer)
 						if probeErr != nil {
 							if errors.Is(probeErr, fingerprint.ErrInvalidSample) {
@@ -142,7 +140,7 @@ func (astraSamplingHandler) Run(parent context.Context, task *model.SystemTask,
 				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 {
+					if err := model.ApplyDetectionScheduleResult(ctx, id, "modeltrace", state.Rounds[index].ID, state.Rounds[index].FinishedAt, service.FingerprintScheduleOutcome(service.AstraSamplingResultStatus(result))); err != nil {
 						persistErr = err
 						cancel()
 					}
@@ -166,22 +164,39 @@ func (astraSamplingHandler) Run(parent context.Context, task *model.SystemTask,
 	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) {
+func runAstraFingerprintProbe(parent context.Context, channel *model.Channel, index int) (fingerprint.ProbeResponse, int, error) {
+	return runModelDetectionProbe(parent, channel, index, "fingerprint")
+}
+
+func runModelDetectionProbe(parent context.Context, channel *model.Channel, index int, kind string) (fingerprint.ProbeResponse, int, error) {
 	ctx, cancel := context.WithTimeout(parent, 180*time.Second)
 	defer cancel()
-	seed, err := rand.Int(rand.Reader, big.NewInt(82))
+	var probe fingerprint.Probe
+	var err error
+	if kind == "fingerprint" {
+		probe, err = fingerprint.ProbeAt(index)
+	}
 	if err != nil {
-		return fingerprint.ProbeResponse{}, 0, err
+		return fingerprint.ProbeResponse{}, index, 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)
+	count, output := probe.Count, fingerprint.OutputTokens
+	messages := []map[string]any{{"role": "user", "content": probe.Prompt}}
+	if kind == "benchmark" {
+		secondary, err := benchmark.ProbeAt(index)
+		if err != nil {
+			return fingerprint.ProbeResponse{}, index, err
+		}
+		count, output = index, benchmark.OutputTokens
+		messages = []map[string]any{{"role": "system", "content": secondary.System}, {"role": "user", "content": secondary.Prompt}}
 	}
-	input, _ := common.Marshal([]map[string]any{{"role": "user", "content": []map[string]string{{"type": "input_text", "text": prompt}}}})
+	if kind == "candy" {
+		count, output = 1, service.CandyOutputTokens
+		messages = []map[string]any{{"role": "user", "content": service.CandyPrompt}}
+	}
+	input, _ := common.Marshal(messages)
 	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"`)}
+	limit := uint(output)
+	request := &dto.OpenAIResponsesRequest{Model: fingerprint.Model, Input: input, Stream: &stream, MaxOutputTokens: &limit, Store: []byte("false"), Reasoning: &dto.Reasoning{Effort: "low"}}
 	encoded, err := common.Marshal(request)
 	if err != nil {
 		return fingerprint.ProbeResponse{}, count, err
@@ -207,7 +222,7 @@ func runAstraFingerprintProbe(parent context.Context, channel *model.Channel, pa
 	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