Skip to content

Commit 1cb6955

Browse files
dancing-uidancing-ui
committed
feat: optimize response header information
1 parent 2372981 commit 1cb6955

File tree

11 files changed

+306
-41
lines changed

11 files changed

+306
-41
lines changed

‎.github/workflows/go.yml‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -100,6 +100,7 @@ jobs:
100100
sleep-duration: 5
101101

102102
- name: Test Redis Cluster
103+
if: matrix.go_version == '1.22.0'
103104
run: |
104105
sudo apt-get install -y redis-tools
105106
docker ps -a
@@ -115,6 +116,8 @@ jobs:
115116
run: |
116117
cd ./pkg/adapters/langchaingo
117118
go test -race -count=1 ./... -coverprofile=coverage.txt -covermode=atomic
119+
cd ../eino
120+
go test -race -count=1 ./... -coverprofile=coverage.txt -covermode=atomic
118121
119122
- name: Coverage
120123
run: bash <(curl -s https://codecov.io/bash)

‎core/llm_token_ratelimit/config.go‎

Lines changed: 16 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -244,30 +244,33 @@ func (c *SafeConfig) GetConfig() *Config {
244244
}
245245

246246
func (c *SafeConfig) IsEnabled() bool {
247-
cfg := c.GetConfig()
248-
if cfg == nil {
247+
if c == nil || c.config == nil {
249248
logging.Error(errors.New("safe config is nil"), "found safe config is nil")
250249
return false
251250
}
252-
return cfg.Enabled
251+
c.mu.RLock()
252+
defer c.mu.RUnlock()
253+
return c.config.Enabled
253254
}
254255

255-
func GetErrorCode() int32 {
256-
cfg := globalConfig.GetConfig()
257-
if cfg == nil {
256+
func (c *SafeConfig) GetErrorCode() int32 {
257+
if c == nil || c.config == nil {
258258
logging.Error(errors.New("safe config is nil"), "found safe config is nil")
259-
return -1
259+
return DefaultErrorCode
260260
}
261-
return cfg.ErrorCode
261+
c.mu.RLock()
262+
defer c.mu.RUnlock()
263+
return c.config.ErrorCode
262264
}
263265

264-
func GetErrorMsg() string {
265-
cfg := globalConfig.GetConfig()
266-
if cfg == nil {
266+
func (c *SafeConfig) GetErrorMsg() string {
267+
if c == nil || c.config == nil {
267268
logging.Error(errors.New("safe config is nil"), "found safe config is nil")
268-
return ""
269+
return DefaultErrorMessage
269270
}
270-
return cfg.ErrorMessage
271+
c.mu.RLock()
272+
defer c.mu.RUnlock()
273+
return c.config.ErrorMessage
271274
}
272275

273276
func (p *TokenEncoderProvider) UnmarshalYAML(unmarshal func(interface{}) error) error {

‎core/llm_token_ratelimit/constant.go‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -48,6 +48,8 @@ const (
4848
KeyMatchedRules string = "SentinelLLMTokenRatelimitMatchedRules"
4949
KeyResponseHeaders string = "SentinelLLMTokenRatelimitResponseHeaders"
5050
KeyRequestID string = "SentinelLLMTokenRatelimitRequestID"
51+
KeyErrorCode string = "SentinelLLMTokenRatelimitErrorCode"
52+
KeyErrorMessage string = "SentinelLLMTokenRatelimitErrorMessage"
5153
)
5254

5355
// ================================= RedisRatelimitKeyFormat ==================
@@ -57,6 +59,7 @@ const (
5759

5860
// ================================= ResponseHeader ==================
5961
const (
62+
ResponseHeaderRequestID string = "X-Sentinel-LLM-Token-Ratelimit-RequestID"
6063
ResponseHeaderRemainingTokens string = "X-Sentinel-LLM-Token-Ratelimit-RemainingTokens"
6164
ResponseHeaderWaitingTime string = "X-Sentinel-LLM-Token-Ratelimit-WaitingTime"
6265
)

‎core/llm_token_ratelimit/ratelimit_checker.go‎

Lines changed: 16 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -85,13 +85,18 @@ func (c *FixedWindowChecker) checkLimitKey(ctx *Context, rule *MatchedRule) bool
8585
defer func() {
8686
ctx.Set(KeyResponseHeaders, responseHeader)
8787
}()
88-
responseHeader.Set(KeyRequestID, ctx.Get(KeyRequestID).(string))
88+
// set response headers
89+
responseHeader.Set(ResponseHeaderRequestID, ctx.Get(KeyRequestID).(string))
90+
responseHeader.Set(ResponseHeaderRemainingTokens, fmt.Sprintf("%d", remaining))
8991
if remaining < 0 {
90-
responseHeader.Set(ResponseHeaderRemainingTokens, fmt.Sprintf("%d", remaining))
92+
// set waiting time in milliseconds
9193
responseHeader.Set(ResponseHeaderWaitingTime, (time.Duration(result[1]) * time.Millisecond).String())
94+
// set error code and message
95+
responseHeader.ErrorCode = globalConfig.GetErrorCode()
96+
responseHeader.ErrorMessage = globalConfig.GetErrorMsg()
97+
// reject the request
9298
return false
9399
}
94-
ctx.Set(KeyResponseHeaders, responseHeader)
95100
return true
96101
}
97102

@@ -183,10 +188,16 @@ func (c *PETAChecker) checkLimitKey(ctx *Context, rule *MatchedRule) bool {
183188
defer func() {
184189
ctx.Set(KeyResponseHeaders, responseHeader)
185190
}()
186-
responseHeader.Set(KeyRequestID, ctx.Get(KeyRequestID).(string))
191+
// set response headers
192+
responseHeader.Set(ResponseHeaderRequestID, ctx.Get(KeyRequestID).(string))
193+
responseHeader.Set(ResponseHeaderRemainingTokens, fmt.Sprintf("%d", result[0]))
187194
if waitingTime != PETANoWaiting {
188-
responseHeader.Set(ResponseHeaderRemainingTokens, fmt.Sprintf("%d", result[0]))
195+
// set waiting time in milliseconds
189196
responseHeader.Set(ResponseHeaderWaitingTime, (time.Duration(waitingTime) * time.Millisecond).String())
197+
// set error code and message
198+
responseHeader.ErrorCode = globalConfig.GetErrorCode()
199+
responseHeader.ErrorMessage = globalConfig.GetErrorMsg()
200+
// reject the request
190201
return false
191202
}
192203
c.cacheEstimatedToken(rule, result[2])

‎core/llm_token_ratelimit/response_header.go‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -15,7 +15,9 @@
1515
package llmtokenratelimit
1616

1717
type ResponseHeader struct {
18-
headers map[string]string
18+
headers map[string]string
19+
ErrorCode int32
20+
ErrorMessage string
1921
}
2022

2123
func NewResponseHeader() *ResponseHeader {

‎example/llm_token_ratelimit/ratelimit/ratelimit.go‎

Lines changed: 13 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -110,28 +110,34 @@ func SentinelMiddleware(opts ...Option) gin.HandlerFunc {
110110
reqInfos = options.requestInfosExtract(c)
111111
}
112112

113-
// check
113+
// Check
114114
entry, err := sentinel.Entry(resource, sentinel.WithTrafficType(base.Inbound), sentinel.WithArgs(reqInfos))
115-
116115
if err != nil {
117116
// Block
118117
if options.blockFallback != nil {
119118
options.blockFallback(c)
120119
} else {
121-
setResponseHeaders(c, err.TriggeredValue().(*llmtokenratelimit.ResponseHeader))
122-
c.AbortWithStatusJSON(int(llmtokenratelimit.GetErrorCode()), gin.H{
123-
"error": llmtokenratelimit.GetErrorMsg(),
120+
responseHeader, ok := err.TriggeredValue().(*llmtokenratelimit.ResponseHeader)
121+
if !ok || responseHeader == nil {
122+
c.AbortWithStatusJSON(500, gin.H{
123+
"error": "internal server error. invalid response header.",
124+
})
125+
return
126+
}
127+
setResponseHeaders(c, responseHeader)
128+
c.AbortWithStatusJSON(int(responseHeader.ErrorCode), gin.H{
129+
"error": responseHeader.ErrorMessage,
124130
})
125131
}
126132
return
127133
}
128-
// Pass or Disabled
129-
c.Next()
130134
// Set response headers
131135
responseHeader, ok := entry.Context().GetPair(llmtokenratelimit.KeyResponseHeaders).(*llmtokenratelimit.ResponseHeader)
132136
if ok && responseHeader != nil {
133137
setResponseHeaders(c, responseHeader)
134138
}
139+
// Pass or Disabled
140+
c.Next()
135141
// Update used token info
136142
usedTokenInfos, exists := c.Get(llmtokenratelimit.KeyUsedTokenInfos)
137143
if exists && usedTokenInfos != nil {

‎pkg/adapters/eino/go.mod‎

Lines changed: 8 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,7 @@ replace github.com/alibaba/sentinel-golang => ../../../
77
require (
88
github.com/alibaba/sentinel-golang v1.0.4
99
github.com/cloudwego/eino v0.4.8
10+
github.com/cloudwego/eino-ext/components/model/openai v0.0.0-20250905035413-86dbae6351d5
1011
)
1112

1213
require (
@@ -18,10 +19,12 @@ require (
1819
github.com/bytedance/sonic/loader v0.3.0 // indirect
1920
github.com/cespare/xxhash/v2 v2.2.0 // indirect
2021
github.com/cloudwego/base64x v0.1.5 // indirect
22+
github.com/cloudwego/eino-ext/libs/acl/openai v0.0.0-20250826113018-8c6f6358d4bb // indirect
2123
github.com/dgryski/go-rendezvous v0.0.0-20200823014737-9f7001d12a5f // indirect
2224
github.com/dlclark/regexp2 v1.10.0 // indirect
2325
github.com/dustin/go-humanize v1.0.1 // indirect
2426
github.com/eino-contrib/jsonschema v1.0.0 // indirect
27+
github.com/evanphx/json-patch v0.5.2 // indirect
2528
github.com/getkin/kin-openapi v0.118.0 // indirect
2629
github.com/go-ole/go-ole v1.2.4 // indirect
2730
github.com/go-openapi/jsonpointer v0.19.5 // indirect
@@ -37,6 +40,7 @@ require (
3740
github.com/kr/text v0.2.0 // indirect
3841
github.com/mailru/easyjson v0.7.7 // indirect
3942
github.com/matttproud/golang_protobuf_extensions v1.0.4 // indirect
43+
github.com/meguminnnnnnnnn/go-openai v0.0.0-20250821095446-07791bea23a0 // indirect
4044
github.com/modern-go/concurrent v0.0.0-20180306012644-bacd9c7ef1dd // indirect
4145
github.com/modern-go/reflect2 v1.0.2 // indirect
4246
github.com/mohae/deepcopy v0.0.0-20170929034955-c48cc78d4826 // indirect
@@ -53,16 +57,18 @@ require (
5357
github.com/sirupsen/logrus v1.9.3 // indirect
5458
github.com/slongfield/pyfmt v0.0.0-20220222012616-ea85ff4c361f // indirect
5559
github.com/spaolacci/murmur3 v1.1.0 // indirect
56-
github.com/stretchr/testify v1.10.0 // indirect
60+
github.com/stretchr/objx v0.5.2 // indirect
5761
github.com/tklauser/go-sysconf v0.3.6 // indirect
5862
github.com/tklauser/numcpus v0.2.2 // indirect
5963
github.com/twitchyliquid64/golang-asm v0.15.1 // indirect
6064
github.com/wk8/go-ordered-map/v2 v2.1.8 // indirect
6165
github.com/yargevad/filepathx v1.0.0 // indirect
6266
golang.org/x/arch v0.11.0 // indirect
63-
golang.org/x/crypto v0.31.0 // indirect
6467
golang.org/x/exp v0.0.0-20230713183714-613f0c0eb8a1 // indirect
68+
golang.org/x/net v0.21.0 // indirect
6569
golang.org/x/sys v0.28.0 // indirect
70+
golang.org/x/term v0.27.0 // indirect
71+
golang.org/x/text v0.21.0 // indirect
6672
google.golang.org/protobuf v1.30.0 // indirect
6773
gopkg.in/yaml.v2 v2.4.0 // indirect
6874
gopkg.in/yaml.v3 v3.0.1 // indirect

‎pkg/adapters/eino/go.sum‎

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,8 @@ github.com/buger/jsonparser v1.1.1 h1:2PnMjfWD7wBILjqQbt530v576A/cAbQvEW9gGIpYMU
1111
github.com/buger/jsonparser v1.1.1/go.mod h1:6RYKKt7H4d4+iWqouImQ9R2FZql3VbhNgx27UK13J/0=
1212
github.com/bugsnag/bugsnag-go v1.4.0/go.mod h1:2oa8nejYd4cQ/b0hMIopN0lCRxU0bueqREvZLWFrtK8=
1313
github.com/bugsnag/panicwrap v1.2.0/go.mod h1:D/8v3kj0zr8ZAKg1AQ6crr+5VwKN5eIywRkfhyM/+dE=
14+
github.com/bytedance/mockey v1.2.14 h1:KZaFgPdiUwW+jOWFieo3Lr7INM1P+6adO3hxZhDswY8=
15+
github.com/bytedance/mockey v1.2.14/go.mod h1:1BPHF9sol5R1ud/+0VEHGQq/+i2lN+GTsr3O2Q9IENY=
1416
github.com/bytedance/sonic v1.14.0 h1:/OfKt8HFw0kh2rj8N0F6C/qPGRESq0BbaNZgcNXXzQQ=
1517
github.com/bytedance/sonic v1.14.0/go.mod h1:WoEbx8WTcFJfzCe0hbmyTGrfjt8PzNEBdxlNUO24NhA=
1618
github.com/bytedance/sonic/loader v0.1.1/go.mod h1:ncP89zfokxS5LZrJxl5z0UJcsk4M4yY2JpfqGeCtNLU=
@@ -24,6 +26,10 @@ github.com/cloudwego/base64x v0.1.5 h1:XPciSp1xaq2VCSt6lF0phncD4koWyULpl5bUxbfCy
2426
github.com/cloudwego/base64x v0.1.5/go.mod h1:0zlkT4Wn5C6NdauXdJRhSKRlJvmclQ1hhJgA0rcu/8w=
2527
github.com/cloudwego/eino v0.4.8 h1:wptTU24tQad1mFCHw0+4zSzH+p8dLEBk6HtggPlcvP0=
2628
github.com/cloudwego/eino v0.4.8/go.mod h1:1TDlOmwGSsbCJaWB92w9YLZi2FL0WRZoRcD4eMvqikg=
29+
github.com/cloudwego/eino-ext/components/model/openai v0.0.0-20250905035413-86dbae6351d5 h1:D04jOL3xKn9CstVRaPkunhPffIbEVp0VfQV5e/26lS8=
30+
github.com/cloudwego/eino-ext/components/model/openai v0.0.0-20250905035413-86dbae6351d5/go.mod h1:QQhCuQxuBAVWvu/YAZBhs/RsR76mUigw59Tl0kh04C8=
31+
github.com/cloudwego/eino-ext/libs/acl/openai v0.0.0-20250826113018-8c6f6358d4bb h1:RMslzyijc3bi9EkqCulpS0hZupTl1y/wayR3+fVRN/c=
32+
github.com/cloudwego/eino-ext/libs/acl/openai v0.0.0-20250826113018-8c6f6358d4bb/go.mod h1:fHn/6OqPPY1iLLx9wzz+MEVT5Dl9gwuZte1oLEnCoYw=
2733
github.com/cloudwego/iasm v0.2.0/go.mod h1:8rXZaNYT2n95jn+zTI1sDr+IgcD2GVs0nlbbQPiEFhY=
2834
github.com/creack/pty v1.1.9/go.mod h1:oKZEueFk5CKHvIhNR5MUki03XCEU+Q6VDXinZuGJ33E=
2935
github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
@@ -37,6 +43,8 @@ github.com/dustin/go-humanize v1.0.1 h1:GzkhY7T5VNhEkwH0PVJgjz+fX1rhBrR7pRT3mDkp
3743
github.com/dustin/go-humanize v1.0.1/go.mod h1:Mu1zIs6XwVuF/gI1OepvI0qD18qycQx+mFykh5fBlto=
3844
github.com/eino-contrib/jsonschema v1.0.0 h1:dXxbhGNZuI3+xNi8x3JT8AGyoXz6Pff6mRvmpjVl5Ww=
3945
github.com/eino-contrib/jsonschema v1.0.0/go.mod h1:cpnX4SyKjWjGC7iN2EbhxaTdLqGjCi0e9DxpLYxddD4=
46+
github.com/evanphx/json-patch v0.5.2 h1:xVCHIVMUu1wtM/VkR9jVZ45N3FhZfYMMYGorLCR8P3k=
47+
github.com/evanphx/json-patch v0.5.2/go.mod h1:ZWS5hhDbVDyob71nXKNL0+PWn6ToqBHMikGIFbs31qQ=
4048
github.com/fsnotify/fsnotify v1.4.7/go.mod h1:jwhsz4b93w/PPRr/qN1Yymfu8t87LnFCMoQvtojpjFo=
4149
github.com/fsnotify/fsnotify v1.4.9 h1:hsms1Qyu0jgnwNXIxa+/V/PDsU6CfLf6CNO8H7IWoS4=
4250
github.com/fsnotify/fsnotify v1.4.9/go.mod h1:znqG4EE+3YCdAaPaxE2ZRY/06pZUdp0tY4IgpuI1SZQ=
@@ -85,6 +93,7 @@ github.com/gorilla/mux v1.8.0/go.mod h1:DVbg23sWSpFRCP0SfiEN6jmj59UnW/n46BH5rLB7
8593
github.com/hpcloud/tail v1.0.0/go.mod h1:ab1qPbhIpdTxEkNHXyeSf5vhxWSCs/tWer42PpOxQnU=
8694
github.com/invopop/yaml v0.1.0 h1:YW3WGUoJEXYfzWBjn00zIlrw7brGVD0fUKRYDPAPhrc=
8795
github.com/invopop/yaml v0.1.0/go.mod h1:2XuRLgs/ouIrW3XNzuNj7J3Nvu/Dig5MXvbCEdiBN3Q=
96+
github.com/jessevdk/go-flags v1.4.0/go.mod h1:4FA24M0QyGHXBuZZK/XkWh8h0e1EYbRYJSGM75WSRxI=
8897
github.com/josharian/intern v1.0.0 h1:vlS4z54oSdjm0bgjRigI+G1HpF+tI+9rE5LLzOg8HmY=
8998
github.com/josharian/intern v1.0.0/go.mod h1:5DoeVV0s6jJacbCEi61lwdGj/aVlrQvzHFFd8Hwg//Y=
9099
github.com/json-iterator/go v1.1.12 h1:PV8peI4a0ysnczrg+LtxykD8LfKY9ML6u2jnxaEnrnM=
@@ -113,6 +122,8 @@ github.com/mattn/go-isatty v0.0.8 h1:HLtExJ+uU2HOZ+wI0Tt5DtUDrx8yhUqDcp7fYERX4CE
113122
github.com/mattn/go-isatty v0.0.8/go.mod h1:Iq45c/XA43vh69/j3iqttzPXn0bhXyGjM0Hdxcsrc5s=
114123
github.com/matttproud/golang_protobuf_extensions v1.0.4 h1:mmDVorXM7PCGKw94cs5zkfA9PSy5pEvNWRP0ET0TIVo=
115124
github.com/matttproud/golang_protobuf_extensions v1.0.4/go.mod h1:BSXmuO+STAnVfrANrmjBb36TMTDstsz7MSK+HVaYKv4=
125+
github.com/meguminnnnnnnnn/go-openai v0.0.0-20250821095446-07791bea23a0 h1:nIohpHs1ViKR0SVgW/cbBstHjmnqFZDM9RqgX9m9Xu8=
126+
github.com/meguminnnnnnnnn/go-openai v0.0.0-20250821095446-07791bea23a0/go.mod h1:qs96ysDmxhE4BZoU45I43zcyfnaYxU3X+aRzLko/htY=
116127
github.com/mgutz/ansi v0.0.0-20170206155736-9520e82c474b h1:j7+1HpAFS1zy5+Q4qx1fWh90gTKwiN4QCGoY9TWyyO4=
117128
github.com/mgutz/ansi v0.0.0-20170206155736-9520e82c474b/go.mod h1:01TrycV0kFyexm33Z7vhZRXopbI8J3TDReVlkTgMUxE=
118129
github.com/modern-go/concurrent v0.0.0-20180228061459-e0a39a4cb421/go.mod h1:6dJC0mAP4ikYIbvyc7fijjWJddQyLn8Ig3JB5CqoB9Q=

‎pkg/adapters/eino/options.go‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -77,6 +77,9 @@ func evaluateOptions(opts ...Option) *options {
7777
for _, part := range msg.MultiContent {
7878
prompts = append(prompts, part.Text)
7979
}
80+
if len(msg.Content) != 0 {
81+
prompts = append(prompts, msg.Content)
82+
}
8083
}
8184
return prompts
8285
},

0 commit comments

Comments
 (0)