superidou
5759c1862e
M0-M7 已完成:核心网关(身份/RBAC/TOTP/OIDC/SAML/Provider/配额/路由/内容策略/审计/定价)+ 资源市场(MCP/Skills/数字员工)。 含 22 个 PostgreSQL 迁移、管理端/门户端前端源码、OpenAPI 契约、部署 compose。 Co-Authored-By: Claude <noreply@anthropic.com>
47 lines
1.2 KiB
Go
47 lines
1.2 KiB
Go
package gateway
|
|
|
|
import (
|
|
"bytes"
|
|
"encoding/json"
|
|
"errors"
|
|
"io"
|
|
"net/http"
|
|
)
|
|
|
|
var errRequestBodyTooLarge = errors.New("request body is too large")
|
|
|
|
func prepareTokenBudget(request *http.Request, maxBody int64) (int64, error) {
|
|
if request.Body == nil || request.Method == http.MethodGet {
|
|
return 0, nil
|
|
}
|
|
limited := io.LimitReader(request.Body, maxBody+1)
|
|
body, err := io.ReadAll(limited)
|
|
if err != nil {
|
|
return 0, err
|
|
}
|
|
_ = request.Body.Close()
|
|
if int64(len(body)) > maxBody {
|
|
return 0, errRequestBodyTooLarge
|
|
}
|
|
restoreRequestBody(request, body)
|
|
|
|
inputEstimate := max(int64((len(body)+3)/4), 1)
|
|
if request.URL.Path == "/v1/embeddings" {
|
|
return inputEstimate, nil
|
|
}
|
|
outputBudget := int64(1024)
|
|
decoder := json.NewDecoder(bytes.NewReader(body))
|
|
decoder.UseNumber()
|
|
var payload map[string]any
|
|
if decoder.Decode(&payload) == nil {
|
|
for _, field := range []string{"max_output_tokens", "max_completion_tokens", "max_tokens"} {
|
|
if value := jsonInt64(payload[field]); value > 0 {
|
|
outputBudget = value
|
|
break
|
|
}
|
|
}
|
|
}
|
|
// A per-request ceiling prevents malicious payloads from reserving unbounded Redis counters.
|
|
return min(inputEstimate+outputBudget, int64(100_000_000)), nil
|
|
}
|