352 lines
13 KiB
Go
352 lines
13 KiB
Go
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"
|
|
"github.com/QuantumNous/new-api/tokenkit"
|
|
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.Other.SetContent(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 = tokenkit.Count(collector.info.OriginModelName, text)
|
|
collector.observed = true
|
|
}
|
|
}
|
|
}
|
|
if strings.HasSuffix(event.Type, ".delta") && event.Delta != "" {
|
|
collector.observed = true
|
|
collector.outputTokens += tokenkit.Count(collector.info.OriginModelName, event.Delta)
|
|
}
|
|
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)
|
|
}
|
|
}()
|
|
}
|