294 lines
10 KiB
Go
294 lines
10 KiB
Go
package controller
|
|
|
|
import (
|
|
"bytes"
|
|
"context"
|
|
"errors"
|
|
"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"
|
|
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"
|
|
"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 < fingerprint.ProbeCount && ctx.Err() == nil; probe++ {
|
|
mu.Lock()
|
|
state.Rounds[index].Sent++
|
|
ok = persist()
|
|
mu.Unlock()
|
|
if !ok {
|
|
break
|
|
}
|
|
answer, count, probeErr := runAstraFingerprintProbe(ctx, channel, probe)
|
|
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 persisted := persist(); persisted {
|
|
service.InvalidateGroupQualityCache()
|
|
publishGroupStatusChange()
|
|
if !payload.Force {
|
|
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()
|
|
}
|
|
}
|
|
}
|
|
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, 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()
|
|
var probe fingerprint.Probe
|
|
var err error
|
|
if kind == "fingerprint" {
|
|
probe, err = fingerprint.ProbeAt(index)
|
|
}
|
|
if err != nil {
|
|
return fingerprint.ProbeResponse{}, index, err
|
|
}
|
|
effort := "low"
|
|
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}}
|
|
}
|
|
if kind == "vision" {
|
|
count, output = 1, service.VisionOutputTokens
|
|
effort = "medium"
|
|
messages = []map[string]any{{"role": "user", "content": []map[string]any{
|
|
{"type": "input_text", "text": service.VisionPrompt},
|
|
{"type": "input_image", "image_url": service.VisionImageURL, "detail": "high"},
|
|
}}}
|
|
}
|
|
if kind == "ciallo" {
|
|
count, output = 1, service.VisionOutputTokens
|
|
effort = "medium"
|
|
messages = []map[string]any{{"role": "user", "content": service.CialloPrompt}}
|
|
}
|
|
input, _ := common.Marshal(messages)
|
|
stream := true
|
|
limit := uint(output)
|
|
request := &dto.OpenAIResponsesRequest{Model: fingerprint.Model, Input: input, Stream: &stream, MaxOutputTokens: &limit, Store: []byte("false"), Reasoning: &dto.Reasoning{Effort: effort}}
|
|
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")
|
|
}
|
|
{
|
|
bounded, boundedCloser, boundErr := boundDetectionBody(body, request)
|
|
if boundErr != nil {
|
|
return fingerprint.ProbeResponse{}, count, boundErr
|
|
}
|
|
defer boundedCloser.Close()
|
|
body = bounded
|
|
}
|
|
started := time.Now()
|
|
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()
|
|
stopClose := context.AfterFunc(ctx, func() { _ = resp.Body.Close() })
|
|
defer stopClose()
|
|
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.
|
|
var firstTokenMS *int64
|
|
answer, err := fingerprint.ReadResponse(resp.Body, resp.Header.Get("Content-Type"), func() {
|
|
elapsed := time.Since(started).Milliseconds()
|
|
firstTokenMS = &elapsed
|
|
})
|
|
answer.FirstTokenMS = firstTokenMS
|
|
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}})
|
|
}
|