cneicy/new-api-custom
1
0 Fork 0
Code Issues 1 Pull Requests Packages Projects Releases Wiki Activity

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

Browse Source
This commit is contained in:
cneicy committed 2026-10-06 00:07:10 +08:00
1 parent d2e03ab7e8
commit ffc8cf56db
154 files changed
+39540 -20509
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

CUSTOMIZATION.md
+2 -2
View File
@@ -1,6 +1,6 @@
# 自定义版本说明
当前版本为 v1.0.0-rc.38-custom.44-20260926，基于官方 v1.0.0-rc.38，对应 2026-09-26 已部署的自定义源码。
当前版本为 v1.0.0-rc.38-custom.67-20261004，基于官方 v1.0.0-rc.38，对应 2026-10-04 已部署的自定义源码。
## 主要定制
@@ -17,7 +17,7 @@
## 从源码构建
docker build -f Dockerfile.local-release -t local/new-api:v1.0.0-rc.38-custom.44-20260926 .
docker build -f Dockerfile.local-release -t local/new-api:v1.0.0-rc.38-custom.67-20261004 .
此入口从 Bun、Go 和 Debian 基础镜像构建前端与服务，不要求预先存在旧版 local/new-api 镜像。VERSION 用于编译版本号；Compose 如果设置了 VERSION 环境变量，也应同步更新。
Dockerfile.affinity-switch
+3
View File
@@ -0,0 +1,3 @@
FROM local/new-api:v1.0.0-rc.38-custom.45-20260926
COPY new-api /new-api
LABEL org.opencontainers.image.version="v1.0.0-rc.38-custom.46-20260927"
Dockerfile.browser-performance
+4
View File
@@ -0,0 +1,4 @@
FROM local/new-api:v1.0.0-rc.38-custom.59-20260929
COPY new-api /new-api
LABEL org.opencontainers.image.version="v1.0.0-rc.38-custom.60-20260929"
LABEL custom.patch="browser-performance"
Dockerfile.candy-test
+3
View File
@@ -0,0 +1,3 @@
FROM local/new-api:v1.0.0-rc.38-custom.66-20261004
COPY new-api /new-api
LABEL org.opencontainers.image.version="v1.0.0-rc.38-custom.67"
Dockerfile.ccswitch
+3
View File
@@ -0,0 +1,3 @@
FROM local/new-api:v1.0.0-rc.38-custom.44-20260926
COPY new-api /new-api
LABEL org.opencontainers.image.version="v1.0.0-rc.38-custom.45-20260926"
Dockerfile.combined
+3
View File
@@ -0,0 +1,3 @@
FROM local/new-api:v1.0.0-rc.38-custom.51-20260928
COPY new-api /new-api
LABEL custom.patch="automatic-routing-pool-failover-model-costs"
Dockerfile.downgrade-rate
+3
View File
@@ -0,0 +1,3 @@
FROM local/new-api:v1.0.0-rc.38-custom.50-20260927
COPY new-api /new-api
LABEL org.opencontainers.image.version="v1.0.0-rc.38-custom.51-20260928"
Dockerfile.fingerprint-color
+3
View File
@@ -0,0 +1,3 @@
FROM local/new-api:v1.0.0-rc.38-custom.65-20261004
COPY new-api /new-api
LABEL org.opencontainers.image.version="v1.0.0-rc.38-custom.66"
Dockerfile.group-degradation
+3
View File
@@ -0,0 +1,3 @@
FROM local/new-api:v1.0.0-rc.38-custom.47-20260927
COPY new-api /new-api
LABEL org.opencontainers.image.version="v1.0.0-rc.38-custom.48-20260927"
Dockerfile.local-release
+1 -1
View File
@@ -25,7 +25,7 @@ RUN apt-get update \
COPY --from=build /build/new-api /
COPY LICENSE NOTICE THIRD-PARTY-LICENSES.md /licenses/
COPY pkg/model_fingerprint/LICENSE-*.txt /licenses/model-fingerprint/
LABEL org.opencontainers.image.version="v1.0.0-rc.38-custom.44-20260926"
LABEL org.opencontainers.image.version="v1.0.0-rc.38-custom.67-20261004"
EXPOSE 3000
WORKDIR /data
ENTRYPOINT ["/new-api"]
Dockerfile.meow
+9
View File
@@ -0,0 +1,9 @@
FROM debian:bookworm-slim@sha256:f06537653ac770703bc45b4b113475bd402f451e85223f0f2837acbf89ab020a
RUN apt-get update && apt-get install -y --no-install-recommends ca-certificates tzdata libasan8 wget && rm -rf /var/lib/apt/lists/* && update-ca-certificates
COPY new-api /new-api
COPY LICENSE NOTICE THIRD-PARTY-LICENSES.md /licenses/
COPY pkg/model_fingerprint/LICENSE-*.txt /licenses/model-fingerprint/
LABEL org.opencontainers.image.version="v1.0.0-rc.38-custom.64"
EXPOSE 3000
WORKDIR /data
ENTRYPOINT ["/new-api"]
Dockerfile.modeltrace
+10
View File
@@ -0,0 +1,10 @@
FROM debian:bookworm-slim@sha256:f06537653ac770703bc45b4b113475bd402f451e85223f0f2837acbf89ab020a
RUN apt-get update && apt-get install -y --no-install-recommends ca-certificates tzdata libasan8 wget && rm -rf /var/lib/apt/lists/* && update-ca-certificates
COPY new-api /new-api
COPY LICENSE NOTICE THIRD-PARTY-LICENSES.md /licenses/
COPY pkg/model_fingerprint/LICENSE-*.txt /licenses/model-fingerprint/
COPY pkg/model_benchmark/LICENSE-*.txt /licenses/model-benchmark/
LABEL org.opencontainers.image.version="v1.0.0-rc.38-custom.65"
EXPOSE 3000
WORKDIR /data
ENTRYPOINT ["/new-api"]
Dockerfile.stream-health
+3
View File
@@ -0,0 +1,3 @@
FROM local/new-api:v1.0.0-rc.38-custom.49-20260927
COPY new-api /new-api
LABEL org.opencontainers.image.version="v1.0.0-rc.38-custom.50-20260927"
THIRD-PARTY-LICENSES.md
+7
View File
@@ -326,3 +326,10 @@ For more information, please refer to https://unlicense.org/
## Numerical model fingerprint assets
The embedded bank and scorer in `pkg/model_fingerprint` are from [gpt-model-check](https://github.com/Katao123/gpt-model-check), commit `3d6716a0903cf061daea001d5d44ee410846634e`, using the ModelTrace bank dated 2026-09-05. The upstream gpt-model-check, ModelTrace, and is-gpt-nerfed MIT notices are included alongside the sources and under `/licenses/model-fingerprint/` in the image. Relative candidate weights are not calibrated identity probabilities.
## Meow model fingerprint
The active short-probe benchmark and predictive scorer are adapted from chen-006/meow-llm-detector, commit 0f31d4a, benchmark 4.5.4-predictive.20260924.2 (quick tier, 32 requests). License: PolyForm Noncommercial 1.0.0; see pkg/model_fingerprint/LICENSE-meow.txt. Required Notice: Copyright 2026 chen-006 and contributors. Original project: https://github.com/chen-006/gpt56_api_detector. Match scores are not identity probabilities.
ModelTrace primary fingerprint: https://github.com/xqy2006/ModelTrace, commit d4131b30243dfa05e70180b5eedde742103f1d73, MIT. Scorer and unified candidate bank are embedded locally; see pkg/model_fingerprint/LICENSE-ModelTrace.txt.
Meow 4.5.4 predictive benchmark remains an auxiliary signal; see pkg/model_benchmark/LICENSE-meow.txt.
VERSION
+1 -1
View File
@@ -1 +1 @@
v1.0.0-rc.38-custom.44-20260926
v1.0.0-rc.38-custom.67
controller/astra_sampling.go
+32 -17
View File
@@ -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
controller/astra_sampling_test.go
+12 -2
View File
@@ -5,11 +5,14 @@ import (
"io"
"net/http"
"net/http/httptest"
"strings"
"testing"
"github.com/QuantumNous/new-api/common"
"github.com/QuantumNous/new-api/model"
fingerprint "github.com/QuantumNous/new-api/pkg/model_fingerprint"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
"github.com/tidwall/gjson"
)
@@ -17,6 +20,13 @@ func TestAstraSamplingUsesIsolatedResponsesWithoutGuard(t *testing.T) {
requests := 0
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
requests++
bodyPreview, readErr := io.ReadAll(r.Body)
require.NoError(t, readErr)
r.Body = io.NopCloser(strings.NewReader(string(bodyPreview)))
assert.EqualValues(t, fingerprint.OutputTokens, gjson.GetBytes(bodyPreview, "max_output_tokens").Int())
assert.Equal(t, "low", gjson.GetBytes(bodyPreview, "reasoning.effort").String())
assert.Contains(t, gjson.GetBytes(bodyPreview, "input.0.content").String(), "355")
body, _ := io.ReadAll(r.Body)
if r.URL.Path != "/v1/responses" || gjson.GetBytes(body, "model").String() != fingerprint.Model || !gjson.GetBytes(body, "stream").Bool() {
t.Errorf("wrong probe request %s %s", r.URL.Path, body)
@@ -30,8 +40,8 @@ func TestAstraSamplingUsesIsolatedResponsesWithoutGuard(t *testing.T) {
defer server.Close()
url := server.URL
channel := &model.Channel{Id: 1, Type: 1, Key: "test-only", BaseURL: &url, Status: common.ChannelStatusEnabled, Models: fingerprint.Model}
r, count, err := runAstraFingerprintProbe(context.Background(), channel)
if err != nil || r.DeclaredModel != "gpt-5.6-luna" || count < 292 || count > 332 || requests != 1 {
r, count, err := runAstraFingerprintProbe(context.Background(), channel, 0)
if err != nil || r.DeclaredModel != "gpt-5.6-luna" || count < 292 || count > 304 || requests != 1 {
t.Fatalf("probe %+v count=%d requests=%d err=%v", r, count, requests, err)
}
}
controller/automatic_model_list_test.go
+80
View File
@@ -0,0 +1,80 @@
package controller
import (
"github.com/QuantumNous/new-api/common"
"github.com/QuantumNous/new-api/constant"
"github.com/QuantumNous/new-api/middleware"
"github.com/QuantumNous/new-api/model"
"github.com/QuantumNous/new-api/setting"
"github.com/QuantumNous/new-api/setting/ratio_setting"
"github.com/gin-gonic/gin"
"github.com/stretchr/testify/require"
"net/http"
"net/http/httptest"
"testing"
)
func TestListModelsAutomaticPoolAllFormats(t *testing.T) {
db := setupModelListControllerTestDB(t)
withSelfUseModeDisabled(t)
oldGroups, oldRatios, oldModelRatios := setting.UserUsableGroups2JSONString(), ratio_setting.GroupRatio2JSONString(), ratio_setting.ModelRatio2JSONString()
oldMemory := common.MemoryCacheEnabled
common.MemoryCacheEnabled = false
t.Cleanup(func() {
common.MemoryCacheEnabled = oldMemory
require.NoError(t, setting.UpdateUserUsableGroupsByJSONString(oldGroups))
require.NoError(t, ratio_setting.UpdateGroupRatioByJSONString(oldRatios))
require.NoError(t, ratio_setting.UpdateModelRatioByJSONString(oldModelRatios))
})
require.NoError(t, setting.UpdateUserUsableGroupsByJSONString(`{"list-a":"","list-b":""}`))
require.NoError(t, ratio_setting.UpdateGroupRatioByJSONString(`{"list-a":1,"list-b":1}`))
require.NoError(t, ratio_setting.UpdateModelRatioByJSONString(`{"zz-list-a":1,"zz-list-b":1,"zz-shared":1,"zz-disabled":1}`))
require.NoError(t, db.Create(&[]model.Channel{{Id: 5601, Type: constant.ChannelTypeOpenAI, Status: 1}, {Id: 5602, Type: constant.ChannelTypeOpenAI, Status: 1}, {Id: 5603, Status: 2}}).Error)
require.NoError(t, db.Create(&[]model.Ability{
{ChannelId: 5601, Group: "list-a", Model: "zz-list-a", Enabled: true},
{ChannelId: 5602, Group: "list-b", Model: "zz-list-b", Enabled: true},
{ChannelId: 5601, Group: "list-a", Model: "zz-shared", Enabled: true},
{ChannelId: 5602, Group: "list-b", Model: "zz-shared", Enabled: true},
{ChannelId: 5603, Group: "list-b", Model: "zz-disabled", Enabled: true},
{ChannelId: 5602, Group: "list-b", Model: "zz-unpriced", Enabled: true},
}).Error)
for _, preference := range []string{"cost", "speed", "stable"} {
for _, limited := range []bool{false, true} {
for _, format := range []int{constant.ChannelTypeOpenAI, constant.ChannelTypeAnthropic, constant.ChannelTypeGemini} {
policy, err := common.Marshal(model.TokenRoutingPolicy{Version: 1, Mode: "smart", AutoPreference: preference})
require.NoError(t, err)
token := &model.Token{Group: "auto", RoutingPolicy: string(policy), ModelLimitsEnabled: limited, ModelLimits: "zz-list-b,zz-disabled,zz-unpriced"}
engine := gin.New()
engine.GET("/v1/models", func(c *gin.Context) {
common.SetContextKey(c, constant.ContextKeyUserGroup, "list-a")
require.NoError(t, middleware.SetupContextForToken(c, token))
ListModels(c, format)
})
res := httptest.NewRecorder()
engine.ServeHTTP(res, httptest.NewRequest(http.MethodGet, "/v1/models", nil))
require.Equal(t, 200, res.Code)
var body struct {
Data []struct {
ID string `json:"id"`
} `json:"data"`
Models []struct {
Name string `json:"name"`
} `json:"models"`
}
require.NoError(t, common.Unmarshal(res.Body.Bytes(), &body))
names := []string{}
for _, item := range body.Data {
names = append(names, item.ID)
}
for _, item := range body.Models {
names = append(names, item.Name)
}
expected := []string{"zz-list-a", "zz-list-b", "zz-shared"}
if limited {
expected = []string{"zz-list-b"}
}
require.ElementsMatch(t, expected, names, "preference=%s limited=%v format=%d", preference, limited, format)
}
}
}
}
controller/channel-test.go
+3 -2
View File
@@ -239,6 +239,7 @@ func testChannel(ctx context.Context, channel *model.Channel, testUserID int, te
}
}
isStream = info.IsStream
info.IsChannelTest = true
info.InitChannelMeta(c)
@@ -666,7 +667,7 @@ func validateTestResponseBody(respBody []byte, isStream bool) error {
}
func shouldUseStreamForAutomaticChannelTest(channel *model.Channel) bool {
return channel != nil && channel.Type == constant.ChannelTypeCodex
return channel != nil
}
func detectErrorMessageFromJSONBytes(jsonBytes []byte) string {
@@ -863,7 +864,7 @@ func TestChannel(c *gin.Context) {
//}()
testModel := c.Query("model")
endpointType := c.Query("endpoint_type")
isStream, _ := strconv.ParseBool(c.Query("stream"))
isStream, _ := strconv.ParseBool(c.DefaultQuery("stream", "true"))
testUserID, err := resolveChannelTestUserID(c)
if err != nil {
common.ApiError(c, err)
controller/detection.go
+12 -13
View File
@@ -16,6 +16,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"
@@ -41,7 +42,7 @@ func detectionInput(c *gin.Context) (detectionRequest, *model.User, service.Dete
if err := c.ShouldBindJSON(&input); err != nil {
return input, nil, service.DetectionPrice{}, 0, errors.New("invalid detection request")
}
if input.Kind != "pelican" && input.Kind != "fingerprint" {
if input.Kind != "benchmark" && input.Kind != "fingerprint" {
return input, nil, service.DetectionPrice{}, 0, errors.New("invalid detection type")
}
user, err := model.GetUserById(c.GetInt("id"), false)
@@ -62,9 +63,9 @@ func EstimateDetection(c *gin.Context) {
return
}
encoded, _ := common.Marshal(price)
calls := 1
calls := benchmark.ProbeCount
if input.Kind == "fingerprint" {
calls = 3
calls = fingerprint.ProbeCount
}
c.JSON(200, gin.H{"success": true, "data": gin.H{"reserved": reserve, "ratio": price.Snapshot.GroupRatio, "calls": calls, "price_version": model.DetectionSecretHash(string(encoded))}})
}
@@ -267,7 +268,7 @@ func detectionProxyRequest(data []byte) (*dto.OpenAIResponsesRequest, error) {
return nil, errors.New("input too large")
}
// Service-selected bounds cannot be overridden by the bearer.
fields = map[string]any{"model": fingerprint.Model, "input": fields["input"], "instructions": instructions, "stream": true, "max_output_tokens": 16384, "reasoning": map[string]string{"effort": "high"}, "store": false, "tools": []any{}}
fields = map[string]any{"model": fingerprint.Model, "input": fields["input"], "instructions": instructions, "stream": true, "max_output_tokens": 16384, "reasoning": map[string]string{"effort": "medium"}, "store": false, "tools": []any{}}
encoded, _ := common.Marshal(fields)
var request dto.OpenAIResponsesRequest
err = common.Unmarshal(encoded, &request)
@@ -342,14 +343,6 @@ func DetectionProxy(c *gin.Context) {
}
}
}
if success {
if html := detectionHTML(raw, contentType); html != "" {
if saveErr := model.DB.Model(&model.DetectionRun{}).Where("id = ?", run.ID).Update("html", html).Error; saveErr != nil {
call.Error = "artwork_save_failed"
success = false
}
}
}
if err = model.FinishDetectionCall(call); err != nil {
success = false
}
@@ -481,7 +474,13 @@ func boundDetectionBody(body common.ReplayableBody, request *dto.OpenAIResponses
prepared["instructions"] = instructions
prepared["max_output_tokens"] = *request.MaxOutputTokens
if *request.MaxOutputTokens == 16384 {
prepared["reasoning"] = map[string]string{"effort": "high"}
prepared["reasoning"] = map[string]string{"effort": "medium"}
}
if *request.MaxOutputTokens == fingerprint.OutputTokens || *request.MaxOutputTokens == benchmark.OutputTokens {
prepared["reasoning"] = map[string]string{"effort": "low"}
delete(prepared, "instructions")
delete(prepared, "temperature")
delete(prepared, "top_p")
}
prepared["tools"] = []any{}
prepared["tool_choice"] = "none"
controller/detection_test.go
+217
View File
@@ -2,6 +2,8 @@ package controller
import (
"bufio"
"context"
"fmt"
"github.com/QuantumNous/new-api/common"
"github.com/QuantumNous/new-api/model"
"github.com/QuantumNous/new-api/pkg/billingexpr"
@@ -11,10 +13,15 @@ import (
"github.com/glebarez/sqlite"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
"github.com/tidwall/gjson"
"gorm.io/driver/mysql"
"gorm.io/driver/postgres"
"gorm.io/gorm"
"gorm.io/gorm/schema"
"io"
"net/http"
"net/http/httptest"
"os"
"strings"
"testing"
"time"
@@ -32,6 +39,10 @@ func TestDetectionProxyFlushesBeforeGenerationAndConsumesCredentialOnce(t *testi
require.NoError(t, db.AutoMigrate(&model.Channel{}, &model.DetectionRun{}, &model.DetectionCall{}))
release := make(chan struct{})
upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
body, readErr := io.ReadAll(r.Body)
require.NoError(t, readErr)
assert.Equal(t, "medium", gjson.GetBytes(body, "reasoning.effort").String())
select {
case <-release:
case <-r.Context().Done():
@@ -169,3 +180,209 @@ func TestDetectionLatestPerChannelStatus(t *testing.T) {
})
}
}
func TestDetectionMeowPaidRoundAndResume(t *testing.T) {
oldRedis := common.RedisEnabled
common.RedisEnabled = false
t.Cleanup(func() { common.RedisEnabled = oldRedis })
db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{})
require.NoError(t, err)
conn, err := db.DB()
require.NoError(t, err)
conn.SetMaxOpenConns(1)
old := model.DB
model.DB = db
t.Cleanup(func() { model.DB = old; _ = conn.Close() })
require.NoError(t, db.AutoMigrate(&model.Channel{}, &model.DetectionRun{}, &model.DetectionCall{}, &model.User{}))
data, err := os.ReadFile("../pkg/model_fingerprint/testdata/meow-reference.json")
require.NoError(t, err)
var fixtures []struct {
Samples []fingerprint.Sample `json:"samples"`
}
require.NoError(t, common.Unmarshal(data, &fixtures))
calls := 0
upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
require.Less(t, calls, 32)
answer := fixtures[0].Samples[calls].Text
calls++
body, encodeErr := common.Marshal(map[string]any{"status": "completed", "usage": map[string]int{"input_tokens": 10, "output_tokens": 2}, "output": []any{map[string]any{"type": "message", "content": []any{map[string]string{"type": "output_text", "text": answer}}}}})
require.NoError(t, encodeErr)
w.Header().Set("Content-Type", "application/json")
_, _ = w.Write(body)
}))
defer upstream.Close()
url := upstream.URL
channel := model.Channel{Type: 1, Key: "test-only", BaseURL: &url, Status: common.ChannelStatusEnabled, Models: fingerprint.Model}
require.NoError(t, db.Create(&channel).Error)
owner := model.User{Username: "meow-test", Quota: 1000, Status: common.UserStatusEnabled}
require.NoError(t, db.Create(&owner).Error)
price, err := common.Marshal(service.DetectionPrice{Snapshot: billingexpr.BillingSnapshot{ExprString: "p+c", ExprHash: billingexpr.ExprHashString("p+c"), GroupRatio: 1, QuotaPerUnit: 500000, ExprVersion: 1}, Request: billingexpr.RequestInput{Body: []byte(`{"max_output_tokens":128,"reasoning":{"effort":"low"}}`)}})
require.NoError(t, err)
run := model.DetectionRun{ID: "meow-test", RequestKey: "meow-test", UserID: owner.Id, Kind: "benchmark", ChannelID: channel.Id, Status: "pending", Deadline: time.Now().Unix() + 60, Reserved: 320, Price: string(price)}
require.NoError(t, model.CreateDetection(&run))
processPaidFingerprint(context.Background(), run)
require.NoError(t, db.First(&run, "id = ?", run.ID).Error)
assert.Equal(t, "succeeded", run.Status)
assert.Equal(t, "normal", run.Quality)
assert.True(t, run.Settled)
assert.Equal(t, 32, calls)
assert.Equal(t, 192, run.Charged)
require.NoError(t, db.First(&owner, owner.Id).Error)
assert.Equal(t, 808, owner.Quota)
processPaidFingerprint(context.Background(), run)
assert.Equal(t, 32, calls, "resume must not resend or recharge")
require.NoError(t, db.First(&owner, owner.Id).Error)
assert.Equal(t, 808, owner.Quota)
}
func TestDetectionModelTracePaidRoundAndResume(t *testing.T) {
oldRedis := common.RedisEnabled
common.RedisEnabled = false
t.Cleanup(func() { common.RedisEnabled = oldRedis })
db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{})
require.NoError(t, err)
conn, err := db.DB()
require.NoError(t, err)
conn.SetMaxOpenConns(1)
old := model.DB
model.DB = db
t.Cleanup(func() { model.DB = old; _ = conn.Close() })
require.NoError(t, db.AutoMigrate(&model.Channel{}, &model.DetectionRun{}, &model.DetectionCall{}, &model.User{}))
data, err := os.ReadFile("../pkg/model_fingerprint/testdata/modeltrace-reference.json")
require.NoError(t, err)
var fixtures []struct {
Samples []fingerprint.Sample `json:"samples"`
}
require.NoError(t, common.Unmarshal(data, &fixtures))
calls := 0
upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
require.Less(t, calls, 3)
answer := fixtures[1].Samples[calls].Text
calls++
body, encodeErr := common.Marshal(map[string]any{"status": "completed", "usage": map[string]int{"input_tokens": 10, "output_tokens": 2}, "output": []any{map[string]any{"type": "message", "content": []any{map[string]string{"type": "output_text", "text": answer}}}}})
require.NoError(t, encodeErr)
w.Header().Set("Content-Type", "application/json")
_, _ = w.Write(body)
}))
defer upstream.Close()
url := upstream.URL
channel := model.Channel{Type: 1, Key: "test-only", BaseURL: &url, Status: common.ChannelStatusEnabled, Models: fingerprint.Model}
require.NoError(t, db.Create(&channel).Error)
owner := model.User{Username: "meow-test", Quota: 1000, Status: common.UserStatusEnabled}
require.NoError(t, db.Create(&owner).Error)
price, err := common.Marshal(service.DetectionPrice{Snapshot: billingexpr.BillingSnapshot{ExprString: "p+c", ExprHash: billingexpr.ExprHashString("p+c"), GroupRatio: 1, QuotaPerUnit: 500000, ExprVersion: 1}, Request: billingexpr.RequestInput{Body: []byte(`{"max_output_tokens":4096,"reasoning":{"effort":"low"}}`)}})
require.NoError(t, err)
run := model.DetectionRun{ID: "meow-test", RequestKey: "meow-test", UserID: owner.Id, Kind: "fingerprint", ChannelID: channel.Id, Status: "pending", Deadline: time.Now().Unix() + 60, Reserved: 320, Price: string(price)}
require.NoError(t, model.CreateDetection(&run))
processPaidFingerprint(context.Background(), run)
require.NoError(t, db.First(&run, "id = ?", run.ID).Error)
assert.Equal(t, "succeeded", run.Status)
assert.Equal(t, "match", run.Quality)
assert.True(t, run.Settled)
assert.Equal(t, 3, calls)
assert.Equal(t, 18, run.Charged)
require.NoError(t, db.First(&owner, owner.Id).Error)
assert.Equal(t, 982, owner.Quota)
processPaidFingerprint(context.Background(), run)
assert.Equal(t, 3, calls, "resume must not resend or recharge")
require.NoError(t, db.First(&owner, owner.Id).Error)
assert.Equal(t, 982, owner.Quota)
}
func TestCandyAutomaticRoundAndResume(t *testing.T) {
for _, dialect := range []string{"sqlite", "mysql", "postgres"} {
t.Run(dialect, func(t *testing.T) {
oldRedis := common.RedisEnabled
common.RedisEnabled = false
t.Cleanup(func() { common.RedisEnabled = oldRedis })
var driver gorm.Dialector = sqlite.Open(":memory:")
if dialect == "mysql" {
if os.Getenv("TEST_MYSQL_DSN") == "" {
t.Skip("TEST_MYSQL_DSN unset")
}
driver = mysql.Open(os.Getenv("TEST_MYSQL_DSN"))
}
if dialect == "postgres" {
if os.Getenv("TEST_POSTGRES_DSN") == "" {
t.Skip("TEST_POSTGRES_DSN unset")
}
driver = postgres.Open(os.Getenv("TEST_POSTGRES_DSN"))
}
db, err := gorm.Open(driver, &gorm.Config{NamingStrategy: schema.NamingStrategy{TablePrefix: fmt.Sprintf("candy_%d_", time.Now().UnixNano())}})
require.NoError(t, err)
conn, err := db.DB()
require.NoError(t, err)
conn.SetMaxOpenConns(1)
oldType, oldLogType := common.MainDatabaseType(), common.LogDatabaseType()
common.SetDatabaseTypes(common.DatabaseType(dialect), common.DatabaseType(dialect))
t.Cleanup(func() { common.SetDatabaseTypes(oldType, oldLogType) })
old := model.DB
model.DB = db
t.Cleanup(func() {
for _, row := range []any{&model.Channel{}, &model.DetectionRun{}, &model.DetectionCall{}, &model.User{}, &model.DetectionSchedule{}, &model.SystemTask{}} {
require.NoError(t, db.Migrator().DropTable(row))
}
model.DB = old
_ = conn.Close()
})
require.NoError(t, db.AutoMigrate(&model.Channel{}, &model.DetectionRun{}, &model.DetectionCall{}, &model.User{}, &model.DetectionSchedule{}, &model.SystemTask{}))
calls := 0
upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
require.Less(t, calls, 1)
answer := "21"
raw, readErr := io.ReadAll(r.Body)
require.NoError(t, readErr)
assert.Contains(t, gjson.GetBytes(raw, "input.0.content").String(), "圆形 7 9 8")
assert.NotContains(t, gjson.GetBytes(raw, "input.0.content").String(), "21")
calls++
body, encodeErr := common.Marshal(map[string]any{"status": "completed", "usage": map[string]int{"input_tokens": 10, "output_tokens": 2}, "output": []any{map[string]any{"type": "message", "content": []any{map[string]string{"type": "output_text", "text": answer}}}}})
require.NoError(t, encodeErr)
w.Header().Set("Content-Type", "application/json")
_, _ = w.Write(body)
}))
defer upstream.Close()
url := upstream.URL
channel := model.Channel{Type: 1, Key: "test-only", BaseURL: &url, Status: common.ChannelStatusEnabled, Models: fingerprint.Model, Group: "candy-group"}
require.NoError(t, db.Create(&channel).Error)
owner := model.User{Username: "meow-test", Quota: 1000, Status: common.UserStatusEnabled}
require.NoError(t, db.Create(&owner).Error)
price, err := common.Marshal(service.DetectionPrice{Snapshot: billingexpr.BillingSnapshot{ExprString: "p+c", ExprHash: billingexpr.ExprHashString("p+c"), GroupRatio: 1, QuotaPerUnit: 500000, ExprVersion: 1}, Request: billingexpr.RequestInput{Body: []byte(`{"max_output_tokens":4096,"reasoning":{"effort":"low"}}`)}})
require.NoError(t, err)
run := model.DetectionRun{ID: "meow-test", RequestKey: "meow-test", Kind: "candy", CreatedAt: time.Now().Unix(), ChannelID: channel.Id, Status: "pending", Deadline: time.Now().Unix() + 60, Reserved: 320, Price: string(price)}
require.NoError(t, model.CreateDetection(&run))
require.NoError(t, model.EnsureDetectionSchedule(context.Background(), model.DetectionSchedule{ChannelID: channel.Id, Kind: "candy", IntervalMinutes: 30}))
processDetection(context.Background(), run)
require.NoError(t, db.First(&run, "id = ?", run.ID).Error)
assert.Equal(t, "succeeded", run.Status)
assert.Equal(t, "normal", run.Quality)
assert.True(t, run.Settled)
var clock model.DetectionSchedule
require.NoError(t, db.First(&clock, "channel_id = ? AND kind = ?", channel.Id, "candy").Error)
assert.Equal(t, "normal", clock.LastValid)
assert.Greater(t, clock.NextDueAt, run.FinishedAt)
assert.Equal(t, 1, calls)
assert.Equal(t, 0, run.Charged, "automatic tests do not charge a customer")
require.NoError(t, db.First(&owner, owner.Id).Error)
assert.Equal(t, 1000, owner.Quota)
processDetection(context.Background(), run)
assert.Equal(t, 1, calls, "resume must not resend or recharge")
require.NoError(t, db.First(&owner, owner.Id).Error)
assert.Equal(t, 1000, owner.Quota)
summaries, err := service.GroupDetectionSummaries(context.Background(), time.Now(), map[string]*service.AstraSamplingSummary{}, map[string]model.StatusPageGroup{})
require.NoError(t, err)
require.NotNil(t, summaries["candy-group"].Candy)
assert.Equal(t, 1, summaries["candy-group"].Candy.Normal)
assert.Equal(t, "21", summaries["candy-group"].Candy.History[0].Closest)
assert.Zero(t, summaries["candy-group"].Benchmark.Normal)
})
}
}
func TestCandyAnswers(t *testing.T) {
for _, tc := range []struct{ answer, quality string }{
{"21", "normal"}, {" 21\n", "normal"}, {"28", "degraded"}, {"121", "degraded"}, {"21或28", "degraded"}, {"答案不是21", "degraded"}, {"", "degraded"},
} {
t.Run(tc.answer, func(t *testing.T) { assert.Equal(t, tc.quality, service.CandyQuality(tc.answer)) })
}
}
controller/detection_worker.go
+59 -187
View File
@@ -1,22 +1,15 @@
package controller
import (
"bytes"
"context"
"crypto/hmac"
"crypto/sha256"
"encoding/hex"
"errors"
"fmt"
"io"
"net/http"
"strconv"
"strings"
"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"
@@ -24,15 +17,9 @@ import (
type detectionHandler struct{}
func pelicanCredential(id string) string {
mac := hmac.New(sha256.New, []byte(common.SessionSecret))
mac.Write([]byte("newapi-pelican-credential-v1:" + id))
return "det_" + hex.EncodeToString(mac.Sum(nil))
}
func (detectionHandler) Type() string { return "group_detection" }
func (detectionHandler) Enabled() bool {
return common.GetEnvOrDefaultBool("PELICAN_DETECTION_ENABLED", true)
return common.GetEnvOrDefaultBool("MODEL_DETECTION_ENABLED", true)
}
func (detectionHandler) Interval() time.Duration { return 15 * time.Second }
func (detectionHandler) NewPayload() any { return nil }
@@ -41,7 +28,10 @@ func (detectionHandler) Run(ctx context.Context, task *model.SystemTask, runner
_ = 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 := schedulePelicanDetections(ctx)
scheduleErr := scheduleBenchmarkDetections(ctx)
if candyErr := scheduleAutomaticDetections(ctx, "candy"); scheduleErr == nil {
scheduleErr = candyErr
}
var err error
{
var runs []model.DetectionRun
@@ -76,8 +66,12 @@ func (detectionHandler) Run(ctx context.Context, task *model.SystemTask, runner
finishSystemTaskHandler(task, runner, status, nil, err)
}
func schedulePelicanDetections(ctx context.Context) error {
channels, schedules, err := service.AutomaticDetectionChannelsDue(ctx, "pelican", time.Now())
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
}
@@ -87,7 +81,7 @@ func schedulePelicanDetections(ctx context.Context) error {
if len(groups) == 0 {
continue
}
price, reserve, priceErr := service.DetectionPricing("", groups[0], "pelican")
price, reserve, priceErr := service.DetectionPricing("", groups[0], kind)
if priceErr != nil {
return priceErr
}
@@ -98,7 +92,7 @@ func schedulePelicanDetections(ctx context.Context) error {
if err != nil {
return err
}
run := &model.DetectionRun{ID: id, RequestKey: fmt.Sprintf("adaptive:%d:%s", channel.Id, schedules[channel.Id].LastResultID), ChannelID: channel.Id, GroupName: groups[0], Kind: "pelican", Status: "pending", CreatedAt: now, Price: string(encoded), Reserved: reserve}
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
}
@@ -116,8 +110,6 @@ func processDetection(ctx context.Context, run model.DetectionRun) {
return
}
}
taskCtx, taskCancel := context.WithDeadline(ctx, time.Unix(run.Deadline, 0))
defer taskCancel()
if time.Now().Unix() >= run.Deadline {
var calls []model.DetectionCall
if model.DB.Where("run_id = ?", run.ID).Find(&calls).Error != nil {
@@ -137,11 +129,6 @@ func processDetection(ctx context.Context, run model.DetectionRun) {
call.Quota = 0
_ = model.FinishDetectionCall(&call)
}
if run.RemoteID != "" {
cancelCtx, cancel := context.WithTimeout(ctx, 5*time.Second)
_, _ = manxueRunRequest(cancelCtx, &run, http.MethodDelete, "/api/v1/tests/"+run.RemoteID, nil, "")
cancel()
}
finishDetection(run, "failed", "detection_timeout")
return
}
@@ -149,107 +136,16 @@ func processDetection(ctx context.Context, run model.DetectionRun) {
settleDetectionAndSchedule(run)
return
}
if run.Kind == "fingerprint" {
processPaidFingerprint(ctx, run)
if run.Kind != "fingerprint" && run.Kind != "benchmark" && run.Kind != "candy" {
finishDetection(run, "failed", "detection_retired")
return
}
if run.RemoteID == localPelicanID {
processLocalPelican(taskCtx, run)
return
}
if run.RemoteID == "" {
var existing int64
if model.DB.Model(&model.DetectionCall{}).Where("run_id = ?", run.ID).Count(&existing).Error != nil {
return
}
// A domain-separated PRF allows safe retry after a submit/commit crash;
// only its digest is stored, and the random run ID prevents enumeration.
secret := pelicanCredential(run.ID)
call := model.DetectionCall{ID: run.ID, RunID: run.ID, SecretHash: model.DetectionSecretHash(secret), ExpiresAt: run.Deadline + 300, State: "ready"}
if existing == 0 && model.DB.Create(&call).Error != nil {
return
}
payload := map[string]any{"benchmark": "pelican", "base_url": "https://ai.shrinken.pw/api/detection-proxy/v1", "api_key": secret, "model": fingerprint.Model, "protocol": "responses", "reasoning_effort": "high"}
var result gjson.Result
var err error
for taskCtx.Err() == nil {
result, err = manxueRunRequest(taskCtx, &run, http.MethodPost, "/api/v1/tests", payload, run.ID)
if errors.Is(err, errDetectionServiceUnavailable) && taskCtx.Err() == nil {
processLocalPelican(taskCtx, run)
return
}
if err == nil {
break
}
select {
case <-taskCtx.Done():
case <-time.After(5 * time.Second):
}
}
if err != nil || taskCtx.Err() != nil {
if ctx.Err() == nil {
processDetection(ctx, run)
}
return
}
remoteID := result.Get("id").String()
if remoteID == "" || len(remoteID) > 128 || strings.ContainsAny(remoteID, "/?#\\") {
finishDetection(run, "failed", "invalid_detection_response")
return
}
if model.DB.Model(&model.DetectionRun{}).Where("id = ?", run.ID).Updates(map[string]any{"remote_id": remoteID, "status": "running"}).Error != nil {
return
}
run.RemoteID = remoteID
}
// Each worker retains its slot until completion, limiting remote generation
// concurrency as well as local submission concurrency.
for time.Now().Unix() < run.Deadline {
result, err := manxueRunRequest(taskCtx, &run, http.MethodGet, "/api/v1/tests/"+run.RemoteID, nil, "")
if errors.Is(err, errDetectionServiceUnavailable) && taskCtx.Err() == nil {
processLocalPelican(taskCtx, run)
return
}
if err != nil {
select {
case <-ctx.Done():
return
case <-time.After(5 * time.Second):
}
continue
}
switch result.Get("status").String() {
case "succeeded":
quality := result.Get("assessment.quality").String()
if quality != "normal" && quality != "degraded" {
quality = "unknown"
}
html := result.Get("result.html").String()
if len(html) > 2<<20 {
html = ""
}
metadata, _ := common.Marshal(map[string]string{"source_url": manxueArtworkURL(result)})
if model.DB.Model(&model.DetectionRun{}).Where("id = ?", run.ID).Updates(map[string]any{"quality": quality, "html": html, "result": string(metadata)}).Error != nil {
return
}
finishDetection(run, "succeeded", "")
return
case "failed", "cancelled":
finishDetection(run, "failed", "detection_service_failed")
return
}
select {
case <-ctx.Done():
return
case <-time.After(5 * time.Second):
}
}
processDetection(ctx, run)
processPaidFingerprint(ctx, run)
}
func finishDetection(run model.DetectionRun, status, reason string) {
updates := map[string]any{"status": status, "error": reason}
if status == "failed" && run.Kind == "pelican" {
if status == "failed" && run.Kind == "benchmark" {
updates["result"] = ""
}
if model.DB.Model(&model.DetectionRun{}).Where("id = ?", run.ID).Updates(updates).Error != nil {
@@ -259,7 +155,7 @@ func finishDetection(run model.DetectionRun, status, reason string) {
}
func settleDetectionAndSchedule(run model.DetectionRun) {
if model.SettleDetection(run.ID) != nil || run.UserID != 0 || run.Kind != "pelican" {
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.
@@ -270,64 +166,7 @@ func settleDetectionAndSchedule(run model.DetectionRun) {
if run.Status == "succeeded" {
outcome = run.Quality
}
_ = model.ApplyDetectionScheduleResult(context.Background(), run.ChannelID, "pelican", run.ID, run.FinishedAt, outcome)
}
func manxueRequest(ctx context.Context, method, path string, payload any, key string) (gjson.Result, error) {
client := &http.Client{Timeout: 35 * time.Second, CheckRedirect: func(*http.Request, []*http.Request) error { return http.ErrUseLastResponse }}
return manxueRequestWithClient(ctx, client, method, path, payload, key)
}
func manxueRequestWithClient(ctx context.Context, client *http.Client, method, path string, payload any, key string) (gjson.Result, error) {
var encoded []byte
if payload != nil {
var err error
encoded, err = common.Marshal(payload)
if err != nil {
return gjson.Result{}, err
}
}
for attempt := 0; attempt < 3; attempt++ {
req, err := http.NewRequestWithContext(ctx, method, "https://manxue.ai"+path, bytes.NewReader(encoded))
if err != nil {
return gjson.Result{}, err
}
req.Header.Set("Content-Type", "application/json")
req.Header.Set("User-Agent", "Mozilla/5.0")
if strings.HasPrefix(path, "/api/visitor/") {
req.Header.Set("Origin", "https://manxue.ai")
}
if key != "" {
req.Header.Set("Idempotency-Key", key)
}
resp, err := client.Do(req)
if err != nil {
return gjson.Result{}, errDetectionServiceUnavailable
}
body, readErr := io.ReadAll(io.LimitReader(resp.Body, (3<<20)+1))
resp.Body.Close()
if readErr != nil || resp.StatusCode >= 500 {
return gjson.Result{}, errDetectionServiceUnavailable
}
if resp.StatusCode == 429 || resp.StatusCode == 503 {
seconds, _ := strconv.Atoi(resp.Header.Get("Retry-After"))
if at, err := http.ParseTime(resp.Header.Get("Retry-After")); err == nil {
seconds = int(time.Until(at).Seconds()) + 1
}
seconds = max(seconds, 2<<attempt)
select {
case <-ctx.Done():
return gjson.Result{}, ctx.Err()
case <-time.After(time.Duration(seconds) * time.Second):
}
continue
}
if readErr != nil || len(body) > 3<<20 || resp.StatusCode < 200 || resp.StatusCode >= 300 || !gjson.ValidBytes(body) {
return gjson.Result{}, errors.New("invalid detection service response")
}
return gjson.ParseBytes(body), nil
}
return gjson.Result{}, errDetectionServiceUnavailable
_ = model.ApplyDetectionScheduleResult(context.Background(), run.ChannelID, run.Kind, run.ID, run.FinishedAt, outcome)
}
func processPaidFingerprint(ctx context.Context, run model.DetectionRun) {
@@ -336,8 +175,20 @@ func processPaidFingerprint(ctx context.Context, run model.DetectionRun) {
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 < 3; index++ {
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 {
@@ -359,7 +210,7 @@ func processPaidFingerprint(ctx context.Context, run model.DetectionRun) {
}
probeCtx, cancel := context.WithDeadline(ctx, time.Unix(run.Deadline-10, 0))
detectionCalls.Store(call.ID, cancel)
answer, count, probeErr := runAstraFingerprintProbe(probeCtx, channel, true)
answer, count, probeErr := runModelDetectionProbe(probeCtx, channel, index, run.Kind)
interrupted := probeCtx.Err() != nil
cancel()
available := probeErr == nil || answer.HasUsage
@@ -374,7 +225,7 @@ func processPaidFingerprint(ctx context.Context, run model.DetectionRun) {
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/3 {
if quotaErr != nil || quota > run.Reserved/calls {
call.Error = "settlement_unavailable"
} else {
call.Quota = quota
@@ -382,7 +233,7 @@ func processPaidFingerprint(ctx context.Context, run model.DetectionRun) {
} else {
call.Error = "usage_unavailable"
}
if probeErr == nil && fingerprint.ValidArray(answer.Text, count) {
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)
@@ -397,13 +248,34 @@ func processPaidFingerprint(ctx context.Context, run model.DetectionRun) {
return
}
}
result, err := fingerprint.Score(samples)
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": result.Status, "result": string(encoded)}).Error != nil {
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", "")
controller/group.go
+20 -2
View File
@@ -34,6 +34,20 @@ func GetUserGroups(c *gin.Context) {
userGroup := ""
userId := c.GetInt("id")
userGroup, _ = model.GetUserGroup(userId, false)
if err := model.LoadPersonalGroupRatios(c, userId); err != nil {
common.ApiError(c, err)
return
}
metadata, err := model.GetGroupMetadata(c.Request.Context())
if err != nil {
common.ApiError(c, err)
return
}
recent, err := service.RecentDegradedGroups(c.Request.Context())
if err != nil {
common.ApiError(c, err)
return
}
userUsableGroups := service.GetUserUsableGroups(userGroup)
for groupName, _ := range ratio_setting.GetGroupRatioCopy() {
if _, ok := enabledGroups[groupName]; !ok {
@@ -42,8 +56,12 @@ func GetUserGroups(c *gin.Context) {
// UserUsableGroups contains the groups that the user can use
if desc, ok := userUsableGroups[groupName]; ok {
usableGroups[groupName] = map[string]any{
"ratio": service.GetUserGroupRatio(userGroup, groupName),
"desc": desc,
"ratio": model.EffectiveGroupRatio(c, userGroup, groupName).GroupRatio,
"base_ratio": ratio_setting.GetGroupRatio(groupName),
"has_personal_ratio": model.HasPersonalGroupRatio(c, groupName),
"tags": metadata[groupName].Tags,
"recent_degraded": recent[groupName] && !metadata[groupName].HideRecentDegraded,
"desc": desc,
}
}
}
controller/group_configuration.go
+86
View File
@@ -0,0 +1,86 @@
package controller
import (
"net/http"
"strconv"
"github.com/QuantumNous/new-api/common"
"github.com/QuantumNous/new-api/model"
"github.com/gin-gonic/gin"
)
func groupRatioTarget(c *gin.Context) (int, bool) {
id, err := strconv.Atoi(c.Param("id"))
if err != nil || id <= 0 {
c.JSON(http.StatusBadRequest, gin.H{"success": false, "message": "Invalid user"})
return 0, false
}
user, err := model.GetUserById(id, false)
if err != nil {
common.ApiError(c, err)
return 0, false
}
if c.GetInt("role") < common.RoleAdminUser || !canManageTargetRole(c.GetInt("role"), user.Role) {
c.JSON(http.StatusForbidden, gin.H{"success": false, "message": "Permission denied"})
return 0, false
}
return id, true
}
func GetUserGroupRatios(c *gin.Context) {
id, ok := groupRatioTarget(c)
if !ok {
return
}
ratios, err := model.GetPersonalGroupRatios(c.Request.Context(), id)
if err != nil {
common.ApiError(c, err)
return
}
c.JSON(http.StatusOK, gin.H{"success": true, "data": ratios})
}
func UpdateUserGroupRatios(c *gin.Context) {
id, ok := groupRatioTarget(c)
if !ok {
return
}
var input struct {
Ratios map[string]float64 `json:"ratios"`
}
if err := c.ShouldBindJSON(&input); err != nil || input.Ratios == nil {
c.JSON(http.StatusBadRequest, gin.H{"success": false, "message": "Invalid ratios"})
return
}
if err := model.SavePersonalGroupRatios(c.Request.Context(), id, input.Ratios); err != nil {
common.ApiError(c, err)
return
}
recordManageAuditFor(c, id, "user.group_ratios.update", map[string]any{"ratios": input.Ratios})
c.JSON(http.StatusOK, gin.H{"success": true})
}
func GetGroupMetadata(c *gin.Context) {
groups, err := model.GetGroupMetadata(c.Request.Context())
if err != nil {
common.ApiError(c, err)
return
}
c.JSON(http.StatusOK, gin.H{"success": true, "data": groups})
}
func UpdateGroupMetadata(c *gin.Context) {
var input struct {
Groups map[string]model.GroupMetadataValue `json:"groups"`
}
if err := c.ShouldBindJSON(&input); err != nil || input.Groups == nil {
c.JSON(http.StatusBadRequest, gin.H{"success": false, "message": "Invalid group metadata"})
return
}
if err := model.SaveGroupMetadata(c.Request.Context(), input.Groups); err != nil {
common.ApiError(c, err)
return
}
recordManageAuditFor(c, 0, "group.metadata.update", map[string]any{"groups": input.Groups})
c.JSON(http.StatusOK, gin.H{"success": true})
}
controller/group_configuration_test.go
+65
View File
@@ -0,0 +1,65 @@
package controller
import (
"bytes"
"context"
"encoding/json"
"fmt"
"net/http/httptest"
"testing"
"github.com/QuantumNous/new-api/common"
"github.com/QuantumNous/new-api/model"
"github.com/gin-gonic/gin"
"github.com/glebarez/sqlite"
"github.com/stretchr/testify/require"
"gorm.io/gorm"
)
func TestPersonalRatioAdminAuthorization(t *testing.T) {
db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{})
require.NoError(t, err)
require.NoError(t, db.AutoMigrate(&model.User{}, &model.PersonalGroupRatio{}, &model.AuditLog{}))
oldLog := model.LOG_DB
model.LOG_DB = db
old, oldRedis := model.DB, common.RedisEnabled
model.DB, common.RedisEnabled = db, false
t.Cleanup(func() {
model.LOG_DB = oldLog
model.DB, common.RedisEnabled = old, oldRedis
conn, _ := db.DB()
_ = conn.Close()
})
user := model.User{Id: 12, Username: "personal-test", Role: common.RoleCommonUser}
require.NoError(t, db.Create(&user).Error)
for _, role := range []int{common.RoleCommonUser, common.RoleAdminUser, common.RoleRootUser} {
w := httptest.NewRecorder()
c, _ := gin.CreateTestContext(w)
c.Set("role", role)
c.Params = gin.Params{{Key: "id", Value: "12"}}
c.Request = httptest.NewRequest("PUT", "/api/user/12/group-ratios", bytes.NewBufferString("{\"ratios\":{\"default\":0}}"))
c.Request.Header.Set("Content-Type", "application/json")
UpdateUserGroupRatios(c)
if role == common.RoleCommonUser {
require.Equal(t, 403, w.Code)
} else {
var body struct{ Success bool }
require.NoError(t, json.Unmarshal(w.Body.Bytes(), &body))
require.True(t, body.Success, w.Body.String())
}
}
var auditCount int64
require.NoError(t, db.Model(&model.AuditLog{}).Where("action = ?", "user.group_ratios.update").Count(&auditCount).Error)
require.EqualValues(t, 2, auditCount)
value, err := model.GetPersonalGroupRatios(context.Background(), 12)
require.NoError(t, err)
require.Equal(t, 0.0, value["default"])
require.NoError(t, db.Model(&user).Update("role", common.RoleAdminUser).Error)
w := httptest.NewRecorder()
c, _ := gin.CreateTestContext(w)
c.Set("role", common.RoleAdminUser)
c.Params = gin.Params{{Key: "id", Value: fmt.Sprint(user.Id)}}
c.Request = httptest.NewRequest("GET", "/api/user/12/group-ratios", nil)
GetUserGroupRatios(c)
require.Equal(t, 403, w.Code)
}
controller/live_status_test.go
+2 -2
View File
@@ -32,13 +32,13 @@ func TestCurrentGroupStatus(t *testing.T) {
{"future evidence is ignored", now + 1, []service.ProbeHistoryPoint{{At: now + 1, Status: 1}}, false, -1, "detection", 0},
} {
t.Run(tc.name, func(t *testing.T) {
m := Monitor{Status: 0, StatusSource: "requests", StatusCheckedAt: now - 100, ChannelDisabled: tc.disabled, ProbeHistory: tc.probes, Uptime: 0.58, Fingerprint: &service.AstraSamplingSummary{Errors: 25}, Pelican: &service.PelicanSummary{Failed: 33}}
m := Monitor{Status: 0, StatusSource: "requests", StatusCheckedAt: now - 100, ChannelDisabled: tc.disabled, ProbeHistory: tc.probes, Uptime: 0.58, Fingerprint: &service.AstraSamplingSummary{Errors: 25}, Benchmark: &service.BenchmarkSummary{Failed: 33}}
applyCurrentStatus(&m, perfmetrics.UserRequestResult{At: tc.success * 1000, Success: true}, now)
assert.Equal(t, tc.want, m.Status)
assert.Equal(t, tc.source, m.StatusSource)
assert.Equal(t, tc.at, m.StatusCheckedAt)
assert.Equal(t, 25, m.Fingerprint.Errors)
assert.Equal(t, 33, m.Pelican.Failed)
assert.Equal(t, 33, m.Benchmark.Failed)
assert.Equal(t, 0.58, m.Uptime)
})
}
controller/manxue_artwork_url_test.go
-19
View File
@@ -1,19 +0,0 @@
package controller
import (
"github.com/stretchr/testify/assert"
"github.com/tidwall/gjson"
"testing"
)
func TestManxueArtworkURL(t *testing.T) {
assert.Empty(t, manxueArtworkURL(gjson.Parse(`{"public":false,"result":{"id":"private"}}`)))
assert.Equal(t, "https://manxue.ai/?result=community-one", manxueArtworkURL(gjson.Parse(`{"public":true,"result":{"id":"community-one"}}`)))
assert.Empty(t, manxueArtworkURL(gjson.Parse(`{"share_url":"https://evil.test/?result=one"}`)))
}
func TestArtworkPreservesScriptDrawnLegs(t *testing.T) {
result := sanitizeDetectionHTML(`<svg><path id="leg"></path></svg><button id="toggle">pause</button><script>document.getElementById('leg').setAttribute('d','M0 0L10 10')</script>`)
assert.Contains(t, result, "<script>")
assert.Contains(t, result, "setAttribute")
assert.Contains(t, result, "<button")
}
controller/manxue_public.go
-64
View File
@@ -1,64 +0,0 @@
package controller
import (
"context"
"errors"
"net/http"
"net/http/cookiejar"
"net/url"
"strings"
"time"
"github.com/QuantumNous/new-api/common"
"github.com/QuantumNous/new-api/model"
"github.com/tidwall/gjson"
)
// Visitor sessions belong to a single detection. Keep them in its private task
// metadata until completion so polling can resume after a process restart.
type manxueVisitorSession struct {
Cookies []*http.Cookie `json:"manxue_cookies,omitempty"`
}
func manxueRunRequest(ctx context.Context, run *model.DetectionRun, method, path string, payload any, key string) (gjson.Result, error) {
var session manxueVisitorSession
_ = common.UnmarshalJsonStr(run.Result, &session)
if len(session.Cookies) == 0 && !(method == http.MethodPost && run.RemoteID == "") {
return manxueRequest(ctx, method, path, payload, key)
}
jar, _ := cookiejar.New(nil)
origin, _ := url.Parse("https://manxue.ai/api/visitor/")
client := &http.Client{Timeout: 35 * time.Second, Jar: jar, CheckRedirect: func(*http.Request, []*http.Request) error { return http.ErrUseLastResponse }}
if len(session.Cookies) == 0 {
if _, err := manxueRequestWithClient(ctx, client, http.MethodGet, "/api/visitor/config", nil, ""); err != nil {
return gjson.Result{}, errDetectionServiceUnavailable
}
session.Cookies = jar.Cookies(origin)
if len(session.Cookies) == 0 {
return gjson.Result{}, errDetectionServiceUnavailable
}
encoded, err := common.Marshal(session)
if err != nil {
return gjson.Result{}, err
}
if err = model.DB.Model(&model.DetectionRun{}).Where("id = ?", run.ID).Update("result", string(encoded)).Error; err != nil {
return gjson.Result{}, err
}
run.Result = string(encoded)
} else {
jar.SetCookies(origin, session.Cookies)
}
if method == http.MethodPost {
fields, ok := payload.(map[string]any)
if !ok {
return gjson.Result{}, errors.New("invalid pelican payload")
}
public := make(map[string]any, len(fields)+1)
for k, v := range fields {
public[k] = v
}
public["public"] = true
payload = public
}
return manxueRequestWithClient(ctx, client, method, strings.Replace(path, "/api/v1/tests", "/api/visitor/tests", 1), payload, key)
}
controller/manxue_public_test.go
-80
View File
@@ -1,80 +0,0 @@
package controller
import (
"context"
"encoding/json"
"net/http"
"net/http/httptest"
"testing"
"github.com/QuantumNous/new-api/model"
"github.com/glebarez/sqlite"
"github.com/stretchr/testify/require"
"github.com/tidwall/gjson"
"gorm.io/gorm"
)
type manxueTestTransport func(*http.Request) (*http.Response, error)
func (f manxueTestTransport) RoundTrip(r *http.Request) (*http.Response, error) { return f(r) }
func TestManxuePublicDetectionResumesAndReturnsPublicURL(t *testing.T) {
db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{})
require.NoError(t, err)
conn, err := db.DB()
require.NoError(t, err)
conn.SetMaxOpenConns(1)
oldDB, oldTransport := model.DB, http.DefaultTransport
model.DB = db
t.Cleanup(func() { model.DB = oldDB; http.DefaultTransport = oldTransport; _ = conn.Close() })
require.NoError(t, db.AutoMigrate(&model.DetectionRun{}))
configCalls, createCalls := 0, 0
http.DefaultTransport = manxueTestTransport(func(r *http.Request) (*http.Response, error) {
w := httptest.NewRecorder()
w.Header().Set("Content-Type", "application/json")
switch r.URL.Path {
case "/api/visitor/config":
configCalls++
http.SetCookie(w, &http.Cookie{Name: "gallery_visitor", Value: "test-session", Path: "/api/visitor/", Secure: true, HttpOnly: true})
_, _ = w.WriteString(`{"prompt":"pelican"}`)
case "/api/visitor/tests":
createCalls++
cookie, err := r.Cookie("gallery_visitor")
require.NoError(t, err)
require.Equal(t, "test-session", cookie.Value)
require.Equal(t, "stable-key", r.Header.Get("Idempotency-Key"))
var body map[string]any
require.NoError(t, json.NewDecoder(r.Body).Decode(&body))
require.Equal(t, true, body["public"])
_, _ = w.WriteString(`{"id":"remote-one","status":"running"}`)
case "/api/visitor/tests/remote-one":
_, err := r.Cookie("gallery_visitor")
require.NoError(t, err)
_, _ = w.WriteString(`{"status":"succeeded","public_id":"community-one","result":{"html":"<html>pelican</html>"}}`)
case "/api/v1/tests/legacy":
require.Empty(t, r.Header.Get("Cookie"))
_, _ = w.WriteString(`{"status":"succeeded"}`)
default:
t.Fatalf("unexpected path: %s", r.URL.Path)
}
return w.Result(), nil
})
run := model.DetectionRun{ID: "public-test", RequestKey: "public-test", Kind: "pelican"}
require.NoError(t, db.Create(&run).Error)
input := map[string]any{"benchmark": "pelican"}
_, err = manxueRunRequest(context.Background(), &run, http.MethodPost, "/api/v1/tests", input, "stable-key")
require.NoError(t, err)
require.NotContains(t, input, "public")
var resumed model.DetectionRun
require.NoError(t, db.First(&resumed, "id = ?", run.ID).Error)
resumed.RemoteID = "remote-one"
result, err := manxueRunRequest(context.Background(), &resumed, http.MethodGet, "/api/v1/tests/remote-one", nil, "")
require.NoError(t, err)
require.Equal(t, "https://manxue.ai/?result=community-one", manxueArtworkURL(result))
require.Equal(t, 1, configCalls)
require.Equal(t, 1, createCalls)
legacy := model.DetectionRun{RemoteID: "legacy"}
_, err = manxueRunRequest(context.Background(), &legacy, http.MethodGet, "/api/v1/tests/legacy", nil, "")
require.NoError(t, err)
require.Empty(t, manxueArtworkURL(gjson.Parse(`{"id":"private","public":false}`)))
}
controller/model.go
+6 -1
View File
@@ -260,7 +260,12 @@ func ListModels(c *gin.Context, modelType int) {
tokenModelLimit = map[string]bool{}
}
}
models := service.GetGroupsEnabledModels(ownerGroups)
var models []string
if policy := service.RequestRoutingPolicy(c); policy != nil && policy.Mode == "smart" {
models = service.AutomaticPoolModels(c)
} else {
models = service.GetGroupsEnabledModels(ownerGroups)
}
for _, modelName := range models {
if modelLimitEnable {
matchingName := ratio_setting.RoutingMatchModelName(modelName)
controller/pelican_fallback.go
-96
View File
@@ -1,96 +0,0 @@
package controller
import (
"bytes"
"context"
"errors"
"github.com/QuantumNous/new-api/common"
"github.com/QuantumNous/new-api/model"
"github.com/gin-gonic/gin"
"github.com/tidwall/gjson"
"net/http"
"net/http/httptest"
"strings"
)
var errDetectionServiceUnavailable = errors.New("detection service unreachable")
const localPelicanID = "local-pelican"
const pelicanPrompt = "Create a complete self-contained HTML page showing a pelican riding a bicycle. Draw the pelican and bicycle with inline SVG and animate pedaling and rotating wheels using CSS. Make the scene visually polished. Do not load external resources or use JavaScript. Return only the complete HTML document."
func detectionHTML(raw []byte, contentType string) string {
result := gjson.ParseBytes(raw)
if strings.Contains(contentType, "text/event-stream") {
result = gjson.Result{}
for line := range strings.SplitSeq(string(raw), "\n") {
if strings.HasPrefix(line, "data:") {
event := gjson.Parse(strings.TrimSpace(strings.TrimPrefix(line, "data:")))
if event.Get("type").String() == "response.completed" {
result = event.Get("response")
}
}
}
}
if result.Get("status").String() != "completed" {
return ""
}
var output strings.Builder
for _, item := range result.Get("output").Array() {
if item.Get("type").String() != "message" {
continue
}
for _, part := range item.Get("content").Array() {
if part.Get("type").String() == "output_text" {
output.WriteString(part.Get("text").String())
}
}
}
html := strings.TrimSpace(output.String())
lower := strings.ToLower(html)
start := strings.Index(lower, "<!doctype html")
if start < 0 {
start = strings.Index(lower, "<html")
}
end := strings.LastIndex(lower, "</html>")
if start < 0 || end < start || len(html) > 2<<20 {
return ""
}
return html[start : end+7]
}
func processLocalPelican(ctx context.Context, run model.DetectionRun) {
if ctx.Err() != nil {
return
}
if model.DB.Model(&model.DetectionRun{}).Where("id = ?", run.ID).Updates(map[string]any{"remote_id": localPelicanID, "quality": "unknown", "result": ""}).Error != nil {
return
}
var call model.DetectionCall
if model.DB.First(&call, "id = ?", run.ID).Error != nil {
return
}
if call.State == "running" {
return
}
if call.State == "ready" {
payload, _ := common.Marshal(map[string]any{"model": "gpt-6-astra", "input": pelicanPrompt})
c, _ := gin.CreateTestContext(httptest.NewRecorder())
c.Request = httptest.NewRequestWithContext(ctx, http.MethodPost, "/api/detection-proxy/v1/responses", bytes.NewReader(payload))
c.Request.Header.Set("Authorization", "Bearer "+pelicanCredential(run.ID))
DetectionProxy(c)
if model.DB.First(&call, "id = ?", run.ID).Error != nil {
return
}
}
if call.State != "finished" {
return
}
if model.DB.First(&run, "id = ?", run.ID).Error != nil {
return
}
if run.HTML != "" {
finishDetection(run, "succeeded", "")
} else {
finishDetection(run, "failed", "local_generation_failed")
}
}
controller/pelican_fallback_test.go
-14
View File
@@ -1,14 +0,0 @@
package controller
import (
"github.com/stretchr/testify/assert"
"testing"
)
func TestDetectionHTMLRequiresCompletedDocument(t *testing.T) {
raw := []byte(`{"status":"completed","output":[{"type":"message","content":[{"type":"output_text","text":"<html><body>pelican</body></html>"}]}]}`)
assert.Equal(t, "<html><body>pelican</body></html>", detectionHTML(raw, "application/json"))
assert.Empty(t, detectionHTML([]byte(`{"status":"incomplete","output":[]}`), "application/json"))
sse := []byte("data: {\"type\":\"response.completed\",\"response\":" + string(raw) + "}\n\n")
assert.Equal(t, detectionHTML(raw, "application/json"), detectionHTML(sse, "text/event-stream"))
}
controller/pelican_local_integration_test.go
-50
View File
@@ -1,50 +0,0 @@
package controller
import (
"context"
"github.com/QuantumNous/new-api/common"
"github.com/QuantumNous/new-api/model"
"github.com/QuantumNous/new-api/pkg/billingexpr"
"github.com/QuantumNous/new-api/service"
"github.com/glebarez/sqlite"
"github.com/stretchr/testify/require"
"gorm.io/gorm"
"net/http"
"net/http/httptest"
"testing"
"time"
)
func TestLocalPelicanGeneratesOnceWithUnknownVerdict(t *testing.T) {
db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{})
require.NoError(t, err)
conn, err := db.DB()
require.NoError(t, err)
conn.SetMaxOpenConns(1)
old := model.DB
model.DB = db
t.Cleanup(func() { model.DB = old; _ = conn.Close() })
require.NoError(t, db.AutoMigrate(&model.Channel{}, &model.DetectionRun{}, &model.DetectionCall{}, &model.DetectionSchedule{}))
calls := 0
upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
calls++
w.Header().Set("Content-Type", "application/json")
_, _ = w.Write([]byte(`{"status":"completed","usage":{"input_tokens":10,"output_tokens":5},"output":[{"type":"message","content":[{"type":"output_text","text":"<html><body>pelican</body></html>"}]}]}`))
}))
defer upstream.Close()
url := upstream.URL
channel := model.Channel{Type: 1, Key: "test-only", BaseURL: &url, Status: common.ChannelStatusEnabled, Models: "gpt-6-astra"}
require.NoError(t, db.Create(&channel).Error)
price, _ := common.Marshal(service.DetectionPrice{Snapshot: billingexpr.BillingSnapshot{ExprString: "p+c", ExprHash: billingexpr.ExprHashString("p+c"), GroupRatio: 1, QuotaPerUnit: 500000, ExprVersion: 1}})
run := model.DetectionRun{ID: "local-test", RequestKey: "local-test", Kind: "pelican", ChannelID: channel.Id, Status: "running", Deadline: time.Now().Unix() + 60, Reserved: 100, Price: string(price)}
require.NoError(t, db.Create(&run).Error)
require.NoError(t, db.Create(&model.DetectionCall{ID: run.ID, RunID: run.ID, SecretHash: model.DetectionSecretHash(pelicanCredential(run.ID)), State: "ready", ExpiresAt: run.Deadline}).Error)
processLocalPelican(context.Background(), run)
require.NoError(t, db.First(&run, "id = ?", run.ID).Error)
require.Equal(t, "succeeded", run.Status)
require.Equal(t, "unknown", run.Quality)
require.Contains(t, run.HTML, "pelican")
require.True(t, run.Settled)
processLocalPelican(context.Background(), run)
require.Equal(t, 1, calls)
}
controller/pricing.go
+5 -4
View File
@@ -52,11 +52,12 @@ func GetPricing(c *gin.Context) {
user, err := model.GetUserCache(userId.(int))
if err == nil {
group = user.Group
if err := model.LoadPersonalGroupRatios(c, user.Id); err != nil {
common.ApiError(c, err)
return
}
for g := range groupRatio {
ratio, ok := ratio_setting.GetGroupGroupRatio(group, g)
if ok {
groupRatio[g] = ratio
}
groupRatio[g] = model.EffectiveGroupRatio(c, group, g).GroupRatio
}
}
}
controller/relay.go
+17 -2
View File
@@ -194,9 +194,12 @@ func Relay(c *gin.Context, relayFormat types.RelayFormat) {
relayInfo.RetryIndex = 0
relayInfo.LastError = nil
maxRetries := common.RetryTimes
poolBudget := service.PoolRetryBudget(c, relayInfo.OriginModelName)
poolFailures := 0
maxRetries = max(maxRetries, poolBudget)
ordinaryRetries := 0
if relayInfo.DurableAttemptBilling {
maxRetries = 2
maxRetries = max(maxRetries, 2)
timeout := constant.StreamingTimeout
if timeout <= 0 {
timeout = 300
@@ -210,7 +213,7 @@ func Relay(c *gin.Context, relayFormat types.RelayFormat) {
newAPIError = types.NewErrorWithStatusCode(errors.New("response model confirmation budget exhausted"), relaycommon.ResponseGuardErrorCode, http.StatusGatewayTimeout, types.ErrOptionWithSkipRetry())
break
}
if relayInfo.ResponseGuard != nil && relayInfo.UpstreamRequests >= relaycommon.ResponseGuardMaxAttempts {
if relayInfo.ResponseGuard != nil && relayInfo.UpstreamRequests-poolFailures >= relaycommon.ResponseGuardMaxAttempts {
newAPIError = types.NewErrorWithStatusCode(errors.New("response model attempt budget exhausted"), relaycommon.ResponseGuardErrorCode, http.StatusBadGateway, types.ErrOptionWithSkipRetry())
break
}
@@ -259,6 +262,9 @@ func Relay(c *gin.Context, relayFormat types.RelayFormat) {
if newAPIError == nil {
service.MarkRequestPolicySuccess(c, relayInfo.StreamStatus)
if service.RequestPolicy(c).Successful {
service.RecordPoolSuccess(c, relayInfo.OriginModelName, currentRequestGroup(c, channel))
}
relayInfo.LastError = nil
service.RecordGroupOutcome(currentRequestGroup(c, channel), true)
return
@@ -276,7 +282,16 @@ func Relay(c *gin.Context, relayFormat types.RelayFormat) {
if relayInfo.DurableAttemptBilling {
remaining = common.RetryTimes - ordinaryRetries
}
poolFailure := service.IsPoolFailoverError(c, newAPIError)
if poolFailure {
remaining = maxRetries - retryParam.GetRetry()
poolFailures++
service.RecordPoolRetryableFailure(c, relayInfo.OriginModelName, currentRequestGroup(c, channel))
}
decision := service.DecideRelayRetry(c, newAPIError, remaining)
if poolFailure && service.PoolRetryExhausted(c, relayInfo.OriginModelName, currentRequestGroup(c, channel)) {
decision = service.PolicyDecision{Action: "stop", Reason: "pool_no_backup", Source: "routing"}
}
if canSwitchRejectedAstra(c, relayInfo, newAPIError) {
decision = service.PolicyDecision{Action: "retry", Reason: "model_mismatch_switch", Source: "model_guard"}
service.ClearPoolAffinity(c)
controller/responses_websocket_test.go
+81 -1
View File
@@ -25,8 +25,10 @@ import (
"github.com/QuantumNous/new-api/relay"
"github.com/QuantumNous/new-api/relaykit/dto"
"github.com/QuantumNous/new-api/relaykit/types"
"github.com/QuantumNous/new-api/service"
"github.com/QuantumNous/new-api/setting"
"github.com/QuantumNous/new-api/setting/config"
"github.com/QuantumNous/new-api/setting/ratio_setting"
"github.com/alicebob/miniredis/v2"
"github.com/gin-gonic/gin"
"github.com/go-redis/redis/v8"
@@ -38,6 +40,84 @@ import (
var responsesWSTestUserSequence atomic.Int64
func TestResponsesHTTPPoolRetriesTwiceThenUsesBackup(t *testing.T) {
fixture := newResponsesWSBillingTest(t, `tier("request", fixed(0.002))`, func(*websocket.Conn, *http.Request) {})
oldRetries := common.RetryTimes
common.RetryTimes = 0
t.Cleanup(func() { common.RetryTimes = oldRetries })
oldGroups := setting.UserUsableGroups2JSONString()
require.NoError(t, setting.UpdateUserUsableGroupsByJSONString(`{"default":"Default","vip":"VIP"}`))
t.Cleanup(func() { require.NoError(t, setting.UpdateUserUsableGroupsByJSONString(oldGroups)) })
require.NoError(t, ratio_setting.UpdateGroupRatioByJSONString(`{"default":1,"vip":1}`))
if os.Getenv("TEST_PERSONAL_RATIO") == "1" {
require.NoError(t, model.SavePersonalGroupRatios(context.Background(), fixture.user.Id, map[string]float64{"default": 0.8, "vip": 0.2}))
}
service.ResetGroupAvailability()
t.Cleanup(service.ResetGroupAvailability)
var primaryCalls, backupCalls atomic.Int64
fixture.httpUpstream = func(w http.ResponseWriter, r *http.Request) {
primaryCalls.Add(1)
w.Header().Set("Content-Type", "application/json")
w.WriteHeader(500)
_, _ = fmt.Fprint(w, `{"error":{"type":"server_error","code":"server_error","message":"test failure"}}`)
}
backup := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
backupCalls.Add(1)
w.Header().Set("Content-Type", "application/json")
_, _ = fmt.Fprint(w, `{"id":"completed","status":"completed","usage":{"input_tokens":1000,"output_tokens":1,"total_tokens":1001}}`)
}))
t.Cleanup(backup.Close)
require.NoError(t, model.DB.Model(fixture.channel).Updates(map[string]any{"key": "single-primary-key", "channel_info": `{}`, "auto_ban": 0}).Error)
backupChannel := &model.Channel{Name: "backup-pool-test", Key: "test-only", Status: common.ChannelStatusEnabled, Type: constant.ChannelTypeOpenAI, Group: "vip", Models: "ws-billing", BaseURL: &backup.URL, AutoBan: common.GetPointer(0)}
require.NoError(t, model.DB.Create(backupChannel).Error)
require.NoError(t, model.DB.Create(&model.Ability{ChannelId: backupChannel.Id, Model: "ws-billing", Group: "vip", Enabled: true}).Error)
t.Cleanup(func() {
require.NoError(t, model.DB.Where("channel_id = ?", backupChannel.Id).Delete(&model.Ability{}).Error)
require.NoError(t, model.DB.Delete(backupChannel).Error)
})
require.NoError(t, model.DB.First(fixture.token, fixture.token.Id).Error)
body, err := common.Marshal(model.TokenRoutingPolicy{Version: 1, Mode: "custom", Groups: []model.TokenGroupEntry{{Group: "default", Priority: 0}, {Group: "vip", Priority: 1}}})
require.NoError(t, err)
fixture.token.RoutingPolicy = string(body)
fixture.token.Group = "default"
fixture.token.CrossGroupRetry = true
require.NoError(t, model.DB.Model(fixture.token).Updates(map[string]any{"routing_policy": string(body), "group": "default", "cross_group_retry": true}).Error)
require.NoError(t, fixture.token.SelectUpdate())
for range 2 {
req, err := http.NewRequest(http.MethodPost, fixture.gatewayURL+"/v1/responses", strings.NewReader(`{"model":"ws-billing","input":"hello"}`))
require.NoError(t, err)
req.Header.Set("Authorization", "Bearer sk-"+fixture.token.Key)
req.Header.Set("Content-Type", "application/json")
response, err := http.DefaultClient.Do(req)
require.NoError(t, err)
data, err := io.ReadAll(response.Body)
require.NoError(t, err)
require.NoError(t, response.Body.Close())
require.Equal(t, 200, response.StatusCode, string(data))
select {
case <-fixture.httpDone:
case <-time.After(3 * time.Second):
t.Fatal("HTTP request did not finish")
}
}
require.Equal(t, int64(2), primaryCalls.Load())
require.Equal(t, int64(2), backupCalls.Load())
var logs []model.Log
require.NoError(t, model.LOG_DB.Where("token_id = ? AND type = ?", fixture.token.Id, model.LogTypeConsume).Find(&logs).Error)
require.Len(t, logs, 2)
for _, log := range logs {
require.Equal(t, "vip", log.Group)
expected := 1000
if os.Getenv("TEST_PERSONAL_RATIO") == "1" {
expected = 200
}
require.Equal(t, expected, log.Quota)
}
if os.Getenv("TEST_PERSONAL_RATIO") == "1" {
assertResponsesWSAccounting(t, fixture, []int{200, 200})
}
}
func setupResponsesWSRequestTest(t *testing.T) (*model.User, *model.Token) {
t.Helper()
previousDB := model.DB
@@ -88,7 +168,7 @@ func setupResponsesWSRequestTest(t *testing.T) (*model.User, *model.Token) {
setting.ModelRequestRateLimitMutex.Unlock()
require.NoError(t, sqlDB.Close())
})
require.NoError(t, db.AutoMigrate(&model.User{}, &model.Token{}))
require.NoError(t, db.AutoMigrate(&model.User{}, &model.Token{}, &model.PersonalGroupRatio{}))
// The shared in-memory limiter outlives each database fixture. Give every
// user a separate quota bucket, including when the tests run with -count.
user := &model.User{Id: 5062000 + int(responsesWSTestUserSequence.Add(1)), Username: "responses-ws-user", Status: common.UserStatusEnabled, Group: "default", Quota: 1000, AuthVersion: 1}
controller/status_page.go
+24
View File
@@ -7,6 +7,7 @@ import (
"github.com/QuantumNous/new-api/common"
"github.com/QuantumNous/new-api/model"
"github.com/QuantumNous/new-api/service"
"github.com/gin-gonic/gin"
)
@@ -60,6 +61,7 @@ func UpdateStatusPageGroup(c *gin.Context) {
common.ApiError(c, err)
return
}
service.InvalidateGroupQualityCache()
c.JSON(http.StatusOK, gin.H{"success": true})
}
@@ -70,12 +72,34 @@ func GetSiteGroupStatus(c *gin.Context) {
common.ApiError(c, err)
return
}
costs, costErr := service.GroupModelCosts(c.Request.Context())
if costErr != nil {
common.SysError("group model costs: " + costErr.Error())
}
metadata, err := model.GetGroupMetadata(c.Request.Context())
if err != nil {
common.ApiError(c, err)
return
}
recent, err := service.RecentDegradedGroups(c.Request.Context())
if err != nil {
common.ApiError(c, err)
return
}
for i := range results {
monitors := make([]Monitor, 0, len(results[i].Monitors))
for _, monitor := range results[i].Monitors {
monitor.Tags = metadata[monitor.Group].Tags
monitor.RecentDegraded = recent[monitor.Group] && !metadata[monitor.Group].HideRecentDegraded
monitor.Name = monitor.Group
monitor.Alias = ""
monitor.Hidden = false
monitor.ModelCosts = map[string]model.GroupModelCost{}
for _, name := range monitor.Models {
if value, ok := costs[monitor.Group][name]; ok {
monitor.ModelCosts[name] = value
}
}
monitors = append(monitors, monitor)
}
results[i].Monitors = monitors
controller/token_routing.go
+5 -6
View File
@@ -2,6 +2,7 @@ package controller
import (
"github.com/QuantumNous/new-api/common"
"github.com/QuantumNous/new-api/constant"
"github.com/QuantumNous/new-api/model"
"github.com/QuantumNous/new-api/service"
"github.com/QuantumNous/new-api/setting"
@@ -29,7 +30,7 @@ func setTokenRoutingPolicy(c *gin.Context, token *model.Token, policy *model.Tok
token.RoutingPolicy = string(body)
token.CrossGroupRetry = true
token.AutoGroups = ""
if policy.Mode == "inherit" {
if policy.Mode == "inherit" || policy.Mode == "smart" {
token.Group = "auto"
} else if policy.Mode == "user" {
token.Group = ""
@@ -57,10 +58,8 @@ func PreviewTokenRouting(c *gin.Context) {
common.ApiError(c, err)
return
}
inherited := service.GetUserAutoGroup(group)
if request.Policy.Mode == "user" {
inherited = []string{group}
}
entries := request.Policy.Resolve(request.Model, inherited, func(candidate string) bool { return service.IsUserSelectableGroup(group, candidate) })
common.SetContextKey(c, constant.ContextKeyUserGroup, group)
c.Set(model.TokenRoutingContextKey, &request.Policy)
entries := service.ResolveRequestPool(c, request.Model)
common.ApiSuccess(c, entries)
}
controller/uptime_kuma.go
+29 -22
View File
@@ -11,6 +11,7 @@ import (
"github.com/QuantumNous/new-api/model"
perfmetrics "github.com/QuantumNous/new-api/pkg/perf_metrics"
"github.com/QuantumNous/new-api/service"
"github.com/QuantumNous/new-api/setting/ratio_setting"
"github.com/gin-gonic/gin"
)
@@ -18,26 +19,31 @@ import (
const requestTimeout = 30 * time.Second
type Monitor struct {
ProbeHistory []service.ProbeHistoryPoint `json:"probe_history"`
ID string `json:"id"`
Models []string `json:"models"`
Alias string `json:"alias,omitempty"`
Hidden bool `json:"hidden,omitempty"`
Pelican *service.PelicanSummary `json:"pelican,omitempty"`
StatusSource string `json:"status_source"`
StatusCheckedAt int64 `json:"status_checked_at,omitempty"`
Fingerprint *service.AstraSamplingSummary `json:"fingerprint,omitempty"`
Metrics *perfmetrics.GroupResult `json:"metrics,omitempty"`
Name string `json:"name"`
Uptime float64 `json:"uptime"`
Status int `json:"status"`
Group string `json:"group,omitempty"`
Ping int `json:"ping,omitempty"`
LastChecked string `json:"last_checked,omitempty"`
History []MonitorSample `json:"history,omitempty"`
CacheStats *model.GroupCacheStats `json:"cache_stats,omitempty"`
ProbeStatus int `json:"probe_status"`
ChannelDisabled bool `json:"channel_disabled"`
Tags []string `json:"tags,omitempty"`
RecentDegraded bool `json:"recent_degraded,omitempty"`
ModelCosts map[string]model.GroupModelCost `json:"model_costs,omitempty"`
GroupRatio float64 `json:"group_ratio"`
ProbeHistory []service.ProbeHistoryPoint `json:"probe_history"`
ID string `json:"id"`
Models []string `json:"models"`
Alias string `json:"alias,omitempty"`
Hidden bool `json:"hidden,omitempty"`
Candy *service.BenchmarkSummary `json:"candy,omitempty"`
Benchmark *service.BenchmarkSummary `json:"benchmark,omitempty"`
StatusSource string `json:"status_source"`
StatusCheckedAt int64 `json:"status_checked_at,omitempty"`
Fingerprint *service.AstraSamplingSummary `json:"fingerprint,omitempty"`
Metrics *perfmetrics.GroupResult `json:"metrics,omitempty"`
Name string `json:"name"`
Uptime float64 `json:"uptime"`
Status int `json:"status"`
Group string `json:"group,omitempty"`
Ping int `json:"ping,omitempty"`
LastChecked string `json:"last_checked,omitempty"`
History []MonitorSample `json:"history,omitempty"`
CacheStats *model.GroupCacheStats `json:"cache_stats,omitempty"`
ProbeStatus int `json:"probe_status"`
ChannelDisabled bool `json:"channel_disabled"`
}
type MonitorSample struct {
@@ -129,7 +135,7 @@ func currentGroupMonitors(ctx context.Context, admin bool) ([]UptimeGroupResult,
userResults := perfmetrics.LatestUserResults(ctx, order)
monitors := make([]Monitor, 0, len(order))
for _, group := range order {
monitor := Monitor{ID: model.DetectionSecretHash(group)[:16], Name: group, Status: -1, Models: models[group]}
monitor := Monitor{ID: model.DetectionSecretHash(group)[:16], Name: group, Status: -1, Models: models[group], GroupRatio: ratio_setting.GetGroupRatio(group)}
if monitor.Models == nil {
monitor.Models = []string{}
}
@@ -157,7 +163,8 @@ func currentGroupMonitors(ctx context.Context, admin bool) ([]UptimeGroupResult,
if !admin && monitor.Metrics != nil {
monitor.Metrics.Group = monitor.Name
}
monitor.Pelican = detections[group].Pelican
monitor.Benchmark = detections[group].Benchmark
monitor.Candy = detections[group].Candy
monitor.ProbeHistory = detections[group].ProbeHistory
applyCurrentStatus(&monitor, userResults[group], now.Unix())
monitors = append(monitors, monitor)
controller/uptime_kuma_test.go
+33 -19
View File
@@ -9,6 +9,7 @@ import (
fingerprint "github.com/QuantumNous/new-api/pkg/model_fingerprint"
"github.com/QuantumNous/new-api/service"
"github.com/QuantumNous/new-api/setting"
"github.com/QuantumNous/new-api/setting/ratio_setting"
"github.com/gin-gonic/gin"
"github.com/glebarez/sqlite"
"github.com/stretchr/testify/assert"
@@ -51,10 +52,11 @@ func TestStatusPageDatabaseMatrix(t *testing.T) {
oldType, oldLogType := common.MainDatabaseType(), common.LogDatabaseType()
oldRedis, oldSecret, oldSecure := common.RedisEnabled, common.SessionSecret, common.SessionCookieSecure
oldGroups := setting.UserUsableGroups2JSONString()
oldRatios := ratio_setting.GroupRatio2JSONString()
model.DB, model.LOG_DB = db, db
common.SetDatabaseTypes(common.DatabaseType(dialect), common.DatabaseType(dialect))
common.RedisEnabled, common.SessionSecret, common.SessionCookieSecure = false, "status-test-only", true
rows := []any{&model.Option{}, &model.Channel{}, &model.Ability{}, &model.SystemTask{}, &model.DetectionRun{}, &model.DetectionCall{}, &model.DetectionSchedule{}, &model.Log{}, &model.PerfMetric{}, &model.User{}, &model.UserSession{}}
rows := []any{&model.GroupMetadata{}, &model.Option{}, &model.Channel{}, &model.Ability{}, &model.SystemTask{}, &model.DetectionRun{}, &model.DetectionCall{}, &model.DetectionSchedule{}, &model.Log{}, &model.PerfMetric{}, &model.User{}, &model.UserSession{}}
require.NoError(t, db.AutoMigrate(rows...))
require.NoError(t, db.Table("audit_logs").AutoMigrate(&model.AuditLog{}))
t.Cleanup(func() {
@@ -66,6 +68,7 @@ func TestStatusPageDatabaseMatrix(t *testing.T) {
common.SetDatabaseTypes(oldType, oldLogType)
common.RedisEnabled, common.SessionSecret, common.SessionCookieSecure = oldRedis, oldSecret, oldSecure
require.NoError(t, setting.UpdateUserUsableGroupsByJSONString(oldGroups))
require.NoError(t, ratio_setting.UpdateGroupRatioByJSONString(oldRatios))
groupCacheSnapshot.updated = time.Time{}
require.NoError(t, conn.Close())
})
@@ -77,15 +80,17 @@ func TestStatusPageDatabaseMatrix(t *testing.T) {
require.NoError(t, db.Raw(q).Scan(&version).Error)
t.Log("database version:", version)
require.NoError(t, setting.UpdateUserUsableGroupsByJSONString(`{"internal":"","shared":"","hidden":""}`))
require.NoError(t, ratio_setting.UpdateGroupRatioByJSONString(`{"internal":0.25,"shared":0,"hidden":1}`))
now := time.Now().Unix()
channels := []model.Channel{{Id: 1, Name: "private-channel", Key: "private-key", Group: "internal,shared", Models: fingerprint.Model, Status: 1}, {Id: 2, Group: "hidden", Models: fingerprint.Model, Status: 1}, {Id: 3, Group: "internal", Models: "disabled-model", Status: 2}}
require.NoError(t, db.Create(&channels).Error)
abilities := []model.Ability{{Group: "internal", Model: fingerprint.Model, ChannelId: 1, Enabled: true}, {Group: "shared", Model: fingerprint.Model, ChannelId: 1, Enabled: true}, {Group: "hidden", Model: fingerprint.Model, ChannelId: 2, Enabled: true}, {Group: "internal", Model: "disabled-model", ChannelId: 3, Enabled: true}, {Group: "internal", Model: "disabled-ability", ChannelId: 1}}
require.NoError(t, db.Create(&abilities).Error)
state, err := common.Marshal(service.AstraSamplingState{Rounds: []service.AstraSamplingRound{{ID: "old-round", ChannelID: 1, StartedAt: now - 20, FinishedAt: now - 10, Result: fingerprint.Result{Status: "match", Closest: fingerprint.Model, ValidSamples: 3}}}})
require.NoError(t, db.Create(&model.Log{Type: model.LogTypeConsume, TokenId: 1, ChannelId: 1, Group: "internal", ModelName: fingerprint.Model, Quota: 50, PromptTokens: 1000, CreatedAt: now - 10, Other: `{}`}).Error)
state, err := common.Marshal(service.AstraSamplingState{Rounds: []service.AstraSamplingRound{{ID: "old-round", ChannelID: 1, StartedAt: now - 20, FinishedAt: now - 10, Result: fingerprint.Result{Engine: fingerprint.Version, Status: "match", Closest: fingerprint.Model, ValidSamples: 3}}}})
require.NoError(t, err)
require.NoError(t, db.Create(&model.SystemTask{TaskID: "old-task", Type: service.AstraSamplingTaskType, Status: model.SystemTaskStatusSucceeded, CreatedAt: now - 20, State: string(state)}).Error)
runs := []model.DetectionRun{{ID: "old-work", RequestKey: "old-work", ChannelID: 1, GroupName: "internal", Kind: "pelican", Status: "succeeded", Quality: "normal", CreatedAt: now - 20, Settled: true, HTML: "<h1>art</h1>", Charged: 17}, {ID: "in-flight", RequestKey: "in-flight", ChannelID: 1, Kind: "pelican", Status: "running", CreatedAt: now - 10}, {ID: "hidden-art", RequestKey: "hidden-art", ChannelID: 2, Kind: "pelican", Status: "succeeded", CreatedAt: now - 20, HTML: "hidden"}}
runs := []model.DetectionRun{{ID: "old-work", RequestKey: "old-work", ChannelID: 1, GroupName: "internal", Kind: "benchmark", Status: "succeeded", Quality: "normal", CreatedAt: now - 20, Settled: true, HTML: "<h1>art</h1>", Charged: 17}, {ID: "in-flight", RequestKey: "in-flight", ChannelID: 1, Kind: "benchmark", Status: "running", CreatedAt: now - 10}, {ID: "hidden-art", RequestKey: "hidden-art", ChannelID: 2, Kind: "benchmark", Status: "succeeded", CreatedAt: now - 20, HTML: "hidden"}}
require.NoError(t, db.Create(&runs).Error)
tokens := map[int]string{}
for _, role := range []int{1, 10} {
@@ -131,23 +136,32 @@ func TestStatusPageDatabaseMatrix(t *testing.T) {
groupCacheSnapshot.updated = time.Time{}
res := request("GET", "/api/uptime/status", "", "", "")
require.Equal(t, 200, res.Code, res.Body.String())
assert.NotContains(t, res.Body.String(), "model_costs")
for _, secret := range []string{"internal", "secret-alias", "private-channel", "private-key", "disabled-model", "disabled-ability"} {
assert.NotContains(t, res.Body.String(), secret)
}
assert.Contains(t, res.Body.String(), "#688")
fallback := model.DetectionRun{ID:"fallback-completed",RequestKey:"fallback-completed",Kind:"pelican",ChannelID:1,GroupName:"internal",Status:"running",Quality:"normal",HTML:"<html>pelican</html>",CreatedAt:now,Deadline:now+600}
require.NoError(t,db.Create(&fallback).Error)
require.NoError(t,db.Create(&model.DetectionCall{ID:fallback.ID,RunID:fallback.ID,SecretHash:"fallback-only-test",State:"finished"}).Error)
processLocalPelican(context.Background(),fallback)
require.NoError(t,db.First(&fallback,"id = ?",fallback.ID).Error)
assert.Equal(t,"succeeded",fallback.Status)
assert.Equal(t,"unknown",fallback.Quality)
assert.Equal(t,"<html>pelican</html>",fallback.HTML)
assert.True(t,fallback.Settled)
require.NoError(t,db.Delete(&model.DetectionCall{},"id = ?",fallback.ID).Error)
require.NoError(t,db.Delete(&model.DetectionRun{},"id = ?",fallback.ID).Error)
assert.Contains(t, res.Body.String(), `"group_ratio":0.25`)
require.NoError(t, model.SaveGroupMetadata(context.Background(), map[string]model.GroupMetadataValue{"internal": {Tags: []string{"打野", "不稳定"}}}))
service.InvalidateGroupQualityCache()
visitorSite := request("GET", "/api/group-status", "", "", "")
require.Equal(t, 200, visitorSite.Code)
var siteData struct {
Data []UptimeGroupResult `json:"data"`
}
require.NoError(t, common.Unmarshal(visitorSite.Body.Bytes(), &siteData))
require.NotEmpty(t, siteData.Data)
for _, monitor := range siteData.Data[0].Monitors {
if monitor.Group == "internal" {
require.Equal(t, []string{"打野", "不稳定"}, monitor.Tags)
require.NotNil(t, monitor.ModelCosts[fingerprint.Model].USDPerMillion)
assert.Equal(t, int64(1), monitor.ModelCosts[fingerprint.Model].Requests)
}
if monitor.Group == "shared" {
assert.Empty(t, monitor.ModelCosts, "shared channel membership must not copy another group's billed usage")
}
}
assert.NotContains(t, request("GET", "/api/uptime/status", "", "", "").Body.String(), "model_costs")
assert.Contains(t, visitorSite.Body.String(), `"group":"hidden"`)
assert.Contains(t, visitorSite.Body.String(), `"name":"internal"`)
assert.NotContains(t, visitorSite.Body.String(), "#688")
@@ -171,9 +185,9 @@ func TestStatusPageDatabaseMatrix(t *testing.T) {
}
assert.Empty(t, byGroup["internal"].Fingerprint.History)
assert.Nil(t, byGroup["internal"].Fingerprint.Rate)
assert.Empty(t, byGroup["internal"].Pelican.History)
assert.Empty(t, byGroup["internal"].Benchmark.History)
assert.Len(t, byGroup["shared"].Fingerprint.History, 1)
assert.Len(t, byGroup["shared"].Pelican.History, 2)
assert.Len(t, byGroup["shared"].Benchmark.History, 2)
assert.True(t, byGroup["hidden"].Hidden)
assert.Equal(t, []string{fingerprint.Model}, byGroup["internal"].Models)
require.NoError(t, model.SaveStatusPageGroup(context.Background(), "internal", &model.StatusPageGroup{Alias: "#689", Hidden: true}))
@@ -187,13 +201,13 @@ func TestStatusPageDatabaseMatrix(t *testing.T) {
require.NoError(t, db.First(&retained, "id = ?", "old-work").Error)
assert.Equal(t, 17, retained.Charged)
assert.Equal(t, "<h1>art</h1>", retained.HTML)
require.NoError(t, db.Create(&model.DetectionRun{ID: "new-work", RequestKey: "new-work", ChannelID: 1, Kind: "pelican", Status: "succeeded", Quality: "normal", CreatedAt: settings["internal"].ResetBefore + 1, Settled: true}).Error)
require.NoError(t, db.Create(&model.DetectionRun{ID: "new-work", RequestKey: "new-work", ChannelID: 1, Kind: "benchmark", Status: "succeeded", Quality: "normal", CreatedAt: settings["internal"].ResetBefore + 1, Settled: true}).Error)
groups, err = currentGroupMonitors(context.Background(), true)
require.NoError(t, err)
for _, m := range groups[0].Monitors {
if m.Group == "internal" {
require.Len(t, m.Pelican.History, 1)
assert.Equal(t, "new-work", m.Pelican.History[0].ID)
require.Len(t, m.Benchmark.History, 1)
assert.Equal(t, "new-work", m.Benchmark.History[0].ID)
}
}
})
docs/error-log-group-custom19.md
+1
View File
@@ -5,5 +5,6 @@
- 历史错误日志在 API 输出时，依据保存的 request_policy 中匹配渠道的分组纠正显示；在用户权限裁剪之前完成，仍不泄露私有路由信息。
- 无证据的旧错误记录不猜测分组。消费日志、金额、token 用量和数据库中的历史记录不改写。
- 历史记录的数据库分组筛选仍按原存储值工作；新的错误记录保存正确分组，筛选正常。
- 用户机器将 features.memories 暂设为 false，原配置已备份，记忆文件未删除；用户需重启 Codex 生效，暂停期间也不自动注入已有记忆。
测试：实际错误分组、重试分组、无渠道、固定分组；历史读时投影及普通用户日志隐私；真实候选 HTTP 回归；旧版 Astra 守卫/逐次结算保持不变。
middleware/auth.go
+5
View File
@@ -450,6 +450,11 @@ func TokenAuth() func(c *gin.Context) {
return
}
if err := model.LoadPersonalGroupRatios(c, userCache.Id); err != nil {
abortWithOpenAiMessage(c, http.StatusServiceUnavailable, "User pricing temporarily unavailable")
return
}
defer service.ReleasePoolProbes(c)
userCache.WriteContext(c)
userGroup := userCache.Group
middleware/distributor.go
+22 -5
View File
@@ -27,13 +27,16 @@ import (
)
type ModelRequest struct {
Model string `json:"model"`
Group string `json:"group,omitempty"`
Model string `json:"model"`
Group string `json:"group,omitempty"`
Stream bool `json:"stream"`
}
func Distribute() func(c *gin.Context) {
return func(c *gin.Context) {
var channel *model.Channel
var requestedStream bool
rejectionMessage := "No available channel for this request."
defer func() {
if c.Writer.Status() >= 400 {
service.RecordRequestPolicyTermination(c, types.NewErrorWithStatusCode(errors.New("request rejected"), types.ErrorCodeInvalidRequest, c.Writer.Status(), types.ErrOptionWithSkipRetry()))
@@ -42,7 +45,7 @@ func Distribute() func(c *gin.Context) {
other := model.NewLogOther()
other.SetPublic("status_code", c.Writer.Status())
service.AppendRelayLogAdminInfo(c, nil, other)
model.RecordErrorLog(c, c.GetInt("id"), 0, c.GetString("original_model"), c.GetString("token_name"), "No available channel for this request.", c.GetInt("token_id"), 0, false, "", other)
model.RecordErrorLog(c, c.GetInt("id"), 0, c.GetString("original_model"), c.GetString("token_name"), rejectionMessage, c.GetInt("token_id"), 0, requestedStream, "", other)
}
}
}()
@@ -57,6 +60,8 @@ func Distribute() func(c *gin.Context) {
abortWithOpenAiMessage(c, http.StatusBadRequest, i18n.T(c, i18n.MsgDistributorInvalidRequest, map[string]any{"Error": err.Error()}))
return
}
requestedStream = modelRequest.Stream || strings.Contains(c.Request.URL.Path, "streamGenerateContent") || strings.EqualFold(c.GetHeader("Upgrade"), "websocket")
c.Set("pool_request_stream", requestedStream)
_, pinned, _ := constraints.ResolvedPin()
if !pinned {
// Select a channel for the user
@@ -124,6 +129,17 @@ func Distribute() func(c *gin.Context) {
} else if selectErr.MessageID != "" {
message = i18n.T(c, selectErr.MessageID, selectErr.Params)
}
if service.RequestRoutingPolicy(c) != nil && selectErr.NoAvailableChannel {
if c.GetBool("pool_degradation_unavailable") {
message = "Group quality checks are temporarily unavailable. Please retry later."
} else if wait := c.GetInt("pool_retry_after"); wait > 0 {
message = fmt.Sprintf("No eligible group is currently available; failed groups are cooling down or being checked. Retry after %d seconds.", wait)
c.Header("Retry-After", fmt.Sprint(wait))
} else {
message = "No eligible group is available for this model under the current pool and quality settings."
}
}
rejectionMessage = message
abortWithOpenAiMessage(c, selectErr.StatusCode, message, selectErr.Code)
return
}
@@ -295,8 +311,9 @@ func getModelFromJSONBody(c *gin.Context) (*ModelRequest, error) {
c.Request.Body = io.NopCloser(storage)
return &ModelRequest{
Model: model,
Group: group,
Model: model,
Group: group,
Stream: gjson.GetBytes(requestBody, "stream").Bool(),
}, nil
}
model/group_metadata.go
+74
View File
@@ -0,0 +1,74 @@
package model
import (
"context"
"encoding/json"
"fmt"
"strings"
"unicode/utf8"
"github.com/QuantumNous/new-api/setting/ratio_setting"
"gorm.io/gorm"
"gorm.io/gorm/clause"
)
type GroupMetadata struct {
GroupName string `gorm:"primaryKey;size:191"`
Value string `gorm:"type:text;not null"`
}
type GroupMetadataValue struct {
Tags []string `json:"tags"`
HideRecentDegraded bool `json:"hide_recent_degraded"`
}
func GetGroupMetadata(ctx context.Context) (map[string]GroupMetadataValue, error) {
var rows []GroupMetadata
if err := DB.WithContext(ctx).Find(&rows).Error; err != nil {
return nil, err
}
result := map[string]GroupMetadataValue{}
for _, row := range rows {
var value GroupMetadataValue
if err := json.Unmarshal([]byte(row.Value), &value); err != nil {
return nil, err
}
result[row.GroupName] = value
}
return result, nil
}
// Updates only submitted groups; absent groups keep their existing metadata.
func SaveGroupMetadata(ctx context.Context, groups map[string]GroupMetadataValue) error {
rows := []GroupMetadata{}
for group, value := range groups {
if !ratio_setting.ContainsGroupRatio(group) || group == "auto" {
return fmt.Errorf("invalid group: %s", group)
}
tags, seen := []string{}, map[string]bool{}
for _, tag := range value.Tags {
tag = strings.TrimSpace(tag)
if tag == "" || seen[tag] {
continue
}
if utf8.RuneCountInString(tag) > 40 || len(tags) >= 20 {
return fmt.Errorf("too many or too long group tags")
}
seen[tag] = true
tags = append(tags, tag)
}
value.Tags = tags
body, err := json.Marshal(value)
if err != nil {
return err
}
rows = append(rows, GroupMetadata{GroupName: group, Value: string(body)})
}
return DB.WithContext(ctx).Transaction(func(tx *gorm.DB) error {
for _, row := range rows {
if err := tx.Clauses(clause.OnConflict{Columns: []clause.Column{{Name: "group_name"}}, DoUpdates: clause.AssignmentColumns([]string{"value"})}).Create(&row).Error; err != nil {
return err
}
}
return nil
})
}
model/group_model_cost.go
+86
View File
@@ -0,0 +1,86 @@
package model
import (
"context"
"time"
"github.com/QuantumNous/new-api/common"
)
type GroupModelCost struct {
USDPerMillion *float64 `json:"usd_per_million"`
Requests int64 `json:"requests"`
Tokens float64 `json:"tokens"`
quota float64
}
// Consumption logs contain the settled charge, not the pre-consumed quota.
// Attribute by the group actually billed, never by today's channel membership.
func GetGroupModelCosts(ctx context.Context, start, end time.Time) (map[string]map[string]GroupModelCost, error) {
result := map[string]map[string]GroupModelCost{}
rows, err := LOG_DB.WithContext(ctx).Model(&Log{}).Select("group", "model_name", "quota", "prompt_tokens", "completion_tokens", "other").
Where("type = ? AND token_id > 0 AND created_at >= ? AND created_at < ?", LogTypeConsume, start.Unix(), end.Unix()).Rows()
if err != nil {
return nil, err
}
defer rows.Close()
for rows.Next() {
var group, name, raw string
var quota, prompt, completion int64
if err := rows.Scan(&group, &name, &quota, &prompt, &completion, &raw); err != nil {
return nil, err
}
if group == "" || name == "" || quota < 0 {
continue
}
tokens, valid := consumptionTokenTotal(prompt, completion, raw)
if !valid || tokens <= 0 {
continue
}
if result[group] == nil {
result[group] = map[string]GroupModelCost{}
}
value := result[group][name]
value.Requests++
value.Tokens += tokens
value.quota += float64(quota)
result[group][name] = value
}
if err := rows.Err(); err != nil {
return nil, err
}
for group, models := range result {
for name, value := range models {
if common.QuotaPerUnit > 0 && value.Tokens > 0 {
rate := value.quota / common.QuotaPerUnit / value.Tokens * 1_000_000
value.USDPerMillion = &rate
}
result[group][name] = value
}
}
return result, nil
}
func consumptionTokenTotal(prompt, completion int64, raw string) (float64, bool) {
var usage struct {
Total int64 `json:"input_tokens_total"`
Cached int64 `json:"cache_tokens"`
Created int64 `json:"cache_creation_tokens"`
Created5m int64 `json:"cache_creation_tokens_5m"`
Created1h int64 `json:"cache_creation_tokens_1h"`
Written int64 `json:"cache_write_tokens"`
Semantic string `json:"usage_semantic"`
Claude bool `json:"claude"`
}
if raw != "" && common.UnmarshalJsonStr(raw, &usage) != nil {
return 0, false
}
input := float64(max(prompt, 0))
if usage.Total > 0 {
input = float64(usage.Total)
} else if usage.Semantic == "anthropic" || (usage.Semantic == "" && usage.Claude) {
created := max(float64(usage.Written), float64(usage.Created), float64(max(usage.Created5m, 0))+float64(max(usage.Created1h, 0)), 0)
input += float64(max(usage.Cached, 0)) + created
}
return input + float64(max(completion, 0)), true
}
model/group_status_test.go
+58
View File
@@ -3,6 +3,7 @@ package model
import (
"context"
"fmt"
"github.com/QuantumNous/new-api/common"
"os"
"testing"
"time"
@@ -16,6 +17,63 @@ import (
"gorm.io/gorm/schema"
)
func TestGroupModelCostsDatabaseMatrix(t *testing.T) {
for _, dialect := range []string{"sqlite", "mysql", "postgres"} {
t.Run(dialect, func(t *testing.T) {
var driver gorm.Dialector
switch dialect {
case "sqlite":
driver = sqlite.Open(":memory:")
case "mysql":
if os.Getenv("TEST_MYSQL_DSN") == "" {
t.Skip("no MySQL test database")
}
driver = mysql.Open(os.Getenv("TEST_MYSQL_DSN"))
case "postgres":
if os.Getenv("TEST_POSTGRES_DSN") == "" {
t.Skip("no PostgreSQL test database")
}
driver = postgres.Open(os.Getenv("TEST_POSTGRES_DSN"))
}
db, err := gorm.Open(driver, &gorm.Config{NamingStrategy: schema.NamingStrategy{TablePrefix: fmt.Sprintf("cost_%d_", time.Now().UnixNano())}})
require.NoError(t, err)
oldDB, oldUnit := LOG_DB, common.QuotaPerUnit
LOG_DB = db
common.QuotaPerUnit = 500000
t.Cleanup(func() {
LOG_DB = oldDB
common.QuotaPerUnit = oldUnit
require.NoError(t, db.Migrator().DropTable(&Log{}))
sqlDB, e := db.DB()
require.NoError(t, e)
require.NoError(t, sqlDB.Close())
})
require.NoError(t, db.AutoMigrate(&Log{}))
logs := []Log{
{Type: 2, TokenId: 1, CreatedAt: 150, Group: "a", ModelName: "m", Quota: 50, PromptTokens: 100, CompletionTokens: 100, Other: `{"cache_tokens":80}`},
{Type: 2, TokenId: 2, CreatedAt: 150, Group: "a", ModelName: "m", Quota: 150, PromptTokens: 200, CompletionTokens: 100, Other: `{"usage_semantic":"anthropic","cache_tokens":100,"cache_creation_tokens":400}`},
{Type: 2, TokenId: 3, CreatedAt: 150, Group: "b", ModelName: "m", Quota: 1000, PromptTokens: 200, CompletionTokens: 100, Other: `{"input_tokens_total":900}`},
{Type: 2, TokenId: 1, CreatedAt: 150, Group: "a", ModelName: "free", Quota: 0, PromptTokens: 100},
{Type: 2, TokenId: 0, CreatedAt: 150, Group: "a", ModelName: "m", Quota: 999999, PromptTokens: 100},
{Type: 2, TokenId: 1, CreatedAt: 99, Group: "a", ModelName: "m", Quota: 999999, PromptTokens: 100},
{Type: 2, TokenId: 1, CreatedAt: 200, Group: "a", ModelName: "m", Quota: 999999, PromptTokens: 100},
{Type: 5, TokenId: 1, CreatedAt: 150, Group: "a", ModelName: "m", Quota: 999999, PromptTokens: 100},
{Type: 2, TokenId: 1, CreatedAt: 150, Group: "a", ModelName: "no-tokens", Quota: 500},
}
require.NoError(t, db.Create(&logs).Error)
result, err := GetGroupModelCosts(context.Background(), time.Unix(100, 0), time.Unix(200, 0))
require.NoError(t, err)
require.NotNil(t, result["a"]["m"].USDPerMillion)
assert.InDelta(t, 0.4, *result["a"]["m"].USDPerMillion, 1e-9)
assert.Equal(t, int64(2), result["a"]["m"].Requests)
assert.Equal(t, float64(1000), result["a"]["m"].Tokens)
assert.InDelta(t, 2, *result["b"]["m"].USDPerMillion, 1e-9)
assert.Zero(t, *result["a"]["free"].USDPerMillion)
assert.NotContains(t, result["a"], "no-tokens")
})
}
}
func TestGroupCacheStatsDatabaseMatrix(t *testing.T) {
for _, dialect := range []string{"sqlite", "mysql", "postgres"} {
t.Run(dialect, func(t *testing.T) {
model/main.go
+2
View File
@@ -352,6 +352,8 @@ func migrateDB() error {
&Token{},
&RelayAttempt{},
&User{},
&PersonalGroupRatio{},
&GroupMetadata{},
&UserSession{},
&AuthFlow{},
&ExternalIdentityClaim{},
model/personal_group_ratio.go
+87
View File
@@ -0,0 +1,87 @@
package model
import (
"context"
"encoding/json"
"fmt"
"math"
"github.com/QuantumNous/new-api/setting/ratio_setting"
"github.com/QuantumNous/new-api/types"
"github.com/gin-gonic/gin"
"gorm.io/gorm/clause"
)
// Stored independently from User and UserSetting to prevent profile mass assignment.
type PersonalGroupRatio struct {
UserID int `gorm:"primaryKey;autoIncrement:false"`
Ratios string `gorm:"type:text;not null"`
}
const personalRatioContextKey = "personal_group_ratios"
func GetPersonalGroupRatios(ctx context.Context, userID int) (map[string]float64, error) {
var rows []PersonalGroupRatio
if err := DB.WithContext(ctx).Where("user_id = ?", userID).Find(&rows).Error; err != nil {
return nil, err
}
result := map[string]float64{}
if len(rows) != 0 {
if err := json.Unmarshal([]byte(rows[0].Ratios), &result); err != nil {
return nil, err
}
}
return result, nil
}
func SavePersonalGroupRatios(ctx context.Context, userID int, ratios map[string]float64) error {
for group, ratio := range ratios {
if !ratio_setting.ContainsGroupRatio(group) || group == "auto" || ratio < 0 || math.IsNaN(ratio) || math.IsInf(ratio, 0) {
return fmt.Errorf("invalid group or ratio: %s", group)
}
}
if ratios == nil {
ratios = map[string]float64{}
}
body, err := json.Marshal(ratios)
if err != nil {
return err
}
row := PersonalGroupRatio{UserID: userID, Ratios: string(body)}
if err = DB.WithContext(ctx).Clauses(clause.OnConflict{Columns: []clause.Column{{Name: "user_id"}}, DoUpdates: clause.AssignmentColumns([]string{"ratios"})}).Create(&row).Error; err != nil {
return err
}
// Pricing is loaded from the database at request admission, never from a stale
// process/Redis cache. The request snapshot survives retries and price updates.
return invalidateUserCache(userID)
}
func LoadPersonalGroupRatios(c *gin.Context, userID int) error {
ratios, err := GetPersonalGroupRatios(c.Request.Context(), userID)
if err == nil {
c.Set(personalRatioContextKey, ratios)
}
return err
}
func EffectiveGroupRatio(c *gin.Context, userGroup, group string) types.GroupRatioInfo {
result := types.GroupRatioInfo{GroupRatio: ratio_setting.GetGroupRatio(group), GroupSpecialRatio: -1}
if ratio, ok := ratio_setting.GetGroupGroupRatio(userGroup, group); ok {
result.GroupRatio, result.GroupSpecialRatio, result.HasSpecialRatio = ratio, ratio, true
}
if value, ok := c.Get(personalRatioContextKey); ok {
if ratios, ok := value.(map[string]float64); ok {
if ratio, ok := ratios[group]; ok {
result.GroupRatio, result.GroupSpecialRatio, result.HasSpecialRatio = ratio, ratio, true
}
}
}
return result
}
func HasPersonalGroupRatio(c *gin.Context, group string) bool {
value, _ := c.Get(personalRatioContextKey)
ratios, _ := value.(map[string]float64)
_, ok := ratios[group]
return ok
}
model/personal_group_ratio_test.go
+96
View File
@@ -0,0 +1,96 @@
package model
import (
"context"
"encoding/json"
"fmt"
"math"
"net/http/httptest"
"os"
"testing"
"time"
"github.com/QuantumNous/new-api/common"
"github.com/QuantumNous/new-api/setting/ratio_setting"
"github.com/gin-gonic/gin"
"github.com/glebarez/sqlite"
"github.com/stretchr/testify/require"
"gorm.io/driver/mysql"
"gorm.io/driver/postgres"
"gorm.io/gorm"
"gorm.io/gorm/schema"
)
func TestPersonalGroupConfigurationDatabaseMatrix(t *testing.T) {
for _, dialect := range []string{"sqlite", "mysql", "postgres"} {
t.Run(dialect, func(t *testing.T) {
var driver gorm.Dialector
switch dialect {
case "sqlite":
driver = sqlite.Open(":memory:")
case "mysql":
if os.Getenv("TEST_MYSQL_DSN") == "" {
t.Skip("DSN unset")
}
driver = mysql.Open(os.Getenv("TEST_MYSQL_DSN"))
case "postgres":
if os.Getenv("TEST_POSTGRES_DSN") == "" {
t.Skip("DSN unset")
}
driver = postgres.Open(os.Getenv("TEST_POSTGRES_DSN"))
}
db, err := gorm.Open(driver, &gorm.Config{NamingStrategy: schema.NamingStrategy{TablePrefix: fmt.Sprintf("personal_%d_", time.Now().UnixNano())}})
require.NoError(t, err)
old, redis := DB, common.RedisEnabled
DB, common.RedisEnabled = db, false
ratios, special := ratio_setting.GroupRatio2JSONString(), ratio_setting.GroupGroupRatio2JSONString()
rows := []any{&PersonalGroupRatio{}, &GroupMetadata{}}
t.Cleanup(func() {
for _, row := range rows {
require.NoError(t, db.Migrator().DropTable(row))
}
DB, common.RedisEnabled = old, redis
require.NoError(t, ratio_setting.UpdateGroupRatioByJSONString(ratios))
require.NoError(t, ratio_setting.UpdateGroupGroupRatioByJSONString(special))
conn, _ := db.DB()
_ = conn.Close()
})
require.NoError(t, db.AutoMigrate(rows...))
require.NoError(t, db.AutoMigrate(rows...))
require.NoError(t, ratio_setting.UpdateGroupRatioByJSONString("{\"default\":1,\"vip\":2}"))
require.NoError(t, ratio_setting.UpdateGroupGroupRatioByJSONString("{\"default\":{\"vip\":0.7}}"))
makeCtx := func() *gin.Context {
c, _ := gin.CreateTestContext(httptest.NewRecorder())
c.Request = httptest.NewRequest("POST", "/v1/responses", nil)
require.NoError(t, LoadPersonalGroupRatios(c, 42))
return c
}
require.Equal(t, 0.7, EffectiveGroupRatio(makeCtx(), "default", "vip").GroupRatio)
require.Equal(t, 1.0, EffectiveGroupRatio(makeCtx(), "default", "default").GroupRatio)
require.NoError(t, SavePersonalGroupRatios(context.Background(), 42, map[string]float64{"vip": 0.2, "default": 0}))
inFlight := makeCtx()
require.Equal(t, 0.2, EffectiveGroupRatio(inFlight, "default", "vip").GroupRatio)
require.Equal(t, 0.0, EffectiveGroupRatio(inFlight, "default", "default").GroupRatio)
require.True(t, HasPersonalGroupRatio(inFlight, "vip"))
for _, bad := range []float64{-1, math.NaN(), math.Inf(1)} {
require.Error(t, SavePersonalGroupRatios(context.Background(), 42, map[string]float64{"vip": bad}))
}
require.NoError(t, SavePersonalGroupRatios(context.Background(), 42, map[string]float64{}))
require.Equal(t, 0.2, EffectiveGroupRatio(inFlight, "default", "vip").GroupRatio)
require.Equal(t, 0.7, EffectiveGroupRatio(makeCtx(), "default", "vip").GroupRatio)
require.False(t, HasPersonalGroupRatio(makeCtx(), "vip"))
require.NoError(t, SaveGroupMetadata(context.Background(), map[string]GroupMetadataValue{"vip": {Tags: []string{"官Key", "官Key", " 打野 ", "不稳定"}, HideRecentDegraded: true}}))
metadata, err := GetGroupMetadata(context.Background())
require.NoError(t, err)
require.Equal(t, []string{"官Key", "打野", "不稳定"}, metadata["vip"].Tags)
require.True(t, metadata["vip"].HideRecentDegraded)
require.NoError(t, db.AutoMigrate(rows...))
metadata, err = GetGroupMetadata(context.Background())
require.NoError(t, err)
require.Len(t, metadata["vip"].Tags, 3)
var u User
require.NoError(t, json.Unmarshal([]byte("{\"personal_group_ratios\":{\"vip\":0},\"ratios\":{\"vip\":0}}"), &u))
require.Equal(t, 0.7, EffectiveGroupRatio(makeCtx(), "default", "vip").GroupRatio)
})
}
}
model/token_routing.go
+12 -5
View File
@@ -16,19 +16,26 @@ type ModelRoutingRule struct {
}
type TokenRoutingPolicy struct {
Version int `json:"version"`
Mode string `json:"mode"`
Groups []TokenGroupEntry `json:"groups"`
ModelRules []ModelRoutingRule `json:"model_rules"`
SkipDegradedNonAstra bool `json:"skip_degraded_non_astra,omitempty"`
AutoPreference string `json:"auto_preference,omitempty"`
AllowDegraded bool `json:"allow_degraded,omitempty"`
SkipDegradedGroups bool `json:"skip_degraded_groups,omitempty"`
Version int `json:"version"`
Mode string `json:"mode"`
Groups []TokenGroupEntry `json:"groups"`
ModelRules []ModelRoutingRule `json:"model_rules"`
}
func (p *TokenRoutingPolicy) Validate(maxGroups int, allowed func(string) bool) error {
if p.Version != 1 || (p.Mode != "custom" && p.Mode != "inherit" && p.Mode != "user") {
if p.Version != 1 || (p.Mode != "custom" && p.Mode != "inherit" && p.Mode != "user" && p.Mode != "smart") {
return fmt.Errorf("invalid routing policy version or mode")
}
if p.Mode == "custom" && len(p.Groups) == 0 {
return fmt.Errorf("a group pool must contain at least one group")
}
if p.Mode == "smart" && (len(p.ModelRules) != 0 || (p.AutoPreference != "" && p.AutoPreference != "cost" && p.AutoPreference != "speed" && p.AutoPreference != "stable")) {
return fmt.Errorf("invalid automatic routing settings")
}
if p.Mode != "custom" && len(p.Groups) != 0 {
return fmt.Errorf("inherited pools cannot contain custom groups")
}
pkg/billingexpr/image_generation_price_test.go
+22
View File
@@ -0,0 +1,22 @@
package billingexpr
import (
"github.com/stretchr/testify/require"
"testing"
)
func TestImageGenerationPriceSizesAndQuantity(t *testing.T) {
expression := `(param("size") != nil && (param("size") == "4k" || (param("size") matches "^[0-9]+x[0-9]+$" && min(float(split(param("size"), "x")[0]), float(split(param("size"), "x")[1])) > 2048)) ? tier("4k", fixed(0.12)) : param("size") != nil && (param("size") == "2k" || (param("size") matches "^[0-9]+x[0-9]+$" && min(float(split(param("size"), "x")[0]), float(split(param("size"), "x")[1])) > 1024)) ? tier("2k", fixed(0.08)) : tier("1k", fixed(0.06))) * image_count`
for _, tc := range []struct {
body, tier string
price float64
}{{`{}`, "1k", 0.06}, {`{"size":"auto"}`, "1k", 0.06}, {`{"size":"1024x1536"}`, "1k", 0.06}, {`{"size":"2048x1152"}`, "2k", 0.08}, {`{"size":"3840x2160"}`, "4k", 0.12}, {`{"size":"2k"}`, "2k", 0.08}} {
for _, n := range []int{1, 2, 3} {
cost, trace, err := RunExprWithRequest(expression, TokenParams{P: 10000, C: 5000}, RequestInput{Body: []byte(tc.body), ImageCount: &n})
require.NoError(t, err)
require.InDelta(t, tc.price*float64(n)*1000000, cost, 0.01)
require.Equal(t, tc.tier, trace.MatchedTier)
require.Equal(t, BillingUnitRequest, trace.BillingUnit)
}
}
}
pkg/model_benchmark/LICENSE-meow.txt
+75
View File
@@ -0,0 +1,75 @@
# PolyForm Noncommercial License 1.0.0
<https://polyformproject.org/licenses/noncommercial/1.0.0>
## Acceptance
In order to get any license under these terms, you must agree to them as both strict obligations and conditions to all your licenses.
## Copyright License
The licensor grants you a copyright license for the software to do everything you might do with the software that would otherwise infringe the licensor's copyright in it for any permitted purpose. However, you may only distribute the software according to [Distribution License](#distribution-license) and make changes or new works based on the software according to [Changes and New Works License](#changes-and-new-works-license).
## Distribution License
The licensor grants you an additional copyright license to distribute copies of the software. Your license to distribute covers distributing the software with changes and new works permitted by [Changes and New Works License](#changes-and-new-works-license).
## Notices
You must ensure that anyone who gets a copy of any part of the software from you also gets a copy of these terms or the URL for them above, as well as copies of any plain-text lines beginning with `Required Notice:` that the licensor provided with the software. For example:
> Required Notice: Copyright Yoyodyne, Inc. (http://example.com)
Required Notice: Copyright 2026 chen-006 and contributors. Original project: https://github.com/chen-006/gpt56_api_detector
## Changes and New Works License
The licensor grants you an additional copyright license to make changes and new works based on the software for any permitted purpose.
## Patent License
The licensor grants you a patent license for the software that covers patent claims the licensor can license, or becomes able to license, that you would infringe by using the software.
## Noncommercial Purposes
Any noncommercial purpose is a permitted purpose.
## Personal Uses
Personal use for research, experiment, and testing for the benefit of public knowledge, personal study, private entertainment, hobby projects, amateur pursuits, or religious observance, without any anticipated commercial application, is use for a permitted purpose.
## Noncommercial Organizations
Use by any charitable organization, educational institution, public research organization, public safety or health organization, environmental protection organization, or government institution is use for a permitted purpose regardless of the source of funding or obligations resulting from the funding.
## Fair Use
You may have "fair use" rights for the software under the law. These terms do not limit them.
## No Other Rights
These terms do not allow you to sublicense or transfer any of your licenses to anyone else, or prevent the licensor from granting licenses to anyone else. These terms do not imply any other licenses.
## Patent Defense
If you make any written claim that the software infringes or contributes to infringement of any patent, your patent license for the software granted under these terms ends immediately. If your company makes such a claim, your patent license ends immediately for work on behalf of your company.
## Violations
The first time you are notified in writing that you have violated any of these terms, or done anything with the software not covered by your licenses, your licenses can nonetheless continue if you come into full compliance with these terms, and take practical steps to correct past violations, within 32 days of receiving notice. Otherwise, all your licenses end immediately.
## No Liability
***As far as the law allows, the software comes as is, without any warranty or condition, and the licensor will not be liable to you for any damages arising out of these terms or the use or nature of the software, under any kind of legal claim.***
## Definitions
The **licensor** is the individual or entity offering these terms, and the **software** is the software the licensor makes available under these terms.
**You** refers to the individual or entity agreeing to these terms.
**Your company** is any legal entity, sole proprietorship, or other kind of organization that you work for, plus all organizations that have control over, are under the control of, or are under common control with that organization. **Control** means ownership of substantially all the assets of an entity, or the power to direct its management and policies by vote, contract, or otherwise. Control can be direct or indirect.
**Your licenses** are all the licenses granted to you for the software under these terms.
**Use** means anything you do with the software requiring one of your licenses.
pkg/model_benchmark/fingerprint.go
+200
View File
@@ -0,0 +1,200 @@
// Meow predictive fingerprint integration, pinned to the public GPT benchmark.
// Required Notice: Copyright 2026 chen-006 and contributors.
// Original project: https://github.com/chen-006/gpt56_api_detector
// See LICENSE-meow.txt for the PolyForm Noncommercial 1.0.0 license.
package model_benchmark
import (
_ "embed"
"errors"
"math"
"sort"
"strings"
"sync"
"unicode/utf8"
"github.com/QuantumNous/new-api/common"
"golang.org/x/text/cases"
)
const Model = "gpt-6-astra"
const Version = "meow-4.5.4-predictive.20260924.2"
const ProbeCount = 32
const OutputTokens = 128
//go:embed meow.json
var benchmarkJSON []byte
type Sample struct {
Text string `json:"text"`
// Ordinal in the fixed quick-tier plan, under the existing storage key.
Count int `json:"expected_count"`
}
type Result struct {
Status string `json:"status"`
ValidSamples int `json:"valid_samples"`
Closest string `json:"closest,omitempty"`
Weight float64 `json:"relative_weight,omitempty"`
Gap float64 `json:"score_gap,omitempty"`
Engine string `json:"engine,omitempty"`
Matches map[string]float64 `json:"matches,omitempty"`
Thresholds map[string]float64 `json:"thresholds,omitempty"`
}
type Probe struct {
ID string `json:"id"`
System string `json:"system"`
Prompt string `json:"prompt"`
}
type fittedCell struct {
Categories []string `json:"categories"`
Alpha map[string][]float64 `json:"alpha"`
}
type benchmark struct {
Plan []Probe `json:"-"`
Probes []struct {
Cells []Probe `json:"cells"`
} `json:"probes"`
Fitted struct {
Models []string `json:"models"`
Sources []string `json:"sources"`
References []string `json:"reference_sources"`
Cells map[string]fittedCell `json:"cells"`
} `json:"fitted"`
Tiers map[string]struct {
Counts map[string]int `json:"counts"`
Thresholds map[string]float64 `json:"thresholds"`
} `json:"tiers"`
}
var load = sync.OnceValues(func() (*benchmark, error) {
var b benchmark
if err := common.Unmarshal(benchmarkJSON, &b); err != nil {
return nil, err
}
var plan []Probe
for _, p := range b.Probes {
for _, c := range p.Cells {
for range b.Tiers["low"].Counts[c.ID] {
plan = append(plan, c)
}
}
}
sort.SliceStable(plan, func(i, j int) bool { return plan[i].ID < plan[j].ID })
if len(plan) != ProbeCount {
return nil, errors.New("invalid meow plan")
}
b.Plan = plan
return &b, nil
})
func ProbeAt(index int) (Probe, error) {
b, err := load()
if err != nil {
return Probe{}, err
}
if index < 0 || index >= len(b.Plan) {
return Probe{}, errors.New("invalid probe index")
}
return b.Plan[index], nil
}
func ValidSample(text string) bool {
if !utf8.ValidString(text) || utf8.RuneCountInString(text) > 65536 {
return false
}
text = cases.Fold().String(strings.TrimSpace(text))
return text != "" && utf8.RuneCountInString(text) <= 4096
}
// Whole-run Dirichlet posterior prediction, using average log evidence versus
// the strongest rival and calibrated thresholds. Closest alone is no verdict.
func Score(samples []Sample) (Result, error) {
b, err := load()
if err != nil {
return Result{}, err
}
r := Result{Status: "inconclusive", Engine: Version, Matches: map[string]float64{}, Thresholds: b.Tiers["low"].Thresholds}
counts := map[string]map[string]int{}
seen := map[int]bool{}
for _, sample := range samples {
if sample.Count < 0 || sample.Count >= len(b.Plan) || seen[sample.Count] {
return r, errors.New("invalid or duplicate meow sample")
}
seen[sample.Count] = true
if !ValidSample(sample.Text) {
continue
}
id := b.Plan[sample.Count].ID
text := cases.Fold().String(strings.TrimSpace(sample.Text))
if counts[id] == nil {
counts[id] = map[string]int{}
}
counts[id][text]++
r.ValidSamples++
}
sufficient := r.ValidSamples >= int(math.Ceil(.6*ProbeCount))
evidence := map[string]float64{}
for id, planned := range b.Tiers["low"].Counts {
cell := b.Fitted.Cells[id]
numbers := make([]int, len(cell.Categories))
positions := map[string]int{}
for i, category := range cell.Categories {
positions[category] = i
}
n := 0
for text, number := range counts[id] {
i, ok := positions[text]
if !ok {
i = positions["__UNSEEN_IN_TRAINING__"]
}
numbers[i] += number
n += number
}
if n < int(math.Ceil(.6*float64(planned))) {
sufficient = false
}
for _, source := range b.Fitted.Sources {
alpha := cell.Alpha[source]
total, value := 0.0, 0.0
for i, a := range alpha {
total += a
left, _ := math.Lgamma(a + float64(numbers[i]))
right, _ := math.Lgamma(a)
value += left - right
}
left, _ := math.Lgamma(total)
right, _ := math.Lgamma(total + float64(n))
evidence[source] += value + left - right
}
}
evidence["other_known_external"] = math.Inf(-1)
for _, source := range b.Fitted.References {
evidence["other_known_external"] = math.Max(evidence["other_known_external"], evidence[source])
}
for _, model := range b.Fitted.Models {
rival := math.Inf(-1)
for _, other := range b.Fitted.Models {
if other != model {
rival = math.Max(rival, evidence[other])
}
}
r.Matches[model] = 1 / (1 + math.Exp((rival-evidence[model])/float64(max(r.ValidSamples, 1))))
}
ordered := append([]string(nil), b.Fitted.Models...)
sort.SliceStable(ordered, func(i, j int) bool { return r.Matches[ordered[i]] > r.Matches[ordered[j]] })
top := ordered[0]
r.Weight = r.Matches[top]
r.Gap = (evidence[top] - evidence[ordered[1]]) / float64(max(r.ValidSamples, 1))
if r.ValidSamples > 0 {
r.Closest = top
}
if sufficient && r.Weight > r.Thresholds[top] && r.Weight > r.Matches[ordered[1]] {
r.Status = "mismatch"
if top == Model {
r.Status = "match"
}
}
return r, nil
}
pkg/model_benchmark/fingerprint_test.go
+73
View File
@@ -0,0 +1,73 @@
package model_benchmark
import (
"os"
"strings"
"testing"
"github.com/QuantumNous/new-api/common"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
)
func TestFingerprintMeowParity(t *testing.T) {
data, err := os.ReadFile("testdata/meow-reference.json")
require.NoError(t, err)
var fixtures []struct {
Name string `json:"name"`
Samples []Sample `json:"samples"`
Expected struct {
Verdict string `json:"verdict"`
Matches map[string]float64 `json:"matches"`
Valid int `json:"valid_samples"`
} `json:"expected"`
}
require.NoError(t, common.Unmarshal(data, &fixtures))
for _, f := range fixtures {
t.Run(f.Name, func(t *testing.T) {
result, err := Score(f.Samples)
require.NoError(t, err)
status := f.Expected.Verdict
if status == "insufficient" {
status = "inconclusive"
}
assert.Equal(t, status, result.Status)
assert.Equal(t, f.Expected.Valid, result.ValidSamples)
for model, score := range f.Expected.Matches {
assert.InDelta(t, score, result.Matches[model], 1e-10, model)
}
})
}
_, err = Score([]Sample{{Text: "42", Count: 0}, {Text: "42", Count: 0}})
require.Error(t, err)
assert.False(t, ValidSample(" "))
assert.False(t, ValidSample(strings.Repeat("a", 4097)))
}
func TestFingerprintResponseCompletionAndTools(t *testing.T) {
stream := "data: {\"type\":\"response.output_text.delta\",\"delta\":\"42\"}\n\ndata: {\"type\":\"response.completed\",\"response\":{\"status\":\"completed\",\"usage\":{\"input_tokens\":10,\"output_tokens\":1}}}\n\n"
answer, err := ReadResponse(strings.NewReader(stream), "text/event-stream")
require.NoError(t, err)
assert.Equal(t, "42", answer.Text)
assert.True(t, answer.HasUsage)
answer, err = ReadResponse(strings.NewReader(strings.Replace(stream, `"status":"completed"`, `"status":"completed","output":[]`, 1)), "text/event-stream")
require.NoError(t, err)
assert.Empty(t, answer.Text)
completed := `{"status":"completed","model":"gpt-6-astra","output":[{"type":"message","content":[{"type":"output_text","text":"[1,2,3]"}]}],"usage":{"input_tokens":21,"output_tokens":12,"input_tokens_details":{"cached_tokens":10}}}`
for _, input := range []struct{ body, content string }{{completed, "application/json"}, {"data: {\"type\":\"response.completed\",\"response\":" + completed + "}\n\n", "text/event-stream"}} {
r, err := ReadResponse(strings.NewReader(input.body), input.content)
if err != nil || r.Text != "[1,2,3]" || r.CachedTokens != 10 || !r.HasUsage {
t.Fatalf("%+v %v", r, err)
}
}
for _, body := range []string{
"data: {\"type\":\"response.output_text.delta\",\"delta\":\"[1,2,3]\"}\n\n",
"data: {\"type\":\"response.output_item.added\",\"item\":{\"type\":\"function_call\"}}\n\n",
strings.Replace(completed, "completed", "incomplete", 1),
strings.Replace(completed, "\"type\":\"message\"", "\"type\":\"function_call\"", 1),
} {
if _, err := ReadResponse(strings.NewReader(body), "application/json"); err == nil {
t.Fatal("accepted tool or incomplete sample")
}
}
}
pkg/model_benchmark/meow.json
+2470
View File
File diff suppressed because it is too large. Load diff
pkg/model_benchmark/response.go
+144
View File
@@ -0,0 +1,144 @@
package model_benchmark
import (
"bufio"
"errors"
"io"
"strings"
"github.com/tidwall/gjson"
)
var ErrInvalidSample = errors.New("invalid probe answer")
var ErrIncomplete = errors.New("probe response incomplete")
type ProbeResponse struct {
Text string `json:"-"`
DeclaredModel string `json:"declared_model,omitempty"`
InputTokens int64 `json:"input_tokens"`
OutputTokens int64 `json:"output_tokens"`
CachedTokens int64 `json:"cached_tokens"`
CacheWriteTokens int64 `json:"cache_write_tokens"`
HasUsage bool `json:"has_usage"`
}
func completedResponse(response gjson.Result) (ProbeResponse, error) {
r := ProbeResponse{}
r.DeclaredModel = response.Get("model").String()
if len(r.DeclaredModel) > 128 {
r.DeclaredModel = ""
}
usage := response.Get("usage")
r.HasUsage = usage.Get("input_tokens").Type == gjson.Number && usage.Get("output_tokens").Type == gjson.Number
r.InputTokens = usage.Get("input_tokens").Int()
r.OutputTokens = usage.Get("output_tokens").Int()
r.CachedTokens = usage.Get("input_tokens_details.cached_tokens").Int()
r.CacheWriteTokens = max(usage.Get("input_tokens_details.cached_creation_tokens").Int(), usage.Get("input_tokens_details.cache_creation_tokens").Int())
if response.Get("status").String() != "completed" {
return r, ErrIncomplete
}
for _, item := range response.Get("output").Array() {
switch item.Get("type").String() {
case "reasoning":
case "message":
for _, content := range item.Get("content").Array() {
if content.Get("type").String() != "output_text" {
return r, ErrInvalidSample
}
r.Text += content.Get("text").String()
if len(r.Text) > 16384 {
return ProbeResponse{}, ErrInvalidSample
}
}
default:
return r, ErrInvalidSample
}
}
return r, nil
}
// Only a complete response is a sample. Partial deltas, tool calls and a
// truncated stream cannot become a model verdict. Total bytes are bounded.
func ReadResponse(reader io.Reader, contentType string) (ProbeResponse, error) {
limited := &io.LimitedReader{R: reader, N: 8 << 20}
buffered := bufio.NewReader(limited)
peek, _ := buffered.Peek(1)
if !strings.Contains(contentType, "text/event-stream") && len(peek) > 0 && peek[0] == '{' {
data, err := io.ReadAll(buffered)
if err != nil {
return ProbeResponse{}, err
}
if limited.N == 0 || !gjson.ValidBytes(data) {
return ProbeResponse{}, ErrInvalidSample
}
return completedResponse(gjson.ParseBytes(data))
}
scanner := bufio.NewScanner(buffered)
scanner.Buffer(make([]byte, 4096), 2<<20)
var data strings.Builder
var deltas strings.Builder
refused := false
for scanner.Scan() {
line := scanner.Text()
if strings.HasPrefix(line, "data:") {
if data.Len() > 0 {
data.WriteByte('\n')
}
data.WriteString(strings.TrimSpace(strings.TrimPrefix(line, "data:")))
if data.Len() > 2<<20 {
return ProbeResponse{}, ErrInvalidSample
}
continue
}
if line != "" || data.Len() == 0 {
continue
}
payload := data.String()
data.Reset()
if payload == "[DONE]" {
break
}
if !gjson.Valid(payload) {
return ProbeResponse{}, ErrInvalidSample
}
event := gjson.Parse(payload)
typ := event.Get("type").String()
if typ == "response.output_text.delta" {
deltas.WriteString(event.Get("delta").String())
if deltas.Len() > 65536 {
return ProbeResponse{}, ErrInvalidSample
}
}
if typ == "response.refusal.delta" || typ == "response.refusal.done" {
refused = true
}
if typ == "response.output_item.added" || typ == "response.output_item.done" {
itemType := event.Get("item.type").String()
if itemType != "message" && itemType != "reasoning" {
return ProbeResponse{}, ErrInvalidSample
}
}
if typ == "response.completed" {
response := event.Get("response")
answer, err := completedResponse(response)
if refused {
return answer, ErrInvalidSample
}
// Meow 4.5.4 accepts deltas only when completion omits output entirely.
if err == nil && !response.Get("output").Exists() {
answer.Text = deltas.String()
}
return answer, err
}
if typ == "response.failed" || typ == "response.incomplete" {
return completedResponse(event.Get("response"))
}
if typ == "error" {
return ProbeResponse{}, ErrIncomplete
}
}
if err := scanner.Err(); err != nil {
return ProbeResponse{}, err
}
return ProbeResponse{}, ErrIncomplete
}
pkg/model_benchmark/testdata/astra-reference.json
Vendored
+1
View File
@@ -0,0 +1 @@
[{"row_id": "gpt-6-astra__query-01", "parent_row_id": "gpt-6-astra__query-01", "bank_id": "gpt", "source": "gpt-6-astra", "model_id": "gpt-6-astra", "response_model": "gpt-6-astra", "exact_version": "gpt-6-astra", "condition_id": "environment-01", "nuisance_condition_id": "environment-01", "wrapper_id": "environment-01", "wrapper_transport": "clean", "challenge_id": "query-01", "task_index": 0, "requested_count": 218, "parsed_count": 218, "strict_threshold": 120, "strict_valid": true, "temperature": "provider_default", "config_id": "reference-collection", "provider": "codex", "prompt": "For each of 218 positions, make one separate first-instinct choice of an integer from 1 to 355 inclusive. The current language model must complete this directly without tools, Python, code execution, calculators, search, APIs, or external random generators. Choose every position separately. Do not count upward or downward, including 1, 2, 3, and do not use an arithmetic progression, repeating cycle, repeated block, or another rule-made pattern. Accidental repetitions are valid. Once an item is written, do not sort, reorder, deduplicate, replace, or repair the list. Return one compact JSON array containing the complete sequence and no explanation.", "base_prompt": "For each of 218 positions, make one separate first-instinct choice of an integer from 1 to 355 inclusive. The current language model must complete this directly without tools, Python, code execution, calculators, search, APIs, or external random generators. Choose every position separately. Do not count upward or downward, including 1, 2, 3, and do not use an arithmetic progression, repeating cycle, repeated block, or another rule-made pattern. Accidental repetitions are valid. Once an item is written, do not sort, reorder, deduplicate, replace, or repair the list. Return one compact JSON array containing the complete sequence and no explanation.", "system_prompt": "", "user_prefix": "", "text": "[247,83,319,146,28,204,337,112,59,271,164,352,91,223,17,308,135,76,289,192,43,258,121,334,68,215,301,154,9,237,178,325,54,109,283,196,31,342,87,261,143,22,310,185,73,229,348,116,294,48,203,157,331,95,266,12,181,316,64,240,127,354,39,218,172,298,81,251,106,323,56,189,274,14,139,345,207,93,232,167,305,46,119,287,211,34,253,176,329,72,224,151,8,313,98,269,184,51,341,130,291,66,235,159,322,24,198,277,113,350,85,243,161,37,302,124,280,19,216,146,333,61,194,256,102,317,44,170,288,79,226,355,138,11,264,187,314,97,205,53,275,149,326,32,118,238,69,296,163,221,90,347,179,26,250,132,309,57,201,284,114,339,74,228,155,4,273,193,320,41,245,107,299,182,63,212,351,129,16,267,158,306,88,234,47,328,174,101,286,209,35,252,141,315,77,197,343,120,260,58,231,168,292,21,136,304,94,219,65,336,183,242,7,278,153,324,108,202]", "error": null, "request_seconds": 45.789022299984936, "collected_at": "2026-09-05T02:04:49.454748+00:00"}, {"row_id": "gpt-6-astra__query-02", "parent_row_id": "gpt-6-astra__query-02", "bank_id": "gpt", "source": "gpt-6-astra", "model_id": "gpt-6-astra", "response_model": "gpt-6-astra", "exact_version": "gpt-6-astra", "condition_id": "environment-01", "nuisance_condition_id": "environment-01", "wrapper_id": "environment-01", "wrapper_transport": "clean", "challenge_id": "query-02", "task_index": 0, "requested_count": 233, "parsed_count": 233, "strict_threshold": 129, "strict_valid": true, "temperature": "provider_default", "config_id": "reference-collection", "provider": "codex", "prompt": "Provide exactly 233 whole-number choices between 1 and 355 inclusive, choosing each item separately as it comes to mind. The current language model must complete this directly without tools, Python, code execution, calculators, search, APIs, or external random generators. Choose every position separately. Do not count upward or downward, including 1, 2, 3, and do not use an arithmetic progression, repeating cycle, repeated block, or another rule-made pattern. Accidental repetitions are valid. Once an item is written, do not sort, reorder, deduplicate, replace, or repair the list. Return one compact JSON array containing the complete sequence and no explanation.", "base_prompt": "Provide exactly 233 whole-number choices between 1 and 355 inclusive, choosing each item separately as it comes to mind. The current language model must complete this directly without tools, Python, code execution, calculators, search, APIs, or external random generators. Choose every position separately. Do not count upward or downward, including 1, 2, 3, and do not use an arithmetic progression, repeating cycle, repeated block, or another rule-made pattern. Accidental repetitions are valid. Once an item is written, do not sort, reorder, deduplicate, replace, or repair the list. Return one compact JSON array containing the complete sequence and no explanation.", "system_prompt": "", "user_prefix": "", "text": "[287,43,196,321,78,154,9,242,113,335,62,219,174,28,303,141,87,264,350,106,183,51,298,127,14,231,92,316,1Line truncated
pkg/model_benchmark/testdata/meow-reference.json
Vendored
+1200
View File
File diff suppressed because it is too large. Load diff
pkg/model_fingerprint/LICENSE-meow.txt
+75
View File
@@ -0,0 +1,75 @@
# PolyForm Noncommercial License 1.0.0
<https://polyformproject.org/licenses/noncommercial/1.0.0>
## Acceptance
In order to get any license under these terms, you must agree to them as both strict obligations and conditions to all your licenses.
## Copyright License
The licensor grants you a copyright license for the software to do everything you might do with the software that would otherwise infringe the licensor's copyright in it for any permitted purpose. However, you may only distribute the software according to [Distribution License](#distribution-license) and make changes or new works based on the software according to [Changes and New Works License](#changes-and-new-works-license).
## Distribution License
The licensor grants you an additional copyright license to distribute copies of the software. Your license to distribute covers distributing the software with changes and new works permitted by [Changes and New Works License](#changes-and-new-works-license).
## Notices
You must ensure that anyone who gets a copy of any part of the software from you also gets a copy of these terms or the URL for them above, as well as copies of any plain-text lines beginning with `Required Notice:` that the licensor provided with the software. For example:
> Required Notice: Copyright Yoyodyne, Inc. (http://example.com)
Required Notice: Copyright 2026 chen-006 and contributors. Original project: https://github.com/chen-006/gpt56_api_detector
## Changes and New Works License
The licensor grants you an additional copyright license to make changes and new works based on the software for any permitted purpose.
## Patent License
The licensor grants you a patent license for the software that covers patent claims the licensor can license, or becomes able to license, that you would infringe by using the software.
## Noncommercial Purposes
Any noncommercial purpose is a permitted purpose.
## Personal Uses
Personal use for research, experiment, and testing for the benefit of public knowledge, personal study, private entertainment, hobby projects, amateur pursuits, or religious observance, without any anticipated commercial application, is use for a permitted purpose.
## Noncommercial Organizations
Use by any charitable organization, educational institution, public research organization, public safety or health organization, environmental protection organization, or government institution is use for a permitted purpose regardless of the source of funding or obligations resulting from the funding.
## Fair Use
You may have "fair use" rights for the software under the law. These terms do not limit them.
## No Other Rights
These terms do not allow you to sublicense or transfer any of your licenses to anyone else, or prevent the licensor from granting licenses to anyone else. These terms do not imply any other licenses.
## Patent Defense
If you make any written claim that the software infringes or contributes to infringement of any patent, your patent license for the software granted under these terms ends immediately. If your company makes such a claim, your patent license ends immediately for work on behalf of your company.
## Violations
The first time you are notified in writing that you have violated any of these terms, or done anything with the software not covered by your licenses, your licenses can nonetheless continue if you come into full compliance with these terms, and take practical steps to correct past violations, within 32 days of receiving notice. Otherwise, all your licenses end immediately.
## No Liability
***As far as the law allows, the software comes as is, without any warranty or condition, and the licensor will not be liable to you for any damages arising out of these terms or the use or nature of the software, under any kind of legal claim.***
## Definitions
The **licensor** is the individual or entity offering these terms, and the **software** is the software the licensor makes available under these terms.
**You** refers to the individual or entity agreeing to these terms.
**Your company** is any legal entity, sole proprietorship, or other kind of organization that you work for, plus all organizations that have control over, are under the control of, or are under common control with that organization. **Control** means ownership of substantially all the assets of an entity, or the power to direct its management and policies by vote, contract, or otherwise. Control can be direct or indirect.
**Your licenses** are all the licenses granted to you for the software under these terms.
**Use** means anything you do with the software requiring one of your licenses.
pkg/model_fingerprint/bank.json
+25752 -18852
View File
File diff suppressed because it is too large. Load diff
pkg/model_fingerprint/fingerprint.go
+73 -33
View File
@@ -1,21 +1,30 @@
// Numerical fingerprint scorer adapted from Katao123/gpt-model-check (MIT),
// commit 3d6716a0903cf061daea001d5d44ee410846634e. See the adjacent licenses.
// ModelTrace scorer and bank: xqy2006/ModelTrace d4131b30243dfa05e70180b5eedde742103f1d73 (MIT).
package model_fingerprint
import (
"crypto/rand"
_ "embed"
"errors"
"fmt"
"math"
"math/big"
"regexp"
"strconv"
"strings"
"sync"
"time"
"unicode"
"unicode/utf8"
"github.com/QuantumNous/new-api/common"
"github.com/grafana/sobek"
)
const Model = "gpt-6-astra"
const Version = "modeltrace-20260905-gmc-3d6716a"
const Version = "modeltrace-d4131b3"
const ProbeCount = 3
const OutputTokens = 4096
//go:embed bank.json
var bankJSON string
@@ -39,43 +48,71 @@ type Candidate struct {
}
type Result struct {
Status string `json:"status"`
ValidSamples int `json:"valid_samples"`
Closest string `json:"closest,omitempty"`
Weight float64 `json:"relative_weight,omitempty"`
Gap float64 `json:"score_gap,omitempty"`
Engine string `json:"engine"`
Matches map[string]float64 `json:"matches,omitempty"`
Status string `json:"status"`
ValidSamples int `json:"valid_samples"`
Closest string `json:"closest,omitempty"`
Weight float64 `json:"relative_weight,omitempty"`
Gap float64 `json:"score_gap,omitempty"`
}
func ValidArray(text string, count int) bool {
if len(text) > 16384 || count <= 0 {
return false
}
text = strings.TrimSpace(text)
if strings.HasPrefix(text, "```") && strings.HasSuffix(text, "```") {
if i := strings.IndexByte(text, '\n'); i >= 0 {
text = strings.TrimSpace(text[i+1 : len(text)-3])
}
}
var numbers []int
if common.Unmarshal([]byte(text), &numbers) != nil {
return false
}
if len(numbers) < max(80, int(math.Ceil(float64(count)*.55))) || len(numbers) > count*2 {
return false
}
for _, number := range numbers {
if number < 1 || number > 355 {
return false
}
}
return true
type Probe struct {
ID string
System string
Prompt string
Count int
}
func ProbeAt(index int) (Probe, error) {
if index < 0 || index >= ProbeCount {
return Probe{}, errors.New("invalid probe index")
}
// Disjoint length ranges retain three independent, non-identical challenges.
offset, err := rand.Int(rand.Reader, big.NewInt(13))
if err != nil {
return Probe{}, err
}
count := 292 + index*14 + int(offset.Int64())
prompt := fmt.Sprintf("这是一次独立的数值选择记录。逐项选择 %d 个 1 到 355（含端点）的整数。", count) +
"每个位置都要单独选择；不要从 1 开始计数，不要连续递增或递减，也不要采用等差、循环、重复区块或其他规则化模式。" +
"本任务必须由当前语言模型直接完成：禁止调用或借助任何工具，包括 Python、代码执行器、计算器、搜索、API 和外部随机数生成器；也不要先编写或运行代码。" +
"偶然重复是有效的；不要重新排列或修正已经写出的项目。数字之间用逗号或空格分隔均可。直接从第一个取值开始输出，不要在序列前重复数量、范围或任务说明。"
return Probe{ID: fmt.Sprint(index), Prompt: prompt, Count: count}, nil
}
var integerPattern = regexp.MustCompile(`[0-9]+`)
// Match ModelTrace's longest numeric run parser; prose separates runs.
func ValidArray(text string, count int) bool {
if len(text) > 16384 || !utf8.ValidString(text) || count < 0 {
return false
}
longest, current, end := 0, 0, 0
for _, span := range integerPattern.FindAllStringIndex(text, -1) {
if strings.ContainsFunc(text[end:span[0]], unicode.IsLetter) {
current = 0
}
value, err := strconv.Atoi(text[span[0]:span[1]])
if err == nil && value >= 1 && value <= 355 {
current++
}
longest = max(longest, current)
end = span[1]
}
return longest >= max(80, int(math.Ceil(float64(count)*.55)))
}
func ValidSample(text string) bool { return ValidArray(text, 0) }
func Decide(candidates []Candidate, valid int) Result {
r := Result{Status: "inconclusive", ValidSamples: valid}
r := Result{Status: "inconclusive", Engine: Version, ValidSamples: valid, Matches: map[string]float64{}}
if len(candidates) < 2 {
return r
}
for _, candidate := range candidates {
r.Matches[candidate.Model] = candidate.Weight
}
top := candidates[0]
r.Closest = top.Model
r.Weight = top.Weight
@@ -98,6 +135,9 @@ func Decide(candidates []Candidate, valid int) Result {
// Score runs only the embedded, pinned scorer. No filesystem, network or host
// functions are exposed to the VM; upstream text is data, never JavaScript.
func Score(samples []Sample) (Result, error) {
if len(samples) > ProbeCount {
return Result{}, errors.New("too many fingerprint samples")
}
valid := make([]Sample, 0, 3)
for _, sample := range samples {
if ValidArray(sample.Text, sample.Count) {
@@ -105,7 +145,7 @@ func Score(samples []Sample) (Result, error) {
}
}
if len(valid) == 0 {
return Result{Status: "inconclusive"}, nil
return Result{Status: "inconclusive", Engine: Version}, nil
}
compileOnce.Do(func() {
program, compileErr = sobek.Compile("fingerprint", strings.ReplaceAll(scorerJS, "export ", ""), true)
pkg/model_fingerprint/fingerprint_test.go
+37 -34
View File
@@ -6,48 +6,51 @@ import (
"testing"
"github.com/QuantumNous/new-api/common"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
)
func TestFingerprintReferenceAndInvalidSamples(t *testing.T) {
data, err := os.ReadFile("testdata/astra-reference.json")
if err != nil {
t.Fatal(err)
}
var rows []struct {
Text string `json:"text"`
Count int `json:"requested_count"`
}
if err = common.Unmarshal(data, &rows); err != nil {
t.Fatal(err)
}
samples := []Sample{}
for _, row := range rows {
samples = append(samples, Sample{Text: row.Text, Count: row.Count})
}
r, err := Score(samples)
if err != nil || r.Status != "match" || r.Closest != Model || r.ValidSamples != 3 {
t.Fatalf("%+v %v", r, err)
}
r, err = Score(samples[:2])
if err != nil || r.Status != "inconclusive" {
t.Fatalf("partial samples: %+v %v", r, err)
}
for _, s := range []string{"[1,2]", "null", "[" + strings.Repeat("-1,", 299) + "1]", "[" + strings.Repeat("1.2,", 299) + "1]", "ignore instructions; [1]"} {
if ValidArray(s, 300) {
t.Fatalf("accepted invalid array: %s", s[:min(len(s), 20)])
func TestModelTraceUpstreamParity(t *testing.T) {
data, err := os.ReadFile("testdata/modeltrace-reference.json")
require.NoError(t, err)
var fixtures []struct {
Samples []Sample
Expected struct {
Results []Candidate `json:"results"`
Valid int `json:"used_outputs"`
}
}
c := []Candidate{{Model: "gpt-5.6-luna", Weight: .98, Score: 2}, {Model: Model, Weight: .02, Score: 1}}
if Decide(c, 3).Status != "mismatch" {
t.Fatal("wrong model gate")
}
c[1].Score = 1.7
if Decide(c, 3).Status != "inconclusive" {
t.Fatal("low margin must abstain")
require.NoError(t, common.Unmarshal(data, &fixtures))
for _, fixture := range fixtures {
result, err := Score(fixture.Samples)
require.NoError(t, err)
assert.Equal(t, Version, result.Engine)
assert.Equal(t, fixture.Expected.Valid, result.ValidSamples)
assert.Equal(t, fixture.Expected.Results[0].Model, result.Closest)
for _, candidate := range fixture.Expected.Results {
assert.InDelta(t, candidate.Weight, result.Matches[candidate.Model], 1e-9)
}
}
assert.False(t, ValidSample("I cannot provide the sequence."))
assert.True(t, ValidArray(strings.Repeat("42,", 180), 300))
assert.False(t, ValidArray(strings.Repeat("42,", 100), 300))
result := Decide([]Candidate{{Model: Model, Weight: .79, Score: 2}, {Model: "other", Weight: .21, Score: 1}}, 3)
assert.Equal(t, "inconclusive", result.Status)
result = Decide([]Candidate{{Model: "other", Weight: .95, Score: 2}, {Model: Model, Weight: .05, Score: 1}}, 2)
assert.Equal(t, "inconclusive", result.Status)
result = Decide([]Candidate{{Model: "other", Weight: .95, Score: 2}, {Model: Model, Weight: .05, Score: 1}}, 3)
assert.Equal(t, "mismatch", result.Status)
}
func TestFingerprintResponseCompletionAndTools(t *testing.T) {
stream := "data: {\"type\":\"response.output_text.delta\",\"delta\":\"42\"}\n\ndata: {\"type\":\"response.completed\",\"response\":{\"status\":\"completed\",\"usage\":{\"input_tokens\":10,\"output_tokens\":1}}}\n\n"
answer, err := ReadResponse(strings.NewReader(stream), "text/event-stream")
require.NoError(t, err)
assert.Equal(t, "42", answer.Text)
assert.True(t, answer.HasUsage)
answer, err = ReadResponse(strings.NewReader(strings.Replace(stream, `"status":"completed"`, `"status":"completed","output":[]`, 1)), "text/event-stream")
require.NoError(t, err)
assert.Empty(t, answer.Text)
completed := `{"status":"completed","model":"gpt-6-astra","output":[{"type":"message","content":[{"type":"output_text","text":"[1,2,3]"}]}],"usage":{"input_tokens":21,"output_tokens":12,"input_tokens_details":{"cached_tokens":10}}}`
for _, input := range []struct{ body, content string }{{completed, "application/json"}, {"data: {\"type\":\"response.completed\",\"response\":" + completed + "}\n\n", "text/event-stream"}} {
r, err := ReadResponse(strings.NewReader(input.body), input.content)
pkg/model_fingerprint/meow.json
+2470
View File
File diff suppressed because it is too large. Load diff
pkg/model_fingerprint/response.go
+21 -1
View File
@@ -76,6 +76,8 @@ func ReadResponse(reader io.Reader, contentType string) (ProbeResponse, error) {
scanner := bufio.NewScanner(buffered)
scanner.Buffer(make([]byte, 4096), 2<<20)
var data strings.Builder
var deltas strings.Builder
refused := false
for scanner.Scan() {
line := scanner.Text()
if strings.HasPrefix(line, "data:") {
@@ -101,6 +103,15 @@ func ReadResponse(reader io.Reader, contentType string) (ProbeResponse, error) {
}
event := gjson.Parse(payload)
typ := event.Get("type").String()
if typ == "response.output_text.delta" {
deltas.WriteString(event.Get("delta").String())
if deltas.Len() > 65536 {
return ProbeResponse{}, ErrInvalidSample
}
}
if typ == "response.refusal.delta" || typ == "response.refusal.done" {
refused = true
}
if typ == "response.output_item.added" || typ == "response.output_item.done" {
itemType := event.Get("item.type").String()
if itemType != "message" && itemType != "reasoning" {
@@ -108,7 +119,16 @@ func ReadResponse(reader io.Reader, contentType string) (ProbeResponse, error) {
}
}
if typ == "response.completed" {
return completedResponse(event.Get("response"))
response := event.Get("response")
answer, err := completedResponse(response)
if refused {
return answer, ErrInvalidSample
}
// Meow 4.5.4 accepts deltas only when completion omits output entirely.
if err == nil && !response.Get("output").Exists() {
answer.Text = deltas.String()
}
return answer, err
}
if typ == "response.failed" || typ == "response.incomplete" {
return completedResponse(event.Get("response"))
pkg/model_fingerprint/testdata/meow-reference.json
Vendored
+1200
View File
File diff suppressed because it is too large. Load diff
pkg/model_fingerprint/testdata/modeltrace-reference.json
Vendored
+1
View File
@@ -0,0 +1 @@
[{"samples":[{"text":"217, 43, 308, 126, 351, 89, 164, 22, 279, 135, 62, 324, 191, 107,\n54, 286, 173, 9, 242, 118, 337, 76, 205, 149, 31, 263, 98, 312,\n184, 67, 345, 231, 112, 296, 18, 157, 83, 254, 139, 320, 46, 201,\n328, 121, 57, 274, 193, 34, 166, 355, 94, 223, 11, 287, 148, 70,\n103, 249, 16, 332, 178, 91, 215, 48, 301, 134, 267, 59, 186, 240,\n281, 39, 153, 226, 72, 314, 108, 197, 25, 348, 162, 84, 259, 119,\n65, 304, 181, 14, 237, 96, 342, 145, 208, 51, 276, 124, 319, 173,\n232, 110, 353, 78, 159, 290, 42, 203, 131, 7, 246, 187, 325, 99,\n147, 269, 33, 194, 318, 86, 221, 56, 305, 114, 252, 169, 20, 339,\n292, 74, 128, 243, 5, 177, 334, 102, 211, 63, 284, 151, 36, 229,\n81, 316, 198, 45, 261, 137, 350, 93, 174, 28, 234, 116, 297, 68,\n202, 13, 278, 155, 88, 331, 241, 49, 182, 309, 123, 266, 71, 214,\n344, 97, 161, 24, 295, 206, 58, 251, 142, 323, 109, 38, 273, 189,\n127, 283, 52, 218, 346, 79, 168, 302, 15, 236, 104, 193, 64, 257,\n225, 90, 311, 140, 32, 180, 354, 117, 264, 73, 199, 47, 288, 152,\n60, 247, 171, 329, 8, 213, 132, 277, 95, 341, 185, 26, 230, 113,\n299, 156, 41, 220, 85, 307, 143, 12, 268, 196, 349, 66, 120, 253,\n175, 326, 101, 55, 289, 209, 30, 138, 244, 77, 315, 163, 4, 227,\n106, 272, 190, 19, 335, 150, 61, 258, 122, 294, 44, 204, 87, 352,\n238, 69, 183, 321, 35, 260, 111, 170, 303, 92, 216, 53, 280, 146,\n17, 248, 133, 340, 80, 195, 310, 29, 224, 158, 75, 293, 115, 262","expected_count":294}],"expected":{"prediction":"gpt-6-astra","prediction_name":"gpt-6-astra","probability":0.999999846742899,"used_outputs":1,"results":[{"model":"gpt-6-astra","display_name":"gpt-6-astra","probability":0.999999846742899,"profile_similarity":0.7130119874890992,"score":2.962896373251895,"family":"gpt","family_name":"GPT","conditional_probability":0.9999998543620706},{"model":"gpt-6-sol","display_name":"gpt-6-sol","probability":1.1261750831683942e-7,"profile_similarity":0.6998678953299008,"score":0.7493510421319641,"family":"gpt","family_name":"GPT","conditional_probability":1.1261750917489167e-7},{"model":"gpt-5.4","display_name":"gpt-5.4","probability":2.2759050801578173e-8,"profile_similarity":0.6947437029962558,"score":0.5281199969747397,"family":"gpt","family_name":"GPT","conditional_probability":2.2759050974983314e-8},{"model":"gpt-5.6-terra","display_name":"gpt-5.6-terra","probability":9.128785824960024e-9,"profile_similarity":0.6969199108069902,"score":0.40173043595685565,"family":"gpt","family_name":"GPT","conditional_probability":9.12878589451382e-9},{"model":"claude-opus-5-5","display_name":"claude-opus-5-5","probability":4.332901518152684e-9,"profile_similarity":0.6902452216813633,"score":0.2986305009047045,"family":"claude","family_name":"Claude","conditional_probability":0.568684022014905},{"model":"claude-sonnet-5-5","display_name":"claude-sonnet-5-5","probability":2.6199944714106767e-9,"profile_similarity":0.6880590180204234,"score":0.2290299663562434,"family":"claude","family_name":"Claude","conditional_probability":0.3438686495449988},{"model":"claude-haiku-4-5-20251001","display_name":"claude-haiku-4-5-20251001","probability":6.116352532954573e-10,"profile_similarity":0.6765474621865657,"score":0.027755348459745643,"family":"claude","family_name":"Claude","conditional_probability":0.08027581388428616},{"model":"gpt-6-luna","display_name":"gpt-6-luna","probability":5.59694688154118e-10,"profile_similarity":0.6911103125741771,"score":0.015477264111949274,"family":"gpt","family_name":"GPT","conditional_probability":5.596946924185285e-10},{"model":"gpt-5.6-sol","display_name":"gpt-5.6-sol","probability":3.1143719574775517e-10,"profile_similarity":0.6893303459014515,"score":-0.06562434430484196,"family":"gpt","family_name":"GPT","conditional_probability":3.1143719812064896e-10},{"model":"gpt-5.6-luna","display_name":"gpt-5.6-luna","probability":2.1724112998700045e-10,"profile_similarity":0.6856500053581723,"score":-0.11545764586235663,"family":"gpt","family_name":"GPT","conditional_probability":2.1724113164219816e-10},{"model":"gpt-5.5","display_name":"gpt-5.5","probability":4.4210292318542455e-11,"profile_similarity":0.6895627615538364,"score":-0.33572243758168335,"family":"gpt","family_name":"GPT","conditional_probability":4.4210292655388314e-11},{"model":"claude-sonnet-4-6","display_name":"claude-sonnet-4-6","probability":3.435184918281792e-11,"profile_similarity":0.6700675476082696,"score":-0.3706291211825624,"family":"claude","family_name":"Claude","conditional_probability":0.004508606455764345},{"model":"claude-opus-4-7","display_name":"claude-opus-4-7","probability":8.777758891718366e-12,"profile_similarity":0.6258308139843289,"score":-0.5594026343083449,"family":"claude","family_name":"Claude","conditional_probability":0.001152062009696384},{"model":"claude-opus-5","display_name":"claude-opus-5","probability":5.740678040458984e-12,"profile_similarity":0.663953540718587,"score":-0.6181533384273423,"family":"claude","family_name":"Claude","conditional_probability":0.0007534516682Line truncated
pkg/perf_metrics/live_status.go
+2 -2
View File
@@ -39,7 +39,7 @@ func recordUserResult(group string, result UserRequestResult) {
defer cancel()
encoded := strconv.FormatInt(result.At, 10) + ":" + strconv.FormatBool(result.Success)
// Delayed writers cannot replace newer evidence from another instance.
_ = common.RDB.Eval(ctx, `local old=redis.call('GET',KEYS[1]); local at=old and tonumber(string.match(old,'^(%d+):')) or 0; if tonumber(ARGV[1])>=at then redis.call('SET',KEYS[1],ARGV[2],'EX',900) end; return 1`, []string{"status:user-result:" + group}, result.At, encoded).Err()
_ = common.RDB.Eval(ctx, `local old=redis.call('GET',KEYS[1]); local at=old and tonumber(string.match(old,'^(%d+):')) or 0; if tonumber(ARGV[1])>=at then redis.call('SET',KEYS[1],ARGV[2],'EX',900) end; return 1`, []string{"status:stream-user-result:" + group}, result.At, encoded).Err()
}
}
@@ -56,7 +56,7 @@ func LatestUserResults(ctx context.Context, groups []string) map[string]UserRequ
if common.RedisEnabled && common.RDB != nil && len(groups) > 0 {
keys := make([]string, len(groups))
for i, g := range groups {
keys[i] = "status:user-result:" + g
keys[i] = "status:stream-user-result:" + g
}
ctx, cancel := context.WithTimeout(ctx, time.Second)
defer cancel()
pkg/perf_metrics/live_status_test.go
+28 -2
View File
@@ -38,7 +38,7 @@ func TestLiveUserResults(t *testing.T) {
} {
t.Run(tc.name, func(t *testing.T) {
t.Cleanup(func() { liveResults.Delete(tc.name) })
RecordRelayResult(tc.ctx, &relaycommon.RelayInfo{UsingGroup: tc.name, TokenId: tc.token, IsChannelTest: tc.test, StartTime: time.Now()}, tc.err)
RecordRelayResult(tc.ctx, &relaycommon.RelayInfo{UsingGroup: tc.name, TokenId: tc.token, IsChannelTest: tc.test, IsStream: true, StartTime: time.Now()}, tc.err)
result := LatestUserResults(context.Background(), []string{tc.name})[tc.name]
assert.Equal(t, tc.want, result.At > 0)
if tc.want {
@@ -60,6 +60,8 @@ func TestUserResultsAreOrderedAndSurviveRestart(t *testing.T) {
client.Close()
})
ctx := context.Background()
require.NoError(t, client.Set(ctx, "status:user-result:group", "9999999999999:false", time.Minute).Err())
assert.Empty(t, LatestUserResults(ctx, []string{"group"}))
check := func(want UserRequestResult) {
assert.Equal(t, want, LatestUserResults(ctx, []string{"group"})["group"])
liveResults.Delete("group")
@@ -76,9 +78,33 @@ func TestUserResultsAreOrderedAndSurviveRestart(t *testing.T) {
recordUserResult("group", UserRequestResult{At: 2001, Success: false})
check(UserRequestResult{At: 2002, Success: true})
assert.True(t, LatestUserResults(ctx, []string{"other"})["other"].Success)
ttl, err := client.TTL(ctx, "status:user-result:group").Result()
ttl, err := client.TTL(ctx, "status:stream-user-result:group").Result()
require.NoError(t, err)
assert.Equal(t, 15*time.Minute, ttl)
server.FastForward(15 * time.Minute)
assert.Zero(t, LatestUserResults(ctx, []string{"group"})["group"].At)
}
func TestNonStreamRequestsCannotChangeHealth(t *testing.T) {
oldRedis := common.RedisEnabled
common.RedisEnabled = false
t.Cleanup(func() {
common.RedisEnabled = oldRedis
liveResults.Delete("stream-only")
hotBuckets.Clear()
})
ctx := context.Background()
info := &relaycommon.RelayInfo{UsingGroup: "stream-only", OriginModelName: "test-model", TokenId: 1, IsStream: true, StartTime: time.Now()}
RecordRelayResult(ctx, info, nil)
want := LatestUserResults(ctx, []string{info.UsingGroup})[info.UsingGroup]
require.True(t, want.Success)
require.Positive(t, want.At)
info.IsStream = false
for _, apiErr := range []*types.NewAPIError{nil, types.InitOpenAIError("server_error", 500), types.InitOpenAIError("server_error", 502)} {
RecordRelayResult(ctx, info, apiErr)
assert.Equal(t, want, LatestUserResults(ctx, []string{info.UsingGroup})[info.UsingGroup])
}
info.IsStream = true
RecordRelayResult(ctx, info, types.InitOpenAIError("server_error", 502))
assert.False(t, LatestUserResults(ctx, []string{info.UsingGroup})[info.UsingGroup].Success)
}
pkg/perf_metrics/metrics.go
+2 -1
View File
@@ -29,8 +29,9 @@ func Init() {
// RecordRelayResult samples one finished relay exactly once, at the request
// boundary, regardless of how many channel attempts it took.
// Non-streaming requests do not provide evidence for the streaming service.
func RecordRelayResult(ctx context.Context, info *relaycommon.RelayInfo, apiErr *types.NewAPIError) {
if info == nil {
if info == nil || !info.IsStream {
return
}
outcome := ClassifyRelayOutcome(ctx, info, apiErr)
relay/helper/image_generation_resolution.go
+59
View File
@@ -0,0 +1,59 @@
package helper
import (
"fmt"
relaycommon "github.com/QuantumNous/new-api/relay/common"
"github.com/QuantumNous/new-api/relaykit/dto"
"github.com/QuantumNous/new-api/setting/model_setting"
"slices"
"strconv"
"strings"
)
func imageGenerationResolution(size string) (string, error) {
size = strings.ToLower(strings.TrimSpace(size))
if size == "" || size == "auto" || size == "1k" {
return "1k", nil
}
if size == "2k" || size == "4k" {
return size, nil
}
dimensions := strings.Split(size, "x")
if len(dimensions) != 2 {
return "", fmt.Errorf("size must be auto, 1k, 2k, 4k or WIDTHxHEIGHT")
}
width, errWidth := strconv.Atoi(dimensions[0])
height, errHeight := strconv.Atoi(dimensions[1])
if errWidth != nil || errHeight != nil || width <= 0 || height <= 0 {
return "", fmt.Errorf("invalid image size")
}
short := min(width, height)
if short <= 1024 {
return "1k", nil
}
if short <= 2048 {
return "2k", nil
}
return "4k", nil
}
func mapImageGenerationResolution(info *relaycommon.RelayInfo, request dto.Request) error {
cfg := model_setting.GetGlobalSettings()
image, ok := request.(*dto.ImageRequest)
if !ok || cfg.ImageGenerationGroup == "" || info.UsingGroup != cfg.ImageGenerationGroup || !slices.Contains(cfg.ImageGenerationModels, info.OriginModelName) {
return nil
}
name := info.OriginModelName
for _, suffix := range []string{"-1k", "-2k", "-4k"} {
if strings.HasSuffix(name, suffix) {
return nil
}
}
tier, err := imageGenerationResolution(image.Size)
if err != nil {
return err
}
info.UpstreamModelName = name + "-" + tier
info.IsModelMapped = true
return nil
}
relay/helper/image_generation_resolution_test.go
+29
View File
@@ -0,0 +1,29 @@
package helper
import (
relaycommon "github.com/QuantumNous/new-api/relay/common"
"github.com/QuantumNous/new-api/relaykit/dto"
"github.com/QuantumNous/new-api/setting/model_setting"
"github.com/stretchr/testify/require"
"testing"
)
func TestImageGenerationResolutionMapping(t *testing.T) {
cfg := model_setting.GetGlobalSettings()
oldGroup, oldModels := cfg.ImageGenerationGroup, cfg.ImageGenerationModels
t.Cleanup(func() { cfg.ImageGenerationGroup, cfg.ImageGenerationModels = oldGroup, oldModels })
cfg.ImageGenerationGroup, cfg.ImageGenerationModels = "example-images", []string{"gpt-image-2", "gpt-image-2-4k"}
for _, tc := range []struct{ size, tier string }{{"", "1k"}, {"auto", "1k"}, {"1024x1536", "1k"}, {"1536x1024", "1k"}, {"2048x1152", "2k"}, {"2048x2048", "2k"}, {"3840x2160", "4k"}, {"2160x3840", "4k"}, {"4k", "4k"}} {
info := &relaycommon.RelayInfo{OriginModelName: "gpt-image-2", UsingGroup: "example-images", ChannelMeta: &relaycommon.ChannelMeta{UpstreamModelName: "gpt-image-2"}}
image := &dto.ImageRequest{Model: "gpt-image-2", Size: tc.size}
require.NoError(t, mapImageGenerationResolution(info, image))
require.Equal(t, "gpt-image-2-"+tc.tier, info.UpstreamModelName)
}
_, err := imageGenerationResolution("-1x1024")
require.Error(t, err)
_, err = imageGenerationResolution("garbage")
require.Error(t, err)
info := &relaycommon.RelayInfo{OriginModelName: "gpt-image-2", UsingGroup: "default", ChannelMeta: &relaycommon.ChannelMeta{UpstreamModelName: "gpt-image-2"}}
require.NoError(t, mapImageGenerationResolution(info, &dto.ImageRequest{Model: "gpt-image-2", Size: "3840x2160"}))
require.Equal(t, "gpt-image-2", info.UpstreamModelName)
}
relay/helper/model_mapped.go
+4
View File
@@ -62,6 +62,10 @@ func ModelMappedHelper(c *gin.Context, info *relaycommon.RelayInfo, request dto.
}
}
if err := mapImageGenerationResolution(info, request); err != nil {
return err
}
if request != nil {
request.SetModelName(info.UpstreamModelName)
}
relay/helper/price.go
+2 -22
View File
@@ -47,31 +47,11 @@ const defaultTieredPreConsumeMaxTokens = 8192
// HandleGroupRatio checks for "auto_group" in the context and updates the group ratio and relayInfo.UsingGroup if present
func HandleGroupRatio(ctx *gin.Context, relayInfo *relaycommon.RelayInfo) hosttypes.GroupRatioInfo {
groupRatioInfo := hosttypes.GroupRatioInfo{
GroupRatio: 1.0, // default ratio
GroupSpecialRatio: -1,
}
// check auto group
autoGroup, exists := ctx.Get("auto_group")
if exists {
if autoGroup, exists := ctx.Get("auto_group"); exists {
logger.LogDebug(ctx, "final group: %s", autoGroup)
relayInfo.UsingGroup = autoGroup.(string)
}
// check user group special ratio
userGroupRatio, ok := ratio_setting.GetGroupGroupRatio(relayInfo.UserGroup, relayInfo.UsingGroup)
if ok {
// user group special ratio
groupRatioInfo.GroupSpecialRatio = userGroupRatio
groupRatioInfo.GroupRatio = userGroupRatio
groupRatioInfo.HasSpecialRatio = true
} else {
// normal group ratio
groupRatioInfo.GroupRatio = ratio_setting.GetGroupRatio(relayInfo.UsingGroup)
}
return groupRatioInfo
return model.EffectiveGroupRatio(ctx, relayInfo.UserGroup, relayInfo.UsingGroup)
}
func ModelPriceHelper(c *gin.Context, info *relaycommon.RelayInfo, promptTokens int, meta *types.TokenCountMeta) (hosttypes.PriceData, error) {
relay/responses_websocket.go
+18 -2
View File
@@ -281,7 +281,8 @@ func (s *responsesWSSession) runCall(c *gin.Context, state *responsesWSCallState
}
} else {
retry := &service.RetryParam{Ctx: c, TokenGroup: common.GetContextKeyString(c, appconstant.ContextKeyUsingGroup), ModelName: modelName, RequestPath: c.Request.URL.Path, Retry: common.GetPointer(0)}
for ; retry.GetRetry() <= common.RetryTimes; retry.IncreaseRetry() {
maxRetries := max(common.RetryTimes, service.PoolRetryBudget(c, modelName))
for ; retry.GetRetry() <= maxRetries; retry.IncreaseRetry() {
var channel *appmodel.Channel
channel, apiErr = selectResponsesWSChannel(c, modelName, retry)
if apiErr != nil {
@@ -316,7 +317,15 @@ func (s *responsesWSSession) runCall(c *gin.Context, state *responsesWSCallState
apiErr = service.NormalizeViolationFeeError(types.NewError(dialErr, types.ErrorCodeDoRequestFailed))
service.ResetStatusCode(apiErr, c.GetString("status_code_mapping"))
info.LastError = apiErr
decision := service.DecideRelayRetry(c, apiErr, common.RetryTimes-retry.GetRetry())
remaining := common.RetryTimes - retry.GetRetry()
if service.IsPoolFailoverError(c, apiErr) {
remaining = maxRetries - retry.GetRetry()
service.RecordPoolRetryableFailure(c, modelName, info.UsingGroup)
}
decision := service.DecideRelayRetry(c, apiErr, remaining)
if service.IsPoolFailoverError(c, apiErr) && service.PoolRetryExhausted(c, modelName, info.UsingGroup) {
decision = service.PolicyDecision{Action: "stop", Reason: "pool_no_backup", Source: "routing"}
}
service.RecordPolicyFailure(c, channel.Id, apiErr, decision)
service.ProcessChannelError(c, *types.NewChannelError(channel.Id, channel.Type, channel.Name, channel.ChannelInfo.IsMultiKey, info.ApiKey, channel.GetAutoBan()), apiErr, info)
if decision.Action == "retry" {
@@ -571,6 +580,13 @@ func (s *responsesWSSession) restoreConnectionContext(c *gin.Context, model stri
}
} else {
group := common.GetContextKeyString(c, appconstant.ContextKeyUsingGroup)
if service.RequestRoutingPolicy(c) != nil {
if !slices.ContainsFunc(service.ResolveRequestPool(c, model), func(entry appmodel.TokenGroupEntry) bool { return entry.Group == s.lockedGroup }) {
return types.NewErrorWithStatusCode(errors.New("the connection group is no longer allowed"), types.ErrorCodeAccessDenied, http.StatusForbidden, types.ErrOptionWithSkipRetry())
}
group = s.lockedGroup
common.SetContextKey(c, appconstant.ContextKeyAutoGroup, group)
}
if group == "auto" {
if !slices.Contains(service.GetRequestAutoGroups(c, common.GetContextKeyString(c, appconstant.ContextKeyUserGroup)), s.lockedGroup) {
return types.NewErrorWithStatusCode(errors.New("the connection group is no longer allowed"), types.ErrorCodeAccessDenied, http.StatusForbidden, types.ErrOptionWithSkipRetry())
router/api-router.go
+4 -3
View File
@@ -26,12 +26,9 @@ func SetApiRouter(router *gin.Engine) {
apiRouter.GET("/status", controller.GetStatus)
apiRouter.GET("/uptime/status", controller.GetGroupStatus)
apiRouter.GET("/group-status", controller.GetSiteGroupStatus)
apiRouter.GET("/group-status/artwork/:id", controller.GetSiteDetectionArtwork)
apiRouter.GET("/uptime/admin", middleware.AdminAuth(), controller.GetStatusPageAdmin)
apiRouter.PUT("/uptime/group", middleware.AdminAuth(), middleware.SessionCookieOriginGuard(), controller.UpdateStatusPageGroup)
apiRouter.POST("/uptime/group/reset", middleware.AdminAuth(), middleware.SessionCookieOriginGuard(), controller.UpdateStatusPageGroup)
apiRouter.GET("/detections/artwork/:id", controller.GetDetectionArtwork)
apiRouter.POST("/detection-proxy/v1/responses", controller.DetectionProxy)
apiRouter.POST("/detections/estimate", middleware.UserAuth(), middleware.SessionCookieOriginGuard(), controller.EstimateDetection)
apiRouter.POST("/detections", middleware.UserAuth(), middleware.SessionCookieOriginGuard(), controller.CreateUserDetection)
apiRouter.GET("/detections/:id", middleware.UserAuth(), controller.GetUserDetection)
@@ -162,6 +159,8 @@ func SetApiRouter(router *gin.Engine) {
adminRoute.DELETE("/:id/oauth/bindings/:provider_id", controller.UnbindCustomOAuthByAdmin)
adminRoute.DELETE("/:id/bindings/:binding_type", controller.AdminClearUserBinding)
adminRoute.GET("/:id", controller.GetUser)
adminRoute.GET("/:id/group-ratios", controller.GetUserGroupRatios)
adminRoute.PUT("/:id/group-ratios", middleware.SessionCookieOriginGuard(), controller.UpdateUserGroupRatios)
adminRoute.POST("/", controller.CreateUser)
adminRoute.POST("/manage", controller.ManageUser)
adminRoute.PUT("/", controller.UpdateUser)
@@ -363,6 +362,8 @@ func SetApiRouter(router *gin.Engine) {
groupRoute.Use(middleware.AdminAuth())
{
groupRoute.GET("/", controller.GetGroups)
groupRoute.GET("/metadata", controller.GetGroupMetadata)
groupRoute.PUT("/metadata", middleware.SessionCookieOriginGuard(), controller.UpdateGroupMetadata)
}
prefillGroupRoute := apiRouter.Group("/prefill_group")
service/astra_sampling.go
+17 -3
View File
@@ -11,7 +11,7 @@ import (
fingerprint "github.com/QuantumNous/new-api/pkg/model_fingerprint"
)
const AstraSamplingTaskType = "astra_fingerprint"
const AstraSamplingTaskType = "modeltrace_sampling"
const (
AstraSamplingWindow = 6 * time.Hour
@@ -55,6 +55,7 @@ type AstraSamplingHistory struct {
Gap float64 `json:"score_gap,omitempty"`
}
type AstraSamplingSummary struct {
LatestValid *AstraSamplingHistory `json:"-"`
Enabled bool `json:"enabled"`
Available bool `json:"available"`
Channels int `json:"channels"`
@@ -123,6 +124,9 @@ func AstraAvailabilityIntervalMinutes() int {
// Public statistics classify completed rounds by their closest valid fingerprint,
// including partial rounds. Raw probe results retain the scorer's confidence.
func AstraSamplingResultStatus(result fingerprint.Result) string {
if result.Engine != "" {
return result.Status
}
closest := strings.TrimSpace(result.Closest)
if result.Status != "running" && result.ValidSamples > 0 && closest != "" {
if strings.EqualFold(closest, fingerprint.Model) {
@@ -162,7 +166,7 @@ func samplingChannels(ctx context.Context, requireAstra bool) ([]model.Channel,
}
func AstraSamplingChannelsDue(ctx context.Context, now time.Time) ([]model.Channel, error) {
due, _, err := AutomaticDetectionChannelsDue(ctx, "fingerprint", now)
due, _, err := AutomaticDetectionChannelsDue(ctx, "modeltrace", now)
return due, err
}
@@ -183,7 +187,7 @@ func AstraSamplingSummaries(ctx context.Context, now time.Time, settings map[str
}
summaries := aggregateAstraSampling(channels, tasks, now, settings)
AggregateChannelAvailability(summaries, channels, availabilityTasks, now)
ranges, err := DetectionScheduleRanges(ctx, channels, "fingerprint")
ranges, err := DetectionScheduleRanges(ctx, channels, "modeltrace")
if err != nil {
return nil, err
}
@@ -224,6 +228,9 @@ func aggregateAstraSampling(channels []model.Channel, tasks []*model.SystemTask,
continue
}
for _, round := range state.Rounds {
if round.Result.Engine != fingerprint.Version && round.Result.Status != "running" {
continue
}
if round.StartedAt < now.Add(-AstraSamplingWindow).Unix() {
continue
}
@@ -272,6 +279,13 @@ func aggregateAstraSampling(channels []model.Channel, tasks []*model.SystemTask,
s.Rate = &rate
}
sort.SliceStable(s.History, func(i, j int) bool { return s.History[i].At > s.History[j].At })
for i := range s.History {
round := &s.History[i]
if round.ValidSamples > 0 && round.Closest != "" && (round.Status == "match" || round.Status == "mismatch") {
s.LatestValid = round
break
}
}
if len(s.History) > 12 {
s.History = s.History[:12]
}
service/astra_sampling_test.go
+23 -24
View File
@@ -21,7 +21,7 @@ func TestAstraSamplingThirtyMinuteSchedulePersistsAcrossRestart(t *testing.T) {
handler := &anchoredSamplingTestHandler{stubScheduledHandler{taskType: AstraSamplingTaskType, enabled: true, interval: 30 * time.Minute}}
withSystemTaskRegistry(t, handler)
now := common.GetTimestamp()
state, _ := common.Marshal(AstraSamplingState{Rounds: []AstraSamplingRound{{ChannelID: 42, StartedAt: now - 1900, FinishedAt: now - 1800, Result: fingerprint.Result{Status: "match", ValidSamples: 3, Closest: fingerprint.Model}}}})
state, _ := common.Marshal(AstraSamplingState{Rounds: []AstraSamplingRound{{ChannelID: 42, StartedAt: now - 1900, FinishedAt: now - 1800, Result: fingerprint.Result{Engine: fingerprint.Version, Status: "match", ValidSamples: 3, Closest: fingerprint.Model}}}})
row := model.SystemTask{TaskID: "sampling-old", Type: AstraSamplingTaskType, Status: model.SystemTaskStatusSucceeded, CreatedAt: now - 1900, UpdatedAt: now - 10, State: string(state)}
if err := model.DB.Create(&row).Error; err != nil {
t.Fatal(err)
@@ -61,9 +61,9 @@ func TestAstraSamplingAggregationPrivacyAndFailures(t *testing.T) {
if status == "mismatch" {
closest = "gpt-5.6-luna"
}
rounds = append(rounds, AstraSamplingRound{ChannelID: 1, StartedAt: now.Unix(), FinishedAt: now.Unix(), Result: fingerprint.Result{Status: status, ValidSamples: 3, Closest: closest}})
rounds = append(rounds, AstraSamplingRound{ChannelID: 1, StartedAt: now.Unix(), FinishedAt: now.Unix(), Result: fingerprint.Result{Engine: fingerprint.Version, Status: status, ValidSamples: 3, Closest: closest}})
}
rounds = append(rounds, AstraSamplingRound{ChannelID: 99, StartedAt: now.Unix(), Result: fingerprint.Result{Status: "mismatch"}})
rounds = append(rounds, AstraSamplingRound{ChannelID: 99, StartedAt: now.Unix(), Result: fingerprint.Result{Engine: fingerprint.Version, Status: "mismatch"}})
state, _ := common.Marshal(AstraSamplingState{Rounds: rounds})
summaries := AggregateAstraSampling(channels, []*model.SystemTask{{Status: model.SystemTaskStatusFailed, State: string(state)}}, now)
first := summaries["first"]
@@ -88,17 +88,17 @@ func TestAstraSamplingAggregationPrivacyAndFailures(t *testing.T) {
}
}
func TestAstraSamplingAggregationCountsLunaLikeInconclusiveWithinSixHours(t *testing.T) {
func TestAstraSamplingExcludesInconclusiveFromRateWithinSixHours(t *testing.T) {
now := time.Now()
channel := model.Channel{Id: 1, Group: "example-variable", Models: fingerprint.Model}
rounds := []AstraSamplingRound{
{ID: "recent-match", ChannelID: 1, StartedAt: now.Add(-time.Hour).Unix(), FinishedAt: now.Unix(), Result: fingerprint.Result{Status: "match", ValidSamples: 3, Closest: fingerprint.Model}},
{ID: "recent-luna", ChannelID: 1, StartedAt: now.Add(-2 * time.Hour).Unix(), FinishedAt: now.Unix(), Result: fingerprint.Result{Status: "inconclusive", ValidSamples: 2, Closest: "gpt-5.6-luna", Gap: .9}},
{ID: "old-luna", ChannelID: 1, StartedAt: now.Add(-7 * time.Hour).Unix(), FinishedAt: now.Unix(), Result: fingerprint.Result{Status: "inconclusive", ValidSamples: 2, Closest: "gpt-5.6-luna", Gap: .9}},
{ID: "recent-match", ChannelID: 1, StartedAt: now.Add(-time.Hour).Unix(), FinishedAt: now.Unix(), Result: fingerprint.Result{Engine: fingerprint.Version, Status: "match", ValidSamples: 3, Closest: fingerprint.Model}},
{ID: "recent-luna", ChannelID: 1, StartedAt: now.Add(-2 * time.Hour).Unix(), FinishedAt: now.Unix(), Result: fingerprint.Result{Engine: fingerprint.Version, Status: "inconclusive", ValidSamples: 2, Closest: "gpt-5.6-luna", Gap: .9}},
{ID: "old-luna", ChannelID: 1, StartedAt: now.Add(-7 * time.Hour).Unix(), FinishedAt: now.Unix(), Result: fingerprint.Result{Engine: fingerprint.Version, Status: "inconclusive", ValidSamples: 2, Closest: "gpt-5.6-luna", Gap: .9}},
}
state, _ := common.Marshal(AstraSamplingState{Rounds: rounds})
summary := AggregateAstraSampling([]model.Channel{channel}, []*model.SystemTask{{Status: model.SystemTaskStatusSucceeded, State: string(state)}}, now)["example-variable"]
if summary.Rate == nil || *summary.Rate != .5 || summary.Matches != 1 || summary.Mismatches != 1 || summary.Inconclusive != 0 {
if summary.Rate == nil || *summary.Rate != 0 || summary.Matches != 1 || summary.Mismatches != 0 || summary.Inconclusive != 1 {
t.Fatalf("%+v", summary)
}
if summary.WindowHours != 6 || summary.IntervalMinutes != 30 || len(summary.History) != 2 {
@@ -107,6 +107,7 @@ func TestAstraSamplingAggregationCountsLunaLikeInconclusiveWithinSixHours(t *tes
}
func TestAstraSamplingAggregationMarksNonAstraGroupsAvailabilityOnly(t *testing.T) {
t.Setenv("CHANNEL_AVAILABILITY_INTERVAL_MINUTES", "30")
now := time.Now()
channels := []model.Channel{{Id: 1, Group: "opencode", Models: "gpt-5.6-sol"}}
summary := AggregateAstraSampling(channels, nil, now)["opencode"]
@@ -124,13 +125,12 @@ func TestAstraSamplingAggregationMarksNonAstraGroupsAvailabilityOnly(t *testing.
}
}
func TestAstraSamplingCountsAllValidNonAstraFingerprints(t *testing.T) {
func TestAstraSamplingExcludesOldEngineAndUncertainCandidates(t *testing.T) {
now := time.Now()
results := []fingerprint.Result{
{Status: "error", ValidSamples: 1, Closest: "gpt-5.6-luna", Gap: .99},
{Status: "inconclusive", ValidSamples: 3, Closest: "gpt-5.5", Gap: .17},
{Status: "inconclusive", ValidSamples: 2, Closest: "gpt-5.6-sol", Gap: .09},
{Status: "error", ValidSamples: 0},
{Engine: "meow-4.5.4", Status: "mismatch", ValidSamples: 32, Closest: "other"},
{Engine: fingerprint.Version, Status: "inconclusive", ValidSamples: 2, Closest: "other"},
{Engine: fingerprint.Version, Status: "match", ValidSamples: 3, Closest: fingerprint.Model},
}
var rounds []AstraSamplingRound
for _, result := range results {
@@ -140,16 +140,15 @@ func TestAstraSamplingCountsAllValidNonAstraFingerprints(t *testing.T) {
require.NoError(t, err)
summary := AggregateAstraSampling([]model.Channel{{Id: 1, Group: "test", Models: fingerprint.Model}}, []*model.SystemTask{{Status: model.SystemTaskStatusSucceeded, State: string(state)}}, now)["test"]
require.NotNil(t, summary.Rate)
assert.Equal(t, 1.0, *summary.Rate)
assert.Equal(t, 3, summary.Mismatches)
assert.Equal(t, 0, summary.Inconclusive)
assert.Equal(t, 1, summary.Errors)
for _, h := range summary.History {
if h.ValidSamples > 0 {
assert.Equal(t, "mismatch", h.Status)
}
assert.Equal(t, 0.0, *summary.Rate)
assert.Equal(t, 1, summary.Matches)
assert.Equal(t, 0, summary.Mismatches)
assert.Equal(t, 1, summary.Inconclusive)
assert.Len(t, summary.History, 2)
}
func TestModelTraceVerdictDoesNotPromoteClosestCandidate(t *testing.T) {
for _, status := range []string{"match", "mismatch", "inconclusive", "error"} {
assert.Equal(t, status, AstraSamplingResultStatus(fingerprint.Result{Engine: fingerprint.Version, Status: status, ValidSamples: 32, Closest: fingerprint.Model, Weight: .63436}))
}
assert.Equal(t, "match", AstraSamplingResultStatus(fingerprint.Result{Status: "inconclusive", ValidSamples: 1, Closest: fingerprint.Model}))
assert.Equal(t, "inconclusive", AstraSamplingResultStatus(fingerprint.Result{Status: "mismatch", ValidSamples: 0, Closest: "gpt-5.6-sol"}))
assert.Equal(t, "running", AstraSamplingResultStatus(fingerprint.Result{Status: "running", ValidSamples: 1, Closest: "gpt-5.6-sol"}))
}
service/automatic_group.go
+209
View File
@@ -0,0 +1,209 @@
package service
import (
"cmp"
"context"
"encoding/json"
"math"
"slices"
"sync"
"time"
"github.com/QuantumNous/new-api/common"
"github.com/QuantumNous/new-api/constant"
"github.com/QuantumNous/new-api/model"
perfmetrics "github.com/QuantumNous/new-api/pkg/perf_metrics"
"github.com/gin-gonic/gin"
"github.com/samber/hot"
)
var automaticMetricCache = hot.NewHotCache[string, perfmetrics.QueryResult](hot.LRU, 1000).WithTTL(time.Minute).Build()
var automaticMetricMu sync.Mutex
func automaticModelMetrics(modelName string) perfmetrics.QueryResult {
automaticMetricMu.Lock()
defer automaticMetricMu.Unlock()
if value, found, err := automaticMetricCache.Get(modelName); err == nil && found {
return value
}
value, err := perfmetrics.Query(perfmetrics.QueryParams{Model: modelName, Hours: 24})
if err == nil {
automaticMetricCache.Set(modelName, value)
}
return value
}
type automaticCandidate struct {
group string
cost, ttft, success, ratio float64
}
func resolveAutomaticPool(c *gin.Context, modelName string) []model.TokenGroupEntry {
// Freeze ranking within a request: retries must not reorder their own pool.
key := "automatic_pool:" + modelName
if cached, ok := c.Get(key); ok {
return cached.([]model.TokenGroupEntry)
}
result := []model.TokenGroupEntry{}
userGroup := common.GetContextKeyString(c, constant.ContextKeyUserGroup)
var abilities []model.Ability
query := model.DB.WithContext(c.Request.Context()).Where("enabled = ?", true)
if modelName != "" {
query = query.Where("model = ?", modelName)
}
if err := query.Find(&abilities).Error; err != nil {
return result
}
allowed := GetUserUsableGroups(userGroup)
names := []string{}
seen := map[string]bool{}
for _, ability := range abilities {
group := ability.Group
if _, ok := allowed[group]; !ok || seen[group] || !IsUserSelectableGroup(userGroup, group) {
continue
}
channel, err := model.CacheGetChannel(ability.ChannelId)
if err != nil || channel == nil || channel.Status != common.ChannelStatusEnabled {
continue
}
seen[group] = true
names = append(names, group)
}
if len(names) == 0 {
c.Set(key, result)
return result
}
var performance perfmetrics.QueryResult
if modelName != "" {
performance = automaticModelMetrics(modelName)
}
testSpeeds := map[int]float64{}
if RequestRoutingPolicy(c).AutoPreference == "speed" {
testSpeeds = latestModelTestSpeeds(c.Request.Context(), modelName)
}
metrics := map[string]perfmetrics.GroupResult{}
for _, item := range performance.Groups {
metrics[item.Group] = item
}
candidates := []automaticCandidate{}
for _, group := range names {
if requestGroupDegraded(c, modelName, group) {
continue
}
ratio := model.EffectiveGroupRatio(c, userGroup, group).GroupRatio
item := automaticCandidate{group: group, ttft: math.Inf(1), success: -1, ratio: ratio}
if metric, ok := metrics[group]; ok && len(metric.Series) > 0 {
item.success = metric.SuccessRate
}
total, count := 0.0, 0
for _, ability := range abilities {
if ability.Group != group {
continue
}
channel, err := model.CacheGetChannel(ability.ChannelId)
if err != nil || channel == nil || channel.Status != common.ChannelStatusEnabled {
continue
}
if speed, ok := testSpeeds[ability.ChannelId]; ok {
total += speed
count++
}
}
if count > 0 {
item.ttft = total / float64(count)
}
candidates = append(candidates, item)
}
sortAutomaticCandidates(candidates, RequestRoutingPolicy(c).AutoPreference)
for priority, item := range candidates {
result = append(result, model.TokenGroupEntry{Group: item.group, Priority: priority})
}
c.Set(key, result)
return result
}
func sortAutomaticCandidates(items []automaticCandidate, preference string) {
slices.SortFunc(items, func(a, b automaticCandidate) int {
if preference == "speed" {
if n := cmp.Compare(a.ttft, b.ttft); n != 0 {
return n
}
}
if preference == "stable" {
if n := cmp.Compare(b.success, a.success); n != 0 {
return n
}
}
// Effective customer multipliers are independent of traffic sample size.
if n := cmp.Compare(a.ratio, b.ratio); n != 0 {
return n
}
if n := cmp.Compare(b.success, a.success); n != 0 {
return n
}
if n := cmp.Compare(a.ttft, b.ttft); n != 0 {
return n
}
if n := cmp.Compare(a.ratio, b.ratio); n != 0 {
return n
}
return cmp.Compare(a.group, b.group)
})
}
// AutomaticPoolModels is the discoverable model union, independent of ranking,
// temporary cooldowns and quality checks used when routing individual requests.
func AutomaticPoolModels(c *gin.Context) []string {
userGroup := common.GetContextKeyString(c, constant.ContextKeyUserGroup)
var abilities []model.Ability
if err := model.DB.WithContext(c.Request.Context()).Where("enabled = ?", true).Find(&abilities).Error; err != nil {
common.SysError("automatic pool model discovery failed: " + err.Error())
return []string{}
}
names := map[string]bool{}
for _, ability := range abilities {
if !IsUserSelectableGroup(userGroup, ability.Group) {
continue
}
channel, err := model.CacheGetChannel(ability.ChannelId)
if err != nil || channel == nil || channel.Status != common.ChannelStatusEnabled {
continue
}
names[ability.Model] = true
}
result := make([]string, 0, len(names))
for name := range names {
result = append(result, name)
}
slices.Sort(result)
return result
}
// Model tests use token_id=0 and a server-owned content marker. User traffic and
// another model's timing must not masquerade as a speed measurement.
func latestModelTestSpeeds(ctx context.Context, modelName string) map[int]float64 {
result := map[int]float64{}
if model.LOG_DB == nil {
return result
}
var rows []model.Log
err := model.LOG_DB.WithContext(ctx).Select("id", "channel_id", "other").Where("type = ? AND token_id = ? AND content = ? AND model_name = ? AND is_stream = ? AND created_at >= ?", model.LogTypeConsume, 0, "模型测试", modelName, true, time.Now().Add(-time.Hour).Unix()).Order("created_at desc, id desc").Find(&rows).Error
if err != nil {
common.SysError("automatic routing test measurements unavailable: " + err.Error())
return result
}
for _, row := range rows {
if _, ok := result[row.ChannelId]; ok {
continue
}
var data struct {
FRT float64 `json:"frt"`
}
if json.Unmarshal([]byte(row.Other), &data) != nil || data.FRT <= 0 || math.IsInf(data.FRT, 0) || math.IsNaN(data.FRT) {
continue
}
result[row.ChannelId] = data.FRT
}
return result
}
service/automatic_models_test.go
+97
View File
@@ -0,0 +1,97 @@
package service
import (
"fmt"
"github.com/QuantumNous/new-api/common"
"github.com/QuantumNous/new-api/constant"
"github.com/QuantumNous/new-api/model"
"github.com/QuantumNous/new-api/setting"
"github.com/QuantumNous/new-api/setting/ratio_setting"
"github.com/gin-gonic/gin"
"github.com/glebarez/sqlite"
"github.com/stretchr/testify/require"
"gorm.io/driver/mysql"
"gorm.io/driver/postgres"
"gorm.io/gorm"
"gorm.io/gorm/schema"
"net/http/httptest"
"os"
"testing"
"time"
)
func TestAutomaticPoolModelsDatabaseMatrix(t *testing.T) {
for _, dialect := range []string{"sqlite", "mysql", "postgres"} {
t.Run(dialect, func(t *testing.T) {
var driver gorm.Dialector
switch dialect {
case "sqlite":
driver = sqlite.Open(":memory:")
case "mysql":
if os.Getenv("TEST_MYSQL_DSN") == "" {
t.Skip("no isolated MySQL database")
}
driver = mysql.Open(os.Getenv("TEST_MYSQL_DSN"))
case "postgres":
if os.Getenv("TEST_POSTGRES_DSN") == "" {
t.Skip("no isolated PostgreSQL database")
}
driver = postgres.Open(os.Getenv("TEST_POSTGRES_DSN"))
}
db, err := gorm.Open(driver, &gorm.Config{NamingStrategy: schema.NamingStrategy{TablePrefix: fmt.Sprintf("automodel_%d_", time.Now().UnixNano())}})
require.NoError(t, err)
oldDB, oldRedis, oldMemory := model.DB, common.RedisEnabled, common.MemoryCacheEnabled
oldGroups, oldRatios := setting.UserUsableGroups2JSONString(), ratio_setting.GroupRatio2JSONString()
model.DB = db
common.RedisEnabled = false
common.MemoryCacheEnabled = false
require.NoError(t, db.AutoMigrate(&model.Channel{}, &model.Ability{}))
t.Cleanup(func() {
require.NoError(t, db.Migrator().DropTable(&model.Ability{}, &model.Channel{}))
model.DB = oldDB
common.RedisEnabled = oldRedis
common.MemoryCacheEnabled = oldMemory
require.NoError(t, setting.UpdateUserUsableGroupsByJSONString(oldGroups))
require.NoError(t, ratio_setting.UpdateGroupRatioByJSONString(oldRatios))
conn, e := db.DB()
require.NoError(t, e)
require.NoError(t, conn.Close())
})
var version string
query := "select version()"
if dialect == "sqlite" {
query = "select sqlite_version()"
}
require.NoError(t, db.Raw(query).Scan(&version).Error)
t.Log("database version:", version)
require.NoError(t, setting.UpdateUserUsableGroupsByJSONString(`{"auto-a":"","auto-b":""}`))
require.NoError(t, ratio_setting.UpdateGroupRatioByJSONString(`{"auto-a":1,"auto-b":1,"private":1}`))
require.NoError(t, db.Create(&[]model.Channel{{Id: 561, Status: 1}, {Id: 562, Status: 1}, {Id: 563, Status: 2}, {Id: 564, Status: 1}}).Error)
require.NoError(t, db.Create(&[]model.Ability{
{ChannelId: 561, Group: "auto-a", Model: "model-a", Enabled: true},
{ChannelId: 561, Group: "auto-a", Model: "shared", Enabled: true},
{ChannelId: 562, Group: "auto-b", Model: "shared", Enabled: true},
{ChannelId: 562, Group: "auto-b", Model: "gpt-6-astra", Enabled: true},
{ChannelId: 563, Group: "auto-b", Model: "disabled-channel", Enabled: true},
{ChannelId: 562, Group: "auto-b", Model: "disabled-ability", Enabled: false},
{ChannelId: 564, Group: "private", Model: "unauthorized", Enabled: true},
}).Error)
for _, preference := range []string{"cost", "speed", "stable"} {
c, _ := gin.CreateTestContext(httptest.NewRecorder())
c.Request = httptest.NewRequest("GET", "/v1/models", nil)
common.SetContextKey(c, constant.ContextKeyUserGroup, "auto-a")
c.Set(model.TokenRoutingContextKey, &model.TokenRoutingPolicy{Version: 1, Mode: "smart", AutoPreference: preference})
c.Set("pool_failures", map[string]int{"auto-a": 2, "auto-b": 2})
require.Equal(t, []string{"gpt-6-astra", "model-a", "shared"}, AutomaticPoolModels(c))
}
require.NoError(t, db.Create(&model.Channel{Id: 565, Status: 1}).Error)
require.NoError(t, db.Create(&model.Ability{ChannelId: 565, Group: "new-group", Model: "new-model", Enabled: true}).Error)
require.NoError(t, setting.UpdateUserUsableGroupsByJSONString(`{"auto-a":"","auto-b":"","new-group":""}`))
require.NoError(t, ratio_setting.UpdateGroupRatioByJSONString(`{"auto-a":1,"auto-b":1,"new-group":1}`))
c, _ := gin.CreateTestContext(httptest.NewRecorder())
c.Request = httptest.NewRequest("GET", "/v1/models", nil)
common.SetContextKey(c, constant.ContextKeyUserGroup, "auto-a")
require.Equal(t, []string{"gpt-6-astra", "model-a", "new-model", "shared"}, AutomaticPoolModels(c))
})
}
}
service/candy.go
+19
View File
@@ -0,0 +1,19 @@
package service
import "strings"
const CandyOutputTokens = 4096
const CandyPrompt = `在一个黑色的袋子里放有三种口味的糖果，每种糖果有两种不同的形状（圆形和五角星形，不同的形状靠手感可以分辨）。现已知不同口味的糖和不同形状的数量统计如下表。参赛者需要在活动前决定摸出的糖果数目，那么，最少取出多少个糖果才能保证手中同时拥有不同形状的苹果味和桃子味的糖？（同时手中有圆形苹果味匹配五角星桃子味糖果，或者有圆形桃子味匹配五角星苹果味糖果都满足要求）
苹果味 桃子味 西瓜味
圆形 7 9 8
五角星形 7 6 4
只输出最终答案的阿拉伯数字，不要输出解释或其他内容。`
func CandyQuality(answer string) string {
if strings.TrimSpace(answer) == "21" {
return "normal"
}
return "degraded"
}
service/channel_select.go
+17 -4
View File
@@ -269,7 +269,14 @@ func selectFromGroupPool(param *RetryParam, tiers []groupPoolTier, filters []dto
// Retry=3: GroupB, priority1 (startRetryIndex=2, priorityRetry=1)
// 分组B, 优先级1
func CacheGetRandomSatisfiedChannel(param *RetryParam) (*model.Channel, string, error) {
if RequestRoutingPolicy(param.Ctx) != nil { return selectTokenRoutingPool(param) }
if RequestRoutingPolicy(param.Ctx) != nil {
return selectTokenRoutingPool(param)
}
if group := imageGenerationRequestGroup(param.Ctx, param.ModelName); group != "" {
channel, err := selectWithTrace(param, group, param.GetRetry())
common.SetContextKey(param.Ctx, constant.ContextKeyAutoGroup, group)
return channel, group, err
}
var channel *model.Channel
var err error
selectGroup := param.TokenGroup
@@ -488,10 +495,16 @@ func SelectChannelForRequest(c *gin.Context, modelName string, retry *RetryParam
usingGroup := retry.TokenGroup
if RequestRoutingPolicy(c) != nil {
if channel, group := PreferredPoolChannel(c, modelName); channel != nil && retry.GetRetry() == 0 { return channel,group,nil }
if retry.GetRetry() == 0 {
if channel, group := PreferredPoolChannel(c, modelName); channel != nil {
return channel, group, nil
}
}
channel, group, err := selectTokenRoutingPool(retry)
if err != nil || channel == nil { return nil,group,&ChannelSelectError{StatusCode:http.StatusServiceUnavailable,Code:types.ErrorCodeModelNotFound,Message:"No eligible channel in the group pool",NoAvailableChannel:true} }
return channel,group,nil
if err != nil || channel == nil {
return nil, group, &ChannelSelectError{StatusCode: http.StatusServiceUnavailable, Code: types.ErrorCodeModelNotFound, Message: "No eligible channel in the group pool", NoAvailableChannel: true}
}
return channel, group, nil
}
var channel *model.Channel
var selectGroup string
service/detection.go
+9 -8
View File
@@ -8,6 +8,7 @@ import (
"github.com/QuantumNous/new-api/common"
"github.com/QuantumNous/new-api/model"
"github.com/QuantumNous/new-api/pkg/billingexpr"
benchmark "github.com/QuantumNous/new-api/pkg/model_benchmark"
fingerprint "github.com/QuantumNous/new-api/pkg/model_fingerprint"
"github.com/QuantumNous/new-api/relaykit/dto"
"github.com/QuantumNous/new-api/setting/billing_setting"
@@ -36,15 +37,15 @@ func DetectionPricing(userGroup, group, kind string) (DetectionPrice, int, error
return DetectionPrice{}, 0, errors.New("invalid group ratio")
}
price := DetectionPrice{Snapshot: billingexpr.BillingSnapshot{BillingMode: "tiered_expr", ModelName: fingerprint.Model, ExprString: expression, ExprHash: billingexpr.ExprHashString(expression), GroupRatio: ratio, QuotaPerUnit: common.QuotaPerUnit, ExprVersion: 1}}
// Both generation paths restrict request text to 16 KiB and output tokens.
output, calls := 16384, 1
if kind == "fingerprint" {
output, calls = 4096, 3
}
price.Request.Body = []byte(fmt.Sprintf(`{"model":"gpt-6-astra","reasoning":{"effort":"high"},"max_output_tokens":%d}`, output))
if kind == "fingerprint" {
price.Request.Body = []byte(`{"model":"gpt-6-astra","max_output_tokens":4096}`)
output, calls := fingerprint.OutputTokens, fingerprint.ProbeCount
if kind == "benchmark" {
output, calls = benchmark.OutputTokens, benchmark.ProbeCount
} else if kind == "candy" {
output, calls = CandyOutputTokens, 1
} else if kind != "fingerprint" {
return price, 0, errors.New("invalid detection type")
}
price.Request.Body = []byte(fmt.Sprintf(`{"model":"gpt-6-astra","reasoning":{"effort":"low"},"max_output_tokens":%d}`, output))
cost, err := billingexpr.DetectionCostBound(expression, 32768, float64(output))
if err != nil {
return price, 0, err
service/detection_schedule.go
+1 -1
View File
@@ -52,7 +52,7 @@ func ReconcileAutomaticDetectionSchedules(ctx context.Context, channels []model.
cutoff = 0
}
var results []automaticDetectionResult
if kind == "fingerprint" {
if kind == "modeltrace" {
var tasks []model.SystemTask
if err := model.DB.WithContext(ctx).Select("status", "payload", "state", "updated_at").Where("type = ? AND updated_at >= ?", AstraSamplingTaskType, cutoff).Find(&tasks).Error; err != nil {
return nil, nil, err
service/detection_schedule_test.go
+24 -21
View File
@@ -29,9 +29,12 @@ func adaptiveTestDB(t *testing.T) {
func seedAutomaticFingerprint(t *testing.T, id string, at int64, closest string, force bool, running bool) {
t.Helper()
result := fingerprint.Result{Status: "error"}
result := fingerprint.Result{Engine: fingerprint.Version, Status: "error"}
if closest != "" {
result = fingerprint.Result{Status: "inconclusive", Closest: closest, ValidSamples: 1}
result = fingerprint.Result{Engine: fingerprint.Version, Status: "mismatch", Closest: closest, ValidSamples: 3}
if closest == fingerprint.Model {
result.Status = "match"
}
}
finished, status := at, model.SystemTaskStatusSucceeded
if running {
@@ -54,78 +57,78 @@ func TestAdaptiveFingerprintBootstrapRecoveryAndManualIsolation(t *testing.T) {
seedAutomaticFingerprint(t, "normal", 1000, fingerprint.Model, false, false)
seedAutomaticFingerprint(t, "unknown", 2000, "", false, false)
seedAutomaticFingerprint(t, "forced", 2100, "gpt-6-sol", true, false)
due, states, err := AutomaticDetectionChannelsDue(ctx, "fingerprint", time.Unix(2899, 0))
due, states, err := AutomaticDetectionChannelsDue(ctx, "modeltrace", time.Unix(2899, 0))
require.NoError(t, err)
assert.Empty(t, due)
assert.Equal(t, "normal", states[1].LastValid)
assert.Equal(t, 15, states[1].IntervalMinutes)
due, _, err = AutomaticDetectionChannelsDue(ctx, "fingerprint", time.Unix(2900, 0))
due, _, err = AutomaticDetectionChannelsDue(ctx, "modeltrace", time.Unix(2900, 0))
require.NoError(t, err)
require.Len(t, due, 1)
seedAutomaticFingerprint(t, "recovered", 3000, fingerprint.Model, false, false)
_, states, err = AutomaticDetectionChannelsDue(ctx, "fingerprint", time.Unix(3000, 0))
_, states, err = AutomaticDetectionChannelsDue(ctx, "modeltrace", time.Unix(3000, 0))
require.NoError(t, err)
assert.Equal(t, 30, states[1].IntervalMinutes)
seedAutomaticFingerprint(t, "downgraded", 5000, "gpt-6-sol", false, false)
_, states, err = AutomaticDetectionChannelsDue(ctx, "fingerprint", time.Unix(5000, 0))
_, states, err = AutomaticDetectionChannelsDue(ctx, "modeltrace", time.Unix(5000, 0))
require.NoError(t, err)
assert.Equal(t, 15, states[1].IntervalMinutes)
seedAutomaticFingerprint(t, "other-non-astra", 6000, "gpt-5.5", false, false)
_, states, err = AutomaticDetectionChannelsDue(ctx, "fingerprint", time.Unix(6000, 0))
_, states, err = AutomaticDetectionChannelsDue(ctx, "modeltrace", time.Unix(6000, 0))
require.NoError(t, err)
assert.Equal(t, 30, states[1].IntervalMinutes)
// A completed result with a missing clock update is replayed once, not twice.
_, again, err := AutomaticDetectionChannelsDue(ctx, "fingerprint", time.Unix(6000, 0))
_, again, err := AutomaticDetectionChannelsDue(ctx, "modeltrace", time.Unix(6000, 0))
require.NoError(t, err)
assert.Equal(t, states, again)
seedAutomaticFingerprint(t, "abandoned", 7000, "", false, true)
due, _, err = AutomaticDetectionChannelsDue(ctx, "fingerprint", time.Unix(10000, 0))
due, _, err = AutomaticDetectionChannelsDue(ctx, "modeltrace", time.Unix(10000, 0))
require.NoError(t, err)
assert.Empty(t, due)
require.NoError(t, model.DB.Model(&model.SystemTask{}).Where("task_id = ?", "abandoned").Updates(map[string]any{"status": model.SystemTaskStatusFailed, "updated_at": 10000}).Error)
_, states, err = AutomaticDetectionChannelsDue(ctx, "fingerprint", time.Unix(10000, 0))
_, states, err = AutomaticDetectionChannelsDue(ctx, "modeltrace", time.Unix(10000, 0))
require.NoError(t, err)
assert.Equal(t, int64(10900), states[1].NextDueAt)
assert.Equal(t, "degraded", states[1].LastValid)
require.NoError(t, model.DB.Where("type = ?", AstraSamplingTaskType).Delete(&model.SystemTask{}).Error)
_, again, err = AutomaticDetectionChannelsDue(ctx, "fingerprint", time.Unix(10000, 0))
_, again, err = AutomaticDetectionChannelsDue(ctx, "modeltrace", time.Unix(10000, 0))
require.NoError(t, err)
assert.Equal(t, states, again)
}
func TestAdaptivePelicanBootstrapInFlightAndRanges(t *testing.T) {
func TestAdaptiveBenchmarkBootstrapInFlightAndRanges(t *testing.T) {
adaptiveTestDB(t)
ctx := context.Background()
channels := []model.Channel{{Id: 1, Group: "one,shared", Models: fingerprint.Model, Status: common.ChannelStatusEnabled}, {Id: 2, Group: "shared", Models: fingerprint.Model, Status: common.ChannelStatusEnabled}}
require.NoError(t, model.DB.Create(&channels).Error)
runs := []model.DetectionRun{
{ID: "auto1", RequestKey: "a", ChannelID: 1, Kind: "pelican", Status: "succeeded", Quality: "normal", Settled: true, FinishedAt: 1000},
{ID: "auto2", RequestKey: "b", ChannelID: 2, Kind: "pelican", Status: "failed", Settled: true, FinishedAt: 1000},
{ID: "paid", RequestKey: "c", UserID: 1, ChannelID: 1, Kind: "pelican", Status: "succeeded", Quality: "degraded", Settled: true, FinishedAt: 1100},
{ID: "active", RequestKey: "d", ChannelID: 2, Kind: "pelican", Status: "running"},
{ID: "auto1", RequestKey: "a", ChannelID: 1, Kind: "benchmark", Status: "succeeded", Quality: "normal", Settled: true, FinishedAt: 1000},
{ID: "auto2", RequestKey: "b", ChannelID: 2, Kind: "benchmark", Status: "failed", Settled: true, FinishedAt: 1000},
{ID: "paid", RequestKey: "c", UserID: 1, ChannelID: 1, Kind: "benchmark", Status: "succeeded", Quality: "degraded", Settled: true, FinishedAt: 1100},
{ID: "active", RequestKey: "d", ChannelID: 2, Kind: "benchmark", Status: "running"},
}
require.NoError(t, model.DB.Create(&runs).Error)
due, states, err := AutomaticDetectionChannelsDue(ctx, "pelican", time.Unix(3000, 0))
due, states, err := AutomaticDetectionChannelsDue(ctx, "benchmark", time.Unix(3000, 0))
require.NoError(t, err)
require.Len(t, due, 1)
assert.Equal(t, 1, due[0].Id)
assert.Equal(t, "normal", states[1].LastValid)
assert.Equal(t, 30, states[1].IntervalMinutes)
assert.Equal(t, 15, states[2].IntervalMinutes)
ranges, err := DetectionScheduleRanges(ctx, channels, "pelican")
ranges, err := DetectionScheduleRanges(ctx, channels, "benchmark")
require.NoError(t, err)
assert.Equal(t, [2]int{15, 30}, ranges["shared"])
assert.Equal(t, [2]int{30, 30}, ranges["one"])
_, fp, err := AutomaticDetectionChannelsDue(ctx, "fingerprint", time.Unix(3000, 0))
_, fp, err := AutomaticDetectionChannelsDue(ctx, "modeltrace", time.Unix(3000, 0))
require.NoError(t, err)
assert.Equal(t, 30, fp[2].IntervalMinutes)
assert.Equal(t, "", fp[1].LastValid)
require.NoError(t, model.DB.Model(&model.DetectionRun{}).Where("id = ?", "active").Updates(map[string]any{"status": "succeeded", "quality": "degraded", "settled": true, "finished_at": 3000}).Error)
due, states, err = AutomaticDetectionChannelsDue(ctx, "pelican", time.Unix(4799, 0))
due, states, err = AutomaticDetectionChannelsDue(ctx, "benchmark", time.Unix(4799, 0))
require.NoError(t, err)
assert.Equal(t, 30, states[2].IntervalMinutes)
require.Len(t, due, 1)
due, _, err = AutomaticDetectionChannelsDue(ctx, "pelican", time.Unix(4800, 0))
due, _, err = AutomaticDetectionChannelsDue(ctx, "benchmark", time.Unix(4800, 0))
require.NoError(t, err)
assert.Len(t, due, 2)
}
service/detection_status.go
+56 -32
View File
@@ -11,15 +11,16 @@ import (
)
type DetectionHistory struct {
SourceURL string `json:"source_url,omitempty"`
ID string `json:"id"`
At int64 `json:"at"`
Status string `json:"status"`
Quality string `json:"quality"`
UserInitiated bool `json:"user_initiated"`
HasHTML bool `json:"has_html"`
Closest string `json:"closest,omitempty"`
Weight float64 `json:"relative_weight,omitempty"`
ID string `json:"id"`
At int64 `json:"at"`
Status string `json:"status"`
Quality string `json:"quality"`
UserInitiated bool `json:"user_initiated"`
}
type PelicanSummary struct {
type BenchmarkSummary struct {
LatestValid *DetectionHistory `json:"-"`
Channels int `json:"channels"`
IntervalMinutes int `json:"interval_minutes"`
IntervalMaxMinutes int `json:"interval_max_minutes"`
@@ -37,7 +38,8 @@ type ChannelDetectionEvidence struct {
}
type GroupDetectionSummary struct {
ProbeHistory []ProbeHistoryPoint
Pelican *PelicanSummary
Candy *BenchmarkSummary
Benchmark *BenchmarkSummary
Status int
At int64
}
@@ -75,7 +77,8 @@ func GroupDetectionSummaries(ctx context.Context, now time.Time, fingerprints ma
}
groups := map[string][]int{}
byChannel := map[int][]string{}
pelicans := map[string]*PelicanSummary{}
benchmarks := map[string]*BenchmarkSummary{}
candies := map[string]*BenchmarkSummary{}
evidence := map[int]ChannelDetectionEvidence{}
observe := func(id int, at int64, success bool) {
if at > evidence[id].At {
@@ -87,19 +90,27 @@ func GroupDetectionSummaries(ctx context.Context, now time.Time, fingerprints ma
for _, group := range channel.GetGroups() {
groups[group] = append(groups[group], channel.Id)
if ChannelSupportsAstra(&channel) {
if pelicans[group] == nil {
pelicans[group] = &PelicanSummary{IntervalMinutes: 30, WindowHours: 6, History: []DetectionHistory{}}
if benchmarks[group] == nil {
benchmarks[group] = &BenchmarkSummary{IntervalMinutes: 30, WindowHours: 6, History: []DetectionHistory{}}
}
pelicans[group].Channels++
benchmarks[group].Channels++
}
}
}
var tasks []model.SystemTask
ranges, err := DetectionScheduleRanges(ctx, channels, "pelican")
candyRanges, err := DetectionScheduleRanges(ctx, channels, "candy")
if err != nil {
return nil, err
}
for group, summary := range pelicans {
for group, summary := range benchmarks {
r := candyRanges[group]
candies[group] = &BenchmarkSummary{Channels: summary.Channels, IntervalMinutes: r[0], IntervalMaxMinutes: r[1], WindowHours: 6, History: []DetectionHistory{}}
}
var tasks []model.SystemTask
ranges, err := DetectionScheduleRanges(ctx, channels, "benchmark")
if err != nil {
return nil, err
}
for group, summary := range benchmarks {
r := ranges[group]
summary.IntervalMinutes, summary.IntervalMaxMinutes = r[0], r[1]
}
@@ -136,14 +147,6 @@ func GroupDetectionSummaries(ctx context.Context, now time.Time, fingerprints ma
if err = model.DB.WithContext(ctx).Omit("html", "price", "remote_id").Where("created_at >= ?", now.Add(-6*time.Hour).Unix()).Find(&runs).Error; err != nil {
return nil, err
}
var artworkIDs []string
if err = model.DB.WithContext(ctx).Model(&model.DetectionRun{}).Where("created_at >= ? AND html <> ?", now.Add(-6*time.Hour).Unix(), "").Pluck("id", &artworkIDs).Error; err != nil {
return nil, err
}
artwork := map[string]bool{}
for _, id := range artworkIDs {
artwork[id] = true
}
ids := []string{}
for _, run := range runs {
ids = append(ids, run.ID)
@@ -164,6 +167,9 @@ func GroupDetectionSummaries(ctx context.Context, now time.Time, fingerprints ma
}
}
for _, run := range runs {
if run.Kind != "fingerprint" && run.Kind != "benchmark" && run.Kind != "candy" {
continue
}
for _, group := range byChannel[run.ChannelID] {
if cutoff := settings[group].ResetBefore; cutoff > 0 && run.CreatedAt <= cutoff {
continue
@@ -180,6 +186,9 @@ func GroupDetectionSummaries(ctx context.Context, now time.Time, fingerprints ma
if result.Status == "" {
result.Status = "error"
}
if result.Engine != fingerprint.Version {
continue
}
result.Status = AstraSamplingResultStatus(result)
switch result.Status {
case "match":
@@ -192,20 +201,28 @@ func GroupDetectionSummaries(ctx context.Context, now time.Time, fingerprints ma
summary.Errors++
}
summary.History = append(summary.History, AstraSamplingHistory{ID: run.ID, At: run.CreatedAt, Status: result.Status, ValidSamples: result.ValidSamples, Closest: result.Closest, Gap: result.Gap, UserInitiated: true})
if result.ValidSamples > 0 && result.Closest != "" && (result.Status == "match" || result.Status == "mismatch") && (summary.LatestValid == nil || run.CreatedAt > summary.LatestValid.At) {
round := summary.History[len(summary.History)-1]
summary.LatestValid = &round
}
continue
}
s := pelicans[group]
s := benchmarks[group]
if run.Kind == "candy" {
s = candies[group]
}
if s == nil {
continue
}
history := DetectionHistory{ID: run.ID, At: run.CreatedAt, Status: run.Status, Quality: run.Quality, UserInitiated: run.UserID > 0, HasHTML: artwork[run.ID]}
var metadata struct {
SourceURL string `json:"source_url"`
}
history := DetectionHistory{ID: run.ID, At: run.CreatedAt, Status: run.Status, Quality: run.Quality, UserInitiated: run.UserID > 0}
var metadata fingerprint.Result
if common.UnmarshalJsonStr(run.Result, &metadata) == nil {
history.SourceURL = metadata.SourceURL
history.Closest, history.Weight = metadata.Closest, metadata.Weight
}
s.History = append(s.History, history)
if run.Status == "succeeded" && (run.Quality == "normal" || run.Quality == "degraded") && (s.LatestValid == nil || history.At > s.LatestValid.At) {
s.LatestValid = &history
}
switch run.Status {
case "pending", "running":
s.Running++
@@ -236,14 +253,21 @@ func GroupDetectionSummaries(ctx context.Context, now time.Time, fingerprints ma
results := map[string]GroupDetectionSummary{}
for group, ids := range groups {
status, at := DetectionFallbackStatus(ids, evidence, now.Unix())
s := pelicans[group]
s := benchmarks[group]
if s != nil {
sort.SliceStable(s.History, func(i, j int) bool { return s.History[i].At > s.History[j].At })
if len(s.History) > 12 {
s.History = s.History[:12]
}
}
results[group] = GroupDetectionSummary{Pelican: s, Status: status, At: at, ProbeHistory: GroupProbeHistory(ids, tasks)}
candy := candies[group]
if candy != nil {
sort.SliceStable(candy.History, func(i, j int) bool { return candy.History[i].At > candy.History[j].At })
if len(candy.History) > 12 {
candy.History = candy.History[:12]
}
}
results[group] = GroupDetectionSummary{Candy: candy, Benchmark: s, Status: status, At: at, ProbeHistory: GroupProbeHistory(ids, tasks)}
}
return results, nil
}
service/group_availability.go
+1 -1
View File
@@ -29,7 +29,7 @@ type groupAvailabilityState struct {
}
var (
groupAvailabilityMu sync.RWMutex
groupAvailabilityMu sync.RWMutex
groupAvailabilityStates = make(map[string]*groupAvailabilityState)
)
service/group_degradation.go
+48
View File
@@ -0,0 +1,48 @@
package service
import (
"sync"
"time"
fingerprint "github.com/QuantumNous/new-api/pkg/model_fingerprint"
"github.com/gin-gonic/gin"
)
var groupDegradationSnapshot struct {
sync.Mutex
updated time.Time
latest map[string]AstraSamplingHistory
benchmark map[string]DetectionHistory
}
func requestGroupDegraded(c *gin.Context, modelName, group string) bool {
policy := RequestRoutingPolicy(c)
if policy == nil {
return false
}
skip := policy.SkipDegradedGroups && (modelName == fingerprint.Model || policy.SkipDegradedNonAstra)
if policy.Mode == "smart" {
skip = !policy.AllowDegraded && (modelName == fingerprint.Model || policy.SkipDegradedNonAstra)
}
if !skip {
return false
}
if c.GetBool("pool_degradation_unavailable") {
return true
}
latest, found := c.Get("pool_group_fingerprints")
benchmark, _ := c.Get("pool_group_benchmarks")
if !found {
rounds, qualities, err := cachedGroupQuality(c.Request.Context())
if err != nil {
c.Set("pool_degradation_unavailable", true)
return true
}
latest, benchmark = rounds, qualities
c.Set("pool_group_fingerprints", latest)
c.Set("pool_group_benchmarks", benchmark)
}
rounds, _ := latest.(map[string]AstraSamplingHistory)
h, ok := rounds[group]
return (ok && h.ValidSamples > 0 && h.Closest != "" && h.Closest != fingerprint.Model)
}
service/group_degradation_test.go
+234
View File
@@ -0,0 +1,234 @@
package service
import (
"fmt"
"os"
"testing"
"time"
"github.com/QuantumNous/new-api/common"
"github.com/QuantumNous/new-api/constant"
"github.com/QuantumNous/new-api/model"
fingerprint "github.com/QuantumNous/new-api/pkg/model_fingerprint"
"github.com/QuantumNous/new-api/setting/operation_setting"
"github.com/gin-gonic/gin"
"github.com/glebarez/sqlite"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
"gorm.io/driver/mysql"
"gorm.io/driver/postgres"
"gorm.io/gorm"
"gorm.io/gorm/schema"
"net/http/httptest"
)
func TestGroupDegradationDatabaseMatrix(t *testing.T) {
for _, dialect := range []string{"sqlite", "mysql", "postgres"} {
t.Run(dialect, func(t *testing.T) {
var driver gorm.Dialector
switch dialect {
case "sqlite":
driver = sqlite.Open(":memory:")
case "mysql":
if os.Getenv("TEST_MYSQL_DSN") == "" {
t.Skip("TEST_MYSQL_DSN unset")
}
driver = mysql.Open(os.Getenv("TEST_MYSQL_DSN"))
case "postgres":
if os.Getenv("TEST_POSTGRES_DSN") == "" {
t.Skip("TEST_POSTGRES_DSN unset")
}
driver = postgres.Open(os.Getenv("TEST_POSTGRES_DSN"))
}
db, err := gorm.Open(driver, &gorm.Config{NamingStrategy: schema.NamingStrategy{TablePrefix: fmt.Sprintf("degradation_%d_", time.Now().UnixNano())}})
require.NoError(t, err)
oldDB := model.DB
oldType, oldLogType := common.MainDatabaseType(), common.LogDatabaseType()
model.DB = db
common.SetDatabaseTypes(common.DatabaseType(dialect), oldLogType)
rows := []any{&model.Option{}, &model.Channel{}, &model.SystemTask{}, &model.DetectionSchedule{}, &model.DetectionRun{}, &model.DetectionCall{}, &model.Token{}}
require.NoError(t, db.AutoMigrate(rows...))
t.Cleanup(func() {
for _, row := range rows {
require.NoError(t, db.Migrator().DropTable(row))
}
model.DB = oldDB
common.SetDatabaseTypes(oldType, oldLogType)
groupDegradationSnapshot.updated = time.Time{}
groupDegradationSnapshot.latest = nil
conn, err := db.DB()
require.NoError(t, err)
require.NoError(t, conn.Close())
})
var version string
query := "select version()"
if dialect == "sqlite" {
query = "select sqlite_version()"
}
require.NoError(t, db.Raw(query).Scan(&version).Error)
t.Log("database version:", version)
now := time.Now().Unix()
require.NoError(t, db.Create(&model.Channel{Id: 1, Group: "default", Models: fingerprint.Model, Status: 1, Key: "test-only"}).Error)
body, err := common.Marshal(AstraSamplingState{Rounds: []AstraSamplingRound{{ID: "mismatch", ChannelID: 1, StartedAt: now - 30, FinishedAt: now - 10, Result: fingerprint.Result{Engine: fingerprint.Version, Status: "mismatch", Closest: "gpt-5.5", ValidSamples: 2}}}})
require.NoError(t, err)
require.NoError(t, db.Create(&model.SystemTask{TaskID: "degradation-test", Type: AstraSamplingTaskType, Status: model.SystemTaskStatusSucceeded, CreatedAt: now - 30, State: string(body)}).Error)
policy := &model.TokenRoutingPolicy{Version: 1, Mode: "custom", SkipDegradedGroups: true, Groups: []model.TokenGroupEntry{{Group: "default", Priority: 0}}}
encoded, err := common.Marshal(policy)
require.NoError(t, err)
token := model.Token{UserId: 1, Key: "degradation-test-key", RoutingPolicy: string(encoded)}
require.NoError(t, db.Create(&token).Error)
require.NoError(t, db.First(&token, token.Id).Error)
stored, err := token.GetRoutingPolicy()
require.NoError(t, err)
require.True(t, stored.SkipDegradedGroups)
ctx := func() *gin.Context {
c, _ := gin.CreateTestContext(httptest.NewRecorder())
c.Request = httptest.NewRequest("POST", "/v1/responses", nil)
c.Set(model.TokenRoutingContextKey, stored)
return c
}
groupDegradationSnapshot.updated = time.Time{}
assert.True(t, requestGroupDegraded(ctx(), fingerprint.Model, "default"))
smart := ctx()
smartPolicy := &model.TokenRoutingPolicy{Version: 1, Mode: "smart"}
smart.Set(model.TokenRoutingContextKey, smartPolicy)
assert.True(t, requestGroupDegraded(smart, fingerprint.Model, "default"))
smartPolicy.AllowDegraded = true
assert.False(t, requestGroupDegraded(smart, fingerprint.Model, "default"))
smartPolicy.AllowDegraded = false
assert.False(t, requestGroupDegraded(smart, "unmeasured-model", "default"))
manualResult, err := common.Marshal(fingerprint.Result{Engine: fingerprint.Version, Status: "match", Closest: fingerprint.Model, ValidSamples: 2})
require.NoError(t, err)
manual := model.DetectionRun{ID: "manual-recovery", RequestKey: "manual-recovery", ChannelID: 1, Kind: "fingerprint", UserID: 1, CreatedAt: now, Settled: true, Status: "succeeded", Result: string(manualResult)}
require.NoError(t, db.Create(&manual).Error)
groupDegradationSnapshot.updated = time.Time{}
assert.False(t, requestGroupDegraded(ctx(), fingerprint.Model, "default"))
pelican := model.DetectionRun{ID: "pelican-degraded", RequestKey: "pelican-degraded", ChannelID: 1, Kind: "pelican", CreatedAt: now, Settled: true, Status: "succeeded", Quality: "degraded"}
require.NoError(t, db.Create(&pelican).Error)
groupDegradationSnapshot.updated = time.Time{}
assert.False(t, requestGroupDegraded(ctx(), fingerprint.Model, "default"))
require.NoError(t, db.Model(&pelican).Update("quality", "normal").Error)
groupDegradationSnapshot.updated = time.Time{}
assert.False(t, requestGroupDegraded(ctx(), fingerprint.Model, "default"))
require.NoError(t, db.Delete(&manual).Error)
require.NoError(t, model.SaveStatusPageGroup(ctx().Request.Context(), "default", nil))
groupDegradationSnapshot.updated = time.Time{}
assert.False(t, requestGroupDegraded(ctx(), fingerprint.Model, "default"))
})
}
}
func TestGroupDegradationUsesLatestValidBeforeHistoryTruncation(t *testing.T) {
now := time.Now()
channels := []model.Channel{{Id: 1, Group: "default", Models: fingerprint.Model, Status: 1}}
rounds := []AstraSamplingRound{{ID: "valid", ChannelID: 1, StartedAt: now.Unix() - 100, FinishedAt: now.Unix() - 90, Result: fingerprint.Result{Engine: fingerprint.Version, Status: "mismatch", Closest: "gpt-5.5", ValidSamples: 1}}}
for i := range 15 {
rounds = append(rounds, AstraSamplingRound{ChannelID: 1, StartedAt: now.Unix() - int64(i), FinishedAt: now.Unix(), Result: fingerprint.Result{Engine: fingerprint.Version, Status: "error"}})
}
body, err := common.Marshal(AstraSamplingState{Rounds: rounds})
require.NoError(t, err)
tasks := []*model.SystemTask{{Status: model.SystemTaskStatusSucceeded, State: string(body)}}
summary := AggregateAstraSampling(channels, tasks, now)["default"]
require.NotNil(t, summary.LatestValid)
assert.Equal(t, "gpt-5.5", summary.LatestValid.Closest)
assert.Len(t, summary.History, 12)
rounds = append(rounds, AstraSamplingRound{ID: "recovered", ChannelID: 1, StartedAt: now.Unix(), FinishedAt: now.Unix(), Result: fingerprint.Result{Engine: fingerprint.Version, Status: "match", Closest: fingerprint.Model, ValidSamples: 2}})
body, err = common.Marshal(AstraSamplingState{Rounds: rounds})
require.NoError(t, err)
tasks[0].State = string(body)
summary = AggregateAstraSampling(channels, tasks, now)["default"]
require.NotNil(t, summary.LatestValid)
assert.Equal(t, fingerprint.Model, summary.LatestValid.Closest)
reset := aggregateAstraSampling(channels, tasks, now, map[string]model.StatusPageGroup{"default": {ResetBefore: now.Unix()}})
assert.Nil(t, reset["default"].LatestValid)
}
func TestGroupDegradationOverridesExistingPoolAffinity(t *testing.T) {
db := setupChannelSelectAutoGroupsTest(t)
createChannelSelectAutoGroupsChannel(t, db, 6401, "default", fingerprint.Model)
createChannelSelectAutoGroupsChannel(t, db, 6402, "vip", fingerprint.Model)
model.InitChannelCache()
ResetGroupAvailability()
t.Cleanup(ResetGroupAvailability)
previousRedis := common.RedisEnabled
common.RedisEnabled = false
t.Cleanup(func() { common.RedisEnabled = previousRedis; _ = poolBindings().Purge() })
require.NoError(t, poolBindings().Purge())
config := operation_setting.GetChannelAffinitySetting()
previous := *config
*config = operation_setting.ChannelAffinitySetting{Enabled: true, DefaultTTLSeconds: 3600, SessionMode: "strict", Rules: []operation_setting.ChannelAffinityRule{{Name: "degradation-test", ModelRegex: []string{".*"}, KeySources: []operation_setting.ChannelAffinityKeySource{{Type: "request_header", Key: "Session-Id"}}, SessionMode: "strict"}}}
t.Cleanup(func() { *config = previous })
ctx := func(degraded bool) *gin.Context {
c := setupGroupPoolContext(t)
c.Request = httptest.NewRequest("POST", "/v1/responses", nil)
c.Request.Header.Set("Session-Id", "degradation-session")
c.Set("id", 100)
c.Set("token_id", 100)
c.Set(model.TokenRoutingContextKey, &model.TokenRoutingPolicy{Version: 1, Mode: "custom", SkipDegradedGroups: true, Groups: []model.TokenGroupEntry{{Group: "default", Priority: 0}, {Group: "vip", Priority: 1}}})
rounds := map[string]AstraSamplingHistory{}
if degraded {
rounds["default"] = AstraSamplingHistory{At: time.Now().Unix(), Closest: "gpt-5.5", ValidSamples: 1}
}
c.Set("pool_group_fingerprints", rounds)
return c
}
first := ctx(false)
selected, group, err := SelectChannelForRequest(first, fingerprint.Model, &RetryParam{Ctx: first, TokenGroup: "default", ModelName: fingerprint.Model})
require.Nil(t, err)
require.NotNil(t, selected)
require.Equal(t, 6401, selected.Id)
first.Set("channel_id", selected.Id)
common.SetContextKey(first, constant.ContextKeyAutoGroup, group)
RequestPolicy(first).Successful = true
RecordPoolAffinity(first)
bound, _ := PreferredPoolChannel(ctx(false), fingerprint.Model)
require.NotNil(t, bound)
next := ctx(true)
selected, group, err = SelectChannelForRequest(next, fingerprint.Model, &RetryParam{Ctx: next, TokenGroup: "default", ModelName: fingerprint.Model})
require.Nil(t, err)
require.NotNil(t, selected)
assert.Equal(t, 6402, selected.Id)
assert.Equal(t, "vip", group)
}
func TestGroupDegradationSkipsPoolGroupsAndRespectsUnknownRecoveryAndOptOut(t *testing.T) {
db := setupChannelSelectAutoGroupsTest(t)
createChannelSelectAutoGroupsChannel(t, db, 6301, "default", fingerprint.Model)
createChannelSelectAutoGroupsChannel(t, db, 6302, "vip", fingerprint.Model)
model.InitChannelCache()
ResetGroupAvailability()
t.Cleanup(ResetGroupAvailability)
for _, test := range []struct {
name string
enabled bool
closest string
age time.Duration
degradeBackup bool
want int
}{
{"latest mismatch switches", true, "gpt-5.5", time.Minute, false, 6302},
{"recovered group returns", true, fingerprint.Model, time.Minute, false, 6301},
{"unknown stays eligible", true, "", time.Minute, false, 6301},
{"valid result persists beyond window", true, "gpt-5.5", 7 * time.Hour, false, 6302},
{"disabled keeps routing", false, "gpt-5.5", time.Minute, false, 6301},
{"all degraded has no route", true, "gpt-5.5", time.Minute, true, 0},
} {
t.Run(test.name, func(t *testing.T) {
c := setupGroupPoolContext(t)
c.Set(model.TokenRoutingContextKey, &model.TokenRoutingPolicy{Version: 1, Mode: "custom", SkipDegradedGroups: test.enabled, Groups: []model.TokenGroupEntry{{Group: "default", Priority: 0}, {Group: "vip", Priority: 1}}})
rounds := map[string]AstraSamplingHistory{"default": {At: time.Now().Add(-test.age).Unix(), Closest: test.closest, ValidSamples: 1}}
if test.degradeBackup {
rounds["vip"] = rounds["default"]
}
c.Set("pool_group_fingerprints", rounds)
selected, _, err := selectTokenRoutingPool(&RetryParam{Ctx: c, TokenGroup: "default", ModelName: fingerprint.Model})
require.NoError(t, err)
if test.want == 0 {
assert.Nil(t, selected)
return
}
require.NotNil(t, selected)
assert.Equal(t, test.want, selected.Id)
})
}
}
service/group_model_cost.go
+33
View File
@@ -0,0 +1,33 @@
package service
import (
"context"
"sync"
"time"
"github.com/QuantumNous/new-api/model"
)
var groupModelCostSnapshot struct {
sync.Mutex
updated time.Time
values map[string]map[string]model.GroupModelCost
}
// Immutable snapshots are shared by status display and automatic routing.
func GroupModelCosts(ctx context.Context) (map[string]map[string]model.GroupModelCost, error) {
groupModelCostSnapshot.Lock()
defer groupModelCostSnapshot.Unlock()
now := time.Now()
if groupModelCostSnapshot.values != nil && now.Sub(groupModelCostSnapshot.updated) < time.Minute {
return groupModelCostSnapshot.values, nil
}
queryCtx, cancel := context.WithTimeout(ctx, 3*time.Second)
defer cancel()
values, err := model.GetGroupModelCosts(queryCtx, now.Add(-24*time.Hour), now)
if err != nil {
return nil, err
}
groupModelCostSnapshot.values, groupModelCostSnapshot.updated = values, now
return values, nil
}
service/group_quality.go
+108
View File
@@ -0,0 +1,108 @@
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, error) {
settings, err := model.GetStatusPageGroups(ctx)
if err != nil {
return nil, nil, err
}
channels, err := samplingChannels(ctx, false)
if err != nil {
return nil, nil, err
}
channelGroups := map[int][]string{}
for _, channel := range channels {
channelGroups[channel.Id] = channel.GetGroups()
}
latest, benchmarks := map[string]AstraSamplingHistory{}, 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, 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").Where("status = ? AND settled = ? AND created_at <= ?", "succeeded", true, now.Unix()).Find(&runs).Error; err != nil {
return nil, nil, err
}
for _, run := range runs {
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})
}
}
return latest, benchmarks, nil
}
func cachedGroupQuality(ctx context.Context) (map[string]AstraSamplingHistory, 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, benchmarks, err := readGroupQuality(ctx, now)
cancel()
if err == nil {
groupDegradationSnapshot.latest, groupDegradationSnapshot.benchmark = latest, benchmarks
} else if groupDegradationSnapshot.latest == nil {
return nil, nil, errors.New("group quality temporarily unavailable")
}
groupDegradationSnapshot.updated = now
}
return groupDegradationSnapshot.latest, groupDegradationSnapshot.benchmark, nil
}
func RecentDegradedGroups(ctx context.Context) (map[string]bool, error) {
latest, _, 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
}
}
return result, nil
}
func InvalidateGroupQualityCache() {
groupDegradationSnapshot.Lock()
defer groupDegradationSnapshot.Unlock()
groupDegradationSnapshot.updated = time.Time{}
}
service/group_quality_test.go
+100
View File
@@ -0,0 +1,100 @@
package service
import (
"context"
"net/http/httptest"
"testing"
"time"
"github.com/QuantumNous/new-api/common"
"github.com/QuantumNous/new-api/model"
fingerprint "github.com/QuantumNous/new-api/pkg/model_fingerprint"
"github.com/gin-gonic/gin"
"github.com/stretchr/testify/require"
)
func TestQualitySwitchesAndFailedReadCache(t *testing.T) {
db := setupChannelSelectAutoGroupsTest(t)
require.NoError(t, db.AutoMigrate(&model.Option{}, &model.SystemTask{}, &model.DetectionRun{}))
createChannelSelectAutoGroupsChannel(t, db, 6501, "default", fingerprint.Model)
old := time.Now().Add(-8 * time.Hour).Unix()
result, err := common.Marshal(fingerprint.Result{Engine: fingerprint.Version, Status: "mismatch", ValidSamples: 3, Closest: "gpt-6-sol"})
require.NoError(t, err)
run := model.DetectionRun{Result: string(result), ID: "old-valid", RequestKey: "old-valid", ChannelID: 6501, Kind: "fingerprint", Status: "succeeded", Settled: true, Quality: "degraded", CreatedAt: old}
require.NoError(t, db.Create(&run).Error)
require.NoError(t, db.Create(&model.DetectionRun{ID: "failed", RequestKey: "failed", ChannelID: 6501, Kind: "pelican", Status: "failed", Settled: true, CreatedAt: time.Now().Unix()}).Error)
groupDegradationSnapshot.latest = nil
InvalidateGroupQualityCache()
t.Cleanup(func() {
groupDegradationSnapshot.latest = nil
groupDegradationSnapshot.benchmark = nil
InvalidateGroupQualityCache()
})
ctx := func(allow, child bool) *gin.Context {
c, _ := gin.CreateTestContext(httptest.NewRecorder())
c.Request = httptest.NewRequest("POST", "/v1/responses", nil)
c.Set(model.TokenRoutingContextKey, &model.TokenRoutingPolicy{Version: 1, Mode: "smart", AllowDegraded: allow, SkipDegradedNonAstra: child})
return c
}
for _, allow := range []bool{false, true} {
for _, child := range []bool{false, true} {
require.Equal(t, !allow, requestGroupDegraded(ctx(allow, child), fingerprint.Model, "default"))
require.Equal(t, !allow && child, requestGroupDegraded(ctx(allow, child), "gpt-6-sol", "default"))
}
}
recent, err := RecentDegradedGroups(context.Background())
require.NoError(t, err)
require.False(t, recent["default"])
require.NoError(t, db.Migrator().DropTable(&model.DetectionRun{}))
InvalidateGroupQualityCache()
require.True(t, requestGroupDegraded(ctx(false, false), fingerprint.Model, "default"), "read failure retains valid conclusion")
groupDegradationSnapshot.latest = nil
InvalidateGroupQualityCache()
cold := ctx(false, false)
require.True(t, requestGroupDegraded(cold, fingerprint.Model, "default"))
require.True(t, cold.GetBool("pool_degradation_unavailable"))
require.NoError(t, db.AutoMigrate(&model.DetectionRun{}))
InvalidateGroupQualityCache()
require.False(t, requestGroupDegraded(ctx(false, false), fingerprint.Model, "default"), "successful empty history remains eligible")
}
func TestAbandonedPoolProbeIsReleasedWithoutClearingFailures(t *testing.T) {
db := setupChannelSelectAutoGroupsTest(t)
createChannelSelectAutoGroupsChannel(t, db, 7001, "vip", "m")
model.InitChannelCache()
c := setupGroupPoolContext(t)
c.Request = httptest.NewRequest("POST", "/v1/responses", nil)
c.Set("id", 701)
c.Set("token_id", 701)
key := poolCooldownKey(c, "m", "default")
require.NoError(t, poolCooldowns().SetWithTTL(key, poolCooldown{Until: time.Now().Add(-time.Second).UnixMilli()}, 0))
t.Cleanup(func() { _, _ = poolCooldowns().DeleteMany([]string{key}) })
require.True(t, claimPoolGroup(c, "m", "default"))
require.True(t, poolGroupSkipped(c, "m", "default"))
ReleasePoolProbes(c)
require.False(t, poolGroupSkipped(c, "m", "default"))
require.True(t, claimPoolGroup(c, "m", "default"))
c.Set(model.TokenRoutingContextKey, &model.TokenRoutingPolicy{Version: 1, Mode: "smart"})
RecordPoolRetryableFailure(c, "m", "default")
ReleasePoolProbes(c)
value, found, err := poolCooldowns().Get(key)
require.NoError(t, err)
require.True(t, found)
require.Greater(t, value.Until, time.Now().UnixMilli())
}
func TestManualPoolNonAstraDegradationSwitch(t *testing.T) {
for _, mode := range []string{"inherit", "custom"} {
for _, enabled := range []bool{false, true} {
for _, nonAstra := range []bool{false, true} {
c, _ := gin.CreateTestContext(httptest.NewRecorder())
c.Set(model.TokenRoutingContextKey, &model.TokenRoutingPolicy{Version: 1, Mode: mode, SkipDegradedGroups: enabled, SkipDegradedNonAstra: nonAstra})
c.Set("pool_group_fingerprints", map[string]AstraSamplingHistory{"default": {Status: "mismatch", ValidSamples: 3, Closest: "gpt-6-sol"}})
c.Set("pool_group_pelicans", map[string]DetectionHistory{"default": {Quality: "degraded"}})
require.Equal(t, enabled, requestGroupDegraded(c, fingerprint.Model, "default"))
require.Equal(t, enabled && nonAstra, requestGroupDegraded(c, "gpt-6-sol", "default"))
require.False(t, requestGroupDegraded(c, "gpt-6-sol", "untested"))
}
}
}
}
service/image_generation_group.go
+58
View File
@@ -0,0 +1,58 @@
package service
import (
"github.com/QuantumNous/new-api/common"
"github.com/QuantumNous/new-api/constant"
"github.com/QuantumNous/new-api/model"
"github.com/QuantumNous/new-api/setting/model_setting"
"github.com/gin-gonic/gin"
"slices"
)
func imageGenerationPoolGroup(c *gin.Context) string {
settings := model_setting.GetGlobalSettings()
group := settings.ImageGenerationGroup
if group == "" || !IsUserSelectableGroup(common.GetContextKeyString(c, constant.ContextKeyUserGroup), group) {
return ""
}
policy := RequestRoutingPolicy(c)
if policy != nil {
if policy.Mode == "user" {
return ""
}
return group
}
if common.GetContextKeyString(c, constant.ContextKeyTokenGroup) == "auto" || common.GetContextKeyBool(c, constant.ContextKeyTokenCrossGroupRetry) {
return group
}
return ""
}
func imageGenerationRequestGroup(c *gin.Context, modelName string) string {
if !slices.Contains(model_setting.GetGlobalSettings().ImageGenerationModels, modelName) {
return ""
}
return imageGenerationPoolGroup(c)
}
// The dedicated route is shared by selection, affinity and the routing preview.
func WithImageGenerationGroup(c *gin.Context, modelName string, entries []model.TokenGroupEntry) []model.TokenGroupEntry {
group := imageGenerationPoolGroup(c)
if group == "" {
return entries
}
if modelName != "" {
if imageGenerationRequestGroup(c, modelName) != "" {
return []model.TokenGroupEntry{{Group: group, Priority: 0}}
}
return entries
}
priority := 0
for _, entry := range entries {
if entry.Group == group {
return entries
}
priority = max(priority, entry.Priority+1)
}
return append(slices.Clone(entries), model.TokenGroupEntry{Group: group, Priority: priority})
}
service/image_generation_group_test.go
+55
View File
@@ -0,0 +1,55 @@
package service
import (
"github.com/QuantumNous/new-api/common"
"github.com/QuantumNous/new-api/constant"
"github.com/QuantumNous/new-api/model"
"github.com/QuantumNous/new-api/setting"
"github.com/QuantumNous/new-api/setting/model_setting"
"github.com/QuantumNous/new-api/setting/ratio_setting"
"github.com/gin-gonic/gin"
"github.com/stretchr/testify/require"
"net/http/httptest"
"testing"
)
func TestDedicatedImageGenerationPoolSelection(t *testing.T) {
db := setupChannelSelectAutoGroupsTest(t)
cfg := model_setting.GetGlobalSettings()
oldGroup, oldModels := cfg.ImageGenerationGroup, cfg.ImageGenerationModels
t.Cleanup(func() { cfg.ImageGenerationGroup, cfg.ImageGenerationModels = oldGroup, oldModels })
cfg.ImageGenerationGroup, cfg.ImageGenerationModels = "example-images", []string{"gpt-image-2"}
require.NoError(t, setting.UpdateUserUsableGroupsByJSONString(`{"default":"Default","vip":"VIP","example-images":"Images"}`))
require.NoError(t, ratio_setting.UpdateGroupRatioByJSONString(`{"default":1,"vip":2,"example-images":1.25}`))
createChannelSelectAutoGroupsChannel(t, db, 8001, "default", "gpt-image-2")
createChannelSelectAutoGroupsChannel(t, db, 8002, "example-images", "gpt-image-2")
model.InitChannelCache()
ResetGroupAvailability()
t.Cleanup(ResetGroupAvailability)
for _, mode := range []string{"smart", "inherit", "custom"} {
t.Run(mode, func(t *testing.T) {
c, _ := gin.CreateTestContext(httptest.NewRecorder())
c.Request = httptest.NewRequest("POST", "/v1/images/generations", nil)
common.SetContextKey(c, constant.ContextKeyUserGroup, "default")
c.Set(model.TokenRoutingContextKey, &model.TokenRoutingPolicy{Version: 1, Mode: mode, Groups: []model.TokenGroupEntry{{Group: "default", Priority: 0}}})
entries := ResolveRequestPool(c, "gpt-image-2")
require.Equal(t, []model.TokenGroupEntry{{Group: "example-images", Priority: 0}}, entries)
channel, group, err := CacheGetRandomSatisfiedChannel(&RetryParam{Ctx: c, TokenGroup: "default", ModelName: "gpt-image-2", Retry: common.GetPointer(0)})
require.NoError(t, err)
require.NotNil(t, channel)
require.Equal(t, 8002, channel.Id)
require.Equal(t, "example-images", group)
})
}
c, _ := gin.CreateTestContext(httptest.NewRecorder())
c.Request = httptest.NewRequest("POST", "/v1/images/generations", nil)
common.SetContextKey(c, constant.ContextKeyUserGroup, "default")
defaults := []model.TokenGroupEntry{{Group: "default", Priority: 0}}
c.Set(model.TokenRoutingContextKey, &model.TokenRoutingPolicy{Version: 1, Mode: "custom", Groups: defaults})
require.Equal(t, defaults, ResolveRequestPool(c, "text-model"))
c.Set(model.TokenRoutingContextKey, &model.TokenRoutingPolicy{Version: 1, Mode: "user"})
require.Equal(t, defaults, ResolveRequestPool(c, "gpt-image-2"))
c.Set(model.TokenRoutingContextKey, &model.TokenRoutingPolicy{Version: 1, Mode: "custom", Groups: defaults})
require.NoError(t, setting.UpdateUserUsableGroupsByJSONString(`{"default":"Default","vip":"VIP"}`))
require.Equal(t, defaults, ResolveRequestPool(c, "gpt-image-2"))
}
service/log_info_generate.go
+3 -1
View File
@@ -100,7 +100,9 @@ func AppendRelayLogAdminInfo(ctx *gin.Context, relayInfo *relaycommon.RelayInfo,
func GenerateTextOtherInfo(ctx *gin.Context, relayInfo *relaycommon.RelayInfo, modelRatio, groupRatio, completionRatio float64,
cacheTokens int, cacheRatio float64, modelPrice float64, userGroupRatio float64) *model.LogOther {
if !ctx.GetBool("attempt_checkpoint") && ctx.GetString("attempt_guard_result") == "" { MarkRequestPolicySuccess(ctx, relayInfo.StreamStatus) }
if !ctx.GetBool("attempt_checkpoint") && ctx.GetString("attempt_guard_result") == "" {
MarkRequestPolicySuccess(ctx, relayInfo.StreamStatus)
}
other := model.NewLogOther()
other.SetPublic("model_ratio", modelRatio)
other.SetPublic("group_ratio", groupRatio)
service/pool_failover.go
+247
View File
@@ -0,0 +1,247 @@
package service
import (
"fmt"
"sync"
"time"
"github.com/QuantumNous/new-api/common"
"github.com/QuantumNous/new-api/model"
"github.com/QuantumNous/new-api/pkg/cachex"
"github.com/QuantumNous/new-api/relaykit/types"
"github.com/QuantumNous/new-api/setting/operation_setting"
"github.com/gin-gonic/gin"
"github.com/samber/hot"
)
const poolFailoverCooldown = 10 * time.Minute
type poolCooldown struct {
Until int64 `json:"until"`
ProbeUntil int64 `json:"probe_until"`
ProbeID string `json:"probe_id"`
}
var poolCooldownOnce sync.Once
var poolCooldownCache *cachex.HybridCache[poolCooldown]
var poolCooldownMu sync.Mutex
func poolCooldowns() *cachex.HybridCache[poolCooldown] {
poolCooldownOnce.Do(func() {
poolCooldownCache = cachex.NewHybridCache[poolCooldown](cachex.HybridCacheConfig[poolCooldown]{
Namespace: cachex.Namespace("new-api:pool-failover:v3"), Redis: common.RDB,
RedisEnabled: func() bool { return common.RedisEnabled && common.RDB != nil }, RedisCodec: cachex.JSONCodec[poolCooldown]{},
Memory: func() *hot.HotCache[string, poolCooldown] {
return hot.NewHotCache[string, poolCooldown](hot.LRU, 100000).Build()
},
})
})
return poolCooldownCache
}
func poolCooldownKey(c *gin.Context, modelName, group string) string {
return common.GenerateHMAC(fmt.Sprintf("%d:%d:%t:%d:%s:%d:%s", c.GetInt("id"), c.GetInt("token_id"), c.GetBool("pool_request_stream"), len(modelName), modelName, len(group), group))
}
func poolFailures(c *gin.Context) map[string]int {
if value, ok := c.Get("pool_failures"); ok {
return value.(map[string]int)
}
value := map[string]int{}
c.Set("pool_failures", value)
return value
}
// A cooldown may exclude a group only while another eligible route exists.
// Inspect alternatives without claiming probes or recursively checking cooldowns.
func hasPoolAlternative(c *gin.Context, modelName, current string) bool {
for _, entry := range ResolveRequestPool(c, modelName) {
group := entry.Group
if group == current || poolFailures(c)[group] >= 2 || requestGroupDegraded(c, modelName, group) {
continue
}
value, found, err := poolCooldowns().Get(poolCooldownKey(c, modelName, group))
now := time.Now().UnixMilli()
if err != nil || (found && (now < value.Until || now < value.ProbeUntil)) {
continue
}
channel, err := model.GetRandomSatisfiedChannel(group, modelName, 0, GetChannelConstraints(c).Filters)
if err == nil && channel != nil {
return true
}
}
return false
}
func poolGroupSkipped(c *gin.Context, modelName, group string) bool {
if poolFailures(c)[group] >= 2 {
return true
}
if RequestRoutingPolicy(c) != nil && !hasPoolAlternative(c, modelName, group) {
return false
}
value, found, err := poolCooldowns().Get(poolCooldownKey(c, modelName, group))
if err != nil || !found {
return false
}
now := time.Now().UnixMilli()
skipped := now < value.Until || now < value.ProbeUntil
if skipped {
wait := int((max(value.Until, value.ProbeUntil) - now + 999) / 1000)
previous := c.GetInt("pool_retry_after")
if previous == 0 || wait < previous {
c.Set("pool_retry_after", wait)
}
}
return skipped
}
// A request rejected after selection or interrupted while streaming must not
// leave its recovery reservation blocking subsequent requests for five minutes.
func ReleasePoolProbes(c *gin.Context) {
raw, ok := c.Get("pool_claimed_groups")
if !ok {
return
}
for group, key := range raw.(map[string]string) {
observed, ok := c.Get("pool_probe:" + group)
if !ok {
continue
}
trial := observed.(poolCooldown)
poolCooldownMu.Lock()
current, found, err := poolCooldowns().Get(key)
if err == nil && found && current.Until == trial.Until && current.ProbeID == trial.ProbeID {
current.ProbeUntil, current.ProbeID = 0, ""
_ = poolCooldowns().SetWithTTL(key, current, 0)
}
poolCooldownMu.Unlock()
}
}
// Only the selected candidate can claim a half-open probe. The expired marker
// survives the ten-minute cooldown so a failed probe never gets a second try.
func claimPoolGroup(c *gin.Context, modelName, group string) bool {
if RequestRoutingPolicy(c) != nil && !hasPoolAlternative(c, modelName, group) {
return true
}
poolCooldownMu.Lock()
defer poolCooldownMu.Unlock()
key := poolCooldownKey(c, modelName, group)
value, found, err := poolCooldowns().Get(key)
if err != nil || !found {
return true
}
now := time.Now().UnixMilli()
if now < value.Until || now < value.ProbeUntil {
return false
}
value.ProbeUntil = now + int64((5*time.Minute)/time.Millisecond)
value.ProbeID = fmt.Sprint(time.Now().UnixNano())
if err := poolCooldowns().SetWithTTL(key, value, 0); err != nil {
return false
}
c.Set("pool_probe:"+group, value)
raw, ok := c.Get("pool_claimed_groups")
claimed := map[string]string{}
if ok {
claimed = raw.(map[string]string)
}
claimed[group] = key
c.Set("pool_claimed_groups", claimed)
return true
}
func IsPoolFailoverError(c *gin.Context, err *types.NewAPIError) bool {
if _, pinned, _ := GetChannelConstraints(c).ResolvedPin(); pinned {
return false
}
if err == nil || RequestRoutingPolicy(c) == nil || GetChannelConstraints(c).SuppressesRetry() || types.IsSkipRetryError(err) {
return false
}
if c.Request != nil && c.Request.Context().Err() != nil {
return false
}
if c.Writer != nil && c.Writer.Written() {
return false
}
if err.StatusCode < 500 || err.StatusCode > 599 {
return false
}
return !operation_setting.IsAlwaysSkipRetryStatusCode(err.StatusCode) && !operation_setting.IsAlwaysSkipRetryCode(err.GetErrorCode()) && operation_setting.ShouldRetryByStatusCode(err.StatusCode)
}
func RecordPoolRetryableFailure(c *gin.Context, modelName, group string) {
if _, pinned, _ := GetChannelConstraints(c).ResolvedPin(); pinned {
return
}
if RequestRoutingPolicy(c) == nil || group == "" || GetChannelConstraints(c).SuppressesRetry() {
return
}
failures := poolFailures(c)
failures[group]++
c.Set("pool_retry_group", group)
if _, trial := c.Get("pool_probe:" + group); trial {
failures[group] = 2
}
if failures[group] < 2 || !hasPoolAlternative(c, modelName, group) {
return
}
poolCooldownMu.Lock()
key := poolCooldownKey(c, modelName, group)
previous, _, _ := poolCooldowns().Get(key)
value := poolCooldown{Until: max(previous.Until+1, time.Now().Add(poolFailoverCooldown).UnixMilli())}
if err := poolCooldowns().SetWithTTL(key, value, 0); err != nil {
common.SysError("pool cooldown write failed: " + err.Error())
}
poolCooldownMu.Unlock()
ClearPoolAffinity(c)
}
func RecordPoolSuccess(c *gin.Context, modelName, group string) {
raw, trial := c.Get("pool_probe:" + group)
if !trial {
return
}
observed := raw.(poolCooldown)
poolCooldownMu.Lock()
defer poolCooldownMu.Unlock()
key := poolCooldownKey(c, modelName, group)
current, found, err := poolCooldowns().Get(key)
if err == nil && found && current.Until == observed.Until && current.ProbeID == observed.ProbeID {
_, _ = poolCooldowns().DeleteMany([]string{key})
}
}
func PoolRetryBudget(c *gin.Context, modelName string) int {
if _, pinned, _ := GetChannelConstraints(c).ResolvedPin(); pinned {
return 0
}
if RequestRoutingPolicy(c) == nil || GetChannelConstraints(c).SuppressesRetry() {
return 0
}
return max(0, 2*len(ResolveRequestPool(c, modelName))-1)
}
func poolHigherPriorityRecovered(c *gin.Context, modelName, boundGroup string) bool {
entries := ResolveRequestPool(c, modelName)
priority := -1
for _, entry := range entries {
if entry.Group == boundGroup {
priority = entry.Priority
break
}
}
for _, entry := range entries {
if entry.Priority >= priority {
break
}
value, found, err := poolCooldowns().Get(poolCooldownKey(c, modelName, entry.Group))
if err == nil && found && time.Now().UnixMilli() >= value.Until && !poolGroupSkipped(c, modelName, entry.Group) && !requestGroupDegraded(c, modelName, entry.Group) {
return true
}
}
return false
}
func PoolRetryExhausted(c *gin.Context, modelName, group string) bool {
return poolFailures(c)[group] >= 2 && !hasPoolAlternative(c, modelName, group)
}
service/pool_no_backup_test.go
+95
View File
@@ -0,0 +1,95 @@
package service
import (
"context"
"github.com/QuantumNous/new-api/common"
"github.com/QuantumNous/new-api/model"
"github.com/gin-gonic/gin"
"github.com/stretchr/testify/require"
"net/http/httptest"
"testing"
"time"
)
func TestPoolWithoutAvailableBackupDoesNotCoolDown(t *testing.T) {
db := setupChannelSelectAutoGroupsTest(t)
createChannelSelectAutoGroupsChannel(t, db, 7801, "default", "m")
createChannelSelectAutoGroupsChannel(t, db, 7802, "vip", "other")
model.InitChannelCache()
ResetGroupAvailability()
t.Cleanup(func() { ResetGroupAvailability(); _ = poolCooldowns().Purge() })
makeCtx := func() *gin.Context {
c := setupGroupPoolContext(t)
c.Request = httptest.NewRequest("POST", "/v1/responses", nil)
c.Set("token_id", 7801)
c.Set(model.TokenRoutingContextKey, &model.TokenRoutingPolicy{Version: 1, Mode: "custom", Groups: []model.TokenGroupEntry{{Group: "default"}, {Group: "vip", Priority: 1}}})
return c
}
c := makeCtx()
RecordPoolRetryableFailure(c, "m", "default")
RecordPoolRetryableFailure(c, "m", "default")
_, found, err := poolCooldowns().Get(poolCooldownKey(c, "m", "default"))
require.NoError(t, err)
require.False(t, found)
// An old cooldown and global health state must not strand a single eligible group.
require.NoError(t, poolCooldowns().SetWithTTL(poolCooldownKey(c, "m", "default"), poolCooldown{Until: time.Now().Add(10 * time.Minute).UnixMilli()}, 0))
for i := 0; i < 4; i++ {
RecordGroupOutcome("default", false)
}
next := makeCtx()
ch, group, err := selectTokenRoutingPool(&RetryParam{Ctx: next, ModelName: "m", Retry: common.GetPointer(0)})
require.NoError(t, err)
require.NotNil(t, ch)
require.Equal(t, "default", group)
AppendUsedChannel(next, ch.Id)
RecordPoolRetryableFailure(next, "m", "default")
ch, _, err = selectTokenRoutingPool(&RetryParam{Ctx: next, ModelName: "m", Retry: common.GetPointer(1)})
require.NoError(t, err)
require.NotNil(t, ch)
AppendUsedChannel(next, ch.Id)
RecordPoolRetryableFailure(next, "m", "default")
ch, _, err = selectTokenRoutingPool(&RetryParam{Ctx: next, ModelName: "m", Retry: common.GetPointer(2)})
require.NoError(t, err)
require.Nil(t, ch, "do not loop indefinitely inside one request")
}
func TestAutomaticSpeedUsesLatestExactModelStreamingTest(t *testing.T) {
db := setupChannelSelectAutoGroupsTest(t)
require.NoError(t, db.AutoMigrate(&model.Log{}))
previous := model.LOG_DB
model.LOG_DB = db
t.Cleanup(func() { model.LOG_DB = previous })
now := time.Now().Unix()
samples := []model.Log{
{ChannelId: 1, Type: 2, TokenId: 0, Content: "模型测试", ModelName: "m", IsStream: true, CreatedAt: now - 20, Other: `{"frt":500}`},
{ChannelId: 1, Type: 2, TokenId: 0, Content: "模型测试", ModelName: "m", IsStream: true, CreatedAt: now - 10, Other: `{"frt":60000}`},
{ChannelId: 1, Type: 2, TokenId: 9, Content: "模型测试", ModelName: "m", IsStream: true, CreatedAt: now, Other: `{"frt":1}`},
{ChannelId: 2, Type: 2, TokenId: 0, Content: "模型测试", ModelName: "m", IsStream: true, CreatedAt: now - 5, Other: `{"frt":1200}`},
{ChannelId: 2, Type: 2, TokenId: 0, Content: "模型测试", ModelName: "other", IsStream: true, CreatedAt: now, Other: `{"frt":10}`},
{ChannelId: 3, Type: 2, TokenId: 0, Content: "模型测试", ModelName: "m", IsStream: false, CreatedAt: now, Other: `{"frt":10}`},
{ChannelId: 4, Type: 2, TokenId: 0, Content: "模型测试", ModelName: "m", IsStream: true, CreatedAt: now - 2*3600, Other: `{"frt":10}`},
}
require.NoError(t, db.Create(&samples).Error)
speeds := latestModelTestSpeeds(context.Background(), "m")
require.Equal(t, map[int]float64{1: 60000, 2: 1200}, speeds)
}
func TestProbeAndStreamingCooldownsAreIsolated(t *testing.T) {
db := setupChannelSelectAutoGroupsTest(t)
createChannelSelectAutoGroupsChannel(t, db, 7901, "default", "m")
createChannelSelectAutoGroupsChannel(t, db, 7902, "vip", "m")
model.InitChannelCache()
t.Cleanup(func() { _ = poolCooldowns().Purge() })
makeCtx := func(stream bool) *gin.Context {
c := setupGroupPoolContext(t)
c.Request = httptest.NewRequest("POST", "/v1/responses", nil)
c.Set("token_id", 7901)
c.Set("pool_request_stream", stream)
c.Set(model.TokenRoutingContextKey, &model.TokenRoutingPolicy{Version: 1, Mode: "custom", Groups: []model.TokenGroupEntry{{Group: "default"}, {Group: "vip", Priority: 1}}})
return c
}
probe := makeCtx(false)
RecordPoolRetryableFailure(probe, "m", "default")
RecordPoolRetryableFailure(probe, "m", "default")
require.True(t, poolGroupSkipped(makeCtx(false), "m", "default"))
require.False(t, poolGroupSkipped(makeCtx(true), "m", "default"))
}
service/quota.go
+1 -15
View File
@@ -7,7 +7,6 @@ import (
"time"
"github.com/QuantumNous/new-api/common"
"github.com/QuantumNous/new-api/constant"
"github.com/QuantumNous/new-api/logger"
"github.com/QuantumNous/new-api/model"
"github.com/QuantumNous/new-api/pkg/billingexpr"
@@ -103,21 +102,8 @@ func PreWssConsumeQuota(ctx *gin.Context, relayInfo *relaycommon.RelayInfo, usag
textOutTokens := usage.OutputTokenDetails.TextTokens
audioInputTokens := usage.InputTokenDetails.AudioTokens
audioOutTokens := usage.OutputTokenDetails.AudioTokens
groupRatio := ratio_setting.GetGroupRatio(relayInfo.UsingGroup)
modelRatio, _, _ := ratio_setting.GetModelRatio(modelName)
autoGroup, exists := common.GetContextKey(ctx, constant.ContextKeyAutoGroup)
if exists {
groupRatio = ratio_setting.GetGroupRatio(autoGroup.(string))
logger.LogDebug(ctx, "final group ratio: %f", groupRatio)
relayInfo.UsingGroup = autoGroup.(string)
}
actualGroupRatio := groupRatio
userGroupRatio, ok := ratio_setting.GetGroupGroupRatio(relayInfo.UserGroup, relayInfo.UsingGroup)
if ok {
actualGroupRatio = userGroupRatio
}
actualGroupRatio := relayInfo.PriceData.GroupRatioInfo.GroupRatio
quotaInfo := QuotaInfo{
InputDetails: TokenDetails{
service/relay_error.go
+7 -1
View File
@@ -22,7 +22,13 @@ func DecideRelayRetry(c *gin.Context, err *types.NewAPIError, retryTimes int) Po
if err == nil {
return PolicyDecision{Action: "stop", Reason: "request_completed", Source: "system"}
}
if ShouldSkipRetryAfterChannelAffinityFailure(c) {
if c.Request != nil && c.Request.Context().Err() != nil {
return PolicyDecision{Action: "stop", Reason: "client_cancelled", Source: "system"}
}
if c.Writer != nil && c.Writer.Written() {
return PolicyDecision{Action: "stop", Reason: "response_started", Source: "system"}
}
if ShouldSkipRetryAfterChannelAffinityFailure(c) && !IsPoolFailoverError(c, err) {
source := RequestPolicy(c).SessionModeSource
if source == "" {
source = "session_rule"
service/request_details.go
+11 -10
View File
@@ -11,19 +11,20 @@ import (
)
type RequestDetails struct {
SessionID string `json:"session_id,omitempty"`
SessionSource string `json:"session_source"`
CorrelationKey string `json:"correlation_key,omitempty"`
RequestedModel string `json:"requested_model,omitempty"`
UpstreamModel string `json:"upstream_model,omitempty"`
ReturnedModel string `json:"returned_model,omitempty"`
ReasoningEffort string `json:"reasoning_effort,omitempty"`
Status int `json:"status"`
ModelGuard *relaycommon.ResponseModelGuard `json:"model_guard,omitempty"`
RoutingPolicyFingerprint string `json:"routing_policy_fingerprint,omitempty"`
SessionID string `json:"session_id,omitempty"`
SessionSource string `json:"session_source"`
CorrelationKey string `json:"correlation_key,omitempty"`
RequestedModel string `json:"requested_model,omitempty"`
UpstreamModel string `json:"upstream_model,omitempty"`
ReturnedModel string `json:"returned_model,omitempty"`
ReasoningEffort string `json:"reasoning_effort,omitempty"`
Status int `json:"status"`
ModelGuard *relaycommon.ResponseModelGuard `json:"model_guard,omitempty"`
}
func appendRequestDetails(c *gin.Context, info *relaycommon.RelayInfo, other *model.LogOther) {
d := RequestDetails{SessionSource: "none", RequestedModel: c.GetString("original_model"), Status: c.Writer.Status()}
d := RequestDetails{RoutingPolicyFingerprint: c.GetString("pool_affinity_policy_fingerprint"), SessionSource: "none", RequestedModel: c.GetString("original_model"), Status: c.Writer.Status()}
if status, ok := other.Snapshot()["status_code"].(int); ok {
d.Status = status
}
service/request_policy.go
+3
View File
@@ -145,6 +145,9 @@ func MarkRequestPolicySuccess(c *gin.Context, stream *relaycommon.StreamStatus)
}
state.OutcomeRecorded = true
state.Successful = stream == nil || stream.IsNormalEnd() && !stream.HasErrors() && (stream.ResponseOutcome() == "" || stream.ResponseOutcome() == "completed")
if state.Successful {
RecordPoolSuccess(c, c.GetString("original_model"), common.GetContextKeyString(c, constant.ContextKeyAutoGroup))
}
decision := PolicyDecision{Action: "success", Reason: "request_completed", Source: "upstream"}
if !state.Successful {
decision = PolicyDecision{Action: "stop", Reason: "stream_not_successful", Source: "system"}
Loaded 100 of 154 files, more files were not shown because too many files have changed in this diff. Show more