Skip to content

Commit 2cfb2bb

Browse files
dancing-uidancing-ui
committed
feat: adapt to eino framework; fix initialization issues
1 parent 11896b8 commit 2cfb2bb

File tree

25 files changed

+1345
-748
lines changed

25 files changed

+1345
-748
lines changed

‎.github/workflows/go.yml‎

Lines changed: 31 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -37,6 +37,7 @@ jobs:
3737
go_version:
3838
- 1.18
3939
- 1.19
40+
- 1.22.0
4041
os:
4142
- ubuntu-latest
4243

@@ -62,6 +63,7 @@ jobs:
6263
restore-keys: ${{ runner.os }}-go-
6364

6465
- name: Test
66+
if: matrix.go_version != '1.22.0'
6567
run: |
6668
go test -v -race ./... -coverprofile=coverage.txt -covermode=atomic
6769
cd ./pkg/datasource/consul
@@ -84,6 +86,35 @@ jobs:
8486
go test -race -count=1 ./... -coverprofile=coverage.txt -covermode=atomic
8587
cd ../micro
8688
go test -race -count=1 ./... -coverprofile=coverage.txt -covermode=atomic
89+
90+
- name: Setup Redis Cluster
91+
if: matrix.go_version == '1.22.0'
92+
uses: vishnudxb/redis-cluster@1.0.9
93+
with:
94+
master1-port: 6379
95+
master2-port: 6380
96+
master3-port: 6381
97+
slave1-port: 6382
98+
slave2-port: 6383
99+
slave3-port: 6384
100+
sleep-duration: 5
101+
102+
- name: Test Redis Cluster
103+
run: |
104+
sudo apt-get install -y redis-tools
105+
docker ps -a
106+
redis-cli -h 127.0.0.1 -p 6379 ping
107+
redis-cli -h 127.0.0.1 -p 6379 cluster nodes
108+
109+
- name: LLM Token Ratelimit Adapter Test
110+
if: matrix.go_version == '1.22.0'
111+
env:
112+
LLM_API_KEY: ${{ secrets.LLM_API_KEY }}
113+
LLM_BASE_URL: ${{ secrets.LLM_BASE_URL }}
114+
LLM_MODEL: ${{ secrets.LLM_MODEL }}
115+
run: |
116+
cd ./pkg/ratelimit/langchaingo
117+
go test -race -count=1 ./... -coverprofile=coverage.txt -covermode=atomic
87118
88119
- name: Coverage
89120
run: bash <(curl -s https://codecov.io/bash)

‎api/init.go‎

Lines changed: 2 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -135,10 +135,8 @@ func initCoreComponents() error {
135135
return nil
136136
}
137137

138-
if config.LLMTokenRateLimit() != nil {
139-
if err := llmtokenratelimit.Init(config.LLMTokenRateLimit()); err != nil {
140-
return err
141-
}
138+
if err := llmtokenratelimit.Init(config.LLMTokenRateLimit()); err != nil {
139+
return err
142140
}
143141

144142
return nil

‎core/llm_token_ratelimit/config.go‎

Lines changed: 99 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -112,12 +112,102 @@ type Redis struct {
112112
}
113113

114114
type Config struct {
115+
Enabled bool `json:"enabled" yaml:"enabled"`
115116
Rules []*Rule `json:"rules" yaml:"rules"`
116117
Redis *Redis `json:"redis" yaml:"redis"`
117118
ErrorCode int32 `json:"errorCode" yaml:"errorCode"`
118119
ErrorMessage string `json:"errorMessage" yaml:"errorMessage"`
119120
}
120121

122+
func NewDefaultRedisConfig() *Redis {
123+
return &Redis{
124+
Addrs: []*RedisAddr{
125+
{
126+
Name: DefaultRedisAddrName,
127+
Port: DefaultRedisAddrPort,
128+
},
129+
},
130+
DialTimeout: DefaultRedisTimeout,
131+
ReadTimeout: DefaultRedisTimeout,
132+
WriteTimeout: DefaultRedisTimeout,
133+
PoolTimeout: DefaultRedisTimeout,
134+
PoolSize: DefaultRedisPoolSize,
135+
MinIdleConns: DefaultRedisMinIdleConns,
136+
MaxRetries: DefaultRedisMaxRetries,
137+
}
138+
}
139+
140+
func NewDefaultConfig() *Config {
141+
return &Config{
142+
Enabled: false,
143+
Rules: nil,
144+
Redis: nil,
145+
ErrorCode: DefaultErrorCode,
146+
ErrorMessage: DefaultErrorMessage,
147+
}
148+
}
149+
150+
func (c *Redis) setDefaultConfigOptions() {
151+
if c == nil {
152+
return
153+
}
154+
if len(c.Addrs) == 0 {
155+
c.Addrs = []*RedisAddr{
156+
{
157+
Name: DefaultRedisAddrName,
158+
Port: DefaultRedisAddrPort,
159+
},
160+
}
161+
}
162+
for i := range c.Addrs {
163+
if c.Addrs[i] == nil {
164+
continue
165+
}
166+
if strings.TrimSpace(c.Addrs[i].Name) == "" {
167+
c.Addrs[i].Name = DefaultRedisAddrName
168+
}
169+
if c.Addrs[i].Port == 0 {
170+
c.Addrs[i].Port = DefaultRedisAddrPort
171+
}
172+
}
173+
if c.DialTimeout == 0 {
174+
c.DialTimeout = DefaultRedisTimeout
175+
}
176+
if c.ReadTimeout == 0 {
177+
c.ReadTimeout = DefaultRedisTimeout
178+
}
179+
if c.WriteTimeout == 0 {
180+
c.WriteTimeout = DefaultRedisTimeout
181+
}
182+
if c.PoolTimeout == 0 {
183+
c.PoolTimeout = DefaultRedisTimeout
184+
}
185+
if c.PoolSize == 0 {
186+
c.PoolSize = DefaultRedisPoolSize
187+
}
188+
if c.MinIdleConns == 0 {
189+
c.MinIdleConns = DefaultRedisMinIdleConns
190+
}
191+
if c.MaxRetries == 0 {
192+
c.MaxRetries = DefaultRedisMaxRetries
193+
}
194+
}
195+
196+
func (c *Config) setDefaultConfigOptions() {
197+
if c == nil {
198+
return
199+
}
200+
if c.ErrorCode == 0 {
201+
c.ErrorCode = DefaultErrorCode
202+
}
203+
if strings.TrimSpace(c.ErrorMessage) == "" {
204+
c.ErrorMessage = DefaultErrorMessage
205+
}
206+
if c.Redis == nil {
207+
c.Redis = NewDefaultRedisConfig()
208+
}
209+
}
210+
121211
type SafeConfig struct {
122212
mu sync.RWMutex
123213
config *Config
@@ -153,6 +243,15 @@ func (c *SafeConfig) GetConfig() *Config {
153243
return c.config
154244
}
155245

246+
func (c *SafeConfig) IsEnabled() bool {
247+
cfg := c.GetConfig()
248+
if cfg == nil {
249+
logging.Error(errors.New("safe config is nil"), "found safe config is nil")
250+
return false
251+
}
252+
return cfg.Enabled
253+
}
254+
156255
func GetErrorCode() int32 {
157256
cfg := globalConfig.GetConfig()
158257
if cfg == nil {

‎core/llm_token_ratelimit/init.go‎

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,8 @@ package llmtokenratelimit
1616

1717
import (
1818
"fmt"
19+
20+
"github.com/alibaba/sentinel-golang/logging"
1921
)
2022

2123
func Init(cfg *Config) error {
@@ -32,6 +34,20 @@ func Init(cfg *Config) error {
3234
return fmt.Errorf("global token calculator is nil")
3335
}
3436

37+
if cfg == nil {
38+
cfg = NewDefaultConfig()
39+
if cfg == nil {
40+
return fmt.Errorf("new default config failed")
41+
}
42+
}
43+
44+
if !cfg.Enabled {
45+
logging.Warn("[LLMTokenRateLimit] llm token rate limit is disabled, please enable it in config if needed")
46+
return nil
47+
}
48+
49+
cfg.setDefaultConfigOptions()
50+
3551
if err := globalConfig.SetConfig(cfg); err != nil {
3652
return err
3753
}

‎core/llm_token_ratelimit/ratelimit_checker.go‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -80,7 +80,7 @@ func (c *FixedWindowChecker) checkLimitKey(ctx *Context, rule *MatchedRule) bool
8080
"failed to create response header in llm_token_ratelimit.FixedWindowChecker.checkLimitKey()",
8181
"requestID", ctx.Get(KeyRequestID),
8282
)
83-
return false
83+
return true
8484
}
8585
defer func() {
8686
ctx.Set(KeyResponseHeaders, responseHeader)
@@ -160,7 +160,7 @@ func (c *PETAChecker) checkLimitKey(ctx *Context, rule *MatchedRule) bool {
160160
)
161161
return true
162162
}
163-
logging.Info("[LLMTokenRateLimit] estimated infos",
163+
logging.Info("[LLMTokenRateLimit] withhold infos",
164164
"limitKey", rule.LimitKey,
165165
"current_capacity", result[0],
166166
"waiting_time(ms)", result[1],
@@ -178,7 +178,7 @@ func (c *PETAChecker) checkLimitKey(ctx *Context, rule *MatchedRule) bool {
178178
"failed to create response header in llm_token_ratelimit.PETAChecker.checkLimitKey()",
179179
"requestID", ctx.Get(KeyRequestID),
180180
)
181-
return false
181+
return true
182182
}
183183
defer func() {
184184
ctx.Set(KeyResponseHeaders, responseHeader)

‎core/llm_token_ratelimit/redis_client.go‎

Lines changed: 6 additions & 32 deletions
Original file line numberDiff line numberDiff line change
@@ -54,29 +54,25 @@ func (c *SafeRedisClient) Init(cfg *Redis) error {
5454
return fmt.Errorf("safe redis client is nil")
5555
}
5656
if cfg == nil {
57-
return fmt.Errorf("config is nil")
57+
cfg = NewDefaultRedisConfig()
58+
if cfg == nil {
59+
return fmt.Errorf("redis config is nil")
60+
}
5861
}
5962

63+
cfg.setDefaultConfigOptions()
64+
6065
addrsMap := make(map[string]struct{}, len(cfg.Addrs))
6166
for _, addr := range cfg.Addrs {
6267
if addr == nil {
6368
continue
6469
}
65-
if len(addr.Name) == 0 {
66-
addr.Name = DefaultRedisAddrName
67-
}
68-
if addr.Port == 0 {
69-
addr.Port = DefaultRedisAddrPort
70-
}
7170
addrsMap[fmt.Sprintf("%s:%d", addr.Name, addr.Port)] = struct{}{}
7271
}
7372
addrs := make([]string, 0, len(addrsMap))
7473
for addr := range addrsMap {
7574
addrs = append(addrs, addr)
7675
}
77-
if len(addrs) == 0 {
78-
addrs = append(addrs, fmt.Sprintf("%s:%d", DefaultRedisAddrName, DefaultRedisAddrPort))
79-
}
8076

8177
dialTimeout := time.Duration(cfg.DialTimeout) * time.Millisecond
8278
readTimeout := time.Duration(cfg.ReadTimeout) * time.Millisecond
@@ -86,28 +82,6 @@ func (c *SafeRedisClient) Init(cfg *Redis) error {
8682
minIdleConns := cfg.MinIdleConns
8783
maxRetries := cfg.MaxRetries
8884

89-
if dialTimeout == 0 {
90-
dialTimeout = time.Duration(DefaultRedisTimeout) * time.Millisecond
91-
}
92-
if readTimeout == 0 {
93-
readTimeout = time.Duration(DefaultRedisTimeout) * time.Millisecond
94-
}
95-
if writeTimeout == 0 {
96-
writeTimeout = time.Duration(DefaultRedisTimeout) * time.Millisecond
97-
}
98-
if poolTimeout == 0 {
99-
poolTimeout = time.Duration(DefaultRedisTimeout) * time.Millisecond
100-
}
101-
if poolSize == 0 {
102-
poolSize = DefaultRedisPoolSize
103-
}
104-
if minIdleConns == 0 {
105-
minIdleConns = DefaultRedisMinIdleConns
106-
}
107-
if maxRetries == 0 {
108-
maxRetries = DefaultRedisMaxRetries
109-
}
110-
11185
newClient := redis.NewClusterClient(
11286
&redis.ClusterOptions{
11387
Addrs: addrs,

‎core/llm_token_ratelimit/slot.go‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -53,6 +53,9 @@ func (s *Slot) Check(ctx *base.EntryContext) *base.TokenResult {
5353
}
5454

5555
func (s *Slot) checkPass(ctx *base.EntryContext) (bool, *Rule, interface{}) {
56+
if !globalConfig.IsEnabled() {
57+
return true, nil, nil
58+
}
5659
requestID := generateUUID()
5760
llmTokenRatelimitCtx, ok := ctx.GetPair(KeyContext).(*Context)
5861
if !ok || llmTokenRatelimitCtx == nil {

‎core/llm_token_ratelimit/stat_slot.go‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -42,6 +42,9 @@ func (c *LLMTokenRatelimitStatSlot) OnEntryBlocked(_ *base.EntryContext, _ *base
4242
}
4343

4444
func (c *LLMTokenRatelimitStatSlot) OnCompleted(ctx *base.EntryContext) {
45+
if !globalConfig.IsEnabled() {
46+
return
47+
}
4548
usedTokenInfos, ok := ctx.GetPair(KeyUsedTokenInfos).(*UsedTokenInfos)
4649
if !ok || usedTokenInfos == nil {
4750
return

‎core/llm_token_ratelimit/token_extractor.go‎

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -28,17 +28,17 @@ func OpenAITokenExtractor(response interface{}) (*UsedTokenInfos, error) {
2828
return nil, fmt.Errorf("response is not map[string]any")
2929
}
3030

31-
inputTokens, ok := resp["PromptTokens"].(int)
31+
inputTokens, ok := resp["prompt_tokens"].(int)
3232
if !ok {
33-
return nil, fmt.Errorf("PromptTokens not found or not int")
33+
return nil, fmt.Errorf("prompt_tokens not found or not int")
3434
}
35-
outputTokens, ok := resp["CompletionTokens"].(int)
35+
outputTokens, ok := resp["completion_tokens"].(int)
3636
if !ok {
37-
return nil, fmt.Errorf("CompletionTokens not found or not int")
37+
return nil, fmt.Errorf("completion_tokens not found or not int")
3838
}
39-
totalTokens, ok := resp["TotalTokens"].(int)
39+
totalTokens, ok := resp["total_tokens"].(int)
4040
if !ok {
41-
return nil, fmt.Errorf("TotalTokens not found or not int")
41+
return nil, fmt.Errorf("total_tokens not found or not int")
4242
}
4343

4444
return GenerateUsedTokenInfos(

‎core/llm_token_ratelimit/token_updater.go‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -112,7 +112,7 @@ func (u *PETAUpdater) updateLimitKey(ctx *Context, rule *MatchedRule, infos *Use
112112
return
113113
}
114114
actualToken := calculator.Calculate(ctx, infos)
115-
logging.Info("[LLMTokenRateLimit] actual infos",
115+
logging.Info("[LLMTokenRateLimit] correct infos",
116116
"limitKey", rule.LimitKey,
117117
"estimated_token", rule.EstimatedToken,
118118
"actual_token", actualToken,

0 commit comments

Comments
 (0)