fix: fix token prediction accuracy and response header issues

Author: dancing-ui

core/llm_token_ratelimit/ratelimit_checker.go

@@ -177,13 +177,15 @@ func (c *PETAChecker) checkLimitKey(ctx *Context, rule *MatchedRule) bool {
 		)
 		return false
 	}
+	defer func() {
+		ctx.Set(KeyResponseHeaders, responseHeader)
+	}()
 	responseHeader.Set(KeyRequestID, ctx.Get(KeyRequestID).(string))
 	if waitingTime != PETANoWaiting {
 		responseHeader.Set(ResponseHeaderRemainingTokens, fmt.Sprintf("%d", result[0]))
 		responseHeader.Set(ResponseHeaderWaitingTime, (time.Duration(waitingTime) * time.Millisecond).String())
 		return false
 	}
-	ctx.Set(KeyResponseHeaders, responseHeader)
 	c.cacheEstimatedToken(rule, result[2])
 	return true
 }

core/llm_token_ratelimit/script/peta/correct.lua

@@ -90,12 +90,7 @@ if released_tokens > 0 then -- Expired tokens exist, attempt to replenish new to
 end
 -- Update the difference from the token encoder
 local difference = actual - estimated
-local ttl = redis.call('PTTL', token_encoder_key)
-if ttl < 0 then
-    redis.call('SET', token_encoder_key, difference, 'PX', window_size + 5000)
-else
-    redis.call('INCRBY', token_encoder_key, difference)
-end
+redis.call('SET', token_encoder_key, difference, 'PX', window_size + 5000)
 -- Correction result for reservation
 local correct_result = 0
 if estimated < 0 or actual < 0 then