121 lines
5.7 KiB
Go
121 lines
5.7 KiB
Go
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) {
|
|
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)
|
|
}
|