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,350 @@
|
||||
package service
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"net/http"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/QuantumNous/new-api/common"
|
||||
"github.com/QuantumNous/new-api/constant"
|
||||
"github.com/QuantumNous/new-api/model"
|
||||
relaycommon "github.com/QuantumNous/new-api/relay/common"
|
||||
"github.com/QuantumNous/new-api/relaykit/dto"
|
||||
"github.com/QuantumNous/new-api/relaykit/types"
|
||||
"github.com/gin-gonic/gin"
|
||||
)
|
||||
|
||||
func BeginRelayAttempt(c *gin.Context, info *relaycommon.RelayInfo) *types.NewAPIError {
|
||||
if !info.DurableAttemptBilling {
|
||||
return nil
|
||||
}
|
||||
reserved := info.PriceData.QuotaToPreConsume
|
||||
if snap := info.TieredBillingSnapshot; snap != nil {
|
||||
reserved = snap.EstimatedQuotaAfterGroup
|
||||
}
|
||||
price, err := common.Marshal(struct {
|
||||
Price any
|
||||
Expression any
|
||||
}{info.PriceData, info.TieredBillingSnapshot})
|
||||
if err != nil {
|
||||
return types.NewError(err, types.ErrorCodeModelPriceError)
|
||||
}
|
||||
a := &model.RelayAttempt{AttemptKey: fmt.Sprintf("%s:%d", info.RequestId, info.UpstreamRequests+1), RequestID: info.RequestId, Number: info.UpstreamRequests + 1, UserID: info.UserId, TokenID: info.TokenId, ChannelID: c.GetInt("channel_id"), GroupName: info.UsingGroup, RequestedModel: info.OriginModelName, Reserved: reserved, EstimateQuota: reserved, PriceJSON: string(price), Deadline: time.Now().Add(10 * time.Minute).Unix(), GuardResult: "unchecked"}
|
||||
if err := model.ReserveRelayAttempt(a, info.UserSetting.BillingPreference); err != nil {
|
||||
return types.NewErrorWithStatusCode(err, types.ErrorCodeInsufficientUserQuota, http.StatusForbidden, types.ErrOptionWithSkipRetry())
|
||||
}
|
||||
info.DurableAttemptID = a.ID
|
||||
info.AttemptDetached = false
|
||||
info.Billing = nil
|
||||
info.BillingSource = a.Funding
|
||||
info.SubscriptionId = a.SubscriptionID
|
||||
info.FinalPreConsumedQuota = a.Reserved
|
||||
c.Set("relay_attempt_number", a.Number)
|
||||
c.Set("attempt_guard_result", "")
|
||||
c.Set("attempt_returned_model", "")
|
||||
c.Set("attempt_usage_source", "")
|
||||
requestContext := c.Request.Context()
|
||||
go func(id int64) {
|
||||
ticker := time.NewTicker(30 * time.Second)
|
||||
defer ticker.Stop()
|
||||
for {
|
||||
select {
|
||||
case <-requestContext.Done():
|
||||
return
|
||||
case <-ticker.C:
|
||||
r := model.DB.Model(&model.RelayAttempt{}).Where("id = ? AND state <> ?", id, "settled").Update("deadline", time.Now().Add(10*time.Minute).Unix())
|
||||
if r.Error == nil && r.RowsAffected == 0 {
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
}(a.ID)
|
||||
return nil
|
||||
}
|
||||
|
||||
// Log and account through one durable transaction instead of the ordinary
|
||||
// single-generation session. A checkpoint calculates price but never debits.
|
||||
func resolveAttemptGuardResult(reason string, info *relaycommon.RelayInfo) string {
|
||||
if reason != "" {
|
||||
return reason
|
||||
}
|
||||
if info != nil && info.ResponseModel != nil && info.ResponseModel.ReturnedModel != "" {
|
||||
if info.ResponseModel.Mismatch {
|
||||
return "mismatch_allowed"
|
||||
}
|
||||
return "accepted"
|
||||
}
|
||||
if info != nil && info.AstraAutoSwitch {
|
||||
return "accepted"
|
||||
}
|
||||
return "unchecked"
|
||||
}
|
||||
|
||||
func recordDurableAttempt(c *gin.Context, info *relaycommon.RelayInfo, usage *dto.Usage, params model.RecordConsumeLogParams) {
|
||||
number := c.GetInt("relay_attempt_number")
|
||||
reason := resolveAttemptGuardResult(c.GetString("attempt_guard_result"), info)
|
||||
if reason == "model_mismatch" {
|
||||
reason = "mismatch_rejected"
|
||||
}
|
||||
returned := c.GetString("attempt_returned_model")
|
||||
if returned == "" && info.ResponseModel != nil {
|
||||
returned = info.ResponseModel.ReturnedModel
|
||||
}
|
||||
source := c.GetString("attempt_usage_source")
|
||||
if source == "" {
|
||||
source = "upstream"
|
||||
if usage != nil && usage.UsageSource == "estimated" {
|
||||
source = "estimated"
|
||||
}
|
||||
}
|
||||
status := "verified"
|
||||
if reason == "unchecked" {
|
||||
status = "unchecked"
|
||||
} else if reason == "mismatch_allowed" {
|
||||
status = "mismatch_allowed"
|
||||
} else if reason == "transport_error" {
|
||||
status = "transport_error"
|
||||
} else if reason != "accepted" {
|
||||
status = "rejected"
|
||||
if params.Quota > 0 {
|
||||
status = "charged_rejection"
|
||||
}
|
||||
} else if number > 1 {
|
||||
status = "switched_success"
|
||||
}
|
||||
params.Other.SetPublic("routing_status", status)
|
||||
params.Other.SetPublic("attempt_usage_source", source)
|
||||
params.Other.SetPublic("routing_attempt", map[string]any{"request_id": info.RequestId, "number": number, "guard_result": reason, "usage_source": source, "requested_model": info.OriginModelName, "returned_model": returned})
|
||||
params.Other.SetPublic("response_model", map[string]any{"requested_model": info.OriginModelName, "upstream_model": info.UpstreamModelName, "returned_model": returned, "mismatch": returned != "" && returned != info.OriginModelName})
|
||||
if reason != "accepted" && reason != "unchecked" {
|
||||
params.Other.SetPublic("model_guard_result", reason)
|
||||
}
|
||||
log := &model.Log{UserId: info.UserId, Username: common.GetContextKeyString(c, constant.ContextKeyUserName), CreatedAt: time.Now().Unix(), Type: model.LogTypeConsume, Content: params.Content, TokenName: params.TokenName, ModelName: params.ModelName, Quota: params.Quota, PromptTokens: params.PromptTokens, CompletionTokens: params.CompletionTokens, UseTime: params.UseTimeSeconds, IsStream: params.IsStream, ChannelId: params.ChannelId, TokenId: params.TokenId, Group: params.Group, RequestId: info.RequestId, Other: params.Other.JSONString(), CorrelationKey: c.GetString("log_correlation_key")}
|
||||
billable := params.Quota > 0 || params.PromptTokens > 0 || params.CompletionTokens > 0
|
||||
usageBody, _ := common.Marshal(usage)
|
||||
logBody, _ := common.Marshal(log)
|
||||
deadline := time.Now().Add(90 * time.Second).Unix()
|
||||
if c.GetBool("attempt_checkpoint") {
|
||||
deadline = time.Now().Add(10 * time.Minute).Unix()
|
||||
}
|
||||
if c.GetBool("attempt_checkpoint") {
|
||||
if err := model.CheckpointRelayAttempt(info.DurableAttemptID, params.Quota, string(usageBody), string(logBody), returned, reason, source, billable, deadline); err != nil {
|
||||
common.SysError("attempt checkpoint: " + err.Error())
|
||||
}
|
||||
return
|
||||
}
|
||||
details := model.RelayAttempt{UsageJSON: string(usageBody), ReturnedModel: returned, GuardResult: reason, UsageSource: source, Billable: billable}
|
||||
if err := model.SettleRelayAttempt(info.DurableAttemptID, params.Quota, log, details); err != nil {
|
||||
common.SysError("attempt settlement: " + err.Error())
|
||||
info.AttemptDetached = true
|
||||
// Retain the authoritative usage in memory during transient DB outages;
|
||||
// restart recovery can independently use the last durable checkpoint.
|
||||
id := info.DurableAttemptID
|
||||
quota := params.Quota
|
||||
go func() {
|
||||
for range 60 {
|
||||
time.Sleep(time.Second)
|
||||
if model.SettleRelayAttempt(id, quota, log, details) == nil {
|
||||
return
|
||||
}
|
||||
}
|
||||
}()
|
||||
return
|
||||
}
|
||||
if err := model.DeliverRelayAttemptLogs(); err != nil {
|
||||
common.SysError("attempt log outbox: " + err.Error())
|
||||
}
|
||||
}
|
||||
|
||||
type GuardUsageCollector struct {
|
||||
ctx *gin.Context
|
||||
info relaycommon.RelayInfo
|
||||
usage dto.Usage
|
||||
outputTokens int
|
||||
observed bool
|
||||
hasUsage bool
|
||||
lastCheckpoint time.Time
|
||||
}
|
||||
|
||||
func NewGuardUsageCollector(c *gin.Context, info *relaycommon.RelayInfo) *GuardUsageCollector {
|
||||
copyInfo := *info
|
||||
if info.ChannelMeta != nil {
|
||||
meta := *info.ChannelMeta
|
||||
copyInfo.ChannelMeta = &meta
|
||||
}
|
||||
if info.TieredBillingSnapshot != nil {
|
||||
snap := *info.TieredBillingSnapshot
|
||||
copyInfo.TieredBillingSnapshot = &snap
|
||||
}
|
||||
copyInfo.ResponseModel = nil
|
||||
copyInfo.ResponseGuard = nil
|
||||
copyInfo.ResponsesUsageInfo = nil
|
||||
copyInfo.StreamStatus = relaycommon.NewStreamStatus()
|
||||
copyContext := c.Copy()
|
||||
state := RequestPolicy(c)
|
||||
copyContext.Set(requestPolicyContextKey, &RequestPolicyState{StartedAt: state.StartedAt, Attempts: state.Attempts, SelectedGroup: state.SelectedGroup, SessionMode: state.SessionMode, SessionModeSource: state.SessionModeSource, RuleName: state.RuleName, events: state.Events()})
|
||||
return &GuardUsageCollector{ctx: copyContext, info: copyInfo}
|
||||
}
|
||||
|
||||
func (collector *GuardUsageCollector) Observe(payload []byte, stream bool) {
|
||||
var event dto.ResponsesStreamResponse
|
||||
if stream {
|
||||
if common.Unmarshal(payload, &event) != nil {
|
||||
return
|
||||
}
|
||||
} else {
|
||||
var response dto.OpenAIResponsesResponse
|
||||
if common.Unmarshal(payload, &response) != nil {
|
||||
return
|
||||
}
|
||||
event.Type = "response.completed"
|
||||
event.Response = &response
|
||||
}
|
||||
if event.Response != nil {
|
||||
collector.info.ObserveResponseModel(event.Response.Model)
|
||||
if event.Response.Model != "" {
|
||||
collector.ctx.Set("attempt_returned_model", event.Response.Model)
|
||||
if event.Response.Model != collector.info.OriginModelName {
|
||||
collector.ctx.Set("attempt_guard_result", "model_mismatch")
|
||||
}
|
||||
}
|
||||
if event.Response.Usage != nil && (event.Type == "response.completed" || event.Type == "response.incomplete" || event.Type == "response.failed" || event.Type == "response.cancelled") {
|
||||
ApplyResponsesUsage(&collector.usage, event.Response.Usage)
|
||||
collector.hasUsage = true
|
||||
}
|
||||
if collector.outputTokens == 0 {
|
||||
text := ExtractOutputTextFromResponses(event.Response)
|
||||
if text != "" {
|
||||
collector.outputTokens = CountTextToken(text, collector.info.OriginModelName)
|
||||
collector.observed = true
|
||||
}
|
||||
}
|
||||
}
|
||||
if strings.HasSuffix(event.Type, ".delta") && event.Delta != "" {
|
||||
collector.observed = true
|
||||
collector.outputTokens += CountTextToken(event.Delta, collector.info.OriginModelName)
|
||||
}
|
||||
if collector.hasUsage || collector.observed || collector.ctx.GetString("attempt_returned_model") != "" && collector.ctx.GetString("attempt_returned_model") != collector.info.OriginModelName {
|
||||
if time.Since(collector.lastCheckpoint) > time.Second || collector.hasUsage {
|
||||
collector.persist(true)
|
||||
collector.lastCheckpoint = time.Now()
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (collector *GuardUsageCollector) persist(checkpoint bool) {
|
||||
usage := collector.usage
|
||||
if !collector.hasUsage {
|
||||
usage.PromptTokens = collector.info.GetEstimatePromptTokens()
|
||||
usage.CompletionTokens = collector.outputTokens
|
||||
usage.TotalTokens = usage.PromptTokens + usage.CompletionTokens
|
||||
collector.ctx.Set("attempt_usage_source", "estimated")
|
||||
} else {
|
||||
collector.ctx.Set("attempt_usage_source", "upstream")
|
||||
}
|
||||
collector.ctx.Set("attempt_checkpoint", checkpoint)
|
||||
PostTextConsumeQuota(collector.ctx, &collector.info, &usage, nil)
|
||||
}
|
||||
|
||||
// Flush the most recent observed output when a stream ends between periodic
|
||||
// checkpoints, including disconnects and late guard rejection.
|
||||
func (collector *GuardUsageCollector) CheckpointFinal(reason string) {
|
||||
if !collector.hasUsage && !collector.observed && reason != "model_mismatch" {
|
||||
return
|
||||
}
|
||||
if reason != "accepted" && reason != "pending" {
|
||||
collector.ctx.Set("attempt_guard_result", reason)
|
||||
}
|
||||
collector.persist(true)
|
||||
}
|
||||
|
||||
func (collector *GuardUsageCollector) FinishRejected(info *relaycommon.RelayInfo, guard *relaycommon.ResponseGuardAttempt, resp *http.Response) bool {
|
||||
if info.DurableAttemptID == 0 {
|
||||
return false
|
||||
}
|
||||
collector.ctx.Set("attempt_guard_result", guard.Reason)
|
||||
if guard.Reason != "model_mismatch" && !collector.hasUsage && !collector.observed {
|
||||
return false
|
||||
}
|
||||
info.AttemptDetached = true
|
||||
collector.persist(true)
|
||||
guard.Detach()
|
||||
go func() {
|
||||
defer func() {
|
||||
if recovered := recover(); recovered != nil {
|
||||
common.SysError(fmt.Sprintf("attempt usage collection panic: %v", recovered))
|
||||
}
|
||||
}()
|
||||
defer guard.Close()
|
||||
if resp != nil && resp.Body != nil {
|
||||
defer resp.Body.Close()
|
||||
}
|
||||
defer collector.persist(false)
|
||||
if guard.Drain != nil {
|
||||
guard.Drain()
|
||||
}
|
||||
}()
|
||||
return true
|
||||
}
|
||||
|
||||
func EndFailedRelayAttempt(c *gin.Context, info *relaycommon.RelayInfo) {
|
||||
if info.DurableAttemptID == 0 || info.AttemptDetached {
|
||||
return
|
||||
}
|
||||
var attempt model.RelayAttempt
|
||||
if err := model.DB.First(&attempt, info.DurableAttemptID).Error; err != nil {
|
||||
common.SysError(err.Error())
|
||||
return
|
||||
}
|
||||
if attempt.State == "settled" {
|
||||
return
|
||||
}
|
||||
if !attempt.Billable {
|
||||
reason := "failed"
|
||||
if guard := info.ResponseGuard.Last(); guard != nil {
|
||||
reason = guard.Reason
|
||||
}
|
||||
_ = model.CheckpointRelayAttempt(attempt.ID, 0, "{}", "", attempt.ReturnedModel, reason, "", false, time.Now().Unix())
|
||||
}
|
||||
quota := 0
|
||||
var log *model.Log
|
||||
if attempt.Billable {
|
||||
quota = attempt.ObservedQuota
|
||||
log = &model.Log{}
|
||||
if common.UnmarshalJsonStr(attempt.LogJSON, log) != nil {
|
||||
log = nil
|
||||
}
|
||||
}
|
||||
if err := model.SettleRelayAttempt(attempt.ID, quota, log); err != nil {
|
||||
common.SysError(err.Error())
|
||||
}
|
||||
}
|
||||
|
||||
func StartRelayAttemptRecovery() {
|
||||
go func() {
|
||||
for {
|
||||
var attempts []model.RelayAttempt
|
||||
if err := model.DB.Where("state <> ? AND deadline < ?", "settled", time.Now().Unix()).Limit(100).Find(&attempts).Error; err == nil {
|
||||
for _, attempt := range attempts {
|
||||
quota := 0
|
||||
var log *model.Log
|
||||
if attempt.Billable {
|
||||
quota = attempt.ObservedQuota
|
||||
log = &model.Log{}
|
||||
if common.UnmarshalJsonStr(attempt.LogJSON, log) != nil {
|
||||
continue
|
||||
}
|
||||
}
|
||||
if err := model.SettleRelayAttempt(attempt.ID, quota, log); err != nil {
|
||||
common.SysError("attempt recovery: " + err.Error())
|
||||
}
|
||||
}
|
||||
}
|
||||
if err := model.DeliverRelayAttemptLogs(); err != nil {
|
||||
common.SysError("attempt outbox: " + err.Error())
|
||||
}
|
||||
time.Sleep(5 * time.Second)
|
||||
}
|
||||
}()
|
||||
}
|
||||
Reference in new issue
Block a user