feat: sync deployed v1.0.0-rc.38-custom.42
Sync production source including routing, billing, Astra detection and public uptime. Update custom build instructions and license packaging.
This commit is contained in:
cneicy committed
1 parent
aa6c3f50f4
commit
ceddfd3d6b
828 files changed
+122253
-13702
No files matched your search
@@ -0,0 +1,121 @@
|
||||
package relay
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"github.com/QuantumNous/new-api/common"
|
||||
"github.com/QuantumNous/new-api/constant"
|
||||
"github.com/QuantumNous/new-api/model"
|
||||
"github.com/QuantumNous/new-api/pkg/billingexpr"
|
||||
"github.com/QuantumNous/new-api/relay/channel"
|
||||
"github.com/QuantumNous/new-api/relay/channel/openai"
|
||||
relaycommon "github.com/QuantumNous/new-api/relay/common"
|
||||
"github.com/QuantumNous/new-api/relaykit/dto"
|
||||
"github.com/QuantumNous/new-api/service"
|
||||
hosttypes "github.com/QuantumNous/new-api/types"
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/glebarez/sqlite"
|
||||
"github.com/stretchr/testify/require"
|
||||
"gorm.io/gorm"
|
||||
"io"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
func TestRelayAttemptMismatchCollectsUsageAndChargesRequestedPrice(t *testing.T) {
|
||||
service.InitTokenEncoders()
|
||||
db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{})
|
||||
require.NoError(t, err)
|
||||
sqlDB, err := db.DB()
|
||||
require.NoError(t, err)
|
||||
sqlDB.SetMaxOpenConns(1)
|
||||
previousDB, previousLog, previousRedis := model.DB, model.LOG_DB, common.RedisEnabled
|
||||
model.DB = db
|
||||
model.LOG_DB = db
|
||||
common.RedisEnabled = false
|
||||
t.Cleanup(func() {
|
||||
model.DB = previousDB
|
||||
model.LOG_DB = previousLog
|
||||
common.RedisEnabled = previousRedis
|
||||
require.NoError(t, sqlDB.Close())
|
||||
})
|
||||
require.NoError(t, db.AutoMigrate(&model.User{}, &model.Token{}, &model.Channel{}, &model.RelayAttempt{}, &model.Log{}, &model.QuotaData{}))
|
||||
user := model.User{Username: "attempt-relay", Quota: 10000}
|
||||
require.NoError(t, db.Create(&user).Error)
|
||||
token := model.Token{UserId: user.Id, Key: "attempt-relay", Status: common.TokenStatusEnabled, RemainQuota: 10000}
|
||||
require.NoError(t, db.Create(&token).Error)
|
||||
expression := `p * 2 + c * 4 + cr * 0.5`
|
||||
for i, ratio := range []float64{1, 3, 1} {
|
||||
selected := model.Channel{Name: fmt.Sprint(i), Status: common.ChannelStatusEnabled}
|
||||
require.NoError(t, db.Create(&selected).Error)
|
||||
// The final usage is buffered behind the rejected prefix: Drain must reuse
|
||||
// the same reader rather than lose bytes already read from the socket.
|
||||
body := "data: {\"type\":\"response.created\",\"response\":{\"model\":\"gpt-5.6-luna\"}}\n\n" +
|
||||
"data: {\"type\":\"response.output_text.delta\",\"delta\":\"hidden generated text\"}\n\n" +
|
||||
"data: {\"type\":\"response.completed\",\"response\":{\"model\":\"gpt-5.6-luna\",\"usage\":{\"input_tokens\":10,\"output_tokens\":2,\"total_tokens\":12,\"input_tokens_details\":{\"cached_tokens\":4}}}}\n\ndata: [DONE]\n\n"
|
||||
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
w.Header().Set("Content-Type", "text/event-stream")
|
||||
_, _ = io.WriteString(w, body)
|
||||
}))
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
c, writer := gin.CreateTestContext(httptest.NewRecorder())
|
||||
_ = writer
|
||||
c.Request = httptest.NewRequest("POST", "/v1/responses", strings.NewReader(`{"model":"gpt-6-astra","input":"secret prompt"}`)).WithContext(ctx)
|
||||
c.Set("channel_id", selected.Id)
|
||||
c.Set("id", user.Id)
|
||||
c.Set("token_id", token.Id)
|
||||
common.SetContextKey(c, constant.ContextKeyUserName, user.Username)
|
||||
info := &relaycommon.RelayInfo{RequestId: fmt.Sprintf("attempt-request-%d", i), UserId: user.Id, TokenId: token.Id, TokenKey: token.Key, OriginModelName: "gpt-6-astra", UsingGroup: fmt.Sprint(ratio), UserSetting: dto.UserSetting{BillingPreference: "wallet_only"}, StartTime: time.Now(), DurableAttemptBilling: true, ResponseGuard: &relaycommon.ResponseModelGuard{}, ChannelMeta: &relaycommon.ChannelMeta{ChannelId: selected.Id, UpstreamModelName: "gpt-6-astra"}, PriceData: hosttypes.PriceData{GroupRatioInfo: hosttypes.GroupRatioInfo{GroupRatio: ratio}}, TieredBillingSnapshot: &billingexpr.BillingSnapshot{BillingMode: "tiered_expr", ExprString: expression, ExprHash: billingexpr.ExprHashString(expression), QuotaPerUnit: common.QuotaPerUnit, GroupRatio: ratio, EstimatedQuotaAfterGroup: 100}}
|
||||
info.SetEstimatePromptTokens(30)
|
||||
info.AstraAutoSwitch = i < 2
|
||||
if !info.AstraAutoSwitch {
|
||||
info.ResponseGuard = nil
|
||||
}
|
||||
require.Nil(t, service.BeginRelayAttempt(c, info))
|
||||
req, err := http.NewRequest("POST", server.URL, strings.NewReader("{}"))
|
||||
require.NoError(t, err)
|
||||
resp, err := channel.DoRequest(c, req, info)
|
||||
if info.AstraAutoSwitch {
|
||||
require.Error(t, err)
|
||||
require.Nil(t, resp)
|
||||
require.False(t, c.Writer.Written())
|
||||
require.True(t, info.AttemptDetached)
|
||||
require.False(t, service.RequestPolicy(c).Successful, "collector cannot mark the client request successful")
|
||||
} else {
|
||||
require.NoError(t, err)
|
||||
usage, apiErr := openai.OaiResponsesHandler(c, info, resp)
|
||||
require.Nil(t, apiErr)
|
||||
require.True(t, c.Writer.Written())
|
||||
require.True(t, info.ResponseModel.Mismatch)
|
||||
ConsumeResponsesQuota(c, info, usage)
|
||||
}
|
||||
var attempt model.RelayAttempt
|
||||
require.Eventually(t, func() bool {
|
||||
return db.First(&attempt, info.DurableAttemptID).Error == nil && attempt.State == "settled"
|
||||
}, time.Second*5, time.Millisecond*10)
|
||||
require.Equal(t, int(11*ratio), attempt.Quota)
|
||||
if info.AstraAutoSwitch {
|
||||
require.Equal(t, "mismatch_rejected", attempt.GuardResult)
|
||||
} else {
|
||||
require.Equal(t, "mismatch_allowed", attempt.GuardResult)
|
||||
require.Contains(t, attempt.LogJSON, "mismatch_allowed")
|
||||
}
|
||||
require.Equal(t, "upstream", attempt.UsageSource)
|
||||
require.NotContains(t, attempt.LogJSON, "secret prompt")
|
||||
require.NotContains(t, attempt.LogJSON, "hidden generated text")
|
||||
require.NoError(t, model.DeliverRelayAttemptLogs())
|
||||
service.EndFailedRelayAttempt(c, info)
|
||||
cancel()
|
||||
server.Close()
|
||||
}
|
||||
require.NoError(t, db.First(&user, user.Id).Error)
|
||||
require.Equal(t, 9945, user.Quota)
|
||||
require.NoError(t, db.First(&token, token.Id).Error)
|
||||
require.Equal(t, 55, token.UsedQuota)
|
||||
var logs int64
|
||||
require.NoError(t, db.Model(&model.Log{}).Count(&logs).Error)
|
||||
require.EqualValues(t, 3, logs)
|
||||
}
|
||||
Reference in new issue
Block a user