diff --git a/.gitignore b/.gitignore index d5f65e836..97bd554a6 100644 --- a/.gitignore +++ b/.gitignore @@ -64,3 +64,9 @@ Temporary Items # coverage file coverage.html +coverage.txt + +pkg/adapters/eino/*_test.go +pkg/adapters/langchaingo/*_test.go + +.env \ No newline at end of file diff --git a/api/api.go b/api/api.go index 529fb327a..d7e729134 100644 --- a/api/api.go +++ b/api/api.go @@ -138,7 +138,9 @@ func Entry(resource string, opts ...EntryOption) (*base.SentinelEntry, *base.Blo }() for _, opt := range opts { - opt(options) + if opt != nil { + opt(options) + } } if options.slotChain == nil { options.slotChain = GlobalSlotChain() diff --git a/api/init.go b/api/init.go index 27229201f..4ee2b900a 100644 --- a/api/init.go +++ b/api/init.go @@ -20,6 +20,7 @@ import ( "net/http" "github.com/alibaba/sentinel-golang/core/config" + llmtokenratelimit "github.com/alibaba/sentinel-golang/core/llm_token_ratelimit" "github.com/alibaba/sentinel-golang/core/log/metric" "github.com/alibaba/sentinel-golang/core/system_metric" metric_exporter "github.com/alibaba/sentinel-golang/exporter/metric" @@ -134,6 +135,21 @@ func initCoreComponents() error { return nil } + if err := llmtokenratelimit.InitMetricLogger(&llmtokenratelimit.MetricLoggerConfig{ + AppName: config.AppName(), + LogDir: config.LogBaseDir(), + MaxFileSize: config.MetricLogSingleFileMaxSize(), + MaxFileAmount: config.MetricLogMaxFileAmount(), + FlushInterval: config.MetricLogFlushIntervalSec(), + UsePid: config.LogUsePid(), + }); err != nil { + return err + } + + if err := llmtokenratelimit.Init(config.LLMTokenRateLimit()); err != nil { + return err + } + return nil } diff --git a/api/slot_chain.go b/api/slot_chain.go index 35ae23340..fe7d56d33 100644 --- a/api/slot_chain.go +++ b/api/slot_chain.go @@ -20,6 +20,7 @@ import ( "github.com/alibaba/sentinel-golang/core/flow" "github.com/alibaba/sentinel-golang/core/hotspot" "github.com/alibaba/sentinel-golang/core/isolation" + llmtokenratelimit "github.com/alibaba/sentinel-golang/core/llm_token_ratelimit" "github.com/alibaba/sentinel-golang/core/log" "github.com/alibaba/sentinel-golang/core/stat" "github.com/alibaba/sentinel-golang/core/system" @@ -40,11 +41,13 @@ func BuildDefaultSlotChain() *base.SlotChain { sc.AddRuleCheckSlot(isolation.DefaultSlot) sc.AddRuleCheckSlot(hotspot.DefaultSlot) sc.AddRuleCheckSlot(circuitbreaker.DefaultSlot) + sc.AddRuleCheckSlot(llmtokenratelimit.DefaultSlot) sc.AddStatSlot(stat.DefaultSlot) sc.AddStatSlot(log.DefaultSlot) sc.AddStatSlot(flow.DefaultStandaloneStatSlot) sc.AddStatSlot(hotspot.DefaultConcurrencyStatSlot) sc.AddStatSlot(circuitbreaker.DefaultMetricStatSlot) + sc.AddStatSlot(llmtokenratelimit.DefaultLLMTokenRatelimitStatSlot) return sc } diff --git a/core/base/result.go b/core/base/result.go index 799e634a7..95584e00d 100644 --- a/core/base/result.go +++ b/core/base/result.go @@ -28,16 +28,18 @@ const ( BlockTypeCircuitBreaking BlockTypeSystemFlow BlockTypeHotSpotParamFlow + BlockTypeLLMTokenRateLimit ) var ( blockTypeMap = map[BlockType]string{ - BlockTypeUnknown: "BlockTypeUnknown", - BlockTypeFlow: "BlockTypeFlowControl", - BlockTypeIsolation: "BlockTypeIsolation", - BlockTypeCircuitBreaking: "BlockTypeCircuitBreaking", - BlockTypeSystemFlow: "BlockTypeSystem", - BlockTypeHotSpotParamFlow: "BlockTypeHotSpotParamFlow", + BlockTypeUnknown: "BlockTypeUnknown", + BlockTypeFlow: "BlockTypeFlowControl", + BlockTypeIsolation: "BlockTypeIsolation", + BlockTypeCircuitBreaking: "BlockTypeCircuitBreaking", + BlockTypeSystemFlow: "BlockTypeSystem", + BlockTypeHotSpotParamFlow: "BlockTypeHotSpotParamFlow", + BlockTypeLLMTokenRateLimit: "BlockTypeLLMTokenRateLimit", } blockTypeExisted = fmt.Errorf("block type existed") ) diff --git a/core/base/result_test.go b/core/base/result_test.go index fcce6b26f..2f94171de 100644 --- a/core/base/result_test.go +++ b/core/base/result_test.go @@ -70,6 +70,8 @@ func (t BlockType) stringSwitch() string { return "System" case BlockTypeHotSpotParamFlow: return "HotSpotParamFlow" + case BlockTypeLLMTokenRateLimit: + return "LLMTokenRateLimit" default: return fmt.Sprintf("%d", t) } @@ -77,12 +79,13 @@ func (t BlockType) stringSwitch() string { var ( blockTypeNames = []string{ - BlockTypeUnknown: "Unknown", - BlockTypeFlow: "FlowControl", - BlockTypeIsolation: "BlockTypeIsolation", - BlockTypeCircuitBreaking: "CircuitBreaking", - BlockTypeSystemFlow: "System", - BlockTypeHotSpotParamFlow: "HotSpotParamFlow", + BlockTypeUnknown: "Unknown", + BlockTypeFlow: "FlowControl", + BlockTypeIsolation: "BlockTypeIsolation", + BlockTypeCircuitBreaking: "CircuitBreaking", + BlockTypeSystemFlow: "System", + BlockTypeHotSpotParamFlow: "HotSpotParamFlow", + BlockTypeLLMTokenRateLimit: "LLMTokenRateLimit", } blockTypeErr = fmt.Errorf("block type err") ) diff --git a/core/config/config.go b/core/config/config.go index cd65f2dec..25eafb823 100644 --- a/core/config/config.go +++ b/core/config/config.go @@ -21,6 +21,7 @@ import ( "strconv" "sync" + llmtokenratelimit "github.com/alibaba/sentinel-golang/core/llm_token_ratelimit" "github.com/alibaba/sentinel-golang/logging" "github.com/alibaba/sentinel-golang/util" "github.com/pkg/errors" @@ -262,3 +263,7 @@ func MetricStatisticIntervalMs() uint32 { func MetricStatisticSampleCount() uint32 { return globalCfg.MetricStatisticSampleCount() } + +func LLMTokenRateLimit() *llmtokenratelimit.Config { + return globalCfg.LLMTokenRateLimit() +} diff --git a/core/config/entity.go b/core/config/entity.go index 2bb389a6e..6559361db 100644 --- a/core/config/entity.go +++ b/core/config/entity.go @@ -19,6 +19,7 @@ import ( "fmt" "github.com/alibaba/sentinel-golang/core/base" + llmtokenratelimit "github.com/alibaba/sentinel-golang/core/llm_token_ratelimit" "github.com/alibaba/sentinel-golang/logging" "github.com/pkg/errors" ) @@ -46,6 +47,8 @@ type SentinelConfig struct { Stat StatConfig // UseCacheTime indicates whether to cache time(ms) UseCacheTime bool `yaml:"useCacheTime"` + // LLMTokenRateLimit represents configuration items related to llm token rate limit. + LLMTokenRateLimit *llmtokenratelimit.Config `yaml:"llmTokenRatelimit"` } // ExporterConfig represents configuration items related to exporter, like metric exporter. @@ -259,3 +262,7 @@ func (entity *Entity) MetricStatisticIntervalMs() uint32 { func (entity *Entity) MetricStatisticSampleCount() uint32 { return entity.Sentinel.Stat.MetricStatisticSampleCount } + +func (entity *Entity) LLMTokenRateLimit() *llmtokenratelimit.Config { + return entity.Sentinel.LLMTokenRateLimit +} diff --git a/core/llm_token_ratelimit/README_en.md b/core/llm_token_ratelimit/README_en.md new file mode 100644 index 000000000..4ba5254e0 --- /dev/null +++ b/core/llm_token_ratelimit/README_en.md @@ -0,0 +1,199 @@ +#### Integration Steps + +From the user's perspective, to integrate the Token rate limiting function provided by Sentinel, the following steps are required: + +1. Prepare a Redis instance + +2. Configure and initialize Sentinel's runtime environment. + 1. Only initialization from a YAML file is supported + +3. Embed points (define resources) with fixed resource type: `ResourceType=ResTypeCommon` and `TrafficType=Inbound` + +4. Load rules according to the configuration file below. The rule configuration items include: resource name, rate limiting strategy, specific rule items, Redis configuration, error code, and error message. The following is an example of rule configuration, with specific field meanings detailed in the "Configuration File Description" below. + + ```go + _, err = llmtokenratelimit.LoadRules([]*llmtokenratelimit.Rule{ + { + + Resource: ".*", + Strategy: llmtokenratelimit.FixedWindow, + SpecificItems: []llmtokenratelimit.SpecificItem{ + { + Identifier: llmtokenratelimit.Identifier{ + Type: llmtokenratelimit.Header, + Value: ".*", + }, + KeyItems: []llmtokenratelimit.KeyItem{ + { + Key: ".*", + Token: llmtokenratelimit.Token{ + Number: 1000, + CountStrategy: llmtokenratelimit.TotalTokens, + }, + Time: llmtokenratelimit.Time{ + Unit: llmtokenratelimit.Second, + Value: 60, + }, + }, + }, + }, + }, + }, + }) + ``` + +5. Optional: Create an LLM instance and embed it into the provided adapter + + +#### Configuration File Description + +Overall rule configuration + +| Configuration Item | Type | Required | Default Value | Description | +| :----------------- | :------------------- | :------- | :------------------ | :----------------------------------------------------------- | +| enabled | bool | No | false | Whether to enable the LLM Token rate limiting function. Values: false (disable), true (enable) | +| rules | array of rule object | No | nil | Rate limiting rules | +| redis | object | No | | Redis instance connection information | +| errorCode | int | No | 429 | Error code. Will be changed to 429 if set to 0 | +| errorMessage | string | No | "Too Many Requests" | Error message | + +rule configuration + +| Configuration Item | Type | Required | Default Value | Description | +| :----------------- | :--------------------------- | :------- | :-------------- | :----------------------------------------------------------- | +| resource | string | No | ".*" | Rule resource name, supporting regular expressions. Values: ".*" (global match), user-defined regular expressions | +| strategy | string | No | "fixed-window" | Rate limiting strategy. Values: fixed-window, peta (predictive error temporal allocation) | +| encoding | object | No | | Token encoding method, **exclusively for peta rate limiting strategy** | +| specificItems | array of specificItem object | Yes | | Specific rule items | + +encoding configuration + +| Configuration Item | Type | Required | Default Value | Description | +| :----------------- | :----- | :------- | :------------ | :-------------------- | +| provider | string | No | "openai" | Model provider | +| model | string | No | "gpt-4" | Model name | + +specificItem configuration + +| Configuration Item | Type | Required | Default Value | Description | +| :----------------- | :---------------------- | :------- | :------------ | :------------------------------------------- | +| identifier | object | No | | Request identifier | +| keyItems | array of keyItem object | Yes | | Key-value information for rule matching | + +identifier configuration + +| Configuration Item | Type | Required | Default Value | Description | +| :----------------- | :----- | :------- | :------------ | :----------------------------------------------------------- | +| type | string | No | "all" | Request identifier type. Values: all (global rate limiting), header | +| value | string | No | ".*" | Request identifier value, supporting regular expressions. Values: ".*" (global match), user-defined regular expressions | + +keyItem configuration + +| Configuration Item | Type | Required | Default Value | Description | +| :----------------- | :----- | :------- | :------------ | :----------------------------------------------------------- | +| key | string | No | ".*" | Specific rule item value, supporting regular expressions. Values: ".*" (global match), user-defined regular expressions | +| token | object | Yes | | Token quantity and calculation strategy configuration | +| time | object | Yes | | Time unit and cycle configuration | + +token configuration + +| Configuration Item | Type | Required | Default Value | Description | +| :----------------- | :----- | :------- | :-------------- | :----------------------------------------------------------- | +| number | int | Yes | | Token quantity, greater than or equal to 0 | +| countStrategy | string | No | "total-tokens" | Token calculation strategy. Values: input-tokens, output-tokens, total-tokens | + +time configuration + +| Configuration Item | Type | Required | Default Value | Description | +| :----------------- | :----- | :------- | :------------ | :----------------------------------------------------------- | +| unit | string | Yes | | Time unit. Values: second, minute, hour, day | +| value | int | Yes | | Time value, greater than or equal to 0 | + +redis configuration + +| Configuration Item | Type | Required | Default Value | Description | +| :----------------- | :------------------- | :------- | :----------------------------------- | :----------------------------------------------------------- | +| addrs | array of addr object | No | [{name: "127.0.0.1", port: 6379}] | Redis node services, **see notes below** | +| username | string | No | Empty string | Redis username | +| password | string | No | Empty string | Redis password | +| dialTimeout | int | No | 0 | Maximum waiting time for establishing a Redis connection, unit: milliseconds | +| readTimeout | int | No | 0 | Maximum waiting time for Redis server response, unit: milliseconds | +| writeTimeout | int | No | 0 | Maximum time for sending command data to the network connection, unit: milliseconds | +| poolTimeout | int | No | 0 | Maximum waiting time for getting an idle connection from the connection pool, unit: milliseconds | +| poolSize | int | No | 10 | Number of connections in the connection pool | +| minIdleConns | int | No | 5 | Minimum number of idle connections in the connection pool | +| maxRetries | int | No | 3 | Maximum number of retries for failed operations | + +addr configuration + +| Configuration Item | Type | Required | Default Value | Description | +| :----------------- | :----- | :------- | :------------- | :----------------------------------------------------------- | +| name | string | No | "127.0.0.1" | Redis node service name, a complete [FQDN](https://en.wikipedia.org/wiki/Fully_qualified_domain_name) with service type, e.g., my-redis.dns, redis.my-ns.svc.cluster.local | +| port | int | No | 6379 | Redis node service port | + + +#### Overall Configuration File Example + +```YAML +version: "v1" +sentinel: + app: + name: sentinel-go-demo + log: + metric: + maxFileCount: 7 + llmTokenRatelimit: + enabled: true, + rules: + - resource: ".*" + strategy: "fixed-window" + specificItems: + - identifier: + type: "header" + value: ".*" + keyItems: + - key: ".*" + token: + number: 1000 + countStrategy: "total-tokens" + time: + unit: "second" + value: 60 + + errorCode: 429 + errorMessage: "Too Many Requests" + + redis: + addrs: + - name: "127.0.0.1" + port: 6379 + username: "redis" + password: "redis" + dialTimeout: 5000 + readTimeout: 5000 + writeTimeout: 5000 + poolTimeout: 5000 + poolSize: 10 + minIdleConns: 5 + maxRetries: 3 +``` + +#### LLM Framework Adaptation +Currently, it supports non-intrusive integration of Langchaingo and Eino frameworks into the Token rate limiting capability provided by Sentinel, which is mainly applicable to text generation scenarios. For usage details, refer to: +- pkg/adapters/langchaingo/wrapper.go +- pkg/adapters/eino/wrapper.go + +#### Notes + +- Since only input tokens can be predicted at present, **it is recommended to use PETA for rate limiting specifically targeting input tokens** +- PETA uses tiktoken to estimate input token consumption but requires downloading or preconfiguring the `Byte Pair Encoding (BPE)` dictionary + - Online mode + - tiktoken needs to download encoding files online for the first use + - Offline mode + - Prepare pre-cached tiktoken encoding files (**not directly downloaded files, but files processed by tiktoken**) in advance, and specify the file directory via the TIKTOKEN_CACHE_DIR environment variable +- Rule deduplication description + - In keyItems, if only the number differs, the latest number will be retained after deduplication + - In specificItems, only deduplicated keyItems will be retained + - In resource, only the latest resource will be retained +- Redis configuration description + - **If the connected Redis is in cluster mode, the number of addresses in addrs must be at least 2; otherwise, it will default to Redis standalone mode, causing rate limiting to fail** \ No newline at end of file diff --git a/core/llm_token_ratelimit/README_zh.md b/core/llm_token_ratelimit/README_zh.md new file mode 100644 index 000000000..4b63cd5a1 --- /dev/null +++ b/core/llm_token_ratelimit/README_zh.md @@ -0,0 +1,199 @@ +#### 接入步骤 + +从用户角度，接入Sentinel提供的Token限流功能，需要以下几步： + +1. 准备Redis实例 + +2. 对 Sentinel 的运行环境进行相关配置并初始化。 + + 1. 仅支持从yaml文件初始化 + +3. 埋点（定义资源），固定`ResourceType=ResTypeCommon`且`TrafficType=Inbound`的资源类型 + +4. 根据下面的配置文件加载规则，规则配置项包括：资源名称、限流策略、具体规则项、redis配置、错误码、错误信息。如下是配置规则的示例，具体字段含义在下文的“配置文件描述”中有具体说明。 + + ```go + _, err = llmtokenratelimit.LoadRules([]*llmtokenratelimit.Rule{ + { + + Resource: ".*", + Strategy: llmtokenratelimit.FixedWindow, + SpecificItems: []llmtokenratelimit.SpecificItem{ + { + Identifier: llmtokenratelimit.Identifier{ + Type: llmtokenratelimit.Header, + Value: ".*", + }, + KeyItems: []llmtokenratelimit.KeyItem{ + { + Key: ".*", + Token: llmtokenratelimit.Token{ + Number: 1000, + CountStrategy: llmtokenratelimit.TotalTokens, + }, + Time: llmtokenratelimit.Time{ + Unit: llmtokenratelimit.Second, + Value: 60, + }, + }, + }, + }, + }, + }, + }) + ``` + +5. 可选：创建LLM实例嵌入到提供的适配器中即可 + +#### 配置文件描述 + +总体规则配置 + +| 配置项 | 类型 | 必填 | 默认值 | 说明 | +| :----------- | :------------------- | :--- | :------------------ | :--------------------------------------------------------- | +| enabled | bool | 否 | false | 是否启用LLM Token限流功能，取值：false(不启用)、true(启用) | +| rules | array of rule object | 否 | nil | 限流规则 | +| redis | object | 否 | | redis实例连接信息 | +| errorCode | int | 否 | 429 | 错误码，设置为0时会修改为429 | +| errorMessage | string | 否 | "Too Many Requests" | 错误信息 | + +rule配置 + +| 配置项 | 类型 | 必填 | 默认值 | 说明 | +| :------------ | :--------------------------- | :--- | :------------- | :----------------------------------------------------------- | +| resource | string | 否 | ".*" | 规则资源名称，支持正则表达式，取值：".*"(全局匹配)、用户自定义正则表达式 | +| strategy | string | 否 | "fixed-window" | 限流策略，取值：fixed-window（固定窗口）、peta（预测误差时序分摊） | +| encoding | object | 否 | | token编码方式，**专用于peta限流策略** | +| specificItems | array of specificItem object | 是 | | 具体规则项 | + +encoding配置 + +| 配置项 | 类型 | 必填 | 默认值 | 说明 | +| :------- | :----- | :--- | :------- | :------- | +| provider | string | 否 | "openai" | 模型厂商 | +| model | string | 否 | "gpt-4" | 模型名称 | + +specificItem配置 + +| 配置项 | 类型 | 必填 | 默认值 | 说明 | +| :--------- | :---------------------- | :--- | :----- | :----------------- | +| identifier | object | 否 | | 请求标识符 | +| keyItems | array of keyItem object | 是 | | 规则匹配的键值信息 | + +identifier配置 + +| 配置项 | 类型 | 必填 | 默认值 | 说明 | +| :----- | :----- | :--- | :----- | :----------------------------------------------------------- | +| type | string | 否 | "all" | 请求标识符类型，取值：all(全局限流)、header | +| value | string | 否 | ".*" | 请求标识符取值，支持正则表达式，取值：".*"(全局匹配)、用户自定义正则表达式 | + +keyItem配置 + +| 配置项 | 类型 | 必填 | 默认值 | 说明 | +| :----- | :----- | :--- | :----- | :----------------------------------------------------------- | +| key | string | 否 | ".*" | 具体规则项取值，支持正则表达式，取值：".*"(全局匹配)、用户自定义正则表达式 | +| token | object | 是 | | token数量和计算策略配置 | +| time | object | 是 | | 时间单位和周期配置 | + +token配置 + +| 配置项 | 类型 | 必填 | 默认值 | 说明 | +| :------------ | :----- | :--- | :------------- | :----------------------------------------------------------- | +| number | int | 是 | | token数量，大于等于0 | +| countStrategy | string | 否 | "total-tokens" | token计算策略，取值：input-tokens、output-tokens、total-tokens | + +time配置 + +| 配置项 | 类型 | 必填 | 默认值 | 说明 | +| :----- | :----- | :--- | :----- | :---------------------------------------- | +| unit | string | 是 | | 时间单位，取值：second、minute、hour、day | +| value | int | 是 | | 时间值，大于等于0 | + +redis配置 + +| 配置项 | 类型 | 必填 | 默认值 | 说明 | +| :----------- | :------------------- | :--- | :-------------------------------- | :------------------------------------------------- | +| addrs | array of addr object | 否 | [{name: "127.0.0.1", port: 6379}] | redis节点服务，**见注意事项说明** | +| username | string | 否 | 空字符串 | redis用户名 | +| password | string | 否 | 空字符串 | redis密码 | +| dialTimeout | int | 否 | 0 | 建立redis连接的最长等待时间，单位：毫秒 | +| readTimeout | int | 否 | 0 | 等待Redis服务器响应的最长时间，单位：毫秒 | +| writeTimeout | int | 否 | 0 | 向网络连接发送命令数据的最长时间，单位：毫秒 | +| poolTimeout | int | 否 | 0 | 从连接池获取一个空闲连接的最大等待时间，单位：毫秒 | +| poolSize | int | 否 | 10 | 连接池中的连接数量 | +| minIdleConns | int | 否 | 5 | 连接池闲置连接的最少数量 | +| maxRetries | int | 否 | 3 | 操作失败，最大尝试次数 | + +addr配置 + +| 配置项 | 类型 | 必填 | 默认值 | 说明 | +| :----- | :----- | :--- | :---------- | :----------------------------------------------------------- | +| name | string | 否 | "127.0.0.1" | redis节点服务名称，带服务类型的完整 [FQDN](https://en.wikipedia.org/wiki/Fully_qualified_domain_name) 名称，例如 my-redis.dns、redis.my-ns.svc.cluster.local | +| port | int | 否 | 6379 | redis节点服务端口 | + +#### 总体配置文件示例 + +```YAML +version: "v1" +sentinel: + app: + name: sentinel-go-demo + log: + metric: + maxFileCount: 7 + llmTokenRatelimit: + enabled: true, + rules: + - resource: ".*" + strategy: "fixed-window" + specificItems: + - identifier: + type: "header" + value: ".*" + keyItems: + - key: ".*" + token: + number: 1000 + countStrategy: "total-tokens" + time: + unit: "second" + value: 60 + + errorCode: 429 + errorMessage: "Too Many Requests" + + redis: + addrs: + - name: "127.0.0.1" + port: 6379 + username: "redis" + password: "redis" + dialTimeout: 5000 + readTimeout: 5000 + writeTimeout: 5000 + poolTimeout: 5000 + poolSize: 10 + minIdleConns: 5 + maxRetries: 3 +``` +#### LLM框架适配 + +目前支持Langchaingo和Eino框架无侵入式接入Sentinel提供的Token限流能力，主要适用于文本生成方面，使用方法详见： + +- pkg/adapters/langchaingo/wrapper.go +- pkg/adapters/eino/wrapper.go + +#### 注意事项 + +- 由于目前仅可预知input tokens，所以**建议使用PETA专对于input tokens进行限流** +- PETA使用tiktoken预估输入消耗token数，但是需要下载或预先配置`字节对编码(Byte Pair Encoding，BPE)`字典 + - 在线模式 + - 首次使用时，tiktoken需要联网下载编码文件 + - 离线模式 + - 预先准备缓存好的tiktoken的编码文件（**非直接下载文件，而是经过tiktoken处理后的文件**），并通过配置TIKTOKEN_CACHE_DIR环境变量指定文件目录位置 +- 规则去重说明 + - keyItems中，若仅number不同，会去重保留最新的number + - specificItems中，仅保留去重后的keyItems + - resource中，仅保留最新的resource +- redis配置说明 + - **若连接的redis是集群模式，那么addrs里面的地址数量必须大于等于2个，否则会默认进入redis单点模式，导致限流失效** diff --git a/core/llm_token_ratelimit/config.go b/core/llm_token_ratelimit/config.go new file mode 100644 index 000000000..69ccf2a5e --- /dev/null +++ b/core/llm_token_ratelimit/config.go @@ -0,0 +1,524 @@ +// Copyright 1999-2020 Alibaba Group Holding Ltd. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package llm_token_ratelimit + +import ( + "errors" + "fmt" + "strings" + "sync" + + "github.com/alibaba/sentinel-golang/logging" +) + +type TokenEncoderProvider uint32 + +const ( + OpenAIEncoderProvider TokenEncoderProvider = iota +) + +type IdentifierType uint32 + +const ( + AllIdentifier IdentifierType = iota + Header +) + +type CountStrategy uint32 + +const ( + TotalTokens CountStrategy = iota + InputTokens + OutputTokens +) + +type TimeUnit uint32 + +const ( + Second TimeUnit = iota + Minute + Hour + Day +) + +type Strategy uint32 + +const ( + FixedWindow Strategy = iota + PETA +) + +type TokenEncoding struct { + Provider TokenEncoderProvider `json:"provider" yaml:"provider"` + Model string `json:"model" yaml:"model"` +} + +type Identifier struct { + Type IdentifierType `json:"type" yaml:"type"` + Value string `json:"value" yaml:"value"` +} + +type Token struct { + Number int64 `json:"number" yaml:"number"` + CountStrategy CountStrategy `json:"countStrategy" yaml:"countStrategy"` +} + +type Time struct { + Unit TimeUnit `json:"unit" yaml:"unit"` + Value int64 `json:"value" yaml:"value"` +} + +type KeyItem struct { + Key string `json:"key" yaml:"key"` + Token Token `json:"token" yaml:"token"` + Time Time `json:"time" yaml:"time"` +} + +type SpecificItem struct { + Identifier Identifier `json:"identifier" yaml:"identifier"` + KeyItems []*KeyItem `json:"keyItems" yaml:"keyItems"` +} + +type RedisAddr struct { + Name string `json:"name" yaml:"name"` + Port int32 `json:"port" yaml:"port"` +} + +type Redis struct { + Addrs []*RedisAddr `json:"addrs" yaml:"addrs"` + Username string `json:"username" yaml:"username"` + Password string `json:"password" yaml:"password"` + + DialTimeout int32 `json:"dialTimeout" yaml:"dialTimeout"` + ReadTimeout int32 `json:"readTimeout" yaml:"readTimeout"` + WriteTimeout int32 `json:"writeTimeout" yaml:"writeTimeout"` + PoolTimeout int32 `json:"poolTimeout" yaml:"poolTimeout"` + + PoolSize int32 `json:"poolSize" yaml:"poolSize"` + MinIdleConns int32 `json:"minIdleConns" yaml:"minIdleConns"` + MaxRetries int32 `json:"maxRetries" yaml:"maxRetries"` +} + +type Config struct { + Enabled bool `json:"enabled" yaml:"enabled"` + Rules []*Rule `json:"rules" yaml:"rules"` + Redis *Redis `json:"redis" yaml:"redis"` + ErrorCode int32 `json:"errorCode" yaml:"errorCode"` + ErrorMessage string `json:"errorMessage" yaml:"errorMessage"` +} + +func NewDefaultRedisConfig() *Redis { + return &Redis{ + Addrs: []*RedisAddr{ + { + Name: DefaultRedisAddrName, + Port: DefaultRedisAddrPort, + }, + }, + DialTimeout: DefaultRedisTimeout, + ReadTimeout: DefaultRedisTimeout, + WriteTimeout: DefaultRedisTimeout, + PoolTimeout: DefaultRedisTimeout, + PoolSize: DefaultRedisPoolSize, + MinIdleConns: DefaultRedisMinIdleConns, + MaxRetries: DefaultRedisMaxRetries, + } +} + +func NewDefaultConfig() *Config { + return &Config{ + Enabled: false, + Rules: nil, + Redis: nil, + ErrorCode: DefaultErrorCode, + ErrorMessage: DefaultErrorMessage, + } +} + +func (c *Redis) setDefaultConfigOptions() { + if c == nil { + return + } + if len(c.Addrs) == 0 { + c.Addrs = []*RedisAddr{ + { + Name: DefaultRedisAddrName, + Port: DefaultRedisAddrPort, + }, + } + } + for i := range c.Addrs { + if c.Addrs[i] == nil { + continue + } + if strings.TrimSpace(c.Addrs[i].Name) == "" { + c.Addrs[i].Name = DefaultRedisAddrName + } + if c.Addrs[i].Port == 0 { + c.Addrs[i].Port = DefaultRedisAddrPort + } + } + if c.DialTimeout == 0 { + c.DialTimeout = DefaultRedisTimeout + } + if c.ReadTimeout == 0 { + c.ReadTimeout = DefaultRedisTimeout + } + if c.WriteTimeout == 0 { + c.WriteTimeout = DefaultRedisTimeout + } + if c.PoolTimeout == 0 { + c.PoolTimeout = DefaultRedisTimeout + } + if c.PoolSize == 0 { + c.PoolSize = DefaultRedisPoolSize + } + if c.MinIdleConns == 0 { + c.MinIdleConns = DefaultRedisMinIdleConns + } + if c.MaxRetries == 0 { + c.MaxRetries = DefaultRedisMaxRetries + } +} + +func (c *Config) setDefaultConfigOptions() { + if c == nil { + return + } + if c.ErrorCode == 0 { + c.ErrorCode = DefaultErrorCode + } + if strings.TrimSpace(c.ErrorMessage) == "" { + c.ErrorMessage = DefaultErrorMessage + } + if c.Redis == nil { + c.Redis = NewDefaultRedisConfig() + } +} + +type SafeConfig struct { + mu sync.RWMutex + config *Config +} + +var globalConfig = NewGlobalConfig() + +func NewGlobalConfig() *SafeConfig { + return &SafeConfig{} +} + +func (c *SafeConfig) SetConfig(newConfig *Config) error { + if c == nil { + return fmt.Errorf("safe config is nil") + } + if newConfig == nil { + return fmt.Errorf("new config cannot be nil") + } + + c.mu.Lock() + defer c.mu.Unlock() + c.config = newConfig + return nil +} + +func (c *SafeConfig) GetConfig() *Config { + if c == nil { + logging.Error(errors.New("safe config is nil"), "found safe config is nil") + return nil + } + c.mu.RLock() + defer c.mu.RUnlock() + return c.config +} + +func (c *SafeConfig) IsEnabled() bool { + if c == nil || c.config == nil { + logging.Error(errors.New("safe config is nil"), "found safe config is nil") + return false + } + c.mu.RLock() + defer c.mu.RUnlock() + return c.config.Enabled +} + +func (c *SafeConfig) GetErrorCode() int32 { + if c == nil || c.config == nil { + logging.Error(errors.New("safe config is nil"), "found safe config is nil") + return DefaultErrorCode + } + c.mu.RLock() + defer c.mu.RUnlock() + return c.config.ErrorCode +} + +func (c *SafeConfig) GetErrorMsg() string { + if c == nil || c.config == nil { + logging.Error(errors.New("safe config is nil"), "found safe config is nil") + return DefaultErrorMessage + } + c.mu.RLock() + defer c.mu.RUnlock() + return c.config.ErrorMessage +} + +func (p *TokenEncoderProvider) UnmarshalYAML(unmarshal func(interface{}) error) error { + if p == nil { + return fmt.Errorf("token encoder provider is nil") + } + var str string + if err := unmarshal(&str); err != nil { + return err + } + switch str { + case "openai": + *p = OpenAIEncoderProvider + default: + return fmt.Errorf("unknown token encoder provider: %s", str) + } + return nil +} + +func (it *IdentifierType) UnmarshalYAML(unmarshal func(interface{}) error) error { + if it == nil { + return fmt.Errorf("identifier type is nil") + } + var str string + if err := unmarshal(&str); err != nil { + return err + } + switch str { + case "all": + *it = AllIdentifier + case "header": + *it = Header + default: + return fmt.Errorf("unknown identifier type: %s", str) + } + return nil +} + +func (ct *CountStrategy) UnmarshalYAML(unmarshal func(interface{}) error) error { + if ct == nil { + return fmt.Errorf("count strategy is nil") + } + var str string + if err := unmarshal(&str); err != nil { + return err + } + switch str { + case "total-tokens": + *ct = TotalTokens + case "input-tokens": + *ct = InputTokens + case "output-tokens": + *ct = OutputTokens + default: + return fmt.Errorf("unknown count strategy: %s", str) + } + return nil +} + +func (tu *TimeUnit) UnmarshalYAML(unmarshal func(interface{}) error) error { + if tu == nil { + return fmt.Errorf("time unit is nil") + } + var str string + if err := unmarshal(&str); err != nil { + return err + } + switch str { + case "second": + *tu = Second + case "minute": + *tu = Minute + case "hour": + *tu = Hour + case "day": + *tu = Day + default: + return fmt.Errorf("unknown time unit: %s", str) + } + return nil +} + +func (s *Strategy) UnmarshalYAML(unmarshal func(interface{}) error) error { + if s == nil { + return fmt.Errorf("strategy is nil") + } + var str string + if err := unmarshal(&str); err != nil { + return err + } + switch str { + case "fixed-window": + *s = FixedWindow + case "peta": + *s = PETA + default: + return fmt.Errorf("unknown strategy: %s", str) + } + return nil +} + +func (e *TokenEncoding) String() string { + if e == nil { + return "TokenEncoding{nil}" + } + return fmt.Sprintf("TokenEncoding{Provider:%s, Model:%s}", e.Provider.String(), e.Model) +} + +func (it IdentifierType) String() string { + switch it { + case AllIdentifier: + return "all" + case Header: + return "header" + default: + return "undefined" + } +} + +func (ct CountStrategy) String() string { + switch ct { + case TotalTokens: + return "total-tokens" + case InputTokens: + return "input-tokens" + case OutputTokens: + return "output-tokens" + default: + return "undefined" + } +} + +func (tu TimeUnit) String() string { + switch tu { + case Second: + return "second" + case Minute: + return "minute" + case Hour: + return "hour" + case Day: + return "day" + default: + return "undefined" + } +} + +func (s Strategy) String() string { + switch s { + case FixedWindow: + return "fixed-window" + case PETA: + return "peta" + default: + return "undefined" + } +} + +func (p TokenEncoderProvider) String() string { + switch p { + case OpenAIEncoderProvider: + return "openai" + default: + return "undefined" + } +} + +func (si *SpecificItem) String() string { + if si == nil { + return "SpecificItem{nil}" + } + + var sb strings.Builder + sb.WriteString("SpecificItem{") + sb.WriteString(fmt.Sprintf("Identifier:%s", si.Identifier.String())) + + if len(si.KeyItems) > 0 { + sb.WriteString(", KeyItems:[") + for i, item := range si.KeyItems { + if i > 0 { + sb.WriteString(", ") + } + sb.WriteString(item.String()) + } + sb.WriteString("]") + } else { + sb.WriteString(", KeyItems:[]") + } + + sb.WriteString("}") + return sb.String() +} + +func (id *Identifier) String() string { + if id == nil { + return "Identifier{nil}" + } + return fmt.Sprintf("Identifier{Type:%s, Value:%s}", id.Type.String(), id.Value) +} + +func (ki *KeyItem) String() string { + if ki == nil { + return "KeyItem{nil}" + } + return fmt.Sprintf("KeyItem{Key:%s, Token:%s, Time:%s}", + ki.Key, ki.Token.String(), ki.Time.String()) +} + +func (t *Token) String() string { + if t == nil { + return "Token{nil}" + } + return fmt.Sprintf("Token{Number:%d, CountStrategy:%s}", + t.Number, t.CountStrategy.String()) +} + +func (t *Time) String() string { + if t == nil { + return "Time{nil}" + } + return fmt.Sprintf("Time{Value:%d second}", t.convertToSeconds()) +} + +func (t *Time) convertToSeconds() int64 { + switch t.Unit { + case Second: + return t.Value + case Minute: + return t.Value * 60 + case Hour: + return t.Value * 3600 + case Day: + return t.Value * 86400 + default: + return ErrorTimeDuration + } +} + +func (r *Redis) String() string { + if r == nil { + return "Redis{nil}" + } + var addrs []string + for _, addr := range r.Addrs { + if addr != nil { + addrs = append(addrs, fmt.Sprintf("%s:%d", addr.Name, addr.Port)) + } + } + return fmt.Sprintf("Redis{Addrs:[%s], Username:%s, Password:%s, DialTimeout:%d, ReadTimeout:%d, WriteTimeout:%d, PoolTimeout:%d, PoolSize:%d, MinIdleConns:%d, MaxRetries:%d}", + strings.Join(addrs, ", "), r.Username, "******", r.DialTimeout, r.ReadTimeout, r.WriteTimeout, r.PoolTimeout, r.PoolSize, r.MinIdleConns, r.MaxRetries) +} diff --git a/core/llm_token_ratelimit/config_test.go b/core/llm_token_ratelimit/config_test.go new file mode 100644 index 000000000..7a2f00fee --- /dev/null +++ b/core/llm_token_ratelimit/config_test.go @@ -0,0 +1,618 @@ +// Copyright 1999-2020 Alibaba Group Holding Ltd. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package llm_token_ratelimit + +import ( + "testing" + + "gopkg.in/yaml.v3" +) + +func TestIdentifierType_UnmarshalYAML(t *testing.T) { + tests := []struct { + name string + yamlData string + expected IdentifierType + wantErr bool + }{ + {"all identifier", `type: all`, AllIdentifier, false}, + {"header identifier", `type: header`, Header, false}, + {"unknown identifier", `type: unknown`, AllIdentifier, true}, + {"empty identifier", `type: ""`, AllIdentifier, true}, + {"number as identifier", `type: 123`, AllIdentifier, true}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + var data struct { + Type IdentifierType `yaml:"type"` + } + + err := yaml.Unmarshal([]byte(tt.yamlData), &data) + if tt.wantErr { + if err == nil { + t.Errorf("Expected error but got none") + } + } else { + if err != nil { + t.Errorf("Unexpected error: %v", err) + } + if data.Type != tt.expected { + t.Errorf("Expected %v, got %v", tt.expected, data.Type) + } + } + }) + } +} + +func TestCountStrategy_UnmarshalYAML(t *testing.T) { + tests := []struct { + name string + yamlData string + expected CountStrategy + wantErr bool + }{ + {"total tokens", `strategy: total-tokens`, TotalTokens, false}, + {"input tokens", `strategy: input-tokens`, InputTokens, false}, + {"output tokens", `strategy: output-tokens`, OutputTokens, false}, + {"unknown strategy", `strategy: unknown-tokens`, TotalTokens, true}, + {"empty strategy", `strategy: ""`, TotalTokens, true}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + var data struct { + Strategy CountStrategy `yaml:"strategy"` + } + + err := yaml.Unmarshal([]byte(tt.yamlData), &data) + if tt.wantErr { + if err == nil { + t.Errorf("Expected error but got none") + } + } else { + if err != nil { + t.Errorf("Unexpected error: %v", err) + } + if data.Strategy != tt.expected { + t.Errorf("Expected %v, got %v", tt.expected, data.Strategy) + } + } + }) + } +} + +func TestTimeUnit_UnmarshalYAML(t *testing.T) { + tests := []struct { + name string + yamlData string + expected TimeUnit + wantErr bool + }{ + {"second unit", `unit: second`, Second, false}, + {"minute unit", `unit: minute`, Minute, false}, + {"hour unit", `unit: hour`, Hour, false}, + {"day unit", `unit: day`, Day, false}, + {"unknown unit", `unit: week`, Second, true}, + {"empty unit", `unit: ""`, Second, true}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + var data struct { + Unit TimeUnit `yaml:"unit"` + } + + err := yaml.Unmarshal([]byte(tt.yamlData), &data) + if tt.wantErr { + if err == nil { + t.Errorf("Expected error but got none") + } + } else { + if err != nil { + t.Errorf("Unexpected error: %v", err) + } + if data.Unit != tt.expected { + t.Errorf("Expected %v, got %v", tt.expected, data.Unit) + } + } + }) + } +} + +func TestStrategy_UnmarshalYAML(t *testing.T) { + tests := []struct { + name string + yamlData string + expected Strategy + wantErr bool + }{ + {"fixed window", `strategy: fixed-window`, FixedWindow, false}, + {"unknown strategy", `strategy: sliding-window`, FixedWindow, true}, + {"empty strategy", `strategy: ""`, FixedWindow, true}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + var data struct { + Strategy Strategy `yaml:"strategy"` + } + + err := yaml.Unmarshal([]byte(tt.yamlData), &data) + if tt.wantErr { + if err == nil { + t.Errorf("Expected error but got none") + } + } else { + if err != nil { + t.Errorf("Unexpected error: %v", err) + } + if data.Strategy != tt.expected { + t.Errorf("Expected %v, got %v", tt.expected, data.Strategy) + } + } + }) + } +} + +func TestIdentifierType_String(t *testing.T) { + tests := []struct { + name string + it IdentifierType + expected string + }{ + {"all identifier", AllIdentifier, "all"}, + {"header identifier", Header, "header"}, + {"undefined identifier", IdentifierType(999), "undefined"}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + result := tt.it.String() + if result != tt.expected { + t.Errorf("Expected %q, got %q", tt.expected, result) + } + }) + } +} + +func TestCountStrategy_String(t *testing.T) { + tests := []struct { + name string + cs CountStrategy + expected string + }{ + {"total tokens", TotalTokens, "total-tokens"}, + {"input tokens", InputTokens, "input-tokens"}, + {"output tokens", OutputTokens, "output-tokens"}, + {"undefined strategy", CountStrategy(999), "undefined"}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + result := tt.cs.String() + if result != tt.expected { + t.Errorf("Expected %q, got %q", tt.expected, result) + } + }) + } +} + +func TestTimeUnit_String(t *testing.T) { + tests := []struct { + name string + tu TimeUnit + expected string + }{ + {"second", Second, "second"}, + {"minute", Minute, "minute"}, + {"hour", Hour, "hour"}, + {"day", Day, "day"}, + {"undefined unit", TimeUnit(999), "undefined"}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + result := tt.tu.String() + if result != tt.expected { + t.Errorf("Expected %q, got %q", tt.expected, result) + } + }) + } +} + +func TestStrategy_String(t *testing.T) { + tests := []struct { + name string + s Strategy + expected string + }{ + {"fixed window", FixedWindow, "fixed-window"}, + {"undefined strategy", Strategy(999), "undefined"}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + result := tt.s.String() + if result != tt.expected { + t.Errorf("Expected %q, got %q", tt.expected, result) + } + }) + } +} + +func TestIdentifier_String(t *testing.T) { + tests := []struct { + name string + id *Identifier + expected string + }{ + {"nil identifier", nil, "Identifier{nil}"}, + {"all identifier", &Identifier{Type: AllIdentifier, Value: ".*"}, "Identifier{Type:all, Value:.*}"}, + {"header identifier", &Identifier{Type: Header, Value: "user-id"}, "Identifier{Type:header, Value:user-id}"}, + {"empty value", &Identifier{Type: Header, Value: ""}, "Identifier{Type:header, Value:}"}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + result := tt.id.String() + if result != tt.expected { + t.Errorf("Expected %q, got %q", tt.expected, result) + } + }) + } +} + +func TestToken_String(t *testing.T) { + tests := []struct { + name string + token *Token + expected string + }{ + {"nil token", nil, "Token{nil}"}, + {"total tokens", &Token{Number: 1000, CountStrategy: TotalTokens}, "Token{Number:1000, CountStrategy:total-tokens}"}, + {"input tokens", &Token{Number: 500, CountStrategy: InputTokens}, "Token{Number:500, CountStrategy:input-tokens}"}, + {"output tokens", &Token{Number: 300, CountStrategy: OutputTokens}, "Token{Number:300, CountStrategy:output-tokens}"}, + {"zero tokens", &Token{Number: 0, CountStrategy: TotalTokens}, "Token{Number:0, CountStrategy:total-tokens}"}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + result := tt.token.String() + if result != tt.expected { + t.Errorf("Expected %q, got %q", tt.expected, result) + } + }) + } +} + +func TestTime_String(t *testing.T) { + tests := []struct { + name string + time *Time + expected string + }{ + {"nil time", nil, "Time{nil}"}, + {"second", &Time{Value: 30, Unit: Second}, "Time{Value:30 second}"}, + {"minute", &Time{Value: 5, Unit: Minute}, "Time{Value:300 second}"}, + {"hour", &Time{Value: 2, Unit: Hour}, "Time{Value:7200 second}"}, + {"day", &Time{Value: 1, Unit: Day}, "Time{Value:86400 second}"}, + {"zero value", &Time{Value: 0, Unit: Hour}, "Time{Value:0 second}"}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + result := tt.time.String() + if result != tt.expected { + t.Errorf("Expected %q, got %q", tt.expected, result) + } + }) + } +} + +func TestKeyItem_String(t *testing.T) { + tests := []struct { + name string + ki *KeyItem + expected string + }{ + {"nil keyitem", nil, "KeyItem{nil}"}, + { + "complete keyitem", + &KeyItem{ + Key: "rate-limit", + Token: Token{Number: 1000, CountStrategy: TotalTokens}, + Time: Time{Value: 1, Unit: Hour}, + }, + "KeyItem{Key:rate-limit, Token:Token{Number:1000, CountStrategy:total-tokens}, Time:Time{Value:3600 second}}", + }, + { + "empty key", + &KeyItem{ + Key: "", + Token: Token{Number: 500, CountStrategy: InputTokens}, + Time: Time{Value: 30, Unit: Second}, + }, + "KeyItem{Key:, Token:Token{Number:500, CountStrategy:input-tokens}, Time:Time{Value:30 second}}", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + result := tt.ki.String() + if result != tt.expected { + t.Errorf("Expected %q, got %q", tt.expected, result) + } + }) + } +} + +func TestSpecificItem_String(t *testing.T) { + tests := []struct { + name string + ri *SpecificItem + expected string + }{ + {"nil specificItem", nil, "SpecificItem{nil}"}, + { + "empty keyitems", + &SpecificItem{ + Identifier: Identifier{Type: AllIdentifier, Value: ".*"}, + KeyItems: []*KeyItem{}, + }, + "SpecificItem{Identifier:Identifier{Type:all, Value:.*}, KeyItems:[]}", + }, + { + "single keyitem", + &SpecificItem{ + Identifier: Identifier{Type: Header, Value: "user-id"}, + KeyItems: []*KeyItem{ + { + Key: "limit1", + Token: Token{Number: 1000, CountStrategy: TotalTokens}, + Time: Time{Value: 1, Unit: Hour}, + }, + }, + }, + "SpecificItem{Identifier:Identifier{Type:header, Value:user-id}, KeyItems:[KeyItem{Key:limit1, Token:Token{Number:1000, CountStrategy:total-tokens}, Time:Time{Value:3600 second}}]}", + }, + { + "multiple keyitems", + &SpecificItem{ + Identifier: Identifier{Type: Header, Value: "api-key"}, + KeyItems: []*KeyItem{ + { + Key: "limit1", + Token: Token{Number: 500, CountStrategy: InputTokens}, + Time: Time{Value: 30, Unit: Minute}, + }, + { + Key: "limit2", + Token: Token{Number: 300, CountStrategy: OutputTokens}, + Time: Time{Value: 1, Unit: Hour}, + }, + }, + }, + "SpecificItem{Identifier:Identifier{Type:header, Value:api-key}, KeyItems:[KeyItem{Key:limit1, Token:Token{Number:500, CountStrategy:input-tokens}, Time:Time{Value:1800 second}}, KeyItem{Key:limit2, Token:Token{Number:300, CountStrategy:output-tokens}, Time:Time{Value:3600 second}}]}", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + result := tt.ri.String() + if result != tt.expected { + t.Errorf("Expected %q, got %q", tt.expected, result) + } + }) + } +} + +func TestCompleteYAMLUnmarshaling(t *testing.T) { + yamlData := ` +rules: + - id: "rule1" + resource: "/api/chat" + strategy: fixed-window + specificItems: + - identifier: + type: header + value: "user-id" + keyItems: + - key: "hourly-limit" + token: + number: 1000 + countStrategy: total-tokens + time: + value: 1 + unit: hour + - key: "daily-limit" + token: + number: 10000 + countStrategy: input-tokens + time: + value: 1 + unit: day +redis: + addrs: + - name: "localhost" + port: 6380 + username: "redis-user" + password: "redis-pass" + timeout: 5 + poolSize: 10 + minIdleConns: 2 + maxRetries: 3 +errorCode: 429 +errorMessage: "Rate limit exceeded" +` + + var config Config + err := yaml.Unmarshal([]byte(yamlData), &config) + if err != nil { + t.Fatalf("Failed to unmarshal YAML: %v", err) + } + + // Validate basic structure + if len(config.Rules) != 1 { + t.Errorf("Expected 1 rule, got %d", len(config.Rules)) + } + + rule := config.Rules[0] + if rule.ID != "rule1" { + t.Errorf("Expected rule ID 'rule1', got %q", rule.ID) + } + + if rule.Strategy != FixedWindow { + t.Errorf("Expected FixedWindow strategy, got %v", rule.Strategy) + } + + if len(rule.SpecificItems) != 1 { + t.Errorf("Expected 1 rule item, got %d", len(rule.SpecificItems)) + } + + specificItem := rule.SpecificItems[0] + if specificItem.Identifier.Type != Header { + t.Errorf("Expected Header identifier type, got %v", specificItem.Identifier.Type) + } + + if len(specificItem.KeyItems) != 2 { + t.Errorf("Expected 2 key items, got %d", len(specificItem.KeyItems)) + } + + // Test Redis config + if config.Redis.Addrs[0].Name != "localhost" { + t.Errorf("Expected Redis addr name 'localhost', got %q", config.Redis.Addrs[0].Name) + } + + if config.Redis.Addrs[0].Port != 6380 { + t.Errorf("Expected Redis addr port 6380, got %d", config.Redis.Addrs[0].Port) + } + + // Test error config + if config.ErrorCode != 429 { + t.Errorf("Expected error code 429, got %d", config.ErrorCode) + } +} + +func TestYAMLUnmarshalingErrors(t *testing.T) { + tests := []struct { + name string + yamlData string + wantErr bool + }{ + { + "invalid identifier type", + ` +rules: + - specificItems: + - identifier: + type: invalid +`, + true, + }, + { + "invalid count strategy", + ` +rules: + - specificItems: + - keyItems: + - token: + countStrategy: invalid-strategy +`, + true, + }, + { + "invalid time unit", + ` +rules: + - specificItems: + - keyItems: + - time: + unit: invalid-unit +`, + true, + }, + { + "invalid strategy", + ` +rules: + - strategy: invalid-strategy +`, + true, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + var config Config + err := yaml.Unmarshal([]byte(tt.yamlData), &config) + if tt.wantErr && err == nil { + t.Error("Expected error but got none") + } else if !tt.wantErr && err != nil { + t.Errorf("Unexpected error: %v", err) + } + }) + } +} + +// Benchmark tests +func BenchmarkIdentifierType_String(b *testing.B) { + it := Header + b.ReportAllocs() + for i := 0; i < b.N; i++ { + _ = it.String() + } +} + +func BenchmarkCountStrategy_String(b *testing.B) { + cs := TotalTokens + b.ReportAllocs() + for i := 0; i < b.N; i++ { + _ = cs.String() + } +} + +func BenchmarkTimeUnit_String(b *testing.B) { + tu := Hour + b.ReportAllocs() + for i := 0; i < b.N; i++ { + _ = tu.String() + } +} + +func BenchmarkToken_String(b *testing.B) { + token := &Token{Number: 1000, CountStrategy: TotalTokens} + b.ReportAllocs() + for i := 0; i < b.N; i++ { + _ = token.String() + } +} + +func BenchmarkSpecificItem_String(b *testing.B) { + ri := &SpecificItem{ + Identifier: Identifier{Type: Header, Value: "user-id"}, + KeyItems: []*KeyItem{ + { + Key: "limit1", + Token: Token{Number: 1000, CountStrategy: TotalTokens}, + Time: Time{Value: 1, Unit: Hour}, + }, + }, + } + b.ReportAllocs() + for i := 0; i < b.N; i++ { + _ = ri.String() + } +} diff --git a/core/llm_token_ratelimit/constant.go b/core/llm_token_ratelimit/constant.go new file mode 100644 index 000000000..dd91bc297 --- /dev/null +++ b/core/llm_token_ratelimit/constant.go @@ -0,0 +1,100 @@ +// Copyright 1999-2020 Alibaba Group Holding Ltd. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package llm_token_ratelimit + +// ================================= Config ==================================== +const ( + DefaultResourcePattern string = ".*" + DefaultIdentifierValuePattern string = ".*" + DefaultKeyPattern string = ".*" + + DefaultRedisAddrName string = "127.0.0.1" + DefaultRedisAddrPort int32 = 6379 + DefaultRedisTimeout int32 = 0 // milliseconds + DefaultRedisPoolSize int32 = 10 + DefaultRedisMinIdleConns int32 = 5 + DefaultRedisMaxRetries int32 = 3 + + DefaultErrorCode int32 = 429 + DefaultErrorMessage string = "Too Many Requests" +) + +var DefaultTokenEncodingModel = map[TokenEncoderProvider]string{ + OpenAIEncoderProvider: "gpt-4", +} + +// ================================= CommonError ============================== +const ( + ErrorTimeDuration int64 = -1 +) + +// ================================= Context ================================== +const ( + KeyContext string = "SentinelLLMTokenRatelimitContext" + KeyRequestInfos string = "SentinelLLMTokenRatelimitReqInfos" + KeyUsedTokenInfos string = "SentinelLLMTokenRatelimitUsedTokenInfos" + KeyMatchedRules string = "SentinelLLMTokenRatelimitMatchedRules" + KeyResponseHeaders string = "SentinelLLMTokenRatelimitResponseHeaders" + KeyRequestID string = "SentinelLLMTokenRatelimitRequestID" + KeyErrorCode string = "SentinelLLMTokenRatelimitErrorCode" + KeyErrorMessage string = "SentinelLLMTokenRatelimitErrorMessage" +) + +// ================================= RedisRatelimitKeyFormat ================== +const ( + RedisRatelimitKeyFormat string = "sentinel-go:llm-token-ratelimit:resource-%s:%s:%s:%d:%s" // hashedResource, strategy, identifierType, timeWindow, tokenCountStrategy +) + +// ================================= ResponseHeader ================== +const ( + ResponseHeaderRequestID string = "X-Sentinel-LLM-Token-Ratelimit-RequestID" + ResponseHeaderRemainingTokens string = "X-Sentinel-LLM-Token-Ratelimit-RemainingTokens" + ResponseHeaderWaitingTime string = "X-Sentinel-LLM-Token-Ratelimit-WaitingTime" +) + +// ================================= FixedWindowStrategy ====================== + +// ================================= PETAStrategy ============================= +const ( + PETANoWaiting int64 = 0 + PETACorrectOK int64 = 0 + PETACorrectUnderestimateError int64 = 1 + PETACorrectOverestimateError int64 = 2 + PETASlidingWindowKeyFormat string = "{shard-%s}:sliding-window:%s" // hashTag, redisRatelimitKey + PETATokenBucketKeyFormat string = "{shard-%s}:token-bucket:%s" // hashTag, redisRatelimitKey + PETARandomStringLength int = 16 +) + +// ================================= Generate Random String =================== +const ( + RandomLetterBytes = "0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz" + RandomLetterIdxBits = 6 + RandomLetterIdxMask = 1< 0 { + _ = cfg.Rules[0].ID // Access should not cause data race + } + if cfg.Redis != nil { + _ = cfg.Redis.Addrs // Access should not cause data race + } + } + } + }(i) + } + + // Wait for completion with timeout + done := make(chan struct{}) + go func() { + wg.Wait() + close(done) + }() + + select { + case <-done: + // Success + case err := <-errors: + t.Fatal(err) + case <-time.After(30 * time.Second): + t.Fatal("Test timed out") + } + + // Check for any remaining errors + close(errors) + for err := range errors { + t.Errorf("Concurrent error: %v", err) + } +} + +func TestSafeConfig_NilHandling(t *testing.T) { + var nilConfig *SafeConfig + + // Test nil SafeConfig + result := nilConfig.GetConfig() + assert.Nil(t, result) + + err := nilConfig.SetConfig(&Config{}) + assert.Error(t, err) + assert.Contains(t, err.Error(), "safe config is nil") + + // Test setting nil config + config := &SafeConfig{} + err = config.SetConfig(nil) + assert.Error(t, err) + assert.Contains(t, err.Error(), "config cannot be nil") +} + +func TestSafeConfig_GetAfterSet(t *testing.T) { + config := &SafeConfig{} + + // Initially nil + assert.Nil(t, config.GetConfig()) + + // Set and get + testConfig := &Config{ + ErrorCode: 500, + ErrorMessage: "test error", + } + + err := config.SetConfig(testConfig) + require.NoError(t, err) + + retrieved := config.GetConfig() + require.NotNil(t, retrieved) + assert.Equal(t, int32(500), retrieved.ErrorCode) + assert.Equal(t, "test error", retrieved.ErrorMessage) +} + +// Test concurrent safety of global operations +func TestInit_ConcurrentSafety(t *testing.T) { + // Save original state + originalConfig := globalConfig.GetConfig() + defer func() { + if originalConfig != nil { + globalConfig.SetConfig(originalConfig) + } + }() + + const numGoroutines = 50 + var wg sync.WaitGroup + errors := make(chan error, numGoroutines) + + // Create different configs for each goroutine + configs := make([]*Config, numGoroutines) + for i := 0; i < numGoroutines; i++ { + configs[i] = &Config{ + Rules: []*Rule{ + { + ID: fmt.Sprintf("concurrent-rule-%d", i), + Resource: fmt.Sprintf("/api/concurrent-%d", i), + Strategy: FixedWindow, + SpecificItems: []*SpecificItem{ + { + Identifier: Identifier{Type: Header, Value: "*"}, + KeyItems: []*KeyItem{ + { + Key: "*", + Token: Token{Number: int64(1000 + i), CountStrategy: TotalTokens}, + Time: Time{Unit: Second, Value: 60}, + }, + }, + }, + }, + }, + }, + Redis: &Redis{ + Addrs: []*RedisAddr{ + {Name: "localhost", Port: 6379}, + }, + DialTimeout: 5000, + ReadTimeout: 5000, + WriteTimeout: 5000, + PoolTimeout: 5000, + PoolSize: 10, + MinIdleConns: 5, + MaxRetries: 3, + }, + ErrorCode: int32(400 + i), + ErrorMessage: fmt.Sprintf("concurrent error %d", i), + } + } + + // Concurrent Init calls + for i := 0; i < numGoroutines; i++ { + wg.Add(1) + go func(id int) { + defer wg.Done() + + // Note: Init might fail due to Redis connection, but should not cause data races + if err := Init(configs[id]); err != nil { + // Redis connection errors are expected in test environment + if !containsRedisError(err) { + errors <- fmt.Errorf("unexpected Init error in goroutine %d: %v", id, err) + } + } + + // Test reading config after init + cfg := globalConfig.GetConfig() + if cfg != nil { + // Basic validation + _ = cfg.ErrorCode + _ = cfg.ErrorMessage + if len(cfg.Rules) > 0 { + _ = cfg.Rules[0].ID + } + } + }(i) + } + + // Wait for completion + done := make(chan struct{}) + go func() { + wg.Wait() + close(done) + }() + + select { + case <-done: + // Success + case err := <-errors: + t.Fatal(err) + case <-time.After(30 * time.Second): + t.Fatal("Test timed out") + } + + // Check for any remaining errors + close(errors) + for err := range errors { + t.Errorf("Concurrent error: %v", err) + } +} + +// Test that concurrent access doesn't cause data corruption +func TestSafeConfig_DataIntegrity(t *testing.T) { + config := &SafeConfig{} + + // Predefined configs with distinct data + configs := []*Config{ + { + ErrorCode: 100, + ErrorMessage: "config-100", + Rules: []*Rule{ + {ID: "rule-100"}, + }, + }, + { + ErrorCode: 200, + ErrorMessage: "config-200", + Rules: []*Rule{ + {ID: "rule-200"}, + }, + }, + { + ErrorCode: 300, + ErrorMessage: "config-300", + Rules: []*Rule{ + {ID: "rule-300"}, + }, + }, + } + + const numIterations = 1000 + var wg sync.WaitGroup + errors := make(chan error, 10) + + // Writer goroutines - each writes a specific config repeatedly + for configIndex, cfg := range configs { + wg.Add(1) + go func(configIdx int, safeConfig *SafeConfig, config *Config) { + defer wg.Done() + for i := 0; i < numIterations; i++ { + if err := safeConfig.SetConfig(config); err != nil { + errors <- fmt.Errorf("SetConfig failed for config %d: %v", configIdx, err) + return + } + } + }(configIndex, config, cfg) + } + + // Reader goroutines - validate data integrity + for i := 0; i < 5; i++ { + wg.Add(1) + go func(readerID int) { + defer wg.Done() + for j := 0; j < numIterations; j++ { + cfg := config.GetConfig() + if cfg != nil { + // Validate data consistency - ErrorCode should match ErrorMessage pattern + expectedMessage := fmt.Sprintf("config-%d", cfg.ErrorCode) + if cfg.ErrorMessage != expectedMessage { + errors <- fmt.Errorf("data corruption detected by reader %d: ErrorCode=%d, ErrorMessage=%s", + readerID, cfg.ErrorCode, cfg.ErrorMessage) + return + } + + // Validate rules consistency + if len(cfg.Rules) > 0 { + expectedRuleID := fmt.Sprintf("rule-%d", cfg.ErrorCode) + if cfg.Rules[0].ID != expectedRuleID { + errors <- fmt.Errorf("rule data corruption detected by reader %d: ErrorCode=%d, RuleID=%s", + readerID, cfg.ErrorCode, cfg.Rules[0].ID) + return + } + } + } + } + }(i) + } + + // Wait for completion + done := make(chan struct{}) + go func() { + wg.Wait() + close(done) + }() + + select { + case <-done: + // Success + case err := <-errors: + t.Fatal(err) + case <-time.After(30 * time.Second): + t.Fatal("Test timed out") + } + + // Check for any remaining errors + close(errors) + for err := range errors { + t.Errorf("Data integrity error: %v", err) + } +} + +// Benchmark concurrent access +func BenchmarkSafeConfig_ConcurrentSetGet(b *testing.B) { + config := &SafeConfig{} + testConfig := &Config{ + ErrorCode: 500, + ErrorMessage: "benchmark config", + Rules: []*Rule{ + { + ID: "benchmark-rule", + Resource: "/api/benchmark", + }, + }, + } + + // Pre-set a config + config.SetConfig(testConfig) + + b.ResetTimer() + b.RunParallel(func(pb *testing.PB) { + for pb.Next() { + // Mix of reads and writes (90% reads, 10% writes) + if b.N%10 == 0 { + config.SetConfig(testConfig) + } else { + config.GetConfig() + } + } + }) +} + +func BenchmarkSafeConfig_ReadOnly(b *testing.B) { + config := &SafeConfig{} + testConfig := &Config{ + ErrorCode: 500, + ErrorMessage: "benchmark config", + } + config.SetConfig(testConfig) + + b.ResetTimer() + b.RunParallel(func(pb *testing.PB) { + for pb.Next() { + config.GetConfig() + } + }) +} + +func BenchmarkSafeConfig_WriteOnly(b *testing.B) { + config := &SafeConfig{} + testConfig := &Config{ + ErrorCode: 500, + ErrorMessage: "benchmark config", + } + + b.ResetTimer() + for i := 0; i < b.N; i++ { + config.SetConfig(testConfig) + } +} + +// Test global config concurrency +func TestGlobalConfig_ConcurrentAccess(t *testing.T) { + // Save original state + originalConfig := globalConfig.GetConfig() + defer func() { + if originalConfig != nil { + globalConfig.SetConfig(originalConfig) + } + }() + + const numGoroutines = 20 + const numOperations = 500 + + var wg sync.WaitGroup + errors := make(chan error, numGoroutines) + + // Test config + testConfig := &Config{ + ErrorCode: 999, + ErrorMessage: "global test config", + Rules: []*Rule{ + { + ID: "global-rule", + Resource: "/api/global", + }, + }, + } + + // Concurrent access to global config + for i := 0; i < numGoroutines; i++ { + wg.Add(1) + go func(id int) { + defer wg.Done() + for j := 0; j < numOperations; j++ { + // Mix operations + switch j % 3 { + case 0: + // Set config + if err := globalConfig.SetConfig(testConfig); err != nil { + errors <- fmt.Errorf("globalConfig.SetConfig failed in goroutine %d: %v", id, err) + return + } + case 1: + // Get config + cfg := globalConfig.GetConfig() + if cfg != nil { + _ = cfg.ErrorCode // Basic access + } + } + } + }(i) + } + + // Wait for completion + done := make(chan struct{}) + go func() { + wg.Wait() + close(done) + }() + + select { + case <-done: + // Success + case err := <-errors: + t.Fatal(err) + case <-time.After(30 * time.Second): + t.Fatal("Test timed out") + } + + // Check for any remaining errors + close(errors) + for err := range errors { + t.Errorf("Global config concurrent error: %v", err) + } +} + +// Test edge cases with nil values +func TestSafeConfig_NilValuesConcurrency(t *testing.T) { + + const numGoroutines = 20 + var wg sync.WaitGroup + errors := make(chan error, numGoroutines) + + // Concurrent operations with nil handling + for i := 0; i < numGoroutines; i++ { + wg.Add(1) + go func(id int) { + defer wg.Done() + + config := &SafeConfig{} + // Try to set nil config (should fail gracefully) + err := config.SetConfig(nil) + if err == nil { + errors <- fmt.Errorf("SetConfig should reject nil config in goroutine %d", id) + return + } + + // Get config when nothing is set + cfg := config.GetConfig() + if cfg != nil { + errors <- fmt.Errorf("GetConfig should return nil when no config is set in goroutine %d", id) + return + } + + // Set valid config + validConfig := &Config{ErrorCode: int32(id)} + if err := config.SetConfig(validConfig); err != nil { + errors <- fmt.Errorf("SetConfig failed for valid config in goroutine %d: %v", id, err) + return + } + + // Get config again + cfg = config.GetConfig() + if cfg == nil { + errors <- fmt.Errorf("GetConfig returned nil after setting valid config in goroutine %d", id) + return + } + }(i) + } + + // Wait for completion + done := make(chan struct{}) + go func() { + wg.Wait() + close(done) + }() + + select { + case <-done: + // Success + case err := <-errors: + t.Fatal(err) + case <-time.After(15 * time.Second): + t.Fatal("Test timed out") + } + + // Check for any remaining errors + close(errors) + for err := range errors { + t.Errorf("Nil values concurrency error: %v", err) + } +} + +// Helper function to check if error is related to Redis connection +func containsRedisError(err error) bool { + if err == nil { + return false + } + errStr := err.Error() + return containsAny(errStr, []string{ + "failed to connect to redis", + "redis client", + "connection refused", + "timeout", + "no route to host", + }) +} + +func containsAny(str string, substrings []string) bool { + for _, substr := range substrings { + if strings.Contains(str, substr) { + return true + } + } + return false +} diff --git a/core/llm_token_ratelimit/metric_logger.go b/core/llm_token_ratelimit/metric_logger.go new file mode 100644 index 000000000..73718f230 --- /dev/null +++ b/core/llm_token_ratelimit/metric_logger.go @@ -0,0 +1,351 @@ +// Copyright 1999-2020 Alibaba Group Holding Ltd. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package llm_token_ratelimit + +import ( + "bufio" + "fmt" + "os" + "path/filepath" + "strconv" + "strings" + "sync" + "time" + + "github.com/alibaba/sentinel-golang/logging" + "github.com/alibaba/sentinel-golang/util" +) + +var ( + metricLogger *MetricLogger + once sync.Once +) + +type MetricLogger struct { + mu sync.RWMutex + writer *bufio.Writer + file *os.File + baseDir string + fileName string + currentFileSize uint64 + maxFileSize uint64 + maxFileAmount uint32 + flushIntervalSec uint32 + buffer []MetricItem + bufferSize int + ticker *time.Ticker + stopChan chan struct{} + usePid bool +} + +type MetricLoggerConfig struct { + AppName string + LogDir string + MaxFileSize uint64 + MaxFileAmount uint32 + FlushInterval uint32 + UsePid bool +} + +type MetricItem struct { + Timestamp uint64 `json:"timestamp"` + RequestID string `json:"request_id"` + LimitKey string `json:"limit_key"` + + // PETA.Withhold + CurrentCapacity int64 `json:"current_capacity"` + EstimatedToken int64 `json:"estimated_token"` + Difference int64 `json:"difference"` + TokenizationLength int `json:"tokenization_length"` + WaitingTime int64 `json:"waiting_time"` + + // PETA.Correct + ActualToken int `json:"actual_token"` + CorrectResult int64 `json:"correct_result"` +} + +func InitMetricLogger(config *MetricLoggerConfig) error { + var err error + once.Do(func() { + metricLogger, err = newMetricLogger(config) + if err != nil { + logging.Error(err, "[LLMTokenRateLimit] failed to initialize the MetricLogger") + return + } + metricLogger.startPeriodicFlush() + }) + return err +} + +func FormMetricFileName(config *MetricLoggerConfig) string { + dot := "." + separator := "-" + serviceName := config.AppName + + if strings.Contains(serviceName, dot) { + serviceName = strings.ReplaceAll(serviceName, dot, separator) + } + filename := serviceName + separator + MetricFileNameSuffix + dot + util.FormatDate(util.CurrentTimeMillis()) + if config.UsePid { + pid := os.Getpid() + filename = filename + ".pid" + strconv.Itoa(pid) + } + return filename +} + +func newMetricLogger(config *MetricLoggerConfig) (*MetricLogger, error) { + logging.Info("[LLMTokenRateLimit] init MetricLogger", + "config", config.String(), + ) + + logDir := config.LogDir + if len(logDir) == 0 { + return nil, fmt.Errorf("log directory cannot be empty") + } + + maxFileSize := config.MaxFileSize + if maxFileSize == 0 { + maxFileSize = DefaultMaxFileSize + } + + maxFileAmount := config.MaxFileAmount + if maxFileAmount == 0 { + maxFileAmount = DefaultMaxFileAmount + } + + flushInterval := config.FlushInterval + if flushInterval == 0 { + flushInterval = DefaultFlushInterval + } + + if err := util.CreateDirIfNotExists(logDir); err != nil { + return nil, fmt.Errorf("failed to create log directory: %v", err) + } + + fileName := FormMetricFileName(config) + + ml := &MetricLogger{ + baseDir: logDir, + fileName: fileName, + maxFileSize: maxFileSize, + maxFileAmount: maxFileAmount, + flushIntervalSec: flushInterval, + bufferSize: DefaultBufferSize, + buffer: make([]MetricItem, 0, DefaultBufferSize), + stopChan: make(chan struct{}), + usePid: config.UsePid, + } + + if err := ml.createOrOpenLogFile(); err != nil { + return nil, err + } + + return ml, nil +} + +func (c *MetricLoggerConfig) String() string { + return fmt.Sprintf("MetricLoggerConfig{AppName: %s, LogDir: %s, MaxFileSize: %d, MaxFileAmount: %d, FlushInterval: %d, UsePid: %t}", + c.AppName, c.LogDir, c.MaxFileSize, c.MaxFileAmount, c.FlushInterval, c.UsePid) +} + +func (ml *MetricLogger) createOrOpenLogFile() error { + filePath := filepath.Join(ml.baseDir, ml.fileName) + + if stat, err := os.Stat(filePath); err == nil { + ml.currentFileSize = uint64(stat.Size()) + if ml.currentFileSize >= ml.maxFileSize { + if err := ml.rotateLogFile(); err != nil { + return fmt.Errorf("failed to rotate log file: %v", err) + } + } + } + + file, err := os.OpenFile(filePath, os.O_CREATE|os.O_WRONLY|os.O_APPEND, 0644) + if err != nil { + return fmt.Errorf("failed to open log file: %v", err) + } + + logging.Info("[LLMTokenRateLimit] new metric log file created", "filename", filePath) + + ml.file = file + ml.writer = bufio.NewWriterSize(file, 8192) // 8KB + + return nil +} + +func (ml *MetricLogger) rotateLogFile() error { + if ml.writer != nil { + ml.writer.Flush() + ml.writer = nil + } + if ml.file != nil { + ml.file.Close() + ml.file = nil + } + + currentPath := filepath.Join(ml.baseDir, ml.fileName) + + for i := int(ml.maxFileAmount) - 1; i >= 1; i-- { + oldPath := fmt.Sprintf("%s.%d", currentPath, i) + newPath := fmt.Sprintf("%s.%d", currentPath, i+1) + + if i == int(ml.maxFileAmount)-1 { + if err := os.Remove(newPath); err != nil && !os.IsNotExist(err) { + return fmt.Errorf("[LLMTokenRateLimit] failed to remove old log file %s: %v", + newPath, err) + } + } + + if _, err := os.Stat(oldPath); err == nil { + if err := os.Rename(oldPath, newPath); err != nil { + return fmt.Errorf("failed to rotate log file from %s to %s: %v", + oldPath, newPath, err) + } + } + } + + if _, err := os.Stat(currentPath); err == nil { + if err := os.Rename(currentPath, currentPath+".1"); err != nil { + return fmt.Errorf("failed to rotate current log file from %s to %s: %v", + currentPath, currentPath+".1", err) + } + } + + ml.currentFileSize = 0 + + return ml.createOrOpenLogFile() +} + +func (ml *MetricLogger) startPeriodicFlush() { + ml.ticker = time.NewTicker(time.Duration(ml.flushIntervalSec) * time.Second) + go func() { + for { + select { + case <-ml.ticker.C: + ml.flush() + case <-ml.stopChan: + ml.ticker.Stop() + ml.flush() + return + } + } + }() +} + +func (ml *MetricLogger) Record(item MetricItem) { + if ml == nil { + return + } + + if item.Timestamp == 0 { + item.Timestamp = util.CurrentTimeMillis() + } + + ml.mu.Lock() + defer ml.mu.Unlock() + + ml.buffer = append(ml.buffer, item) + + if len(ml.buffer) > ml.bufferSize { + ml.flushUnsafe() + } +} + +func (ml *MetricLogger) flush() { + if ml == nil { + return + } + + ml.mu.Lock() + defer ml.mu.Unlock() + + ml.flushUnsafe() +} + +func (ml *MetricLogger) flushUnsafe() { + if len(ml.buffer) == 0 || ml.writer == nil { + return + } + + for _, item := range ml.buffer { + line := ml.formatLogLine(item) + n, err := ml.writer.WriteString(line) + if err != nil { + logging.Error(err, "[LLMTokenRateLimit] failed to write log line") + continue + } + + ml.currentFileSize += uint64(n) + + if ml.currentFileSize >= ml.maxFileSize { + ml.writer.Flush() + if err := ml.rotateLogFile(); err != nil { + logging.Error(err, "[LLMTokenRateLimit] failed to rotate log file") + } + } + } + + ml.writer.Flush() + ml.buffer = ml.buffer[:0] +} + +func (ml *MetricLogger) formatLogLine(item MetricItem) string { + return fmt.Sprintf("%s|%s|%s|%d|%d|%d|%d|%d|%d|%d\n", + util.FormatTimeMillis(item.Timestamp), + item.RequestID, + item.LimitKey, + item.CurrentCapacity, + item.WaitingTime, + item.EstimatedToken, + item.Difference, + item.TokenizationLength, + item.ActualToken, + item.CorrectResult, + ) +} + +func (ml *MetricLogger) Stop() { + if ml == nil { + return + } + + close(ml.stopChan) + + ml.mu.Lock() + defer ml.mu.Unlock() + + ml.flushUnsafe() + + if ml.writer != nil { + ml.writer = nil + } + + if ml.file != nil { + ml.file.Close() + ml.file = nil + } +} + +func RecordMetric(item MetricItem) { + if metricLogger != nil { + metricLogger.Record(item) + } +} + +func StopMetricLogger() { + if metricLogger != nil { + metricLogger.Stop() + } +} diff --git a/core/llm_token_ratelimit/metric_logger_test.go b/core/llm_token_ratelimit/metric_logger_test.go new file mode 100644 index 000000000..b82204b80 --- /dev/null +++ b/core/llm_token_ratelimit/metric_logger_test.go @@ -0,0 +1,974 @@ +// Copyright 1999-2020 Alibaba Group Holding Ltd. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package llm_token_ratelimit + +import ( + "fmt" + "io/ioutil" + "os" + "path/filepath" + "strconv" + "strings" + "sync" + "testing" + + "github.com/alibaba/sentinel-golang/util" +) + +// Test constants and helpers +const ( + TestLogDir = "./test_logs" + TestAppName = "test-app" + TestMaxFileSize = 1024 // 1KB for testing + TestMaxFileAmount = 3 + TestFlushInterval = 1 +) + +// Helper function to create test config +func createTestConfig() *MetricLoggerConfig { + return &MetricLoggerConfig{ + AppName: TestAppName, + LogDir: TestLogDir, + MaxFileSize: TestMaxFileSize, + MaxFileAmount: TestMaxFileAmount, + FlushInterval: TestFlushInterval, + UsePid: false, + } +} + +// Helper function to clean up test directory +func cleanupTestDir(t *testing.T) { + if err := os.RemoveAll(TestLogDir); err != nil { + t.Logf("Failed to remove test directory: %v", err) + } +} + +// Helper function to create test metric item +func createTestMetricItem() MetricItem { + return MetricItem{ + Timestamp: util.CurrentTimeMillis(), + RequestID: "test-request-123", + LimitKey: "test-limit-key", + CurrentCapacity: 1000, + EstimatedToken: 50, + Difference: -50, + TokenizationLength: 25, + WaitingTime: 100, + ActualToken: 45, + CorrectResult: 0, + } +} + +// Test MetricLoggerConfig.String function +func TestMetricLoggerConfig_String(t *testing.T) { + tests := []struct { + name string + config *MetricLoggerConfig + expect string + }{ + { + name: "Complete config", + config: &MetricLoggerConfig{ + AppName: "test-app", + LogDir: "/logs", + MaxFileSize: 1024, + MaxFileAmount: 5, + FlushInterval: 2, + UsePid: true, + }, + expect: "MetricLoggerConfig{AppName: test-app, LogDir: /logs, MaxFileSize: 1024, MaxFileAmount: 5, FlushInterval: 2, UsePid: true}", + }, + { + name: "Empty config", + config: &MetricLoggerConfig{ + AppName: "", + LogDir: "", + MaxFileSize: 0, + MaxFileAmount: 0, + FlushInterval: 0, + UsePid: false, + }, + expect: "MetricLoggerConfig{AppName: , LogDir: , MaxFileSize: 0, MaxFileAmount: 0, FlushInterval: 0, UsePid: false}", + }, + { + name: "Config with special characters", + config: &MetricLoggerConfig{ + AppName: "app-with.dots", + LogDir: "/path/with spaces", + MaxFileSize: 999999, + MaxFileAmount: 10, + FlushInterval: 5, + UsePid: true, + }, + expect: "MetricLoggerConfig{AppName: app-with.dots, LogDir: /path/with spaces, MaxFileSize: 999999, MaxFileAmount: 10, FlushInterval: 5, UsePid: true}", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + result := tt.config.String() + if result != tt.expect { + t.Errorf("Expected %q, got %q", tt.expect, result) + } + }) + } +} + +// Test FormMetricFileName function +func TestFormMetricFileName(t *testing.T) { + tests := []struct { + name string + config *MetricLoggerConfig + verify func(string) bool + }{ + { + name: "Basic app name without dots", + config: &MetricLoggerConfig{ + AppName: "testapp", + UsePid: false, + }, + verify: func(filename string) bool { + return strings.Contains(filename, "testapp-"+MetricFileNameSuffix) && + strings.Contains(filename, util.FormatDate(util.CurrentTimeMillis())) && + !strings.Contains(filename, ".pid") + }, + }, + { + name: "App name with dots (should be replaced with dashes)", + config: &MetricLoggerConfig{ + AppName: "test.app.name", + UsePid: false, + }, + verify: func(filename string) bool { + return strings.Contains(filename, "test-app-name-"+MetricFileNameSuffix) && + !strings.Contains(filename, "test.app.name") + }, + }, + { + name: "With PID enabled", + config: &MetricLoggerConfig{ + AppName: "testapp", + UsePid: true, + }, + verify: func(filename string) bool { + pid := os.Getpid() + expectedPidSuffix := ".pid" + strconv.Itoa(pid) + return strings.Contains(filename, "testapp-"+MetricFileNameSuffix) && + strings.Contains(filename, expectedPidSuffix) + }, + }, + { + name: "Empty app name", + config: &MetricLoggerConfig{ + AppName: "", + UsePid: false, + }, + verify: func(filename string) bool { + return strings.HasPrefix(filename, "-"+MetricFileNameSuffix) + }, + }, + { + name: "App name with multiple consecutive dots", + config: &MetricLoggerConfig{ + AppName: "test..app...name", + UsePid: false, + }, + verify: func(filename string) bool { + return strings.Contains(filename, "test--app---name-"+MetricFileNameSuffix) + }, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + filename := FormMetricFileName(tt.config) + + if filename == "" { + t.Error("Expected non-empty filename, got empty string") + } + + if !tt.verify(filename) { + t.Errorf("Filename %q did not meet verification criteria", filename) + } + }) + } +} + +// Test InitMetricLogger function +func TestInitMetricLogger(t *testing.T) { + defer cleanupTestDir(t) + + // Reset global state + metricLogger = nil + once = sync.Once{} + + tests := []struct { + name string + config *MetricLoggerConfig + expectErr bool + }{ + { + name: "Valid config", + config: createTestConfig(), + expectErr: false, + }, + { + name: "Empty log directory", + config: &MetricLoggerConfig{ + AppName: TestAppName, + LogDir: "", + MaxFileSize: TestMaxFileSize, + MaxFileAmount: TestMaxFileAmount, + FlushInterval: TestFlushInterval, + }, + expectErr: true, + }, + { + name: "Config with default values", + config: &MetricLoggerConfig{ + AppName: TestAppName, + LogDir: TestLogDir, + // Other fields will use defaults + }, + expectErr: false, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + // Reset global state for each test + metricLogger = nil + once = sync.Once{} + + err := InitMetricLogger(tt.config) + + if tt.expectErr { + if err == nil { + t.Error("Expected error, got nil") + } + } else { + if err != nil { + t.Errorf("Expected no error, got %v", err) + } + + if metricLogger == nil { + t.Error("Expected metricLogger to be initialized, got nil") + } + } + + // Clean up after each test + if metricLogger != nil { + metricLogger.Stop() + } + cleanupTestDir(t) + }) + } +} + +// Test InitMetricLogger multiple calls (should only initialize once) +func TestInitMetricLogger_MultipleCallsOnlyInitializeOnce(t *testing.T) { + defer cleanupTestDir(t) + + // Reset global state + metricLogger = nil + once = sync.Once{} + + config := createTestConfig() + + // First call + err1 := InitMetricLogger(config) + if err1 != nil { + t.Fatalf("First initialization failed: %v", err1) + } + + firstLogger := metricLogger + + // Second call + err2 := InitMetricLogger(config) + if err2 != nil { + t.Errorf("Second call should not return error, got %v", err2) + } + + // Should be the same instance + if metricLogger != firstLogger { + t.Error("Expected same logger instance on multiple calls") + } + + // Clean up + if metricLogger != nil { + metricLogger.Stop() + } +} + +// Test newMetricLogger function +func TestNewMetricLogger(t *testing.T) { + defer cleanupTestDir(t) + + tests := []struct { + name string + config *MetricLoggerConfig + expectErr bool + verify func(*MetricLogger) bool + }{ + { + name: "Valid config", + config: createTestConfig(), + expectErr: false, + verify: func(ml *MetricLogger) bool { + return ml.baseDir == TestLogDir && + ml.maxFileSize == TestMaxFileSize && + ml.maxFileAmount == TestMaxFileAmount && + ml.flushIntervalSec == TestFlushInterval && + ml.bufferSize == DefaultBufferSize + }, + }, + { + name: "Config with zero values (should use defaults)", + config: &MetricLoggerConfig{ + AppName: TestAppName, + LogDir: TestLogDir, + // Other fields are zero, should use defaults + }, + expectErr: false, + verify: func(ml *MetricLogger) bool { + return ml.maxFileSize == DefaultMaxFileSize && + ml.maxFileAmount == DefaultMaxFileAmount && + ml.flushIntervalSec == DefaultFlushInterval + }, + }, + { + name: "Empty log directory", + config: &MetricLoggerConfig{ + AppName: TestAppName, + LogDir: "", + MaxFileSize: TestMaxFileSize, + MaxFileAmount: TestMaxFileAmount, + FlushInterval: TestFlushInterval, + }, + expectErr: true, + verify: nil, + }, + { + name: "Invalid log directory path", + config: &MetricLoggerConfig{ + AppName: TestAppName, + LogDir: "/invalid/\x00/path", // Contains null byte + MaxFileSize: TestMaxFileSize, + MaxFileAmount: TestMaxFileAmount, + FlushInterval: TestFlushInterval, + }, + expectErr: true, + verify: nil, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + ml, err := newMetricLogger(tt.config) + + if tt.expectErr { + if err == nil { + t.Error("Expected error, got nil") + } + if ml != nil { + t.Error("Expected nil MetricLogger on error, got non-nil") + } + } else { + if err != nil { + t.Errorf("Expected no error, got %v", err) + } + if ml == nil { + t.Error("Expected non-nil MetricLogger, got nil") + } else { + if tt.verify != nil && !tt.verify(ml) { + t.Error("MetricLogger verification failed") + } + ml.Stop() + } + } + + cleanupTestDir(t) + }) + } +} + +// Test MetricLogger.createOrOpenLogFile function +func TestMetricLogger_CreateOrOpenLogFile(t *testing.T) { + defer cleanupTestDir(t) + + config := createTestConfig() + ml, err := newMetricLogger(config) + if err != nil { + t.Fatalf("Failed to create MetricLogger: %v", err) + } + defer ml.Stop() + + // Verify file was created + filePath := filepath.Join(ml.baseDir, ml.fileName) + if _, err := os.Stat(filePath); os.IsNotExist(err) { + t.Error("Expected log file to be created, but it doesn't exist") + } + + // Verify file handles are set + if ml.file == nil { + t.Error("Expected file handle to be set, got nil") + } + if ml.writer == nil { + t.Error("Expected writer to be set, got nil") + } +} + +// Test MetricLogger.Record function +func TestMetricLogger_Record(t *testing.T) { + defer cleanupTestDir(t) + + config := createTestConfig() + ml, err := newMetricLogger(config) + if err != nil { + t.Fatalf("Failed to create MetricLogger: %v", err) + } + defer ml.Stop() + + tests := []struct { + name string + item MetricItem + }{ + { + name: "Complete metric item", + item: createTestMetricItem(), + }, + { + name: "Metric item with zero timestamp (should be set automatically)", + item: MetricItem{ + RequestID: "test-request-456", + LimitKey: "test-limit-key-2", + CurrentCapacity: 500, + EstimatedToken: 25, + Difference: -25, + TokenizationLength: 12, + WaitingTime: 50, + ActualToken: 20, + CorrectResult: 1, + }, + }, + { + name: "Metric item with negative values", + item: MetricItem{ + Timestamp: util.CurrentTimeMillis(), + RequestID: "test-request-789", + LimitKey: "test-limit-key-3", + CurrentCapacity: -100, + EstimatedToken: -10, + Difference: 10, + TokenizationLength: 5, + WaitingTime: 0, + ActualToken: -5, + CorrectResult: 2, + }, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + initialBufferLen := len(ml.buffer) + + ml.Record(tt.item) + + // Check if item was added to buffer + if len(ml.buffer) != initialBufferLen+1 { + t.Errorf("Expected buffer length to increase by 1, got %d -> %d", + initialBufferLen, len(ml.buffer)) + } + + // Check if timestamp was set for zero timestamp + if tt.item.Timestamp == 0 { + lastItem := ml.buffer[len(ml.buffer)-1] + if lastItem.Timestamp == 0 { + t.Error("Expected timestamp to be set automatically, got 0") + } + } + }) + } +} + +// Test MetricLogger.Record with nil logger +func TestMetricLogger_Record_NilLogger(t *testing.T) { + var ml *MetricLogger = nil + item := createTestMetricItem() + + // Should not panic + ml.Record(item) +} + +// Test MetricLogger.Record buffer overflow +func TestMetricLogger_Record_BufferOverflow(t *testing.T) { + defer cleanupTestDir(t) + + config := createTestConfig() + ml, err := newMetricLogger(config) + if err != nil { + t.Fatalf("Failed to create MetricLogger: %v", err) + } + defer ml.Stop() + + // Fill buffer to capacity + for i := 0; i < DefaultBufferSize; i++ { + item := createTestMetricItem() + item.RequestID = fmt.Sprintf("request-%d", i) + ml.Record(item) + } + + // Buffer should be full but not flushed yet + if len(ml.buffer) != DefaultBufferSize { + t.Errorf("Expected buffer length to be %d, got %d", DefaultBufferSize, len(ml.buffer)) + } + + // Add one more item, should trigger flush + overflowItem := createTestMetricItem() + overflowItem.RequestID = "overflow-request" + ml.Record(overflowItem) + + // Buffer should be empty after flush + if len(ml.buffer) != 0 { + t.Errorf("Expected buffer to be empty after overflow flush, got length %d", len(ml.buffer)) + } +} + +// Test MetricLogger.flush function +func TestMetricLogger_Flush(t *testing.T) { + defer cleanupTestDir(t) + + config := createTestConfig() + ml, err := newMetricLogger(config) + if err != nil { + t.Fatalf("Failed to create MetricLogger: %v", err) + } + defer ml.Stop() + + // Add some items to buffer + for i := 0; i < 5; i++ { + item := createTestMetricItem() + item.RequestID = fmt.Sprintf("request-%d", i) + ml.Record(item) + } + + // Verify buffer is not empty + if len(ml.buffer) == 0 { + t.Error("Expected buffer to have items before flush") + } + + // Flush + ml.flush() + + // Verify buffer is empty after flush + if len(ml.buffer) != 0 { + t.Errorf("Expected buffer to be empty after flush, got length %d", len(ml.buffer)) + } + + // Verify data was written to file + filePath := filepath.Join(ml.baseDir, ml.fileName) + content, err := ioutil.ReadFile(filePath) + if err != nil { + t.Fatalf("Failed to read log file: %v", err) + } + + if len(content) == 0 { + t.Error("Expected log file to contain data after flush, got empty file") + } + + // Verify log format + lines := strings.Split(string(content), "\n") + nonEmptyLines := 0 + for _, line := range lines { + if strings.TrimSpace(line) != "" { + nonEmptyLines++ + // Check if line contains expected separators + if !strings.Contains(line, "|") { + t.Errorf("Expected log line to contain '|' separators, got: %s", line) + } + } + } + + if nonEmptyLines != 5 { + t.Errorf("Expected 5 log lines, got %d", nonEmptyLines) + } +} + +// Test MetricLogger.flush with nil logger +func TestMetricLogger_Flush_NilLogger(t *testing.T) { + var ml *MetricLogger = nil + + // Should not panic + ml.flush() +} + +// Test MetricLogger.formatLogLine function +func TestMetricLogger_FormatLogLine(t *testing.T) { + defer cleanupTestDir(t) + + config := createTestConfig() + ml, err := newMetricLogger(config) + if err != nil { + t.Fatalf("Failed to create MetricLogger: %v", err) + } + defer ml.Stop() + + item := MetricItem{ + Timestamp: 1234567890123, + RequestID: "test-request", + LimitKey: "test-key", + CurrentCapacity: 1000, + EstimatedToken: 50, + Difference: -50, + TokenizationLength: 25, + WaitingTime: 100, + ActualToken: 45, + CorrectResult: 0, + } + + line := ml.formatLogLine(item) + + // Verify line is not empty + if line == "" { + t.Error("Expected non-empty log line, got empty string") + } + + // Verify line ends with newline + if !strings.HasSuffix(line, "\n") { + t.Error("Expected log line to end with newline") + } + + // Verify line contains expected separators + separatorCount := strings.Count(line, "|") + expectedSeparators := 9 // Based on the format string + if separatorCount != expectedSeparators { + t.Errorf("Expected %d separators, got %d", expectedSeparators, separatorCount) + } + + // Verify specific fields are present + expectedFields := []string{ + "test-request", + "test-key", + "1000", + "100", + "50", + "-50", + "25", + "45", + "0", + } + + for _, field := range expectedFields { + if !strings.Contains(line, field) { + t.Errorf("Expected log line to contain %q, got: %s", field, line) + } + } +} + +// Test MetricLogger.rotateLogFile function +func TestMetricLogger_RotateLogFile(t *testing.T) { + defer cleanupTestDir(t) + + config := createTestConfig() + config.MaxFileAmount = 3 // Keep it small for testing + + ml, err := newMetricLogger(config) + if err != nil { + t.Fatalf("Failed to create MetricLogger: %v", err) + } + defer ml.Stop() + + // Create initial log file with some content + initialContent := "initial log content\n" + filePath := filepath.Join(ml.baseDir, ml.fileName) + + if err := ioutil.WriteFile(filePath, []byte(initialContent), 0644); err != nil { + t.Fatalf("Failed to write initial content: %v", err) + } + + // Force rotation + err = ml.rotateLogFile() + if err != nil { + t.Errorf("Expected no error during rotation, got %v", err) + } + + // Verify original file was rotated to .1 + rotatedPath := filePath + ".1" + if _, err := os.Stat(rotatedPath); os.IsNotExist(err) { + t.Error("Expected rotated file to exist at .1 suffix") + } + + // Verify new file was created + if _, err := os.Stat(filePath); os.IsNotExist(err) { + t.Error("Expected new log file to be created after rotation") + } + + // Verify rotated file contains original content + rotatedContent, err := ioutil.ReadFile(rotatedPath) + if err != nil { + t.Fatalf("Failed to read rotated file: %v", err) + } + + if string(rotatedContent) != initialContent { + t.Errorf("Expected rotated file to contain %q, got %q", initialContent, string(rotatedContent)) + } + + // Verify current file size was reset + if ml.currentFileSize != 0 { + t.Errorf("Expected currentFileSize to be reset to 0, got %d", ml.currentFileSize) + } +} + +// Test MetricLogger.Stop function +func TestMetricLogger_Stop(t *testing.T) { + defer cleanupTestDir(t) + + config := createTestConfig() + ml, err := newMetricLogger(config) + if err != nil { + t.Fatalf("Failed to create MetricLogger: %v", err) + } + + // Start periodic flush to have ticker running + ml.startPeriodicFlush() + + // Add some items to buffer + for i := 0; i < 3; i++ { + item := createTestMetricItem() + item.RequestID = fmt.Sprintf("stop-test-request-%d", i) + ml.Record(item) + } + + // Stop the logger + ml.Stop() + + // Verify resources are cleaned up + if ml.writer != nil { + t.Error("Expected writer to be nil after stop") + } + if ml.file != nil { + t.Error("Expected file to be nil after stop") + } + + // Verify final flush occurred (buffer should be empty) + if len(ml.buffer) != 0 { + t.Errorf("Expected buffer to be empty after stop, got length %d", len(ml.buffer)) + } +} + +// Test MetricLogger.Stop with nil logger +func TestMetricLogger_Stop_NilLogger(t *testing.T) { + var ml *MetricLogger = nil + + // Should not panic + ml.Stop() +} + +// Test RecordMetric function +func TestRecordMetric(t *testing.T) { + defer cleanupTestDir(t) + + // Reset global state + metricLogger = nil + once = sync.Once{} + + // Initialize global logger + config := createTestConfig() + err := InitMetricLogger(config) + if err != nil { + t.Fatalf("Failed to initialize global metric logger: %v", err) + } + defer StopMetricLogger() + + item := createTestMetricItem() + + // Record metric + RecordMetric(item) + + // Verify item was recorded + if len(metricLogger.buffer) == 0 { + t.Error("Expected buffer to contain recorded item, got empty buffer") + } +} + +// Test RecordMetric with nil global logger +func TestRecordMetric_NilGlobalLogger(t *testing.T) { + // Ensure global logger is nil + metricLogger = nil + + item := createTestMetricItem() + + // Should not panic + RecordMetric(item) +} + +// Test StopMetricLogger function +func TestStopMetricLogger(t *testing.T) { + defer cleanupTestDir(t) + + // Reset global state + metricLogger = nil + once = sync.Once{} + + // Initialize global logger + config := createTestConfig() + err := InitMetricLogger(config) + if err != nil { + t.Fatalf("Failed to initialize global metric logger: %v", err) + } + + // Add some items + for i := 0; i < 3; i++ { + item := createTestMetricItem() + item.RequestID = fmt.Sprintf("global-stop-request-%d", i) + RecordMetric(item) + } + + // Stop global logger + StopMetricLogger() + + // Verify logger was stopped (resources cleaned up) + if metricLogger.writer != nil { + t.Error("Expected global logger writer to be nil after stop") + } + if metricLogger.file != nil { + t.Error("Expected global logger file to be nil after stop") + } +} + +// Test StopMetricLogger with nil global logger +func TestStopMetricLogger_NilGlobalLogger(t *testing.T) { + // Ensure global logger is nil + metricLogger = nil + + // Should not panic + StopMetricLogger() +} + +// Test concurrent access to MetricLogger +func TestMetricLogger_ConcurrentAccess(t *testing.T) { + defer cleanupTestDir(t) + + config := createTestConfig() + ml, err := newMetricLogger(config) + if err != nil { + t.Fatalf("Failed to create MetricLogger: %v", err) + } + defer ml.Stop() + + const numGoroutines = 6 + const itemsPerGoroutine = 9 + + var wg sync.WaitGroup + wg.Add(numGoroutines) + + // Start multiple goroutines writing concurrently + for i := 0; i < numGoroutines; i++ { + go func(goroutineID int) { + defer wg.Done() + for j := 0; j < itemsPerGoroutine; j++ { + item := createTestMetricItem() + item.RequestID = fmt.Sprintf("concurrent-%d-%d", goroutineID, j) + ml.Record(item) + } + }(i) + } + + wg.Wait() + + // Flush to ensure all data is written + ml.flush() + + // Verify data was written to file + totalLines := 0 + + // Count lines in main log file + mainFilePath := filepath.Join(ml.baseDir, ml.fileName) + if content, err := ioutil.ReadFile(mainFilePath); err == nil { + lines := strings.Split(string(content), "\n") + for _, line := range lines { + if strings.TrimSpace(line) != "" { + totalLines++ + } + } + } + + // Count lines in rotated log files (.1, .2, .3, etc.) + for i := 1; i <= int(TestMaxFileAmount); i++ { + rotatedFilePath := fmt.Sprintf("%s.%d", mainFilePath, i) + if content, err := ioutil.ReadFile(rotatedFilePath); err == nil { + lines := strings.Split(string(content), "\n") + for _, line := range lines { + if strings.TrimSpace(line) != "" { + totalLines++ + } + } + t.Logf("Found rotated file %s with %d lines", rotatedFilePath, + len(strings.Split(string(content), "\n"))-1) // -1 for empty last line + } + } + + expectedLines := numGoroutines * itemsPerGoroutine + if totalLines != expectedLines { + t.Errorf("Expected %d log lines, got %d", expectedLines, totalLines) + } +} + +// Benchmark tests +func BenchmarkMetricLogger_Record(b *testing.B) { + defer cleanupTestDir(&testing.T{}) + + config := createTestConfig() + ml, err := newMetricLogger(config) + if err != nil { + b.Fatalf("Failed to create MetricLogger: %v", err) + } + defer ml.Stop() + + item := createTestMetricItem() + + b.ResetTimer() + for i := 0; i < b.N; i++ { + item.RequestID = fmt.Sprintf("bench-request-%d", i) + ml.Record(item) + } +} + +func BenchmarkMetricLogger_FormatLogLine(b *testing.B) { + defer cleanupTestDir(&testing.T{}) + + config := createTestConfig() + ml, err := newMetricLogger(config) + if err != nil { + b.Fatalf("Failed to create MetricLogger: %v", err) + } + defer ml.Stop() + + item := createTestMetricItem() + + b.ResetTimer() + for i := 0; i < b.N; i++ { + ml.formatLogLine(item) + } +} + +func BenchmarkFormMetricFileName(b *testing.B) { + config := createTestConfig() + + b.ResetTimer() + for i := 0; i < b.N; i++ { + FormMetricFileName(config) + } +} diff --git a/core/llm_token_ratelimit/ratelimit_checker.go b/core/llm_token_ratelimit/ratelimit_checker.go new file mode 100644 index 000000000..a314962ab --- /dev/null +++ b/core/llm_token_ratelimit/ratelimit_checker.go @@ -0,0 +1,239 @@ +// Copyright 1999-2020 Alibaba Group Holding Ltd. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package llm_token_ratelimit + +import ( + _ "embed" + "fmt" + "time" + + "errors" + + "github.com/alibaba/sentinel-golang/logging" + "github.com/alibaba/sentinel-golang/util" +) + +// ================================= FixedWindowChecker ==================================== + +//go:embed script/fixed_window/query.lua +var globalFixedWindowQueryScript string + +type FixedWindowChecker struct{} + +func (c *FixedWindowChecker) Check(ctx *Context, rules []*MatchedRule) bool { + if c == nil { + return true + } + + if len(rules) == 0 { + return true + } + + for _, rule := range rules { + if !c.checkLimitKey(ctx, rule) { + return false + } + } + return true +} + +func (c *FixedWindowChecker) checkLimitKey(ctx *Context, rule *MatchedRule) bool { + if c == nil { + return true + } + + keys := []string{rule.LimitKey} + args := []interface{}{rule.TokenSize, rule.TimeWindow * 1000} + response, err := globalRedisClient.Eval(globalFixedWindowQueryScript, keys, args...) + if err != nil { + logging.Error(err, "failed to execute redis script in llm_token_ratelimit.FixedWindowChecker.checkLimitKey()", + "requestID", ctx.Get(KeyRequestID), + ) + return true + } + result := parseRedisResponse(ctx, response) + if result == nil || len(result) != 2 { + logging.Error(errors.New("invalid redis response"), + "invalid redis response in llm_token_ratelimit.FixedWindowChecker.checkLimitKey()", + "response", response, + "requestID", ctx.Get(KeyRequestID), + ) + return true + } + + remaining := result[0] + responseHeader := NewResponseHeader() + if responseHeader == nil { + logging.Error(errors.New("failed to create response header"), + "failed to create response header in llm_token_ratelimit.FixedWindowChecker.checkLimitKey()", + "requestID", ctx.Get(KeyRequestID), + ) + return true + } + defer func() { + ctx.Set(KeyResponseHeaders, responseHeader) + }() + // set response headers + responseHeader.Set(ResponseHeaderRequestID, ctx.Get(KeyRequestID).(string)) + responseHeader.Set(ResponseHeaderRemainingTokens, fmt.Sprintf("%d", remaining)) + if remaining < 0 { + // set waiting time in milliseconds + responseHeader.Set(ResponseHeaderWaitingTime, (time.Duration(result[1]) * time.Millisecond).String()) + // set error code and message + responseHeader.ErrorCode = globalConfig.GetErrorCode() + responseHeader.ErrorMessage = globalConfig.GetErrorMsg() + // reject the request + return false + } + return true +} + +// ================================= PETAChecker ==================================== + +//go:embed script/peta/withhold.lua +var globalPETAWithholdScript string + +type PETAChecker struct{} + +func (c *PETAChecker) Check(ctx *Context, rules []*MatchedRule) bool { + if c == nil || ctx == nil { + return true + } + + if len(rules) == 0 { + return true + } + + for _, rule := range rules { + if !c.checkLimitKey(ctx, rule) { + return false + } + } + return true +} + +func (c *PETAChecker) checkLimitKey(ctx *Context, rule *MatchedRule) bool { + if c == nil || ctx == nil || rule == nil { + return true + } + + prompts := []string{} + reqInfos := extractRequestInfos(ctx) + if reqInfos != nil { + prompts = reqInfos.Prompts + } + + length, err := c.countTokens(ctx, prompts, rule) + if err != nil { + logging.Error(err, "failed to count tokens in llm_token_ratelimit.PETAChecker.checkLimitKey()", + "requestID", ctx.Get(KeyRequestID), + ) + return true + } + + slidingWindowKey := fmt.Sprintf(PETASlidingWindowKeyFormat, generateHash(rule.LimitKey), rule.LimitKey) + tokenBucketKey := fmt.Sprintf(PETATokenBucketKeyFormat, generateHash(rule.LimitKey), rule.LimitKey) + tokenEncoderKey := fmt.Sprintf(TokenEncoderKeyFormat, generateHash(rule.LimitKey), rule.Encoding.Provider.String(), rule.Encoding.Model, rule.LimitKey) + + keys := []string{slidingWindowKey, tokenBucketKey, tokenEncoderKey} + args := []interface{}{length, util.CurrentTimeMillis(), rule.TokenSize, rule.TimeWindow * 1000, generateRandomString(PETARandomStringLength)} + response, err := globalRedisClient.Eval(globalPETAWithholdScript, keys, args...) + if err != nil { + logging.Error(err, "failed to execute redis script in llm_token_ratelimit.PETAChecker.checkLimitKey()", + "requestID", ctx.Get(KeyRequestID), + ) + return true + } + result := parseRedisResponse(ctx, response) + if result == nil || len(result) != 4 { + logging.Error(errors.New("invalid redis response"), + "invalid redis response in llm_token_ratelimit.PETAChecker.checkLimitKey()", + "response", response, + "requestID", ctx.Get(KeyRequestID), + ) + return true + } + + RecordMetric(MetricItem{ + Timestamp: util.CurrentTimeMillis(), + RequestID: ctx.Get(KeyRequestID).(string), + LimitKey: rule.LimitKey, + CurrentCapacity: result[0], + WaitingTime: result[1], + EstimatedToken: result[2], + Difference: result[3], + TokenizationLength: length, + }) + + // TODO: add waiting and timeout callback + waitingTime := result[1] + responseHeader := NewResponseHeader() + if responseHeader == nil { + logging.Error(errors.New("failed to create response header"), + "failed to create response header in llm_token_ratelimit.PETAChecker.checkLimitKey()", + "requestID", ctx.Get(KeyRequestID), + ) + return true + } + defer func() { + ctx.Set(KeyResponseHeaders, responseHeader) + }() + // set response headers + responseHeader.Set(ResponseHeaderRequestID, ctx.Get(KeyRequestID).(string)) + responseHeader.Set(ResponseHeaderRemainingTokens, fmt.Sprintf("%d", result[0])) + if waitingTime != PETANoWaiting { + // set waiting time in milliseconds + responseHeader.Set(ResponseHeaderWaitingTime, (time.Duration(waitingTime) * time.Millisecond).String()) + // set error code and message + responseHeader.ErrorCode = globalConfig.GetErrorCode() + responseHeader.ErrorMessage = globalConfig.GetErrorMsg() + // reject the request + return false + } + c.cacheEstimatedToken(rule, result[2]) + return true +} + +func (c *PETAChecker) countTokens(ctx *Context, prompts []string, rule *MatchedRule) (int, error) { + if c == nil { + return 0, fmt.Errorf("PETAChecker is nil") + } + + switch rule.CountStrategy { + case OutputTokens: // cannot predict output tokens + return 0, nil + case InputTokens, TotalTokens: + encoder := LookupTokenEncoder(ctx, rule.Encoding) // try to get cached encoder + if encoder == nil { + encoder = NewTokenEncoder(ctx, rule.Encoding) // create a new encoder + if encoder == nil { + return 0, fmt.Errorf("failed to create token encoder for encoding") + } + } + length, err := encoder.CountTokens(ctx, prompts, rule) + if err != nil { + return 0, fmt.Errorf("failed to count tokens: %v", err) + } + return length, nil + } + return 0, fmt.Errorf("unknown count strategy: %s", rule.CountStrategy.String()) +} + +func (c *PETAChecker) cacheEstimatedToken(rule *MatchedRule, count int64) { + if c == nil || rule == nil { + return + } + rule.EstimatedToken = count +} diff --git a/core/llm_token_ratelimit/redis_client.go b/core/llm_token_ratelimit/redis_client.go new file mode 100644 index 000000000..48757edc2 --- /dev/null +++ b/core/llm_token_ratelimit/redis_client.go @@ -0,0 +1,136 @@ +// Copyright 1999-2020 Alibaba Group Holding Ltd. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package llm_token_ratelimit + +import ( + "context" + "fmt" + "sync" + "time" + + redis "github.com/go-redis/redis/v8" +) + +type SafeRedisClient struct { + mu sync.RWMutex + client redis.UniversalClient +} + +var globalRedisClient = NewGlobalRedisClient() + +func NewGlobalRedisClient() *SafeRedisClient { + return &SafeRedisClient{} +} + +func (c *SafeRedisClient) SetRedisClient(client redis.UniversalClient) error { + if c == nil { + return fmt.Errorf("safe redis client is nil") + } + c.mu.Lock() + defer c.mu.Unlock() + + if c.client != nil { + c.client.Close() + } + + c.client = client + return nil +} + +func (c *SafeRedisClient) Init(cfg *Redis) error { + if c == nil { + return fmt.Errorf("safe redis client is nil") + } + if cfg == nil { + cfg = NewDefaultRedisConfig() + if cfg == nil { + return fmt.Errorf("redis config is nil") + } + } + + cfg.setDefaultConfigOptions() + + addrsMap := make(map[string]struct{}, len(cfg.Addrs)) + for _, addr := range cfg.Addrs { + if addr == nil { + continue + } + addrsMap[fmt.Sprintf("%s:%d", addr.Name, addr.Port)] = struct{}{} + } + addrs := make([]string, 0, len(addrsMap)) + for addr := range addrsMap { + addrs = append(addrs, addr) + } + + dialTimeout := time.Duration(cfg.DialTimeout) * time.Millisecond + readTimeout := time.Duration(cfg.ReadTimeout) * time.Millisecond + writeTimeout := time.Duration(cfg.WriteTimeout) * time.Millisecond + poolTimeout := time.Duration(cfg.PoolTimeout) * time.Millisecond + poolSize := cfg.PoolSize + minIdleConns := cfg.MinIdleConns + maxRetries := cfg.MaxRetries + + newClient := redis.NewUniversalClient( + &redis.UniversalOptions{ + Addrs: addrs, + + Username: cfg.Username, + Password: cfg.Password, + + DialTimeout: dialTimeout, + ReadTimeout: readTimeout, + WriteTimeout: writeTimeout, + PoolTimeout: poolTimeout, + + PoolSize: int(poolSize), + MinIdleConns: int(minIdleConns), + MaxRetries: int(maxRetries), + }, + ) + + if _, err := newClient.Ping(context.TODO()).Result(); err != nil { + return fmt.Errorf("failed to connect to redis cluster: %v", err) + } + // Perform lock replacement only after the new client successfully connects; + // otherwise, a deadlock will occur if the connection fails + return c.SetRedisClient(newClient) +} + +func (c *SafeRedisClient) Eval(script string, keys []string, args ...interface{}) (interface{}, error) { + if c == nil { + return nil, fmt.Errorf("safe redis client is nil") + } + c.mu.RLock() + defer c.mu.RUnlock() + + if c.client == nil { + return nil, fmt.Errorf("redis client is not initialized") + } + + return c.client.Eval(context.TODO(), script, keys, args...).Result() +} + +func (c *SafeRedisClient) Close() error { + if c == nil { + return fmt.Errorf("safe redis client is nil") + } + c.mu.Lock() + defer c.mu.Unlock() + + if c.client != nil { + return c.client.Close() + } + return nil +} diff --git a/core/llm_token_ratelimit/request_info.go b/core/llm_token_ratelimit/request_info.go new file mode 100644 index 000000000..7135e3b25 --- /dev/null +++ b/core/llm_token_ratelimit/request_info.go @@ -0,0 +1,61 @@ +// Copyright 1999-2020 Alibaba Group Holding Ltd. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package llm_token_ratelimit + +type RequestInfos struct { + Headers map[string][]string `json:"headers"` + Prompts []string `json:"prompts"` +} + +type RequestInfo func(*RequestInfos) + +func WithHeader(headers map[string][]string) RequestInfo { + return func(infos *RequestInfos) { + infos.Headers = headers + } +} + +func WithPrompts(prompts []string) RequestInfo { + return func(infos *RequestInfos) { + infos.Prompts = prompts + } +} + +func GenerateRequestInfos(ri ...RequestInfo) *RequestInfos { + infos := new(RequestInfos) + for _, info := range ri { + if info != nil { + info(infos) + } + } + return infos +} + +func extractRequestInfos(ctx *Context) *RequestInfos { + if ctx == nil { + return nil + } + + reqInfosRaw := ctx.Get(KeyRequestInfos) + if reqInfosRaw == nil { + return nil + } + + if reqInfos, ok := reqInfosRaw.(*RequestInfos); ok && reqInfos != nil { + return reqInfos + } + + return nil +} diff --git a/core/llm_token_ratelimit/request_info_test.go b/core/llm_token_ratelimit/request_info_test.go new file mode 100644 index 000000000..11e1dd717 --- /dev/null +++ b/core/llm_token_ratelimit/request_info_test.go @@ -0,0 +1,769 @@ +// Copyright 1999-2020 Alibaba Group Holding Ltd. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package llm_token_ratelimit + +import ( + "fmt" + "reflect" + "sync" + "testing" + "time" +) + +func TestWithHeader(t *testing.T) { + tests := []struct { + name string + headers map[string][]string + expected map[string][]string + }{ + { + name: "normal headers", + headers: map[string][]string{ + "User-Id": {"test123"}, + "App-Id": {"app456"}, + }, + expected: map[string][]string{ + "User-Id": {"test123"}, + "App-Id": {"app456"}, + }, + }, + { + name: "empty headers", + headers: map[string][]string{}, + expected: map[string][]string{}, + }, + { + name: "nil headers", + headers: nil, + expected: nil, + }, + { + name: "single header with multiple values", + headers: map[string][]string{ + "Accept": {"application/json", "text/html"}, + }, + expected: map[string][]string{ + "Accept": {"application/json", "text/html"}, + }, + }, + { + name: "header with empty value", + headers: map[string][]string{ + "Empty-Header": {""}, + }, + expected: map[string][]string{ + "Empty-Header": {""}, + }, + }, + { + name: "header with nil slice", + headers: map[string][]string{ + "Nil-Values": nil, + }, + expected: map[string][]string{ + "Nil-Values": nil, + }, + }, + { + name: "complex headers", + headers: map[string][]string{ + "Content-Type": {"application/json; charset=utf-8"}, + "Authorization": {"Bearer token123"}, + "Cache-Control": {"no-cache", "no-store"}, + "X-Custom-Info": {"custom_value"}, + }, + expected: map[string][]string{ + "Content-Type": {"application/json; charset=utf-8"}, + "Authorization": {"Bearer token123"}, + "Cache-Control": {"no-cache", "no-store"}, + "X-Custom-Info": {"custom_value"}, + }, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + infos := &RequestInfos{} + headerFunc := WithHeader(tt.headers) + headerFunc(infos) + + if !reflect.DeepEqual(infos.Headers, tt.expected) { + t.Errorf("WithHeader() failed, got %v, want %v", infos.Headers, tt.expected) + } + }) + } +} + +func TestWithPrompts(t *testing.T) { + tests := []struct { + name string + prompts []string + expected []string + }{ + { + name: "normal prompts", + prompts: []string{"prompt1", "prompt2", "prompt3"}, + expected: []string{"prompt1", "prompt2", "prompt3"}, + }, + { + name: "single prompt", + prompts: []string{"single_prompt"}, + expected: []string{"single_prompt"}, + }, + { + name: "empty prompts", + prompts: []string{}, + expected: []string{}, + }, + { + name: "nil prompts", + prompts: nil, + expected: nil, + }, + { + name: "prompts with empty strings", + prompts: []string{"", "non_empty", ""}, + expected: []string{"", "non_empty", ""}, + }, + { + name: "long prompts", + prompts: []string{ + "This is a very long prompt that contains multiple words and sentences to test the handling of lengthy input data.", + "Another long prompt with different content to ensure proper storage and retrieval of extended text data.", + }, + expected: []string{ + "This is a very long prompt that contains multiple words and sentences to test the handling of lengthy input data.", + "Another long prompt with different content to ensure proper storage and retrieval of extended text data.", + }, + }, + { + name: "prompts with special characters", + prompts: []string{ + "Prompt with special chars: !@#$%^&*()", + "Unicode prompt: 你好世界", + "Newline\nand\ttab characters", + }, + expected: []string{ + "Prompt with special chars: !@#$%^&*()", + "Unicode prompt: 你好世界", + "Newline\nand\ttab characters", + }, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + infos := &RequestInfos{} + promptFunc := WithPrompts(tt.prompts) + promptFunc(infos) + + if !reflect.DeepEqual(infos.Prompts, tt.expected) { + t.Errorf("WithPrompts() failed, got %v, want %v", infos.Prompts, tt.expected) + } + }) + } +} + +func TestGenerateRequestInfos(t *testing.T) { + tests := []struct { + name string + funcs []RequestInfo + expected *RequestInfos + }{ + { + name: "empty funcs", + funcs: []RequestInfo{}, + expected: &RequestInfos{}, + }, + { + name: "nil funcs", + funcs: nil, + expected: &RequestInfos{}, + }, + { + name: "with headers only", + funcs: []RequestInfo{ + WithHeader(map[string][]string{ + "User-Id": {"test123"}, + }), + }, + expected: &RequestInfos{ + Headers: map[string][]string{ + "User-Id": {"test123"}, + }, + }, + }, + { + name: "with prompts only", + funcs: []RequestInfo{ + WithPrompts([]string{"prompt1", "prompt2"}), + }, + expected: &RequestInfos{ + Prompts: []string{"prompt1", "prompt2"}, + }, + }, + { + name: "with both headers and prompts", + funcs: []RequestInfo{ + WithHeader(map[string][]string{ + "User-Id": {"test123"}, + "App-Id": {"app456"}, + }), + WithPrompts([]string{"prompt1", "prompt2"}), + }, + expected: &RequestInfos{ + Headers: map[string][]string{ + "User-Id": {"test123"}, + "App-Id": {"app456"}, + }, + Prompts: []string{"prompt1", "prompt2"}, + }, + }, + { + name: "with nil function in slice", + funcs: []RequestInfo{ + WithHeader(map[string][]string{ + "User-Id": {"test123"}, + }), + nil, + WithPrompts([]string{"prompt1"}), + }, + expected: &RequestInfos{ + Headers: map[string][]string{ + "User-Id": {"test123"}, + }, + Prompts: []string{"prompt1"}, + }, + }, + { + name: "overwrite headers", + funcs: []RequestInfo{ + WithHeader(map[string][]string{ + "User-Id": {"test123"}, + }), + WithHeader(map[string][]string{ + "App-Id": {"app456"}, + }), + }, + expected: &RequestInfos{ + Headers: map[string][]string{ + "App-Id": {"app456"}, + }, + }, + }, + { + name: "overwrite prompts", + funcs: []RequestInfo{ + WithPrompts([]string{"prompt1", "prompt2"}), + WithPrompts([]string{"prompt3"}), + }, + expected: &RequestInfos{ + Prompts: []string{"prompt3"}, + }, + }, + { + name: "multiple operations", + funcs: []RequestInfo{ + WithHeader(map[string][]string{ + "Initial-Header": {"initial"}, + }), + WithPrompts([]string{"initial_prompt"}), + WithHeader(map[string][]string{ + "Final-Header": {"final"}, + }), + WithPrompts([]string{"final_prompt1", "final_prompt2"}), + }, + expected: &RequestInfos{ + Headers: map[string][]string{ + "Final-Header": {"final"}, + }, + Prompts: []string{"final_prompt1", "final_prompt2"}, + }, + }, + { + name: "all nil functions", + funcs: []RequestInfo{nil, nil, nil}, + expected: &RequestInfos{}, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + got := GenerateRequestInfos(tt.funcs...) + if !reflect.DeepEqual(got, tt.expected) { + t.Errorf("GenerateRequestInfos() = %v, want %v", got, tt.expected) + } + }) + } +} + +func TestExtractRequestInfos(t *testing.T) { + tests := []struct { + name string + ctx *Context + expected *RequestInfos + }{ + { + name: "nil context", + ctx: nil, + expected: nil, + }, + { + name: "context with no request infos", + ctx: NewContext(), + expected: nil, + }, + { + name: "context with valid request infos", + ctx: func() *Context { + ctx := NewContext() + reqInfos := &RequestInfos{ + Headers: map[string][]string{ + "User-Id": {"test123"}, + }, + Prompts: []string{"prompt1"}, + } + ctx.Set(KeyRequestInfos, reqInfos) + return ctx + }(), + expected: &RequestInfos{ + Headers: map[string][]string{ + "User-Id": {"test123"}, + }, + Prompts: []string{"prompt1"}, + }, + }, + { + name: "context with invalid type for request infos", + ctx: func() *Context { + ctx := NewContext() + ctx.Set(KeyRequestInfos, "invalid_type") + return ctx + }(), + expected: nil, + }, + { + name: "context with nil request infos", + ctx: func() *Context { + ctx := NewContext() + ctx.Set(KeyRequestInfos, (*RequestInfos)(nil)) + return ctx + }(), + expected: nil, + }, + { + name: "context with empty request infos", + ctx: func() *Context { + ctx := NewContext() + reqInfos := &RequestInfos{} + ctx.Set(KeyRequestInfos, reqInfos) + return ctx + }(), + expected: &RequestInfos{}, + }, + { + name: "context with complex request infos", + ctx: func() *Context { + ctx := NewContext() + reqInfos := &RequestInfos{ + Headers: map[string][]string{ + "Content-Type": {"application/json"}, + "Authorization": {"Bearer token"}, + "User-Agent": {"test-client/1.0"}, + }, + Prompts: []string{ + "Tell me about AI", + "What is machine learning?", + "Explain neural networks", + }, + } + ctx.Set(KeyRequestInfos, reqInfos) + return ctx + }(), + expected: &RequestInfos{ + Headers: map[string][]string{ + "Content-Type": {"application/json"}, + "Authorization": {"Bearer token"}, + "User-Agent": {"test-client/1.0"}, + }, + Prompts: []string{ + "Tell me about AI", + "What is machine learning?", + "Explain neural networks", + }, + }, + }, + { + name: "context with different key", + ctx: func() *Context { + ctx := NewContext() + reqInfos := &RequestInfos{ + Headers: map[string][]string{ + "Test": {"value"}, + }, + } + ctx.Set("different_key", reqInfos) + return ctx + }(), + expected: nil, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + got := extractRequestInfos(tt.ctx) + if !reflect.DeepEqual(got, tt.expected) { + t.Errorf("extractRequestInfos() = %v, want %v", got, tt.expected) + } + }) + } +} + +// Edge cases and error handling tests +func TestGenerateRequestInfos_EdgeCases(t *testing.T) { + // Test with variadic function edge cases + t.Run("no arguments", func(t *testing.T) { + result := GenerateRequestInfos() + if result == nil { + t.Fatal("GenerateRequestInfos() should not return nil") + } + if result.Headers != nil || result.Prompts != nil { + t.Fatal("GenerateRequestInfos() should return empty RequestInfos") + } + }) + + t.Run("large number of functions", func(t *testing.T) { + funcs := make([]RequestInfo, 1000) + for i := 0; i < 1000; i++ { + if i%2 == 0 { + funcs[i] = WithHeader(map[string][]string{ + fmt.Sprintf("Header-%d", i): {fmt.Sprintf("value-%d", i)}, + }) + } else { + funcs[i] = WithPrompts([]string{fmt.Sprintf("prompt-%d", i)}) + } + } + + result := GenerateRequestInfos(funcs...) + if result == nil { + t.Fatal("GenerateRequestInfos() should not return nil") + } + // Last header and prompt should be set + if result.Headers == nil || result.Prompts == nil { + t.Fatal("Headers and Prompts should be set") + } + }) +} + +// Concurrency tests +func TestWithHeader_Concurrency(t *testing.T) { + const numGoroutines = 100 + const numOperations = 100 + + var wg sync.WaitGroup + wg.Add(numGoroutines) + + for i := 0; i < numGoroutines; i++ { + go func(id int) { + defer wg.Done() + for j := 0; j < numOperations; j++ { + headers := map[string][]string{ + fmt.Sprintf("Header-%d-%d", id, j): {fmt.Sprintf("value-%d-%d", id, j)}, + } + infos := &RequestInfos{} + headerFunc := WithHeader(headers) + headerFunc(infos) + + if infos.Headers == nil { + t.Errorf("Headers should not be nil") + } + } + }(i) + } + + done := make(chan bool) + go func() { + wg.Wait() + done <- true + }() + + select { + case <-done: + // Test passed + case <-time.After(10 * time.Second): + t.Fatal("Test timeout") + } +} + +func TestGenerateRequestInfos_Concurrency(t *testing.T) { + const numGoroutines = 50 + const numOperations = 200 + + var wg sync.WaitGroup + wg.Add(numGoroutines) + + for i := 0; i < numGoroutines; i++ { + go func(id int) { + defer wg.Done() + for j := 0; j < numOperations; j++ { + funcs := []RequestInfo{ + WithHeader(map[string][]string{ + fmt.Sprintf("User-Id-%d", id): {fmt.Sprintf("user-%d-%d", id, j)}, + }), + WithPrompts([]string{fmt.Sprintf("prompt-%d-%d", id, j)}), + } + + result := GenerateRequestInfos(funcs...) + if result == nil { + t.Error("GenerateRequestInfos should not return nil") + return + } + if result.Headers == nil || result.Prompts == nil { + t.Error("Headers and Prompts should be set") + return + } + } + }(i) + } + + wg.Wait() +} + +// Stress tests +func TestGenerateRequestInfos_StressTest(t *testing.T) { + const numFunctions = 10000 + + funcs := make([]RequestInfo, numFunctions) + for i := 0; i < numFunctions; i++ { + if i%2 == 0 { + funcs[i] = WithHeader(map[string][]string{ + fmt.Sprintf("Header-%d", i): {fmt.Sprintf("value-%d", i)}, + }) + } else { + funcs[i] = WithPrompts([]string{fmt.Sprintf("prompt-%d", i)}) + } + } + + result := GenerateRequestInfos(funcs...) + if result == nil { + t.Fatal("GenerateRequestInfos should not return nil") + } + + // Verify that the last values are set correctly + expectedHeaderKey := fmt.Sprintf("Header-%d", numFunctions-2) // Last even index + if result.Headers == nil || result.Headers[expectedHeaderKey] == nil { + t.Fatal("Last header should be set") + } + + expectedPrompt := fmt.Sprintf("prompt-%d", numFunctions-1) // Last odd index + if len(result.Prompts) == 0 || result.Prompts[0] != expectedPrompt { + t.Fatal("Last prompt should be set") + } +} + +// Performance tests +func BenchmarkWithHeader(b *testing.B) { + headers := map[string][]string{ + "User-Id": {"user123"}, + "App-Id": {"app456"}, + "X-Request": {"req789"}, + "X-Version": {"v1.0"}, + } + + b.ResetTimer() + for i := 0; i < b.N; i++ { + infos := &RequestInfos{} + headerFunc := WithHeader(headers) + headerFunc(infos) + } +} + +func BenchmarkWithPrompts(b *testing.B) { + prompts := []string{ + "Tell me about artificial intelligence", + "What is machine learning?", + "Explain deep learning concepts", + "How do neural networks work?", + } + + b.ResetTimer() + for i := 0; i < b.N; i++ { + infos := &RequestInfos{} + promptFunc := WithPrompts(prompts) + promptFunc(infos) + } +} + +func BenchmarkGenerateRequestInfos(b *testing.B) { + headers := map[string][]string{ + "User-Id": {"user123"}, + "App-Id": {"app456"}, + } + prompts := []string{"prompt1", "prompt2"} + + b.ResetTimer() + for i := 0; i < b.N; i++ { + GenerateRequestInfos( + WithHeader(headers), + WithPrompts(prompts), + ) + } +} + +func BenchmarkExtractRequestInfos(b *testing.B) { + ctx := NewContext() + reqInfos := &RequestInfos{ + Headers: map[string][]string{ + "User-Id": {"test123"}, + }, + Prompts: []string{"prompt1"}, + } + ctx.Set(KeyRequestInfos, reqInfos) + + b.ResetTimer() + for i := 0; i < b.N; i++ { + extractRequestInfos(ctx) + } +} + +func BenchmarkGenerateRequestInfos_Complex(b *testing.B) { + headers := map[string][]string{ + "Content-Type": {"application/json; charset=utf-8"}, + "Authorization": {"Bearer eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9"}, + "User-Agent": {"test-client/1.0"}, + "Accept": {"application/json", "text/plain"}, + "Cache-Control": {"no-cache"}, + "X-Request-ID": {"req-12345"}, + "X-User-ID": {"user-67890"}, + "X-Session-ID": {"sess-abcdef"}, + } + prompts := []string{ + "Explain the fundamentals of quantum computing and its potential applications in cryptography", + "Describe the differences between supervised and unsupervised machine learning algorithms", + "What are the key components of a neural network and how do they work together?", + "Discuss the ethical implications of artificial intelligence in autonomous vehicles", + } + + b.ResetTimer() + for i := 0; i < b.N; i++ { + GenerateRequestInfos( + WithHeader(headers), + WithPrompts(prompts), + ) + } +} + +// Parallel benchmarks +func BenchmarkWithHeader_Parallel(b *testing.B) { + headers := map[string][]string{ + "User-Id": {"user123"}, + "App-Id": {"app456"}, + } + + b.ResetTimer() + b.RunParallel(func(pb *testing.PB) { + for pb.Next() { + infos := &RequestInfos{} + headerFunc := WithHeader(headers) + headerFunc(infos) + } + }) +} + +func BenchmarkGenerateRequestInfos_Parallel(b *testing.B) { + headers := map[string][]string{ + "User-Id": {"user123"}, + "App-Id": {"app456"}, + } + prompts := []string{"prompt1", "prompt2"} + + b.ResetTimer() + b.RunParallel(func(pb *testing.PB) { + for pb.Next() { + GenerateRequestInfos( + WithHeader(headers), + WithPrompts(prompts), + ) + } + }) +} + +func BenchmarkExtractRequestInfos_Parallel(b *testing.B) { + ctx := NewContext() + reqInfos := &RequestInfos{ + Headers: map[string][]string{ + "User-Id": {"test123"}, + }, + Prompts: []string{"prompt1"}, + } + ctx.Set(KeyRequestInfos, reqInfos) + + b.ResetTimer() + b.RunParallel(func(pb *testing.PB) { + for pb.Next() { + extractRequestInfos(ctx) + } + }) +} + +// Memory allocation benchmarks +func BenchmarkWithHeader_Memory(b *testing.B) { + headers := map[string][]string{ + "User-Id": {"user123"}, + "App-Id": {"app456"}, + } + + b.ReportAllocs() + b.ResetTimer() + for i := 0; i < b.N; i++ { + infos := &RequestInfos{} + headerFunc := WithHeader(headers) + headerFunc(infos) + } +} + +func BenchmarkGenerateRequestInfos_Memory(b *testing.B) { + headers := map[string][]string{ + "User-Id": {"user123"}, + "App-Id": {"app456"}, + } + prompts := []string{"prompt1", "prompt2"} + + b.ReportAllocs() + b.ResetTimer() + for i := 0; i < b.N; i++ { + GenerateRequestInfos( + WithHeader(headers), + WithPrompts(prompts), + ) + } +} + +func BenchmarkExtractRequestInfos_Memory(b *testing.B) { + ctx := NewContext() + reqInfos := &RequestInfos{ + Headers: map[string][]string{ + "User-Id": {"test123"}, + }, + Prompts: []string{"prompt1"}, + } + ctx.Set(KeyRequestInfos, reqInfos) + + b.ReportAllocs() + b.ResetTimer() + for i := 0; i < b.N; i++ { + extractRequestInfos(ctx) + } +} diff --git a/core/llm_token_ratelimit/response_header.go b/core/llm_token_ratelimit/response_header.go new file mode 100644 index 000000000..325ddb189 --- /dev/null +++ b/core/llm_token_ratelimit/response_header.go @@ -0,0 +1,57 @@ +// Copyright 1999-2020 Alibaba Group Holding Ltd. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package llm_token_ratelimit + +type ResponseHeader struct { + headers map[string]string + ErrorCode int32 + ErrorMessage string +} + +func NewResponseHeader() *ResponseHeader { + return &ResponseHeader{ + headers: make(map[string]string), + } +} + +func (rh *ResponseHeader) Set(key, value string) { + if rh == nil { + return + } + + if rh.headers == nil { + rh.headers = make(map[string]string) + } + rh.headers[key] = value +} + +func (rh *ResponseHeader) Get(key string) string { + if rh == nil { + return "" + } + + if rh.headers == nil { + return "" + } + return rh.headers[key] +} + +func (rh *ResponseHeader) GetAll() map[string]string { + if rh == nil { + return nil + } + + return rh.headers +} diff --git a/core/llm_token_ratelimit/response_header_test.go b/core/llm_token_ratelimit/response_header_test.go new file mode 100644 index 000000000..af61fa889 --- /dev/null +++ b/core/llm_token_ratelimit/response_header_test.go @@ -0,0 +1,603 @@ +// Copyright 1999-2020 Alibaba Group Holding Ltd. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package llm_token_ratelimit + +import ( + "fmt" + "reflect" + "testing" +) + +func TestNewResponseHeader(t *testing.T) { + tests := []struct { + name string + }{ + { + name: "create new response header", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + rh := NewResponseHeader() + if rh == nil { + t.Fatal("NewResponseHeader() should not return nil") + } + if rh.headers == nil { + t.Fatal("headers should be initialized") + } + if len(rh.headers) != 0 { + t.Fatal("headers should be empty initially") + } + if rh.ErrorCode != 0 { + t.Fatal("ErrorCode should be 0 initially") + } + if rh.ErrorMessage != "" { + t.Fatal("ErrorMessage should be empty initially") + } + }) + } +} + +func TestResponseHeader_Set(t *testing.T) { + tests := []struct { + name string + rh *ResponseHeader + key string + value string + verify func(*testing.T, *ResponseHeader) + }{ + { + name: "nil response header", + rh: nil, + key: "test", + value: "value", + verify: func(t *testing.T, rh *ResponseHeader) { + // Should not panic, and rh remains nil + if rh != nil { + t.Error("nil ResponseHeader should remain nil") + } + }, + }, + { + name: "normal set", + rh: NewResponseHeader(), + key: "Content-Type", + value: "application/json", + verify: func(t *testing.T, rh *ResponseHeader) { + if rh.Get("Content-Type") != "application/json" { + t.Error("header should be set correctly") + } + }, + }, + { + name: "set with nil headers map", + rh: &ResponseHeader{headers: nil}, + key: "Authorization", + value: "Bearer token", + verify: func(t *testing.T, rh *ResponseHeader) { + if rh.headers == nil { + t.Error("headers map should be initialized") + } + if rh.Get("Authorization") != "Bearer token" { + t.Error("header should be set correctly") + } + }, + }, + { + name: "overwrite existing key", + rh: func() *ResponseHeader { + rh := NewResponseHeader() + rh.Set("User-Id", "old_value") + return rh + }(), + key: "User-Id", + value: "new_value", + verify: func(t *testing.T, rh *ResponseHeader) { + if rh.Get("User-Id") != "new_value" { + t.Error("header should be overwritten") + } + }, + }, + { + name: "empty key", + rh: NewResponseHeader(), + key: "", + value: "empty_key_value", + verify: func(t *testing.T, rh *ResponseHeader) { + if rh.Get("") != "empty_key_value" { + t.Error("empty key should be handled") + } + }, + }, + { + name: "empty value", + rh: NewResponseHeader(), + key: "Empty-Value", + value: "", + verify: func(t *testing.T, rh *ResponseHeader) { + if rh.Get("Empty-Value") != "" { + t.Error("empty value should be set") + } + }, + }, + { + name: "special characters in key", + rh: NewResponseHeader(), + key: "X-Custom-Header_123!@#", + value: "special_value", + verify: func(t *testing.T, rh *ResponseHeader) { + if rh.Get("X-Custom-Header_123!@#") != "special_value" { + t.Error("special character key should be handled") + } + }, + }, + { + name: "unicode characters", + rh: NewResponseHeader(), + key: "Unicode-Header", + value: "测试值", + verify: func(t *testing.T, rh *ResponseHeader) { + if rh.Get("Unicode-Header") != "测试值" { + t.Error("unicode value should be handled") + } + }, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + originalRh := tt.rh + tt.rh.Set(tt.key, tt.value) + tt.verify(t, originalRh) + }) + } +} + +func TestResponseHeader_Get(t *testing.T) { + tests := []struct { + name string + rh *ResponseHeader + key string + want string + }{ + { + name: "nil response header", + rh: nil, + key: "test", + want: "", + }, + { + name: "nil headers map", + rh: &ResponseHeader{headers: nil}, + key: "test", + want: "", + }, + { + name: "existing key", + rh: func() *ResponseHeader { + rh := NewResponseHeader() + rh.Set("Content-Type", "application/json") + return rh + }(), + key: "Content-Type", + want: "application/json", + }, + { + name: "non-existing key", + rh: NewResponseHeader(), + key: "nonexistent", + want: "", + }, + { + name: "empty key", + rh: func() *ResponseHeader { + rh := NewResponseHeader() + rh.Set("", "empty_key_value") + return rh + }(), + key: "", + want: "empty_key_value", + }, + { + name: "empty value", + rh: func() *ResponseHeader { + rh := NewResponseHeader() + rh.Set("Empty-Value", "") + return rh + }(), + key: "Empty-Value", + want: "", + }, + { + name: "case sensitive key", + rh: func() *ResponseHeader { + rh := NewResponseHeader() + rh.Set("Content-Type", "application/json") + return rh + }(), + key: "content-type", + want: "", + }, + { + name: "special characters", + rh: func() *ResponseHeader { + rh := NewResponseHeader() + rh.Set("X-Special!@#", "special_value") + return rh + }(), + key: "X-Special!@#", + want: "special_value", + }, + { + name: "unicode characters", + rh: func() *ResponseHeader { + rh := NewResponseHeader() + rh.Set("Unicode", "测试值") + return rh + }(), + key: "Unicode", + want: "测试值", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + got := tt.rh.Get(tt.key) + if got != tt.want { + t.Errorf("ResponseHeader.Get() = %v, want %v", got, tt.want) + } + }) + } +} + +func TestResponseHeader_GetAll(t *testing.T) { + tests := []struct { + name string + rh *ResponseHeader + want map[string]string + }{ + { + name: "nil response header", + rh: nil, + want: nil, + }, + { + name: "nil headers map", + rh: &ResponseHeader{headers: nil}, + want: nil, + }, + { + name: "empty headers", + rh: NewResponseHeader(), + want: map[string]string{}, + }, + { + name: "single header", + rh: func() *ResponseHeader { + rh := NewResponseHeader() + rh.Set("Content-Type", "application/json") + return rh + }(), + want: map[string]string{ + "Content-Type": "application/json", + }, + }, + { + name: "multiple headers", + rh: func() *ResponseHeader { + rh := NewResponseHeader() + rh.Set("Content-Type", "application/json") + rh.Set("Authorization", "Bearer token") + rh.Set("User-Agent", "test-client") + return rh + }(), + want: map[string]string{ + "Content-Type": "application/json", + "Authorization": "Bearer token", + "User-Agent": "test-client", + }, + }, + { + name: "headers with empty values", + rh: func() *ResponseHeader { + rh := NewResponseHeader() + rh.Set("Empty-Header", "") + rh.Set("Normal-Header", "value") + return rh + }(), + want: map[string]string{ + "Empty-Header": "", + "Normal-Header": "value", + }, + }, + { + name: "headers with special characters", + rh: func() *ResponseHeader { + rh := NewResponseHeader() + rh.Set("X-Custom!@#", "special_value") + rh.Set("Unicode-Header", "测试值") + return rh + }(), + want: map[string]string{ + "X-Custom!@#": "special_value", + "Unicode-Header": "测试值", + }, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + got := tt.rh.GetAll() + if !reflect.DeepEqual(got, tt.want) { + t.Errorf("ResponseHeader.GetAll() = %v, want %v", got, tt.want) + } + }) + } +} + +// Test ResponseHeader fields +func TestResponseHeader_Fields(t *testing.T) { + tests := []struct { + name string + errorCode int32 + errorMessage string + }{ + { + name: "default values", + errorCode: 0, + errorMessage: "", + }, + { + name: "custom error values", + errorCode: 404, + errorMessage: "Not Found", + }, + { + name: "negative error code", + errorCode: -1, + errorMessage: "Internal Error", + }, + { + name: "large error code", + errorCode: 999999, + errorMessage: "Custom Error", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + rh := NewResponseHeader() + rh.ErrorCode = tt.errorCode + rh.ErrorMessage = tt.errorMessage + + if rh.ErrorCode != tt.errorCode { + t.Errorf("ErrorCode = %v, want %v", rh.ErrorCode, tt.errorCode) + } + if rh.ErrorMessage != tt.errorMessage { + t.Errorf("ErrorMessage = %v, want %v", rh.ErrorMessage, tt.errorMessage) + } + }) + } +} + +func TestResponseHeader_ErrorCodeErrorMessage(t *testing.T) { + rh := NewResponseHeader() + + // Test setting and getting error fields + rh.ErrorCode = 500 + rh.ErrorMessage = "Internal Server Error" + + if rh.ErrorCode != 500 { + t.Errorf("ErrorCode = %v, want 500", rh.ErrorCode) + } + if rh.ErrorMessage != "Internal Server Error" { + t.Errorf("ErrorMessage = %v, want 'Internal Server Error'", rh.ErrorMessage) + } + + // Test that headers still work + rh.Set("Content-Type", "application/json") + if rh.Get("Content-Type") != "application/json" { + t.Error("Headers should still work when error fields are set") + } +} + +// Edge cases +func TestResponseHeader_EdgeCases(t *testing.T) { + t.Run("very long key and value", func(t *testing.T) { + rh := NewResponseHeader() + longKey := generateLargeString(1000) + longValue := generateLargeString(2000) + + rh.Set(longKey, longValue) + retrieved := rh.Get(longKey) + + if retrieved != longValue { + t.Error("Long key and value should be handled correctly") + } + }) + + t.Run("key with newlines and tabs", func(t *testing.T) { + rh := NewResponseHeader() + specialKey := "Header\nWith\tSpecial\rChars" + specialValue := "Value\nWith\tSpecial\rChars" + + rh.Set(specialKey, specialValue) + retrieved := rh.Get(specialKey) + + if retrieved != specialValue { + t.Error("Special characters should be handled correctly") + } + }) + + t.Run("rapid set and get operations", func(t *testing.T) { + rh := NewResponseHeader() + + for i := 0; i < 1000; i++ { + key := fmt.Sprintf("rapid_%d", i) + value := fmt.Sprintf("value_%d", i) + rh.Set(key, value) + + if rh.Get(key) != value { + t.Errorf("Rapid operation failed at iteration %d", i) + } + } + }) +} + +// Test to ensure map reference safety +func TestResponseHeader_MapReferenceSafety(t *testing.T) { + rh := NewResponseHeader() + rh.Set("test", "original") + + // Get reference to internal map + allHeaders := rh.GetAll() + + // Modify the returned map + if allHeaders != nil { + allHeaders["test"] = "modified" + allHeaders["new_key"] = "new_value" + } + + // Check if original is affected (it should be, since we return direct reference) + if rh.Get("test") != "modified" { + t.Log("Map reference is shared (this is current behavior)") + } + + if rh.Get("new_key") != "new_value" { + t.Log("Map reference is shared (this is current behavior)") + } +} + +// Performance tests +func BenchmarkNewResponseHeader(b *testing.B) { + b.ResetTimer() + for i := 0; i < b.N; i++ { + NewResponseHeader() + } +} + +func BenchmarkResponseHeader_Set(b *testing.B) { + rh := NewResponseHeader() + b.ResetTimer() + for i := 0; i < b.N; i++ { + rh.Set("test_key", "test_value") + } +} + +func BenchmarkResponseHeader_Get(b *testing.B) { + rh := NewResponseHeader() + rh.Set("test_key", "test_value") + b.ResetTimer() + for i := 0; i < b.N; i++ { + rh.Get("test_key") + } +} + +func BenchmarkResponseHeader_GetAll(b *testing.B) { + rh := NewResponseHeader() + rh.Set("Content-Type", "application/json") + rh.Set("Authorization", "Bearer token") + rh.Set("User-Agent", "test-client") + b.ResetTimer() + for i := 0; i < b.N; i++ { + rh.GetAll() + } +} + +func BenchmarkResponseHeader_SetGet(b *testing.B) { + rh := NewResponseHeader() + b.ResetTimer() + for i := 0; i < b.N; i++ { + key := fmt.Sprintf("key_%d", i%100) // Reuse keys to simulate real usage + rh.Set(key, "value") + rh.Get(key) + } +} + +// Memory allocation benchmarks +func BenchmarkNewResponseHeader_Memory(b *testing.B) { + b.ReportAllocs() + b.ResetTimer() + for i := 0; i < b.N; i++ { + NewResponseHeader() + } +} + +func BenchmarkResponseHeader_Set_Memory(b *testing.B) { + rh := NewResponseHeader() + b.ReportAllocs() + b.ResetTimer() + for i := 0; i < b.N; i++ { + rh.Set("test_key", "test_value") + } +} + +func BenchmarkResponseHeader_Get_Memory(b *testing.B) { + rh := NewResponseHeader() + rh.Set("test_key", "test_value") + b.ReportAllocs() + b.ResetTimer() + for i := 0; i < b.N; i++ { + rh.Get("test_key") + } +} + +func BenchmarkResponseHeader_GetAll_Memory(b *testing.B) { + rh := NewResponseHeader() + rh.Set("Content-Type", "application/json") + rh.Set("Authorization", "Bearer token") + rh.Set("User-Agent", "test-client") + b.ReportAllocs() + b.ResetTimer() + for i := 0; i < b.N; i++ { + rh.GetAll() + } +} + +// Large data tests +func BenchmarkResponseHeader_LargeHeaders(b *testing.B) { + rh := NewResponseHeader() + + // Set up many headers + for i := 0; i < 1000; i++ { + rh.Set(fmt.Sprintf("Header-%d", i), fmt.Sprintf("Value-%d", i)) + } + + b.ResetTimer() + for i := 0; i < b.N; i++ { + rh.Get(fmt.Sprintf("Header-%d", i%1000)) + } +} + +func BenchmarkResponseHeader_LargeValues(b *testing.B) { + rh := NewResponseHeader() + largeValue := generateLargeString(1000) + + b.ResetTimer() + for i := 0; i < b.N; i++ { + rh.Set("large_header", largeValue) + rh.Get("large_header") + } +} + +// Helper function for generating large strings +func generateLargeString(length int) string { + result := make([]byte, length) + for i := range result { + result[i] = byte('a' + (i % 26)) + } + return string(result) +} diff --git a/core/llm_token_ratelimit/rule.go b/core/llm_token_ratelimit/rule.go new file mode 100644 index 000000000..5f78b0029 --- /dev/null +++ b/core/llm_token_ratelimit/rule.go @@ -0,0 +1,144 @@ +// Copyright 1999-2020 Alibaba Group Holding Ltd. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package llm_token_ratelimit + +import ( + "fmt" + "strings" +) + +type Rule struct { + // ID represents the unique ID of the rule (optional). + ID string `json:"id,omitempty" yaml:"id,omitempty"` + // Rule resource name, supporting regular expressions, Must be global match. + Resource string `json:"resource" yaml:"resource"` + // Rate limiting strategy. + Strategy Strategy `json:"strategy" yaml:"strategy"` + // Token encoding method, exclusively for peta rate limiting strategy + Encoding TokenEncoding `json:"encoding" yaml:"encoding"` + // Specific rule items + SpecificItems []*SpecificItem `json:"specificItems" yaml:"specificItems"` +} + +func (r *Rule) ResourceName() string { + if r == nil { + return "Rule{nil}" + } + return r.Resource +} + +func (r *Rule) String() string { + if r == nil { + return "Rule{nil}" + } + + var sb strings.Builder + sb.WriteString("Rule{") + + if r.ID != "" { + sb.WriteString(fmt.Sprintf("ID:%s, ", r.ID)) + } + + sb.WriteString(fmt.Sprintf("Resource:%s, ", r.Resource)) + sb.WriteString(fmt.Sprintf("Strategy:%s, ", r.Strategy.String())) + sb.WriteString(fmt.Sprintf("Encoding:%s", r.Encoding.String())) + + if len(r.SpecificItems) > 0 { + sb.WriteString(", SpecificItems:[") + for i, item := range r.SpecificItems { + if i > 0 { + sb.WriteString(", ") + } + sb.WriteString(item.String()) + } + sb.WriteString("]") + } else { + sb.WriteString(", SpecificItems:[]") + } + + sb.WriteString("}") + + return sb.String() +} + +func (r *Rule) setDefaultRuleOption() { + if r == nil { + return + } + + if len(r.Resource) == 0 { + r.Resource = DefaultResourcePattern + } + + if len(r.Encoding.Model) == 0 { + r.Encoding.Model = DefaultTokenEncodingModel[r.Encoding.Provider] + } + + for idx1, specificItem := range r.SpecificItems { + if specificItem == nil { + continue + } + if len(specificItem.Identifier.Value) == 0 { + r.SpecificItems[idx1].Identifier.Value = DefaultIdentifierValuePattern + } + for idx2, keyItem := range specificItem.KeyItems { + if keyItem == nil { + continue + } + if len(keyItem.Key) == 0 { + r.SpecificItems[idx1].KeyItems[idx2].Key = DefaultKeyPattern + } + } + } +} + +func (r *Rule) filterDuplicatedItem() { + if r == nil { + return + } + + occuredKeyItem := make(map[string]struct{}) + var specificItems []*SpecificItem + for idx1 := len(r.SpecificItems) - 1; idx1 >= 0; idx1-- { + if r.SpecificItems[idx1] == nil { + continue + } + var keyItems []*KeyItem + for idx2 := len(r.SpecificItems[idx1].KeyItems) - 1; idx2 >= 0; idx2-- { + if r.SpecificItems[idx1].KeyItems[idx2] == nil { + continue + } + hash := generateHash( + r.SpecificItems[idx1].Identifier.String(), + r.SpecificItems[idx1].KeyItems[idx2].Key, + r.SpecificItems[idx1].KeyItems[idx2].Token.CountStrategy.String(), + r.SpecificItems[idx1].KeyItems[idx2].Time.String(), + ) + if _, exists := occuredKeyItem[hash]; exists { + continue + } + occuredKeyItem[hash] = struct{}{} + keyItems = append(keyItems, r.SpecificItems[idx1].KeyItems[idx2]) + } + if len(keyItems) == 0 { + continue + } + specificItems = append(specificItems, &SpecificItem{ + Identifier: r.SpecificItems[idx1].Identifier, + KeyItems: keyItems, + }) + } + r.SpecificItems = specificItems +} diff --git a/core/llm_token_ratelimit/rule_collector.go b/core/llm_token_ratelimit/rule_collector.go new file mode 100644 index 000000000..6dafdb268 --- /dev/null +++ b/core/llm_token_ratelimit/rule_collector.go @@ -0,0 +1,101 @@ +// Copyright 1999-2020 Alibaba Group Holding Ltd. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package llm_token_ratelimit + +import ( + "errors" + "fmt" + + "github.com/alibaba/sentinel-golang/logging" +) + +type BaseRuleCollector struct{} + +func (c *BaseRuleCollector) Collect(ctx *Context, rule *Rule) []*MatchedRule { + if c == nil || rule == nil { + return nil + } + + reqInfos := extractRequestInfos(ctx) // allow nil for global rate limit + + resourceHash := generateHash(rule.Resource) + ruleStrategy := rule.Strategy.String() + + estimatedSize := 0 + for _, item := range rule.SpecificItems { + if item == nil || item.KeyItems == nil { + continue + } + estimatedSize += len(item.KeyItems) + } + + ruleMap := make(map[string]*MatchedRule, estimatedSize) + + for _, specificItem := range rule.SpecificItems { + if specificItem == nil || specificItem.KeyItems == nil { + continue + } + + identifierChecker := globalRuleMatcher.getIdentifierChecker(specificItem.Identifier.Type) + if identifierChecker == nil { + logging.Error(errors.New("unknown identifier.type"), + "unknown identifier.type in llm_token_ratelimit.BaseRuleCollector.Collect()", + "identifier.type", specificItem.Identifier.Type.String(), + "requestID", ctx.Get(KeyRequestID), + ) + continue + } + + identifierType := specificItem.Identifier.Type.String() + + for _, keyItem := range specificItem.KeyItems { + if !identifierChecker.Check(ctx, reqInfos, specificItem.Identifier, keyItem.Key) { + continue + } + + timeWindow := keyItem.Time.convertToSeconds() + if timeWindow == ErrorTimeDuration { + logging.Error(errors.New("error time window"), + "error time window in llm_token_ratelimit.BaseRuleCollector.Collect()", + "requestID", ctx.Get(KeyRequestID), + ) + continue + } + + limitKey := fmt.Sprintf(RedisRatelimitKeyFormat, + resourceHash, + ruleStrategy, + identifierType, + timeWindow, + keyItem.Token.CountStrategy.String(), + ) + ruleMap[limitKey] = &MatchedRule{ + Strategy: rule.Strategy, + LimitKey: limitKey, + TimeWindow: timeWindow, + TokenSize: keyItem.Token.Number, + CountStrategy: keyItem.Token.CountStrategy, + // PETA + Encoding: rule.Encoding, + } + } + } + + rules := make([]*MatchedRule, 0, len(ruleMap)) + for _, rule := range ruleMap { + rules = append(rules, rule) + } + return rules +} diff --git a/core/llm_token_ratelimit/rule_collector_test.go b/core/llm_token_ratelimit/rule_collector_test.go new file mode 100644 index 000000000..f401246e6 --- /dev/null +++ b/core/llm_token_ratelimit/rule_collector_test.go @@ -0,0 +1,598 @@ +// Copyright 1999-2020 Alibaba Group Holding Ltd. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package llm_token_ratelimit + +import ( + "fmt" + "testing" + "time" +) + +func TestBaseRuleCollector_Collect(t *testing.T) { + tests := []struct { + name string + ctx *Context + rule *Rule + expected int // Expected number of returned rules + wantNil bool + }{ + { + name: "nil collector", + ctx: NewContext(), + rule: &Rule{}, + wantNil: true, + }, + { + name: "nil rule", + ctx: NewContext(), + rule: nil, + expected: 0, + }, + { + name: "rule with nil SpecificItems", + ctx: NewContext(), + rule: &Rule{ + Resource: "test-resource", + Strategy: PETA, + SpecificItems: nil, + }, + expected: 0, + }, + { + name: "valid rule with header identifier", + ctx: func() *Context { + ctx := NewContext() + ctx.Set(KeyRequestInfos, &RequestInfos{ + Headers: map[string][]string{ + "User-Id": {"user123"}, + }, + }) + return ctx + }(), + rule: &Rule{ + Resource: "test-resource", + Strategy: PETA, + Encoding: TokenEncoding{ + Provider: OpenAIEncoderProvider, + Model: "gpt-3.5-turbo", + }, + SpecificItems: []*SpecificItem{ + { + Identifier: Identifier{ + Type: Header, + Value: "User-Id", + }, + KeyItems: []*KeyItem{ + { + Key: "user123", + Token: Token{ + Number: 1000, + CountStrategy: TotalTokens, + }, + Time: Time{ + Unit: Second, + Value: 60, + }, + }, + }, + }, + }, + }, + expected: 1, + }, + { + name: "rule with multiple key items", + ctx: func() *Context { + ctx := NewContext() + ctx.Set(KeyRequestInfos, &RequestInfos{ + Headers: map[string][]string{ + "User-Id": {"user123"}, + "App-Id": {"app456"}, + }, + }) + return ctx + }(), + rule: &Rule{ + Resource: "test-resource", + Strategy: PETA, + Encoding: TokenEncoding{ + Provider: OpenAIEncoderProvider, + Model: "gpt-3.5-turbo", + }, + SpecificItems: []*SpecificItem{ + { + Identifier: Identifier{ + Type: Header, + Value: "User-Id", + }, + KeyItems: []*KeyItem{ + { + Key: "user123", + Token: Token{ + Number: 1000, + CountStrategy: TotalTokens, + }, + Time: Time{ + Unit: Second, + Value: 60, + }, + }, + { + Key: "user123", + Token: Token{ + Number: 500, + CountStrategy: InputTokens, + }, + Time: Time{ + Unit: Minute, + Value: 5, + }, + }, + }, + }, + }, + }, + expected: 2, + }, + { + name: "rule with duplicate limit keys (should deduplicate)", + ctx: func() *Context { + ctx := NewContext() + ctx.Set(KeyRequestInfos, &RequestInfos{ + Headers: map[string][]string{ + "User-Id": {"user123"}, + }, + }) + return ctx + }(), + rule: &Rule{ + Resource: "test-resource", + Strategy: PETA, + Encoding: TokenEncoding{ + Provider: OpenAIEncoderProvider, + Model: "gpt-3.5-turbo", + }, + SpecificItems: []*SpecificItem{ + { + Identifier: Identifier{ + Type: Header, + Value: "User-Id", + }, + KeyItems: []*KeyItem{ + { + Key: "user123", + Token: Token{ + Number: 1000, + CountStrategy: TotalTokens, + }, + Time: Time{ + Unit: Second, + Value: 60, + }, + }, + { + Key: "user123", + Token: Token{ + Number: 1000, // Same configuration, should be deduplicated + CountStrategy: TotalTokens, + }, + Time: Time{ + Unit: Second, + Value: 60, + }, + }, + }, + }, + }, + }, + expected: 1, // Should be deduplicated, only return 1 + }, + { + name: "rule with invalid time window", + ctx: func() *Context { + ctx := NewContext() + ctx.Set(KeyRequestInfos, &RequestInfos{ + Headers: map[string][]string{ + "User-Id": {"user123"}, + }, + }) + return ctx + }(), + rule: &Rule{ + Resource: "test-resource", + Strategy: PETA, + SpecificItems: []*SpecificItem{ + { + Identifier: Identifier{ + Type: Header, + Value: "User-Id", + }, + KeyItems: []*KeyItem{ + { + Key: "user123", + Token: Token{ + Number: 1000, + CountStrategy: TotalTokens, + }, + Time: Time{ + Unit: TimeUnit(999), // Invalid time unit + Value: 60, + }, + }, + }, + }, + }, + }, + expected: 0, // Invalid time window, should be skipped + }, + { + name: "rule with unmatched identifier", + ctx: func() *Context { + ctx := NewContext() + ctx.Set(KeyRequestInfos, &RequestInfos{ + Headers: map[string][]string{ + "User-Id": {"user123"}, + }, + }) + return ctx + }(), + rule: &Rule{ + Resource: "test-resource", + Strategy: PETA, + SpecificItems: []*SpecificItem{ + { + Identifier: Identifier{ + Type: Header, + Value: "User-Id", + }, + KeyItems: []*KeyItem{ + { + Key: "user456", // Unmatched user ID + Token: Token{ + Number: 1000, + CountStrategy: TotalTokens, + }, + Time: Time{ + Unit: Second, + Value: 60, + }, + }, + }, + }, + }, + }, + expected: 0, // Identifier not matched, should return 0 rules + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + var collector *BaseRuleCollector + if tt.wantNil { + collector = nil + } else { + collector = &BaseRuleCollector{} + } + + result := collector.Collect(tt.ctx, tt.rule) + + if tt.wantNil { + if result != nil { + t.Errorf("Expected nil result, got %v", result) + } + return + } + + if len(result) != tt.expected { + t.Errorf("Expected %d rules, got %d", tt.expected, len(result)) + } + + // Validate basic fields of returned rules + for _, rule := range result { + if rule.LimitKey == "" { + t.Error("LimitKey should not be empty") + } + if rule.TimeWindow <= 0 { + t.Error("TimeWindow should be positive") + } + if rule.TokenSize <= 0 { + t.Error("TokenSize should be positive") + } + } + }) + } +} + +// Concurrency safety test +func TestBaseRuleCollector_Collect_Concurrency(t *testing.T) { + collector := &BaseRuleCollector{} + + ctx := NewContext() + ctx.Set(KeyRequestInfos, &RequestInfos{ + Headers: map[string][]string{ + "User-Id": {"user123"}, + }, + }) + + rule := &Rule{ + Resource: "test-resource", + Strategy: PETA, + SpecificItems: []*SpecificItem{ + { + Identifier: Identifier{ + Type: Header, + Value: "User-Id", + }, + KeyItems: []*KeyItem{ + { + Key: "user123", + Token: Token{ + Number: 1000, + CountStrategy: TotalTokens, + }, + Time: Time{ + Unit: Second, + Value: 60, + }, + }, + }, + }, + }, + } + + // Execute concurrently + done := make(chan bool, 10) + for i := 0; i < 10; i++ { + go func() { + defer func() { done <- true }() + for j := 0; j < 100; j++ { + result := collector.Collect(ctx, rule) + if len(result) != 1 { + t.Errorf("Expected 1 rule, got %d", len(result)) + } + } + }() + } + + // Wait for all goroutines to complete + for i := 0; i < 10; i++ { + select { + case <-done: + case <-time.After(5 * time.Second): + t.Fatal("Test timeout") + } + } +} + +// Edge cases test +func TestBaseRuleCollector_Collect_EdgeCases(t *testing.T) { + collector := &BaseRuleCollector{} + + tests := []struct { + name string + ctx *Context + rule *Rule + expected int + }{ + { + name: "empty specific items", + ctx: NewContext(), + rule: &Rule{ + Resource: "test", + Strategy: PETA, + SpecificItems: []*SpecificItem{}, + }, + expected: 0, + }, + { + name: "specific item with empty key items", + ctx: NewContext(), + rule: &Rule{ + Resource: "test", + Strategy: PETA, + SpecificItems: []*SpecificItem{ + { + Identifier: Identifier{Type: AllIdentifier}, + KeyItems: []*KeyItem{}, + }, + }, + }, + expected: 0, + }, + { + name: "specific item with nil key items", + ctx: NewContext(), + rule: &Rule{ + Resource: "test", + Strategy: PETA, + SpecificItems: []*SpecificItem{ + { + Identifier: Identifier{Type: AllIdentifier}, + KeyItems: nil, + }, + }, + }, + expected: 0, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + result := collector.Collect(tt.ctx, tt.rule) + if len(result) != tt.expected { + t.Errorf("Expected %d rules, got %d", tt.expected, len(result)) + } + }) + } +} + +// Performance test +func BenchmarkBaseRuleCollector_Collect(b *testing.B) { + collector := &BaseRuleCollector{} + + // Prepare test data + ctx := NewContext() + ctx.Set(KeyRequestInfos, &RequestInfos{ + Headers: map[string][]string{ + "User-Id": {"user123"}, + "App-Id": {"app456"}, + "Version": {"v1.0"}, + "Platform": {"web"}, + }, + }) + + rule := &Rule{ + Resource: "test-resource", + Strategy: PETA, + Encoding: TokenEncoding{ + Provider: OpenAIEncoderProvider, + Model: "gpt-3.5-turbo", + }, + SpecificItems: []*SpecificItem{ + { + Identifier: Identifier{ + Type: Header, + Value: "User-Id", + }, + KeyItems: []*KeyItem{ + { + Key: "user123", + Token: Token{ + Number: 1000, + CountStrategy: TotalTokens, + }, + Time: Time{ + Unit: Second, + Value: 60, + }, + }, + { + Key: "user123", + Token: Token{ + Number: 500, + CountStrategy: InputTokens, + }, + Time: Time{ + Unit: Minute, + Value: 5, + }, + }, + }, + }, + }, + } + + b.ResetTimer() + for i := 0; i < b.N; i++ { + result := collector.Collect(ctx, rule) + _ = result // Avoid compiler optimization + } +} + +// Large rules performance test +func BenchmarkBaseRuleCollector_Collect_LargeRules(b *testing.B) { + collector := &BaseRuleCollector{} + + // Prepare large amount of test data + ctx := NewContext() + headers := make(map[string][]string) + for i := 0; i < 100; i++ { + headers[fmt.Sprintf("Header-%d", i)] = []string{fmt.Sprintf("value-%d", i)} + } + ctx.Set(KeyRequestInfos, &RequestInfos{Headers: headers}) + + // Create Rule with large amount of rules + var specificItems []*SpecificItem + for i := 0; i < 50; i++ { + var keyItems []*KeyItem + for j := 0; j < 10; j++ { + keyItems = append(keyItems, &KeyItem{ + Key: fmt.Sprintf("value-%d", i), + Token: Token{ + Number: int64(1000 + j), + CountStrategy: CountStrategy(j % 3), + }, + Time: Time{ + Unit: TimeUnit(j % 4), + Value: int64(60 + j), + }, + }) + } + specificItems = append(specificItems, &SpecificItem{ + Identifier: Identifier{ + Type: Header, + Value: fmt.Sprintf("Header-%d", i), + }, + KeyItems: keyItems, + }) + } + + rule := &Rule{ + Resource: "large-test-resource", + Strategy: PETA, + SpecificItems: specificItems, + } + + b.ResetTimer() + for i := 0; i < b.N; i++ { + result := collector.Collect(ctx, rule) + _ = result + } +} + +// Memory allocation test +func BenchmarkBaseRuleCollector_Collect_Memory(b *testing.B) { + collector := &BaseRuleCollector{} + + ctx := NewContext() + ctx.Set(KeyRequestInfos, &RequestInfos{ + Headers: map[string][]string{ + "User-Id": {"user123"}, + }, + }) + + rule := &Rule{ + Resource: "test-resource", + Strategy: PETA, + SpecificItems: []*SpecificItem{ + { + Identifier: Identifier{ + Type: Header, + Value: "User-Id", + }, + KeyItems: []*KeyItem{ + { + Key: "user123", + Token: Token{ + Number: 1000, + CountStrategy: TotalTokens, + }, + Time: Time{ + Unit: Second, + Value: 60, + }, + }, + }, + }, + }, + } + + b.ReportAllocs() + b.ResetTimer() + for i := 0; i < b.N; i++ { + result := collector.Collect(ctx, rule) + _ = result + } +} diff --git a/core/llm_token_ratelimit/rule_filter.go b/core/llm_token_ratelimit/rule_filter.go new file mode 100644 index 000000000..1c9eabeff --- /dev/null +++ b/core/llm_token_ratelimit/rule_filter.go @@ -0,0 +1,55 @@ +// Copyright 1999-2020 Alibaba Group Holding Ltd. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package llm_token_ratelimit + +import "github.com/alibaba/sentinel-golang/logging" + +func FilterRules(rules []*Rule) []*Rule { + if rules == nil { + logging.Warn("[LLMTokenRateLimit] no rules to filter, returning empty slice") + return []*Rule{} + } + var copiedRules = make([]*Rule, len(rules)) + if err := deepCopyByCopier(&rules, &copiedRules); err != nil { + logging.Warn("[LLMTokenRateLimit] failed to deep copy rules, returning empty slice", + "error", err.Error(), + ) + return []*Rule{} + } + // 1. First, filter out invalid rules + // 2. Retain the latest rule corresponding to each unique resource + // 3. For each individual rule, retain the latest keyItems, so the specificItem is unique + ruleMap := make(map[string]*Rule, 16) + for _, rule := range copiedRules { + if rule == nil { + continue + } + rule.setDefaultRuleOption() + if err := IsValidRule(rule); err != nil { + logging.Warn("[LLMTokenRateLimit] ignoring invalid llm_token_ratelimit rule", + "rule", rule.String(), + "reason", err.Error(), + ) + continue + } + ruleMap[rule.Resource] = rule + } + resRules := make([]*Rule, 0, len(ruleMap)) + for _, rule := range ruleMap { + rule.filterDuplicatedItem() + resRules = append(resRules, rule) + } + return resRules +} diff --git a/core/llm_token_ratelimit/rule_filter_test.go b/core/llm_token_ratelimit/rule_filter_test.go new file mode 100644 index 000000000..836edaa72 --- /dev/null +++ b/core/llm_token_ratelimit/rule_filter_test.go @@ -0,0 +1,1045 @@ +// Copyright 1999-2020 Alibaba Group Holding Ltd. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package llm_token_ratelimit + +import ( + "fmt" + "reflect" + "testing" +) + +func TestFilterRules(t *testing.T) { + tests := []struct { + name string + rules []*Rule + expected []*Rule + validate func([]*Rule) bool + }{ + { + name: "nil rules", + rules: nil, + expected: []*Rule{}, + validate: func(result []*Rule) bool { + return len(result) == 0 + }, + }, + { + name: "empty rules", + rules: []*Rule{}, + expected: []*Rule{}, + validate: func(result []*Rule) bool { + return len(result) == 0 + }, + }, + { + name: "single valid rule", + rules: []*Rule{ + { + ID: "rule-1", + Resource: "test-resource", + Strategy: FixedWindow, + Encoding: TokenEncoding{ + Provider: OpenAIEncoderProvider, + Model: "gpt-3.5-turbo", + }, + SpecificItems: []*SpecificItem{ + { + Identifier: Identifier{ + Type: Header, + Value: "user-id", + }, + KeyItems: []*KeyItem{ + { + Key: "rate-limit", + Token: Token{ + Number: 1000, + CountStrategy: TotalTokens, + }, + Time: Time{ + Unit: Minute, + Value: 1, + }, + }, + }, + }, + }, + }, + }, + expected: nil, // Will be validated by validate function + validate: func(result []*Rule) bool { + return len(result) == 1 && + result[0].ID == "rule-1" && + result[0].Resource == "test-resource" + }, + }, + { + name: "multiple valid rules with different resources", + rules: []*Rule{ + { + ID: "rule-1", + Resource: "resource-1", + Strategy: FixedWindow, + Encoding: TokenEncoding{ + Provider: OpenAIEncoderProvider, + Model: "gpt-3.5-turbo", + }, + SpecificItems: []*SpecificItem{ + { + Identifier: Identifier{ + Type: Header, + Value: "user-id", + }, + KeyItems: []*KeyItem{ + { + Key: "rate-limit", + Token: Token{ + Number: 1000, + CountStrategy: TotalTokens, + }, + Time: Time{ + Unit: Minute, + Value: 1, + }, + }, + }, + }, + }, + }, + { + ID: "rule-2", + Resource: "resource-2", + Strategy: PETA, + Encoding: TokenEncoding{ + Provider: OpenAIEncoderProvider, + Model: "gpt-4", + }, + SpecificItems: []*SpecificItem{ + { + Identifier: Identifier{ + Type: AllIdentifier, + Value: ".*", + }, + KeyItems: []*KeyItem{ + { + Key: "global-limit", + Token: Token{ + Number: 5000, + CountStrategy: InputTokens, + }, + Time: Time{ + Unit: Hour, + Value: 1, + }, + }, + }, + }, + }, + }, + }, + expected: nil, + validate: func(result []*Rule) bool { + if len(result) != 2 { + return false + } + resources := make(map[string]bool) + for _, rule := range result { + resources[rule.Resource] = true + } + return resources["resource-1"] && resources["resource-2"] + }, + }, + { + name: "duplicate resources - should keep latest", + rules: []*Rule{ + { + ID: "rule-1", + Resource: "same-resource", + Strategy: FixedWindow, + Encoding: TokenEncoding{ + Provider: OpenAIEncoderProvider, + Model: "gpt-3.5-turbo", + }, + SpecificItems: []*SpecificItem{ + { + Identifier: Identifier{ + Type: Header, + Value: "user-id", + }, + KeyItems: []*KeyItem{ + { + Key: "old-limit", + Token: Token{ + Number: 1000, + CountStrategy: TotalTokens, + }, + Time: Time{ + Unit: Minute, + Value: 1, + }, + }, + }, + }, + }, + }, + { + ID: "rule-2", + Resource: "same-resource", + Strategy: PETA, + Encoding: TokenEncoding{ + Provider: OpenAIEncoderProvider, + Model: "gpt-4", + }, + SpecificItems: []*SpecificItem{ + { + Identifier: Identifier{ + Type: AllIdentifier, + Value: ".*", + }, + KeyItems: []*KeyItem{ + { + Key: "new-limit", + Token: Token{ + Number: 2000, + CountStrategy: InputTokens, + }, + Time: Time{ + Unit: Hour, + Value: 1, + }, + }, + }, + }, + }, + }, + }, + expected: nil, + validate: func(result []*Rule) bool { + if len(result) != 1 { + return false + } + // Should keep the latest rule (rule-2) + return result[0].ID == "rule-2" && + result[0].Strategy == PETA && + result[0].SpecificItems[0].KeyItems[0].Key == "new-limit" + }, + }, + { + name: "rules with nil elements", + rules: []*Rule{ + { + ID: "rule-1", + Resource: "resource-1", + Strategy: FixedWindow, + Encoding: TokenEncoding{ + Provider: OpenAIEncoderProvider, + Model: "gpt-3.5-turbo", + }, + SpecificItems: []*SpecificItem{ + { + Identifier: Identifier{ + Type: Header, + Value: "user-id", + }, + KeyItems: []*KeyItem{ + { + Key: "rate-limit", + Token: Token{ + Number: 1000, + CountStrategy: TotalTokens, + }, + Time: Time{ + Unit: Minute, + Value: 1, + }, + }, + }, + }, + }, + }, + nil, + { + ID: "rule-3", + Resource: "resource-3", + Strategy: PETA, + Encoding: TokenEncoding{ + Provider: OpenAIEncoderProvider, + Model: "gpt-4", + }, + SpecificItems: []*SpecificItem{ + { + Identifier: Identifier{ + Type: AllIdentifier, + Value: ".*", + }, + KeyItems: []*KeyItem{ + { + Key: "global-limit", + Token: Token{ + Number: 5000, + CountStrategy: OutputTokens, + }, + Time: Time{ + Unit: Day, + Value: 1, + }, + }, + }, + }, + }, + }, + }, + expected: nil, + validate: func(result []*Rule) bool { + if len(result) != 2 { + return false + } + ids := make(map[string]bool) + for _, rule := range result { + ids[rule.ID] = true + } + return ids["rule-1"] && ids["rule-3"] + }, + }, + { + name: "invalid rules - no specific items", + rules: []*Rule{ + { + ID: "invalid-rule", + Resource: "test-resource", + Strategy: FixedWindow, + Encoding: TokenEncoding{Provider: OpenAIEncoderProvider, Model: "gpt-3.5-turbo"}, + SpecificItems: nil, // Invalid: no specific items + }, + }, + expected: []*Rule{}, + validate: func(result []*Rule) bool { + return len(result) == 0 + }, + }, + { + name: "invalid rules - negative token number", + rules: []*Rule{ + { + ID: "invalid-rule", + Resource: "test-resource", + Strategy: FixedWindow, + Encoding: TokenEncoding{ + Provider: OpenAIEncoderProvider, + Model: "gpt-3.5-turbo", + }, + SpecificItems: []*SpecificItem{ + { + Identifier: Identifier{ + Type: Header, + Value: "user-id", + }, + KeyItems: []*KeyItem{ + { + Key: "rate-limit", + Token: Token{ + Number: -1000, // Invalid: negative number + CountStrategy: TotalTokens, + }, + Time: Time{ + Unit: Minute, + Value: 1, + }, + }, + }, + }, + }, + }, + }, + expected: []*Rule{}, + validate: func(result []*Rule) bool { + return len(result) == 0 + }, + }, + { + name: "mixed valid and invalid rules", + rules: []*Rule{ + { + ID: "valid-rule", + Resource: "valid-resource", + Strategy: FixedWindow, + Encoding: TokenEncoding{ + Provider: OpenAIEncoderProvider, + Model: "gpt-3.5-turbo", + }, + SpecificItems: []*SpecificItem{ + { + Identifier: Identifier{ + Type: Header, + Value: "user-id", + }, + KeyItems: []*KeyItem{ + { + Key: "rate-limit", + Token: Token{ + Number: 1000, + CountStrategy: TotalTokens, + }, + Time: Time{ + Unit: Minute, + Value: 1, + }, + }, + }, + }, + }, + }, + { + ID: "invalid-rule", + Resource: "invalid-resource", + Strategy: FixedWindow, + Encoding: TokenEncoding{Provider: OpenAIEncoderProvider, Model: "gpt-3.5-turbo"}, + SpecificItems: nil, // Invalid + }, + }, + expected: nil, + validate: func(result []*Rule) bool { + return len(result) == 1 && result[0].ID == "valid-rule" + }, + }, + { + name: "rules with duplicate key items", + rules: []*Rule{ + { + ID: "rule-with-duplicates", + Resource: "test-resource", + Strategy: FixedWindow, + Encoding: TokenEncoding{ + Provider: OpenAIEncoderProvider, + Model: "gpt-3.5-turbo", + }, + SpecificItems: []*SpecificItem{ + { + Identifier: Identifier{ + Type: Header, + Value: "user-id", + }, + KeyItems: []*KeyItem{ + { + Key: "rate-limit", + Token: Token{ + Number: 1000, + CountStrategy: TotalTokens, + }, + Time: Time{ + Unit: Minute, + Value: 1, + }, + }, + { + Key: "rate-limit", // Duplicate key + Token: Token{ + Number: 1000, + CountStrategy: TotalTokens, + }, + Time: Time{ + Unit: Minute, + Value: 1, + }, + }, + { + Key: "different-limit", + Token: Token{ + Number: 2000, + CountStrategy: InputTokens, + }, + Time: Time{ + Unit: Hour, + Value: 1, + }, + }, + }, + }, + }, + }, + }, + expected: nil, + validate: func(result []*Rule) bool { + if len(result) != 1 { + return false + } + // Should filter out duplicate key items + keyItems := result[0].SpecificItems[0].KeyItems + if len(keyItems) != 2 { + return false + } + // Check that both unique key items are present + keys := make(map[string]bool) + for _, item := range keyItems { + keys[item.Key] = true + } + return keys["rate-limit"] && keys["different-limit"] + }, + }, + { + name: "same token, different time - should not be considered duplicate", + rules: []*Rule{ + { + SpecificItems: []*SpecificItem{ + { + KeyItems: []*KeyItem{ + { + Token: Token{ + Number: 1000, + CountStrategy: TotalTokens, + }, + Time: Time{ + Unit: Minute, + Value: 1, + }, + }, + { + Token: Token{ + Number: 1000, + CountStrategy: TotalTokens, + }, + Time: Time{ + Unit: Minute, + Value: 2, + }, + }, + }, + }, + }, + }, + }, + expected: nil, + validate: func(result []*Rule) bool { + if len(result) != 1 { + return false + } + // Should filter out duplicate key items + keyItems := result[0].SpecificItems[0].KeyItems + if len(keyItems) != 2 { + return false + } + return keyItems[0].Time.Value != keyItems[1].Time.Value + }, + }, + { + name: "same time, different token - should be considered duplicate", + rules: []*Rule{ + { + SpecificItems: []*SpecificItem{ + { + KeyItems: []*KeyItem{ + { + Token: Token{ + Number: 1000, + CountStrategy: TotalTokens, + }, + Time: Time{ + Unit: Minute, + Value: 1, + }, + }, + { + Token: Token{ + Number: 1000, + CountStrategy: TotalTokens, + }, + Time: Time{ + Unit: Minute, + Value: 1, + }, + }, + }, + }, + }, + }, + }, + expected: nil, + validate: func(result []*Rule) bool { + if len(result) != 1 { + return false + } + // Should filter out duplicate key items + keyItems := result[0].SpecificItems[0].KeyItems + return len(keyItems) == 1 + }, + }, + { + name: "same SpecificItem and KeyItem", + rules: []*Rule{ + { + SpecificItems: []*SpecificItem{ + { + KeyItems: []*KeyItem{ + { + Token: Token{ + Number: 1000, + CountStrategy: TotalTokens, + }, + Time: Time{ + Unit: Minute, + Value: 1, + }, + }, + { + Token: Token{ + Number: 1000, + CountStrategy: TotalTokens, + }, + Time: Time{ + Unit: Minute, + Value: 1, + }, + }, + }, + }, + { + KeyItems: []*KeyItem{ + { + Token: Token{ + Number: 1000, + CountStrategy: TotalTokens, + }, + Time: Time{ + Unit: Minute, + Value: 1, + }, + }, + }, + }, + }, + }, + }, + expected: nil, + validate: func(result []*Rule) bool { + if len(result) != 1 { + return false + } + if len(result[0].SpecificItems) != 1 { + return false + } + keyItems := result[0].SpecificItems[0].KeyItems + return len(keyItems) == 1 + }, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + result := FilterRules(tt.rules) + + if tt.expected != nil { + if !reflect.DeepEqual(result, tt.expected) { + t.Errorf("FilterRules() = %v, want %v", result, tt.expected) + } + } + + if tt.validate != nil && !tt.validate(result) { + t.Errorf("FilterRules() validation failed for result: %v", result) + } + }) + } +} + +func TestFilterRules_DeepCopyFailure(t *testing.T) { + // Test case where deepCopyByCopier might fail + // This is harder to test without mocking, but we can test the error handling path + + // Create a rule that might cause copy issues + rules := []*Rule{ + { + ID: "test-rule", + Resource: "test-resource", + Strategy: FixedWindow, + Encoding: TokenEncoding{ + Provider: OpenAIEncoderProvider, + Model: "gpt-3.5-turbo", + }, + SpecificItems: []*SpecificItem{ + { + Identifier: Identifier{ + Type: Header, + Value: "user-id", + }, + KeyItems: []*KeyItem{ + { + Key: "rate-limit", + Token: Token{ + Number: 1000, + CountStrategy: TotalTokens, + }, + Time: Time{ + Unit: Minute, + Value: 1, + }, + }, + }, + }, + }, + }, + } + + result := FilterRules(rules) + + // Even if deep copy works, the result should be valid + if len(result) != 1 { + t.Errorf("Expected 1 rule, got %d", len(result)) + } +} + +func TestFilterRules_DefaultOptions(t *testing.T) { + // Test that default options are set correctly + rules := []*Rule{ + { + SpecificItems: []*SpecificItem{ + { + KeyItems: []*KeyItem{ + { + Token: Token{ + Number: 1000, + CountStrategy: TotalTokens, + }, + Time: Time{ + Unit: Minute, + Value: 1, + }, + }, + }, + }, + }, + }, + } + + result := FilterRules(rules) + + if len(result) != 1 { + t.Fatalf("Expected 1 rule, got %d", len(result)) + } + + rule := result[0] + + // Check that defaults were applied + if rule.Resource != DefaultResourcePattern { + t.Errorf("Expected default resource pattern %s, got %s", DefaultResourcePattern, rule.Resource) + } + if rule.Strategy != FixedWindow { + t.Errorf("Expected default strategy %s, got %s", FixedWindow, rule.Strategy) + } + + if rule.Encoding.Provider != OpenAIEncoderProvider { + t.Errorf("Expected default provider %s, got %s", OpenAIEncoderProvider, rule.Encoding.Provider) + } + + if rule.Encoding.Model != DefaultTokenEncodingModel[OpenAIEncoderProvider] { + t.Errorf("Expected default model %s, got %s", DefaultTokenEncodingModel[OpenAIEncoderProvider], rule.Encoding.Model) + } + + if rule.SpecificItems[0].Identifier.Type != AllIdentifier { + t.Errorf("Expected default identifier type %s, got %s", AllIdentifier, rule.SpecificItems[0].Identifier.Type) + } + + if rule.SpecificItems[0].Identifier.Value != DefaultIdentifierValuePattern { + t.Errorf("Expected default identifier value %s, got %s", DefaultIdentifierValuePattern, rule.SpecificItems[0].Identifier.Value) + } + + if rule.SpecificItems[0].KeyItems[0].Key != DefaultKeyPattern { + t.Errorf("Expected default key pattern %s, got %s", DefaultKeyPattern, rule.SpecificItems[0].KeyItems[0].Key) + } +} + +func TestFilterRules_LargeNumberOfRules(t *testing.T) { + // Stress test with a large number of rules + const numRules = 1000 + rules := make([]*Rule, numRules) + + for i := 0; i < numRules; i++ { + rules[i] = &Rule{ + ID: fmt.Sprintf("rule-%d", i), + Resource: fmt.Sprintf("resource-%d", i%100), // Create some duplicates + Strategy: Strategy(i % 2), // Alternate between strategies + Encoding: TokenEncoding{ + Provider: OpenAIEncoderProvider, + Model: "gpt-3.5-turbo", + }, + SpecificItems: []*SpecificItem{ + { + Identifier: Identifier{ + Type: IdentifierType(i % 2), + Value: fmt.Sprintf("identifier-%d", i), + }, + KeyItems: []*KeyItem{ + { + Key: fmt.Sprintf("key-%d", i), + Token: Token{ + Number: int64(i + 1), + CountStrategy: CountStrategy(i % 3), + }, + Time: Time{ + Unit: TimeUnit(i % 4), + Value: int64(i%10 + 1), + }, + }, + }, + }, + }, + } + } + + result := FilterRules(rules) + + // Should have 100 unique resources (0-99) + if len(result) != 100 { + t.Errorf("Expected 100 unique rules, got %d", len(result)) + } + + // Verify all results are valid + resources := make(map[string]bool) + for _, rule := range result { + if resources[rule.Resource] { + t.Errorf("Found duplicate resource: %s", rule.Resource) + } + resources[rule.Resource] = true + + if err := IsValidRule(rule); err != nil { + t.Errorf("Invalid rule after filtering: %v", err) + } + } +} + +func TestFilterRules_EdgeCases(t *testing.T) { + tests := []struct { + name string + rules []*Rule + expected int + }{ + { + name: "all nil rules", + rules: []*Rule{nil, nil, nil}, + expected: 0, + }, + { + name: "rule with invalid regex", + rules: []*Rule{ + { + ID: "invalid-regex-rule", + Resource: "[invalid-regex", // Invalid regex + Strategy: FixedWindow, + Encoding: TokenEncoding{ + Provider: OpenAIEncoderProvider, + Model: "gpt-3.5-turbo", + }, + SpecificItems: []*SpecificItem{ + { + Identifier: Identifier{ + Type: Header, + Value: "user-id", + }, + KeyItems: []*KeyItem{ + { + Key: "rate-limit", + Token: Token{ + Number: 1000, + CountStrategy: TotalTokens, + }, + Time: Time{ + Unit: Minute, + Value: 1, + }, + }, + }, + }, + }, + }, + }, + expected: 0, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + result := FilterRules(tt.rules) + if len(result) != tt.expected { + t.Errorf("FilterRules() returned %d rules, expected %d", len(result), tt.expected) + } + }) + } +} + +// Benchmark tests +func BenchmarkFilterRules_Small(b *testing.B) { + rules := []*Rule{ + { + ID: "rule-1", + Resource: "resource-1", + Strategy: FixedWindow, + Encoding: TokenEncoding{ + Provider: OpenAIEncoderProvider, + Model: "gpt-3.5-turbo", + }, + SpecificItems: []*SpecificItem{ + { + Identifier: Identifier{ + Type: Header, + Value: "user-id", + }, + KeyItems: []*KeyItem{ + { + Key: "rate-limit", + Token: Token{ + Number: 1000, + CountStrategy: TotalTokens, + }, + Time: Time{ + Unit: Minute, + Value: 1, + }, + }, + }, + }, + }, + }, + } + + b.ResetTimer() + for i := 0; i < b.N; i++ { + FilterRules(rules) + } +} + +func BenchmarkFilterRules_Large(b *testing.B) { + const numRules = 1000 + rules := make([]*Rule, numRules) + + for i := 0; i < numRules; i++ { + rules[i] = &Rule{ + ID: fmt.Sprintf("rule-%d", i), + Resource: fmt.Sprintf("resource-%d", i), + Strategy: Strategy(i % 2), + Encoding: TokenEncoding{ + Provider: OpenAIEncoderProvider, + Model: "gpt-3.5-turbo", + }, + SpecificItems: []*SpecificItem{ + { + Identifier: Identifier{ + Type: IdentifierType(i % 2), + Value: fmt.Sprintf("identifier-%d", i), + }, + KeyItems: []*KeyItem{ + { + Key: fmt.Sprintf("key-%d", i), + Token: Token{ + Number: int64(i + 1), + CountStrategy: CountStrategy(i % 3), + }, + Time: Time{ + Unit: TimeUnit(i % 4), + Value: int64(i%10 + 1), + }, + }, + }, + }, + }, + } + } + + b.ResetTimer() + for i := 0; i < b.N; i++ { + FilterRules(rules) + } +} + +func BenchmarkFilterRules_WithDuplicates(b *testing.B) { + const numRules = 500 + rules := make([]*Rule, numRules) + + for i := 0; i < numRules; i++ { + rules[i] = &Rule{ + ID: fmt.Sprintf("rule-%d", i), + Resource: fmt.Sprintf("resource-%d", i%50), // Many duplicates + Strategy: Strategy(i % 2), + Encoding: TokenEncoding{ + Provider: OpenAIEncoderProvider, + Model: "gpt-3.5-turbo", + }, + SpecificItems: []*SpecificItem{ + { + Identifier: Identifier{ + Type: IdentifierType(i % 2), + Value: fmt.Sprintf("identifier-%d", i), + }, + KeyItems: []*KeyItem{ + { + Key: fmt.Sprintf("key-%d", i), + Token: Token{ + Number: int64(i + 1), + CountStrategy: CountStrategy(i % 3), + }, + Time: Time{ + Unit: TimeUnit(i % 4), + Value: int64(i%10 + 1), + }, + }, + }, + }, + }, + } + } + + b.ResetTimer() + for i := 0; i < b.N; i++ { + FilterRules(rules) + } +} + +// Memory allocation benchmark +func BenchmarkFilterRules_Memory(b *testing.B) { + rules := []*Rule{ + { + ID: "rule-1", + Resource: "resource-1", + Strategy: FixedWindow, + Encoding: TokenEncoding{ + Provider: OpenAIEncoderProvider, + Model: "gpt-3.5-turbo", + }, + SpecificItems: []*SpecificItem{ + { + Identifier: Identifier{ + Type: Header, + Value: "user-id", + }, + KeyItems: []*KeyItem{ + { + Key: "rate-limit", + Token: Token{ + Number: 1000, + CountStrategy: TotalTokens, + }, + Time: Time{ + Unit: Minute, + Value: 1, + }, + }, + }, + }, + }, + }, + } + + b.ReportAllocs() + b.ResetTimer() + for i := 0; i < b.N; i++ { + FilterRules(rules) + } +} diff --git a/core/llm_token_ratelimit/rule_manager.go b/core/llm_token_ratelimit/rule_manager.go new file mode 100644 index 000000000..0d30aba13 --- /dev/null +++ b/core/llm_token_ratelimit/rule_manager.go @@ -0,0 +1,210 @@ +// Copyright 1999-2020 Alibaba Group Holding Ltd. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package llm_token_ratelimit + +import ( + "errors" + "reflect" + "strings" + "sync" + + "github.com/alibaba/sentinel-golang/logging" + "github.com/alibaba/sentinel-golang/util" +) + +var ( + ruleMap = make(map[string][]*Rule) + rwMux = &sync.RWMutex{} + currentRules = make(map[string][]*Rule, 0) + updateRuleMux = new(sync.Mutex) +) + +func LoadRules(rules []*Rule) (bool, error) { + filteredRules := FilterRules(rules) + + resRulesMap := make(map[string][]*Rule, 16) + for _, rule := range filteredRules { + resRules, exist := resRulesMap[rule.Resource] + if !exist { + resRules = make([]*Rule, 0, 1) + } + resRulesMap[rule.Resource] = append(resRules, rule) + } + + updateRuleMux.Lock() + defer updateRuleMux.Unlock() + isEqual := reflect.DeepEqual(currentRules, resRulesMap) + if isEqual { + logging.Info("[LLMTokenRateLimit] load rules is the same with current rules, so ignore load operation") + return false, nil + } + err := onRuleUpdate(resRulesMap) + return true, err +} + +func onRuleUpdate(rawResRulesMap map[string][]*Rule) (err error) { + validResRulesMap := make(map[string][]*Rule, len(rawResRulesMap)) + for res, rules := range rawResRulesMap { + if len(rules) > 0 { + validResRulesMap[res] = rules + } + } + + start := util.CurrentTimeNano() + rwMux.Lock() + ruleMap = validResRulesMap + rwMux.Unlock() + currentRules = rawResRulesMap + + logging.Debug("[LLMTokenRateLimit] time statistic(ns) for updating llm_token_ratelimit rule", + "timeCost", util.CurrentTimeNano()-start, + ) + logRuleUpdate(validResRulesMap) + return nil +} + +func LoadRulesOfResource(res string, rules []*Rule) (bool, error) { + if len(res) == 0 { + return false, errors.New("empty resource") + } + filteredRules := FilterRules(rules) + + updateRuleMux.Lock() + defer updateRuleMux.Unlock() + if len(filteredRules) == 0 { + delete(currentRules, res) + rwMux.Lock() + delete(ruleMap, res) + rwMux.Unlock() + logging.Info("[LLMTokenRateLimit] clear resource level rules", + "resource", res, + ) + return true, nil + } + + isEqual := reflect.DeepEqual(currentRules[res], filteredRules) + if isEqual { + logging.Info("[LLMTokenRateLimit] load resource level rules is the same with current resource level rules, so ignore load operation") + return false, nil + } + + err := onResourceRuleUpdate(res, filteredRules) + return true, err +} + +func onResourceRuleUpdate(res string, rawResRules []*Rule) (err error) { + start := util.CurrentTimeNano() + rwMux.Lock() + if len(rawResRules) == 0 { + delete(ruleMap, res) + } else { + ruleMap[res] = rawResRules + } + rwMux.Unlock() + currentRules[res] = rawResRules + logging.Debug("[LLMTokenRateLimit] time statistic(ns) for updating llm_token_ratelimit rule", + "timeCost", util.CurrentTimeNano()-start, + ) + logging.Info("[LLMTokenRateLimit] load resource level rules", + "resource", res, + "filteredRules", rawResRules, + ) + return nil +} + +func ClearRules() error { + _, err := LoadRules(nil) + return err +} + +func ClearRulesOfResource(res string) error { + _, err := LoadRulesOfResource(res, nil) + return err +} + +func GetRules() []Rule { + rules := getRules() + ret := make([]Rule, 0, len(rules)) + for _, rule := range rules { + ret = append(ret, *rule) + } + return ret +} + +func GetRulesOfResource(res string) []Rule { + rules := getRulesOfResource(res) + ret := make([]Rule, 0, len(rules)) + for _, rule := range rules { + ret = append(ret, *rule) + } + return ret +} + +func getRules() []*Rule { + rwMux.RLock() + defer rwMux.RUnlock() + + return rulesFrom(ruleMap) +} + +func getRulesOfResource(res string) []*Rule { + rwMux.RLock() + defer rwMux.RUnlock() + + ret := make([]*Rule, 0) + for resource, rules := range ruleMap { + if util.RegexMatch(resource, res) { + for _, rule := range rules { + if rule != nil { + ret = append(ret, rule) + } + } + } + } + return ret +} + +func rulesFrom(m map[string][]*Rule) []*Rule { + rules := make([]*Rule, 0, 8) + if len(m) == 0 { + return rules + } + for _, rs := range m { + for _, r := range rs { + if r != nil { + rules = append(rules, r) + } + } + } + return rules +} + +func logRuleUpdate(m map[string][]*Rule) { + rs := rulesFrom(m) + if len(rs) == 0 { + logging.Info("[LLMTokenRateLimit] rules were cleared") + } else { + var builder strings.Builder + for i, r := range rs { + builder.WriteString(r.String()) + if i != len(rs)-1 { + builder.WriteString(", ") + } + } + logging.Info("[LLMTokenRateLimit] rules were loaded", + "rules", builder.String(), + ) + } +} diff --git a/core/llm_token_ratelimit/rule_manager_test.go b/core/llm_token_ratelimit/rule_manager_test.go new file mode 100644 index 000000000..90ea72197 --- /dev/null +++ b/core/llm_token_ratelimit/rule_manager_test.go @@ -0,0 +1,1061 @@ +// Copyright 1999-2020 Alibaba Group Holding Ltd. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package llm_token_ratelimit + +import ( + "reflect" + "strings" + "sync" + "testing" +) + +// Test helper functions to create test data +func createTestRuleForRuleManager(resource string) *Rule { + return &Rule{ + ID: "test-rule-1", + Resource: resource, + Strategy: PETA, + Encoding: TokenEncoding{ + Provider: OpenAIEncoderProvider, + Model: "gpt-3.5-turbo", + }, + SpecificItems: []*SpecificItem{ + { + Identifier: Identifier{ + Type: AllIdentifier, + Value: ".*", + }, + KeyItems: []*KeyItem{ + { + Key: "default", + Token: Token{ + Number: 1000, + CountStrategy: TotalTokens, + }, + Time: Time{ + Unit: Minute, + Value: 1, + }, + }, + }, + }, + }, + } +} + +func createTestRuleWithID(resource, id string) *Rule { + rule := createTestRuleForRuleManager(resource) + rule.ID = id + return rule +} + +func resetGlobalState() { + ruleMap = make(map[string][]*Rule) + currentRules = make(map[string][]*Rule, 0) + updateRuleMux = new(sync.Mutex) + rwMux = &sync.RWMutex{} +} + +// Test LoadRules function +func TestLoadRules(t *testing.T) { + defer resetGlobalState() + + tests := []struct { + name string + rules []*Rule + expectUpdated bool + expectError bool + expectedRules int + setupExisting bool + existingRules []*Rule + }{ + { + name: "Load valid rules", + rules: []*Rule{ + createTestRuleForRuleManager("test-resource-1"), + createTestRuleForRuleManager("test-resource-2"), + }, + expectUpdated: true, + expectError: false, + expectedRules: 2, + }, + { + name: "Load empty rules", + rules: []*Rule{}, + setupExisting: true, + existingRules: []*Rule{ + createTestRuleForRuleManager("test-resource-1"), + }, + expectUpdated: true, + expectError: false, + expectedRules: 0, + }, + { + name: "Load nil rules", + rules: nil, + setupExisting: true, + existingRules: []*Rule{ + createTestRuleForRuleManager("test-resource-1"), + }, + expectUpdated: true, + expectError: false, + expectedRules: 0, + }, + { + name: "Load same rules twice (should not update)", + rules: []*Rule{ + createTestRuleForRuleManager("test-resource-1"), + }, + setupExisting: true, + existingRules: []*Rule{ + createTestRuleForRuleManager("test-resource-1"), + }, + expectUpdated: false, + expectError: false, + expectedRules: 1, + }, + { + name: "Load invalid rules (should be filtered out)", + rules: []*Rule{ + createTestRuleForRuleManager("test-resource-1"), + nil, // Invalid rule + { + Resource: "", // Invalid resource + Strategy: PETA, + }, + }, + expectUpdated: true, + expectError: false, + expectedRules: 1, + }, + { + name: "Load rules with duplicate resources (last one wins)", + rules: []*Rule{ + createTestRuleWithID("same-resource", "rule-1"), + createTestRuleWithID("same-resource", "rule-2"), + }, + expectUpdated: true, + expectError: false, + expectedRules: 1, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + resetGlobalState() + + // Setup existing rules if needed + if tt.setupExisting { + _, err := LoadRules(tt.existingRules) + if err != nil { + t.Fatalf("Failed to setup existing rules: %v", err) + } + } + + // Load test rules + updated, err := LoadRules(tt.rules) + + // Check error expectation + if tt.expectError { + if err == nil { + t.Error("Expected error, got nil") + } + } else { + if err != nil { + t.Errorf("Expected no error, got %v", err) + } + } + + // Check update expectation + if updated != tt.expectUpdated { + t.Errorf("Expected updated=%v, got %v", tt.expectUpdated, updated) + } + + // Check rules count + allRules := GetRules() + if len(allRules) != tt.expectedRules { + t.Errorf("Expected %d rules, got %d", tt.expectedRules, len(allRules)) + } + }) + } +} + +// Test onRuleUpdate function +func TestOnRuleUpdate(t *testing.T) { + defer resetGlobalState() + + tests := []struct { + name string + rawRulesMap map[string][]*Rule + expectError bool + expectedSize int + }{ + { + name: "Valid rules map", + rawRulesMap: map[string][]*Rule{ + "resource-1": {createTestRuleForRuleManager("resource-1")}, + "resource-2": {createTestRuleForRuleManager("resource-2")}, + }, + expectError: false, + expectedSize: 2, + }, + { + name: "Empty rules map", + rawRulesMap: map[string][]*Rule{}, + expectError: false, + expectedSize: 0, + }, + { + name: "Rules map with empty rule slices (should be filtered out)", + rawRulesMap: map[string][]*Rule{ + "resource-1": {createTestRuleForRuleManager("resource-1")}, + "resource-2": {}, // Empty slice + }, + expectError: false, + expectedSize: 1, + }, + { + name: "Nil rules map", + rawRulesMap: nil, + expectError: false, + expectedSize: 0, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + resetGlobalState() + + err := onRuleUpdate(tt.rawRulesMap) + + // Check error expectation + if tt.expectError { + if err == nil { + t.Error("Expected error, got nil") + } + } else { + if err != nil { + t.Errorf("Expected no error, got %v", err) + } + } + + // Check rules map size + rwMux.RLock() + actualSize := len(ruleMap) + rwMux.RUnlock() + + if actualSize != tt.expectedSize { + t.Errorf("Expected ruleMap size %d, got %d", tt.expectedSize, actualSize) + } + + // Verify currentRules is updated + if !reflect.DeepEqual(currentRules, tt.rawRulesMap) { + t.Error("currentRules was not updated correctly") + } + }) + } +} + +// Test LoadRulesOfResource function +func TestLoadRulesOfResource(t *testing.T) { + defer resetGlobalState() + + tests := []struct { + name string + resource string + rules []*Rule + expectUpdated bool + expectError bool + setupExisting bool + existingRules []*Rule + }{ + { + name: "Load valid rules for resource", + resource: "test-resource", + rules: []*Rule{ + createTestRuleForRuleManager("test-resource"), + }, + expectUpdated: true, + expectError: false, + }, + { + name: "Empty resource name", + resource: "", + rules: []*Rule{createTestRuleForRuleManager("test-resource")}, + expectUpdated: false, + expectError: true, + }, + { + name: "Clear rules for resource", + resource: "test-resource", + rules: []*Rule{}, + expectUpdated: true, + expectError: false, + }, + { + name: "Clear rules for resource with nil", + resource: "test-resource", + rules: nil, + expectUpdated: true, + expectError: false, + }, + { + name: "Load same rules for resource (should not update)", + resource: "test-resource", + rules: []*Rule{ + createTestRuleForRuleManager("test-resource"), + }, + setupExisting: true, + existingRules: []*Rule{ + createTestRuleForRuleManager("test-resource"), + }, + expectUpdated: false, + expectError: false, + }, + { + name: "Load different rules for resource (should update)", + resource: "test-resource", + rules: []*Rule{ + createTestRuleWithID("test-resource", "new-rule"), + }, + setupExisting: true, + existingRules: []*Rule{ + createTestRuleWithID("test-resource", "old-rule"), + }, + expectUpdated: true, + expectError: false, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + resetGlobalState() + + // Setup existing rules if needed + if tt.setupExisting { + _, err := LoadRulesOfResource(tt.resource, tt.existingRules) + if err != nil { + t.Fatalf("Failed to setup existing rules: %v", err) + } + } + + // Load test rules + updated, err := LoadRulesOfResource(tt.resource, tt.rules) + + // Check error expectation + if tt.expectError { + if err == nil { + t.Error("Expected error, got nil") + } + } else { + if err != nil { + t.Errorf("Expected no error, got %v", err) + } + } + + // Check update expectation + if updated != tt.expectUpdated { + t.Errorf("Expected updated=%v, got %v", tt.expectUpdated, updated) + } + + // If no error and rules were provided, verify they were loaded + if !tt.expectError && len(tt.rules) > 0 { + resourceRules := GetRulesOfResource(tt.resource) + if len(resourceRules) == 0 { + t.Error("Expected rules to be loaded for resource, got none") + } + } + }) + } +} + +// Test onResourceRuleUpdate function +func TestOnResourceRuleUpdate(t *testing.T) { + defer resetGlobalState() + + tests := []struct { + name string + resource string + rawResRules []*Rule + expectError bool + expectExists bool + }{ + { + name: "Update resource with valid rules", + resource: "test-resource", + rawResRules: []*Rule{ + createTestRuleForRuleManager("test-resource"), + }, + expectError: false, + expectExists: true, + }, + { + name: "Update resource with empty rules (should delete)", + resource: "test-resource", + rawResRules: []*Rule{}, + expectError: false, + expectExists: false, + }, + { + name: "Update resource with nil rules (should delete)", + resource: "test-resource", + rawResRules: nil, + expectError: false, + expectExists: false, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + resetGlobalState() + + err := onResourceRuleUpdate(tt.resource, tt.rawResRules) + + // Check error expectation + if tt.expectError { + if err == nil { + t.Error("Expected error, got nil") + } + } else { + if err != nil { + t.Errorf("Expected no error, got %v", err) + } + } + + // Check if resource exists in ruleMap + rwMux.RLock() + _, exists := ruleMap[tt.resource] + rwMux.RUnlock() + + if exists != tt.expectExists { + t.Errorf("Expected resource exists=%v, got %v", tt.expectExists, exists) + } + + // Check currentRules + currentResourceRules, exists := currentRules[tt.resource] + if tt.expectExists { + if !exists { + t.Error("Expected resource to exist in currentRules") + } + if !reflect.DeepEqual(currentResourceRules, tt.rawResRules) { + t.Error("currentRules was not updated correctly") + } + } + }) + } +} + +// Test ClearRules function +func TestClearRules(t *testing.T) { + defer resetGlobalState() + + // Setup some rules first + rules := []*Rule{ + createTestRuleForRuleManager("resource-1"), + createTestRuleForRuleManager("resource-2"), + } + _, err := LoadRules(rules) + if err != nil { + t.Fatalf("Failed to setup test rules: %v", err) + } + + // Verify rules exist + if len(GetRules()) == 0 { + t.Fatal("Expected rules to be loaded, got none") + } + + // Clear rules + err = ClearRules() + if err != nil { + t.Errorf("Expected no error clearing rules, got %v", err) + } + + // Verify rules are cleared + if len(GetRules()) != 0 { + t.Error("Expected all rules to be cleared, but some remain") + } +} + +// Test ClearRulesOfResource function +func TestClearRulesOfResource(t *testing.T) { + defer resetGlobalState() + + tests := []struct { + name string + resource string + expectError bool + setupRules []*Rule + otherResource string + }{ + { + name: "Clear existing resource", + resource: "test-resource", + setupRules: []*Rule{ + createTestRuleForRuleManager("test-resource"), + createTestRuleForRuleManager("other-resource"), + }, + otherResource: "other-resource", + expectError: false, + }, + { + name: "Clear non-existing resource", + resource: "non-existing", + setupRules: []*Rule{createTestRuleForRuleManager("test-resource")}, + expectError: false, + }, + { + name: "Clear with empty resource name", + resource: "", + expectError: true, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + resetGlobalState() + + // Setup rules if provided + if len(tt.setupRules) > 0 { + _, err := LoadRules(tt.setupRules) + if err != nil { + t.Fatalf("Failed to setup test rules: %v", err) + } + } + + // Clear specific resource + err := ClearRulesOfResource(tt.resource) + + // Check error expectation + if tt.expectError { + if err == nil { + t.Error("Expected error, got nil") + } + return + } else { + if err != nil { + t.Errorf("Expected no error, got %v", err) + } + } + + // Verify target resource is cleared + resourceRules := GetRulesOfResource(tt.resource) + if len(resourceRules) != 0 { + t.Errorf("Expected resource rules to be cleared, got %d", len(resourceRules)) + } + + // Verify other resources remain if they exist + if tt.otherResource != "" { + otherRules := GetRulesOfResource(tt.otherResource) + if len(otherRules) == 0 { + t.Error("Expected other resource rules to remain, but they were cleared") + } + } + }) + } +} + +// Test GetRules function +func TestGetRules(t *testing.T) { + defer resetGlobalState() + + tests := []struct { + name string + setupRules []*Rule + expectedCount int + }{ + { + name: "Get rules when none exist", + setupRules: nil, + expectedCount: 0, + }, + { + name: "Get rules when some exist", + setupRules: []*Rule{ + createTestRuleForRuleManager("resource-1"), + createTestRuleForRuleManager("resource-2"), + }, + expectedCount: 2, + }, + { + name: "Get rules with duplicate resources (should return unique)", + setupRules: []*Rule{ + createTestRuleWithID("same-resource", "rule-1"), + createTestRuleWithID("same-resource", "rule-2"), + }, + expectedCount: 1, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + resetGlobalState() + + // Setup rules if provided + if len(tt.setupRules) > 0 { + _, err := LoadRules(tt.setupRules) + if err != nil { + t.Fatalf("Failed to setup test rules: %v", err) + } + } + + // Get rules + rules := GetRules() + + // Check count + if len(rules) != tt.expectedCount { + t.Errorf("Expected %d rules, got %d", tt.expectedCount, len(rules)) + } + + // Verify returned rules are copies (not pointers) + for _, rule := range rules { + // This should not be a pointer dereference but a struct + if reflect.TypeOf(rule).Kind() == reflect.Ptr { + t.Error("Expected rule structs, got pointers") + } + } + }) + } +} + +// Test GetRulesOfResource function +func TestGetRulesOfResource(t *testing.T) { + defer resetGlobalState() + + tests := []struct { + name string + setupRules []*Rule + queryResource string + expectedCount int + }{ + { + name: "Get rules for non-existing resource", + setupRules: []*Rule{createTestRuleForRuleManager("other-resource")}, + queryResource: "test-resource", + expectedCount: 0, + }, + { + name: "Get rules for existing resource", + setupRules: []*Rule{createTestRuleForRuleManager("test-resource")}, + queryResource: "test-resource", + expectedCount: 1, + }, + { + name: "Get rules with pattern matching", + setupRules: []*Rule{ + createTestRuleForRuleManager("test-.*"), + createTestRuleForRuleManager("other-resource"), + }, + queryResource: "test-service", + expectedCount: 1, + }, + { + name: "Get rules for multiple matching patterns", + setupRules: []*Rule{ + createTestRuleForRuleManager("test-.*"), + createTestRuleForRuleManager(".*-service"), + }, + queryResource: "test-service", + expectedCount: 2, + }, + { + name: "Get rules when no rules exist", + setupRules: nil, + queryResource: "any-resource", + expectedCount: 0, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + resetGlobalState() + + // Setup rules if provided + if len(tt.setupRules) > 0 { + _, err := LoadRules(tt.setupRules) + if err != nil { + t.Fatalf("Failed to setup test rules: %v", err) + } + } + + // Get rules for resource + rules := GetRulesOfResource(tt.queryResource) + + // Check count + if len(rules) != tt.expectedCount { + t.Errorf("Expected %d rules, got %d", tt.expectedCount, len(rules)) + } + + // Verify returned rules are copies (not pointers) + for _, rule := range rules { + if reflect.TypeOf(rule).Kind() == reflect.Ptr { + t.Error("Expected rule structs, got pointers") + } + } + }) + } +} + +// Test getRules function +func TestGetRulesInternal(t *testing.T) { + defer resetGlobalState() + + tests := []struct { + name string + setupRules map[string][]*Rule + expectedCount int + }{ + { + name: "Get rules when ruleMap is empty", + setupRules: map[string][]*Rule{}, + expectedCount: 0, + }, + { + name: "Get rules when ruleMap has entries", + setupRules: map[string][]*Rule{ + "resource-1": {createTestRuleForRuleManager("resource-1")}, + "resource-2": {createTestRuleForRuleManager("resource-2")}, + }, + expectedCount: 2, + }, + { + name: "Get rules with nil entries (should be filtered)", + setupRules: map[string][]*Rule{ + "resource-1": {createTestRuleForRuleManager("resource-1"), nil}, + "resource-2": {nil, createTestRuleForRuleManager("resource-2")}, + }, + expectedCount: 2, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + resetGlobalState() + + // Setup ruleMap directly + rwMux.Lock() + ruleMap = tt.setupRules + rwMux.Unlock() + + // Get rules + rules := getRules() + + // Check count + if len(rules) != tt.expectedCount { + t.Errorf("Expected %d rules, got %d", tt.expectedCount, len(rules)) + } + + // Verify all returned rules are non-nil + for _, rule := range rules { + if rule == nil { + t.Error("Found nil rule in results") + } + } + }) + } +} + +// Test getRulesOfResource function +func TestGetRulesOfResourceInternal(t *testing.T) { + defer resetGlobalState() + + tests := []struct { + name string + setupRules map[string][]*Rule + queryResource string + expectedCount int + }{ + { + name: "Get rules from empty ruleMap", + setupRules: map[string][]*Rule{}, + queryResource: "test-resource", + expectedCount: 0, + }, + { + name: "Get rules with exact match", + setupRules: map[string][]*Rule{ + "test-resource": {createTestRuleForRuleManager("test-resource")}, + "other": {createTestRuleForRuleManager("other")}, + }, + queryResource: "test-resource", + expectedCount: 1, + }, + { + name: "Get rules with regex pattern match", + setupRules: map[string][]*Rule{ + "test-.*": {createTestRuleForRuleManager("test-.*")}, + "other": {createTestRuleForRuleManager("other")}, + }, + queryResource: "test-service", + expectedCount: 1, + }, + { + name: "Get rules with nil entries (should be filtered)", + setupRules: map[string][]*Rule{ + "test-resource": {createTestRuleForRuleManager("test-resource"), nil}, + }, + queryResource: "test-resource", + expectedCount: 1, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + resetGlobalState() + + // Setup ruleMap directly + rwMux.Lock() + ruleMap = tt.setupRules + rwMux.Unlock() + + // Get rules for resource + rules := getRulesOfResource(tt.queryResource) + + // Check count + if len(rules) != tt.expectedCount { + t.Errorf("Expected %d rules, got %d", tt.expectedCount, len(rules)) + } + + // Verify all returned rules are non-nil + for _, rule := range rules { + if rule == nil { + t.Error("Found nil rule in results") + } + } + }) + } +} + +// Test rulesFrom function +func TestRulesFrom(t *testing.T) { + tests := []struct { + name string + rulesMap map[string][]*Rule + expectedCount int + }{ + { + name: "Empty map", + rulesMap: map[string][]*Rule{}, + expectedCount: 0, + }, + { + name: "Nil map", + rulesMap: nil, + expectedCount: 0, + }, + { + name: "Map with valid rules", + rulesMap: map[string][]*Rule{ + "resource-1": {createTestRuleForRuleManager("resource-1")}, + "resource-2": {createTestRuleForRuleManager("resource-2")}, + }, + expectedCount: 2, + }, + { + name: "Map with nil rules (should be filtered)", + rulesMap: map[string][]*Rule{ + "resource-1": {createTestRuleForRuleManager("resource-1"), nil}, + "resource-2": {nil, createTestRuleForRuleManager("resource-2")}, + }, + expectedCount: 2, + }, + { + name: "Map with empty rule slices", + rulesMap: map[string][]*Rule{ + "resource-1": {}, + "resource-2": {createTestRuleForRuleManager("resource-2")}, + }, + expectedCount: 1, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + rules := rulesFrom(tt.rulesMap) + + // Check count + if len(rules) != tt.expectedCount { + t.Errorf("Expected %d rules, got %d", tt.expectedCount, len(rules)) + } + + // Verify all returned rules are non-nil + for _, rule := range rules { + if rule == nil { + t.Error("Found nil rule in results") + } + } + }) + } +} + +// Test logRuleUpdate function +func TestLogRuleUpdate(t *testing.T) { + tests := []struct { + name string + rulesMap map[string][]*Rule + verify func() bool // Function to verify log output + }{ + { + name: "Log with empty rules (should log cleared message)", + rulesMap: map[string][]*Rule{}, + verify: func() bool { + // In a real implementation, you might want to capture log output + // For now, we just verify the function doesn't panic + return true + }, + }, + { + name: "Log with rules (should log loaded message)", + rulesMap: map[string][]*Rule{ + "resource-1": {createTestRuleForRuleManager("resource-1")}, + }, + verify: func() bool { + return true + }, + }, + { + name: "Log with nil map", + rulesMap: nil, + verify: func() bool { + return true + }, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + // This function primarily logs, so we test it doesn't panic + defer func() { + if r := recover(); r != nil { + t.Errorf("logRuleUpdate panicked: %v", r) + } + }() + + logRuleUpdate(tt.rulesMap) + + if !tt.verify() { + t.Error("Log verification failed") + } + }) + } +} + +// Test concurrent access to rule manager +func TestRuleManager_ConcurrentAccess(t *testing.T) { + defer resetGlobalState() + + const numGoroutines = 10 + const operationsPerGoroutine = 100 + + var wg sync.WaitGroup + wg.Add(numGoroutines * 3) // 3 types of operations + + // Concurrent rule loading + for i := 0; i < numGoroutines; i++ { + go func(id int) { + defer wg.Done() + for j := 0; j < operationsPerGoroutine; j++ { + rule := createTestRuleWithID("concurrent-resource", + strings.Join([]string{"rule", string(rune(id)), string(rune(j))}, "-")) + LoadRules([]*Rule{rule}) + } + }(i) + } + + // Concurrent rule reading + for i := 0; i < numGoroutines; i++ { + go func() { + defer wg.Done() + for j := 0; j < operationsPerGoroutine; j++ { + GetRules() + GetRulesOfResource("concurrent-resource") + } + }() + } + + // Concurrent resource-specific operations + for i := 0; i < numGoroutines; i++ { + go func(id int) { + defer wg.Done() + for j := 0; j < operationsPerGoroutine; j++ { + resource := strings.Join([]string{"resource", string(rune(id))}, "-") + rule := createTestRuleForRuleManager(resource) + LoadRulesOfResource(resource, []*Rule{rule}) + } + }(i) + } + + wg.Wait() +} + +// Benchmark tests +func BenchmarkLoadRules(b *testing.B) { + defer resetGlobalState() + + rules := []*Rule{ + createTestRuleForRuleManager("benchmark-resource-1"), + createTestRuleForRuleManager("benchmark-resource-2"), + } + + b.ResetTimer() + for i := 0; i < b.N; i++ { + LoadRules(rules) + } +} + +func BenchmarkGetRules(b *testing.B) { + defer resetGlobalState() + + // Setup some rules + rules := []*Rule{ + createTestRuleForRuleManager("benchmark-resource-1"), + createTestRuleForRuleManager("benchmark-resource-2"), + createTestRuleForRuleManager("benchmark-resource-3"), + } + LoadRules(rules) + + b.ResetTimer() + for i := 0; i < b.N; i++ { + GetRules() + } +} + +func BenchmarkGetRulesOfResource(b *testing.B) { + defer resetGlobalState() + + // Setup some rules + rules := []*Rule{ + createTestRuleForRuleManager("benchmark-.*"), + createTestRuleForRuleManager("other-resource"), + } + LoadRules(rules) + + b.ResetTimer() + for i := 0; i < b.N; i++ { + GetRulesOfResource("benchmark-service") + } +} + +func BenchmarkRulesFrom(b *testing.B) { + rulesMap := map[string][]*Rule{ + "resource-1": {createTestRuleForRuleManager("resource-1")}, + "resource-2": {createTestRuleForRuleManager("resource-2")}, + "resource-3": {createTestRuleForRuleManager("resource-3")}, + } + + b.ResetTimer() + for i := 0; i < b.N; i++ { + rulesFrom(rulesMap) + } +} diff --git a/core/llm_token_ratelimit/rule_matcher.go b/core/llm_token_ratelimit/rule_matcher.go new file mode 100644 index 000000000..2f0ce80ee --- /dev/null +++ b/core/llm_token_ratelimit/rule_matcher.go @@ -0,0 +1,204 @@ +// Copyright 1999-2020 Alibaba Group Holding Ltd. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package llm_token_ratelimit + +import ( + "errors" + + "github.com/alibaba/sentinel-golang/logging" +) + +type MatchedRule struct { + Strategy Strategy + LimitKey string + TimeWindow int64 // seconds + TokenSize int64 + CountStrategy CountStrategy + // PETA + Encoding TokenEncoding + EstimatedToken int64 +} + +type MatchedRuleCollector interface { + Collect(ctx *Context, rule *Rule) []*MatchedRule +} + +type RateLimitChecker interface { + Check(ctx *Context, rules []*MatchedRule) bool +} + +type IdentifierChecker interface { + Check(ctx *Context, infos *RequestInfos, identifier Identifier, pattern string) bool +} + +type TokenUpdater interface { + Update(ctx *Context, rule *MatchedRule) +} + +var globalRuleMatcher = NewDefaultRuleMatcher() + +type RuleMatcher struct { + MatchedRuleCollectors map[Strategy]MatchedRuleCollector + RateLimitCheckers map[Strategy]RateLimitChecker + IdentifierCheckers map[IdentifierType]IdentifierChecker + TokenUpdaters map[Strategy]TokenUpdater +} + +func NewDefaultRuleMatcher() *RuleMatcher { + return &RuleMatcher{ + MatchedRuleCollectors: map[Strategy]MatchedRuleCollector{ + FixedWindow: &BaseRuleCollector{}, + PETA: &BaseRuleCollector{}, + }, + RateLimitCheckers: map[Strategy]RateLimitChecker{ + FixedWindow: &FixedWindowChecker{}, + PETA: &PETAChecker{}, + }, + TokenUpdaters: map[Strategy]TokenUpdater{ + FixedWindow: &FixedWindowUpdater{}, + PETA: &PETAUpdater{}, + }, + IdentifierCheckers: map[IdentifierType]IdentifierChecker{ + AllIdentifier: &AllIdentifierChecker{}, + Header: &HeaderChecker{}, + }, + } +} + +func (m *RuleMatcher) getMatchedRuleCollector(strategy Strategy) MatchedRuleCollector { + if m == nil { + return nil + } + collector, exists := m.MatchedRuleCollectors[strategy] + if !exists { + return nil + } + return collector +} + +func (m *RuleMatcher) getRateLimitChecker(strategy Strategy) RateLimitChecker { + if m == nil { + return nil + } + checker, exists := m.RateLimitCheckers[strategy] + if !exists { + return nil + } + return checker +} + +func (m *RuleMatcher) getTokenUpdater(strategy Strategy) TokenUpdater { + if m == nil { + return nil + } + updater, exists := m.TokenUpdaters[strategy] + if !exists { + return nil + } + return updater +} + +func (m *RuleMatcher) getIdentifierChecker(identifier IdentifierType) IdentifierChecker { + if m == nil { + return nil + } + checker, exists := m.IdentifierCheckers[identifier] + if !exists { + return nil + } + return checker +} + +func (m *RuleMatcher) checkPass(ctx *Context, rule *Rule) bool { + if m == nil { + return true + } + collector := m.getMatchedRuleCollector(rule.Strategy) + if collector == nil { + logging.Error(errors.New("unknown strategy"), + "unknown strategy in llm_token_ratelimit.RuleMatcher.checkPass() when get collector", + "strategy", rule.Strategy.String(), + "requestID", ctx.Get(KeyRequestID), + ) + return true + } + + rules := collector.Collect(ctx, rule) + if len(rules) == 0 { + return true + } + + checker := m.getRateLimitChecker(rule.Strategy) + if checker == nil { + logging.Error(errors.New("unknown strategy"), + "unknown strategy in llm_token_ratelimit.RuleMatcher.checkPass() when get checker", + "strategy", rule.Strategy.String(), + "requestID", ctx.Get(KeyRequestID), + ) + return true + } + + if passed := checker.Check(ctx, rules); !passed { + return false + } + + m.cacheMatchedRules(ctx, rules) + return true +} + +func (m *RuleMatcher) cacheMatchedRules(ctx *Context, newRules []*MatchedRule) { + if m == nil { + return + } + + if len(newRules) == 0 { + return + } + + existingValue := ctx.Get(KeyMatchedRules) + if existingValue == nil { + ctx.Set(KeyMatchedRules, newRules) + } else { + existingRules, ok := existingValue.([]*MatchedRule) + if !ok || existingRules == nil { + ctx.Set(KeyMatchedRules, newRules) + } else { + allRules := make([]*MatchedRule, 0, len(existingRules)+len(newRules)) + allRules = append(allRules, existingRules...) + allRules = append(allRules, newRules...) + ctx.Set(KeyMatchedRules, allRules) + } + } +} + +func (m *RuleMatcher) update(ctx *Context, rule *MatchedRule) { + if m == nil { + return + } + + if rule == nil { + return + } + updater := m.getTokenUpdater(rule.Strategy) + if updater == nil { + logging.Error(errors.New("unknown strategy"), + "unknown strategy in llm_token_ratelimit.RuleMatcher.update() when get updater", + "strategy", rule.Strategy.String(), + "requestID", ctx.Get(KeyRequestID), + ) + return + } + updater.Update(ctx, rule) +} diff --git a/core/llm_token_ratelimit/rule_test.go b/core/llm_token_ratelimit/rule_test.go new file mode 100644 index 000000000..277cf37f4 --- /dev/null +++ b/core/llm_token_ratelimit/rule_test.go @@ -0,0 +1,974 @@ +// Copyright 1999-2020 Alibaba Group Holding Ltd. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package llm_token_ratelimit + +import ( + "strings" + "testing" +) + +// Test helper functions to create test data +func createTestRuleForRule() *Rule { + return &Rule{ + ID: "test-rule-1", + Resource: "test-resource", + Strategy: PETA, + Encoding: TokenEncoding{ + Provider: OpenAIEncoderProvider, + Model: "gpt-3.5-turbo", + }, + SpecificItems: []*SpecificItem{ + { + Identifier: Identifier{ + Type: AllIdentifier, + Value: ".*", + }, + KeyItems: []*KeyItem{ + { + Key: "default", + Token: Token{ + Number: 1000, + CountStrategy: TotalTokens, + }, + Time: Time{ + Unit: Minute, + Value: 1, + }, + }, + }, + }, + }, + } +} + +func createEmptyRule() *Rule { + return &Rule{} +} + +func createRuleWithMultipleItems() *Rule { + return &Rule{ + ID: "multi-item-rule", + Resource: "multi-resource", + Strategy: FixedWindow, + Encoding: TokenEncoding{ + Provider: OpenAIEncoderProvider, + Model: "gpt-4", + }, + SpecificItems: []*SpecificItem{ + { + Identifier: Identifier{ + Type: Header, + Value: "user-id", + }, + KeyItems: []*KeyItem{ + { + Key: "premium", + Token: Token{ + Number: 5000, + CountStrategy: TotalTokens, + }, + Time: Time{ + Unit: Hour, + Value: 1, + }, + }, + { + Key: "basic", + Token: Token{ + Number: 1000, + CountStrategy: InputTokens, + }, + Time: Time{ + Unit: Minute, + Value: 30, + }, + }, + }, + }, + { + Identifier: Identifier{ + Type: AllIdentifier, + Value: "fallback", + }, + KeyItems: []*KeyItem{ + { + Key: "default", + Token: Token{ + Number: 500, + CountStrategy: OutputTokens, + }, + Time: Time{ + Unit: Day, + Value: 1, + }, + }, + }, + }, + }, + } +} + +// Test Rule.ResourceName function +func TestRule_ResourceName(t *testing.T) { + tests := []struct { + name string + rule *Rule + expected string + }{ + { + name: "Valid rule with resource name", + rule: &Rule{Resource: "test-service"}, + expected: "test-service", + }, + { + name: "Rule with empty resource name", + rule: &Rule{Resource: ""}, + expected: "", + }, + { + name: "Nil rule", + rule: nil, + expected: "Rule{nil}", + }, + { + name: "Rule with special characters in resource", + rule: &Rule{Resource: "test-service-*-regex"}, + expected: "test-service-*-regex", + }, + { + name: "Rule with spaces in resource", + rule: &Rule{Resource: "test service with spaces"}, + expected: "test service with spaces", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + result := tt.rule.ResourceName() + if result != tt.expected { + t.Errorf("Expected %q, got %q", tt.expected, result) + } + }) + } +} + +// Test Rule.String function +func TestRule_String(t *testing.T) { + tests := []struct { + name string + rule *Rule + contains []string // Expected substrings in the output + notContains []string // Substrings that should not be in the output + }{ + { + name: "Complete rule with all fields", + rule: createTestRuleForRule(), + contains: []string{ + "Rule{", + "ID:test-rule-1", + "Resource:test-resource", + "Strategy:peta", + "Encoding:TokenEncoding{Provider:openai, Model:gpt-3.5-turbo}", + "SpecificItems:[", + "}", + }, + }, + { + name: "Rule without ID", + rule: &Rule{ + Resource: "test-resource", + Strategy: FixedWindow, + Encoding: TokenEncoding{ + Provider: OpenAIEncoderProvider, + Model: "gpt-4", + }, + SpecificItems: []*SpecificItem{}, + }, + contains: []string{ + "Rule{", + "Resource:test-resource", + "Strategy:fixed-window", + "SpecificItems:[]", + }, + notContains: []string{ + "ID:", + }, + }, + { + name: "Rule with empty specific items", + rule: &Rule{ + ID: "empty-items", + Resource: "empty-resource", + Strategy: PETA, + Encoding: TokenEncoding{Provider: OpenAIEncoderProvider, Model: "gpt-3.5-turbo"}, + SpecificItems: []*SpecificItem{}, + }, + contains: []string{ + "SpecificItems:[]", + }, + }, + { + name: "Nil rule", + rule: nil, + contains: []string{"Rule{nil}"}, + }, + { + name: "Rule with multiple specific items", + rule: createRuleWithMultipleItems(), + contains: []string{ + "Rule{", + "ID:multi-item-rule", + "Resource:multi-resource", + "Strategy:fixed-window", + "SpecificItems:[", + ", ", // Should contain comma separator between items + }, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + result := tt.rule.String() + + // Check that expected substrings are present + for _, substr := range tt.contains { + if !strings.Contains(result, substr) { + t.Errorf("Expected string to contain %q, got: %s", substr, result) + } + } + + // Check that unwanted substrings are not present + for _, substr := range tt.notContains { + if strings.Contains(result, substr) { + t.Errorf("Expected string to NOT contain %q, got: %s", substr, result) + } + } + }) + } +} + +// Test Rule.setDefaultRuleOption function +func TestRule_SetDefaultRuleOption(t *testing.T) { + tests := []struct { + name string + rule *Rule + expectedRule *Rule + description string + }{ + { + name: "Empty rule should get default values", + rule: &Rule{ + Encoding: TokenEncoding{Provider: OpenAIEncoderProvider}, + }, + expectedRule: &Rule{ + Resource: DefaultResourcePattern, + Encoding: TokenEncoding{ + Provider: OpenAIEncoderProvider, + Model: DefaultTokenEncodingModel[OpenAIEncoderProvider], + }, + }, + description: "Should set default resource pattern and model", + }, + { + name: "Rule with empty resource should get default", + rule: &Rule{ + Resource: "", + Encoding: TokenEncoding{ + Provider: OpenAIEncoderProvider, + Model: "custom-model", + }, + }, + expectedRule: &Rule{ + Resource: DefaultResourcePattern, + Encoding: TokenEncoding{ + Provider: OpenAIEncoderProvider, + Model: "custom-model", + }, + }, + description: "Should set default resource but keep custom model", + }, + { + name: "Rule with empty model should get default", + rule: &Rule{ + Resource: "custom-resource", + Encoding: TokenEncoding{ + Provider: OpenAIEncoderProvider, + Model: "", + }, + }, + expectedRule: &Rule{ + Resource: "custom-resource", + Encoding: TokenEncoding{ + Provider: OpenAIEncoderProvider, + Model: DefaultTokenEncodingModel[OpenAIEncoderProvider], + }, + }, + description: "Should keep custom resource but set default model", + }, + { + name: "Rule with specific items having empty values", + rule: &Rule{ + Resource: "test-resource", + Encoding: TokenEncoding{ + Provider: OpenAIEncoderProvider, + Model: "gpt-4", + }, + SpecificItems: []*SpecificItem{ + { + Identifier: Identifier{ + Type: AllIdentifier, + Value: "", + }, + KeyItems: []*KeyItem{ + { + Key: "", + Token: Token{Number: 1000, CountStrategy: TotalTokens}, + Time: Time{Unit: Minute, Value: 1}, + }, + }, + }, + }, + }, + expectedRule: &Rule{ + Resource: "test-resource", + Encoding: TokenEncoding{ + Provider: OpenAIEncoderProvider, + Model: "gpt-4", + }, + SpecificItems: []*SpecificItem{ + { + Identifier: Identifier{ + Type: AllIdentifier, + Value: DefaultIdentifierValuePattern, + }, + KeyItems: []*KeyItem{ + { + Key: DefaultKeyPattern, + Token: Token{Number: 1000, CountStrategy: TotalTokens}, + Time: Time{Unit: Minute, Value: 1}, + }, + }, + }, + }, + }, + description: "Should set default identifier value and key pattern", + }, + { + name: "Rule with multiple specific items and key items", + rule: &Rule{ + Resource: "", + Encoding: TokenEncoding{Provider: OpenAIEncoderProvider}, + SpecificItems: []*SpecificItem{ + { + Identifier: Identifier{Type: AllIdentifier, Value: ""}, + KeyItems: []*KeyItem{ + {Key: "", Token: Token{Number: 100}, Time: Time{Unit: Second, Value: 1}}, + {Key: "valid-key", Token: Token{Number: 200}, Time: Time{Unit: Minute, Value: 1}}, + }, + }, + { + Identifier: Identifier{Type: Header, Value: "user-id"}, + KeyItems: []*KeyItem{ + {Key: "", Token: Token{Number: 300}, Time: Time{Unit: Hour, Value: 1}}, + }, + }, + }, + }, + expectedRule: &Rule{ + Resource: DefaultResourcePattern, + Encoding: TokenEncoding{ + Provider: OpenAIEncoderProvider, + Model: DefaultTokenEncodingModel[OpenAIEncoderProvider], + }, + SpecificItems: []*SpecificItem{ + { + Identifier: Identifier{Type: AllIdentifier, Value: DefaultIdentifierValuePattern}, + KeyItems: []*KeyItem{ + {Key: DefaultKeyPattern, Token: Token{Number: 100}, Time: Time{Unit: Second, Value: 1}}, + {Key: "valid-key", Token: Token{Number: 200}, Time: Time{Unit: Minute, Value: 1}}, + }, + }, + { + Identifier: Identifier{Type: Header, Value: "user-id"}, + KeyItems: []*KeyItem{ + {Key: DefaultKeyPattern, Token: Token{Number: 300}, Time: Time{Unit: Hour, Value: 1}}, + }, + }, + }, + }, + description: "Should set defaults for multiple items correctly", + }, + { + name: "Nil rule should not panic", + rule: nil, + expectedRule: nil, + description: "Should handle nil rule gracefully", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + // Make a copy of the original rule to avoid modifying test data + var testRule *Rule + if tt.rule != nil { + testRule = &Rule{ + ID: tt.rule.ID, + Resource: tt.rule.Resource, + Strategy: tt.rule.Strategy, + Encoding: tt.rule.Encoding, + SpecificItems: make([]*SpecificItem, len(tt.rule.SpecificItems)), + } + for i, item := range tt.rule.SpecificItems { + if item != nil { + testRule.SpecificItems[i] = &SpecificItem{ + Identifier: item.Identifier, + KeyItems: make([]*KeyItem, len(item.KeyItems)), + } + for j, keyItem := range item.KeyItems { + if keyItem != nil { + testRule.SpecificItems[i].KeyItems[j] = &KeyItem{ + Key: keyItem.Key, + Token: keyItem.Token, + Time: keyItem.Time, + } + } + } + } + } + } + + // Apply defaults + testRule.setDefaultRuleOption() + + // Verify results + if tt.expectedRule == nil { + if testRule != nil { + t.Error("Expected nil rule after setDefaultRuleOption") + } + return + } + + if testRule == nil { + t.Fatal("Expected non-nil rule after setDefaultRuleOption") + } + + // Check basic fields + if testRule.Resource != tt.expectedRule.Resource { + t.Errorf("Resource: expected %q, got %q", tt.expectedRule.Resource, testRule.Resource) + } + + if testRule.Encoding.Model != tt.expectedRule.Encoding.Model { + t.Errorf("Model: expected %q, got %q", tt.expectedRule.Encoding.Model, testRule.Encoding.Model) + } + + // Check specific items + if len(testRule.SpecificItems) != len(tt.expectedRule.SpecificItems) { + t.Errorf("SpecificItems length: expected %d, got %d", + len(tt.expectedRule.SpecificItems), len(testRule.SpecificItems)) + return + } + + for i, expectedItem := range tt.expectedRule.SpecificItems { + actualItem := testRule.SpecificItems[i] + + if actualItem.Identifier.Value != expectedItem.Identifier.Value { + t.Errorf("SpecificItem[%d].Identifier.Value: expected %q, got %q", + i, expectedItem.Identifier.Value, actualItem.Identifier.Value) + } + + if len(actualItem.KeyItems) != len(expectedItem.KeyItems) { + t.Errorf("SpecificItem[%d].KeyItems length: expected %d, got %d", + i, len(expectedItem.KeyItems), len(actualItem.KeyItems)) + continue + } + + for j, expectedKeyItem := range expectedItem.KeyItems { + actualKeyItem := actualItem.KeyItems[j] + + if actualKeyItem.Key != expectedKeyItem.Key { + t.Errorf("SpecificItem[%d].KeyItems[%d].Key: expected %q, got %q", + i, j, expectedKeyItem.Key, actualKeyItem.Key) + } + } + } + }) + } +} + +// Test Rule.filterDuplicatedItem function +func TestRule_FilterDuplicatedItem(t *testing.T) { + tests := []struct { + name string + rule *Rule + expectedLen int + description string + verify func(*Rule) bool + }{ + { + name: "Rule with duplicate key items", + rule: &Rule{ + SpecificItems: []*SpecificItem{ + { + Identifier: Identifier{Type: AllIdentifier, Value: "test"}, + KeyItems: []*KeyItem{ + { + Key: "duplicate-key", + Token: Token{Number: 1000, CountStrategy: TotalTokens}, + Time: Time{Unit: Minute, Value: 1}, + }, + { + Key: "duplicate-key", + Token: Token{Number: 1000, CountStrategy: TotalTokens}, + Time: Time{Unit: Minute, Value: 1}, + }, + { + Key: "unique-key", + Token: Token{Number: 500, CountStrategy: InputTokens}, + Time: Time{Unit: Second, Value: 30}, + }, + }, + }, + }, + }, + expectedLen: 1, + description: "Should remove duplicate key items", + verify: func(r *Rule) bool { + if len(r.SpecificItems) != 1 { + return false + } + if len(r.SpecificItems[0].KeyItems) != 2 { + return false + } + // Should keep the last occurrence (reverse order processing) + return r.SpecificItems[0].KeyItems[0].Key == "unique-key" + }, + }, + { + name: "Rule with duplicate specific items", + rule: &Rule{ + SpecificItems: []*SpecificItem{ + { + Identifier: Identifier{Type: AllIdentifier, Value: "same"}, + KeyItems: []*KeyItem{ + { + Key: "key1", + Token: Token{Number: 1000, CountStrategy: TotalTokens}, + Time: Time{Unit: Minute, Value: 1}, + }, + }, + }, + { + Identifier: Identifier{Type: AllIdentifier, Value: "same"}, + KeyItems: []*KeyItem{ + { + Key: "key1", + Token: Token{Number: 500, CountStrategy: TotalTokens}, + Time: Time{Unit: Second, Value: 60}, + }, + }, + }, + }, + }, + expectedLen: 1, + description: "Should handle duplicate specific items", + verify: func(r *Rule) bool { + return len(r.SpecificItems) == 1 && len(r.SpecificItems[0].KeyItems) >= 1 + }, + }, + { + name: "Rule with no duplicates", + rule: &Rule{ + SpecificItems: []*SpecificItem{ + { + Identifier: Identifier{Type: AllIdentifier, Value: "test1"}, + KeyItems: []*KeyItem{ + { + Key: "key1", + Token: Token{Number: 1000, CountStrategy: TotalTokens}, + Time: Time{Unit: Minute, Value: 1}, + }, + }, + }, + { + Identifier: Identifier{Type: Header, Value: "test2"}, + KeyItems: []*KeyItem{ + { + Key: "key2", + Token: Token{Number: 500, CountStrategy: InputTokens}, + Time: Time{Unit: Second, Value: 30}, + }, + }, + }, + }, + }, + expectedLen: 2, + description: "Should preserve all unique items", + verify: func(r *Rule) bool { + return len(r.SpecificItems) == 2 + }, + }, + { + name: "Rule with empty specific items", + rule: &Rule{ + SpecificItems: []*SpecificItem{}, + }, + expectedLen: 0, + description: "Should handle empty specific items", + verify: func(r *Rule) bool { + return len(r.SpecificItems) == 0 + }, + }, + { + name: "Rule with specific item having empty key items", + rule: &Rule{ + SpecificItems: []*SpecificItem{ + { + Identifier: Identifier{Type: AllIdentifier, Value: "test"}, + KeyItems: []*KeyItem{}, + }, + { + Identifier: Identifier{Type: Header, Value: "valid"}, + KeyItems: []*KeyItem{ + { + Key: "valid-key", + Token: Token{Number: 1000, CountStrategy: TotalTokens}, + Time: Time{Unit: Minute, Value: 1}, + }, + }, + }, + }, + }, + expectedLen: 1, + description: "Should remove specific items with empty key items", + verify: func(r *Rule) bool { + return len(r.SpecificItems) == 1 && + r.SpecificItems[0].Identifier.Value == "valid" + }, + }, + { + name: "Complex rule with mixed duplicates", + rule: &Rule{ + SpecificItems: []*SpecificItem{ + { + Identifier: Identifier{Type: AllIdentifier, Value: "id1"}, + KeyItems: []*KeyItem{ + {Key: "key1", Token: Token{Number: 100, CountStrategy: TotalTokens}, Time: Time{Unit: Second, Value: 1}}, + {Key: "key2", Token: Token{Number: 200, CountStrategy: InputTokens}, Time: Time{Unit: Minute, Value: 1}}, + {Key: "key1", Token: Token{Number: 100, CountStrategy: TotalTokens}, Time: Time{Unit: Second, Value: 1}}, // Duplicate + }, + }, + { + Identifier: Identifier{Type: Header, Value: "id2"}, + KeyItems: []*KeyItem{ + {Key: "key3", Token: Token{Number: 300, CountStrategy: OutputTokens}, Time: Time{Unit: Hour, Value: 1}}, + }, + }, + { + Identifier: Identifier{Type: AllIdentifier, Value: "id1"}, + KeyItems: []*KeyItem{ + {Key: "key4", Token: Token{Number: 400, CountStrategy: TotalTokens}, Time: Time{Unit: Day, Value: 1}}, + }, + }, + }, + }, + expectedLen: 3, + description: "Should handle complex duplicate scenarios", + verify: func(r *Rule) bool { + if len(r.SpecificItems) != 3 { + return false + } + // Due to reverse processing, the order might be different + // Just verify we have the expected identifiers + hasId1 := false + hasId2 := false + hasKey4 := false + hasKey1AndKey2 := false + for _, item := range r.SpecificItems { + if item.Identifier.Value == "id1" { + hasId1 = true + if len(item.KeyItems) == 1 { + if item.KeyItems[0].Key == "key4" { + hasKey4 = true + } + } else if len(item.KeyItems) == 2 { + if item.KeyItems[0].Key == "key1" && item.KeyItems[1].Key == "key2" { + hasKey1AndKey2 = true + } + } else { + return false + } + } + if item.Identifier.Value == "id2" { + hasId2 = true + } + } + return hasId1 && hasId2 && hasKey4 && hasKey1AndKey2 + }, + }, + { + name: "Nil rule should not panic", + rule: nil, + expectedLen: 0, + description: "Should handle nil rule gracefully", + verify: func(r *Rule) bool { + return r == nil + }, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + // Make a deep copy to avoid modifying test data + var testRule *Rule + if tt.rule != nil { + testRule = &Rule{ + ID: tt.rule.ID, + Resource: tt.rule.Resource, + Strategy: tt.rule.Strategy, + Encoding: tt.rule.Encoding, + SpecificItems: make([]*SpecificItem, len(tt.rule.SpecificItems)), + } + for i, item := range tt.rule.SpecificItems { + if item != nil { + testRule.SpecificItems[i] = &SpecificItem{ + Identifier: item.Identifier, + KeyItems: make([]*KeyItem, len(item.KeyItems)), + } + for j, keyItem := range item.KeyItems { + if keyItem != nil { + testRule.SpecificItems[i].KeyItems[j] = &KeyItem{ + Key: keyItem.Key, + Token: keyItem.Token, + Time: keyItem.Time, + } + } + } + } + } + } + + // Apply filter + testRule.filterDuplicatedItem() + + // Basic verification + if tt.rule != nil && testRule != nil { + if len(testRule.SpecificItems) > tt.expectedLen { + t.Errorf("Expected at most %d specific items, got %d", + tt.expectedLen, len(testRule.SpecificItems)) + } + } + + // Custom verification + if tt.verify != nil && !tt.verify(testRule) { + t.Errorf("Custom verification failed: %s", tt.description) + } + }) + } +} + +// Test filterDuplicatedItem with edge cases +func TestRule_FilterDuplicatedItem_EdgeCases(t *testing.T) { + t.Run("Rule with nil specific items", func(t *testing.T) { + rule := &Rule{ + SpecificItems: []*SpecificItem{nil, nil}, + } + + // Should not panic + rule.filterDuplicatedItem() + + // Should result in empty specific items + if len(rule.SpecificItems) != 0 { + t.Errorf("Expected empty specific items, got %d", len(rule.SpecificItems)) + } + }) + + t.Run("Rule with nil key items", func(t *testing.T) { + rule := &Rule{ + SpecificItems: []*SpecificItem{ + { + Identifier: Identifier{Type: AllIdentifier, Value: "test"}, + KeyItems: []*KeyItem{nil, nil}, + }, + }, + } + + // Should not panic + rule.filterDuplicatedItem() + + // Should remove the specific item with no valid key items + if len(rule.SpecificItems) != 0 { + t.Errorf("Expected no specific items, got %d", len(rule.SpecificItems)) + } + }) +} + +// Test the hash generation consistency +func TestRule_FilterDuplicatedItem_HashConsistency(t *testing.T) { + // Create two rules with identical key items + rule1 := &Rule{ + SpecificItems: []*SpecificItem{ + { + Identifier: Identifier{Type: AllIdentifier, Value: "test"}, + KeyItems: []*KeyItem{ + { + Key: "same-key", + Token: Token{Number: 1000, CountStrategy: TotalTokens}, + Time: Time{Unit: Minute, Value: 1}, + }, + }, + }, + }, + } + + rule2 := &Rule{ + SpecificItems: []*SpecificItem{ + { + Identifier: Identifier{Type: AllIdentifier, Value: "test"}, + KeyItems: []*KeyItem{ + { + Key: "same-key", + Token: Token{Number: 1000, CountStrategy: TotalTokens}, + Time: Time{Unit: Minute, Value: 1}, + }, + { + Key: "same-key", + Token: Token{Number: 1000, CountStrategy: TotalTokens}, + Time: Time{Unit: Minute, Value: 1}, + }, + }, + }, + }, + } + + rule1.filterDuplicatedItem() + rule2.filterDuplicatedItem() + + // Both should have the same result + if len(rule1.SpecificItems) != len(rule2.SpecificItems) { + t.Error("Hash consistency failed: different number of specific items") + } + + if len(rule1.SpecificItems) > 0 && len(rule2.SpecificItems) > 0 { + if len(rule1.SpecificItems[0].KeyItems) != len(rule2.SpecificItems[0].KeyItems) { + t.Error("Hash consistency failed: different number of key items") + } + } +} + +// Test integration between functions +func TestRule_Integration(t *testing.T) { + // Test that setDefaultRuleOption and filterDuplicatedItem work together + rule := &Rule{ + Encoding: TokenEncoding{Provider: OpenAIEncoderProvider}, + SpecificItems: []*SpecificItem{ + { + Identifier: Identifier{Type: AllIdentifier, Value: ""}, + KeyItems: []*KeyItem{ + {Key: "", Token: Token{Number: 1000}, Time: Time{Unit: Minute, Value: 1}}, + {Key: "", Token: Token{Number: 1000}, Time: Time{Unit: Minute, Value: 1}}, + }, + }, + }, + } + + // Apply defaults first + rule.setDefaultRuleOption() + + // Then filter duplicates + rule.filterDuplicatedItem() + + // Verify the result + if len(rule.SpecificItems) != 1 { + t.Errorf("Expected 1 specific item, got %d", len(rule.SpecificItems)) + } + + if len(rule.SpecificItems[0].KeyItems) != 1 { + t.Errorf("Expected 1 key item after filtering duplicates, got %d", + len(rule.SpecificItems[0].KeyItems)) + } + + if rule.Resource != DefaultResourcePattern { + t.Errorf("Expected default resource pattern %q, got %q", + DefaultResourcePattern, rule.Resource) + } + + if rule.SpecificItems[0].KeyItems[0].Key != DefaultKeyPattern { + t.Errorf("Expected default key pattern %q, got %q", + DefaultKeyPattern, rule.SpecificItems[0].KeyItems[0].Key) + } +} + +// Benchmark tests +func BenchmarkRule_ResourceName(b *testing.B) { + rule := createTestRuleForRule() + + b.ResetTimer() + for i := 0; i < b.N; i++ { + rule.ResourceName() + } +} + +func BenchmarkRule_String(b *testing.B) { + rule := createRuleWithMultipleItems() + + b.ResetTimer() + for i := 0; i < b.N; i++ { + _ = rule.String() + } +} + +func BenchmarkRule_SetDefaultRuleOption(b *testing.B) { + for i := 0; i < b.N; i++ { + rule := createEmptyRule() + rule.setDefaultRuleOption() + } +} + +func BenchmarkRule_FilterDuplicatedItem(b *testing.B) { + // Create a rule with many duplicates for benchmarking + rule := &Rule{ + SpecificItems: []*SpecificItem{ + { + Identifier: Identifier{Type: AllIdentifier, Value: "test"}, + KeyItems: make([]*KeyItem, 100), + }, + }, + } + + // Fill with duplicate items + for i := 0; i < 100; i++ { + rule.SpecificItems[0].KeyItems[i] = &KeyItem{ + Key: "duplicate-key", + Token: Token{Number: 1000, CountStrategy: TotalTokens}, + Time: Time{Unit: Minute, Value: 1}, + } + } + + b.ResetTimer() + for i := 0; i < b.N; i++ { + // Create a fresh copy for each iteration + testRule := &Rule{ + SpecificItems: []*SpecificItem{ + { + Identifier: rule.SpecificItems[0].Identifier, + KeyItems: make([]*KeyItem, len(rule.SpecificItems[0].KeyItems)), + }, + }, + } + copy(testRule.SpecificItems[0].KeyItems, rule.SpecificItems[0].KeyItems) + + testRule.filterDuplicatedItem() + } +} diff --git a/core/llm_token_ratelimit/rule_validater.go b/core/llm_token_ratelimit/rule_validater.go new file mode 100644 index 000000000..5a44ea2e4 --- /dev/null +++ b/core/llm_token_ratelimit/rule_validater.go @@ -0,0 +1,199 @@ +// Copyright 1999-2020 Alibaba Group Holding Ltd. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package llm_token_ratelimit + +import ( + "fmt" + "regexp" +) + +func IsValidRule(r *Rule) error { + if r == nil { + return fmt.Errorf("rule cannot be nil") + } + + // Validate Resource + if err := validateResource(r.Resource); err != nil { + return fmt.Errorf("invalid resource: %w", err) + } + + // Validate Strategy + if err := validateStrategy(r.Strategy); err != nil { + return fmt.Errorf("invalid strategy: %w", err) + } + + // Validate Encoding + if err := validateEncoding(r.Encoding); err != nil { + return fmt.Errorf("invalid token encoding: %w", err) + } + + // Validate SpecificItems (required) + if len(r.SpecificItems) == 0 { + return fmt.Errorf("specificItems cannot be empty") + } + + for i, specificItem := range r.SpecificItems { + if err := validateSpecificItem(specificItem); err != nil { + return fmt.Errorf("invalid specificItem[%d]: %w", i, err) + } + } + + return nil +} + +func validateResource(resource string) error { + if resource == "" { + return fmt.Errorf("resource pattern cannot be empty") + } + + if _, err := regexp.Compile(resource); err != nil { + return fmt.Errorf("resource pattern is not a valid regex: %w", err) + } + + return nil +} + +func validateStrategy(strategy Strategy) error { + switch strategy { + case FixedWindow, PETA: + return nil + default: + return fmt.Errorf("unsupported strategy: %s", strategy.String()) + } +} + +func validateEncoding(encoding TokenEncoding) error { + // Validate TokenEncoding + switch encoding.Provider { + case OpenAIEncoderProvider: + return nil + default: + return fmt.Errorf("unsupported token encoding provider: %v", encoding.Provider.String()) + } +} + +func validateSpecificItem(specificItem *SpecificItem) error { + if specificItem == nil { + return fmt.Errorf("specificItem cannot be nil") + } + + // Validate Identifier + if err := validateIdentifier(&specificItem.Identifier); err != nil { + return fmt.Errorf("invalid identifier: %w", err) + } + + // Validate KeyItems (required) + if len(specificItem.KeyItems) == 0 { + return fmt.Errorf("keyItems cannot be empty") + } + + for i, keyItem := range specificItem.KeyItems { + if err := validateKeyItem(keyItem); err != nil { + return fmt.Errorf("invalid keyItem[%d]: %w", i, err) + } + } + + return nil +} + +func validateIdentifier(identifier *Identifier) error { + if identifier == nil { + return fmt.Errorf("identifier cannot be nil") + } + + // Validate Type + switch identifier.Type { + case AllIdentifier, Header: + // Valid types + default: + return fmt.Errorf("unsupported identifier type: %v", identifier.Type) + } + + if identifier.Value == "" { + return fmt.Errorf("identifier value pattern cannot be empty") + } + if _, err := regexp.Compile(identifier.Value); err != nil { + return fmt.Errorf("identifier value is not a valid regex: %w", err) + } + + return nil +} + +func validateKeyItem(keyItem *KeyItem) error { + if keyItem == nil { + return fmt.Errorf("keyItem cannot be nil") + } + + if keyItem.Key == "" { + return fmt.Errorf("key pattern cannot be empty") + } + if _, err := regexp.Compile(keyItem.Key); err != nil { + return fmt.Errorf("key pattern is not a valid regex: %w", err) + } + + // Validate Token (required) + if err := validateToken(&keyItem.Token); err != nil { + return fmt.Errorf("invalid token: %w", err) + } + + // Validate Time (required) + if err := validateTime(&keyItem.Time); err != nil { + return fmt.Errorf("invalid time: %w", err) + } + + return nil +} + +func validateToken(token *Token) error { + if token == nil { + return fmt.Errorf("token cannot be nil") + } + + // Validate Number (required, must be positive) + if token.Number < 0 { + return fmt.Errorf("token number must be positive, got: %d", token.Number) + } + + // Validate CountStrategy + switch token.CountStrategy { + case TotalTokens, InputTokens, OutputTokens: + // Valid strategies + default: + return fmt.Errorf("unsupported count strategy: %v", token.CountStrategy) + } + + return nil +} + +func validateTime(time *Time) error { + if time == nil { + return fmt.Errorf("time cannot be nil") + } + + // Validate Unit (required) + switch time.Unit { + case Second, Minute, Hour, Day: + // Valid units + default: + return fmt.Errorf("unsupported time unit: %v", time.Unit) + } + + // Validate Value (required, must be positive) + if time.Value < 0 { + return fmt.Errorf("time value must be positive, got: %d", time.Value) + } + + return nil +} diff --git a/core/llm_token_ratelimit/rule_validater_test.go b/core/llm_token_ratelimit/rule_validater_test.go new file mode 100644 index 000000000..987bd5696 --- /dev/null +++ b/core/llm_token_ratelimit/rule_validater_test.go @@ -0,0 +1,645 @@ +// Copyright 1999-2020 Alibaba Group Holding Ltd. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package llm_token_ratelimit + +import ( + "testing" +) + +func TestIsValidRule(t *testing.T) { + tests := []struct { + name string + rule *Rule + wantError bool + errorMsg string + }{ + { + name: "nil rule", + rule: nil, + wantError: true, + errorMsg: "rule cannot be nil", + }, + { + name: "valid rule", + rule: &Rule{ + Resource: "test-resource", + Strategy: FixedWindow, + SpecificItems: []*SpecificItem{ + { + Identifier: Identifier{ + Type: Header, + Value: "api-key", + }, + KeyItems: []*KeyItem{ + { + Key: "user-*", + Token: Token{ + Number: 1000, + CountStrategy: TotalTokens, + }, + Time: Time{ + Unit: Minute, + Value: 1, + }, + }, + }, + }, + }, + }, + wantError: false, + }, + { + name: "empty rule items", + rule: &Rule{ + Resource: "test-resource", + Strategy: FixedWindow, + SpecificItems: []*SpecificItem{}, + }, + wantError: true, + errorMsg: "specificItems cannot be empty", + }, + { + name: "invalid resource", + rule: &Rule{ + Resource: "", + Strategy: FixedWindow, + SpecificItems: []*SpecificItem{ + { + Identifier: Identifier{Type: Header, Value: "api-key"}, + KeyItems: []*KeyItem{ + { + Key: "user-*", + Token: Token{Number: 1000, CountStrategy: TotalTokens}, + Time: Time{Unit: Minute, Value: 1}, + }, + }, + }, + }, + }, + wantError: true, + errorMsg: "invalid resource", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + err := IsValidRule(tt.rule) + if tt.wantError { + if err == nil { + t.Errorf("IsValidRule() expected error but got nil") + } else if tt.errorMsg != "" && !contains(err.Error(), tt.errorMsg) { + t.Errorf("IsValidRule() error = %v, want error containing %v", err, tt.errorMsg) + } + } else { + if err != nil { + t.Errorf("IsValidRule() unexpected error = %v", err) + } + } + }) + } +} + +func TestValidateResource(t *testing.T) { + tests := []struct { + name string + resource string + wantError bool + errorMsg string + }{ + { + name: "empty resource", + resource: "", + wantError: true, + errorMsg: "resource pattern cannot be empty", + }, + { + name: "valid resource pattern", + resource: "api-.*", + wantError: false, + }, + { + name: "valid wildcard pattern", + resource: ".*", + wantError: false, + }, + { + name: "invalid regex pattern", + resource: "[invalid", + wantError: true, + errorMsg: "resource pattern is not a valid regex", + }, + { + name: "complex valid pattern", + resource: "^(api|service)-.+$", + wantError: false, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + err := validateResource(tt.resource) + if tt.wantError { + if err == nil { + t.Errorf("validateResource() expected error but got nil") + } else if tt.errorMsg != "" && !contains(err.Error(), tt.errorMsg) { + t.Errorf("validateResource() error = %v, want error containing %v", err, tt.errorMsg) + } + } else { + if err != nil { + t.Errorf("validateResource() unexpected error = %v", err) + } + } + }) + } +} + +func TestValidateStrategy(t *testing.T) { + tests := []struct { + name string + strategy Strategy + wantError bool + errorMsg string + }{ + { + name: "valid fixed window strategy", + strategy: FixedWindow, + wantError: false, + }, + { + name: "invalid strategy", + strategy: Strategy(999), + wantError: true, + errorMsg: "unsupported strategy", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + err := validateStrategy(tt.strategy) + if tt.wantError { + if err == nil { + t.Errorf("validateStrategy() expected error but got nil") + } else if tt.errorMsg != "" && !contains(err.Error(), tt.errorMsg) { + t.Errorf("validateStrategy() error = %v, want error containing %v", err, tt.errorMsg) + } + } else { + if err != nil { + t.Errorf("validateStrategy() unexpected error = %v", err) + } + } + }) + } +} + +func TestValidateIdentifier(t *testing.T) { + tests := []struct { + name string + identifier *Identifier + wantError bool + errorMsg string + }{ + { + name: "nil identifier", + identifier: nil, + wantError: true, + errorMsg: "identifier cannot be nil", + }, + { + name: "valid all identifier", + identifier: &Identifier{ + Type: AllIdentifier, + Value: ".*", + }, + wantError: false, + }, + { + name: "valid header identifier", + identifier: &Identifier{ + Type: Header, + Value: "api-key", + }, + wantError: false, + }, + { + name: "empty value", + identifier: &Identifier{ + Type: Header, + Value: "", + }, + wantError: true, + }, + { + name: "invalid identifier type", + identifier: &Identifier{ + Type: IdentifierType(999), + Value: "api-key", + }, + wantError: true, + errorMsg: "unsupported identifier type", + }, + { + name: "invalid regex value", + identifier: &Identifier{ + Type: Header, + Value: "[invalid", + }, + wantError: true, + errorMsg: "identifier value is not a valid regex", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + err := validateIdentifier(tt.identifier) + if tt.wantError { + if err == nil { + t.Errorf("validateIdentifier() expected error but got nil") + } else if tt.errorMsg != "" && !contains(err.Error(), tt.errorMsg) { + t.Errorf("validateIdentifier() error = %v, want error containing %v", err, tt.errorMsg) + } + } else { + if err != nil { + t.Errorf("validateIdentifier() unexpected error = %v", err) + } + } + }) + } +} + +func TestValidateToken(t *testing.T) { + tests := []struct { + name string + token *Token + wantError bool + errorMsg string + }{ + { + name: "nil token", + token: nil, + wantError: true, + errorMsg: "token cannot be nil", + }, + { + name: "valid token", + token: &Token{ + Number: 1000, + CountStrategy: TotalTokens, + }, + wantError: false, + }, + { + name: "zero token number", + token: &Token{ + Number: 0, + CountStrategy: TotalTokens, + }, + wantError: false, + }, + { + name: "negative token number", + token: &Token{ + Number: -100, + CountStrategy: TotalTokens, + }, + wantError: true, + errorMsg: "token number must be positive", + }, + { + name: "maximum valid token number", + token: &Token{ + Number: 1000000000, // exactly 1 billion + CountStrategy: TotalTokens, + }, + wantError: false, + }, + { + name: "valid input tokens strategy", + token: &Token{ + Number: 1000, + CountStrategy: InputTokens, + }, + wantError: false, + }, + { + name: "valid output tokens strategy", + token: &Token{ + Number: 1000, + CountStrategy: OutputTokens, + }, + wantError: false, + }, + { + name: "invalid count strategy", + token: &Token{ + Number: 1000, + CountStrategy: CountStrategy(999), + }, + wantError: true, + errorMsg: "unsupported count strategy", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + err := validateToken(tt.token) + if tt.wantError { + if err == nil { + t.Errorf("validateToken() expected error but got nil") + } else if tt.errorMsg != "" && !contains(err.Error(), tt.errorMsg) { + t.Errorf("validateToken() error = %v, want error containing %v", err, tt.errorMsg) + } + } else { + if err != nil { + t.Errorf("validateToken() unexpected error = %v", err) + } + } + }) + } +} + +func TestValidateTime(t *testing.T) { + tests := []struct { + name string + time *Time + wantError bool + errorMsg string + }{ + { + name: "nil time", + time: nil, + wantError: true, + errorMsg: "time cannot be nil", + }, + { + name: "valid second time unit", + time: &Time{ + Unit: Second, + Value: 30, + }, + wantError: false, + }, + { + name: "valid minute time unit", + time: &Time{ + Unit: Minute, + Value: 5, + }, + wantError: false, + }, + { + name: "valid hour time unit", + time: &Time{ + Unit: Hour, + Value: 2, + }, + wantError: false, + }, + { + name: "valid day time unit", + time: &Time{ + Unit: Day, + Value: 1, + }, + wantError: false, + }, + { + name: "invalid time unit", + time: &Time{ + Unit: TimeUnit(999), + Value: 1, + }, + wantError: true, + errorMsg: "unsupported time unit", + }, + { + name: "zero time value", + time: &Time{ + Unit: Minute, + Value: 0, + }, + wantError: false, + }, + { + name: "negative time value", + time: &Time{ + Unit: Minute, + Value: -5, + }, + wantError: true, + errorMsg: "time value must be positive", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + err := validateTime(tt.time) + if tt.wantError { + if err == nil { + t.Errorf("validateTime() expected error but got nil") + } else if tt.errorMsg != "" && !contains(err.Error(), tt.errorMsg) { + t.Errorf("validateTime() error = %v, want error containing %v", err, tt.errorMsg) + } + } else { + if err != nil { + t.Errorf("validateTime() unexpected error = %v", err) + } + } + }) + } +} + +func TestValidateKeyItem(t *testing.T) { + tests := []struct { + name string + keyItem *KeyItem + wantError bool + errorMsg string + }{ + { + name: "nil key item", + keyItem: nil, + wantError: true, + errorMsg: "keyItem cannot be nil", + }, + { + name: "valid key item", + keyItem: &KeyItem{ + Key: "user-*", + Token: Token{ + Number: 1000, + CountStrategy: TotalTokens, + }, + Time: Time{ + Unit: Minute, + Value: 1, + }, + }, + wantError: false, + }, + { + name: "empty key", + keyItem: &KeyItem{ + Key: "", + Token: Token{ + Number: 1000, + CountStrategy: TotalTokens, + }, + Time: Time{ + Unit: Minute, + Value: 1, + }, + }, + wantError: true, + }, + { + name: "invalid key regex", + keyItem: &KeyItem{ + Key: "[invalid", + Token: Token{ + Number: 1000, + CountStrategy: TotalTokens, + }, + Time: Time{ + Unit: Minute, + Value: 1, + }, + }, + wantError: true, + errorMsg: "key pattern is not a valid regex", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + err := validateKeyItem(tt.keyItem) + if tt.wantError { + if err == nil { + t.Errorf("validateKeyItem() expected error but got nil") + } else if tt.errorMsg != "" && !contains(err.Error(), tt.errorMsg) { + t.Errorf("validateKeyItem() error = %v, want error containing %v", err, tt.errorMsg) + } + } else { + if err != nil { + t.Errorf("validateKeyItem() unexpected error = %v", err) + } + } + }) + } +} + +func TestValidateSpecificItem(t *testing.T) { + tests := []struct { + name string + specificItem *SpecificItem + wantError bool + errorMsg string + }{ + { + name: "nil rule item", + specificItem: nil, + wantError: true, + errorMsg: "specificItem cannot be nil", + }, + { + name: "valid rule item", + specificItem: &SpecificItem{ + Identifier: Identifier{ + Type: Header, + Value: "api-key", + }, + KeyItems: []*KeyItem{ + { + Key: "user-*", + Token: Token{ + Number: 1000, + CountStrategy: TotalTokens, + }, + Time: Time{ + Unit: Minute, + Value: 1, + }, + }, + }, + }, + wantError: false, + }, + { + name: "empty key items", + specificItem: &SpecificItem{ + Identifier: Identifier{ + Type: Header, + Value: "api-key", + }, + KeyItems: []*KeyItem{}, + }, + wantError: true, + errorMsg: "keyItems cannot be empty", + }, + { + name: "invalid identifier", + specificItem: &SpecificItem{ + Identifier: Identifier{ + Type: IdentifierType(999), + Value: "api-key", + }, + KeyItems: []*KeyItem{ + { + Key: "user-*", + Token: Token{ + Number: 1000, + CountStrategy: TotalTokens, + }, + Time: Time{ + Unit: Minute, + Value: 1, + }, + }, + }, + }, + wantError: true, + errorMsg: "invalid identifier", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + err := validateSpecificItem(tt.specificItem) + if tt.wantError { + if err == nil { + t.Errorf("validateSpecificItem() expected error but got nil") + } else if tt.errorMsg != "" && !contains(err.Error(), tt.errorMsg) { + t.Errorf("validateSpecificItem() error = %v, want error containing %v", err, tt.errorMsg) + } + } else { + if err != nil { + t.Errorf("validateSpecificItem() unexpected error = %v", err) + } + } + }) + } +} + +// Helper function to check if a string contains a substring +func contains(str, substr string) bool { + return len(str) >= len(substr) && (str == substr || + func() bool { + for i := 0; i <= len(str)-len(substr); i++ { + if str[i:i+len(substr)] == substr { + return true + } + } + return false + }()) +} diff --git a/core/llm_token_ratelimit/script/fixed_window/query.lua b/core/llm_token_ratelimit/script/fixed_window/query.lua new file mode 100644 index 000000000..fe6233072 --- /dev/null +++ b/core/llm_token_ratelimit/script/fixed_window/query.lua @@ -0,0 +1,27 @@ +-- Copyright 1999-2020 Alibaba Group Holding Ltd. +-- +-- Licensed under the Apache License, Version 2.0 (the "License"); +-- you may not use this file except in compliance with the License. +-- You may obtain a copy of the License at +-- +-- http://www.apache.org/licenses/LICENSE-2.0 +-- +-- Unless required by applicable law or agreed to in writing, software +-- distributed under the License is distributed on an "AS IS" BASIS, +-- WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +-- See the License for the specific language governing permissions and +-- limitations under the License. +-- KEYS[1]: Fixed Window Key ("") +-- ARGV[1]: Maximum Token capacity +-- ARGV[2]: Window size (milliseconds) +local fixed_window_key = KEYS[1] + +local max_token_capacity = tonumber(ARGV[1]) +local window_size = tonumber(ARGV[2]) + +local ttl = redis.call('PTTL', fixed_window_key) +if ttl < 0 then + redis.call('SET', fixed_window_key, max_token_capacity, 'PX', window_size) + return {max_token_capacity, window_size} +end +return {tonumber(redis.call('GET', fixed_window_key)), ttl} diff --git a/core/llm_token_ratelimit/script/fixed_window/update.lua b/core/llm_token_ratelimit/script/fixed_window/update.lua new file mode 100644 index 000000000..2d98a1a17 --- /dev/null +++ b/core/llm_token_ratelimit/script/fixed_window/update.lua @@ -0,0 +1,29 @@ +-- Copyright 1999-2020 Alibaba Group Holding Ltd. +-- +-- Licensed under the Apache License, Version 2.0 (the "License"); +-- you may not use this file except in compliance with the License. +-- You may obtain a copy of the License at +-- +-- http://www.apache.org/licenses/LICENSE-2.0 +-- +-- Unless required by applicable law or agreed to in writing, software +-- distributed under the License is distributed on an "AS IS" BASIS, +-- WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +-- See the License for the specific language governing permissions and +-- limitations under the License. +-- KEYS[1]: Fixed Window Key ("") +-- ARGV[1]: Maximum Token capacity +-- ARGV[2]: Window size (milliseconds) +-- ARGV[3]: Actual token consumption +local fixed_window_key = KEYS[1] + +local max_token_capacity = tonumber(ARGV[1]) +local window_size = tonumber(ARGV[2]) +local actual = tonumber(ARGV[3]) + +local ttl = redis.call('PTTL', fixed_window_key) +if ttl < 0 then + redis.call('SET', fixed_window_key, max_token_capacity - actual, 'PX', window_size) + return {max_token_capacity - actual, window_size} +end +return {tonumber(redis.call('DECRBY', fixed_window_key, actual)), ttl} diff --git a/core/llm_token_ratelimit/script/peta/correct.lua b/core/llm_token_ratelimit/script/peta/correct.lua new file mode 100644 index 000000000..4a8a976c7 --- /dev/null +++ b/core/llm_token_ratelimit/script/peta/correct.lua @@ -0,0 +1,139 @@ +-- Copyright 1999-2020 Alibaba Group Holding Ltd. +-- +-- Licensed under the Apache License, Version 2.0 (the "License"); +-- you may not use this file except in compliance with the License. +-- You may obtain a copy of the License at +-- +-- http://www.apache.org/licenses/LICENSE-2.0 +-- +-- Unless required by applicable law or agreed to in writing, software +-- distributed under the License is distributed on an "AS IS" BASIS, +-- WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +-- See the License for the specific language governing permissions and +-- limitations under the License. +-- KEYS[1]: Sliding Window Key ("{shard-}:sliding-window:") +-- KEYS[2]: Token Bucket Key ("{shard-}:token-bucket:") +-- KEYS[3]: Token Encoder Key ("{shard-}:token-encoder:::") +-- ARGV[1]: Estimated token consumption +-- ARGV[2]: Current timestamp (milliseconds) +-- ARGV[3]: Token bucket capacity +-- ARGV[4]: Window size (milliseconds) +-- ARGV[5]: Actual token consumption +-- ARGV[6]: Random string for sliding window value (length less than or equal to 255) +local MAX_SEARCH_ITRATIONS = 64 + +local function calculate_tokens_in_range(key, start_time, end_time) + local valid_list = redis.call('ZRANGEBYSCORE', key, start_time, end_time) + local valid_tokens = 0 + for _, v in ipairs(valid_list) do + local _, tokens = struct.unpack('Bc0L', v) + valid_tokens = valid_tokens + tokens + end + return valid_tokens +end + +local function binary_search_compensation_time(key, L, R, window_size, max_capacity, predicted_error) + local iter = 0 + while L < R and iter < MAX_SEARCH_ITRATIONS do + iter = iter + 1 + local mid = math.floor((L + R) / 2) + local valid_tokens = calculate_tokens_in_range(key, mid - window_size, mid) + if valid_tokens + predicted_error <= max_capacity then + R = mid + else + L = mid + 1 + end + end + return L +end + +local sliding_window_key = tostring(KEYS[1]) +local token_bucket_key = tostring(KEYS[2]) +local token_encoder_key = tostring(KEYS[3]) + +local estimated = tonumber(ARGV[1]) +local current_timestamp = tonumber(ARGV[2]) +local bucket_capacity = tonumber(ARGV[3]) +local window_size = tonumber(ARGV[4]) +local actual = tonumber(ARGV[5]) +local random_string = tostring(ARGV[6]) + +-- Valid window start time +local window_start = current_timestamp - window_size +-- Get bucket +local bucket = redis.call('HMGET', token_bucket_key, 'capacity', 'max_capacity') +local current_capacity = tonumber(bucket[1]) +local max_capacity = tonumber(bucket[2]) +-- Initialize bucket manually if it doesn't exist +if not current_capacity then + current_capacity = bucket_capacity + max_capacity = bucket_capacity + redis.call('HMSET', token_bucket_key, 'capacity', bucket_capacity, 'max_capacity', bucket_capacity) + redis.call('ZADD', sliding_window_key, current_timestamp, + struct.pack('Bc0L', string.len(random_string), random_string, 0)) +end +-- Calculate expired tokens +local released_tokens = calculate_tokens_in_range(sliding_window_key, 0, window_start) +if released_tokens > 0 then -- Expired tokens exist, attempt to replenish new tokens + -- Clean up expired data + redis.call('ZREMRANGEBYSCORE', sliding_window_key, 0, window_start) + -- Calculate valid tokens + local valid_tokens = calculate_tokens_in_range(sliding_window_key, '-inf', '+inf') + -- Update token count + if current_capacity + released_tokens > max_capacity then -- If current capacity plus released tokens exceeds max capacity, reset to max capacity minus valid tokens + current_capacity = max_capacity - valid_tokens + else -- Otherwise, directly add the released tokens + current_capacity = current_capacity + released_tokens + end + -- Immediately replenish new tokens + redis.call('HSET', token_bucket_key, 'capacity', current_capacity) +end +-- Update the difference from the token encoder +local difference = actual - estimated +redis.call('SET', token_encoder_key, difference) +-- Correction result for reservation +local correct_result = 0 +if estimated < 0 or actual < 0 then + correct_result = 3 -- Invalid value +elseif estimated < actual then -- Underestimation + -- Mainly handle underestimation cases to properly limit actual usage; overestimation may reject requests but won't affect downstream services + -- Calculate prediction error + local predicted_error = math.abs(actual - estimated) + -- directly deduct all underestimated tokens + current_capacity = current_capacity - predicted_error + redis.call('HSET', token_bucket_key, 'capacity', current_capacity) + -- Get the latest valid timestamp + local last_valid_window = redis.call('ZRANGE', sliding_window_key, -1, -1, 'WITHSCORES') + local compensation_start = tonumber(last_valid_window[2]) + if not compensation_start then -- Possibly all data just expired, use current timestamp minus window size as start + compensation_start = current_timestamp + end + while predicted_error ~= 0 do -- Distribute to future windows until all error is distributed + if max_capacity >= predicted_error then + local compensation_time = binary_search_compensation_time(sliding_window_key, compensation_start, + compensation_start + window_size, window_size, max_capacity, predicted_error) + if calculate_tokens_in_range(sliding_window_key, compensation_time - window_size, compensation_time) + + predicted_error > max_capacity then + correct_result = 1 -- If the compensation time exceeds max capacity, return 1 to indicate failure + break + end + redis.call('ZADD', sliding_window_key, compensation_time, + struct.pack('Bc0L', string.len(random_string), random_string, predicted_error)) + predicted_error = 0 + else + redis.call('ZADD', sliding_window_key, compensation_start, + struct.pack('Bc0L', string.len(random_string), random_string, max_capacity)) + predicted_error = predicted_error - max_capacity + compensation_start = compensation_start + window_size + end + end +elseif estimated > actual then -- Overestimation + correct_result = 2 +end + +-- Set expiration time to window size plus 5 seconds buffer +redis.call('PEXPIRE', sliding_window_key, window_size + 5000) +redis.call('PEXPIRE', token_bucket_key, window_size + 5000) +redis.call('PEXPIRE', token_encoder_key, window_size + 5000) + +return {correct_result} diff --git a/core/llm_token_ratelimit/script/peta/withhold.lua b/core/llm_token_ratelimit/script/peta/withhold.lua new file mode 100644 index 000000000..2ce60678a --- /dev/null +++ b/core/llm_token_ratelimit/script/peta/withhold.lua @@ -0,0 +1,110 @@ +-- Copyright 1999-2020 Alibaba Group Holding Ltd. +-- +-- Licensed under the Apache License, Version 2.0 (the "License"); +-- you may not use this file except in compliance with the License. +-- You may obtain a copy of the License at +-- +-- http://www.apache.org/licenses/LICENSE-2.0 +-- +-- Unless required by applicable law or agreed to in writing, software +-- distributed under the License is distributed on an "AS IS" BASIS, +-- WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +-- See the License for the specific language governing permissions and +-- limitations under the License. +-- KEYS[1]: Sliding Window Key ("{shard-}:sliding-window:") +-- KEYS[2]: Token Bucket Key ("{shard-}:token-bucket:") +-- KEYS[3]: Token Encoder Key ("{shard-}:token-encoder:::") +-- ARGV[1]: Estimated token consumption +-- ARGV[2]: Current timestamp (milliseconds) +-- ARGV[3]: Token bucket capacity +-- ARGV[4]: Window size (milliseconds) +-- ARGV[5]: Random string for sliding window unique value (length less than or equal to 255) +local function calculate_tokens_in_range(key, start_time, end_time) + local valid_list = redis.call('ZRANGEBYSCORE', key, start_time, end_time) + local valid_tokens = 0 + for _, v in ipairs(valid_list) do + local _, tokens = struct.unpack('Bc0L', v) + valid_tokens = valid_tokens + tokens + end + return valid_tokens +end + +local sliding_window_key = tostring(KEYS[1]) +local token_bucket_key = tostring(KEYS[2]) +local token_encoder_key = tostring(KEYS[3]) + +local estimated = tonumber(ARGV[1]) +local current_timestamp = tonumber(ARGV[2]) +local bucket_capacity = tonumber(ARGV[3]) +local window_size = tonumber(ARGV[4]) +local random_string = tostring(ARGV[5]) + +-- Valid window start time +local window_start = current_timestamp - window_size +-- Waiting time +local waiting_time = 0 +-- Get bucket +local bucket = redis.call('HMGET', token_bucket_key, 'capacity', 'max_capacity') +local current_capacity = tonumber(bucket[1]) +local max_capacity = tonumber(bucket[2]) +-- Initialize bucket manually if it doesn't exist +if not current_capacity then + current_capacity = bucket_capacity + max_capacity = bucket_capacity + redis.call('HMSET', token_bucket_key, 'capacity', bucket_capacity, 'max_capacity', bucket_capacity) + redis.call('ZADD', sliding_window_key, current_timestamp, + struct.pack('Bc0L', string.len(random_string), random_string, 0)) +end +-- Calculate expired tokens +local released_tokens = calculate_tokens_in_range(sliding_window_key, 0, window_start) +if released_tokens > 0 then -- Expired tokens exist, attempt to replenish new tokens + -- Clean up expired data + redis.call('ZREMRANGEBYSCORE', sliding_window_key, 0, window_start) + -- Calculate valid tokens + local valid_tokens = calculate_tokens_in_range(sliding_window_key, '-inf', '+inf') + -- Update token count + if current_capacity + released_tokens > max_capacity then -- If current capacity plus released tokens exceeds max capacity, reset to max capacity minus valid tokens + current_capacity = max_capacity - valid_tokens + else -- Otherwise, directly add the released tokens + current_capacity = current_capacity + released_tokens + end + -- Immediately replenish new tokens + redis.call('HSET', token_bucket_key, 'capacity', current_capacity) +end +-- Plus the difference from the token encoder if it exists +local ttl = redis.call('PTTL', token_encoder_key) +local difference = tonumber(redis.call('GET', token_encoder_key)) +if ttl < 0 then + difference = 0 +else + if difference + estimated >= 0 then + estimated = estimated + difference + else + redis.call('SET', token_encoder_key, 0) + end +end +-- Check if the request can be satisfied +if max_capacity < estimated or estimated <= 0 then -- If max capacity is less than estimated consumption or estimated is less than or equal to 0, return -1 indicating rejection + waiting_time = -1 +elseif current_capacity < estimated then -- If current capacity is insufficient to satisfy estimated consumption, calculate waiting time + -- Get the earliest valid timestamp + local first_valid_window = redis.call('ZRANGE', sliding_window_key, 0, 0, 'WITHSCORES') + local first_valid_start = tonumber(first_valid_window[2]) + if not first_valid_start then + first_valid_start = current_timestamp + end + -- Waiting time = fixed delay + window size - valid window interval + waiting_time = 3 + window_size - (current_timestamp - first_valid_start) +else -- Otherwise, capacity satisfies estimated consumption, no waiting required, update data + redis.call('ZADD', sliding_window_key, current_timestamp, + struct.pack('Bc0L', string.len(random_string), random_string, estimated)) + current_capacity = current_capacity - estimated + redis.call('HSET', token_bucket_key, 'capacity', current_capacity) +end + +-- Set expiration time to window size plus 5 seconds buffer +redis.call('PEXPIRE', sliding_window_key, window_size + 5000) +redis.call('PEXPIRE', token_bucket_key, window_size + 5000) +redis.call('PEXPIRE', token_encoder_key, window_size + 5000) + +return {current_capacity, waiting_time, estimated, difference} diff --git a/core/llm_token_ratelimit/slot.go b/core/llm_token_ratelimit/slot.go new file mode 100644 index 000000000..ff2a2de9b --- /dev/null +++ b/core/llm_token_ratelimit/slot.go @@ -0,0 +1,80 @@ +// Copyright 1999-2020 Alibaba Group Holding Ltd. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package llm_token_ratelimit + +import ( + "github.com/alibaba/sentinel-golang/core/base" + "github.com/alibaba/sentinel-golang/logging" +) + +const ( + RuleCheckSlotOrder uint32 = 6000 +) + +var ( + DefaultSlot = &Slot{} +) + +type Slot struct{} + +func (s *Slot) Order() uint32 { + return RuleCheckSlotOrder +} + +func (s *Slot) Check(ctx *base.EntryContext) *base.TokenResult { + resource := ctx.Resource.Name() + result := ctx.RuleCheckResult + if len(resource) == 0 { + return result + } + + if passed, rule, snapshot := s.checkPass(ctx); !passed { + msg := "llm token ratelimit check blocked" + if result == nil { + result = base.NewTokenResultBlockedWithCause(base.BlockTypeLLMTokenRateLimit, msg, rule, snapshot) + } else { + result.ResetToBlockedWithCause(base.BlockTypeLLMTokenRateLimit, msg, rule, snapshot) + } + } + + return result +} + +func (s *Slot) checkPass(ctx *base.EntryContext) (bool, *Rule, interface{}) { + if !globalConfig.IsEnabled() { + return true, nil, nil + } + requestID := generateUUID() + llmTokenRatelimitCtx, ok := ctx.GetPair(KeyContext).(*Context) + if !ok || llmTokenRatelimitCtx == nil { + llmTokenRatelimitCtx = NewContext() + if llmTokenRatelimitCtx == nil { + logging.Warn("[LLMTokenRateLimit] failed to create llm token ratelimit context", + "requestID", requestID, + ) + return true, nil, nil + } + llmTokenRatelimitCtx.extractArgs(ctx) + llmTokenRatelimitCtx.Set(KeyRequestID, requestID) + ctx.SetPair(KeyContext, llmTokenRatelimitCtx) + } + for _, rule := range getRulesOfResource(ctx.Resource.Name()) { + if !globalRuleMatcher.checkPass(llmTokenRatelimitCtx, rule) { + return false, rule, llmTokenRatelimitCtx.Get(KeyResponseHeaders) + } + } + ctx.SetPair(KeyResponseHeaders, llmTokenRatelimitCtx.Get(KeyResponseHeaders)) + return true, nil, nil +} diff --git a/core/llm_token_ratelimit/stat_slot.go b/core/llm_token_ratelimit/stat_slot.go new file mode 100644 index 000000000..dc896ebed --- /dev/null +++ b/core/llm_token_ratelimit/stat_slot.go @@ -0,0 +1,71 @@ +// Copyright 1999-2020 Alibaba Group Holding Ltd. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package llm_token_ratelimit + +import "github.com/alibaba/sentinel-golang/core/base" + +const ( + StatSlotOrder uint32 = 6000 +) + +var ( + DefaultLLMTokenRatelimitStatSlot = &LLMTokenRatelimitStatSlot{} +) + +type LLMTokenRatelimitStatSlot struct { +} + +func (s *LLMTokenRatelimitStatSlot) Order() uint32 { + return StatSlotOrder +} + +func (c *LLMTokenRatelimitStatSlot) OnEntryPassed(_ *base.EntryContext) { + // Do nothing + return +} + +func (c *LLMTokenRatelimitStatSlot) OnEntryBlocked(_ *base.EntryContext, _ *base.BlockError) { + // Do nothing + return +} + +func (c *LLMTokenRatelimitStatSlot) OnCompleted(ctx *base.EntryContext) { + if !globalConfig.IsEnabled() { + return + } + usedTokenInfos, ok := ctx.GetPair(KeyUsedTokenInfos).(*UsedTokenInfos) + if !ok || usedTokenInfos == nil { + return + } + llmTokenRatelimitCtx, ok := ctx.GetPair(KeyContext).(*Context) + if !ok || llmTokenRatelimitCtx == nil { + return + } + llmTokenRatelimitCtx.Set(KeyUsedTokenInfos, usedTokenInfos) + + rulesInterface := llmTokenRatelimitCtx.Get(KeyMatchedRules) + if rulesInterface == nil { + return + } + + rules, ok := rulesInterface.([]*MatchedRule) + if !ok || rules == nil { + return + } + + for _, rule := range rules { + globalRuleMatcher.update(llmTokenRatelimitCtx, rule) + } +} diff --git a/core/llm_token_ratelimit/token_calculator.go b/core/llm_token_ratelimit/token_calculator.go new file mode 100644 index 000000000..3a218a541 --- /dev/null +++ b/core/llm_token_ratelimit/token_calculator.go @@ -0,0 +1,72 @@ +// Copyright 1999-2020 Alibaba Group Holding Ltd. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package llm_token_ratelimit + +var globalTokenCalculator = NewDefaultTokenCalculator() + +type TokenCalculator interface { + Calculate(ctx *Context, infos *UsedTokenInfos) int +} + +type TokenCalculatorManager struct { + calculators map[CountStrategy]TokenCalculator +} + +func NewDefaultTokenCalculator() *TokenCalculatorManager { + return &TokenCalculatorManager{ + calculators: map[CountStrategy]TokenCalculator{ + TotalTokens: &TotalTokensCalculator{}, + InputTokens: &InputTokensCalculator{}, + OutputTokens: &OutputTokensCalculator{}, + }} +} + +func (m *TokenCalculatorManager) getCalculator(strategy CountStrategy) TokenCalculator { + if m == nil || m.calculators == nil { + return nil + } + calculator, exists := m.calculators[strategy] + if !exists { + return nil + } + return calculator +} + +type InputTokensCalculator struct{} + +func (c *InputTokensCalculator) Calculate(ctx *Context, infos *UsedTokenInfos) int { + if c == nil || infos == nil { + return 0 + } + return infos.InputTokens +} + +type OutputTokensCalculator struct{} + +func (c *OutputTokensCalculator) Calculate(ctx *Context, infos *UsedTokenInfos) int { + if c == nil || infos == nil { + return 0 + } + return infos.OutputTokens +} + +type TotalTokensCalculator struct{} + +func (c *TotalTokensCalculator) Calculate(ctx *Context, infos *UsedTokenInfos) int { + if c == nil || infos == nil { + return 0 + } + return infos.TotalTokens +} diff --git a/core/llm_token_ratelimit/token_calculator_test.go b/core/llm_token_ratelimit/token_calculator_test.go new file mode 100644 index 000000000..82eb4a3ae --- /dev/null +++ b/core/llm_token_ratelimit/token_calculator_test.go @@ -0,0 +1,691 @@ +// Copyright 1999-2020 Alibaba Group Holding Ltd. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package llm_token_ratelimit + +import ( + "testing" +) + +// Test NewDefaultTokenCalculator function +func TestNewDefaultTokenCalculator(t *testing.T) { + // Test creating a new default token calculator + calculator := NewDefaultTokenCalculator() + + // Verify the calculator is not nil + if calculator == nil { + t.Fatal("Expected non-nil TokenCalculatorManager, got nil") + } + + // Verify the calculators map is initialized + if calculator.calculators == nil { + t.Fatal("Expected non-nil calculators map, got nil") + } + + // Verify all expected calculators are present + expectedStrategies := []CountStrategy{TotalTokens, InputTokens, OutputTokens} + for _, strategy := range expectedStrategies { + calc := calculator.getCalculator(strategy) + if calc == nil { + t.Errorf("Expected calculator for strategy %v, got nil", strategy) + } + } + + // Verify specific calculator types + if _, ok := calculator.getCalculator(TotalTokens).(*TotalTokensCalculator); !ok { + t.Error("Expected TotalTokensCalculator for TotalTokens strategy") + } + if _, ok := calculator.getCalculator(InputTokens).(*InputTokensCalculator); !ok { + t.Error("Expected InputTokensCalculator for InputTokens strategy") + } + if _, ok := calculator.getCalculator(OutputTokens).(*OutputTokensCalculator); !ok { + t.Error("Expected OutputTokensCalculator for OutputTokens strategy") + } +} + +// Test TokenCalculatorManager.getCalculator function +func TestTokenCalculatorManager_getCalculator(t *testing.T) { + tests := []struct { + name string + manager *TokenCalculatorManager + strategy CountStrategy + expectNil bool + expectType interface{} + }{ + { + name: "Get TotalTokens calculator", + manager: NewDefaultTokenCalculator(), + strategy: TotalTokens, + expectNil: false, + expectType: &TotalTokensCalculator{}, + }, + { + name: "Get InputTokens calculator", + manager: NewDefaultTokenCalculator(), + strategy: InputTokens, + expectNil: false, + expectType: &InputTokensCalculator{}, + }, + { + name: "Get OutputTokens calculator", + manager: NewDefaultTokenCalculator(), + strategy: OutputTokens, + expectNil: false, + expectType: &OutputTokensCalculator{}, + }, + { + name: "Get non-existent calculator", + manager: NewDefaultTokenCalculator(), + strategy: CountStrategy(999), // Invalid strategy + expectNil: true, + }, + { + name: "Nil manager", + manager: nil, + strategy: TotalTokens, + expectNil: true, + }, + { + name: "Manager with nil calculators map", + manager: &TokenCalculatorManager{ + calculators: nil, + }, + strategy: TotalTokens, + expectNil: true, + }, + { + name: "Manager with empty calculators map", + manager: &TokenCalculatorManager{ + calculators: make(map[CountStrategy]TokenCalculator), + }, + strategy: TotalTokens, + expectNil: true, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + calc := tt.manager.getCalculator(tt.strategy) + if tt.expectNil { + if calc != nil { + t.Errorf("Expected nil calculator, got %T", calc) + } + } else { + if calc == nil { + t.Errorf("Expected non-nil calculator, got nil") + } + // Check specific type if provided + if tt.expectType != nil { + expectedType := tt.expectType + switch expectedType.(type) { + case *TotalTokensCalculator: + if _, ok := calc.(*TotalTokensCalculator); !ok { + t.Errorf("Expected TotalTokensCalculator, got %T", calc) + } + case *InputTokensCalculator: + if _, ok := calc.(*InputTokensCalculator); !ok { + t.Errorf("Expected InputTokensCalculator, got %T", calc) + } + case *OutputTokensCalculator: + if _, ok := calc.(*OutputTokensCalculator); !ok { + t.Errorf("Expected OutputTokensCalculator, got %T", calc) + } + } + } + } + }) + } +} + +// Test InputTokensCalculator.Calculate function +func TestInputTokensCalculator_Calculate(t *testing.T) { + tests := []struct { + name string + calc *InputTokensCalculator + ctx *Context + infos *UsedTokenInfos + expected int + }{ + { + name: "Valid input tokens", + calc: &InputTokensCalculator{}, + ctx: NewContext(), + infos: &UsedTokenInfos{ + InputTokens: 100, + OutputTokens: 50, + TotalTokens: 150, + }, + expected: 100, + }, + { + name: "Zero input tokens", + calc: &InputTokensCalculator{}, + ctx: NewContext(), + infos: &UsedTokenInfos{ + InputTokens: 0, + OutputTokens: 50, + TotalTokens: 50, + }, + expected: 0, + }, + { + name: "Negative input tokens", + calc: &InputTokensCalculator{}, + ctx: NewContext(), + infos: &UsedTokenInfos{ + InputTokens: -10, + OutputTokens: 50, + TotalTokens: 40, + }, + expected: -10, + }, + { + name: "Large input tokens", + calc: &InputTokensCalculator{}, + ctx: NewContext(), + infos: &UsedTokenInfos{ + InputTokens: 999999, + OutputTokens: 1, + TotalTokens: 1000000, + }, + expected: 999999, + }, + { + name: "Nil calculator", + calc: nil, + ctx: NewContext(), + infos: &UsedTokenInfos{InputTokens: 100}, + expected: 0, + }, + { + name: "Nil token infos", + calc: &InputTokensCalculator{}, + ctx: NewContext(), + infos: nil, + expected: 0, + }, + { + name: "Nil context (should still work)", + calc: &InputTokensCalculator{}, + ctx: nil, + infos: &UsedTokenInfos{InputTokens: 100}, + expected: 100, + }, + { + name: "Both calculator and infos are nil", + calc: nil, + ctx: NewContext(), + infos: nil, + expected: 0, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + result := tt.calc.Calculate(tt.ctx, tt.infos) + if result != tt.expected { + t.Errorf("Expected %d, got %d", tt.expected, result) + } + }) + } +} + +// Test OutputTokensCalculator.Calculate function +func TestOutputTokensCalculator_Calculate(t *testing.T) { + tests := []struct { + name string + calc *OutputTokensCalculator + ctx *Context + infos *UsedTokenInfos + expected int + }{ + { + name: "Valid output tokens", + calc: &OutputTokensCalculator{}, + ctx: NewContext(), + infos: &UsedTokenInfos{ + InputTokens: 100, + OutputTokens: 75, + TotalTokens: 175, + }, + expected: 75, + }, + { + name: "Zero output tokens", + calc: &OutputTokensCalculator{}, + ctx: NewContext(), + infos: &UsedTokenInfos{ + InputTokens: 100, + OutputTokens: 0, + TotalTokens: 100, + }, + expected: 0, + }, + { + name: "Negative output tokens", + calc: &OutputTokensCalculator{}, + ctx: NewContext(), + infos: &UsedTokenInfos{ + InputTokens: 100, + OutputTokens: -25, + TotalTokens: 75, + }, + expected: -25, + }, + { + name: "Large output tokens", + calc: &OutputTokensCalculator{}, + ctx: NewContext(), + infos: &UsedTokenInfos{ + InputTokens: 1, + OutputTokens: 888888, + TotalTokens: 888889, + }, + expected: 888888, + }, + { + name: "Nil calculator", + calc: nil, + ctx: NewContext(), + infos: &UsedTokenInfos{OutputTokens: 75}, + expected: 0, + }, + { + name: "Nil token infos", + calc: &OutputTokensCalculator{}, + ctx: NewContext(), + infos: nil, + expected: 0, + }, + { + name: "Nil context (should still work)", + calc: &OutputTokensCalculator{}, + ctx: nil, + infos: &UsedTokenInfos{OutputTokens: 75}, + expected: 75, + }, + { + name: "Both calculator and infos are nil", + calc: nil, + ctx: NewContext(), + infos: nil, + expected: 0, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + result := tt.calc.Calculate(tt.ctx, tt.infos) + if result != tt.expected { + t.Errorf("Expected %d, got %d", tt.expected, result) + } + }) + } +} + +// Test TotalTokensCalculator.Calculate function +func TestTotalTokensCalculator_Calculate(t *testing.T) { + tests := []struct { + name string + calc *TotalTokensCalculator + ctx *Context + infos *UsedTokenInfos + expected int + }{ + { + name: "Valid total tokens from TotalTokens field", + calc: &TotalTokensCalculator{}, + ctx: NewContext(), + infos: &UsedTokenInfos{ + InputTokens: 100, + OutputTokens: 50, + TotalTokens: 200, // This value is used directly + }, + expected: 200, + }, + { + name: "Zero total tokens", + calc: &TotalTokensCalculator{}, + ctx: NewContext(), + infos: &UsedTokenInfos{ + InputTokens: 100, + OutputTokens: 50, + TotalTokens: 0, + }, + expected: 0, + }, + { + name: "Negative total tokens", + calc: &TotalTokensCalculator{}, + ctx: NewContext(), + infos: &UsedTokenInfos{ + InputTokens: 100, + OutputTokens: 50, + TotalTokens: -30, + }, + expected: -30, + }, + { + name: "Large total tokens", + calc: &TotalTokensCalculator{}, + ctx: NewContext(), + infos: &UsedTokenInfos{ + InputTokens: 500000, + OutputTokens: 300000, + TotalTokens: 1000000, + }, + expected: 1000000, + }, + { + name: "Total tokens field differs from sum of input and output", + calc: &TotalTokensCalculator{}, + ctx: NewContext(), + infos: &UsedTokenInfos{ + InputTokens: 100, + OutputTokens: 50, + TotalTokens: 200, // Different from 100+50=150 + }, + expected: 200, // Uses TotalTokens field directly + }, + { + name: "Nil calculator", + calc: nil, + ctx: NewContext(), + infos: &UsedTokenInfos{TotalTokens: 150}, + expected: 0, + }, + { + name: "Nil token infos", + calc: &TotalTokensCalculator{}, + ctx: NewContext(), + infos: nil, + expected: 0, + }, + { + name: "Nil context (should still work)", + calc: &TotalTokensCalculator{}, + ctx: nil, + infos: &UsedTokenInfos{ + InputTokens: 100, + OutputTokens: 50, + TotalTokens: 150, + }, + expected: 150, + }, + { + name: "Both calculator and infos are nil", + calc: nil, + ctx: NewContext(), + infos: nil, + expected: 0, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + result := tt.calc.Calculate(tt.ctx, tt.infos) + if result != tt.expected { + t.Errorf("Expected %d, got %d", tt.expected, result) + } + }) + } +} + +// Test integration with TokenCalculatorManager +func TestTokenCalculatorManager_Integration(t *testing.T) { + manager := NewDefaultTokenCalculator() + ctx := NewContext() + infos := &UsedTokenInfos{ + InputTokens: 100, + OutputTokens: 75, + TotalTokens: 200, // Different from sum to test TotalTokens behavior + } + + tests := []struct { + name string + strategy CountStrategy + expected int + }{ + { + name: "Calculate input tokens", + strategy: InputTokens, + expected: 100, + }, + { + name: "Calculate output tokens", + strategy: OutputTokens, + expected: 75, + }, + { + name: "Calculate total tokens (uses TotalTokens field)", + strategy: TotalTokens, + expected: 200, // Uses TotalTokens field, not sum + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + calc := manager.getCalculator(tt.strategy) + if calc == nil { + t.Fatalf("Failed to get calculator for strategy %v", tt.strategy) + } + + result := calc.Calculate(ctx, infos) + if result != tt.expected { + t.Errorf("Strategy %v: expected %d, got %d", tt.strategy, tt.expected, result) + } + }) + } +} + +// Test edge cases with extreme values +func TestTokenCalculators_ExtremeValues(t *testing.T) { + ctx := NewContext() + + tests := []struct { + name string + calculator TokenCalculator + infos *UsedTokenInfos + expected int + }{ + { + name: "InputTokensCalculator with max int", + calculator: &InputTokensCalculator{}, + infos: &UsedTokenInfos{ + InputTokens: 2147483647, // Max int32 + OutputTokens: 0, + TotalTokens: 2147483647, + }, + expected: 2147483647, + }, + { + name: "OutputTokensCalculator with min int", + calculator: &OutputTokensCalculator{}, + infos: &UsedTokenInfos{ + InputTokens: 0, + OutputTokens: -2147483648, // Min int32 + TotalTokens: -2147483648, + }, + expected: -2147483648, + }, + { + name: "TotalTokensCalculator with zero", + calculator: &TotalTokensCalculator{}, + infos: &UsedTokenInfos{ + InputTokens: 1000, + OutputTokens: 500, + TotalTokens: 0, // Zero total despite non-zero components + }, + expected: 0, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + result := tt.calculator.Calculate(ctx, tt.infos) + if result != tt.expected { + t.Errorf("Expected %d, got %d", tt.expected, result) + } + }) + } +} + +// Test context interactions (though current implementation doesn't use context) +func TestTokenCalculators_ContextInteraction(t *testing.T) { + tests := []struct { + name string + ctx *Context + }{ + { + name: "Empty context", + ctx: NewContext(), + }, + { + name: "Context with data", + ctx: func() *Context { + ctx := NewContext() + ctx.Set("test-key", "test-value") + return ctx + }(), + }, + { + name: "Nil context", + ctx: nil, + }, + } + + infos := &UsedTokenInfos{ + InputTokens: 100, + OutputTokens: 50, + TotalTokens: 150, + } + + calculators := []struct { + name string + calc TokenCalculator + }{ + {"InputTokensCalculator", &InputTokensCalculator{}}, + {"OutputTokensCalculator", &OutputTokensCalculator{}}, + {"TotalTokensCalculator", &TotalTokensCalculator{}}, + } + + for _, ctxTest := range tests { + for _, calcTest := range calculators { + t.Run(ctxTest.name+"_"+calcTest.name, func(t *testing.T) { + // Should not panic and should return consistent results + result := calcTest.calc.Calculate(ctxTest.ctx, infos) + + // Verify expected results based on calculator type + switch calcTest.calc.(type) { + case *InputTokensCalculator: + if result != 100 { + t.Errorf("Expected 100, got %d", result) + } + case *OutputTokensCalculator: + if result != 50 { + t.Errorf("Expected 50, got %d", result) + } + case *TotalTokensCalculator: + if result != 150 { + t.Errorf("Expected 150, got %d", result) + } + } + }) + } + } +} + +// Test global token calculator variable +func TestGlobalTokenCalculator(t *testing.T) { + // Test that global variable is initialized + if globalTokenCalculator == nil { + t.Fatal("Expected globalTokenCalculator to be initialized, got nil") + } + + // Test that global calculator functions correctly + calc := globalTokenCalculator.getCalculator(InputTokens) + if calc == nil { + t.Fatal("Expected to get InputTokensCalculator from global instance, got nil") + } + + // Test calculation with global instance + ctx := NewContext() + infos := &UsedTokenInfos{InputTokens: 42} + result := calc.Calculate(ctx, infos) + if result != 42 { + t.Errorf("Expected 42, got %d", result) + } +} + +// Benchmark tests for performance +func BenchmarkTokenCalculators(b *testing.B) { + ctx := NewContext() + infos := &UsedTokenInfos{ + InputTokens: 1000, + OutputTokens: 500, + TotalTokens: 1500, + } + + calculators := map[string]TokenCalculator{ + "InputTokens": &InputTokensCalculator{}, + "OutputTokens": &OutputTokensCalculator{}, + "TotalTokens": &TotalTokensCalculator{}, + } + + for name, calc := range calculators { + b.Run(name, func(b *testing.B) { + for i := 0; i < b.N; i++ { + calc.Calculate(ctx, infos) + } + }) + } +} + +// Benchmark TokenCalculatorManager getCalculator +func BenchmarkTokenCalculatorManager_getCalculator(b *testing.B) { + manager := NewDefaultTokenCalculator() + strategies := []CountStrategy{InputTokens, OutputTokens, TotalTokens} + + for _, strategy := range strategies { + b.Run(strategy.String(), func(b *testing.B) { + for i := 0; i < b.N; i++ { + manager.getCalculator(strategy) + } + }) + } +} + +// Benchmark complete calculation flow +func BenchmarkTokenCalculatorManager_CompleteFlow(b *testing.B) { + manager := NewDefaultTokenCalculator() + ctx := NewContext() + infos := &UsedTokenInfos{ + InputTokens: 1000, + OutputTokens: 500, + TotalTokens: 1500, + } + + strategies := []CountStrategy{InputTokens, OutputTokens, TotalTokens} + + for _, strategy := range strategies { + b.Run(strategy.String()+"_CompleteFlow", func(b *testing.B) { + for i := 0; i < b.N; i++ { + calc := manager.getCalculator(strategy) + if calc != nil { + calc.Calculate(ctx, infos) + } + } + }) + } +} diff --git a/core/llm_token_ratelimit/token_encoder.go b/core/llm_token_ratelimit/token_encoder.go new file mode 100644 index 000000000..bf03ed0a8 --- /dev/null +++ b/core/llm_token_ratelimit/token_encoder.go @@ -0,0 +1,105 @@ +// Copyright 1999-2020 Alibaba Group Holding Ltd. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package llm_token_ratelimit + +import ( + _ "embed" + "fmt" + "strings" + "sync" + + "github.com/alibaba/sentinel-golang/logging" + "github.com/pkoukk/tiktoken-go" +) + +// ================================= TokenEncoder ==================================== +var ( + tokenEncoderMap = make(map[TokenEncoding]TokenEncoder) + tokenEncoderMapRWMux = &sync.RWMutex{} +) + +type TokenEncoder interface { + CountTokens(ctx *Context, prompts []string, rule *MatchedRule) (int, error) +} + +func NewTokenEncoder(ctx *Context, encoding TokenEncoding) TokenEncoder { + var encoder TokenEncoder + switch encoding.Provider { + case OpenAIEncoderProvider: + encoder = NewOpenAIEncoder(ctx, encoding) + default: + logging.Warn("[LLMTokenRateLimit] unsupported token encoder provider, falling back to OpenAIEncoder", + "unsupported encoder prodier", encoding.Provider, + "requestID", ctx.Get(KeyRequestID), + ) + encoder = NewOpenAIEncoder(ctx, encoding) // Fallback to OpenAIEncoder for unsupported providers + } + tokenEncoderMapRWMux.Lock() + defer tokenEncoderMapRWMux.Unlock() + tokenEncoderMap[encoding] = encoder + return encoder +} + +func LookupTokenEncoder(ctx *Context, encoding TokenEncoding) TokenEncoder { + tokenEncoderMapRWMux.RLock() + defer tokenEncoderMapRWMux.RUnlock() + return tokenEncoderMap[encoding] +} + +// ================================= OpenAIEncoder ==================================== + +type OpenAIEncoder struct { + Model string + Encoder *tiktoken.Tiktoken +} + +func NewOpenAIEncoder(ctx *Context, encoding TokenEncoding) *OpenAIEncoder { + encoder, err := tiktoken.EncodingForModel(encoding.Model) + actualModel := encoding.Model + + if err != nil { + actualModel = DefaultTokenEncodingModel[OpenAIEncoderProvider] + logging.Warn("[LLMTokenRateLimit] model not supported, falling back to default model", + "unsupported model", encoding.Model, + "default model", actualModel, + "requestID", ctx.Get(KeyRequestID), + ) + encoder, _ = tiktoken.EncodingForModel(actualModel) + } + + return &OpenAIEncoder{ + Model: actualModel, + Encoder: encoder, + } +} + +func (e *OpenAIEncoder) CountTokens(ctx *Context, prompts []string, rule *MatchedRule) (int, error) { + if e == nil { + return 0, fmt.Errorf("OpenAIEncoder is nil") + } + if e.Encoder == nil { + return 0, fmt.Errorf("OpenAIEncoder's encoder is nil for model: %s", e.Model) + } + if len(prompts) == 0 { + return 0, nil // No prompts to count tokens + } + // Concatenate prompts + var builder strings.Builder + for _, prompt := range prompts { + builder.WriteString(prompt) + } + token := e.Encoder.Encode(builder.String(), nil, nil) + return len(token), nil +} diff --git a/core/llm_token_ratelimit/token_encoder_test.go b/core/llm_token_ratelimit/token_encoder_test.go new file mode 100644 index 000000000..0b20c7bab --- /dev/null +++ b/core/llm_token_ratelimit/token_encoder_test.go @@ -0,0 +1,686 @@ +// Copyright 1999-2020 Alibaba Group Holding Ltd. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package llm_token_ratelimit + +import ( + "reflect" + "strings" + "testing" +) + +// Test helper functions to create test data +func createTestTokenEncoding(provider TokenEncoderProvider, model string) TokenEncoding { + return TokenEncoding{ + Provider: provider, + Model: model, + } +} + +func createTestContextForEncoder() *Context { + ctx := NewContext() + ctx.Set(KeyRequestID, "test-request-123") + return ctx +} + +func createTestMatchedRuleForEncoder() *MatchedRule { + return &MatchedRule{ + Strategy: PETA, + LimitKey: "test-limit-key", + TimeWindow: 60, // 60 seconds + TokenSize: 1000, + CountStrategy: TotalTokens, + // PETA specific fields + Encoding: TokenEncoding{ + Provider: OpenAIEncoderProvider, + Model: "gpt-3.5-turbo", + }, + EstimatedToken: 100, + } +} + +func resetTokenEncoderGlobalState() { + // Clear the global token encoder map for clean testing + tokenEncoderMapRWMux.Lock() + defer tokenEncoderMapRWMux.Unlock() + tokenEncoderMap = make(map[TokenEncoding]TokenEncoder) +} + +// Test NewTokenEncoder function +func TestNewTokenEncoder(t *testing.T) { + defer resetTokenEncoderGlobalState() + + tests := []struct { + name string + ctx *Context + encoding TokenEncoding + description string + }{ + { + name: "OpenAI provider should create OpenAIEncoder", + ctx: createTestContextForEncoder(), + encoding: createTestTokenEncoding(OpenAIEncoderProvider, "gpt-3.5-turbo"), + description: "Should create OpenAIEncoder for OpenAI provider", + }, + { + name: "OpenAI provider with different model", + ctx: createTestContextForEncoder(), + encoding: createTestTokenEncoding(OpenAIEncoderProvider, "gpt-4"), + description: "Should create OpenAIEncoder for different OpenAI model", + }, + { + name: "Unknown provider should fallback to OpenAIEncoder", + ctx: createTestContextForEncoder(), + encoding: TokenEncoding{Provider: TokenEncoderProvider(999), Model: "unknown-model"}, + description: "Should fallback to OpenAIEncoder for unknown provider", + }, + { + name: "Empty model should use default", + ctx: createTestContextForEncoder(), + encoding: createTestTokenEncoding(OpenAIEncoderProvider, ""), + description: "Should handle empty model gracefully", + }, + { + name: "Nil context should not panic", + ctx: nil, + encoding: createTestTokenEncoding(OpenAIEncoderProvider, "gpt-3.5-turbo"), + description: "Should handle nil context gracefully", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + // Create new encoder + encoder := NewTokenEncoder(tt.ctx, tt.encoding) + + // Verify encoder is not nil + if encoder == nil { + t.Error("Expected non-nil encoder") + return + } + + if openAIEncoder, ok := encoder.(*OpenAIEncoder); ok { + if openAIEncoder.Model == "" { + t.Error("Expected OpenAIEncoder to have a model set") + } + } else { + t.Error("Expected encoder to be OpenAIEncoder type") + } + + // Verify encoder is stored in global map + storedEncoder := LookupTokenEncoder(tt.ctx, tt.encoding) + if storedEncoder == nil { + t.Error("Expected encoder to be stored in global map") + } + + if !reflect.DeepEqual(storedEncoder, encoder) { + t.Error("Expected stored encoder to be the same instance") + } + }) + } +} + +// Test LookupTokenEncoder function +func TestLookupTokenEncoder(t *testing.T) { + defer resetTokenEncoderGlobalState() + + tests := []struct { + name string + ctx *Context + encoding TokenEncoding + setupFunc func(TokenEncoding) TokenEncoder + expectNil bool + description string + }{ + { + name: "Existing encoder should be returned", + ctx: createTestContextForEncoder(), + encoding: createTestTokenEncoding(OpenAIEncoderProvider, "gpt-3.5-turbo"), + setupFunc: func(encoding TokenEncoding) TokenEncoder { + return NewTokenEncoder(createTestContextForEncoder(), encoding) + }, + expectNil: false, + description: "Should return existing encoder from map", + }, + { + name: "Non-existing encoder should return nil", + ctx: createTestContextForEncoder(), + encoding: createTestTokenEncoding(OpenAIEncoderProvider, "non-existing-model"), + setupFunc: nil, + expectNil: true, + description: "Should return nil for non-existing encoder", + }, + { + name: "Different encodings should return different encoders", + ctx: createTestContextForEncoder(), + encoding: createTestTokenEncoding(OpenAIEncoderProvider, "gpt-4"), + setupFunc: func(encoding TokenEncoding) TokenEncoder { + // Create encoder for gpt-3.5-turbo first + NewTokenEncoder(createTestContextForEncoder(), createTestTokenEncoding(OpenAIEncoderProvider, "gpt-3.5-turbo")) + // Then create for the test encoding + return NewTokenEncoder(createTestContextForEncoder(), encoding) + }, + expectNil: false, + description: "Should handle multiple different encoders", + }, + { + name: "Nil context should not affect lookup", + ctx: nil, + encoding: createTestTokenEncoding(OpenAIEncoderProvider, "gpt-3.5-turbo"), + setupFunc: func(encoding TokenEncoding) TokenEncoder { + // Create encoder for gpt-3.5-turbo first + NewTokenEncoder(createTestContextForEncoder(), createTestTokenEncoding(OpenAIEncoderProvider, "gpt-3.5-turbo")) + // Then create for the test encoding + return NewTokenEncoder(createTestContextForEncoder(), encoding) + }, + expectNil: false, + description: "Should handle nil context in lookup", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + // Setup encoder if needed + var expectedEncoder TokenEncoder + if tt.setupFunc != nil { + expectedEncoder = tt.setupFunc(tt.encoding) + } + + // Lookup encoder + encoder := LookupTokenEncoder(tt.ctx, tt.encoding) + + // Verify result + if tt.expectNil { + if encoder != nil { + t.Error("Expected nil encoder, got non-nil") + } + } else { + if encoder == nil { + t.Error("Expected non-nil encoder, got nil") + } + if expectedEncoder != nil && !reflect.DeepEqual(expectedEncoder, encoder) { + t.Error("Expected returned encoder to match the created encoder") + } + } + }) + } +} + +// Test NewOpenAIEncoder function +func TestNewOpenAIEncoder(t *testing.T) { + tests := []struct { + name string + ctx *Context + encoding TokenEncoding + expectModel string + expectNilEnc bool + description string + }{ + { + name: "Valid model should create encoder successfully", + ctx: createTestContextForEncoder(), + encoding: createTestTokenEncoding(OpenAIEncoderProvider, "gpt-3.5-turbo"), + expectModel: "gpt-3.5-turbo", + expectNilEnc: false, + description: "Should create encoder with correct model", + }, + { + name: "GPT-4 model should work", + ctx: createTestContextForEncoder(), + encoding: createTestTokenEncoding(OpenAIEncoderProvider, "gpt-4"), + expectModel: "gpt-4", + expectNilEnc: false, + description: "Should create encoder for GPT-4 model", + }, + { + name: "Invalid model should fallback to default", + ctx: createTestContextForEncoder(), + encoding: createTestTokenEncoding(OpenAIEncoderProvider, "invalid-model-12345"), + expectModel: DefaultTokenEncodingModel[OpenAIEncoderProvider], + expectNilEnc: false, + description: "Should fallback to default model for invalid model", + }, + { + name: "Empty model should fallback to default", + ctx: createTestContextForEncoder(), + encoding: createTestTokenEncoding(OpenAIEncoderProvider, ""), + expectModel: DefaultTokenEncodingModel[OpenAIEncoderProvider], + expectNilEnc: false, + description: "Should fallback to default model for empty model", + }, + { + name: "Nil context should not panic", + ctx: nil, + encoding: createTestTokenEncoding(OpenAIEncoderProvider, "gpt-3.5-turbo"), + expectModel: "gpt-3.5-turbo", + expectNilEnc: false, + description: "Should handle nil context gracefully", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + // Create encoder + encoder := NewOpenAIEncoder(tt.ctx, tt.encoding) + + // Verify encoder is not nil + if encoder == nil { + t.Error("Expected non-nil OpenAIEncoder") + return + } + + // Verify model + if encoder.Model != tt.expectModel { + t.Errorf("Expected model %s, got %s", tt.expectModel, encoder.Model) + } + + // Verify internal tiktoken encoder + if tt.expectNilEnc { + if encoder.Encoder != nil { + t.Error("Expected nil internal encoder") + } + } else { + if encoder.Encoder == nil { + t.Error("Expected non-nil internal encoder") + } + } + }) + } +} + +// Test OpenAIEncoder.CountTokens function +func TestOpenAIEncoder_CountTokens(t *testing.T) { + tests := []struct { + name string + encoder *OpenAIEncoder + ctx *Context + prompts []string + rule *MatchedRule + expectCount int + expectError bool + description string + }{ + { + name: "Valid encoder with simple prompts", + encoder: NewOpenAIEncoder(createTestContextForEncoder(), createTestTokenEncoding(OpenAIEncoderProvider, "gpt-3.5-turbo")), + ctx: createTestContextForEncoder(), + prompts: []string{"Hello", "World"}, + rule: createTestMatchedRuleForEncoder(), + expectCount: -1, // We can't predict exact count, just verify > 0 + expectError: false, + description: "Should count tokens for valid prompts", + }, + { + name: "Empty prompts should return zero", + encoder: NewOpenAIEncoder(createTestContextForEncoder(), createTestTokenEncoding(OpenAIEncoderProvider, "gpt-3.5-turbo")), + ctx: createTestContextForEncoder(), + prompts: []string{}, + rule: createTestMatchedRuleForEncoder(), + expectCount: 0, + expectError: false, + description: "Should return 0 for empty prompts", + }, + { + name: "Single prompt should work", + encoder: NewOpenAIEncoder(createTestContextForEncoder(), createTestTokenEncoding(OpenAIEncoderProvider, "gpt-3.5-turbo")), + ctx: createTestContextForEncoder(), + prompts: []string{"How are you?"}, + rule: createTestMatchedRuleForEncoder(), + expectCount: -1, // Verify > 0 + expectError: false, + description: "Should count tokens for single prompt", + }, + { + name: "Long prompt should work", + encoder: NewOpenAIEncoder(createTestContextForEncoder(), createTestTokenEncoding(OpenAIEncoderProvider, "gpt-3.5-turbo")), + ctx: createTestContextForEncoder(), + prompts: []string{"This is a very long prompt that contains many words and should result in multiple tokens being generated by the tokenizer"}, + rule: createTestMatchedRuleForEncoder(), + expectCount: -1, // Verify > 0 + expectError: false, + description: "Should count tokens for long prompt", + }, + { + name: "Multiple prompts should concatenate", + encoder: NewOpenAIEncoder(createTestContextForEncoder(), createTestTokenEncoding(OpenAIEncoderProvider, "gpt-3.5-turbo")), + ctx: createTestContextForEncoder(), + prompts: []string{"First prompt.", "Second prompt.", "Third prompt."}, + rule: createTestMatchedRuleForEncoder(), + expectCount: -1, // Verify > 0 + expectError: false, + description: "Should concatenate and count tokens for multiple prompts", + }, + { + name: "Nil encoder should return error", + encoder: nil, + ctx: createTestContextForEncoder(), + prompts: []string{"Hello"}, + rule: createTestMatchedRuleForEncoder(), + expectCount: 0, + expectError: true, + description: "Should return error for nil encoder", + }, + { + name: "Encoder with nil internal encoder should return error", + encoder: &OpenAIEncoder{ + Model: "test-model", + Encoder: nil, + }, + ctx: createTestContextForEncoder(), + prompts: []string{"Hello"}, + rule: createTestMatchedRuleForEncoder(), + expectCount: 0, + expectError: true, + description: "Should return error for encoder with nil internal encoder", + }, + { + name: "Nil context should not affect counting", + encoder: NewOpenAIEncoder(createTestContextForEncoder(), createTestTokenEncoding(OpenAIEncoderProvider, "gpt-3.5-turbo")), + ctx: nil, + prompts: []string{"Hello"}, + rule: createTestMatchedRuleForEncoder(), + expectCount: -1, // Verify > 0 + expectError: false, + description: "Should handle nil context gracefully", + }, + { + name: "Nil rule should not affect counting", + encoder: NewOpenAIEncoder(createTestContextForEncoder(), createTestTokenEncoding(OpenAIEncoderProvider, "gpt-3.5-turbo")), + ctx: createTestContextForEncoder(), + prompts: []string{"Hello"}, + rule: nil, + expectCount: -1, // Verify > 0 + expectError: false, + description: "Should handle nil rule gracefully", + }, + { + name: "Empty strings in prompts should work", + encoder: NewOpenAIEncoder(createTestContextForEncoder(), createTestTokenEncoding(OpenAIEncoderProvider, "gpt-3.5-turbo")), + ctx: createTestContextForEncoder(), + prompts: []string{"", "Hello", ""}, + rule: createTestMatchedRuleForEncoder(), + expectCount: -1, // Verify > 0 + expectError: false, + description: "Should handle empty strings in prompts", + }, + { + name: "Special characters should work", + encoder: NewOpenAIEncoder(createTestContextForEncoder(), createTestTokenEncoding(OpenAIEncoderProvider, "gpt-3.5-turbo")), + ctx: createTestContextForEncoder(), + prompts: []string{"Hello! @#$%^&*()"}, + rule: createTestMatchedRuleForEncoder(), + expectCount: -1, // Verify > 0 + expectError: false, + description: "Should handle special characters in prompts", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + // Count tokens + count, err := tt.encoder.CountTokens(tt.ctx, tt.prompts, tt.rule) + + // Verify error expectation + if tt.expectError { + if err == nil { + t.Error("Expected error, got nil") + } + if count != tt.expectCount { + t.Errorf("Expected count %d when error occurs, got %d", tt.expectCount, count) + } + return + } + + // Verify no error when not expected + if err != nil { + t.Errorf("Expected no error, got: %v", err) + return + } + + // Verify count + if tt.expectCount >= 0 { + if count != tt.expectCount { + t.Errorf("Expected count %d, got %d", tt.expectCount, count) + } + } else { + // For cases where we expect > 0 but can't predict exact count + if len(tt.prompts) > 0 && strings.Join(tt.prompts, "") != "" { + if count <= 0 { + t.Errorf("Expected positive token count for non-empty prompts, got %d", count) + } + } + } + }) + } +} + +// Test OpenAIEncoder.CountTokens with different models +func TestOpenAIEncoder_CountTokens_DifferentModels(t *testing.T) { + testPrompts := []string{"Hello, how are you today?"} + ctx := createTestContextForEncoder() + rule := createTestMatchedRuleForEncoder() + + models := []string{ + "gpt-3.5-turbo", + "gpt-4", + "text-davinci-003", // If supported + } + + for _, model := range models { + t.Run("Model_"+model, func(t *testing.T) { + encoder := NewOpenAIEncoder(ctx, createTestTokenEncoding(OpenAIEncoderProvider, model)) + if encoder == nil { + t.Error("Expected non-nil encoder") + return + } + + count, err := encoder.CountTokens(ctx, testPrompts, rule) + if err != nil { + t.Errorf("Expected no error for model %s, got: %v", model, err) + return + } + + if count <= 0 { + t.Errorf("Expected positive token count for model %s, got %d", model, count) + } + + t.Logf("Model %s token count: %d", encoder.Model, count) + }) + } +} + +// Test concurrent access to token encoder map +func TestTokenEncoder_ConcurrentAccess(t *testing.T) { + defer resetTokenEncoderGlobalState() + + const numGoroutines = 10 + const numOperations = 20 + + encoding := createTestTokenEncoding(OpenAIEncoderProvider, "gpt-3.5-turbo") + ctx := createTestContextForEncoder() + + // Channel to collect results + results := make(chan TokenEncoder, numGoroutines*numOperations) + errors := make(chan error, numGoroutines*numOperations) + + // Start multiple goroutines performing concurrent operations + for i := 0; i < numGoroutines; i++ { + go func() { + for j := 0; j < numOperations; j++ { + // Alternate between creating and looking up encoders + if j%2 == 0 { + encoder := NewTokenEncoder(ctx, encoding) + results <- encoder + } else { + encoder := LookupTokenEncoder(ctx, encoding) + results <- encoder + } + } + }() + } + + // Collect results + encoderCount := 0 + nilCount := 0 + for i := 0; i < numGoroutines*numOperations; i++ { + select { + case encoder := <-results: + if encoder != nil { + encoderCount++ + } else { + nilCount++ + } + case err := <-errors: + t.Errorf("Unexpected error in concurrent access: %v", err) + } + } + + // Verify we got some non-nil encoders + if encoderCount == 0 { + t.Error("Expected some non-nil encoders from concurrent operations") + } + + // Verify final state + finalEncoder := LookupTokenEncoder(ctx, encoding) + if finalEncoder == nil { + t.Error("Expected final lookup to return non-nil encoder") + } +} + +// Test edge cases and error conditions +func TestTokenEncoder_EdgeCases(t *testing.T) { + defer resetTokenEncoderGlobalState() + + t.Run("Multiple encodings should coexist", func(t *testing.T) { + ctx := createTestContextForEncoder() + + encoding1 := createTestTokenEncoding(OpenAIEncoderProvider, "gpt-3.5-turbo") + encoding2 := createTestTokenEncoding(OpenAIEncoderProvider, "gpt-4") + + encoder1 := NewTokenEncoder(ctx, encoding1) + encoder2 := NewTokenEncoder(ctx, encoding2) + + if encoder1 == nil || encoder2 == nil { + t.Error("Expected both encoders to be non-nil") + } + + if reflect.DeepEqual(encoder1, encoder2) { + t.Error("Expected different encoders for different encodings") + } + + // Verify both can be looked up + lookup1 := LookupTokenEncoder(ctx, encoding1) + lookup2 := LookupTokenEncoder(ctx, encoding2) + + if !reflect.DeepEqual(lookup1, encoder1) || !reflect.DeepEqual(lookup2, encoder2) { + t.Error("Expected lookups to return original encoders") + } + }) + + t.Run("TokenEncoder interface compliance", func(t *testing.T) { + ctx := createTestContextForEncoder() + encoding := createTestTokenEncoding(OpenAIEncoderProvider, "gpt-3.5-turbo") + + encoder := NewTokenEncoder(ctx, encoding) + + // Verify it implements TokenEncoder interface + var _ TokenEncoder = encoder + + // Test the interface method + count, err := encoder.CountTokens(ctx, []string{"test"}, createTestMatchedRuleForEncoder()) + if err != nil { + t.Errorf("Expected no error from interface method, got: %v", err) + } + if count <= 0 { + t.Error("Expected positive token count from interface method") + } + }) +} + +// Benchmark tests +func BenchmarkNewTokenEncoder(b *testing.B) { + ctx := createTestContextForEncoder() + encoding := createTestTokenEncoding(OpenAIEncoderProvider, "gpt-3.5-turbo") + + b.ResetTimer() + for i := 0; i < b.N; i++ { + NewTokenEncoder(ctx, encoding) + } +} + +func BenchmarkLookupTokenEncoder(b *testing.B) { + ctx := createTestContextForEncoder() + encoding := createTestTokenEncoding(OpenAIEncoderProvider, "gpt-3.5-turbo") + + // Setup encoder + NewTokenEncoder(ctx, encoding) + + b.ResetTimer() + for i := 0; i < b.N; i++ { + LookupTokenEncoder(ctx, encoding) + } +} + +func BenchmarkNewOpenAIEncoder(b *testing.B) { + ctx := createTestContextForEncoder() + encoding := createTestTokenEncoding(OpenAIEncoderProvider, "gpt-3.5-turbo") + + b.ResetTimer() + for i := 0; i < b.N; i++ { + NewOpenAIEncoder(ctx, encoding) + } +} + +func BenchmarkOpenAIEncoder_CountTokens(b *testing.B) { + ctx := createTestContextForEncoder() + encoder := NewOpenAIEncoder(ctx, createTestTokenEncoding(OpenAIEncoderProvider, "gpt-3.5-turbo")) + prompts := []string{"This is a test prompt for benchmarking token counting performance."} + rule := createTestMatchedRuleForEncoder() + + b.ResetTimer() + for i := 0; i < b.N; i++ { + encoder.CountTokens(ctx, prompts, rule) + } +} + +func BenchmarkOpenAIEncoder_CountTokens_LongPrompt(b *testing.B) { + ctx := createTestContextForEncoder() + encoder := NewOpenAIEncoder(ctx, createTestTokenEncoding(OpenAIEncoderProvider, "gpt-3.5-turbo")) + longPrompt := strings.Repeat("This is a long prompt with many repeated words for testing performance. ", 100) + prompts := []string{longPrompt} + rule := createTestMatchedRuleForEncoder() + + b.ResetTimer() + for i := 0; i < b.N; i++ { + encoder.CountTokens(ctx, prompts, rule) + } +} + +func BenchmarkOpenAIEncoder_CountTokens_MultiplePrompts(b *testing.B) { + ctx := createTestContextForEncoder() + encoder := NewOpenAIEncoder(ctx, createTestTokenEncoding(OpenAIEncoderProvider, "gpt-3.5-turbo")) + prompts := []string{ + "First prompt for testing.", + "Second prompt for testing.", + "Third prompt for testing.", + "Fourth prompt for testing.", + "Fifth prompt for testing.", + } + rule := createTestMatchedRuleForEncoder() + + b.ResetTimer() + for i := 0; i < b.N; i++ { + encoder.CountTokens(ctx, prompts, rule) + } +} diff --git a/core/llm_token_ratelimit/token_extractor.go b/core/llm_token_ratelimit/token_extractor.go new file mode 100644 index 000000000..91d59c18c --- /dev/null +++ b/core/llm_token_ratelimit/token_extractor.go @@ -0,0 +1,49 @@ +// Copyright 1999-2020 Alibaba Group Holding Ltd. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package llm_token_ratelimit + +import "fmt" + +// ================================= OpenAITokenExtractor ============================== + +func OpenAITokenExtractor(response interface{}) (*UsedTokenInfos, error) { + if response == nil { + return nil, fmt.Errorf("response is nil") + } + + resp, ok := response.(map[string]any) + if !ok { + return nil, fmt.Errorf("response is not map[string]any") + } + + inputTokens, ok := resp["prompt_tokens"].(int) + if !ok { + return nil, fmt.Errorf("prompt_tokens not found or not int") + } + outputTokens, ok := resp["completion_tokens"].(int) + if !ok { + return nil, fmt.Errorf("completion_tokens not found or not int") + } + totalTokens, ok := resp["total_tokens"].(int) + if !ok { + return nil, fmt.Errorf("total_tokens not found or not int") + } + + return GenerateUsedTokenInfos( + WithInputTokens(inputTokens), + WithOutputTokens(outputTokens), + WithTotalTokens(totalTokens), + ), nil +} diff --git a/core/llm_token_ratelimit/token_extractor_test.go b/core/llm_token_ratelimit/token_extractor_test.go new file mode 100644 index 000000000..8cc9932bc --- /dev/null +++ b/core/llm_token_ratelimit/token_extractor_test.go @@ -0,0 +1,593 @@ +// Copyright 1999-2020 Alibaba Group Holding Ltd. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package llm_token_ratelimit + +import ( + "testing" +) + +// Test OpenAITokenExtractor function with valid responses +func TestOpenAITokenExtractor_ValidResponses(t *testing.T) { + tests := []struct { + name string + response interface{} + expected *UsedTokenInfos + }{ + { + name: "Standard OpenAI response with all tokens", + response: map[string]any{ + "prompt_tokens": 100, + "completion_tokens": 50, + "total_tokens": 150, + }, + expected: &UsedTokenInfos{ + InputTokens: 100, + OutputTokens: 50, + TotalTokens: 150, + }, + }, + { + name: "Response with zero tokens", + response: map[string]any{ + "prompt_tokens": 0, + "completion_tokens": 0, + "total_tokens": 0, + }, + expected: &UsedTokenInfos{ + InputTokens: 0, + OutputTokens: 0, + TotalTokens: 0, + }, + }, + { + name: "Response with large token counts", + response: map[string]any{ + "prompt_tokens": 10000, + "completion_tokens": 5000, + "total_tokens": 15000, + }, + expected: &UsedTokenInfos{ + InputTokens: 10000, + OutputTokens: 5000, + TotalTokens: 15000, + }, + }, + { + name: "Response with mixed token values", + response: map[string]any{ + "prompt_tokens": 1, + "completion_tokens": 999, + "total_tokens": 1000, + }, + expected: &UsedTokenInfos{ + InputTokens: 1, + OutputTokens: 999, + TotalTokens: 1000, + }, + }, + { + name: "Response with additional fields (should be ignored)", + response: map[string]any{ + "prompt_tokens": 200, + "completion_tokens": 100, + "total_tokens": 300, + "model": "gpt-3.5-turbo", + "id": "chatcmpl-123", + "object": "chat.completion", + }, + expected: &UsedTokenInfos{ + InputTokens: 200, + OutputTokens: 100, + TotalTokens: 300, + }, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + result, err := OpenAITokenExtractor(tt.response) + if err != nil { + t.Errorf("Expected no error, got %v", err) + return + } + + if result == nil { + t.Fatal("Expected non-nil result, got nil") + } + + if result.InputTokens != tt.expected.InputTokens { + t.Errorf("InputTokens: expected %d, got %d", tt.expected.InputTokens, result.InputTokens) + } + if result.OutputTokens != tt.expected.OutputTokens { + t.Errorf("OutputTokens: expected %d, got %d", tt.expected.OutputTokens, result.OutputTokens) + } + if result.TotalTokens != tt.expected.TotalTokens { + t.Errorf("TotalTokens: expected %d, got %d", tt.expected.TotalTokens, result.TotalTokens) + } + }) + } +} + +// Test OpenAITokenExtractor function with nil input +func TestOpenAITokenExtractor_NilInput(t *testing.T) { + result, err := OpenAITokenExtractor(nil) + + // Should return error for nil input + if err == nil { + t.Error("Expected error for nil input, got nil") + } + + // Should return nil result + if result != nil { + t.Errorf("Expected nil result for nil input, got %v", result) + } + + // Check error message + expectedErrorMsg := "response is nil" + if err.Error() != expectedErrorMsg { + t.Errorf("Expected error message '%s', got '%s'", expectedErrorMsg, err.Error()) + } +} + +// Test OpenAITokenExtractor function with invalid response types +func TestOpenAITokenExtractor_InvalidResponseTypes(t *testing.T) { + tests := []struct { + name string + response interface{} + expectedError string + }{ + { + name: "String response", + response: "not a map", + expectedError: "response is not map[string]any", + }, + { + name: "Integer response", + response: 42, + expectedError: "response is not map[string]any", + }, + { + name: "Slice response", + response: []string{"test"}, + expectedError: "response is not map[string]any", + }, + { + name: "Boolean response", + response: true, + expectedError: "response is not map[string]any", + }, + { + name: "Float response", + response: 3.14, + expectedError: "response is not map[string]any", + }, + { + name: "Map with wrong key type", + response: map[int]any{ + 1: 100, + 2: 50, + 3: 150, + }, + expectedError: "response is not map[string]any", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + result, err := OpenAITokenExtractor(tt.response) + + // Should return error + if err == nil { + t.Error("Expected error, got nil") + } + + // Should return nil result + if result != nil { + t.Errorf("Expected nil result, got %v", result) + } + + // Check error message + if err.Error() != tt.expectedError { + t.Errorf("Expected error message '%s', got '%s'", tt.expectedError, err.Error()) + } + }) + } +} + +// Test OpenAITokenExtractor function with missing required fields +func TestOpenAITokenExtractor_MissingFields(t *testing.T) { + tests := []struct { + name string + response map[string]any + expectedError string + }{ + { + name: "Missing prompt_tokens", + response: map[string]any{ + "completion_tokens": 50, + "total_tokens": 150, + }, + expectedError: "prompt_tokens not found or not int", + }, + { + name: "Missing completion_tokens", + response: map[string]any{ + "prompt_tokens": 100, + "total_tokens": 150, + }, + expectedError: "completion_tokens not found or not int", + }, + { + name: "Missing total_tokens", + response: map[string]any{ + "prompt_tokens": 100, + "completion_tokens": 50, + }, + expectedError: "total_tokens not found or not int", + }, + { + name: "Missing all token fields", + response: map[string]any{}, + expectedError: "prompt_tokens not found or not int", + }, + { + name: "Empty map", + response: map[string]any{ + "other_field": "value", + }, + expectedError: "prompt_tokens not found or not int", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + result, err := OpenAITokenExtractor(tt.response) + + // Should return error + if err == nil { + t.Error("Expected error, got nil") + } + + // Should return nil result + if result != nil { + t.Errorf("Expected nil result, got %v", result) + } + + // Check error message + if err.Error() != tt.expectedError { + t.Errorf("Expected error message '%s', got '%s'", tt.expectedError, err.Error()) + } + }) + } +} + +// Test OpenAITokenExtractor function with invalid token value types +func TestOpenAITokenExtractor_InvalidTokenTypes(t *testing.T) { + tests := []struct { + name string + response map[string]any + expectedError string + }{ + { + name: "prompt_tokens as string", + response: map[string]any{ + "prompt_tokens": "100", + "completion_tokens": 50, + "total_tokens": 150, + }, + expectedError: "prompt_tokens not found or not int", + }, + { + name: "completion_tokens as float", + response: map[string]any{ + "prompt_tokens": 100, + "completion_tokens": 50.5, + "total_tokens": 150, + }, + expectedError: "completion_tokens not found or not int", + }, + { + name: "total_tokens as boolean", + response: map[string]any{ + "prompt_tokens": 100, + "completion_tokens": 50, + "total_tokens": true, + }, + expectedError: "total_tokens not found or not int", + }, + { + name: "prompt_tokens as nil", + response: map[string]any{ + "prompt_tokens": nil, + "completion_tokens": 50, + "total_tokens": 150, + }, + expectedError: "prompt_tokens not found or not int", + }, + { + name: "completion_tokens as slice", + response: map[string]any{ + "prompt_tokens": 100, + "completion_tokens": []int{50}, + "total_tokens": 150, + }, + expectedError: "completion_tokens not found or not int", + }, + { + name: "total_tokens as map", + response: map[string]any{ + "prompt_tokens": 100, + "completion_tokens": 50, + "total_tokens": map[string]int{"value": 150}, + }, + expectedError: "total_tokens not found or not int", + }, + { + name: "All tokens as wrong types", + response: map[string]any{ + "prompt_tokens": "100", + "completion_tokens": 50.0, + "total_tokens": false, + }, + expectedError: "prompt_tokens not found or not int", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + result, err := OpenAITokenExtractor(tt.response) + + // Should return error + if err == nil { + t.Error("Expected error, got nil") + } + + // Should return nil result + if result != nil { + t.Errorf("Expected nil result, got %v", result) + } + + // Check error message + if err.Error() != tt.expectedError { + t.Errorf("Expected error message '%s', got '%s'", tt.expectedError, err.Error()) + } + }) + } +} + +// Test OpenAITokenExtractor function with negative token values +func TestOpenAITokenExtractor_NegativeTokenValues(t *testing.T) { + tests := []struct { + name string + response map[string]any + expected *UsedTokenInfos + }{ + { + name: "Negative prompt_tokens", + response: map[string]any{ + "prompt_tokens": -10, + "completion_tokens": 50, + "total_tokens": 40, + }, + expected: &UsedTokenInfos{ + InputTokens: -10, + OutputTokens: 50, + TotalTokens: 40, + }, + }, + { + name: "Negative completion_tokens", + response: map[string]any{ + "prompt_tokens": 100, + "completion_tokens": -25, + "total_tokens": 75, + }, + expected: &UsedTokenInfos{ + InputTokens: 100, + OutputTokens: -25, + TotalTokens: 75, + }, + }, + { + name: "Negative total_tokens", + response: map[string]any{ + "prompt_tokens": 100, + "completion_tokens": 50, + "total_tokens": -10, + }, + expected: &UsedTokenInfos{ + InputTokens: 100, + OutputTokens: 50, + TotalTokens: -10, + }, + }, + { + name: "All negative values", + response: map[string]any{ + "prompt_tokens": -100, + "completion_tokens": -50, + "total_tokens": -150, + }, + expected: &UsedTokenInfos{ + InputTokens: -100, + OutputTokens: -50, + TotalTokens: -150, + }, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + result, err := OpenAITokenExtractor(tt.response) + + // Should not return error for negative values + if err != nil { + t.Errorf("Expected no error for negative values, got %v", err) + return + } + + if result == nil { + t.Fatal("Expected non-nil result, got nil") + } + + if result.InputTokens != tt.expected.InputTokens { + t.Errorf("InputTokens: expected %d, got %d", tt.expected.InputTokens, result.InputTokens) + } + if result.OutputTokens != tt.expected.OutputTokens { + t.Errorf("OutputTokens: expected %d, got %d", tt.expected.OutputTokens, result.OutputTokens) + } + if result.TotalTokens != tt.expected.TotalTokens { + t.Errorf("TotalTokens: expected %d, got %d", tt.expected.TotalTokens, result.TotalTokens) + } + }) + } +} + +// Test OpenAITokenExtractor function with extreme integer values +func TestOpenAITokenExtractor_ExtremeValues(t *testing.T) { + tests := []struct { + name string + response map[string]any + expected *UsedTokenInfos + }{ + { + name: "Maximum integer values", + response: map[string]any{ + "prompt_tokens": 2147483647, // max int32 + "completion_tokens": 2147483647, + "total_tokens": 2147483647, + }, + expected: &UsedTokenInfos{ + InputTokens: 2147483647, + OutputTokens: 2147483647, + TotalTokens: 2147483647, + }, + }, + { + name: "Minimum integer values", + response: map[string]any{ + "prompt_tokens": -2147483648, // min int32 + "completion_tokens": -2147483648, + "total_tokens": -2147483648, + }, + expected: &UsedTokenInfos{ + InputTokens: -2147483648, + OutputTokens: -2147483648, + TotalTokens: -2147483648, + }, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + result, err := OpenAITokenExtractor(tt.response) + + if err != nil { + t.Errorf("Expected no error for extreme values, got %v", err) + return + } + + if result == nil { + t.Fatal("Expected non-nil result, got nil") + } + + if result.InputTokens != tt.expected.InputTokens { + t.Errorf("InputTokens: expected %d, got %d", tt.expected.InputTokens, result.InputTokens) + } + if result.OutputTokens != tt.expected.OutputTokens { + t.Errorf("OutputTokens: expected %d, got %d", tt.expected.OutputTokens, result.OutputTokens) + } + if result.TotalTokens != tt.expected.TotalTokens { + t.Errorf("TotalTokens: expected %d, got %d", tt.expected.TotalTokens, result.TotalTokens) + } + }) + } +} + +// Test OpenAITokenExtractor integration with GenerateUsedTokenInfos +func TestOpenAITokenExtractor_Integration(t *testing.T) { + response := map[string]any{ + "prompt_tokens": 123, + "completion_tokens": 456, + "total_tokens": 579, + } + + result, err := OpenAITokenExtractor(response) + if err != nil { + t.Fatalf("Expected no error, got %v", err) + } + + // Verify the result uses GenerateUsedTokenInfos correctly + expected := GenerateUsedTokenInfos( + WithInputTokens(123), + WithOutputTokens(456), + WithTotalTokens(579), + ) + + if result.InputTokens != expected.InputTokens { + t.Errorf("InputTokens: expected %d, got %d", expected.InputTokens, result.InputTokens) + } + if result.OutputTokens != expected.OutputTokens { + t.Errorf("OutputTokens: expected %d, got %d", expected.OutputTokens, result.OutputTokens) + } + if result.TotalTokens != expected.TotalTokens { + t.Errorf("TotalTokens: expected %d, got %d", expected.TotalTokens, result.TotalTokens) + } +} + +// Benchmark OpenAITokenExtractor performance +func BenchmarkOpenAITokenExtractor(b *testing.B) { + response := map[string]any{ + "prompt_tokens": 1000, + "completion_tokens": 500, + "total_tokens": 1500, + "model": "gpt-3.5-turbo", + "id": "chatcmpl-123", + "object": "chat.completion", + } + + b.ResetTimer() + for i := 0; i < b.N; i++ { + _, err := OpenAITokenExtractor(response) + if err != nil { + b.Fatalf("Unexpected error in benchmark: %v", err) + } + } +} + +// Benchmark OpenAITokenExtractor with error cases +func BenchmarkOpenAITokenExtractor_ErrorCases(b *testing.B) { + errorCases := []struct { + name string + response interface{} + }{ + {"nil_response", nil}, + {"invalid_type", "not a map"}, + {"missing_fields", map[string]any{"other": "value"}}, + {"wrong_field_type", map[string]any{ + "prompt_tokens": "100", + "completion_tokens": 50, + "total_tokens": 150, + }}, + } + + for _, ec := range errorCases { + b.Run(ec.name, func(b *testing.B) { + for i := 0; i < b.N; i++ { + _, _ = OpenAITokenExtractor(ec.response) + } + }) + } +} diff --git a/core/llm_token_ratelimit/token_info.go b/core/llm_token_ratelimit/token_info.go new file mode 100644 index 000000000..a09cb3c2d --- /dev/null +++ b/core/llm_token_ratelimit/token_info.go @@ -0,0 +1,68 @@ +// Copyright 1999-2020 Alibaba Group Holding Ltd. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package llm_token_ratelimit + +type UsedTokenInfos struct { + InputTokens int `json:"inputTokens"` + OutputTokens int `json:"outputTokens"` + TotalTokens int `json:"totalTokens"` +} + +type UsedTokenInfo func(*UsedTokenInfos) + +func WithInputTokens(inputTokens int) UsedTokenInfo { + return func(infos *UsedTokenInfos) { + infos.InputTokens = inputTokens + } +} + +func WithOutputTokens(outputTokens int) UsedTokenInfo { + return func(infos *UsedTokenInfos) { + infos.OutputTokens = outputTokens + } +} + +func WithTotalTokens(totalTokens int) UsedTokenInfo { + return func(infos *UsedTokenInfos) { + infos.TotalTokens = totalTokens + } +} + +func GenerateUsedTokenInfos(uti ...UsedTokenInfo) *UsedTokenInfos { + infos := new(UsedTokenInfos) + for _, info := range uti { + if info != nil { + info(infos) + } + } + return infos +} + +func extractUsedTokenInfos(ctx *Context) *UsedTokenInfos { + if ctx == nil { + return nil + } + + usedTokenInfosRaw := ctx.Get(KeyUsedTokenInfos) + if usedTokenInfosRaw == nil { + return nil + } + + if usedTokenInfos, ok := usedTokenInfosRaw.(*UsedTokenInfos); ok && usedTokenInfos != nil { + return usedTokenInfos + } + + return nil +} diff --git a/core/llm_token_ratelimit/token_info_test.go b/core/llm_token_ratelimit/token_info_test.go new file mode 100644 index 000000000..ba4ccafa3 --- /dev/null +++ b/core/llm_token_ratelimit/token_info_test.go @@ -0,0 +1,648 @@ +// Copyright 1999-2020 Alibaba Group Holding Ltd. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package llm_token_ratelimit + +import ( + "testing" +) + +// Test UsedTokenInfos struct initialization +func TestUsedTokenInfos_Initialization(t *testing.T) { + // Test zero value initialization + var infos UsedTokenInfos + if infos.InputTokens != 0 { + t.Errorf("Expected InputTokens to be 0, got %d", infos.InputTokens) + } + if infos.OutputTokens != 0 { + t.Errorf("Expected OutputTokens to be 0, got %d", infos.OutputTokens) + } + if infos.TotalTokens != 0 { + t.Errorf("Expected TotalTokens to be 0, got %d", infos.TotalTokens) + } + + // Test struct literal initialization + infos = UsedTokenInfos{ + InputTokens: 100, + OutputTokens: 50, + TotalTokens: 150, + } + if infos.InputTokens != 100 { + t.Errorf("Expected InputTokens to be 100, got %d", infos.InputTokens) + } + if infos.OutputTokens != 50 { + t.Errorf("Expected OutputTokens to be 50, got %d", infos.OutputTokens) + } + if infos.TotalTokens != 150 { + t.Errorf("Expected TotalTokens to be 150, got %d", infos.TotalTokens) + } +} + +// Test WithInputTokens function +func TestWithInputTokens(t *testing.T) { + tests := []struct { + name string + inputTokens int + }{ + { + name: "Positive input tokens", + inputTokens: 100, + }, + { + name: "Zero input tokens", + inputTokens: 0, + }, + { + name: "Negative input tokens", + inputTokens: -50, + }, + { + name: "Large input tokens", + inputTokens: 999999, + }, + { + name: "Maximum integer value", + inputTokens: 2147483647, + }, + { + name: "Minimum integer value", + inputTokens: -2147483648, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + // Create the function + tokenInfoFunc := WithInputTokens(tt.inputTokens) + + // Verify the function is not nil + if tokenInfoFunc == nil { + t.Fatal("Expected non-nil UsedTokenInfo function, got nil") + } + + // Test applying the function to a UsedTokenInfos struct + infos := &UsedTokenInfos{} + tokenInfoFunc(infos) + + // Verify the InputTokens field is set correctly + if infos.InputTokens != tt.inputTokens { + t.Errorf("Expected InputTokens to be %d, got %d", tt.inputTokens, infos.InputTokens) + } + + // Verify other fields remain unchanged (zero values) + if infos.OutputTokens != 0 { + t.Errorf("Expected OutputTokens to remain 0, got %d", infos.OutputTokens) + } + if infos.TotalTokens != 0 { + t.Errorf("Expected TotalTokens to remain 0, got %d", infos.TotalTokens) + } + }) + } +} + +// Test WithOutputTokens function +func TestWithOutputTokens(t *testing.T) { + tests := []struct { + name string + outputTokens int + }{ + { + name: "Positive output tokens", + outputTokens: 75, + }, + { + name: "Zero output tokens", + outputTokens: 0, + }, + { + name: "Negative output tokens", + outputTokens: -25, + }, + { + name: "Large output tokens", + outputTokens: 888888, + }, + { + name: "Maximum integer value", + outputTokens: 2147483647, + }, + { + name: "Minimum integer value", + outputTokens: -2147483648, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + // Create the function + tokenInfoFunc := WithOutputTokens(tt.outputTokens) + + // Verify the function is not nil + if tokenInfoFunc == nil { + t.Fatal("Expected non-nil UsedTokenInfo function, got nil") + } + + // Test applying the function to a UsedTokenInfos struct + infos := &UsedTokenInfos{} + tokenInfoFunc(infos) + + // Verify the OutputTokens field is set correctly + if infos.OutputTokens != tt.outputTokens { + t.Errorf("Expected OutputTokens to be %d, got %d", tt.outputTokens, infos.OutputTokens) + } + + // Verify other fields remain unchanged (zero values) + if infos.InputTokens != 0 { + t.Errorf("Expected InputTokens to remain 0, got %d", infos.InputTokens) + } + if infos.TotalTokens != 0 { + t.Errorf("Expected TotalTokens to remain 0, got %d", infos.TotalTokens) + } + }) + } +} + +// Test WithTotalTokens function +func TestWithTotalTokens(t *testing.T) { + tests := []struct { + name string + totalTokens int + }{ + { + name: "Positive total tokens", + totalTokens: 150, + }, + { + name: "Zero total tokens", + totalTokens: 0, + }, + { + name: "Negative total tokens", + totalTokens: -100, + }, + { + name: "Large total tokens", + totalTokens: 1000000, + }, + { + name: "Maximum integer value", + totalTokens: 2147483647, + }, + { + name: "Minimum integer value", + totalTokens: -2147483648, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + // Create the function + tokenInfoFunc := WithTotalTokens(tt.totalTokens) + + // Verify the function is not nil + if tokenInfoFunc == nil { + t.Fatal("Expected non-nil UsedTokenInfo function, got nil") + } + + // Test applying the function to a UsedTokenInfos struct + infos := &UsedTokenInfos{} + tokenInfoFunc(infos) + + // Verify the TotalTokens field is set correctly + if infos.TotalTokens != tt.totalTokens { + t.Errorf("Expected TotalTokens to be %d, got %d", tt.totalTokens, infos.TotalTokens) + } + + // Verify other fields remain unchanged (zero values) + if infos.InputTokens != 0 { + t.Errorf("Expected InputTokens to remain 0, got %d", infos.InputTokens) + } + if infos.OutputTokens != 0 { + t.Errorf("Expected OutputTokens to remain 0, got %d", infos.OutputTokens) + } + }) + } +} + +// Test GenerateUsedTokenInfos function with no arguments +func TestGenerateUsedTokenInfos_NoArguments(t *testing.T) { + // Test with no arguments + infos := GenerateUsedTokenInfos() + + // Should return a valid pointer + if infos == nil { + t.Fatal("Expected non-nil UsedTokenInfos, got nil") + } + + // All fields should be zero values + if infos.InputTokens != 0 { + t.Errorf("Expected InputTokens to be 0, got %d", infos.InputTokens) + } + if infos.OutputTokens != 0 { + t.Errorf("Expected OutputTokens to be 0, got %d", infos.OutputTokens) + } + if infos.TotalTokens != 0 { + t.Errorf("Expected TotalTokens to be 0, got %d", infos.TotalTokens) + } +} + +// Test GenerateUsedTokenInfos function with single arguments +func TestGenerateUsedTokenInfos_SingleArguments(t *testing.T) { + tests := []struct { + name string + function UsedTokenInfo + expected UsedTokenInfos + }{ + { + name: "Only input tokens", + function: WithInputTokens(100), + expected: UsedTokenInfos{InputTokens: 100, OutputTokens: 0, TotalTokens: 0}, + }, + { + name: "Only output tokens", + function: WithOutputTokens(50), + expected: UsedTokenInfos{InputTokens: 0, OutputTokens: 50, TotalTokens: 0}, + }, + { + name: "Only total tokens", + function: WithTotalTokens(150), + expected: UsedTokenInfos{InputTokens: 0, OutputTokens: 0, TotalTokens: 150}, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + infos := GenerateUsedTokenInfos(tt.function) + + if infos == nil { + t.Fatal("Expected non-nil UsedTokenInfos, got nil") + } + + if infos.InputTokens != tt.expected.InputTokens { + t.Errorf("Expected InputTokens to be %d, got %d", tt.expected.InputTokens, infos.InputTokens) + } + if infos.OutputTokens != tt.expected.OutputTokens { + t.Errorf("Expected OutputTokens to be %d, got %d", tt.expected.OutputTokens, infos.OutputTokens) + } + if infos.TotalTokens != tt.expected.TotalTokens { + t.Errorf("Expected TotalTokens to be %d, got %d", tt.expected.TotalTokens, infos.TotalTokens) + } + }) + } +} + +// Test GenerateUsedTokenInfos function with multiple arguments +func TestGenerateUsedTokenInfos_MultipleArguments(t *testing.T) { + tests := []struct { + name string + functions []UsedTokenInfo + expected UsedTokenInfos + }{ + { + name: "All token types", + functions: []UsedTokenInfo{ + WithInputTokens(100), + WithOutputTokens(50), + WithTotalTokens(150), + }, + expected: UsedTokenInfos{InputTokens: 100, OutputTokens: 50, TotalTokens: 150}, + }, + { + name: "Overlapping input tokens (last one wins)", + functions: []UsedTokenInfo{ + WithInputTokens(100), + WithInputTokens(200), + WithOutputTokens(50), + }, + expected: UsedTokenInfos{InputTokens: 200, OutputTokens: 50, TotalTokens: 0}, + }, + { + name: "Mixed order", + functions: []UsedTokenInfo{ + WithTotalTokens(300), + WithInputTokens(200), + WithOutputTokens(100), + }, + expected: UsedTokenInfos{InputTokens: 200, OutputTokens: 100, TotalTokens: 300}, + }, + { + name: "Duplicate functions", + functions: []UsedTokenInfo{ + WithInputTokens(100), + WithOutputTokens(50), + WithInputTokens(150), // This should override the first one + WithTotalTokens(200), + }, + expected: UsedTokenInfos{InputTokens: 150, OutputTokens: 50, TotalTokens: 200}, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + infos := GenerateUsedTokenInfos(tt.functions...) + + if infos == nil { + t.Fatal("Expected non-nil UsedTokenInfos, got nil") + } + + if infos.InputTokens != tt.expected.InputTokens { + t.Errorf("Expected InputTokens to be %d, got %d", tt.expected.InputTokens, infos.InputTokens) + } + if infos.OutputTokens != tt.expected.OutputTokens { + t.Errorf("Expected OutputTokens to be %d, got %d", tt.expected.OutputTokens, infos.OutputTokens) + } + if infos.TotalTokens != tt.expected.TotalTokens { + t.Errorf("Expected TotalTokens to be %d, got %d", tt.expected.TotalTokens, infos.TotalTokens) + } + }) + } +} + +// Test GenerateUsedTokenInfos function with nil functions +func TestGenerateUsedTokenInfos_WithNilFunctions(t *testing.T) { + tests := []struct { + name string + functions []UsedTokenInfo + expected UsedTokenInfos + }{ + { + name: "Single nil function", + functions: []UsedTokenInfo{nil}, + expected: UsedTokenInfos{InputTokens: 0, OutputTokens: 0, TotalTokens: 0}, + }, + { + name: "Mixed nil and valid functions", + functions: []UsedTokenInfo{ + WithInputTokens(100), + nil, + WithOutputTokens(50), + nil, + WithTotalTokens(150), + }, + expected: UsedTokenInfos{InputTokens: 100, OutputTokens: 50, TotalTokens: 150}, + }, + { + name: "All nil functions", + functions: []UsedTokenInfo{nil, nil, nil}, + expected: UsedTokenInfos{InputTokens: 0, OutputTokens: 0, TotalTokens: 0}, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + infos := GenerateUsedTokenInfos(tt.functions...) + + if infos == nil { + t.Fatal("Expected non-nil UsedTokenInfos, got nil") + } + + if infos.InputTokens != tt.expected.InputTokens { + t.Errorf("Expected InputTokens to be %d, got %d", tt.expected.InputTokens, infos.InputTokens) + } + if infos.OutputTokens != tt.expected.OutputTokens { + t.Errorf("Expected OutputTokens to be %d, got %d", tt.expected.OutputTokens, infos.OutputTokens) + } + if infos.TotalTokens != tt.expected.TotalTokens { + t.Errorf("Expected TotalTokens to be %d, got %d", tt.expected.TotalTokens, infos.TotalTokens) + } + }) + } +} + +// Test extractUsedTokenInfos function with nil context +func TestExtractUsedTokenInfos_NilContext(t *testing.T) { + result := extractUsedTokenInfos(nil) + + if result != nil { + t.Errorf("Expected nil result for nil context, got %v", result) + } +} + +// Test extractUsedTokenInfos function with empty context +func TestExtractUsedTokenInfos_EmptyContext(t *testing.T) { + ctx := NewContext() + result := extractUsedTokenInfos(ctx) + + if result != nil { + t.Errorf("Expected nil result for empty context, got %v", result) + } +} + +// Test extractUsedTokenInfos function with valid token infos +func TestExtractUsedTokenInfos_ValidTokenInfos(t *testing.T) { + ctx := NewContext() + expectedInfos := &UsedTokenInfos{ + InputTokens: 100, + OutputTokens: 50, + TotalTokens: 150, + } + + // Set the token infos in context + ctx.Set(KeyUsedTokenInfos, expectedInfos) + + result := extractUsedTokenInfos(ctx) + + if result == nil { + t.Fatal("Expected non-nil result, got nil") + } + + if result.InputTokens != expectedInfos.InputTokens { + t.Errorf("Expected InputTokens to be %d, got %d", expectedInfos.InputTokens, result.InputTokens) + } + if result.OutputTokens != expectedInfos.OutputTokens { + t.Errorf("Expected OutputTokens to be %d, got %d", expectedInfos.OutputTokens, result.OutputTokens) + } + if result.TotalTokens != expectedInfos.TotalTokens { + t.Errorf("Expected TotalTokens to be %d, got %d", expectedInfos.TotalTokens, result.TotalTokens) + } +} + +// Test extractUsedTokenInfos function with invalid data types +func TestExtractUsedTokenInfos_InvalidDataTypes(t *testing.T) { + tests := []struct { + name string + value interface{} + }{ + { + name: "String value", + value: "not token infos", + }, + { + name: "Integer value", + value: 42, + }, + { + name: "Map value", + value: map[string]int{"tokens": 100}, + }, + { + name: "Slice value", + value: []int{100, 50, 150}, + }, + { + name: "Boolean value", + value: true, + }, + { + name: "Nil pointer to UsedTokenInfos", + value: (*UsedTokenInfos)(nil), + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + ctx := NewContext() + ctx.Set(KeyUsedTokenInfos, tt.value) + + result := extractUsedTokenInfos(ctx) + + if result != nil { + t.Errorf("Expected nil result for invalid data type, got %v", result) + } + }) + } +} + +// Test extractUsedTokenInfos function with different context states +func TestExtractUsedTokenInfos_ContextStates(t *testing.T) { + // Test with context that has other data but not token infos + t.Run("Context with other data", func(t *testing.T) { + ctx := NewContext() + ctx.Set("other_key", "other_value") + ctx.Set("another_key", 123) + + result := extractUsedTokenInfos(ctx) + + if result != nil { + t.Errorf("Expected nil result when token infos not present, got %v", result) + } + }) + + // Test with context that has token infos and other data + t.Run("Context with token infos and other data", func(t *testing.T) { + ctx := NewContext() + ctx.Set("other_key", "other_value") + + expectedInfos := &UsedTokenInfos{ + InputTokens: 200, + OutputTokens: 100, + TotalTokens: 300, + } + ctx.Set(KeyUsedTokenInfos, expectedInfos) + ctx.Set("another_key", 456) + + result := extractUsedTokenInfos(ctx) + + if result == nil { + t.Fatal("Expected non-nil result, got nil") + } + + if result.InputTokens != expectedInfos.InputTokens { + t.Errorf("Expected InputTokens to be %d, got %d", expectedInfos.InputTokens, result.InputTokens) + } + if result.OutputTokens != expectedInfos.OutputTokens { + t.Errorf("Expected OutputTokens to be %d, got %d", expectedInfos.OutputTokens, result.OutputTokens) + } + if result.TotalTokens != expectedInfos.TotalTokens { + t.Errorf("Expected TotalTokens to be %d, got %d", expectedInfos.TotalTokens, result.TotalTokens) + } + }) +} + +// Test integration between all functions +func TestTokenInfo_Integration(t *testing.T) { + // Create token infos using the builder pattern + infos := GenerateUsedTokenInfos( + WithInputTokens(500), + WithOutputTokens(300), + WithTotalTokens(800), + ) + + // Store in context + ctx := NewContext() + ctx.Set(KeyUsedTokenInfos, infos) + + // Extract from context + extracted := extractUsedTokenInfos(ctx) + + // Verify they match + if extracted == nil { + t.Fatal("Expected non-nil extracted infos, got nil") + } + + if extracted.InputTokens != 500 { + t.Errorf("Expected InputTokens to be 500, got %d", extracted.InputTokens) + } + if extracted.OutputTokens != 300 { + t.Errorf("Expected OutputTokens to be 300, got %d", extracted.OutputTokens) + } + if extracted.TotalTokens != 800 { + t.Errorf("Expected TotalTokens to be 800, got %d", extracted.TotalTokens) + } + + // Verify they are the same instance + if extracted != infos { + t.Error("Expected extracted infos to be the same instance as original") + } +} + +// Benchmark tests for performance +func BenchmarkWithInputTokens(b *testing.B) { + for i := 0; i < b.N; i++ { + fn := WithInputTokens(100) + infos := &UsedTokenInfos{} + fn(infos) + } +} + +func BenchmarkWithOutputTokens(b *testing.B) { + for i := 0; i < b.N; i++ { + fn := WithOutputTokens(50) + infos := &UsedTokenInfos{} + fn(infos) + } +} + +func BenchmarkWithTotalTokens(b *testing.B) { + for i := 0; i < b.N; i++ { + fn := WithTotalTokens(150) + infos := &UsedTokenInfos{} + fn(infos) + } +} + +func BenchmarkGenerateUsedTokenInfos(b *testing.B) { + for i := 0; i < b.N; i++ { + GenerateUsedTokenInfos( + WithInputTokens(100), + WithOutputTokens(50), + WithTotalTokens(150), + ) + } +} + +func BenchmarkExtractUsedTokenInfos(b *testing.B) { + ctx := NewContext() + tokenInfos := &UsedTokenInfos{ + InputTokens: 100, + OutputTokens: 50, + TotalTokens: 150, + } + ctx.Set(KeyUsedTokenInfos, tokenInfos) + + b.ResetTimer() + for i := 0; i < b.N; i++ { + extractUsedTokenInfos(ctx) + } +} diff --git a/core/llm_token_ratelimit/token_updater.go b/core/llm_token_ratelimit/token_updater.go new file mode 100644 index 000000000..4c532ec26 --- /dev/null +++ b/core/llm_token_ratelimit/token_updater.go @@ -0,0 +1,155 @@ +// Copyright 1999-2020 Alibaba Group Holding Ltd. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package llm_token_ratelimit + +import ( + _ "embed" + "errors" + "fmt" + + "github.com/alibaba/sentinel-golang/logging" + "github.com/alibaba/sentinel-golang/util" +) + +// ================================= FixedWindowUpdater ==================================== + +//go:embed script/fixed_window/update.lua +var globalFixedWindowUpdateScript string + +type FixedWindowUpdater struct{} + +func (u *FixedWindowUpdater) Update(ctx *Context, rule *MatchedRule) { + if u == nil || ctx == nil || rule == nil { + return + } + + usedTokenInfos := extractUsedTokenInfos(ctx) + if usedTokenInfos == nil { + return + } + + u.updateLimitKey(ctx, rule, usedTokenInfos) +} + +func (u *FixedWindowUpdater) updateLimitKey(ctx *Context, rule *MatchedRule, infos *UsedTokenInfos) { + if u == nil || ctx == nil || rule == nil || infos == nil { + return + } + calculator := globalTokenCalculator.getCalculator(rule.CountStrategy) + if calculator == nil { + logging.Error(errors.New("unknown strategy"), + "unknown strategy in llm_token_ratelimit.FixedWindowUpdater.updateLimitKey() when get calculator", + "strategy", rule.CountStrategy.String(), + "requestID", ctx.Get(KeyRequestID), + ) + return + } + usedToken := calculator.Calculate(ctx, infos) + keys := []string{rule.LimitKey} + args := []interface{}{rule.TokenSize, rule.TimeWindow * 1000, usedToken} + response, err := globalRedisClient.Eval(globalFixedWindowUpdateScript, keys, args...) + if err != nil { + logging.Error(err, "failed to execute redis script in llm_token_ratelimit.FixedWindowUpdater.updateLimitKey()", + "requestID", ctx.Get(KeyRequestID), + ) + return + } + result := parseRedisResponse(ctx, response) + if result == nil || len(result) != 2 { + logging.Error(errors.New("invalid redis response"), + "invalid redis response in llm_token_ratelimit.FixedWindowUpdater.updateLimitKey()", + "response", response, + "requestID", ctx.Get(KeyRequestID), + ) + return + } +} + +// ================================= PETAUpdater ==================================== + +//go:embed script/peta/correct.lua +var globalPETACorrectScript string + +type PETAUpdater struct{} + +func (u *PETAUpdater) Update(ctx *Context, rule *MatchedRule) { + if u == nil || ctx == nil || rule == nil { + return + } + + usedTokenInfos := extractUsedTokenInfos(ctx) + if usedTokenInfos == nil { + return + } + + u.updateLimitKey(ctx, rule, usedTokenInfos) +} + +func (u *PETAUpdater) updateLimitKey(ctx *Context, rule *MatchedRule, infos *UsedTokenInfos) { + if u == nil || ctx == nil || rule == nil || infos == nil { + return + } + + calculator := globalTokenCalculator.getCalculator(rule.CountStrategy) + if calculator == nil { + logging.Error(errors.New("unknown strategy"), + "unknown strategy in llm_token_ratelimit.PETAUpdater.updateLimitKey() when get calculator", + "strategy", rule.CountStrategy.String(), + "requestID", ctx.Get(KeyRequestID), + ) + return + } + actualToken := calculator.Calculate(ctx, infos) + + slidingWindowKey := fmt.Sprintf(PETASlidingWindowKeyFormat, generateHash(rule.LimitKey), rule.LimitKey) + tokenBucketKey := fmt.Sprintf(PETATokenBucketKeyFormat, generateHash(rule.LimitKey), rule.LimitKey) + tokenEncoderKey := fmt.Sprintf(TokenEncoderKeyFormat, generateHash(rule.LimitKey), rule.Encoding.Provider.String(), rule.Encoding.Model, rule.LimitKey) + + keys := []string{slidingWindowKey, tokenBucketKey, tokenEncoderKey} + args := []interface{}{rule.EstimatedToken, util.CurrentTimeMillis(), rule.TokenSize, rule.TimeWindow * 1000, actualToken, generateRandomString(PETARandomStringLength)} + response, err := globalRedisClient.Eval(globalPETACorrectScript, keys, args...) + if err != nil { + logging.Error(err, "failed to execute redis script in llm_token_ratelimit.PETAUpdater.updateLimitKey()", + "requestID", ctx.Get(KeyRequestID), + ) + return + } + result := parseRedisResponse(ctx, response) + if result == nil || len(result) != 1 { + logging.Error(errors.New("invalid redis response"), + "invalid redis response in llm_token_ratelimit.PETAUpdater.updateLimitKey()", + "response", response, + "requestID", ctx.Get(KeyRequestID), + ) + return + } + + correctResult := result[0] + if correctResult != PETACorrectOK && correctResult != PETACorrectOverestimateError { // Temporarily unable to handle overestimation cases + logging.Warn("[LLMTokenRateLimit] failed to update the limit key", + "correct_result", correctResult, + "requestID", ctx.Get(KeyRequestID), + ) + return + } + + RecordMetric(MetricItem{ + Timestamp: util.CurrentTimeMillis(), + RequestID: ctx.Get(KeyRequestID).(string), + EstimatedToken: rule.EstimatedToken, + ActualToken: actualToken, + CorrectResult: result[0], + }) +} diff --git a/core/llm_token_ratelimit/util.go b/core/llm_token_ratelimit/util.go new file mode 100644 index 000000000..afd4eafd2 --- /dev/null +++ b/core/llm_token_ratelimit/util.go @@ -0,0 +1,121 @@ +// Copyright 1999-2020 Alibaba Group Holding Ltd. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package llm_token_ratelimit + +import ( + "errors" + "fmt" + "math/rand" + "strconv" + "unsafe" + + "github.com/alibaba/sentinel-golang/logging" + "github.com/google/uuid" + "github.com/jinzhu/copier" + "github.com/spaolacci/murmur3" +) + +func generateHash(parts ...string) string { + h := murmur3.New64() + for i, part := range parts { + if i > 0 { + h.Write([]byte{0}) // separator + } + h.Write([]byte(part)) + } + return strconv.FormatUint(h.Sum64(), 16) +} + +func parseRedisResponse(ctx *Context, response interface{}) []int64 { + if response == nil { + return nil + } + + resultSlice, ok := response.([]interface{}) + if !ok || resultSlice == nil { + return nil + } + + result := make([]int64, len(resultSlice)) + for i, v := range resultSlice { + switch val := v.(type) { + case int64: + result[i] = val + case string: + num, err := strconv.ParseInt(val, 10, 64) + if err != nil { + logging.Error(err, "failed to parse redis response element in llm_token_ratelimit.parseRedisResponse()", + "index", i, + "value", val, + "error", err.Error(), + "requestID", ctx.Get(KeyRequestID), + ) + return nil + } + result[i] = num + case int: + result[i] = int64(val) + case float64: + result[i] = int64(val) + default: + logging.Error(errors.New("unknown error"), "unexpected redis response element type in llm_token_ratelimit.parseRedisResponse()", + "index", i, + "value", v, + "type", fmt.Sprintf("%T", v), + "requestID", ctx.Get(KeyRequestID), + ) + return nil + } + } + return result +} + +func generateRandomString(n int) string { + if n <= 0 { + return "" + } + + b := make([]byte, n) + + for i, cache, remain := n-1, rand.Int63(), RandomLetterIdxMax; i >= 0; { + if remain == 0 { + cache, remain = rand.Int63(), RandomLetterIdxMax + } + if idx := int(cache & RandomLetterIdxMask); idx < len(RandomLetterBytes) { + b[i] = RandomLetterBytes[idx] + i-- + } + cache >>= RandomLetterIdxBits + remain-- + } + + return *(*string)(unsafe.Pointer(&b)) +} + +func generateUUID() string { + return uuid.NewString() +} + +func deepCopyByCopier(src, dest interface{}) error { + if src == nil { + return errors.New("src is nil") + } + if dest == nil { + return errors.New("dest is nil") + } + return copier.CopyWithOption(dest, src, copier.Option{ + DeepCopy: true, + }) +} diff --git a/core/llm_token_ratelimit/util_test.go b/core/llm_token_ratelimit/util_test.go new file mode 100644 index 000000000..64b288e69 --- /dev/null +++ b/core/llm_token_ratelimit/util_test.go @@ -0,0 +1,1406 @@ +// Copyright 1999-2020 Alibaba Group Holding Ltd. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package llm_token_ratelimit + +import ( + "fmt" + "reflect" + "strings" + "sync" + "testing" + "unicode/utf8" + + "github.com/google/uuid" +) + +func TestGenerateHash(t *testing.T) { + tests := []struct { + name string + parts []string + validate func(string) bool + }{ + { + name: "empty parts", + parts: []string{}, + validate: func(hash string) bool { + return len(hash) > 0 && isHexString(hash) + }, + }, + { + name: "single part", + parts: []string{"test"}, + validate: func(hash string) bool { + return len(hash) > 0 && isHexString(hash) + }, + }, + { + name: "multiple parts", + parts: []string{"part1", "part2", "part3"}, + validate: func(hash string) bool { + return len(hash) > 0 && isHexString(hash) + }, + }, + { + name: "empty string parts", + parts: []string{"", "", ""}, + validate: func(hash string) bool { + return len(hash) > 0 && isHexString(hash) + }, + }, + { + name: "mixed empty and non-empty", + parts: []string{"", "test", "", "data"}, + validate: func(hash string) bool { + return len(hash) > 0 && isHexString(hash) + }, + }, + { + name: "special characters", + parts: []string{"test!@#$%^&*()", "unicode测试", "newline\nand\ttab"}, + validate: func(hash string) bool { + return len(hash) > 0 && isHexString(hash) + }, + }, + { + name: "long strings", + parts: []string{strings.Repeat("a", 1000), strings.Repeat("b", 2000)}, + validate: func(hash string) bool { + return len(hash) > 0 && isHexString(hash) + }, + }, + { + name: "numeric strings", + parts: []string{"123", "456.789", "-999"}, + validate: func(hash string) bool { + return len(hash) > 0 && isHexString(hash) + }, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + hash := generateHash(tt.parts...) + if !tt.validate(hash) { + t.Errorf("generateHash() = %v, validation failed", hash) + } + }) + } + + // Test consistency - same input should produce same hash + t.Run("consistency", func(t *testing.T) { + parts := []string{"test", "consistency", "check"} + hash1 := generateHash(parts...) + hash2 := generateHash(parts...) + if hash1 != hash2 { + t.Errorf("generateHash() inconsistent: %v != %v", hash1, hash2) + } + }) + + // Test uniqueness - different inputs should produce different hashes + t.Run("uniqueness", func(t *testing.T) { + hash1 := generateHash("test1") + hash2 := generateHash("test2") + if hash1 == hash2 { + t.Errorf("generateHash() not unique: %v == %v", hash1, hash2) + } + }) + + // Test order sensitivity + t.Run("order sensitivity", func(t *testing.T) { + hash1 := generateHash("a", "b") + hash2 := generateHash("b", "a") + if hash1 == hash2 { + t.Errorf("generateHash() order insensitive: %v == %v", hash1, hash2) + } + }) +} + +func TestParseRedisResponse(t *testing.T) { + ctx := NewContext() + ctx.Set(KeyRequestID, "test-request-123") + + tests := []struct { + name string + response interface{} + expected []int64 + }{ + { + name: "nil response", + response: nil, + expected: nil, + }, + { + name: "non-slice response", + response: "not a slice", + expected: nil, + }, + { + name: "empty slice", + response: []interface{}{}, + expected: []int64{}, + }, + { + name: "int64 values", + response: []interface{}{int64(1), int64(2), int64(3)}, + expected: []int64{1, 2, 3}, + }, + { + name: "string values", + response: []interface{}{"123", "456", "789"}, + expected: []int64{123, 456, 789}, + }, + { + name: "int values", + response: []interface{}{1, 2, 3}, + expected: []int64{1, 2, 3}, + }, + { + name: "float64 values", + response: []interface{}{1.0, 2.0, 3.0}, + expected: []int64{1, 2, 3}, + }, + { + name: "mixed valid types", + response: []interface{}{int64(1), "2", 3, 4.0}, + expected: []int64{1, 2, 3, 4}, + }, + { + name: "negative numbers", + response: []interface{}{-1, "-2", int64(-3), -4.0}, + expected: []int64{-1, -2, -3, -4}, + }, + { + name: "zero values", + response: []interface{}{0, "0", int64(0), 0.0}, + expected: []int64{0, 0, 0, 0}, + }, + { + name: "large numbers", + response: []interface{}{"9223372036854775807", int64(9223372036854775807)}, + expected: []int64{9223372036854775807, 9223372036854775807}, + }, + { + name: "invalid string number", + response: []interface{}{"not_a_number"}, + expected: nil, + }, + { + name: "invalid type", + response: []interface{}{[]string{"nested", "array"}}, + expected: nil, + }, + { + name: "mixed valid and invalid", + response: []interface{}{1, "invalid", 3}, + expected: nil, + }, + { + name: "float with decimals", + response: []interface{}{1.7, 2.9, 3.1}, + expected: []int64{1, 2, 3}, + }, + { + name: "string with leading/trailing spaces", + response: []interface{}{" 123 ", "456"}, + expected: nil, // strconv.ParseInt doesn't handle spaces + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + result := parseRedisResponse(ctx, tt.response) + if !reflect.DeepEqual(result, tt.expected) { + t.Errorf("parseRedisResponse() = %v, want %v", result, tt.expected) + } + }) + } +} + +func TestParseRedisResponse_NilContext(t *testing.T) { + response := []interface{}{1, 2, 3} + result := parseRedisResponse(nil, response) + expected := []int64{1, 2, 3} + if !reflect.DeepEqual(result, expected) { + t.Errorf("parseRedisResponse() with nil context = %v, want %v", result, expected) + } +} + +func TestGenerateRandomString(t *testing.T) { + tests := []struct { + name string + n int + validate func(string, int) bool + }{ + { + name: "zero length", + n: 0, + validate: func(s string, n int) bool { + return s == "" + }, + }, + { + name: "negative length", + n: -5, + validate: func(s string, n int) bool { + return s == "" + }, + }, + { + name: "length 1", + n: 1, + validate: func(s string, n int) bool { + return len(s) == 1 && isValidRandomChar(s[0]) + }, + }, + { + name: "length 10", + n: 10, + validate: func(s string, n int) bool { + return len(s) == 10 && allValidRandomChars(s) + }, + }, + { + name: "length 100", + n: 100, + validate: func(s string, n int) bool { + return len(s) == 100 && allValidRandomChars(s) + }, + }, + { + name: "length 1000", + n: 1000, + validate: func(s string, n int) bool { + return len(s) == 1000 && allValidRandomChars(s) + }, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + result := generateRandomString(tt.n) + if !tt.validate(result, tt.n) { + t.Errorf("generateRandomString(%d) = %v, validation failed", tt.n, result) + } + }) + } + + // Test uniqueness + t.Run("uniqueness", func(t *testing.T) { + const iterations = 1000 + const length = 20 + seen := make(map[string]bool) + + for i := 0; i < iterations; i++ { + s := generateRandomString(length) + if seen[s] { + t.Errorf("generateRandomString() produced duplicate: %s", s) + } + seen[s] = true + } + }) + + // Test character distribution + t.Run("character distribution", func(t *testing.T) { + const length = 10000 + s := generateRandomString(length) + charCount := make(map[byte]int) + + for i := 0; i < len(s); i++ { + charCount[s[i]]++ + } + + // Should have reasonable distribution (not perfect, but not too skewed) + if len(charCount) < 10 { // Should use more than 10 different characters + t.Errorf("generateRandomString() poor character distribution: only %d unique chars", len(charCount)) + } + }) +} + +func TestGenerateUUID(t *testing.T) { + tests := []struct { + name string + validate func(string) bool + }{ + { + name: "valid uuid format", + validate: func(s string) bool { + _, err := uuid.Parse(s) + return err == nil + }, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + result := generateUUID() + if !tt.validate(result) { + t.Errorf("generateUUID() = %v, validation failed", result) + } + }) + } + + // Test uniqueness + t.Run("uniqueness", func(t *testing.T) { + const iterations = 1000 + seen := make(map[string]bool) + + for i := 0; i < iterations; i++ { + uuid := generateUUID() + if seen[uuid] { + t.Errorf("generateUUID() produced duplicate: %s", uuid) + } + seen[uuid] = true + } + }) + + // Test format consistency + t.Run("format consistency", func(t *testing.T) { + const iterations = 100 + + for i := 0; i < iterations; i++ { + uuid := generateUUID() + if len(uuid) != 36 { // Standard UUID length + t.Errorf("generateUUID() invalid length: %d, expected 36", len(uuid)) + } + + // Check hyphen positions + if uuid[8] != '-' || uuid[13] != '-' || uuid[18] != '-' || uuid[23] != '-' { + t.Errorf("generateUUID() invalid format: %s", uuid) + } + } + }) +} + +func TestDeepCopyByCopier(t *testing.T) { + tests := []struct { + name string + src interface{} + dest interface{} + wantErr bool + verify func(interface{}, interface{}) bool + }{ + { + name: "simple struct", + src: &Rule{ + ID: "test-rule-1", + Resource: "test-resource", + Strategy: FixedWindow, + Encoding: TokenEncoding{ + Provider: OpenAIEncoderProvider, + Model: "gpt-3.5-turbo", + }, + }, + dest: &Rule{}, + wantErr: false, + verify: func(src, dest interface{}) bool { + srcRule := src.(*Rule) + destRule := dest.(*Rule) + return srcRule.ID == destRule.ID && + srcRule.Resource == destRule.Resource && + srcRule.Strategy == destRule.Strategy && + srcRule.Encoding.Provider == destRule.Encoding.Provider && + srcRule.Encoding.Model == destRule.Encoding.Model + }, + }, + { + name: "rule slice - simple rules", + src: []*Rule{ + { + ID: "rule-1", + Resource: "resource-1", + Strategy: FixedWindow, + Encoding: TokenEncoding{ + Provider: OpenAIEncoderProvider, + Model: "gpt-3.5-turbo", + }, + }, + { + ID: "rule-2", + Resource: "resource-2", + Strategy: PETA, + Encoding: TokenEncoding{ + Provider: OpenAIEncoderProvider, + Model: "gpt-4", + }, + }, + }, + dest: &[]*Rule{}, + wantErr: false, + verify: func(src, dest interface{}) bool { + srcRules := src.([]*Rule) + destRules := *(dest.(*[]*Rule)) + + if len(srcRules) != len(destRules) { + return false + } + + for i, srcRule := range srcRules { + destRule := destRules[i] + if srcRule.ID != destRule.ID || + srcRule.Resource != destRule.Resource || + srcRule.Strategy != destRule.Strategy || + srcRule.Encoding.Provider != destRule.Encoding.Provider || + srcRule.Encoding.Model != destRule.Encoding.Model { + return false + } + } + return true + }, + }, + { + name: "rule slice - complex rules with specific items", + src: []*Rule{ + { + ID: "complex-rule-1", + Resource: "llm-api", + Strategy: FixedWindow, + Encoding: TokenEncoding{ + Provider: OpenAIEncoderProvider, + Model: "gpt-3.5-turbo", + }, + SpecificItems: []*SpecificItem{ + { + Identifier: Identifier{ + Type: Header, + Value: "user-id", + }, + KeyItems: []*KeyItem{ + { + Key: "tokens-per-minute", + Token: Token{ + Number: 1000, + CountStrategy: TotalTokens, + }, + Time: Time{ + Unit: Minute, + Value: 1, + }, + }, + { + Key: "tokens-per-hour", + Token: Token{ + Number: 10000, + CountStrategy: InputTokens, + }, + Time: Time{ + Unit: Hour, + Value: 1, + }, + }, + }, + }, + { + Identifier: Identifier{ + Type: AllIdentifier, + Value: "*", + }, + KeyItems: []*KeyItem{ + { + Key: "global-limit", + Token: Token{ + Number: 50000, + CountStrategy: OutputTokens, + }, + Time: Time{ + Unit: Day, + Value: 1, + }, + }, + }, + }, + }, + }, + { + ID: "complex-rule-2", + Resource: "chat-api", + Strategy: PETA, + Encoding: TokenEncoding{ + Provider: OpenAIEncoderProvider, + Model: "gpt-4", + }, + SpecificItems: []*SpecificItem{ + { + Identifier: Identifier{ + Type: Header, + Value: "app-id", + }, + KeyItems: []*KeyItem{ + { + Key: "premium-limit", + Token: Token{ + Number: 100000, + CountStrategy: TotalTokens, + }, + Time: Time{ + Unit: Hour, + Value: 24, + }, + }, + }, + }, + }, + }, + }, + dest: &[]*Rule{}, + wantErr: false, + verify: func(src, dest interface{}) bool { + srcRules := src.([]*Rule) + destRules := *(dest.(*[]*Rule)) + + if len(srcRules) != len(destRules) { + return false + } + + for i, srcRule := range srcRules { + destRule := destRules[i] + + // Verify basic fields + if srcRule.ID != destRule.ID || + srcRule.Resource != destRule.Resource || + srcRule.Strategy != destRule.Strategy || + srcRule.Encoding.Provider != destRule.Encoding.Provider || + srcRule.Encoding.Model != destRule.Encoding.Model { + return false + } + + // Verify SpecificItems + if len(srcRule.SpecificItems) != len(destRule.SpecificItems) { + return false + } + + for j, srcItem := range srcRule.SpecificItems { + destItem := destRule.SpecificItems[j] + + // Verify Identifier + if srcItem.Identifier.Type != destItem.Identifier.Type || + srcItem.Identifier.Value != destItem.Identifier.Value { + return false + } + + // Verify KeyItems + if len(srcItem.KeyItems) != len(destItem.KeyItems) { + return false + } + + for k, srcKeyItem := range srcItem.KeyItems { + destKeyItem := destItem.KeyItems[k] + + if srcKeyItem.Key != destKeyItem.Key || + srcKeyItem.Token.Number != destKeyItem.Token.Number || + srcKeyItem.Token.CountStrategy != destKeyItem.Token.CountStrategy || + srcKeyItem.Time.Unit != destKeyItem.Time.Unit || + srcKeyItem.Time.Value != destKeyItem.Time.Value { + return false + } + } + } + } + return true + }, + }, + { + name: "empty rule slice", + src: []*Rule{}, + dest: &[]*Rule{}, + wantErr: false, + verify: func(src, dest interface{}) bool { + srcRules := src.([]*Rule) + destRules := *(dest.(*[]*Rule)) + return len(srcRules) == 0 && len(destRules) == 0 + }, + }, + { + name: "nil-rule-slice-1", + src: nil, + dest: &[]*Rule{}, + wantErr: true, + verify: nil, + }, + { + name: "nil-rule-slice-2", + src: []*Rule{nil}, + dest: &[]*Rule{}, + wantErr: false, + verify: func(src, dest interface{}) bool { + srcRules := src.([]*Rule) + destRules := *(dest.(*[]*Rule)) + return srcRules[0] == nil && destRules[0] == nil + }, + }, + { + name: "rule slice with nil elements", + src: []*Rule{ + { + ID: "rule-1", + Resource: "resource-1", + Strategy: FixedWindow, + }, + nil, + { + ID: "rule-3", + Resource: "resource-3", + Strategy: PETA, + }, + }, + dest: &[]*Rule{}, + wantErr: false, + verify: func(src, dest interface{}) bool { + srcRules := src.([]*Rule) + destRules := *(dest.(*[]*Rule)) + + if len(srcRules) != len(destRules) { + return false + } + + // Check that nil elements are preserved + if destRules[1] != nil { + return false + } + + // Check non-nil elements + if srcRules[0].ID != destRules[0].ID || + srcRules[2].ID != destRules[2].ID { + return false + } + + return true + }, + }, + { + name: "map to map", + src: map[string]interface{}{"key1": "value1", "key2": 42}, + dest: &map[string]interface{}{}, + wantErr: false, + verify: func(src, dest interface{}) bool { + return reflect.DeepEqual(src, *(dest.(*map[string]interface{}))) + }, + }, + { + name: "slice to slice", + src: []int{1, 2, 3, 4, 5}, + dest: &[]int{}, + wantErr: false, + verify: func(src, dest interface{}) bool { + return reflect.DeepEqual(src, *(dest.(*[]int))) + }, + }, + { + name: "invalid dest type", + src: "test", + dest: "not_a_pointer", + wantErr: true, + verify: nil, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + err := deepCopyByCopier(tt.src, tt.dest) + if (err != nil) != tt.wantErr { + t.Errorf("deepCopyByCopier() error = %v, wantErr %v", err, tt.wantErr) + return + } + + if !tt.wantErr && tt.verify != nil && !tt.verify(tt.src, tt.dest) { + t.Errorf("deepCopyByCopier() verification failed") + } + }) + } +} + +func TestDeepCopyByCopier_RuleSliceMemoryIsolation(t *testing.T) { + // Test that the copied slice is independent from the original + original := []*Rule{ + { + ID: "rule-1", + Resource: "resource-1", + Strategy: FixedWindow, + Encoding: TokenEncoding{ + Provider: OpenAIEncoderProvider, + Model: "gpt-3.5-turbo", + }, + SpecificItems: []*SpecificItem{ + { + Identifier: Identifier{ + Type: Header, + Value: "user-id", + }, + KeyItems: []*KeyItem{ + { + Key: "test-key", + Token: Token{ + Number: 1000, + CountStrategy: TotalTokens, + }, + Time: Time{ + Unit: Minute, + Value: 1, + }, + }, + }, + }, + }, + }, + } + + copied := &[]*Rule{} + err := deepCopyByCopier(original, copied) + if err != nil { + t.Fatalf("deepCopyByCopier() error = %v", err) + } + + // Modify the original + original[0].ID = "modified-rule-1" + original[0].Resource = "modified-resource" + original[0].SpecificItems[0].Identifier.Value = "modified-user-id" + original[0].SpecificItems[0].KeyItems[0].Key = "modified-key" + + // Check that the copy is not affected + copiedRules := *copied + if copiedRules[0].ID != "rule-1" { + t.Errorf("Copy was affected by original modification: ID = %v", copiedRules[0].ID) + } + if copiedRules[0].Resource != "resource-1" { + t.Errorf("Copy was affected by original modification: Resource = %v", copiedRules[0].Resource) + } + if copiedRules[0].SpecificItems[0].Identifier.Value != "user-id" { + t.Errorf("Copy was affected by original modification: Identifier.Value = %v", + copiedRules[0].SpecificItems[0].Identifier.Value) + } + if copiedRules[0].SpecificItems[0].KeyItems[0].Key != "test-key" { + t.Errorf("Copy was affected by original modification: KeyItem.Key = %v", + copiedRules[0].SpecificItems[0].KeyItems[0].Key) + } +} + +func TestDeepCopyByCopier_LargeRuleSlice(t *testing.T) { + // Test with a large number of rules + const numRules = 1000 + original := make([]*Rule, numRules) + + for i := 0; i < numRules; i++ { + original[i] = &Rule{ + ID: fmt.Sprintf("rule-%d", i), + Resource: fmt.Sprintf("resource-%d", i), + Strategy: Strategy(i % 2), // Alternate between FixedWindow and PETA + Encoding: TokenEncoding{ + Provider: OpenAIEncoderProvider, + Model: fmt.Sprintf("model-%d", i), + }, + SpecificItems: []*SpecificItem{ + { + Identifier: Identifier{ + Type: IdentifierType(i % 2), // Alternate between AllIdentifier and Header + Value: fmt.Sprintf("identifier-%d", i), + }, + KeyItems: []*KeyItem{ + { + Key: fmt.Sprintf("key-%d", i), + Token: Token{ + Number: int64(i * 100), + CountStrategy: CountStrategy(i % 3), // Cycle through count strategies + }, + Time: Time{ + Unit: TimeUnit(i % 4), // Cycle through time units + Value: int64(i + 1), + }, + }, + }, + }, + }, + } + } + + copied := &[]*Rule{} + err := deepCopyByCopier(original, copied) + if err != nil { + t.Fatalf("deepCopyByCopier() error = %v", err) + } + + copiedRules := *copied + if len(copiedRules) != numRules { + t.Fatalf("Expected %d rules, got %d", numRules, len(copiedRules)) + } + + // Verify all rules are correctly copied + for i := 0; i < numRules; i++ { + if copiedRules[i].ID != fmt.Sprintf("rule-%d", i) { + t.Errorf("Rule %d ID mismatch: got %v", i, copiedRules[i].ID) + } + if copiedRules[i].Resource != fmt.Sprintf("resource-%d", i) { + t.Errorf("Rule %d Resource mismatch: got %v", i, copiedRules[i].Resource) + } + } +} + +// Concurrency tests +func TestGenerateHash_Concurrency(t *testing.T) { + const numGoroutines = 100 + const numOperations = 100 + + var wg sync.WaitGroup + wg.Add(numGoroutines) + + results := make([][]string, numGoroutines) + + for i := 0; i < numGoroutines; i++ { + go func(id int) { + defer wg.Done() + results[id] = make([]string, numOperations) + + for j := 0; j < numOperations; j++ { + hash := generateHash(fmt.Sprintf("test_%d_%d", id, j)) + results[id][j] = hash + } + }(i) + } + + wg.Wait() + + // Verify all hashes are unique + seen := make(map[string]bool) + for _, goroutineResults := range results { + for _, hash := range goroutineResults { + if seen[hash] { + t.Errorf("Duplicate hash found: %s", hash) + } + seen[hash] = true + } + } +} + +func TestGenerateRandomString_Concurrency(t *testing.T) { + const numGoroutines = 50 + const numOperations = 200 + const stringLength = 20 + + var wg sync.WaitGroup + wg.Add(numGoroutines) + + results := make([][]string, numGoroutines) + + for i := 0; i < numGoroutines; i++ { + go func(id int) { + defer wg.Done() + results[id] = make([]string, numOperations) + + for j := 0; j < numOperations; j++ { + s := generateRandomString(stringLength) + results[id][j] = s + } + }(i) + } + + wg.Wait() + + // Verify all strings are unique and valid + seen := make(map[string]bool) + for _, goroutineResults := range results { + for _, s := range goroutineResults { + if len(s) != stringLength { + t.Errorf("Invalid string length: %d, expected %d", len(s), stringLength) + } + if !allValidRandomChars(s) { + t.Errorf("Invalid characters in string: %s", s) + } + if seen[s] { + t.Errorf("Duplicate string found: %s", s) + } + seen[s] = true + } + } +} + +func TestDeepCopyByCopier_Concurrency(t *testing.T) { + const numGoroutines = 50 + const numOperations = 10 + + originalRules := []*Rule{ + { + ID: "concurrent-rule", + Resource: "concurrent-resource", + Strategy: FixedWindow, + Encoding: TokenEncoding{ + Provider: OpenAIEncoderProvider, + Model: "gpt-3.5-turbo", + }, + }, + } + + var wg sync.WaitGroup + wg.Add(numGoroutines) + + for i := 0; i < numGoroutines; i++ { + go func(id int) { + defer wg.Done() + + for j := 0; j < numOperations; j++ { + copied := &[]*Rule{} + err := deepCopyByCopier(originalRules, copied) + if err != nil { + t.Errorf("Goroutine %d, operation %d: deepCopyByCopier() error = %v", id, j, err) + return + } + + copiedRules := *copied + if len(copiedRules) != 1 { + t.Errorf("Goroutine %d, operation %d: expected 1 rule, got %d", id, j, len(copiedRules)) + return + } + + if copiedRules[0].ID != "concurrent-rule" { + t.Errorf("Goroutine %d, operation %d: wrong ID: %v", id, j, copiedRules[0].ID) + } + } + }(i) + } + + wg.Wait() +} + +// Stress tests +func TestGenerateRandomString_Stress(t *testing.T) { + // Generate many strings and check for basic properties + const numStrings = 10000 + const stringLength = 50 + + strings := make([]string, numStrings) + for i := 0; i < numStrings; i++ { + strings[i] = generateRandomString(stringLength) + } + + // Check all strings are valid length + for i, s := range strings { + if len(s) != stringLength { + t.Errorf("String %d has invalid length: %d", i, len(s)) + } + if !utf8.ValidString(s) { + t.Errorf("String %d is not valid UTF-8", i) + } + } + + // Check for reasonable uniqueness (allowing some duplicates due to randomness) + unique := make(map[string]bool) + for _, s := range strings { + unique[s] = true + } + + uniqueRatio := float64(len(unique)) / float64(numStrings) + if uniqueRatio < 0.99 { // Expect at least 99% uniqueness + t.Errorf("Poor uniqueness ratio: %f", uniqueRatio) + } +} + +func TestParseRedisResponse_Stress(t *testing.T) { + ctx := NewContext() + + // Test with very large response + const size = 100000 + response := make([]interface{}, size) + for i := 0; i < size; i++ { + response[i] = int64(i) + } + + result := parseRedisResponse(ctx, response) + if len(result) != size { + t.Errorf("Expected %d results, got %d", size, len(result)) + } + + for i, val := range result { + if val != int64(i) { + t.Errorf("Unexpected value at index %d: %d", i, val) + } + } +} + +func TestDeepCopyByCopier_RuleSlice_Stress(t *testing.T) { + // Create a very large and complex rule slice + const numRules = 5000 + original := make([]*Rule, numRules) + + for i := 0; i < numRules; i++ { + original[i] = &Rule{ + ID: fmt.Sprintf("stress-rule-%d", i), + Resource: fmt.Sprintf("stress-resource-%d", i), + Strategy: Strategy(i % 2), + Encoding: TokenEncoding{ + Provider: OpenAIEncoderProvider, + Model: fmt.Sprintf("stress-model-%d", i), + }, + SpecificItems: []*SpecificItem{ + { + Identifier: Identifier{ + Type: IdentifierType(i % 2), + Value: fmt.Sprintf("stress-identifier-%d", i), + }, + KeyItems: []*KeyItem{ + { + Key: fmt.Sprintf("stress-key-%d", i), + Token: Token{ + Number: int64(i * 100), + CountStrategy: CountStrategy(i % 3), + }, + Time: Time{ + Unit: TimeUnit(i % 4), + Value: int64(i + 1), + }, + }, + }, + }, + }, + } + } + + copied := &[]*Rule{} + err := deepCopyByCopier(original, copied) + if err != nil { + t.Fatalf("deepCopyByCopier() error = %v", err) + } + + copiedRules := *copied + if len(copiedRules) != numRules { + t.Fatalf("Expected %d rules, got %d", numRules, len(copiedRules)) + } + + // Spot check some rules + checkIndices := []int{0, numRules / 4, numRules / 2, 3 * numRules / 4, numRules - 1} + for _, i := range checkIndices { + if copiedRules[i].ID != fmt.Sprintf("stress-rule-%d", i) { + t.Errorf("Rule %d ID mismatch: got %v", i, copiedRules[i].ID) + } + if copiedRules[i].Resource != fmt.Sprintf("stress-resource-%d", i) { + t.Errorf("Rule %d Resource mismatch: got %v", i, copiedRules[i].Resource) + } + } +} + +// Test edge cases for unsafe pointer usage in generateRandomString +func TestGenerateRandomString_UnsafePointer(t *testing.T) { + // Test that the unsafe pointer conversion doesn't cause issues + for i := 1; i <= 100; i++ { + s := generateRandomString(i) + if len(s) != i { + t.Errorf("Length mismatch for size %d: got %d", i, len(s)) + } + + // Verify string is properly null-terminated and readable + for j, r := range s { + if r == 0 { + t.Errorf("Unexpected null byte at position %d in string of length %d", j, i) + } + } + } +} + +// Performance tests +func BenchmarkGenerateHash_Single(b *testing.B) { + b.ResetTimer() + for i := 0; i < b.N; i++ { + generateHash("test_string") + } +} + +func BenchmarkGenerateHash_Multiple(b *testing.B) { + parts := []string{"part1", "part2", "part3", "part4"} + b.ResetTimer() + for i := 0; i < b.N; i++ { + generateHash(parts...) + } +} + +func BenchmarkGenerateHash_Large(b *testing.B) { + largePart := strings.Repeat("a", 1000) + b.ResetTimer() + for i := 0; i < b.N; i++ { + generateHash(largePart, "part2", largePart) + } +} + +func BenchmarkParseRedisResponse(b *testing.B) { + ctx := NewContext() + response := []interface{}{int64(1), "2", 3, 4.0, int64(5)} + + b.ResetTimer() + for i := 0; i < b.N; i++ { + parseRedisResponse(ctx, response) + } +} + +func BenchmarkParseRedisResponse_Large(b *testing.B) { + ctx := NewContext() + response := make([]interface{}, 1000) + for i := range response { + response[i] = int64(i) + } + + b.ResetTimer() + for i := 0; i < b.N; i++ { + parseRedisResponse(ctx, response) + } +} + +func BenchmarkGenerateRandomString_Small(b *testing.B) { + b.ResetTimer() + for i := 0; i < b.N; i++ { + generateRandomString(10) + } +} + +func BenchmarkGenerateRandomString_Medium(b *testing.B) { + b.ResetTimer() + for i := 0; i < b.N; i++ { + generateRandomString(100) + } +} + +func BenchmarkGenerateRandomString_Large(b *testing.B) { + b.ResetTimer() + for i := 0; i < b.N; i++ { + generateRandomString(1000) + } +} + +func BenchmarkGenerateUUID(b *testing.B) { + b.ResetTimer() + for i := 0; i < b.N; i++ { + generateUUID() + } +} + +func BenchmarkDeepCopyByCopier_SimpleRule(b *testing.B) { + src := &Rule{ + ID: "benchmark-rule", + Resource: "benchmark-resource", + Strategy: FixedWindow, + Encoding: TokenEncoding{ + Provider: OpenAIEncoderProvider, + Model: "gpt-3.5-turbo", + }, + } + + b.ResetTimer() + for i := 0; i < b.N; i++ { + dest := &Rule{} + deepCopyByCopier(src, dest) + } +} + +func BenchmarkDeepCopyByCopier_RuleSlice_Small(b *testing.B) { + src := []*Rule{ + { + ID: "rule-1", + Resource: "resource-1", + Strategy: FixedWindow, + }, + { + ID: "rule-2", + Resource: "resource-2", + Strategy: PETA, + }, + } + + b.ResetTimer() + for i := 0; i < b.N; i++ { + dest := &[]*Rule{} + deepCopyByCopier(src, dest) + } +} + +func BenchmarkDeepCopyByCopier_RuleSlice_Large(b *testing.B) { + src := make([]*Rule, 100) + for i := 0; i < 100; i++ { + src[i] = &Rule{ + ID: fmt.Sprintf("rule-%d", i), + Resource: fmt.Sprintf("resource-%d", i), + Strategy: Strategy(i % 2), + Encoding: TokenEncoding{ + Provider: OpenAIEncoderProvider, + Model: fmt.Sprintf("model-%d", i), + }, + } + } + + b.ResetTimer() + for i := 0; i < b.N; i++ { + dest := &[]*Rule{} + deepCopyByCopier(src, dest) + } +} + +func BenchmarkDeepCopyByCopier_ComplexRuleSlice(b *testing.B) { + src := []*Rule{ + { + ID: "complex-rule", + Resource: "complex-resource", + Strategy: FixedWindow, + Encoding: TokenEncoding{ + Provider: OpenAIEncoderProvider, + Model: "gpt-4", + }, + SpecificItems: []*SpecificItem{ + { + Identifier: Identifier{ + Type: Header, + Value: "user-id", + }, + KeyItems: []*KeyItem{ + { + Key: "tokens-per-minute", + Token: Token{ + Number: 1000, + CountStrategy: TotalTokens, + }, + Time: Time{ + Unit: Minute, + Value: 1, + }, + }, + { + Key: "tokens-per-hour", + Token: Token{ + Number: 10000, + CountStrategy: InputTokens, + }, + Time: Time{ + Unit: Hour, + Value: 1, + }, + }, + }, + }, + }, + }, + } + + b.ResetTimer() + for i := 0; i < b.N; i++ { + dest := &[]*Rule{} + deepCopyByCopier(src, dest) + } +} + +// Parallel benchmarks +func BenchmarkGenerateHash_Parallel(b *testing.B) { + b.RunParallel(func(pb *testing.PB) { + for pb.Next() { + generateHash("test_string", "part2") + } + }) +} + +func BenchmarkGenerateRandomString_Parallel(b *testing.B) { + b.RunParallel(func(pb *testing.PB) { + for pb.Next() { + generateRandomString(20) + } + }) +} + +func BenchmarkGenerateUUID_Parallel(b *testing.B) { + b.RunParallel(func(pb *testing.PB) { + for pb.Next() { + generateUUID() + } + }) +} + +func BenchmarkDeepCopyByCopier_RuleSlice_Parallel(b *testing.B) { + src := []*Rule{ + { + ID: "parallel-rule", + Resource: "parallel-resource", + Strategy: FixedWindow, + }, + } + + b.RunParallel(func(pb *testing.PB) { + for pb.Next() { + dest := &[]*Rule{} + deepCopyByCopier(src, dest) + } + }) +} + +// Memory allocation benchmarks +func BenchmarkGenerateHash_Memory(b *testing.B) { + b.ReportAllocs() + b.ResetTimer() + for i := 0; i < b.N; i++ { + generateHash("test", "memory", "benchmark") + } +} + +func BenchmarkGenerateRandomString_Memory(b *testing.B) { + b.ReportAllocs() + b.ResetTimer() + for i := 0; i < b.N; i++ { + generateRandomString(50) + } +} + +func BenchmarkParseRedisResponse_Memory(b *testing.B) { + ctx := NewContext() + response := []interface{}{int64(1), "2", 3, 4.0} + + b.ReportAllocs() + b.ResetTimer() + for i := 0; i < b.N; i++ { + parseRedisResponse(ctx, response) + } +} + +func BenchmarkDeepCopyByCopier_Memory(b *testing.B) { + src := []*Rule{ + { + ID: "memory-rule", + Resource: "memory-resource", + }, + } + + b.ReportAllocs() + b.ResetTimer() + for i := 0; i < b.N; i++ { + dest := &[]*Rule{} + deepCopyByCopier(src, dest) + } +} + +// Helper functions +func isHexString(s string) bool { + for _, c := range s { + if !((c >= '0' && c <= '9') || (c >= 'a' && c <= 'f') || (c >= 'A' && c <= 'F')) { + return false + } + } + return true +} + +func isValidRandomChar(b byte) bool { + for i := 0; i < len(RandomLetterBytes); i++ { + if b == RandomLetterBytes[i] { + return true + } + } + return false +} + +func allValidRandomChars(s string) bool { + for i := 0; i < len(s); i++ { + if !isValidRandomChar(s[i]) { + return false + } + } + return true +} diff --git a/example/llm_token_ratelimit/go.mod b/example/llm_token_ratelimit/go.mod new file mode 100644 index 000000000..2ed380602 --- /dev/null +++ b/example/llm_token_ratelimit/go.mod @@ -0,0 +1,88 @@ +module llm_token_ratelimit + +go 1.22.0 + +replace github.com/alibaba/sentinel-golang => ../../ + +require ( + github.com/alibaba/sentinel-golang v1.0.4 + github.com/cloudwego/eino v0.4.7 + github.com/cloudwego/eino-ext/components/model/openai v0.0.0-20250904121005-ad78ed3e5e49 + github.com/gin-gonic/gin v1.10.1 + github.com/tmc/langchaingo v0.1.13 +) + +require ( + github.com/bahlo/generic-list-go v0.2.0 // indirect + github.com/beorn7/perks v1.0.1 // indirect + github.com/buger/jsonparser v1.1.1 // indirect + github.com/bytedance/sonic v1.14.0 // indirect + github.com/bytedance/sonic/loader v0.3.0 // indirect + github.com/cespare/xxhash/v2 v2.2.0 // indirect + github.com/cloudwego/base64x v0.1.5 // indirect + github.com/cloudwego/eino-ext/libs/acl/openai v0.0.0-20250826113018-8c6f6358d4bb // indirect + github.com/dgryski/go-rendezvous v0.0.0-20200823014737-9f7001d12a5f // indirect + github.com/dlclark/regexp2 v1.10.0 // indirect + github.com/dustin/go-humanize v1.0.1 // indirect + github.com/eino-contrib/jsonschema v1.0.0 // indirect + github.com/evanphx/json-patch v0.5.2 // indirect + github.com/gabriel-vasile/mimetype v1.4.3 // indirect + github.com/getkin/kin-openapi v0.118.0 // indirect + github.com/gin-contrib/sse v0.1.0 // indirect + github.com/go-ole/go-ole v1.2.6 // indirect + github.com/go-openapi/jsonpointer v0.19.6 // indirect + github.com/go-openapi/swag v0.22.4 // indirect + github.com/go-playground/locales v0.14.1 // indirect + github.com/go-playground/universal-translator v0.18.1 // indirect + github.com/go-playground/validator/v10 v10.20.0 // indirect + github.com/go-redis/redis/v8 v8.11.0 // indirect + github.com/goccy/go-json v0.10.2 // indirect + github.com/golang/protobuf v1.5.4 // indirect + github.com/google/uuid v1.6.0 // indirect + github.com/goph/emperror v0.17.2 // indirect + github.com/invopop/yaml v0.1.0 // indirect + github.com/jinzhu/copier v0.4.0 // indirect + github.com/josharian/intern v1.0.0 // indirect + github.com/json-iterator/go v1.1.12 // indirect + github.com/klauspost/cpuid/v2 v2.2.7 // indirect + github.com/leodido/go-urn v1.4.0 // indirect + github.com/lufia/plan9stats v0.0.0-20211012122336-39d0f177ccd0 // indirect + github.com/mailru/easyjson v0.7.7 // indirect + github.com/mattn/go-isatty v0.0.20 // indirect + github.com/matttproud/golang_protobuf_extensions v1.0.4 // indirect + github.com/meguminnnnnnnnn/go-openai v0.0.0-20250821095446-07791bea23a0 // indirect + github.com/modern-go/concurrent v0.0.0-20180306012644-bacd9c7ef1dd // indirect + github.com/modern-go/reflect2 v1.0.2 // indirect + github.com/mohae/deepcopy v0.0.0-20170929034955-c48cc78d4826 // indirect + github.com/nikolalohinski/gonja v1.5.3 // indirect + github.com/pelletier/go-toml/v2 v2.2.2 // indirect + github.com/perimeterx/marshmallow v1.1.4 // indirect + github.com/pkg/errors v0.9.1 // indirect + github.com/pkoukk/tiktoken-go v0.1.7 // indirect + github.com/power-devops/perfstat v0.0.0-20210106213030-5aafc221ea8c // indirect + github.com/prometheus/client_golang v1.16.0 // indirect + github.com/prometheus/client_model v0.3.0 // indirect + github.com/prometheus/common v0.42.0 // indirect + github.com/prometheus/procfs v0.10.1 // indirect + github.com/shirou/gopsutil/v3 v3.23.12 // indirect + github.com/shoenig/go-m1cpu v0.1.6 // indirect + github.com/sirupsen/logrus v1.9.3 // indirect + github.com/slongfield/pyfmt v0.0.0-20220222012616-ea85ff4c361f // indirect + github.com/spaolacci/murmur3 v1.1.0 // indirect + github.com/tklauser/go-sysconf v0.3.12 // indirect + github.com/tklauser/numcpus v0.6.1 // indirect + github.com/twitchyliquid64/golang-asm v0.15.1 // indirect + github.com/ugorji/go/codec v1.2.12 // indirect + github.com/wk8/go-ordered-map/v2 v2.1.8 // indirect + github.com/yargevad/filepathx v1.0.0 // indirect + github.com/yusufpapurcu/wmi v1.2.3 // indirect + golang.org/x/arch v0.11.0 // indirect + golang.org/x/crypto v0.31.0 // indirect + golang.org/x/exp v0.0.0-20230713183714-613f0c0eb8a1 // indirect + golang.org/x/net v0.25.0 // indirect + golang.org/x/sys v0.28.0 // indirect + golang.org/x/text v0.21.0 // indirect + google.golang.org/protobuf v1.34.1 // indirect + gopkg.in/yaml.v2 v2.4.0 // indirect + gopkg.in/yaml.v3 v3.0.1 // indirect +) diff --git a/example/llm_token_ratelimit/go.sum b/example/llm_token_ratelimit/go.sum new file mode 100644 index 000000000..dad035217 --- /dev/null +++ b/example/llm_token_ratelimit/go.sum @@ -0,0 +1,337 @@ +github.com/airbrake/gobrake v3.6.1+incompatible/go.mod h1:wM4gu3Cn0W0K7GUuVWnlXZU11AGBXMILnrdOU8Kn00o= +github.com/bahlo/generic-list-go v0.2.0 h1:5sz/EEAK+ls5wF+NeqDpk5+iNdMDXrh3z3nPnH1Wvgk= +github.com/bahlo/generic-list-go v0.2.0/go.mod h1:2KvAjgMlE5NNynlg/5iLrrCCZ2+5xWbdbCW3pNTGyYg= +github.com/beorn7/perks v1.0.1 h1:VlbKKnNfV8bJzeqoa4cOKqO6bYr3WgKZxO8Z16+hsOM= +github.com/beorn7/perks v1.0.1/go.mod h1:G2ZrVWU2WbWT9wwq4/hrbKbnv/1ERSJQ0ibhJ6rlkpw= +github.com/bitly/go-simplejson v0.5.0/go.mod h1:cXHtHw4XUPsvGaxgjIAn8PhEWG9NfngEKAMDJEczWVA= +github.com/bmizerany/assert v0.0.0-20160611221934-b7ed37b82869/go.mod h1:Ekp36dRnpXw/yCqJaO+ZrUyxD+3VXMFFr56k5XYrpB4= +github.com/buger/jsonparser v1.1.1 h1:2PnMjfWD7wBILjqQbt530v576A/cAbQvEW9gGIpYMUs= +github.com/buger/jsonparser v1.1.1/go.mod h1:6RYKKt7H4d4+iWqouImQ9R2FZql3VbhNgx27UK13J/0= +github.com/bugsnag/bugsnag-go v1.4.0/go.mod h1:2oa8nejYd4cQ/b0hMIopN0lCRxU0bueqREvZLWFrtK8= +github.com/bugsnag/panicwrap v1.2.0/go.mod h1:D/8v3kj0zr8ZAKg1AQ6crr+5VwKN5eIywRkfhyM/+dE= +github.com/bytedance/mockey v1.2.14 h1:KZaFgPdiUwW+jOWFieo3Lr7INM1P+6adO3hxZhDswY8= +github.com/bytedance/mockey v1.2.14/go.mod h1:1BPHF9sol5R1ud/+0VEHGQq/+i2lN+GTsr3O2Q9IENY= +github.com/bytedance/sonic v1.14.0 h1:/OfKt8HFw0kh2rj8N0F6C/qPGRESq0BbaNZgcNXXzQQ= +github.com/bytedance/sonic v1.14.0/go.mod h1:WoEbx8WTcFJfzCe0hbmyTGrfjt8PzNEBdxlNUO24NhA= +github.com/bytedance/sonic/loader v0.1.1/go.mod h1:ncP89zfokxS5LZrJxl5z0UJcsk4M4yY2JpfqGeCtNLU= +github.com/bytedance/sonic/loader v0.3.0 h1:dskwH8edlzNMctoruo8FPTJDF3vLtDT0sXZwvZJyqeA= +github.com/bytedance/sonic/loader v0.3.0/go.mod h1:N8A3vUdtUebEY2/VQC0MyhYeKUFosQU6FxH2JmUe6VI= +github.com/certifi/gocertifi v0.0.0-20190105021004-abcd57078448/go.mod h1:GJKEexRPVJrBSOjoqN5VNOIKJ5Q3RViH6eu3puDRwx4= +github.com/cespare/xxhash/v2 v2.1.1/go.mod h1:VGX0DQ3Q6kWi7AoAeZDth3/j3BFtOZR5XLFGgcrjCOs= +github.com/cespare/xxhash/v2 v2.2.0 h1:DC2CZ1Ep5Y4k3ZQ899DldepgrayRUGE6BBZ/cd9Cj44= +github.com/cespare/xxhash/v2 v2.2.0/go.mod h1:VGX0DQ3Q6kWi7AoAeZDth3/j3BFtOZR5XLFGgcrjCOs= +github.com/cloudwego/base64x v0.1.5 h1:XPciSp1xaq2VCSt6lF0phncD4koWyULpl5bUxbfCyP4= +github.com/cloudwego/base64x v0.1.5/go.mod h1:0zlkT4Wn5C6NdauXdJRhSKRlJvmclQ1hhJgA0rcu/8w= +github.com/cloudwego/eino v0.4.7 h1:wwqsFWCuzCQuhw1dYKqHjGWULzjDjFfN9sTn/cezYV4= +github.com/cloudwego/eino v0.4.7/go.mod h1:1TDlOmwGSsbCJaWB92w9YLZi2FL0WRZoRcD4eMvqikg= +github.com/cloudwego/eino-ext/components/model/openai v0.0.0-20250904121005-ad78ed3e5e49 h1:Mb7GWb/YkjdyJz/b6XxsN5Wlgq1SHjngFGI47NAZa6Y= +github.com/cloudwego/eino-ext/components/model/openai v0.0.0-20250904121005-ad78ed3e5e49/go.mod h1:QQhCuQxuBAVWvu/YAZBhs/RsR76mUigw59Tl0kh04C8= +github.com/cloudwego/eino-ext/libs/acl/openai v0.0.0-20250826113018-8c6f6358d4bb h1:RMslzyijc3bi9EkqCulpS0hZupTl1y/wayR3+fVRN/c= +github.com/cloudwego/eino-ext/libs/acl/openai v0.0.0-20250826113018-8c6f6358d4bb/go.mod h1:fHn/6OqPPY1iLLx9wzz+MEVT5Dl9gwuZte1oLEnCoYw= +github.com/cloudwego/iasm v0.2.0/go.mod h1:8rXZaNYT2n95jn+zTI1sDr+IgcD2GVs0nlbbQPiEFhY= +github.com/creack/pty v1.1.9/go.mod h1:oKZEueFk5CKHvIhNR5MUki03XCEU+Q6VDXinZuGJ33E= +github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= +github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c= +github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= +github.com/dgryski/go-rendezvous v0.0.0-20200823014737-9f7001d12a5f h1:lO4WD4F/rVNCu3HqELle0jiPLLBs70cWOduZpkS1E78= +github.com/dgryski/go-rendezvous v0.0.0-20200823014737-9f7001d12a5f/go.mod h1:cuUVRXasLTGF7a8hSLbxyZXjz+1KgoB3wDUb6vlszIc= +github.com/dlclark/regexp2 v1.10.0 h1:+/GIL799phkJqYW+3YbOd8LCcbHzT0Pbo8zl70MHsq0= +github.com/dlclark/regexp2 v1.10.0/go.mod h1:DHkYz0B9wPfa6wondMfaivmHpzrQ3v9q8cnmRbL6yW8= +github.com/dustin/go-humanize v1.0.1 h1:GzkhY7T5VNhEkwH0PVJgjz+fX1rhBrR7pRT3mDkpeCY= +github.com/dustin/go-humanize v1.0.1/go.mod h1:Mu1zIs6XwVuF/gI1OepvI0qD18qycQx+mFykh5fBlto= +github.com/eino-contrib/jsonschema v1.0.0 h1:dXxbhGNZuI3+xNi8x3JT8AGyoXz6Pff6mRvmpjVl5Ww= +github.com/eino-contrib/jsonschema v1.0.0/go.mod h1:cpnX4SyKjWjGC7iN2EbhxaTdLqGjCi0e9DxpLYxddD4= +github.com/evanphx/json-patch v0.5.2 h1:xVCHIVMUu1wtM/VkR9jVZ45N3FhZfYMMYGorLCR8P3k= +github.com/evanphx/json-patch v0.5.2/go.mod h1:ZWS5hhDbVDyob71nXKNL0+PWn6ToqBHMikGIFbs31qQ= +github.com/fsnotify/fsnotify v1.4.7/go.mod h1:jwhsz4b93w/PPRr/qN1Yymfu8t87LnFCMoQvtojpjFo= +github.com/fsnotify/fsnotify v1.4.9 h1:hsms1Qyu0jgnwNXIxa+/V/PDsU6CfLf6CNO8H7IWoS4= +github.com/fsnotify/fsnotify v1.4.9/go.mod h1:znqG4EE+3YCdAaPaxE2ZRY/06pZUdp0tY4IgpuI1SZQ= +github.com/gabriel-vasile/mimetype v1.4.3 h1:in2uUcidCuFcDKtdcBxlR0rJ1+fsokWf+uqxgUFjbI0= +github.com/gabriel-vasile/mimetype v1.4.3/go.mod h1:d8uq/6HKRL6CGdk+aubisF/M5GcPfT7nKyLpA0lbSSk= +github.com/getkin/kin-openapi v0.118.0 h1:z43njxPmJ7TaPpMSCQb7PN0dEYno4tyBPQcrFdHoLuM= +github.com/getkin/kin-openapi v0.118.0/go.mod h1:l5e9PaFUo9fyLJCPGQeXI2ML8c3P8BHOEV2VaAVf/pc= +github.com/getsentry/raven-go v0.2.0/go.mod h1:KungGk8q33+aIAZUIVWZDr2OfAEBsO49PX4NzFV5kcQ= +github.com/gin-contrib/sse v0.1.0 h1:Y/yl/+YNO8GZSjAhjMsSuLt29uWRFHdHYUb5lYOV9qE= +github.com/gin-contrib/sse v0.1.0/go.mod h1:RHrZQHXnP2xjPF+u1gW/2HnVO7nvIa9PG3Gm+fLHvGI= +github.com/gin-gonic/gin v1.10.1 h1:T0ujvqyCSqRopADpgPgiTT63DUQVSfojyME59Ei63pQ= +github.com/gin-gonic/gin v1.10.1/go.mod h1:4PMNQiOhvDRa013RKVbsiNwoyezlm2rm0uX/T7kzp5Y= +github.com/go-check/check v0.0.0-20180628173108-788fd7840127 h1:0gkP6mzaMqkmpcJYCFOLkIBwI7xFExG03bbkOkCvUPI= +github.com/go-check/check v0.0.0-20180628173108-788fd7840127/go.mod h1:9ES+weclKsC9YodN5RgxqK/VD9HM9JsCSh7rNhMZE98= +github.com/go-ole/go-ole v1.2.6 h1:/Fpf6oFPoeFik9ty7siob0G6Ke8QvQEuVcuChpwXzpY= +github.com/go-ole/go-ole v1.2.6/go.mod h1:pprOEPIfldk/42T2oK7lQ4v4JSDwmV0As9GaiUsvbm0= +github.com/go-openapi/jsonpointer v0.19.5/go.mod h1:Pl9vOtqEWErmShwVjC8pYs9cog34VGT37dQOVbmoatg= +github.com/go-openapi/jsonpointer v0.19.6 h1:eCs3fxoIi3Wh6vtgmLTOjdhSpiqphQ+DaPn38N2ZdrE= +github.com/go-openapi/jsonpointer v0.19.6/go.mod h1:osyAmYz/mB/C3I+WsTTSgw1ONzaLJoLCyoi6/zppojs= +github.com/go-openapi/swag v0.19.5/go.mod h1:POnQmlKehdgb5mhVOsnJFsivZCEZ/vjK9gh66Z9tfKk= +github.com/go-openapi/swag v0.22.3/go.mod h1:UzaqsxGiab7freDnrUUra0MwWfN/q7tE4j+VcZ0yl14= +github.com/go-openapi/swag v0.22.4 h1:QLMzNJnMGPRNDCbySlcj1x01tzU8/9LTTL9hZZZogBU= +github.com/go-openapi/swag v0.22.4/go.mod h1:UzaqsxGiab7freDnrUUra0MwWfN/q7tE4j+VcZ0yl14= +github.com/go-playground/assert/v2 v2.2.0 h1:JvknZsQTYeFEAhQwI4qEt9cyV5ONwRHC+lYKSsYSR8s= +github.com/go-playground/assert/v2 v2.2.0/go.mod h1:VDjEfimB/XKnb+ZQfWdccd7VUvScMdVu0Titje2rxJ4= +github.com/go-playground/locales v0.14.1 h1:EWaQ/wswjilfKLTECiXz7Rh+3BjFhfDFKv/oXslEjJA= +github.com/go-playground/locales v0.14.1/go.mod h1:hxrqLVvrK65+Rwrd5Fc6F2O76J/NuW9t0sjnWqG1slY= +github.com/go-playground/universal-translator v0.18.1 h1:Bcnm0ZwsGyWbCzImXv+pAJnYK9S473LQFuzCbDbfSFY= +github.com/go-playground/universal-translator v0.18.1/go.mod h1:xekY+UJKNuX9WP91TpwSH2VMlDf28Uj24BCp08ZFTUY= +github.com/go-playground/validator/v10 v10.20.0 h1:K9ISHbSaI0lyB2eWMPJo+kOS/FBExVwjEviJTixqxL8= +github.com/go-playground/validator/v10 v10.20.0/go.mod h1:dbuPbCMFw/DrkbEynArYaCwl3amGuJotoKCe95atGMM= +github.com/go-redis/redis/v8 v8.11.0 h1:O1Td0mQ8UFChQ3N9zFQqo6kTU2cJ+/it88gDB+zg0wo= +github.com/go-redis/redis/v8 v8.11.0/go.mod h1:DLomh7y2e3ggQXQLd1YgmvIfecPJoFl7WU5SOQ/r06M= +github.com/go-test/deep v1.0.8 h1:TDsG77qcSprGbC6vTN8OuXp5g+J+b5Pcguhf7Zt61VM= +github.com/go-test/deep v1.0.8/go.mod h1:5C2ZWiW0ErCdrYzpqxLbTX7MG14M9iiw8DgHncVwcsE= +github.com/goccy/go-json v0.10.2 h1:CrxCmQqYDkv1z7lO7Wbh2HN93uovUHgrECaO5ZrCXAU= +github.com/goccy/go-json v0.10.2/go.mod h1:6MelG93GURQebXPDq3khkgXZkazVtN9CRI+MGFi0w8I= +github.com/gofrs/uuid v3.2.0+incompatible/go.mod h1:b2aQJv3Z4Fp6yNu3cdSllBxTCLRxnplIgP/c0N/04lM= +github.com/golang/protobuf v1.2.0/go.mod h1:6lQm79b+lXiMfvg/cZm0SGofjICqVBUtrP5yJMmIC1U= +github.com/golang/protobuf v1.3.5/go.mod h1:6O5/vntMXwX2lRkT1hjjk0nAC1IDOTvTlVgjlRvqsdk= +github.com/golang/protobuf v1.4.0-rc.1/go.mod h1:ceaxUfeHdC40wWswd/P6IGgMaK3YpKi5j83Wpe3EHw8= +github.com/golang/protobuf v1.4.0-rc.1.0.20200221234624-67d41d38c208/go.mod h1:xKAWHe0F5eneWXFV3EuXVDTCmh+JuBKY0li0aMyXATA= +github.com/golang/protobuf v1.4.0-rc.2/go.mod h1:LlEzMj4AhA7rCAGe4KMBDvJI+AwstrUpVNzEA03Pprs= +github.com/golang/protobuf v1.4.0-rc.4.0.20200313231945-b860323f09d0/go.mod h1:WU3c8KckQ9AFe+yFwt9sWVRKCVIyN9cPHBJSNnbL67w= +github.com/golang/protobuf v1.4.0/go.mod h1:jodUvKwWbYaEsadDk5Fwe5c77LiNKVO9IDvqG2KuDX0= +github.com/golang/protobuf v1.4.2/go.mod h1:oDoupMAO8OvCJWAcko0GGGIgR6R6ocIYbsSw735rRwI= +github.com/golang/protobuf v1.5.4 h1:i7eJL8qZTpSEXOPTxNKhASYpMn+8e5Q6AdndVa1dWek= +github.com/golang/protobuf v1.5.4/go.mod h1:lnTiLA8Wa4RWRcIUkrtSVa5nRhsEGBg48fD6rSs7xps= +github.com/google/go-cmp v0.3.0/go.mod h1:8QqcDgzrUqlUb/G2PQTWiueGozuR1884gddMywk6iLU= +github.com/google/go-cmp v0.3.1/go.mod h1:8QqcDgzrUqlUb/G2PQTWiueGozuR1884gddMywk6iLU= +github.com/google/go-cmp v0.4.0/go.mod h1:v8dTdLbMG2kIc/vJvl+f65V22dbkXbowE6jgT/gNBxE= +github.com/google/go-cmp v0.5.6/go.mod h1:v8dTdLbMG2kIc/vJvl+f65V22dbkXbowE6jgT/gNBxE= +github.com/google/go-cmp v0.5.9/go.mod h1:17dUlkBOakJ0+DkrSSNjCkIjxS6bF9zb3elmeNGIjoY= +github.com/google/go-cmp v0.6.0 h1:ofyhxvXcZhMsU5ulbFiLKl/XBFqE1GSq7atu8tAmTRI= +github.com/google/go-cmp v0.6.0/go.mod h1:17dUlkBOakJ0+DkrSSNjCkIjxS6bF9zb3elmeNGIjoY= +github.com/google/gofuzz v1.0.0/go.mod h1:dBl0BpW6vV/+mYPU4Po3pmUjxk6FQPldtuIdl/M65Eg= +github.com/google/uuid v1.6.0 h1:NIvaJDMOsjHA8n1jAhLSgzrAzy1Hgr+hNrb57e+94F0= +github.com/google/uuid v1.6.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo= +github.com/goph/emperror v0.17.2 h1:yLapQcmEsO0ipe9p5TaN22djm3OFV/TfM/fcYP0/J18= +github.com/goph/emperror v0.17.2/go.mod h1:+ZbQ+fUNO/6FNiUo0ujtMjhgad9Xa6fQL9KhH4LNHic= +github.com/gopherjs/gopherjs v1.17.2 h1:fQnZVsXk8uxXIStYb0N4bGk7jeyTalG/wsZjQ25dO0g= +github.com/gopherjs/gopherjs v1.17.2/go.mod h1:pRRIvn/QzFLrKfvEz3qUuEhtE/zLCWfreZ6J5gM2i+k= +github.com/gorilla/mux v1.8.0/go.mod h1:DVbg23sWSpFRCP0SfiEN6jmj59UnW/n46BH5rLB71So= +github.com/hpcloud/tail v1.0.0/go.mod h1:ab1qPbhIpdTxEkNHXyeSf5vhxWSCs/tWer42PpOxQnU= +github.com/invopop/yaml v0.1.0 h1:YW3WGUoJEXYfzWBjn00zIlrw7brGVD0fUKRYDPAPhrc= +github.com/invopop/yaml v0.1.0/go.mod h1:2XuRLgs/ouIrW3XNzuNj7J3Nvu/Dig5MXvbCEdiBN3Q= +github.com/jessevdk/go-flags v1.4.0/go.mod h1:4FA24M0QyGHXBuZZK/XkWh8h0e1EYbRYJSGM75WSRxI= +github.com/jinzhu/copier v0.4.0 h1:w3ciUoD19shMCRargcpm0cm91ytaBhDvuRpz1ODO/U8= +github.com/jinzhu/copier v0.4.0/go.mod h1:DfbEm0FYsaqBcKcFuvmOZb218JkPGtvSHsKg8S8hyyg= +github.com/josharian/intern v1.0.0 h1:vlS4z54oSdjm0bgjRigI+G1HpF+tI+9rE5LLzOg8HmY= +github.com/josharian/intern v1.0.0/go.mod h1:5DoeVV0s6jJacbCEi61lwdGj/aVlrQvzHFFd8Hwg//Y= +github.com/json-iterator/go v1.1.12 h1:PV8peI4a0ysnczrg+LtxykD8LfKY9ML6u2jnxaEnrnM= +github.com/json-iterator/go v1.1.12/go.mod h1:e30LSqwooZae/UwlEbR2852Gd8hjQvJoHmT4TnhNGBo= +github.com/jtolds/gls v4.20.0+incompatible h1:xdiiI2gbIgH/gLH7ADydsJ1uDOEzR8yvV7C0MuV77Wo= +github.com/jtolds/gls v4.20.0+incompatible/go.mod h1:QJZ7F/aHp+rZTRtaJ1ow/lLfFfVYBRgL+9YlvaHOwJU= +github.com/kardianos/osext v0.0.0-20190222173326-2bc1f35cddc0/go.mod h1:1NbS8ALrpOvjt0rHPNLyCIeMtbizbir8U//inJ+zuB8= +github.com/klauspost/cpuid/v2 v2.0.9/go.mod h1:FInQzS24/EEf25PyTYn52gqo7WaD8xa0213Md/qVLRg= +github.com/klauspost/cpuid/v2 v2.2.7 h1:ZWSB3igEs+d0qvnxR/ZBzXVmxkgt8DdzP6m9pfuVLDM= +github.com/klauspost/cpuid/v2 v2.2.7/go.mod h1:Lcz8mBdAVJIBVzewtcLocK12l3Y+JytZYpaMropDUws= +github.com/knz/go-libedit v1.10.1/go.mod h1:MZTVkCWyz0oBc7JOWP3wNAzd002ZbM/5hgShxwh4x8M= +github.com/konsorten/go-windows-terminal-sequences v1.0.1/go.mod h1:T0+1ngSBFLxvqU3pZ+m/2kptfBszLMUkC4ZK/EgS/cQ= +github.com/kr/pretty v0.1.0/go.mod h1:dAy3ld7l9f0ibDNOQOHHMYYIIbhfbHSm3C4ZsoJORNo= +github.com/kr/pretty v0.2.1/go.mod h1:ipq/a2n7PKx3OHsz4KJII5eveXtPO4qwEXGdVfWzfnI= +github.com/kr/pretty v0.3.1 h1:flRD4NNwYAUpkphVc1HcthR4KEIFJ65n8Mw5qdRn3LE= +github.com/kr/pretty v0.3.1/go.mod h1:hoEshYVHaxMs3cyo3Yncou5ZscifuDolrwPKZanG3xk= +github.com/kr/pty v1.1.1/go.mod h1:pFQYn66WHrOpPYNljwOMqo10TkYh1fy3cYio2l3bCsQ= +github.com/kr/text v0.1.0/go.mod h1:4Jbv+DJW3UT/LiOwJeYQe1efqtUx/iVham/4vfdArNI= +github.com/kr/text v0.2.0 h1:5Nx0Ya0ZqY2ygV366QzturHI13Jq95ApcVaJBhpS+AY= +github.com/kr/text v0.2.0/go.mod h1:eLer722TekiGuMkidMxC/pM04lWEeraHUUmBw8l2grE= +github.com/leodido/go-urn v1.4.0 h1:WT9HwE9SGECu3lg4d/dIA+jxlljEa1/ffXKmRjqdmIQ= +github.com/leodido/go-urn v1.4.0/go.mod h1:bvxc+MVxLKB4z00jd1z+Dvzr47oO32F/QSNjSBOlFxI= +github.com/lufia/plan9stats v0.0.0-20211012122336-39d0f177ccd0 h1:6E+4a0GO5zZEnZ81pIr0yLvtUWk2if982qA3F3QD6H4= +github.com/lufia/plan9stats v0.0.0-20211012122336-39d0f177ccd0/go.mod h1:zJYVVT2jmtg6P3p1VtQj7WsuWi/y4VnjVBn7F8KPB3I= +github.com/mailru/easyjson v0.0.0-20190614124828-94de47d64c63/go.mod h1:C1wdFJiN94OJF2b5HbByQZoLdCWB1Yqtg26g4irojpc= +github.com/mailru/easyjson v0.0.0-20190626092158-b2ccc519800e/go.mod h1:C1wdFJiN94OJF2b5HbByQZoLdCWB1Yqtg26g4irojpc= +github.com/mailru/easyjson v0.7.7 h1:UGYAvKxe3sBsEDzO8ZeWOSlIQfWFlxbzLZe7hwFURr0= +github.com/mailru/easyjson v0.7.7/go.mod h1:xzfreul335JAWq5oZzymOObrkdz5UnU4kGfJJLY9Nlc= +github.com/mattn/go-colorable v0.1.13 h1:fFA4WZxdEF4tXPZVKMLwD8oUnCTTo08duU7wxecdEvA= +github.com/mattn/go-colorable v0.1.13/go.mod h1:7S9/ev0klgBDR4GtXTXX8a3vIGJpMovkB8vQcUbaXHg= +github.com/mattn/go-isatty v0.0.20 h1:xfD0iDuEKnDkl03q4limB+vH+GxLEtL/jb4xVJSWWEY= +github.com/mattn/go-isatty v0.0.20/go.mod h1:W+V8PltTTMOvKvAeJH7IuucS94S2C6jfK/D7dTCTo3Y= +github.com/matttproud/golang_protobuf_extensions v1.0.4 h1:mmDVorXM7PCGKw94cs5zkfA9PSy5pEvNWRP0ET0TIVo= +github.com/matttproud/golang_protobuf_extensions v1.0.4/go.mod h1:BSXmuO+STAnVfrANrmjBb36TMTDstsz7MSK+HVaYKv4= +github.com/meguminnnnnnnnn/go-openai v0.0.0-20250821095446-07791bea23a0 h1:nIohpHs1ViKR0SVgW/cbBstHjmnqFZDM9RqgX9m9Xu8= +github.com/meguminnnnnnnnn/go-openai v0.0.0-20250821095446-07791bea23a0/go.mod h1:qs96ysDmxhE4BZoU45I43zcyfnaYxU3X+aRzLko/htY= +github.com/mgutz/ansi v0.0.0-20170206155736-9520e82c474b h1:j7+1HpAFS1zy5+Q4qx1fWh90gTKwiN4QCGoY9TWyyO4= +github.com/mgutz/ansi v0.0.0-20170206155736-9520e82c474b/go.mod h1:01TrycV0kFyexm33Z7vhZRXopbI8J3TDReVlkTgMUxE= +github.com/modern-go/concurrent v0.0.0-20180228061459-e0a39a4cb421/go.mod h1:6dJC0mAP4ikYIbvyc7fijjWJddQyLn8Ig3JB5CqoB9Q= +github.com/modern-go/concurrent v0.0.0-20180306012644-bacd9c7ef1dd h1:TRLaZ9cD/w8PVh93nsPXa1VrQ6jlwL5oN8l14QlcNfg= +github.com/modern-go/concurrent v0.0.0-20180306012644-bacd9c7ef1dd/go.mod h1:6dJC0mAP4ikYIbvyc7fijjWJddQyLn8Ig3JB5CqoB9Q= +github.com/modern-go/reflect2 v1.0.2 h1:xBagoLtFs94CBntxluKeaWgTMpvLxC4ur3nMaC9Gz0M= +github.com/modern-go/reflect2 v1.0.2/go.mod h1:yWuevngMOJpCy52FWWMvUC8ws7m/LJsjYzDa0/r8luk= +github.com/mohae/deepcopy v0.0.0-20170929034955-c48cc78d4826 h1:RWengNIwukTxcDr9M+97sNutRR1RKhG96O6jWumTTnw= +github.com/mohae/deepcopy v0.0.0-20170929034955-c48cc78d4826/go.mod h1:TaXosZuwdSHYgviHp1DAtfrULt5eUgsSMsZf+YrPgl8= +github.com/nikolalohinski/gonja v1.5.3 h1:GsA+EEaZDZPGJ8JtpeGN78jidhOlxeJROpqMT9fTj9c= +github.com/nikolalohinski/gonja v1.5.3/go.mod h1:RmjwxNiXAEqcq1HeK5SSMmqFJvKOfTfXhkJv6YBtPa4= +github.com/nxadm/tail v1.4.4 h1:DQuhQpB1tVlglWS2hLQ5OV6B5r8aGxSrPc5Qo6uTN78= +github.com/nxadm/tail v1.4.4/go.mod h1:kenIhsEOeOJmVchQTgglprH7qJGnHDVpk1VPCcaMI8A= +github.com/onsi/ginkgo v1.6.0/go.mod h1:lLunBs/Ym6LB5Z9jYTR76FiuTmxDTDusOGeTQH+WWjE= +github.com/onsi/ginkgo v1.8.0/go.mod h1:lLunBs/Ym6LB5Z9jYTR76FiuTmxDTDusOGeTQH+WWjE= +github.com/onsi/ginkgo v1.12.1/go.mod h1:zj2OWP4+oCPe1qIXoGWkgMRwljMUYCdkwsT2108oapk= +github.com/onsi/ginkgo v1.15.0/go.mod h1:hF8qUzuuC8DJGygJH3726JnCZX4MYbRB8yFfISqnKUg= +github.com/onsi/ginkgo v1.16.5 h1:8xi0RTUf59SOSfEtZMvwTvXYMzG4gV23XVHOZiXNtnE= +github.com/onsi/ginkgo v1.16.5/go.mod h1:+E8gABHa3K6zRBolWtd+ROzc/U5bkGt0FwiG042wbpU= +github.com/onsi/gomega v1.5.0/go.mod h1:ex+gbHU/CVuBBDIJjb2X0qEXbFg53c61hWP/1CpauHY= +github.com/onsi/gomega v1.7.1/go.mod h1:XdKZgCCFLUoM/7CFJVPcG8C1xQ1AJ0vpAezJrB7JYyY= +github.com/onsi/gomega v1.10.1/go.mod h1:iN09h71vgCQne3DLsj+A5owkum+a2tYe+TOCB1ybHNo= +github.com/onsi/gomega v1.10.5/go.mod h1:gza4q3jKQJijlu05nKWRCW/GavJumGt8aNRxWg7mt48= +github.com/onsi/gomega v1.27.3 h1:5VwIwnBY3vbBDOJrNtA4rVdiTZCsq9B5F12pvy1Drmk= +github.com/onsi/gomega v1.27.3/go.mod h1:5vG284IBtfDAmDyrK+eGyZmUgUlmi+Wngqo557cZ6Gw= +github.com/pelletier/go-toml/v2 v2.2.2 h1:aYUidT7k73Pcl9nb2gScu7NSrKCSHIDE89b3+6Wq+LM= +github.com/pelletier/go-toml/v2 v2.2.2/go.mod h1:1t835xjRzz80PqgE6HHgN2JOsmgYu/h4qDAS4n929Rs= +github.com/perimeterx/marshmallow v1.1.4 h1:pZLDH9RjlLGGorbXhcaQLhfuV0pFMNfPO55FuFkxqLw= +github.com/perimeterx/marshmallow v1.1.4/go.mod h1:dsXbUu8CRzfYP5a87xpp0xq9S3u0Vchtcl8we9tYaXw= +github.com/pkg/errors v0.8.0/go.mod h1:bwawxfHBFNV+L2hUp1rHADufV3IMtnDRdf1r5NINEl0= +github.com/pkg/errors v0.9.1 h1:FEBLx1zS214owpjy7qsBeixbURkuhQAwrK5UwLGTwt4= +github.com/pkg/errors v0.9.1/go.mod h1:bwawxfHBFNV+L2hUp1rHADufV3IMtnDRdf1r5NINEl0= +github.com/pkoukk/tiktoken-go v0.1.7 h1:qOBHXX4PHtvIvmOtyg1EeKlwFRiMKAcoMp4Q+bLQDmw= +github.com/pkoukk/tiktoken-go v0.1.7/go.mod h1:9NiV+i9mJKGj1rYOT+njbv+ZwA/zJxYdewGl6qVatpg= +github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM= +github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4= +github.com/power-devops/perfstat v0.0.0-20210106213030-5aafc221ea8c h1:ncq/mPwQF4JjgDlrVEn3C11VoGHZN7m8qihwgMEtzYw= +github.com/power-devops/perfstat v0.0.0-20210106213030-5aafc221ea8c/go.mod h1:OmDBASR4679mdNQnz2pUhc2G8CO2JrUAVFDRBDP/hJE= +github.com/prometheus/client_golang v1.16.0 h1:yk/hx9hDbrGHovbci4BY+pRMfSuuat626eFsHb7tmT8= +github.com/prometheus/client_golang v1.16.0/go.mod h1:Zsulrv/L9oM40tJ7T815tM89lFEugiJ9HzIqaAx4LKc= +github.com/prometheus/client_model v0.3.0 h1:UBgGFHqYdG/TPFD1B1ogZywDqEkwp3fBMvqdiQ7Xew4= +github.com/prometheus/client_model v0.3.0/go.mod h1:LDGWKZIo7rky3hgvBe+caln+Dr3dPggB5dvjtD7w9+w= +github.com/prometheus/common v0.42.0 h1:EKsfXEYo4JpWMHH5cg+KOUWeuJSov1Id8zGR8eeI1YM= +github.com/prometheus/common v0.42.0/go.mod h1:xBwqVerjNdUDjgODMpudtOMwlOwf2SaTr1yjz4b7Zbc= +github.com/prometheus/procfs v0.10.1 h1:kYK1Va/YMlutzCGazswoHKo//tZVlFpKYh+PymziUAg= +github.com/prometheus/procfs v0.10.1/go.mod h1:nwNm2aOCAYw8uTR/9bWRREkZFxAUcWzPHWJq+XBB/FM= +github.com/rogpeppe/go-internal v1.13.1 h1:KvO1DLK/DRN07sQ1LQKScxyZJuNnedQ5/wKSR38lUII= +github.com/rogpeppe/go-internal v1.13.1/go.mod h1:uMEvuHeurkdAXX61udpOXGD/AzZDWNMNyH2VO9fmH0o= +github.com/rollbar/rollbar-go v1.0.2/go.mod h1:AcFs5f0I+c71bpHlXNNDbOWJiKwjFDtISeXco0L5PKQ= +github.com/shirou/gopsutil/v3 v3.23.12 h1:z90NtUkp3bMtmICZKpC4+WaknU1eXtp5vtbQ11DgpE4= +github.com/shirou/gopsutil/v3 v3.23.12/go.mod h1:1FrWgea594Jp7qmjHUUPlJDTPgcsb9mGnXDxavtikzM= +github.com/shoenig/go-m1cpu v0.1.6 h1:nxdKQNcEB6vzgA2E2bvzKIYRuNj7XNJ4S/aRSwKzFtM= +github.com/shoenig/go-m1cpu v0.1.6/go.mod h1:1JJMcUBvfNwpq05QDQVAnx3gUHr9IYF7GNg9SUEw2VQ= +github.com/shoenig/test v0.6.4 h1:kVTaSd7WLz5WZ2IaoM0RSzRsUD+m8wRR+5qvntpn4LU= +github.com/shoenig/test v0.6.4/go.mod h1:byHiCGXqrVaflBLAMq/srcZIHynQPQgeyvkvXnjqq0k= +github.com/sirupsen/logrus v1.2.0/go.mod h1:LxeOpSwHxABJmUn/MG1IvRgCAasNZTLOkJPxbbu5VWo= +github.com/sirupsen/logrus v1.9.3 h1:dueUQJ1C2q9oE3F7wvmSGAaVtTmUizReu6fjN8uqzbQ= +github.com/sirupsen/logrus v1.9.3/go.mod h1:naHLuLoDiP4jHNo9R0sCBMtWGeIprob74mVsIT4qYEQ= +github.com/slongfield/pyfmt v0.0.0-20220222012616-ea85ff4c361f h1:Z2cODYsUxQPofhpYRMQVwWz4yUVpHF+vPi+eUdruUYI= +github.com/slongfield/pyfmt v0.0.0-20220222012616-ea85ff4c361f/go.mod h1:JqzWyvTuI2X4+9wOHmKSQCYxybB/8j6Ko43qVmXDuZg= +github.com/smarty/assertions v1.15.0 h1:cR//PqUBUiQRakZWqBiFFQ9wb8emQGDb0HeGdqGByCY= +github.com/smarty/assertions v1.15.0/go.mod h1:yABtdzeQs6l1brC900WlRNwj6ZR55d7B+E8C6HtKdec= +github.com/smartystreets/goconvey v1.8.1 h1:qGjIddxOk4grTu9JPOU31tVfq3cNdBlNa5sSznIX1xY= +github.com/smartystreets/goconvey v1.8.1/go.mod h1:+/u4qLyY6x1jReYOp7GOM2FSt8aP9CzCZL03bI28W60= +github.com/spaolacci/murmur3 v1.1.0 h1:7c1g84S4BPRrfL5Xrdp6fOJ206sU9y293DDHaoy0bLI= +github.com/spaolacci/murmur3 v1.1.0/go.mod h1:JwIasOWyU6f++ZhiEuf87xNszmSA2myDM2Kzu9HwQUA= +github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME= +github.com/stretchr/objx v0.1.1/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME= +github.com/stretchr/objx v0.4.0/go.mod h1:YvHI0jy2hoMjB+UWwv71VJQ9isScKT/TqJzVSSt89Yw= +github.com/stretchr/objx v0.5.0/go.mod h1:Yh+to48EsGEfYuaHDzXPcE3xhTkx73EhmCGUpEOglKo= +github.com/stretchr/objx v0.5.2 h1:xuMeJ0Sdp5ZMRXx/aWO6RZxdr3beISkG5/G/aIRr3pY= +github.com/stretchr/objx v0.5.2/go.mod h1:FRsXN1f5AsAjCGJKqEizvkpNtU+EGNCLh3NxZ/8L+MA= +github.com/stretchr/testify v1.2.2/go.mod h1:a8OnRcib4nhh0OaRAV+Yts87kKdq0PP7pXfy6kDkUVs= +github.com/stretchr/testify v1.3.0/go.mod h1:M5WIy9Dh21IEIfnGCwXGc5bZfKNJtfHm1UVUgZn+9EI= +github.com/stretchr/testify v1.7.0/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg= +github.com/stretchr/testify v1.7.1/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg= +github.com/stretchr/testify v1.8.0/go.mod h1:yNjHg4UonilssWZ8iaSj1OCr/vHnekPRkoO+kdMU+MU= +github.com/stretchr/testify v1.8.1/go.mod h1:w2LPCIKwWwSfY2zedu0+kehJoqGctiVI29o6fzry7u4= +github.com/stretchr/testify v1.8.4/go.mod h1:sz/lmYIOXD/1dqDmKjjqLyZ2RngseejIcXlSw2iwfAo= +github.com/stretchr/testify v1.9.0/go.mod h1:r2ic/lqez/lEtzL7wO/rwa5dbSLXVDPFyf8C91i36aY= +github.com/stretchr/testify v1.10.0 h1:Xv5erBjTwe/5IxqUQTdXv5kgmIvbHo3QQyRwhJsOfJA= +github.com/stretchr/testify v1.10.0/go.mod h1:r2ic/lqez/lEtzL7wO/rwa5dbSLXVDPFyf8C91i36aY= +github.com/tklauser/go-sysconf v0.3.12 h1:0QaGUFOdQaIVdPgfITYzaTegZvdCjmYO52cSFAEVmqU= +github.com/tklauser/go-sysconf v0.3.12/go.mod h1:Ho14jnntGE1fpdOqQEEaiKRpvIavV0hSfmBq8nJbHYI= +github.com/tklauser/numcpus v0.6.1 h1:ng9scYS7az0Bk4OZLvrNXNSAO2Pxr1XXRAPyjhIx+Fk= +github.com/tklauser/numcpus v0.6.1/go.mod h1:1XfjsgE2zo8GVw7POkMbHENHzVg3GzmoZ9fESEdAacY= +github.com/tmc/langchaingo v0.1.13 h1:rcpMWBIi2y3B90XxfE4Ao8dhCQPVDMaNPnN5cGB1CaA= +github.com/tmc/langchaingo v0.1.13/go.mod h1:vpQ5NOIhpzxDfTZK9B6tf2GM/MoaHewPWM5KXXGh7hg= +github.com/twitchyliquid64/golang-asm v0.15.1 h1:SU5vSMR7hnwNxj24w34ZyCi/FmDZTkS4MhqMhdFk5YI= +github.com/twitchyliquid64/golang-asm v0.15.1/go.mod h1:a1lVb/DtPvCB8fslRZhAngC2+aY1QWCk3Cedj/Gdt08= +github.com/ugorji/go v1.2.7/go.mod h1:nF9osbDWLy6bDVv/Rtoh6QgnvNDpmCalQV5urGCCS6M= +github.com/ugorji/go/codec v1.2.7/go.mod h1:WGN1fab3R1fzQlVQTkfxVtIBhWDRqOviHU95kRgeqEY= +github.com/ugorji/go/codec v1.2.12 h1:9LC83zGrHhuUA9l16C9AHXAqEV/2wBQ4nkvumAE65EE= +github.com/ugorji/go/codec v1.2.12/go.mod h1:UNopzCgEMSXjBc6AOMqYvWC1ktqTAfzJZUZgYf6w6lg= +github.com/wk8/go-ordered-map/v2 v2.1.8 h1:5h/BUHu93oj4gIdvHHHGsScSTMijfx5PeYkE/fJgbpc= +github.com/wk8/go-ordered-map/v2 v2.1.8/go.mod h1:5nJHM5DyteebpVlHnWMV0rPz6Zp7+xBAnxjb1X5vnTw= +github.com/x-cray/logrus-prefixed-formatter v0.5.2 h1:00txxvfBM9muc0jiLIEAkAcIMJzfthRT6usrui8uGmg= +github.com/x-cray/logrus-prefixed-formatter v0.5.2/go.mod h1:2duySbKsL6M18s5GU7VPsoEPHyzalCE06qoARUCeBBE= +github.com/yargevad/filepathx v1.0.0 h1:SYcT+N3tYGi+NvazubCNlvgIPbzAk7i7y2dwg3I5FYc= +github.com/yargevad/filepathx v1.0.0/go.mod h1:BprfX/gpYNJHJfc35GjRRpVcwWXS89gGulUIU5tK3tA= +github.com/yuin/goldmark v1.2.1/go.mod h1:3hX8gzYuyVAZsxl0MRgGTJEmQBFcNTphYh9decYSb74= +github.com/yusufpapurcu/wmi v1.2.3 h1:E1ctvB7uKFMOJw3fdOW32DwGE9I7t++CRUEMKvFoFiw= +github.com/yusufpapurcu/wmi v1.2.3/go.mod h1:SBZ9tNy3G9/m5Oi98Zks0QjeHVDvuK0qfxQmPyzfmi0= +golang.org/x/arch v0.11.0 h1:KXV8WWKCXm6tRpLirl2szsO5j/oOODwZf4hATmGVNs4= +golang.org/x/arch v0.11.0/go.mod h1:FEVrYAQjsQXMVJ1nsMoVVXPZg6p2JE2mx8psSWTDQys= +golang.org/x/crypto v0.0.0-20180904163835-0709b304e793/go.mod h1:6SG95UA2DQfeDnfUPMdvaQW0Q7yPrPDi9nlGo2tz2b4= +golang.org/x/crypto v0.0.0-20190308221718-c2843e01d9a2/go.mod h1:djNgcEr1/C05ACkg1iLfiJU5Ep61QUkGW8qpdssI0+w= +golang.org/x/crypto v0.0.0-20191011191535-87dc89f01550/go.mod h1:yigFU9vqHzYiE8UmvKecakEJjdnWj3jj499lnFckfCI= +golang.org/x/crypto v0.0.0-20200622213623-75b288015ac9/go.mod h1:LzIPMQfyMNhhGPhUkYOs5KpL4U8rLKemX1yGLhDgUto= +golang.org/x/crypto v0.31.0 h1:ihbySMvVjLAeSH1IbfcRTkD/iNscyz8rGzjF/E5hV6U= +golang.org/x/crypto v0.31.0/go.mod h1:kDsLvtWBEx7MV9tJOj9bnXsPbxwJQ6csT/x4KIN4Ssk= +golang.org/x/exp v0.0.0-20230713183714-613f0c0eb8a1 h1:MGwJjxBy0HJshjDNfLsYO8xppfqWlA5ZT9OhtUUhTNw= +golang.org/x/exp v0.0.0-20230713183714-613f0c0eb8a1/go.mod h1:FXUEEKJgO7OQYeo8N01OfiKP8RXMtf6e8aTskBGqWdc= +golang.org/x/mod v0.3.0/go.mod h1:s0Qsj1ACt9ePp/hMypM3fl4fZqREWJwdYDEqhRiZZUA= +golang.org/x/net v0.0.0-20180906233101-161cd47e91fd/go.mod h1:mL1N/T3taQHkDXs73rZJwtUhF3w3ftmwwsq0BUmARs4= +golang.org/x/net v0.0.0-20190404232315-eb5bcb51f2a3/go.mod h1:t9HGtf8HONx5eT2rtn7q6eTqICYqUVnKs3thJo3Qplg= +golang.org/x/net v0.0.0-20190620200207-3b0461eec859/go.mod h1:z5CRVTTTmAJ677TzLLGU+0bjPO0LkuOLi4/5GtJWs/s= +golang.org/x/net v0.0.0-20200520004742-59133d7f0dd7/go.mod h1:qpuaurCH72eLCgpAm/N6yyVIVM9cpaDIP3A8BGJEC5A= +golang.org/x/net v0.0.0-20201021035429-f5854403a974/go.mod h1:sp8m0HH+o8qH0wwXwYZr8TS3Oi6o0r6Gce1SSxlDquU= +golang.org/x/net v0.0.0-20201202161906-c7110b5ffcbb/go.mod h1:sp8m0HH+o8qH0wwXwYZr8TS3Oi6o0r6Gce1SSxlDquU= +golang.org/x/net v0.25.0 h1:d/OCCoBEUq33pjydKrGQhw7IlUPI2Oylr+8qLx49kac= +golang.org/x/net v0.25.0/go.mod h1:JkAGAh7GEvH74S6FOH42FLoXpXbE/aqXSrIQjXgsiwM= +golang.org/x/sync v0.0.0-20180314180146-1d60e4601c6f/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM= +golang.org/x/sync v0.0.0-20181221193216-37e7f081c4d4/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM= +golang.org/x/sync v0.0.0-20190423024810-112230192c58/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM= +golang.org/x/sync v0.0.0-20201020160332-67f06af15bc9/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM= +golang.org/x/sys v0.0.0-20180905080454-ebe1bf3edb33/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY= +golang.org/x/sys v0.0.0-20180909124046-d0be0721c37e/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY= +golang.org/x/sys v0.0.0-20190215142949-d0b11bdaac8a/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY= +golang.org/x/sys v0.0.0-20190412213103-97732733099d/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs= +golang.org/x/sys v0.0.0-20190904154756-749cb33beabd/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs= +golang.org/x/sys v0.0.0-20190916202348-b4ddaad3f8a3/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs= +golang.org/x/sys v0.0.0-20191005200804-aed5e4c7ecf9/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs= +golang.org/x/sys v0.0.0-20191120155948-bd437916bb0e/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs= +golang.org/x/sys v0.0.0-20200323222414-85ca7c5b95cd/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs= +golang.org/x/sys v0.0.0-20200930185726-fdedc70b468f/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs= +golang.org/x/sys v0.0.0-20201204225414-ed752295db88/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs= +golang.org/x/sys v0.0.0-20210112080510-489259a85091/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs= +golang.org/x/sys v0.0.0-20220715151400-c0bba94af5f8/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= +golang.org/x/sys v0.5.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= +golang.org/x/sys v0.6.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= +golang.org/x/sys v0.8.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= +golang.org/x/sys v0.11.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= +golang.org/x/sys v0.15.0/go.mod h1:/VUhepiaJMQUp4+oa/7Zr1D23ma6VTLIYjOOTFZPUcA= +golang.org/x/sys v0.28.0 h1:Fksou7UEQUWlKvIdsqzJmUmCX3cZuD2+P3XyyzwMhlA= +golang.org/x/sys v0.28.0/go.mod h1:/VUhepiaJMQUp4+oa/7Zr1D23ma6VTLIYjOOTFZPUcA= +golang.org/x/term v0.27.0 h1:WP60Sv1nlK1T6SupCHbXzSaN0b9wUmsPoRS9b61A23Q= +golang.org/x/term v0.27.0/go.mod h1:iMsnZpn0cago0GOrHO2+Y7u7JPn5AylBrcoWkElMTSM= +golang.org/x/text v0.3.0/go.mod h1:NqM8EUOU14njkJ3fqMW+pc6Ldnwhi/IjpwHt7yyuwOQ= +golang.org/x/text v0.3.3/go.mod h1:5Zoc/QRtKVWzQhOtBMvqHzDpF6irO9z98xDceosuGiQ= +golang.org/x/text v0.21.0 h1:zyQAAkrwaneQ066sspRyJaG9VNi/YJ1NfzcGB3hZ/qo= +golang.org/x/text v0.21.0/go.mod h1:4IBbMaMmOPCJ8SecivzSH54+73PCFmPWxNTLm+vZkEQ= +golang.org/x/tools v0.0.0-20180917221912-90fa682c2a6e/go.mod h1:n7NCudcB/nEzxVGmLbDWY5pfWTLqBcC2KZ6jyYvM4mQ= +golang.org/x/tools v0.0.0-20191119224855-298f0cb1881e/go.mod h1:b+2E5dAYhXwXZwtnZ6UAqBI28+e2cm9otk0dWdXHAEo= +golang.org/x/tools v0.0.0-20201224043029-2b0845dc783e/go.mod h1:emZCQorbCU4vsT4fOWvOPXz4eW1wZW4PmDk9uLelYpA= +golang.org/x/xerrors v0.0.0-20190717185122-a985d3407aa7/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0= +golang.org/x/xerrors v0.0.0-20191011141410-1b5146add898/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0= +golang.org/x/xerrors v0.0.0-20191204190536-9bdfabe68543/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0= +golang.org/x/xerrors v0.0.0-20200804184101-5ec99f83aff1/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0= +google.golang.org/protobuf v0.0.0-20200109180630-ec00e32a8dfd/go.mod h1:DFci5gLYBciE7Vtevhsrf46CRTquxDuWsQurQQe4oz8= +google.golang.org/protobuf v0.0.0-20200221191635-4d8936d0db64/go.mod h1:kwYJMbMJ01Woi6D6+Kah6886xMZcty6N08ah7+eCXa0= +google.golang.org/protobuf v0.0.0-20200228230310-ab0ca4ff8a60/go.mod h1:cfTl7dwQJ+fmap5saPgwCLgHXTUD7jkjRqWcaiX5VyM= +google.golang.org/protobuf v1.20.1-0.20200309200217-e05f789c0967/go.mod h1:A+miEFZTKqfCUM6K7xSMQL9OKL/b6hQv+e19PK+JZNE= +google.golang.org/protobuf v1.21.0/go.mod h1:47Nbq4nVaFHyn7ilMalzfO3qCViNmqZ2kzikPIcrTAo= +google.golang.org/protobuf v1.23.0/go.mod h1:EGpADcykh3NcUnDUJcl1+ZksZNG86OlYog2l/sGQquU= +google.golang.org/protobuf v1.34.1 h1:9ddQBjfCyZPOHPUiPxpYESBLc+T8P3E+Vo4IbKZgFWg= +google.golang.org/protobuf v1.34.1/go.mod h1:c6P6GXX6sHbq/GpV6MGZEdwhWPcYBgnhAHhKbcUYpos= +gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0= +gopkg.in/check.v1 v1.0.0-20180628173108-788fd7840127/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0= +gopkg.in/check.v1 v1.0.0-20201130134442-10cb98267c6c h1:Hei/4ADfdWqJk1ZMxUNpqntNwaWcugrBjAiHlqqRiVk= +gopkg.in/check.v1 v1.0.0-20201130134442-10cb98267c6c/go.mod h1:JHkPIbrfpd72SG/EVd6muEfDQjcINNoR0C8j2r3qZ4Q= +gopkg.in/fsnotify.v1 v1.4.7/go.mod h1:Tz8NjZHkW78fSQdbUxIjBTcgA1z1m8ZHf0WmKUhAMys= +gopkg.in/tomb.v1 v1.0.0-20141024135613-dd632973f1e7 h1:uRGJdciOHaEIrze2W8Q3AKkepLTh2hOroT7a+7czfdQ= +gopkg.in/tomb.v1 v1.0.0-20141024135613-dd632973f1e7/go.mod h1:dt/ZhP58zS4L8KSrWDmTeBkI65Dw0HsyUHuEVlX15mw= +gopkg.in/yaml.v2 v2.2.1/go.mod h1:hI93XBmqTisBFMUTm0b8Fm+jr3Dg1NNxqwp+5A1VGuI= +gopkg.in/yaml.v2 v2.2.2/go.mod h1:hI93XBmqTisBFMUTm0b8Fm+jr3Dg1NNxqwp+5A1VGuI= +gopkg.in/yaml.v2 v2.2.4/go.mod h1:hI93XBmqTisBFMUTm0b8Fm+jr3Dg1NNxqwp+5A1VGuI= +gopkg.in/yaml.v2 v2.3.0/go.mod h1:hI93XBmqTisBFMUTm0b8Fm+jr3Dg1NNxqwp+5A1VGuI= +gopkg.in/yaml.v2 v2.4.0 h1:D8xgwECY7CYvx+Y2n4sBz93Jn9JRvxdiyyo8CTfuKaY= +gopkg.in/yaml.v2 v2.4.0/go.mod h1:RDklbk79AGWmwhnvt/jBztapEOGDOx6ZbXqjP6csGnQ= +gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM= +gopkg.in/yaml.v3 v3.0.0/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM= +gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA= +gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM= +nullprogram.com/x/optparse v1.0.0/go.mod h1:KdyPE+Igbe0jQUrVfMqDMeJQIJZEuyV7pjYmp6pbG50= +sigs.k8s.io/yaml v1.3.0 h1:a2VclLzOGrwOHDiV8EfBGhvjHvP46CtW5j6POvhYGGo= +sigs.k8s.io/yaml v1.3.0/go.mod h1:GeOyir5tyXNByN85N/dRIT9es5UQNerPYEKK56eTBm8= diff --git a/example/llm_token_ratelimit/llm_client/llm_client.go b/example/llm_token_ratelimit/llm_client/llm_client.go new file mode 100644 index 000000000..4412ceee8 --- /dev/null +++ b/example/llm_token_ratelimit/llm_client/llm_client.go @@ -0,0 +1,227 @@ +// Copyright 1999-2020 Alibaba Group Holding Ltd. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package llm_client + +import ( + "context" + "fmt" + "os" + + langchaingo_llms "github.com/tmc/langchaingo/llms" + langchaingo_openai "github.com/tmc/langchaingo/llms/openai" + + eino_openai "github.com/cloudwego/eino-ext/components/model/openai" + eino_model "github.com/cloudwego/eino/components/model" + eino_schema "github.com/cloudwego/eino/schema" +) + +type LLMMessage struct { + Role string `json:"role"` + Content string `json:"content"` +} + +type LLMRequestInfos struct { + Provider LLMProvider `json:"provider"` + Messages []LLMMessage `json:"messages"` + Model string `json:"model"` +} + +type LLMResponseInfos struct { + Content string `json:"content"` + Usage map[string]any `json:"usage,omitempty"` +} + +type LLMProvider int32 + +const ( + LangChain LLMProvider = iota + Eino +) + +func (p LLMProvider) String() string { + switch p { + case LangChain: + return "langchain" + case Eino: + return "eino" + default: + return "unknown" + } +} + +func ParseLLMProvider(s string) (LLMProvider, error) { + switch s { + case "langchain": + return LangChain, nil + case "eino": + return Eino, nil + default: + return 0, fmt.Errorf("unknown LLM provider: %s", s) + } +} + +func (p LLMProvider) MarshalJSON() ([]byte, error) { + return []byte(`"` + p.String() + `"`), nil +} + +func (p *LLMProvider) UnmarshalJSON(data []byte) error { + // 去除引号 + s := string(data) + if len(s) >= 2 && s[0] == '"' && s[len(s)-1] == '"' { + s = s[1 : len(s)-1] + } + + provider, err := ParseLLMProvider(s) + if err != nil { + return err + } + *p = provider + return nil +} + +// ================================= LLMClient ==================================== + +type LLMClient interface { + GenerateContent(infos *LLMRequestInfos) (*LLMResponseInfos, error) + GetProvider() LLMProvider +} + +func NewLLMClient(infos *LLMRequestInfos) (LLMClient, error) { + switch infos.Provider { + case LangChain: + return NewLangChainClient(infos.Model) + case Eino: + return NewEinoClient(infos.Model) + default: + return nil, fmt.Errorf("unsupported provider: %v", infos.Provider) + } +} + +// ================================= LangChainClient ================================ + +type LangChainClient struct { + llm langchaingo_llms.Model +} + +func NewLangChainClient(model string) (LLMClient, error) { + apiKey := os.Getenv("LLM_API_KEY") + if apiKey == "" { + return nil, fmt.Errorf("LLM_API_KEY environment variable is not set") + } + + baseURL := os.Getenv("LLM_BASE_URL") + if baseURL == "" { + return nil, fmt.Errorf("LLM_BASE_URL environment variable is not set") + } + + llm, err := langchaingo_openai.New( + langchaingo_openai.WithToken(apiKey), + langchaingo_openai.WithBaseURL(baseURL), + langchaingo_openai.WithModel(model), + ) + if err != nil { + return nil, fmt.Errorf("failed to create LangChain LLM client: %w", err) + } + + return &LangChainClient{ + llm: llm, + }, nil +} + +func (c *LangChainClient) GenerateContent(infos *LLMRequestInfos) (*LLMResponseInfos, error) { + if infos == nil || infos.Messages == nil { + return nil, fmt.Errorf("invalid request infos") + } + content := make([]langchaingo_llms.MessageContent, len(infos.Messages)) + for i, msg := range infos.Messages { + content[i] = langchaingo_llms.TextParts(langchaingo_llms.ChatMessageType(msg.Role), msg.Content) + } + completion, err := c.llm.GenerateContent(context.Background(), content) + if err != nil { + return nil, fmt.Errorf("failed to generate content: %w", err) + } + return &LLMResponseInfos{ + Content: completion.Choices[0].Content, + Usage: map[string]any{ + "prompt_tokens": completion.Choices[0].GenerationInfo["PromptTokens"], + "completion_tokens": completion.Choices[0].GenerationInfo["CompletionTokens"], + "total_tokens": completion.Choices[0].GenerationInfo["TotalTokens"], + }, + }, nil +} + +func (c *LangChainClient) GetProvider() LLMProvider { + return LangChain +} + +// ================================= EinoClient ==================================== + +type EinoClient struct { + llm eino_model.BaseChatModel +} + +func NewEinoClient(model string) (LLMClient, error) { + apiKey := os.Getenv("LLM_API_KEY") + if apiKey == "" { + return nil, fmt.Errorf("LLM_API_KEY environment variable is not set") + } + + baseURL := os.Getenv("LLM_BASE_URL") + if baseURL == "" { + return nil, fmt.Errorf("LLM_BASE_URL environment variable is not set") + } + + llm, err := eino_openai.NewChatModel(context.Background(), &eino_openai.ChatModelConfig{ + BaseURL: baseURL, + APIKey: apiKey, + Model: model, + }) + if err != nil { + return nil, fmt.Errorf("failed to create Eino LLM client: %w", err) + } + + return &EinoClient{ + llm: llm, + }, nil +} + +func (c *EinoClient) GenerateContent(infos *LLMRequestInfos) (*LLMResponseInfos, error) { + if infos == nil || infos.Messages == nil { + return nil, fmt.Errorf("invalid request infos") + } + content := make([]*eino_schema.Message, len(infos.Messages)) + for i, msg := range infos.Messages { + content[i] = &eino_schema.Message{ + Role: eino_schema.RoleType(msg.Role), + Content: msg.Content, + } + } + completion, err := c.llm.Generate(context.Background(), content) + if err != nil { + return nil, fmt.Errorf("failed to generate content: %w", err) + } + return &LLMResponseInfos{ + Content: completion.Content, + Usage: map[string]any{ + "prompt_tokens": completion.ResponseMeta.Usage.PromptTokens, + "completion_tokens": completion.ResponseMeta.Usage.CompletionTokens, + "total_tokens": completion.ResponseMeta.Usage.TotalTokens, + }, + }, nil +} + +func (c *EinoClient) GetProvider() LLMProvider { + return Eino +} diff --git a/example/llm_token_ratelimit/main.go b/example/llm_token_ratelimit/main.go new file mode 100644 index 000000000..485be1d5f --- /dev/null +++ b/example/llm_token_ratelimit/main.go @@ -0,0 +1,41 @@ +// Copyright 1999-2020 Alibaba Group Holding Ltd. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package main + +import ( + "fmt" + "llm_token_ratelimit/ratelimit" + "llm_token_ratelimit/server" + "net/http" + _ "net/http/pprof" + + "github.com/gin-gonic/gin" +) + +func StartMonitor() { + fmt.Println("pprof is running on http://127.0.0.1:6060/debug/pprof/") + fmt.Println(http.ListenAndServe("127.0.0.1:6060", nil)) +} + +func StartSerivce() { + gin.SetMode(gin.ReleaseMode) + ratelimit.InitSentinel() + server.StartServer("127.0.0.1", 9527) +} + +func main() { + go StartMonitor() + StartSerivce() +} diff --git a/example/llm_token_ratelimit/ratelimit/options.go b/example/llm_token_ratelimit/ratelimit/options.go new file mode 100644 index 000000000..c0ecc8d5d --- /dev/null +++ b/example/llm_token_ratelimit/ratelimit/options.go @@ -0,0 +1,62 @@ +// Copyright 1999-2020 Alibaba Group Holding Ltd. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package ratelimit + +import ( + llmtokenratelimit "github.com/alibaba/sentinel-golang/core/llm_token_ratelimit" + "github.com/gin-gonic/gin" +) + +type Option func(*options) + +type options struct { + resourceExtract func(*gin.Context) string + blockFallback func(*gin.Context) + requestInfosExtract func(*gin.Context) *llmtokenratelimit.RequestInfos + promptsExtract func(*gin.Context) []string +} + +func WithBlockFallback(fn func(ctx *gin.Context)) Option { + return func(opts *options) { + opts.blockFallback = fn + } +} + +func WithResourceExtractor(fn func(*gin.Context) string) Option { + return func(opts *options) { + opts.resourceExtract = fn + } +} + +func WithRequestInfosExtractor(fn func(*gin.Context) *llmtokenratelimit.RequestInfos) Option { + return func(opts *options) { + opts.requestInfosExtract = fn + } +} + +func WithPromptsExtractor(fn func(*gin.Context) []string) Option { + return func(opts *options) { + opts.promptsExtract = fn + } +} + +func evaluateOptions(opts []Option) *options { + optCopy := &options{} + for _, opt := range opts { + opt(optCopy) + } + + return optCopy +} diff --git a/example/llm_token_ratelimit/ratelimit/ratelimit.go b/example/llm_token_ratelimit/ratelimit/ratelimit.go new file mode 100644 index 000000000..0d4b084da --- /dev/null +++ b/example/llm_token_ratelimit/ratelimit/ratelimit.go @@ -0,0 +1,158 @@ +// Copyright 1999-2020 Alibaba Group Holding Ltd. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package ratelimit + +import ( + sentinel "github.com/alibaba/sentinel-golang/api" + "github.com/alibaba/sentinel-golang/core/base" + llmtokenratelimit "github.com/alibaba/sentinel-golang/core/llm_token_ratelimit" + "github.com/gin-gonic/gin" +) + +func InitSentinel() { + if err := sentinel.InitDefault(); err != nil { + panic(err) + } + + _, err := llmtokenratelimit.LoadRules([]*llmtokenratelimit.Rule{ + { + + Resource: "POST:/v1/chat/completion/fixed_window", + Strategy: llmtokenratelimit.FixedWindow, + SpecificItems: []*llmtokenratelimit.SpecificItem{ + { + Identifier: llmtokenratelimit.Identifier{ + Type: llmtokenratelimit.Header, + Value: ".*", + }, + KeyItems: []*llmtokenratelimit.KeyItem{ + { + Key: ".*", + Token: llmtokenratelimit.Token{ + Number: 65, + CountStrategy: llmtokenratelimit.TotalTokens, + }, + Time: llmtokenratelimit.Time{ + Unit: llmtokenratelimit.Second, + Value: 10, + }, + }, + }, + }, + }, + }, + { + + Resource: "POST:/v1/chat/completion/peta", + Strategy: llmtokenratelimit.PETA, + Encoding: llmtokenratelimit.TokenEncoding{}, + SpecificItems: []*llmtokenratelimit.SpecificItem{ + { + Identifier: llmtokenratelimit.Identifier{ + Type: llmtokenratelimit.Header, + Value: ".*", + }, + KeyItems: []*llmtokenratelimit.KeyItem{ + { + Key: ".*", + Token: llmtokenratelimit.Token{ + Number: 65, + CountStrategy: llmtokenratelimit.TotalTokens, + }, + Time: llmtokenratelimit.Time{ + Unit: llmtokenratelimit.Second, + Value: 10, + }, + }, + }, + }, + }, + }, + }) + + if err != nil { + panic(err) + } +} + +func SentinelMiddleware(opts ...Option) gin.HandlerFunc { + options := evaluateOptions(opts) + return func(c *gin.Context) { + resource := c.Request.Method + ":" + c.FullPath() + + if options.resourceExtract != nil { + resource = options.resourceExtract(c) + } + + prompts := []string{} + if options.promptsExtract != nil { + prompts = options.promptsExtract(c) + } + + reqInfos := llmtokenratelimit.GenerateRequestInfos( + llmtokenratelimit.WithHeader(c.Request.Header), + llmtokenratelimit.WithPrompts(prompts), + ) + + if options.requestInfosExtract != nil { + reqInfos = options.requestInfosExtract(c) + } + + // Check + entry, err := sentinel.Entry(resource, sentinel.WithTrafficType(base.Inbound), sentinel.WithArgs(reqInfos)) + if err != nil { + // Block + if options.blockFallback != nil { + options.blockFallback(c) + } else { + responseHeader, ok := err.TriggeredValue().(*llmtokenratelimit.ResponseHeader) + if !ok || responseHeader == nil { + c.AbortWithStatusJSON(500, gin.H{ + "error": "internal server error. invalid response header.", + }) + return + } + setResponseHeaders(c, responseHeader) + c.AbortWithStatusJSON(int(responseHeader.ErrorCode), gin.H{ + "error": responseHeader.ErrorMessage, + }) + } + return + } + // Set response headers + responseHeader, ok := entry.Context().GetPair(llmtokenratelimit.KeyResponseHeaders).(*llmtokenratelimit.ResponseHeader) + if ok && responseHeader != nil { + setResponseHeaders(c, responseHeader) + } + // Pass or Disabled + c.Next() + // Update used token info + usedTokenInfos, exists := c.Get(llmtokenratelimit.KeyUsedTokenInfos) + if exists && usedTokenInfos != nil { + entry.SetPair(llmtokenratelimit.KeyUsedTokenInfos, usedTokenInfos) + } + entry.Exit() // Must be executed immediately after the SetPair function + } +} + +func setResponseHeaders(c *gin.Context, header *llmtokenratelimit.ResponseHeader) { + if c == nil || header == nil { + return + } + + for key, value := range header.GetAll() { + c.Header(key, value) + } +} diff --git a/example/llm_token_ratelimit/server/route.go b/example/llm_token_ratelimit/server/route.go new file mode 100644 index 000000000..b771f14c7 --- /dev/null +++ b/example/llm_token_ratelimit/server/route.go @@ -0,0 +1,94 @@ +// Copyright 1999-2020 Alibaba Group Holding Ltd. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package server + +import ( + "llm_token_ratelimit/llm_client" + "net/http" + "time" + + llmtokenratelimit "github.com/alibaba/sentinel-golang/core/llm_token_ratelimit" + "github.com/gin-gonic/gin" +) + +func (s *Server) chatCompletion(c *gin.Context) { + infos, err := bindJSONFromCache[llm_client.LLMRequestInfos](c) + if err != nil { + c.JSON(http.StatusBadRequest, gin.H{ + "error": gin.H{ + "message": "failed to process LLM request", + "details": err.Error(), + "type": "api_error", + }, + }) + return + } + client, err := llm_client.NewLLMClient(infos) + if err != nil { + c.JSON(http.StatusInternalServerError, gin.H{ + "error": gin.H{ + "message": "failed to create LLM client", + "details": err.Error(), + "type": "client_error", + }, + }) + return + } + response, err := client.GenerateContent(infos) + if err != nil { + c.JSON(http.StatusInternalServerError, gin.H{ + "error": gin.H{ + "message": "failed to generate content", + "details": err.Error(), + "type": "generation_error", + }, + }) + return + } + + if response == nil { + c.JSON(http.StatusInternalServerError, gin.H{ + "error": gin.H{ + "message": "received nil response from LLM", + "type": "response_error", + }, + }) + return + } + + usedTokenInfos, err := llmtokenratelimit.OpenAITokenExtractor(response.Usage) + if err != nil { + c.JSON(http.StatusInternalServerError, gin.H{ + "error": gin.H{ + "message": "failed to extract token usage info", + "details": err.Error(), + "type": "extraction_error", + }, + }) + return + } + c.Set(llmtokenratelimit.KeyUsedTokenInfos, usedTokenInfos) + + c.JSON(http.StatusOK, gin.H{ + "message": "success", + "timestamp": time.Now().Unix(), + "choices": response.Content, + "usage": gin.H{ + "input_tokens": usedTokenInfos.InputTokens, + "output_tokens": usedTokenInfos.OutputTokens, + "total_tokens": usedTokenInfos.TotalTokens, + }, + }) +} diff --git a/example/llm_token_ratelimit/server/server.go b/example/llm_token_ratelimit/server/server.go new file mode 100644 index 000000000..490633084 --- /dev/null +++ b/example/llm_token_ratelimit/server/server.go @@ -0,0 +1,156 @@ +// Copyright 1999-2020 Alibaba Group Holding Ltd. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package server + +import ( + "bytes" + "fmt" + "io" + "llm_token_ratelimit/llm_client" + "llm_token_ratelimit/ratelimit" + "net/http" + + "github.com/gin-gonic/gin" +) + +const ( + KeyRawBody = "rawBody" +) + +func corsMiddleware() gin.HandlerFunc { + return func(c *gin.Context) { + c.Header("Access-Control-Allow-Origin", "*") + c.Header("Access-Control-Allow-Methods", "GET, POST, PUT, DELETE, OPTIONS") + c.Header("Access-Control-Allow-Headers", "Content-Type, Authorization") + + if c.Request.Method == "OPTIONS" { + c.AbortWithStatus(http.StatusNoContent) + return + } + + c.Next() + } +} + +func cacheBodyMiddleware() gin.HandlerFunc { + return func(c *gin.Context) { + if c.Request.Body == nil { + c.Next() + return + } + + bodyBytes, err := io.ReadAll(c.Request.Body) + if err != nil { + c.JSON(http.StatusBadRequest, gin.H{ + "error": "failed to read request body", + }) + c.Abort() + return + } + + c.Request.Body.Close() + c.Request.Body = io.NopCloser(bytes.NewBuffer(bodyBytes)) + + c.Set(KeyRawBody, bodyBytes) + + c.Next() + } +} + +func bindJSONFromCache[T any](c *gin.Context) (*T, error) { + rawBody, exists := c.Get(KeyRawBody) + if !exists { + return nil, fmt.Errorf("raw body not found in context") + } + + bodyBytes, ok := rawBody.([]byte) + if !ok { + return nil, fmt.Errorf("invalid raw body type") + } + + tempBody := io.NopCloser(bytes.NewBuffer(bodyBytes)) + originalBody := c.Request.Body + c.Request.Body = tempBody + + var result T + err := c.ShouldBindJSON(&result) + + c.Request.Body = originalBody + + if err != nil { + return nil, err + } + + return &result, nil +} + +type Server struct { + engine *gin.Engine + ip string + port uint16 +} + +func NewServer(ip string, port uint16) *Server { + engine := gin.New() + + engine.Use(gin.Logger()) + engine.Use(gin.Recovery()) + engine.Use(corsMiddleware()) + engine.Use(cacheBodyMiddleware()) + engine.Use(ratelimit.SentinelMiddleware( + ratelimit.WithPromptsExtractor(func(c *gin.Context) []string { + infos, err := bindJSONFromCache[llm_client.LLMRequestInfos](c) + if err != nil { + c.JSON(http.StatusBadRequest, gin.H{ + "error": gin.H{ + "message": "failed to process LLM request", + "details": err.Error(), + "type": "api_error", + }, + }) + return nil + } + prompts := make([]string, 0, len(infos.Messages)) + for _, msg := range infos.Messages { + prompts = append(prompts, msg.Content) + } + return prompts + }), + )) + + return &Server{ + engine: engine, + ip: ip, + port: port, + } +} + +func (s *Server) setupRoutes() { + s.engine.POST("/v1/chat/completion/fixed_window", s.chatCompletion) + s.engine.POST("/v1/chat/completion/peta", s.chatCompletion) +} + +func (s *Server) Start() error { + s.setupRoutes() + addr := fmt.Sprintf("%s:%d", s.ip, s.port) + return s.engine.Run(addr) +} + +func StartServer(ip string, port uint16) { + server := NewServer(ip, port) + if err := server.Start(); err != nil { + panic(err) + } +} diff --git a/ext/datasource/helper_test.go b/ext/datasource/helper_test.go index 3a4daf78c..d119342cb 100644 --- a/ext/datasource/helper_test.go +++ b/ext/datasource/helper_test.go @@ -404,10 +404,83 @@ func TestHotSpotParamRuleJsonArrayParser(t *testing.T) { for _, r := range rules { fmt.Println(r) } - assert.True(t, strings.Contains(rules[0].String(), "Resource:abc, MetricType:Concurrency, ControlBehavior:Reject, ParamIndex:0, ParamKey:, Threshold:1000, MaxQueueingTimeMs:1, BurstCount:10, DurationInSec:1, ParamsMaxCapacity:10000, SpecificItems:map[true:10003 1000:10001 ximu:10002]")) - assert.True(t, strings.Contains(rules[1].String(), "Resource:abc, MetricType:Concurrency, ControlBehavior:Throttling, ParamIndex:1, ParamKey:, Threshold:2000, MaxQueueingTimeMs:2, BurstCount:20, DurationInSec:2, ParamsMaxCapacity:20000, SpecificItems:map[true:20003 1000:20001 ximu:20002")) - assert.True(t, strings.Contains(rules[2].String(), "Resource:abc, MetricType:QPS, ControlBehavior:Reject, ParamIndex:2, ParamKey:, Threshold:3000, MaxQueueingTimeMs:3, BurstCount:30, DurationInSec:3, ParamsMaxCapacity:30000, SpecificItems:map[true:30003 1000:30001 ximu:30002")) - assert.True(t, strings.Contains(rules[3].String(), "Resource:abc, MetricType:QPS, ControlBehavior:Throttling, ParamIndex:3, ParamKey:, Threshold:4000, MaxQueueingTimeMs:4, BurstCount:40, DurationInSec:4, ParamsMaxCapacity:40000, SpecificItems:map[true:40003 1000:40001 ximu:40002")) + + expectedRule0 := &hotspot.Rule{ + Resource: "abc", + MetricType: hotspot.Concurrency, + ControlBehavior: hotspot.Reject, + ParamIndex: 0, + ParamKey: "", + Threshold: 1000, + MaxQueueingTimeMs: 1, + BurstCount: 10, + DurationInSec: 1, + ParamsMaxCapacity: 10000, + SpecificItems: map[interface{}]int64{ + true: 10003, + 1000: 10001, + "ximu": 10002, + }, + } + + expectedRule1 := &hotspot.Rule{ + Resource: "abc", + MetricType: hotspot.Concurrency, + ControlBehavior: hotspot.Throttling, + ParamIndex: 1, + ParamKey: "", + Threshold: 2000, + MaxQueueingTimeMs: 2, + BurstCount: 20, + DurationInSec: 2, + ParamsMaxCapacity: 20000, + SpecificItems: map[interface{}]int64{ + true: 20003, + 1000: 20001, + "ximu": 20002, + }, + } + + expectedRule2 := &hotspot.Rule{ + Resource: "abc", + MetricType: hotspot.QPS, + ControlBehavior: hotspot.Reject, + ParamIndex: 2, + ParamKey: "", + Threshold: 3000, + MaxQueueingTimeMs: 3, + BurstCount: 30, + DurationInSec: 3, + ParamsMaxCapacity: 30000, + SpecificItems: map[interface{}]int64{ + true: 30003, + 1000: 30001, + "ximu": 30002, + }, + } + + expectedRule3 := &hotspot.Rule{ + Resource: "abc", + MetricType: hotspot.QPS, + ControlBehavior: hotspot.Throttling, + ParamIndex: 3, + ParamKey: "", + Threshold: 4000, + MaxQueueingTimeMs: 4, + BurstCount: 40, + DurationInSec: 4, + ParamsMaxCapacity: 40000, + SpecificItems: map[interface{}]int64{ + true: 40003, + 1000: 40001, + "ximu": 40002, + }, + } + + assert.True(t, reflect.DeepEqual(rules[0], expectedRule0)) + assert.True(t, reflect.DeepEqual(rules[1], expectedRule1)) + assert.True(t, reflect.DeepEqual(rules[2], expectedRule2)) + assert.True(t, reflect.DeepEqual(rules[3], expectedRule3)) }) t.Run("TestHotSpotParamRuleJsonArrayParser_Nil", func(t *testing.T) { diff --git a/go.mod b/go.mod index 988afcf07..ecc73fa3c 100644 --- a/go.mod +++ b/go.mod @@ -3,14 +3,19 @@ module github.com/alibaba/sentinel-golang go 1.18 require ( - github.com/fsnotify/fsnotify v1.4.7 - github.com/google/uuid v1.1.1 + github.com/fsnotify/fsnotify v1.4.9 + github.com/go-redis/redis/v8 v8.11.0 + github.com/google/uuid v1.3.0 + github.com/jinzhu/copier v0.4.0 github.com/pkg/errors v0.9.1 + github.com/pkoukk/tiktoken-go v0.1.7 github.com/prometheus/client_golang v1.16.0 github.com/shirou/gopsutil/v3 v3.21.6 - github.com/stretchr/testify v1.8.0 + github.com/spaolacci/murmur3 v1.1.0 + github.com/stretchr/testify v1.8.2 go.uber.org/multierr v1.5.0 gopkg.in/yaml.v2 v2.4.0 + gopkg.in/yaml.v3 v3.0.1 ) require ( @@ -18,6 +23,8 @@ require ( github.com/beorn7/perks v1.0.1 // indirect github.com/cespare/xxhash/v2 v2.2.0 // indirect github.com/davecgh/go-spew v1.1.1 // indirect + github.com/dgryski/go-rendezvous v0.0.0-20200823014737-9f7001d12a5f // indirect + github.com/dlclark/regexp2 v1.10.0 // indirect github.com/go-ole/go-ole v1.2.4 // indirect github.com/golang/protobuf v1.5.3 // indirect github.com/google/go-cmp v0.6.0 // indirect @@ -27,12 +34,11 @@ require ( github.com/prometheus/common v0.42.0 // indirect github.com/prometheus/procfs v0.10.1 // indirect github.com/rogpeppe/go-internal v1.13.1 // indirect - github.com/stretchr/objx v0.4.0 // indirect + github.com/stretchr/objx v0.5.0 // indirect github.com/tklauser/go-sysconf v0.3.6 // indirect github.com/tklauser/numcpus v0.2.2 // indirect go.uber.org/atomic v1.6.0 // indirect golang.org/x/sys v0.21.0 // indirect golang.org/x/tools v0.22.0 // indirect google.golang.org/protobuf v1.30.0 // indirect - gopkg.in/yaml.v3 v3.0.1 // indirect ) diff --git a/go.sum b/go.sum index 0717c9b37..83d562e8c 100644 --- a/go.sum +++ b/go.sum @@ -4,26 +4,47 @@ github.com/StackExchange/wmi v0.0.0-20190523213315-cbe66965904d h1:G0m3OIz70MZUW github.com/StackExchange/wmi v0.0.0-20190523213315-cbe66965904d/go.mod h1:3eOhrUMpNV+6aFIbp5/iudMxNCF27Vw2OZgy4xEx0Fg= github.com/beorn7/perks v1.0.1 h1:VlbKKnNfV8bJzeqoa4cOKqO6bYr3WgKZxO8Z16+hsOM= github.com/beorn7/perks v1.0.1/go.mod h1:G2ZrVWU2WbWT9wwq4/hrbKbnv/1ERSJQ0ibhJ6rlkpw= +github.com/cespare/xxhash/v2 v2.1.1/go.mod h1:VGX0DQ3Q6kWi7AoAeZDth3/j3BFtOZR5XLFGgcrjCOs= github.com/cespare/xxhash/v2 v2.2.0 h1:DC2CZ1Ep5Y4k3ZQ899DldepgrayRUGE6BBZ/cd9Cj44= github.com/cespare/xxhash/v2 v2.2.0/go.mod h1:VGX0DQ3Q6kWi7AoAeZDth3/j3BFtOZR5XLFGgcrjCOs= github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c= github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= -github.com/fsnotify/fsnotify v1.4.7 h1:IXs+QLmnXW2CcXuY+8Mzv/fWEsPGWxqefPtCP5CnV9I= +github.com/dgryski/go-rendezvous v0.0.0-20200823014737-9f7001d12a5f h1:lO4WD4F/rVNCu3HqELle0jiPLLBs70cWOduZpkS1E78= +github.com/dgryski/go-rendezvous v0.0.0-20200823014737-9f7001d12a5f/go.mod h1:cuUVRXasLTGF7a8hSLbxyZXjz+1KgoB3wDUb6vlszIc= +github.com/dlclark/regexp2 v1.10.0 h1:+/GIL799phkJqYW+3YbOd8LCcbHzT0Pbo8zl70MHsq0= +github.com/dlclark/regexp2 v1.10.0/go.mod h1:DHkYz0B9wPfa6wondMfaivmHpzrQ3v9q8cnmRbL6yW8= github.com/fsnotify/fsnotify v1.4.7/go.mod h1:jwhsz4b93w/PPRr/qN1Yymfu8t87LnFCMoQvtojpjFo= +github.com/fsnotify/fsnotify v1.4.9 h1:hsms1Qyu0jgnwNXIxa+/V/PDsU6CfLf6CNO8H7IWoS4= +github.com/fsnotify/fsnotify v1.4.9/go.mod h1:znqG4EE+3YCdAaPaxE2ZRY/06pZUdp0tY4IgpuI1SZQ= github.com/go-ole/go-ole v1.2.4 h1:nNBDSCOigTSiarFpYE9J/KtEA1IOW4CNeqT9TQDqCxI= github.com/go-ole/go-ole v1.2.4/go.mod h1:XCwSNxSkXRo4vlyPy93sltvi/qJq0jqQhjqQNIwKuxM= +github.com/go-redis/redis/v8 v8.11.0 h1:O1Td0mQ8UFChQ3N9zFQqo6kTU2cJ+/it88gDB+zg0wo= +github.com/go-redis/redis/v8 v8.11.0/go.mod h1:DLomh7y2e3ggQXQLd1YgmvIfecPJoFl7WU5SOQ/r06M= github.com/golang/protobuf v1.2.0/go.mod h1:6lQm79b+lXiMfvg/cZm0SGofjICqVBUtrP5yJMmIC1U= github.com/golang/protobuf v1.3.5/go.mod h1:6O5/vntMXwX2lRkT1hjjk0nAC1IDOTvTlVgjlRvqsdk= +github.com/golang/protobuf v1.4.0-rc.1/go.mod h1:ceaxUfeHdC40wWswd/P6IGgMaK3YpKi5j83Wpe3EHw8= +github.com/golang/protobuf v1.4.0-rc.1.0.20200221234624-67d41d38c208/go.mod h1:xKAWHe0F5eneWXFV3EuXVDTCmh+JuBKY0li0aMyXATA= +github.com/golang/protobuf v1.4.0-rc.2/go.mod h1:LlEzMj4AhA7rCAGe4KMBDvJI+AwstrUpVNzEA03Pprs= +github.com/golang/protobuf v1.4.0-rc.4.0.20200313231945-b860323f09d0/go.mod h1:WU3c8KckQ9AFe+yFwt9sWVRKCVIyN9cPHBJSNnbL67w= +github.com/golang/protobuf v1.4.0/go.mod h1:jodUvKwWbYaEsadDk5Fwe5c77LiNKVO9IDvqG2KuDX0= +github.com/golang/protobuf v1.4.2/go.mod h1:oDoupMAO8OvCJWAcko0GGGIgR6R6ocIYbsSw735rRwI= github.com/golang/protobuf v1.5.0/go.mod h1:FsONVRAS9T7sI+LIUmWTfcYkHO4aIWwzhcaSAoJOfIk= github.com/golang/protobuf v1.5.3 h1:KhyjKVUg7Usr/dYsdSqoFveMYd5ko72D+zANwlG1mmg= github.com/golang/protobuf v1.5.3/go.mod h1:XVQd3VNwM+JqD3oG2Ue2ip4fOMUkwXdXDdiuN0vRsmY= +github.com/google/go-cmp v0.3.0/go.mod h1:8QqcDgzrUqlUb/G2PQTWiueGozuR1884gddMywk6iLU= +github.com/google/go-cmp v0.3.1/go.mod h1:8QqcDgzrUqlUb/G2PQTWiueGozuR1884gddMywk6iLU= +github.com/google/go-cmp v0.4.0/go.mod h1:v8dTdLbMG2kIc/vJvl+f65V22dbkXbowE6jgT/gNBxE= github.com/google/go-cmp v0.5.5/go.mod h1:v8dTdLbMG2kIc/vJvl+f65V22dbkXbowE6jgT/gNBxE= +github.com/google/go-cmp v0.5.6/go.mod h1:v8dTdLbMG2kIc/vJvl+f65V22dbkXbowE6jgT/gNBxE= github.com/google/go-cmp v0.6.0 h1:ofyhxvXcZhMsU5ulbFiLKl/XBFqE1GSq7atu8tAmTRI= github.com/google/go-cmp v0.6.0/go.mod h1:17dUlkBOakJ0+DkrSSNjCkIjxS6bF9zb3elmeNGIjoY= github.com/google/renameio v0.1.0/go.mod h1:KWCgfxg9yswjAJkECMjeO8J8rahYeXnNhOm40UhjYkI= -github.com/google/uuid v1.1.1 h1:Gkbcsh/GbpXz7lPftLA3P6TYMwjCLYm83jiFQZF/3gY= -github.com/google/uuid v1.1.1/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo= +github.com/google/uuid v1.3.0 h1:t6JiXgmwXMjEs8VusXIJk2BXHsn+wx8BZdTaoZ5fu7I= +github.com/google/uuid v1.3.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo= +github.com/hpcloud/tail v1.0.0/go.mod h1:ab1qPbhIpdTxEkNHXyeSf5vhxWSCs/tWer42PpOxQnU= +github.com/jinzhu/copier v0.4.0 h1:w3ciUoD19shMCRargcpm0cm91ytaBhDvuRpz1ODO/U8= +github.com/jinzhu/copier v0.4.0/go.mod h1:DfbEm0FYsaqBcKcFuvmOZb218JkPGtvSHsKg8S8hyyg= github.com/kisielk/gotool v1.0.0/go.mod h1:XhKaO+MFFWcvkIS/tQcRk01m1F5IRFswLeQ+oQHNcck= github.com/kr/pretty v0.1.0/go.mod h1:dAy3ld7l9f0ibDNOQOHHMYYIIbhfbHSm3C4ZsoJORNo= github.com/kr/pretty v0.3.1 h1:flRD4NNwYAUpkphVc1HcthR4KEIFJ65n8Mw5qdRn3LE= @@ -32,8 +53,20 @@ github.com/kr/text v0.1.0 h1:45sCR5RtlFHMR4UwH9sdQ5TC8v0qDQCHnXt+kaKSTVE= github.com/kr/text v0.1.0/go.mod h1:4Jbv+DJW3UT/LiOwJeYQe1efqtUx/iVham/4vfdArNI= github.com/matttproud/golang_protobuf_extensions v1.0.4 h1:mmDVorXM7PCGKw94cs5zkfA9PSy5pEvNWRP0ET0TIVo= github.com/matttproud/golang_protobuf_extensions v1.0.4/go.mod h1:BSXmuO+STAnVfrANrmjBb36TMTDstsz7MSK+HVaYKv4= +github.com/nxadm/tail v1.4.4 h1:DQuhQpB1tVlglWS2hLQ5OV6B5r8aGxSrPc5Qo6uTN78= +github.com/nxadm/tail v1.4.4/go.mod h1:kenIhsEOeOJmVchQTgglprH7qJGnHDVpk1VPCcaMI8A= +github.com/onsi/ginkgo v1.6.0/go.mod h1:lLunBs/Ym6LB5Z9jYTR76FiuTmxDTDusOGeTQH+WWjE= +github.com/onsi/ginkgo v1.12.1/go.mod h1:zj2OWP4+oCPe1qIXoGWkgMRwljMUYCdkwsT2108oapk= +github.com/onsi/ginkgo v1.15.0 h1:1V1NfVQR87RtWAgp1lv9JZJ5Jap+XFGKPi00andXGi4= +github.com/onsi/ginkgo v1.15.0/go.mod h1:hF8qUzuuC8DJGygJH3726JnCZX4MYbRB8yFfISqnKUg= +github.com/onsi/gomega v1.7.1/go.mod h1:XdKZgCCFLUoM/7CFJVPcG8C1xQ1AJ0vpAezJrB7JYyY= +github.com/onsi/gomega v1.10.1/go.mod h1:iN09h71vgCQne3DLsj+A5owkum+a2tYe+TOCB1ybHNo= +github.com/onsi/gomega v1.10.5 h1:7n6FEkpFmfCoo2t+YYqXH0evK+a9ICQz0xcAy9dYcaQ= +github.com/onsi/gomega v1.10.5/go.mod h1:gza4q3jKQJijlu05nKWRCW/GavJumGt8aNRxWg7mt48= github.com/pkg/errors v0.9.1 h1:FEBLx1zS214owpjy7qsBeixbURkuhQAwrK5UwLGTwt4= github.com/pkg/errors v0.9.1/go.mod h1:bwawxfHBFNV+L2hUp1rHADufV3IMtnDRdf1r5NINEl0= +github.com/pkoukk/tiktoken-go v0.1.7 h1:qOBHXX4PHtvIvmOtyg1EeKlwFRiMKAcoMp4Q+bLQDmw= +github.com/pkoukk/tiktoken-go v0.1.7/go.mod h1:9NiV+i9mJKGj1rYOT+njbv+ZwA/zJxYdewGl6qVatpg= github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM= github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4= github.com/prometheus/client_golang v1.16.0 h1:yk/hx9hDbrGHovbci4BY+pRMfSuuat626eFsHb7tmT8= @@ -49,18 +82,23 @@ github.com/rogpeppe/go-internal v1.13.1 h1:KvO1DLK/DRN07sQ1LQKScxyZJuNnedQ5/wKSR github.com/rogpeppe/go-internal v1.13.1/go.mod h1:uMEvuHeurkdAXX61udpOXGD/AzZDWNMNyH2VO9fmH0o= github.com/shirou/gopsutil/v3 v3.21.6 h1:vU7jrp1Ic/2sHB7w6UNs7MIkn7ebVtTb5D9j45o9VYE= github.com/shirou/gopsutil/v3 v3.21.6/go.mod h1:JfVbDpIBLVzT8oKbvMg9P3wEIMDDpVn+LwHTKj0ST88= +github.com/spaolacci/murmur3 v1.1.0 h1:7c1g84S4BPRrfL5Xrdp6fOJ206sU9y293DDHaoy0bLI= +github.com/spaolacci/murmur3 v1.1.0/go.mod h1:JwIasOWyU6f++ZhiEuf87xNszmSA2myDM2Kzu9HwQUA= github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME= -github.com/stretchr/objx v0.4.0 h1:M2gUjqZET1qApGOWNSnZ49BAIMX4F/1plDv3+l31EJ4= github.com/stretchr/objx v0.4.0/go.mod h1:YvHI0jy2hoMjB+UWwv71VJQ9isScKT/TqJzVSSt89Yw= +github.com/stretchr/objx v0.5.0 h1:1zr/of2m5FGMsad5YfcqgdqdWrIhu+EBEJRhR1U7z/c= +github.com/stretchr/objx v0.5.0/go.mod h1:Yh+to48EsGEfYuaHDzXPcE3xhTkx73EhmCGUpEOglKo= github.com/stretchr/testify v1.3.0/go.mod h1:M5WIy9Dh21IEIfnGCwXGc5bZfKNJtfHm1UVUgZn+9EI= github.com/stretchr/testify v1.6.1/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg= github.com/stretchr/testify v1.7.1/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg= -github.com/stretchr/testify v1.8.0 h1:pSgiaMZlXftHpm5L7V1+rVB+AZJydKsMxsQBIJw4PKk= github.com/stretchr/testify v1.8.0/go.mod h1:yNjHg4UonilssWZ8iaSj1OCr/vHnekPRkoO+kdMU+MU= +github.com/stretchr/testify v1.8.2 h1:+h33VjcLVPDHtOdpUCuF+7gSuG3yGIftsP1YvFihtJ8= +github.com/stretchr/testify v1.8.2/go.mod h1:w2LPCIKwWwSfY2zedu0+kehJoqGctiVI29o6fzry7u4= github.com/tklauser/go-sysconf v0.3.6 h1:oc1sJWvKkmvIxhDHeKWvZS4f6AW+YcoguSfRF2/Hmo4= github.com/tklauser/go-sysconf v0.3.6/go.mod h1:MkWzOF4RMCshBAMXuhXJs64Rte09mITnppBXY/rYEFI= github.com/tklauser/numcpus v0.2.2 h1:oyhllyrScuYI6g+h/zUvNXNp1wy7x8qQy3t/piefldA= github.com/tklauser/numcpus v0.2.2/go.mod h1:x3qojaO3uyYt0i56EW/VUYs7uBvdl2fkfZFu0T9wgjM= +github.com/yuin/goldmark v1.2.1/go.mod h1:3hX8gzYuyVAZsxl0MRgGTJEmQBFcNTphYh9decYSb74= go.uber.org/atomic v1.6.0 h1:Ezj3JGmsOnG1MoRWQkPBsKLe9DwWD9QeXzTRzzldNVk= go.uber.org/atomic v1.6.0/go.mod h1:sABNBOSYdrvTF6hTgEIbc7YasKWGhgEQZyfxyTvoXHQ= go.uber.org/multierr v1.5.0 h1:KCa4XfM8CWFCpxXRGok+Q0SS/0XBhMDbHHGABQLvD2A= @@ -69,30 +107,60 @@ go.uber.org/tools v0.0.0-20190618225709-2cfd321de3ee h1:0mgffUl7nfd+FpvXMVz4IDEa go.uber.org/tools v0.0.0-20190618225709-2cfd321de3ee/go.mod h1:vJERXedbb3MVM5f9Ejo0C68/HhF8uaILCdgjnY+goOA= golang.org/x/crypto v0.0.0-20190308221718-c2843e01d9a2/go.mod h1:djNgcEr1/C05ACkg1iLfiJU5Ep61QUkGW8qpdssI0+w= golang.org/x/crypto v0.0.0-20190510104115-cbcb75029529/go.mod h1:yigFU9vqHzYiE8UmvKecakEJjdnWj3jj499lnFckfCI= +golang.org/x/crypto v0.0.0-20191011191535-87dc89f01550/go.mod h1:yigFU9vqHzYiE8UmvKecakEJjdnWj3jj499lnFckfCI= +golang.org/x/crypto v0.0.0-20200622213623-75b288015ac9/go.mod h1:LzIPMQfyMNhhGPhUkYOs5KpL4U8rLKemX1yGLhDgUto= golang.org/x/lint v0.0.0-20190930215403-16217165b5de h1:5hukYrvBGR8/eNkX5mdUezrA6JiaEZDtJb9Ei+1LlBs= golang.org/x/lint v0.0.0-20190930215403-16217165b5de/go.mod h1:6SW0HCj/g11FgYtHlgUYUwCkIfeOF89ocIRzGO/8vkc= golang.org/x/mod v0.0.0-20190513183733-4bf6d317e70e/go.mod h1:mXi4GBBbnImb6dmsKGUJ2LatrhH/nqhxcFungHvyanc= +golang.org/x/mod v0.3.0/go.mod h1:s0Qsj1ACt9ePp/hMypM3fl4fZqREWJwdYDEqhRiZZUA= golang.org/x/mod v0.18.0 h1:5+9lSbEzPSdWkH32vYPBwEpX8KwDbM52Ud9xBUvNlb0= +golang.org/x/net v0.0.0-20180906233101-161cd47e91fd/go.mod h1:mL1N/T3taQHkDXs73rZJwtUhF3w3ftmwwsq0BUmARs4= golang.org/x/net v0.0.0-20190311183353-d8887717615a/go.mod h1:t9HGtf8HONx5eT2rtn7q6eTqICYqUVnKs3thJo3Qplg= golang.org/x/net v0.0.0-20190404232315-eb5bcb51f2a3/go.mod h1:t9HGtf8HONx5eT2rtn7q6eTqICYqUVnKs3thJo3Qplg= golang.org/x/net v0.0.0-20190620200207-3b0461eec859/go.mod h1:z5CRVTTTmAJ677TzLLGU+0bjPO0LkuOLi4/5GtJWs/s= +golang.org/x/net v0.0.0-20200520004742-59133d7f0dd7/go.mod h1:qpuaurCH72eLCgpAm/N6yyVIVM9cpaDIP3A8BGJEC5A= +golang.org/x/net v0.0.0-20201021035429-f5854403a974/go.mod h1:sp8m0HH+o8qH0wwXwYZr8TS3Oi6o0r6Gce1SSxlDquU= +golang.org/x/net v0.0.0-20201202161906-c7110b5ffcbb/go.mod h1:sp8m0HH+o8qH0wwXwYZr8TS3Oi6o0r6Gce1SSxlDquU= +golang.org/x/net v0.26.0 h1:soB7SVo0PWrY4vPW/+ay0jKDNScG2X9wFeYlXIvJsOQ= +golang.org/x/sync v0.0.0-20180314180146-1d60e4601c6f/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM= golang.org/x/sync v0.0.0-20181221193216-37e7f081c4d4/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM= golang.org/x/sync v0.0.0-20190423024810-112230192c58/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM= +golang.org/x/sync v0.0.0-20201020160332-67f06af15bc9/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM= golang.org/x/sync v0.7.0 h1:YsImfSBoP9QPYL0xyKJPq0gcaJdG3rInoqxTWbfQu9M= +golang.org/x/sys v0.0.0-20180909124046-d0be0721c37e/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY= golang.org/x/sys v0.0.0-20190215142949-d0b11bdaac8a/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY= golang.org/x/sys v0.0.0-20190412213103-97732733099d/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs= +golang.org/x/sys v0.0.0-20190904154756-749cb33beabd/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs= +golang.org/x/sys v0.0.0-20191005200804-aed5e4c7ecf9/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs= +golang.org/x/sys v0.0.0-20191120155948-bd437916bb0e/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs= +golang.org/x/sys v0.0.0-20200323222414-85ca7c5b95cd/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs= +golang.org/x/sys v0.0.0-20200930185726-fdedc70b468f/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs= +golang.org/x/sys v0.0.0-20210112080510-489259a85091/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs= golang.org/x/sys v0.0.0-20210316164454-77fc1eacc6aa/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs= golang.org/x/sys v0.21.0 h1:rF+pYz3DAGSQAxAu1CbC7catZg4ebC4UIeIhKxBZvws= golang.org/x/sys v0.21.0/go.mod h1:/VUhepiaJMQUp4+oa/7Zr1D23ma6VTLIYjOOTFZPUcA= golang.org/x/text v0.3.0/go.mod h1:NqM8EUOU14njkJ3fqMW+pc6Ldnwhi/IjpwHt7yyuwOQ= +golang.org/x/text v0.3.3/go.mod h1:5Zoc/QRtKVWzQhOtBMvqHzDpF6irO9z98xDceosuGiQ= +golang.org/x/text v0.7.0 h1:4BRB4x83lYWy72KwLD/qYDuTu7q9PjSagHvijDw7cLo= +golang.org/x/tools v0.0.0-20180917221912-90fa682c2a6e/go.mod h1:n7NCudcB/nEzxVGmLbDWY5pfWTLqBcC2KZ6jyYvM4mQ= golang.org/x/tools v0.0.0-20190311212946-11955173bddd/go.mod h1:LCzVGOaR6xXOjkQ3onu1FJEFr0SW1gC7cKk1uF8kGRs= golang.org/x/tools v0.0.0-20190621195816-6e04913cbbac/go.mod h1:/rFqwRUd4F7ZHNgwSSTFct+R/Kf4OFW1sUzUTQQTgfc= golang.org/x/tools v0.0.0-20191029041327-9cc4af7d6b2c/go.mod h1:b+2E5dAYhXwXZwtnZ6UAqBI28+e2cm9otk0dWdXHAEo= golang.org/x/tools v0.0.0-20191029190741-b9c20aec41a5/go.mod h1:b+2E5dAYhXwXZwtnZ6UAqBI28+e2cm9otk0dWdXHAEo= +golang.org/x/tools v0.0.0-20191119224855-298f0cb1881e/go.mod h1:b+2E5dAYhXwXZwtnZ6UAqBI28+e2cm9otk0dWdXHAEo= +golang.org/x/tools v0.0.0-20201224043029-2b0845dc783e/go.mod h1:emZCQorbCU4vsT4fOWvOPXz4eW1wZW4PmDk9uLelYpA= golang.org/x/tools v0.22.0 h1:gqSGLZqv+AI9lIQzniJ0nZDRG5GBPsSi+DRNHWNz6yA= golang.org/x/tools v0.22.0/go.mod h1:aCwcsjqvq7Yqt6TNyX7QMU2enbQ/Gt0bo6krSeEri+c= golang.org/x/xerrors v0.0.0-20190717185122-a985d3407aa7/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0= +golang.org/x/xerrors v0.0.0-20191011141410-1b5146add898/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0= golang.org/x/xerrors v0.0.0-20191204190536-9bdfabe68543/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0= +golang.org/x/xerrors v0.0.0-20200804184101-5ec99f83aff1/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0= +google.golang.org/protobuf v0.0.0-20200109180630-ec00e32a8dfd/go.mod h1:DFci5gLYBciE7Vtevhsrf46CRTquxDuWsQurQQe4oz8= +google.golang.org/protobuf v0.0.0-20200221191635-4d8936d0db64/go.mod h1:kwYJMbMJ01Woi6D6+Kah6886xMZcty6N08ah7+eCXa0= +google.golang.org/protobuf v0.0.0-20200228230310-ab0ca4ff8a60/go.mod h1:cfTl7dwQJ+fmap5saPgwCLgHXTUD7jkjRqWcaiX5VyM= +google.golang.org/protobuf v1.20.1-0.20200309200217-e05f789c0967/go.mod h1:A+miEFZTKqfCUM6K7xSMQL9OKL/b6hQv+e19PK+JZNE= +google.golang.org/protobuf v1.21.0/go.mod h1:47Nbq4nVaFHyn7ilMalzfO3qCViNmqZ2kzikPIcrTAo= +google.golang.org/protobuf v1.23.0/go.mod h1:EGpADcykh3NcUnDUJcl1+ZksZNG86OlYog2l/sGQquU= google.golang.org/protobuf v1.26.0-rc.1/go.mod h1:jlhhOSvTdKEhbULTjvd4ARK9grFBp09yW+WbY/TyQbw= google.golang.org/protobuf v1.26.0/go.mod h1:9q0QmTI4eRPtz6boOQmLYwt+qCgq0jsYwAQnmE0givc= google.golang.org/protobuf v1.30.0 h1:kPPoIgf3TsEvrm0PFe15JQ+570QVxYzEvvHqChK+cng= @@ -101,6 +169,11 @@ gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8 gopkg.in/check.v1 v1.0.0-20180628173108-788fd7840127/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0= gopkg.in/check.v1 v1.0.0-20201130134442-10cb98267c6c h1:Hei/4ADfdWqJk1ZMxUNpqntNwaWcugrBjAiHlqqRiVk= gopkg.in/errgo.v2 v2.1.0/go.mod h1:hNsd1EY+bozCKY1Ytp96fpM3vjJbqLJn88ws8XvfDNI= +gopkg.in/fsnotify.v1 v1.4.7/go.mod h1:Tz8NjZHkW78fSQdbUxIjBTcgA1z1m8ZHf0WmKUhAMys= +gopkg.in/tomb.v1 v1.0.0-20141024135613-dd632973f1e7 h1:uRGJdciOHaEIrze2W8Q3AKkepLTh2hOroT7a+7czfdQ= +gopkg.in/tomb.v1 v1.0.0-20141024135613-dd632973f1e7/go.mod h1:dt/ZhP58zS4L8KSrWDmTeBkI65Dw0HsyUHuEVlX15mw= +gopkg.in/yaml.v2 v2.2.4/go.mod h1:hI93XBmqTisBFMUTm0b8Fm+jr3Dg1NNxqwp+5A1VGuI= +gopkg.in/yaml.v2 v2.3.0/go.mod h1:hI93XBmqTisBFMUTm0b8Fm+jr3Dg1NNxqwp+5A1VGuI= gopkg.in/yaml.v2 v2.4.0 h1:D8xgwECY7CYvx+Y2n4sBz93Jn9JRvxdiyyo8CTfuKaY= gopkg.in/yaml.v2 v2.4.0/go.mod h1:RDklbk79AGWmwhnvt/jBztapEOGDOx6ZbXqjP6csGnQ= gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM= diff --git a/pkg/adapters/eino/go.mod b/pkg/adapters/eino/go.mod new file mode 100644 index 000000000..4adef1719 --- /dev/null +++ b/pkg/adapters/eino/go.mod @@ -0,0 +1,76 @@ +module github.com/alibaba/sentinel-golang/pkg/adapters/eino + +go 1.22.0 + +replace github.com/alibaba/sentinel-golang => ../../../ + +require ( + github.com/alibaba/sentinel-golang v1.0.4 + github.com/cloudwego/eino v0.4.8 + github.com/cloudwego/eino-ext/components/model/openai v0.0.0-20250905035413-86dbae6351d5 +) + +require ( + github.com/StackExchange/wmi v0.0.0-20190523213315-cbe66965904d // indirect + github.com/bahlo/generic-list-go v0.2.0 // indirect + github.com/beorn7/perks v1.0.1 // indirect + github.com/buger/jsonparser v1.1.1 // indirect + github.com/bytedance/sonic v1.14.0 // indirect + github.com/bytedance/sonic/loader v0.3.0 // indirect + github.com/cespare/xxhash/v2 v2.2.0 // indirect + github.com/cloudwego/base64x v0.1.5 // indirect + github.com/cloudwego/eino-ext/libs/acl/openai v0.0.0-20250826113018-8c6f6358d4bb // indirect + github.com/dgryski/go-rendezvous v0.0.0-20200823014737-9f7001d12a5f // indirect + github.com/dlclark/regexp2 v1.10.0 // indirect + github.com/dustin/go-humanize v1.0.1 // indirect + github.com/eino-contrib/jsonschema v1.0.0 // indirect + github.com/evanphx/json-patch v0.5.2 // indirect + github.com/getkin/kin-openapi v0.118.0 // indirect + github.com/go-ole/go-ole v1.2.4 // indirect + github.com/go-openapi/jsonpointer v0.19.5 // indirect + github.com/go-openapi/swag v0.19.5 // indirect + github.com/go-redis/redis/v8 v8.11.0 // indirect + github.com/golang/protobuf v1.5.3 // indirect + github.com/google/uuid v1.3.0 // indirect + github.com/goph/emperror v0.17.2 // indirect + github.com/invopop/yaml v0.1.0 // indirect + github.com/jinzhu/copier v0.4.0 // indirect + github.com/josharian/intern v1.0.0 // indirect + github.com/json-iterator/go v1.1.12 // indirect + github.com/klauspost/cpuid/v2 v2.0.9 // indirect + github.com/kr/text v0.2.0 // indirect + github.com/mailru/easyjson v0.7.7 // indirect + github.com/matttproud/golang_protobuf_extensions v1.0.4 // indirect + github.com/meguminnnnnnnnn/go-openai v0.0.0-20250821095446-07791bea23a0 // indirect + github.com/modern-go/concurrent v0.0.0-20180306012644-bacd9c7ef1dd // indirect + github.com/modern-go/reflect2 v1.0.2 // indirect + github.com/mohae/deepcopy v0.0.0-20170929034955-c48cc78d4826 // indirect + github.com/nikolalohinski/gonja v1.5.3 // indirect + github.com/pelletier/go-toml/v2 v2.0.9 // indirect + github.com/perimeterx/marshmallow v1.1.4 // indirect + github.com/pkg/errors v0.9.1 // indirect + github.com/pkoukk/tiktoken-go v0.1.7 // indirect + github.com/prometheus/client_golang v1.16.0 // indirect + github.com/prometheus/client_model v0.3.0 // indirect + github.com/prometheus/common v0.42.0 // indirect + github.com/prometheus/procfs v0.10.1 // indirect + github.com/shirou/gopsutil/v3 v3.21.6 // indirect + github.com/sirupsen/logrus v1.9.3 // indirect + github.com/slongfield/pyfmt v0.0.0-20220222012616-ea85ff4c361f // indirect + github.com/spaolacci/murmur3 v1.1.0 // indirect + github.com/stretchr/objx v0.5.2 // indirect + github.com/tklauser/go-sysconf v0.3.6 // indirect + github.com/tklauser/numcpus v0.2.2 // indirect + github.com/twitchyliquid64/golang-asm v0.15.1 // indirect + github.com/wk8/go-ordered-map/v2 v2.1.8 // indirect + github.com/yargevad/filepathx v1.0.0 // indirect + golang.org/x/arch v0.11.0 // indirect + golang.org/x/exp v0.0.0-20230713183714-613f0c0eb8a1 // indirect + golang.org/x/net v0.21.0 // indirect + golang.org/x/sys v0.28.0 // indirect + golang.org/x/term v0.27.0 // indirect + golang.org/x/text v0.21.0 // indirect + google.golang.org/protobuf v1.30.0 // indirect + gopkg.in/yaml.v2 v2.4.0 // indirect + gopkg.in/yaml.v3 v3.0.1 // indirect +) diff --git a/pkg/adapters/eino/go.sum b/pkg/adapters/eino/go.sum new file mode 100644 index 000000000..f42e3371a --- /dev/null +++ b/pkg/adapters/eino/go.sum @@ -0,0 +1,299 @@ +github.com/StackExchange/wmi v0.0.0-20190523213315-cbe66965904d h1:G0m3OIz70MZUWq3EgK3CesDbo8upS2Vm9/P3FtgI+Jk= +github.com/StackExchange/wmi v0.0.0-20190523213315-cbe66965904d/go.mod h1:3eOhrUMpNV+6aFIbp5/iudMxNCF27Vw2OZgy4xEx0Fg= +github.com/airbrake/gobrake v3.6.1+incompatible/go.mod h1:wM4gu3Cn0W0K7GUuVWnlXZU11AGBXMILnrdOU8Kn00o= +github.com/bahlo/generic-list-go v0.2.0 h1:5sz/EEAK+ls5wF+NeqDpk5+iNdMDXrh3z3nPnH1Wvgk= +github.com/bahlo/generic-list-go v0.2.0/go.mod h1:2KvAjgMlE5NNynlg/5iLrrCCZ2+5xWbdbCW3pNTGyYg= +github.com/beorn7/perks v1.0.1 h1:VlbKKnNfV8bJzeqoa4cOKqO6bYr3WgKZxO8Z16+hsOM= +github.com/beorn7/perks v1.0.1/go.mod h1:G2ZrVWU2WbWT9wwq4/hrbKbnv/1ERSJQ0ibhJ6rlkpw= +github.com/bitly/go-simplejson v0.5.0/go.mod h1:cXHtHw4XUPsvGaxgjIAn8PhEWG9NfngEKAMDJEczWVA= +github.com/bmizerany/assert v0.0.0-20160611221934-b7ed37b82869/go.mod h1:Ekp36dRnpXw/yCqJaO+ZrUyxD+3VXMFFr56k5XYrpB4= +github.com/buger/jsonparser v1.1.1 h1:2PnMjfWD7wBILjqQbt530v576A/cAbQvEW9gGIpYMUs= +github.com/buger/jsonparser v1.1.1/go.mod h1:6RYKKt7H4d4+iWqouImQ9R2FZql3VbhNgx27UK13J/0= +github.com/bugsnag/bugsnag-go v1.4.0/go.mod h1:2oa8nejYd4cQ/b0hMIopN0lCRxU0bueqREvZLWFrtK8= +github.com/bugsnag/panicwrap v1.2.0/go.mod h1:D/8v3kj0zr8ZAKg1AQ6crr+5VwKN5eIywRkfhyM/+dE= +github.com/bytedance/mockey v1.2.14 h1:KZaFgPdiUwW+jOWFieo3Lr7INM1P+6adO3hxZhDswY8= +github.com/bytedance/mockey v1.2.14/go.mod h1:1BPHF9sol5R1ud/+0VEHGQq/+i2lN+GTsr3O2Q9IENY= +github.com/bytedance/sonic v1.14.0 h1:/OfKt8HFw0kh2rj8N0F6C/qPGRESq0BbaNZgcNXXzQQ= +github.com/bytedance/sonic v1.14.0/go.mod h1:WoEbx8WTcFJfzCe0hbmyTGrfjt8PzNEBdxlNUO24NhA= +github.com/bytedance/sonic/loader v0.1.1/go.mod h1:ncP89zfokxS5LZrJxl5z0UJcsk4M4yY2JpfqGeCtNLU= +github.com/bytedance/sonic/loader v0.3.0 h1:dskwH8edlzNMctoruo8FPTJDF3vLtDT0sXZwvZJyqeA= +github.com/bytedance/sonic/loader v0.3.0/go.mod h1:N8A3vUdtUebEY2/VQC0MyhYeKUFosQU6FxH2JmUe6VI= +github.com/certifi/gocertifi v0.0.0-20190105021004-abcd57078448/go.mod h1:GJKEexRPVJrBSOjoqN5VNOIKJ5Q3RViH6eu3puDRwx4= +github.com/cespare/xxhash/v2 v2.1.1/go.mod h1:VGX0DQ3Q6kWi7AoAeZDth3/j3BFtOZR5XLFGgcrjCOs= +github.com/cespare/xxhash/v2 v2.2.0 h1:DC2CZ1Ep5Y4k3ZQ899DldepgrayRUGE6BBZ/cd9Cj44= +github.com/cespare/xxhash/v2 v2.2.0/go.mod h1:VGX0DQ3Q6kWi7AoAeZDth3/j3BFtOZR5XLFGgcrjCOs= +github.com/cloudwego/base64x v0.1.5 h1:XPciSp1xaq2VCSt6lF0phncD4koWyULpl5bUxbfCyP4= +github.com/cloudwego/base64x v0.1.5/go.mod h1:0zlkT4Wn5C6NdauXdJRhSKRlJvmclQ1hhJgA0rcu/8w= +github.com/cloudwego/eino v0.4.8 h1:wptTU24tQad1mFCHw0+4zSzH+p8dLEBk6HtggPlcvP0= +github.com/cloudwego/eino v0.4.8/go.mod h1:1TDlOmwGSsbCJaWB92w9YLZi2FL0WRZoRcD4eMvqikg= +github.com/cloudwego/eino-ext/components/model/openai v0.0.0-20250905035413-86dbae6351d5 h1:D04jOL3xKn9CstVRaPkunhPffIbEVp0VfQV5e/26lS8= +github.com/cloudwego/eino-ext/components/model/openai v0.0.0-20250905035413-86dbae6351d5/go.mod h1:QQhCuQxuBAVWvu/YAZBhs/RsR76mUigw59Tl0kh04C8= +github.com/cloudwego/eino-ext/libs/acl/openai v0.0.0-20250826113018-8c6f6358d4bb h1:RMslzyijc3bi9EkqCulpS0hZupTl1y/wayR3+fVRN/c= +github.com/cloudwego/eino-ext/libs/acl/openai v0.0.0-20250826113018-8c6f6358d4bb/go.mod h1:fHn/6OqPPY1iLLx9wzz+MEVT5Dl9gwuZte1oLEnCoYw= +github.com/cloudwego/iasm v0.2.0/go.mod h1:8rXZaNYT2n95jn+zTI1sDr+IgcD2GVs0nlbbQPiEFhY= +github.com/creack/pty v1.1.9/go.mod h1:oKZEueFk5CKHvIhNR5MUki03XCEU+Q6VDXinZuGJ33E= +github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= +github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c= +github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= +github.com/dgryski/go-rendezvous v0.0.0-20200823014737-9f7001d12a5f h1:lO4WD4F/rVNCu3HqELle0jiPLLBs70cWOduZpkS1E78= +github.com/dgryski/go-rendezvous v0.0.0-20200823014737-9f7001d12a5f/go.mod h1:cuUVRXasLTGF7a8hSLbxyZXjz+1KgoB3wDUb6vlszIc= +github.com/dlclark/regexp2 v1.10.0 h1:+/GIL799phkJqYW+3YbOd8LCcbHzT0Pbo8zl70MHsq0= +github.com/dlclark/regexp2 v1.10.0/go.mod h1:DHkYz0B9wPfa6wondMfaivmHpzrQ3v9q8cnmRbL6yW8= +github.com/dustin/go-humanize v1.0.1 h1:GzkhY7T5VNhEkwH0PVJgjz+fX1rhBrR7pRT3mDkpeCY= +github.com/dustin/go-humanize v1.0.1/go.mod h1:Mu1zIs6XwVuF/gI1OepvI0qD18qycQx+mFykh5fBlto= +github.com/eino-contrib/jsonschema v1.0.0 h1:dXxbhGNZuI3+xNi8x3JT8AGyoXz6Pff6mRvmpjVl5Ww= +github.com/eino-contrib/jsonschema v1.0.0/go.mod h1:cpnX4SyKjWjGC7iN2EbhxaTdLqGjCi0e9DxpLYxddD4= +github.com/evanphx/json-patch v0.5.2 h1:xVCHIVMUu1wtM/VkR9jVZ45N3FhZfYMMYGorLCR8P3k= +github.com/evanphx/json-patch v0.5.2/go.mod h1:ZWS5hhDbVDyob71nXKNL0+PWn6ToqBHMikGIFbs31qQ= +github.com/fsnotify/fsnotify v1.4.7/go.mod h1:jwhsz4b93w/PPRr/qN1Yymfu8t87LnFCMoQvtojpjFo= +github.com/fsnotify/fsnotify v1.4.9 h1:hsms1Qyu0jgnwNXIxa+/V/PDsU6CfLf6CNO8H7IWoS4= +github.com/fsnotify/fsnotify v1.4.9/go.mod h1:znqG4EE+3YCdAaPaxE2ZRY/06pZUdp0tY4IgpuI1SZQ= +github.com/getkin/kin-openapi v0.118.0 h1:z43njxPmJ7TaPpMSCQb7PN0dEYno4tyBPQcrFdHoLuM= +github.com/getkin/kin-openapi v0.118.0/go.mod h1:l5e9PaFUo9fyLJCPGQeXI2ML8c3P8BHOEV2VaAVf/pc= +github.com/getsentry/raven-go v0.2.0/go.mod h1:KungGk8q33+aIAZUIVWZDr2OfAEBsO49PX4NzFV5kcQ= +github.com/go-check/check v0.0.0-20180628173108-788fd7840127 h1:0gkP6mzaMqkmpcJYCFOLkIBwI7xFExG03bbkOkCvUPI= +github.com/go-check/check v0.0.0-20180628173108-788fd7840127/go.mod h1:9ES+weclKsC9YodN5RgxqK/VD9HM9JsCSh7rNhMZE98= +github.com/go-ole/go-ole v1.2.4 h1:nNBDSCOigTSiarFpYE9J/KtEA1IOW4CNeqT9TQDqCxI= +github.com/go-ole/go-ole v1.2.4/go.mod h1:XCwSNxSkXRo4vlyPy93sltvi/qJq0jqQhjqQNIwKuxM= +github.com/go-openapi/jsonpointer v0.19.5 h1:gZr+CIYByUqjcgeLXnQu2gHYQC9o73G2XUeOFYEICuY= +github.com/go-openapi/jsonpointer v0.19.5/go.mod h1:Pl9vOtqEWErmShwVjC8pYs9cog34VGT37dQOVbmoatg= +github.com/go-openapi/swag v0.19.5 h1:lTz6Ys4CmqqCQmZPBlbQENR1/GucA2bzYTE12Pw4tFY= +github.com/go-openapi/swag v0.19.5/go.mod h1:POnQmlKehdgb5mhVOsnJFsivZCEZ/vjK9gh66Z9tfKk= +github.com/go-redis/redis/v8 v8.11.0 h1:O1Td0mQ8UFChQ3N9zFQqo6kTU2cJ+/it88gDB+zg0wo= +github.com/go-redis/redis/v8 v8.11.0/go.mod h1:DLomh7y2e3ggQXQLd1YgmvIfecPJoFl7WU5SOQ/r06M= +github.com/go-test/deep v1.0.8 h1:TDsG77qcSprGbC6vTN8OuXp5g+J+b5Pcguhf7Zt61VM= +github.com/go-test/deep v1.0.8/go.mod h1:5C2ZWiW0ErCdrYzpqxLbTX7MG14M9iiw8DgHncVwcsE= +github.com/gofrs/uuid v3.2.0+incompatible/go.mod h1:b2aQJv3Z4Fp6yNu3cdSllBxTCLRxnplIgP/c0N/04lM= +github.com/golang/protobuf v1.2.0/go.mod h1:6lQm79b+lXiMfvg/cZm0SGofjICqVBUtrP5yJMmIC1U= +github.com/golang/protobuf v1.3.5/go.mod h1:6O5/vntMXwX2lRkT1hjjk0nAC1IDOTvTlVgjlRvqsdk= +github.com/golang/protobuf v1.4.0-rc.1/go.mod h1:ceaxUfeHdC40wWswd/P6IGgMaK3YpKi5j83Wpe3EHw8= +github.com/golang/protobuf v1.4.0-rc.1.0.20200221234624-67d41d38c208/go.mod h1:xKAWHe0F5eneWXFV3EuXVDTCmh+JuBKY0li0aMyXATA= +github.com/golang/protobuf v1.4.0-rc.2/go.mod h1:LlEzMj4AhA7rCAGe4KMBDvJI+AwstrUpVNzEA03Pprs= +github.com/golang/protobuf v1.4.0-rc.4.0.20200313231945-b860323f09d0/go.mod h1:WU3c8KckQ9AFe+yFwt9sWVRKCVIyN9cPHBJSNnbL67w= +github.com/golang/protobuf v1.4.0/go.mod h1:jodUvKwWbYaEsadDk5Fwe5c77LiNKVO9IDvqG2KuDX0= +github.com/golang/protobuf v1.4.2/go.mod h1:oDoupMAO8OvCJWAcko0GGGIgR6R6ocIYbsSw735rRwI= +github.com/golang/protobuf v1.5.0/go.mod h1:FsONVRAS9T7sI+LIUmWTfcYkHO4aIWwzhcaSAoJOfIk= +github.com/golang/protobuf v1.5.3 h1:KhyjKVUg7Usr/dYsdSqoFveMYd5ko72D+zANwlG1mmg= +github.com/golang/protobuf v1.5.3/go.mod h1:XVQd3VNwM+JqD3oG2Ue2ip4fOMUkwXdXDdiuN0vRsmY= +github.com/google/go-cmp v0.3.0/go.mod h1:8QqcDgzrUqlUb/G2PQTWiueGozuR1884gddMywk6iLU= +github.com/google/go-cmp v0.3.1/go.mod h1:8QqcDgzrUqlUb/G2PQTWiueGozuR1884gddMywk6iLU= +github.com/google/go-cmp v0.4.0/go.mod h1:v8dTdLbMG2kIc/vJvl+f65V22dbkXbowE6jgT/gNBxE= +github.com/google/go-cmp v0.5.5/go.mod h1:v8dTdLbMG2kIc/vJvl+f65V22dbkXbowE6jgT/gNBxE= +github.com/google/go-cmp v0.5.6/go.mod h1:v8dTdLbMG2kIc/vJvl+f65V22dbkXbowE6jgT/gNBxE= +github.com/google/go-cmp v0.6.0 h1:ofyhxvXcZhMsU5ulbFiLKl/XBFqE1GSq7atu8tAmTRI= +github.com/google/go-cmp v0.6.0/go.mod h1:17dUlkBOakJ0+DkrSSNjCkIjxS6bF9zb3elmeNGIjoY= +github.com/google/gofuzz v1.0.0/go.mod h1:dBl0BpW6vV/+mYPU4Po3pmUjxk6FQPldtuIdl/M65Eg= +github.com/google/uuid v1.3.0 h1:t6JiXgmwXMjEs8VusXIJk2BXHsn+wx8BZdTaoZ5fu7I= +github.com/google/uuid v1.3.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo= +github.com/goph/emperror v0.17.2 h1:yLapQcmEsO0ipe9p5TaN22djm3OFV/TfM/fcYP0/J18= +github.com/goph/emperror v0.17.2/go.mod h1:+ZbQ+fUNO/6FNiUo0ujtMjhgad9Xa6fQL9KhH4LNHic= +github.com/gopherjs/gopherjs v1.17.2 h1:fQnZVsXk8uxXIStYb0N4bGk7jeyTalG/wsZjQ25dO0g= +github.com/gopherjs/gopherjs v1.17.2/go.mod h1:pRRIvn/QzFLrKfvEz3qUuEhtE/zLCWfreZ6J5gM2i+k= +github.com/gorilla/mux v1.8.0/go.mod h1:DVbg23sWSpFRCP0SfiEN6jmj59UnW/n46BH5rLB71So= +github.com/hpcloud/tail v1.0.0/go.mod h1:ab1qPbhIpdTxEkNHXyeSf5vhxWSCs/tWer42PpOxQnU= +github.com/invopop/yaml v0.1.0 h1:YW3WGUoJEXYfzWBjn00zIlrw7brGVD0fUKRYDPAPhrc= +github.com/invopop/yaml v0.1.0/go.mod h1:2XuRLgs/ouIrW3XNzuNj7J3Nvu/Dig5MXvbCEdiBN3Q= +github.com/jessevdk/go-flags v1.4.0/go.mod h1:4FA24M0QyGHXBuZZK/XkWh8h0e1EYbRYJSGM75WSRxI= +github.com/jinzhu/copier v0.4.0 h1:w3ciUoD19shMCRargcpm0cm91ytaBhDvuRpz1ODO/U8= +github.com/jinzhu/copier v0.4.0/go.mod h1:DfbEm0FYsaqBcKcFuvmOZb218JkPGtvSHsKg8S8hyyg= +github.com/josharian/intern v1.0.0 h1:vlS4z54oSdjm0bgjRigI+G1HpF+tI+9rE5LLzOg8HmY= +github.com/josharian/intern v1.0.0/go.mod h1:5DoeVV0s6jJacbCEi61lwdGj/aVlrQvzHFFd8Hwg//Y= +github.com/json-iterator/go v1.1.12 h1:PV8peI4a0ysnczrg+LtxykD8LfKY9ML6u2jnxaEnrnM= +github.com/json-iterator/go v1.1.12/go.mod h1:e30LSqwooZae/UwlEbR2852Gd8hjQvJoHmT4TnhNGBo= +github.com/jtolds/gls v4.20.0+incompatible h1:xdiiI2gbIgH/gLH7ADydsJ1uDOEzR8yvV7C0MuV77Wo= +github.com/jtolds/gls v4.20.0+incompatible/go.mod h1:QJZ7F/aHp+rZTRtaJ1ow/lLfFfVYBRgL+9YlvaHOwJU= +github.com/kardianos/osext v0.0.0-20190222173326-2bc1f35cddc0/go.mod h1:1NbS8ALrpOvjt0rHPNLyCIeMtbizbir8U//inJ+zuB8= +github.com/klauspost/cpuid/v2 v2.0.9 h1:lgaqFMSdTdQYdZ04uHyN2d/eKdOMyi2YLSvlQIBFYa4= +github.com/klauspost/cpuid/v2 v2.0.9/go.mod h1:FInQzS24/EEf25PyTYn52gqo7WaD8xa0213Md/qVLRg= +github.com/knz/go-libedit v1.10.1/go.mod h1:MZTVkCWyz0oBc7JOWP3wNAzd002ZbM/5hgShxwh4x8M= +github.com/konsorten/go-windows-terminal-sequences v1.0.1/go.mod h1:T0+1ngSBFLxvqU3pZ+m/2kptfBszLMUkC4ZK/EgS/cQ= +github.com/kr/pretty v0.1.0/go.mod h1:dAy3ld7l9f0ibDNOQOHHMYYIIbhfbHSm3C4ZsoJORNo= +github.com/kr/pretty v0.3.1 h1:flRD4NNwYAUpkphVc1HcthR4KEIFJ65n8Mw5qdRn3LE= +github.com/kr/pretty v0.3.1/go.mod h1:hoEshYVHaxMs3cyo3Yncou5ZscifuDolrwPKZanG3xk= +github.com/kr/pty v1.1.1/go.mod h1:pFQYn66WHrOpPYNljwOMqo10TkYh1fy3cYio2l3bCsQ= +github.com/kr/text v0.1.0/go.mod h1:4Jbv+DJW3UT/LiOwJeYQe1efqtUx/iVham/4vfdArNI= +github.com/kr/text v0.2.0 h1:5Nx0Ya0ZqY2ygV366QzturHI13Jq95ApcVaJBhpS+AY= +github.com/kr/text v0.2.0/go.mod h1:eLer722TekiGuMkidMxC/pM04lWEeraHUUmBw8l2grE= +github.com/mailru/easyjson v0.0.0-20190614124828-94de47d64c63/go.mod h1:C1wdFJiN94OJF2b5HbByQZoLdCWB1Yqtg26g4irojpc= +github.com/mailru/easyjson v0.0.0-20190626092158-b2ccc519800e/go.mod h1:C1wdFJiN94OJF2b5HbByQZoLdCWB1Yqtg26g4irojpc= +github.com/mailru/easyjson v0.7.7 h1:UGYAvKxe3sBsEDzO8ZeWOSlIQfWFlxbzLZe7hwFURr0= +github.com/mailru/easyjson v0.7.7/go.mod h1:xzfreul335JAWq5oZzymOObrkdz5UnU4kGfJJLY9Nlc= +github.com/mattn/go-colorable v0.1.2 h1:/bC9yWikZXAL9uJdulbSfyVNIR3n3trXl+v8+1sx8mU= +github.com/mattn/go-colorable v0.1.2/go.mod h1:U0ppj6V5qS13XJ6of8GYAs25YV2eR4EVcfRqFIhoBtE= +github.com/mattn/go-isatty v0.0.8 h1:HLtExJ+uU2HOZ+wI0Tt5DtUDrx8yhUqDcp7fYERX4CE= +github.com/mattn/go-isatty v0.0.8/go.mod h1:Iq45c/XA43vh69/j3iqttzPXn0bhXyGjM0Hdxcsrc5s= +github.com/matttproud/golang_protobuf_extensions v1.0.4 h1:mmDVorXM7PCGKw94cs5zkfA9PSy5pEvNWRP0ET0TIVo= +github.com/matttproud/golang_protobuf_extensions v1.0.4/go.mod h1:BSXmuO+STAnVfrANrmjBb36TMTDstsz7MSK+HVaYKv4= +github.com/meguminnnnnnnnn/go-openai v0.0.0-20250821095446-07791bea23a0 h1:nIohpHs1ViKR0SVgW/cbBstHjmnqFZDM9RqgX9m9Xu8= +github.com/meguminnnnnnnnn/go-openai v0.0.0-20250821095446-07791bea23a0/go.mod h1:qs96ysDmxhE4BZoU45I43zcyfnaYxU3X+aRzLko/htY= +github.com/mgutz/ansi v0.0.0-20170206155736-9520e82c474b h1:j7+1HpAFS1zy5+Q4qx1fWh90gTKwiN4QCGoY9TWyyO4= +github.com/mgutz/ansi v0.0.0-20170206155736-9520e82c474b/go.mod h1:01TrycV0kFyexm33Z7vhZRXopbI8J3TDReVlkTgMUxE= +github.com/modern-go/concurrent v0.0.0-20180228061459-e0a39a4cb421/go.mod h1:6dJC0mAP4ikYIbvyc7fijjWJddQyLn8Ig3JB5CqoB9Q= +github.com/modern-go/concurrent v0.0.0-20180306012644-bacd9c7ef1dd h1:TRLaZ9cD/w8PVh93nsPXa1VrQ6jlwL5oN8l14QlcNfg= +github.com/modern-go/concurrent v0.0.0-20180306012644-bacd9c7ef1dd/go.mod h1:6dJC0mAP4ikYIbvyc7fijjWJddQyLn8Ig3JB5CqoB9Q= +github.com/modern-go/reflect2 v1.0.2 h1:xBagoLtFs94CBntxluKeaWgTMpvLxC4ur3nMaC9Gz0M= +github.com/modern-go/reflect2 v1.0.2/go.mod h1:yWuevngMOJpCy52FWWMvUC8ws7m/LJsjYzDa0/r8luk= +github.com/mohae/deepcopy v0.0.0-20170929034955-c48cc78d4826 h1:RWengNIwukTxcDr9M+97sNutRR1RKhG96O6jWumTTnw= +github.com/mohae/deepcopy v0.0.0-20170929034955-c48cc78d4826/go.mod h1:TaXosZuwdSHYgviHp1DAtfrULt5eUgsSMsZf+YrPgl8= +github.com/nikolalohinski/gonja v1.5.3 h1:GsA+EEaZDZPGJ8JtpeGN78jidhOlxeJROpqMT9fTj9c= +github.com/nikolalohinski/gonja v1.5.3/go.mod h1:RmjwxNiXAEqcq1HeK5SSMmqFJvKOfTfXhkJv6YBtPa4= +github.com/nxadm/tail v1.4.4 h1:DQuhQpB1tVlglWS2hLQ5OV6B5r8aGxSrPc5Qo6uTN78= +github.com/nxadm/tail v1.4.4/go.mod h1:kenIhsEOeOJmVchQTgglprH7qJGnHDVpk1VPCcaMI8A= +github.com/onsi/ginkgo v1.6.0/go.mod h1:lLunBs/Ym6LB5Z9jYTR76FiuTmxDTDusOGeTQH+WWjE= +github.com/onsi/ginkgo v1.8.0/go.mod h1:lLunBs/Ym6LB5Z9jYTR76FiuTmxDTDusOGeTQH+WWjE= +github.com/onsi/ginkgo v1.12.1/go.mod h1:zj2OWP4+oCPe1qIXoGWkgMRwljMUYCdkwsT2108oapk= +github.com/onsi/ginkgo v1.15.0/go.mod h1:hF8qUzuuC8DJGygJH3726JnCZX4MYbRB8yFfISqnKUg= +github.com/onsi/ginkgo v1.16.5 h1:8xi0RTUf59SOSfEtZMvwTvXYMzG4gV23XVHOZiXNtnE= +github.com/onsi/ginkgo v1.16.5/go.mod h1:+E8gABHa3K6zRBolWtd+ROzc/U5bkGt0FwiG042wbpU= +github.com/onsi/gomega v1.5.0/go.mod h1:ex+gbHU/CVuBBDIJjb2X0qEXbFg53c61hWP/1CpauHY= +github.com/onsi/gomega v1.7.1/go.mod h1:XdKZgCCFLUoM/7CFJVPcG8C1xQ1AJ0vpAezJrB7JYyY= +github.com/onsi/gomega v1.10.1/go.mod h1:iN09h71vgCQne3DLsj+A5owkum+a2tYe+TOCB1ybHNo= +github.com/onsi/gomega v1.10.5/go.mod h1:gza4q3jKQJijlu05nKWRCW/GavJumGt8aNRxWg7mt48= +github.com/onsi/gomega v1.27.3 h1:5VwIwnBY3vbBDOJrNtA4rVdiTZCsq9B5F12pvy1Drmk= +github.com/onsi/gomega v1.27.3/go.mod h1:5vG284IBtfDAmDyrK+eGyZmUgUlmi+Wngqo557cZ6Gw= +github.com/pelletier/go-toml/v2 v2.0.9 h1:uH2qQXheeefCCkuBBSLi7jCiSmj3VRh2+Goq2N7Xxu0= +github.com/pelletier/go-toml/v2 v2.0.9/go.mod h1:tJU2Z3ZkXwnxa4DPO899bsyIoywizdUvyaeZurnPPDc= +github.com/perimeterx/marshmallow v1.1.4 h1:pZLDH9RjlLGGorbXhcaQLhfuV0pFMNfPO55FuFkxqLw= +github.com/perimeterx/marshmallow v1.1.4/go.mod h1:dsXbUu8CRzfYP5a87xpp0xq9S3u0Vchtcl8we9tYaXw= +github.com/pkg/errors v0.8.0/go.mod h1:bwawxfHBFNV+L2hUp1rHADufV3IMtnDRdf1r5NINEl0= +github.com/pkg/errors v0.9.1 h1:FEBLx1zS214owpjy7qsBeixbURkuhQAwrK5UwLGTwt4= +github.com/pkg/errors v0.9.1/go.mod h1:bwawxfHBFNV+L2hUp1rHADufV3IMtnDRdf1r5NINEl0= +github.com/pkoukk/tiktoken-go v0.1.7 h1:qOBHXX4PHtvIvmOtyg1EeKlwFRiMKAcoMp4Q+bLQDmw= +github.com/pkoukk/tiktoken-go v0.1.7/go.mod h1:9NiV+i9mJKGj1rYOT+njbv+ZwA/zJxYdewGl6qVatpg= +github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM= +github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4= +github.com/prometheus/client_golang v1.16.0 h1:yk/hx9hDbrGHovbci4BY+pRMfSuuat626eFsHb7tmT8= +github.com/prometheus/client_golang v1.16.0/go.mod h1:Zsulrv/L9oM40tJ7T815tM89lFEugiJ9HzIqaAx4LKc= +github.com/prometheus/client_model v0.3.0 h1:UBgGFHqYdG/TPFD1B1ogZywDqEkwp3fBMvqdiQ7Xew4= +github.com/prometheus/client_model v0.3.0/go.mod h1:LDGWKZIo7rky3hgvBe+caln+Dr3dPggB5dvjtD7w9+w= +github.com/prometheus/common v0.42.0 h1:EKsfXEYo4JpWMHH5cg+KOUWeuJSov1Id8zGR8eeI1YM= +github.com/prometheus/common v0.42.0/go.mod h1:xBwqVerjNdUDjgODMpudtOMwlOwf2SaTr1yjz4b7Zbc= +github.com/prometheus/procfs v0.10.1 h1:kYK1Va/YMlutzCGazswoHKo//tZVlFpKYh+PymziUAg= +github.com/prometheus/procfs v0.10.1/go.mod h1:nwNm2aOCAYw8uTR/9bWRREkZFxAUcWzPHWJq+XBB/FM= +github.com/rogpeppe/go-internal v1.13.1 h1:KvO1DLK/DRN07sQ1LQKScxyZJuNnedQ5/wKSR38lUII= +github.com/rogpeppe/go-internal v1.13.1/go.mod h1:uMEvuHeurkdAXX61udpOXGD/AzZDWNMNyH2VO9fmH0o= +github.com/rollbar/rollbar-go v1.0.2/go.mod h1:AcFs5f0I+c71bpHlXNNDbOWJiKwjFDtISeXco0L5PKQ= +github.com/shirou/gopsutil/v3 v3.21.6 h1:vU7jrp1Ic/2sHB7w6UNs7MIkn7ebVtTb5D9j45o9VYE= +github.com/shirou/gopsutil/v3 v3.21.6/go.mod h1:JfVbDpIBLVzT8oKbvMg9P3wEIMDDpVn+LwHTKj0ST88= +github.com/sirupsen/logrus v1.2.0/go.mod h1:LxeOpSwHxABJmUn/MG1IvRgCAasNZTLOkJPxbbu5VWo= +github.com/sirupsen/logrus v1.9.3 h1:dueUQJ1C2q9oE3F7wvmSGAaVtTmUizReu6fjN8uqzbQ= +github.com/sirupsen/logrus v1.9.3/go.mod h1:naHLuLoDiP4jHNo9R0sCBMtWGeIprob74mVsIT4qYEQ= +github.com/slongfield/pyfmt v0.0.0-20220222012616-ea85ff4c361f h1:Z2cODYsUxQPofhpYRMQVwWz4yUVpHF+vPi+eUdruUYI= +github.com/slongfield/pyfmt v0.0.0-20220222012616-ea85ff4c361f/go.mod h1:JqzWyvTuI2X4+9wOHmKSQCYxybB/8j6Ko43qVmXDuZg= +github.com/smarty/assertions v1.15.0 h1:cR//PqUBUiQRakZWqBiFFQ9wb8emQGDb0HeGdqGByCY= +github.com/smarty/assertions v1.15.0/go.mod h1:yABtdzeQs6l1brC900WlRNwj6ZR55d7B+E8C6HtKdec= +github.com/smartystreets/goconvey v1.8.1 h1:qGjIddxOk4grTu9JPOU31tVfq3cNdBlNa5sSznIX1xY= +github.com/smartystreets/goconvey v1.8.1/go.mod h1:+/u4qLyY6x1jReYOp7GOM2FSt8aP9CzCZL03bI28W60= +github.com/spaolacci/murmur3 v1.1.0 h1:7c1g84S4BPRrfL5Xrdp6fOJ206sU9y293DDHaoy0bLI= +github.com/spaolacci/murmur3 v1.1.0/go.mod h1:JwIasOWyU6f++ZhiEuf87xNszmSA2myDM2Kzu9HwQUA= +github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME= +github.com/stretchr/objx v0.1.1/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME= +github.com/stretchr/objx v0.4.0/go.mod h1:YvHI0jy2hoMjB+UWwv71VJQ9isScKT/TqJzVSSt89Yw= +github.com/stretchr/objx v0.5.0/go.mod h1:Yh+to48EsGEfYuaHDzXPcE3xhTkx73EhmCGUpEOglKo= +github.com/stretchr/objx v0.5.2 h1:xuMeJ0Sdp5ZMRXx/aWO6RZxdr3beISkG5/G/aIRr3pY= +github.com/stretchr/objx v0.5.2/go.mod h1:FRsXN1f5AsAjCGJKqEizvkpNtU+EGNCLh3NxZ/8L+MA= +github.com/stretchr/testify v1.2.2/go.mod h1:a8OnRcib4nhh0OaRAV+Yts87kKdq0PP7pXfy6kDkUVs= +github.com/stretchr/testify v1.3.0/go.mod h1:M5WIy9Dh21IEIfnGCwXGc5bZfKNJtfHm1UVUgZn+9EI= +github.com/stretchr/testify v1.6.1/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg= +github.com/stretchr/testify v1.7.0/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg= +github.com/stretchr/testify v1.7.1/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg= +github.com/stretchr/testify v1.8.0/go.mod h1:yNjHg4UonilssWZ8iaSj1OCr/vHnekPRkoO+kdMU+MU= +github.com/stretchr/testify v1.8.1/go.mod h1:w2LPCIKwWwSfY2zedu0+kehJoqGctiVI29o6fzry7u4= +github.com/stretchr/testify v1.8.4/go.mod h1:sz/lmYIOXD/1dqDmKjjqLyZ2RngseejIcXlSw2iwfAo= +github.com/stretchr/testify v1.10.0 h1:Xv5erBjTwe/5IxqUQTdXv5kgmIvbHo3QQyRwhJsOfJA= +github.com/stretchr/testify v1.10.0/go.mod h1:r2ic/lqez/lEtzL7wO/rwa5dbSLXVDPFyf8C91i36aY= +github.com/tklauser/go-sysconf v0.3.6 h1:oc1sJWvKkmvIxhDHeKWvZS4f6AW+YcoguSfRF2/Hmo4= +github.com/tklauser/go-sysconf v0.3.6/go.mod h1:MkWzOF4RMCshBAMXuhXJs64Rte09mITnppBXY/rYEFI= +github.com/tklauser/numcpus v0.2.2 h1:oyhllyrScuYI6g+h/zUvNXNp1wy7x8qQy3t/piefldA= +github.com/tklauser/numcpus v0.2.2/go.mod h1:x3qojaO3uyYt0i56EW/VUYs7uBvdl2fkfZFu0T9wgjM= +github.com/twitchyliquid64/golang-asm v0.15.1 h1:SU5vSMR7hnwNxj24w34ZyCi/FmDZTkS4MhqMhdFk5YI= +github.com/twitchyliquid64/golang-asm v0.15.1/go.mod h1:a1lVb/DtPvCB8fslRZhAngC2+aY1QWCk3Cedj/Gdt08= +github.com/ugorji/go v1.2.7 h1:qYhyWUUd6WbiM+C6JZAUkIJt/1WrjzNHY9+KCIjVqTo= +github.com/ugorji/go v1.2.7/go.mod h1:nF9osbDWLy6bDVv/Rtoh6QgnvNDpmCalQV5urGCCS6M= +github.com/ugorji/go/codec v1.2.7 h1:YPXUKf7fYbp/y8xloBqZOw2qaVggbfwMlI8WM3wZUJ0= +github.com/ugorji/go/codec v1.2.7/go.mod h1:WGN1fab3R1fzQlVQTkfxVtIBhWDRqOviHU95kRgeqEY= +github.com/wk8/go-ordered-map/v2 v2.1.8 h1:5h/BUHu93oj4gIdvHHHGsScSTMijfx5PeYkE/fJgbpc= +github.com/wk8/go-ordered-map/v2 v2.1.8/go.mod h1:5nJHM5DyteebpVlHnWMV0rPz6Zp7+xBAnxjb1X5vnTw= +github.com/x-cray/logrus-prefixed-formatter v0.5.2 h1:00txxvfBM9muc0jiLIEAkAcIMJzfthRT6usrui8uGmg= +github.com/x-cray/logrus-prefixed-formatter v0.5.2/go.mod h1:2duySbKsL6M18s5GU7VPsoEPHyzalCE06qoARUCeBBE= +github.com/yargevad/filepathx v1.0.0 h1:SYcT+N3tYGi+NvazubCNlvgIPbzAk7i7y2dwg3I5FYc= +github.com/yargevad/filepathx v1.0.0/go.mod h1:BprfX/gpYNJHJfc35GjRRpVcwWXS89gGulUIU5tK3tA= +github.com/yuin/goldmark v1.2.1/go.mod h1:3hX8gzYuyVAZsxl0MRgGTJEmQBFcNTphYh9decYSb74= +golang.org/x/arch v0.11.0 h1:KXV8WWKCXm6tRpLirl2szsO5j/oOODwZf4hATmGVNs4= +golang.org/x/arch v0.11.0/go.mod h1:FEVrYAQjsQXMVJ1nsMoVVXPZg6p2JE2mx8psSWTDQys= +golang.org/x/crypto v0.0.0-20180904163835-0709b304e793/go.mod h1:6SG95UA2DQfeDnfUPMdvaQW0Q7yPrPDi9nlGo2tz2b4= +golang.org/x/crypto v0.0.0-20190308221718-c2843e01d9a2/go.mod h1:djNgcEr1/C05ACkg1iLfiJU5Ep61QUkGW8qpdssI0+w= +golang.org/x/crypto v0.0.0-20191011191535-87dc89f01550/go.mod h1:yigFU9vqHzYiE8UmvKecakEJjdnWj3jj499lnFckfCI= +golang.org/x/crypto v0.0.0-20200622213623-75b288015ac9/go.mod h1:LzIPMQfyMNhhGPhUkYOs5KpL4U8rLKemX1yGLhDgUto= +golang.org/x/crypto v0.31.0 h1:ihbySMvVjLAeSH1IbfcRTkD/iNscyz8rGzjF/E5hV6U= +golang.org/x/crypto v0.31.0/go.mod h1:kDsLvtWBEx7MV9tJOj9bnXsPbxwJQ6csT/x4KIN4Ssk= +golang.org/x/exp v0.0.0-20230713183714-613f0c0eb8a1 h1:MGwJjxBy0HJshjDNfLsYO8xppfqWlA5ZT9OhtUUhTNw= +golang.org/x/exp v0.0.0-20230713183714-613f0c0eb8a1/go.mod h1:FXUEEKJgO7OQYeo8N01OfiKP8RXMtf6e8aTskBGqWdc= +golang.org/x/mod v0.3.0/go.mod h1:s0Qsj1ACt9ePp/hMypM3fl4fZqREWJwdYDEqhRiZZUA= +golang.org/x/net v0.0.0-20180906233101-161cd47e91fd/go.mod h1:mL1N/T3taQHkDXs73rZJwtUhF3w3ftmwwsq0BUmARs4= +golang.org/x/net v0.0.0-20190404232315-eb5bcb51f2a3/go.mod h1:t9HGtf8HONx5eT2rtn7q6eTqICYqUVnKs3thJo3Qplg= +golang.org/x/net v0.0.0-20190620200207-3b0461eec859/go.mod h1:z5CRVTTTmAJ677TzLLGU+0bjPO0LkuOLi4/5GtJWs/s= +golang.org/x/net v0.0.0-20200520004742-59133d7f0dd7/go.mod h1:qpuaurCH72eLCgpAm/N6yyVIVM9cpaDIP3A8BGJEC5A= +golang.org/x/net v0.0.0-20201021035429-f5854403a974/go.mod h1:sp8m0HH+o8qH0wwXwYZr8TS3Oi6o0r6Gce1SSxlDquU= +golang.org/x/net v0.0.0-20201202161906-c7110b5ffcbb/go.mod h1:sp8m0HH+o8qH0wwXwYZr8TS3Oi6o0r6Gce1SSxlDquU= +golang.org/x/net v0.21.0 h1:AQyQV4dYCvJ7vGmJyKki9+PBdyvhkSd8EIx/qb0AYv4= +golang.org/x/net v0.21.0/go.mod h1:bIjVDfnllIU7BJ2DNgfnXvpSvtn8VRwhlsaeUTyUS44= +golang.org/x/sync v0.0.0-20180314180146-1d60e4601c6f/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM= +golang.org/x/sync v0.0.0-20181221193216-37e7f081c4d4/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM= +golang.org/x/sync v0.0.0-20190423024810-112230192c58/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM= +golang.org/x/sync v0.0.0-20201020160332-67f06af15bc9/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM= +golang.org/x/sys v0.0.0-20180905080454-ebe1bf3edb33/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY= +golang.org/x/sys v0.0.0-20180909124046-d0be0721c37e/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY= +golang.org/x/sys v0.0.0-20190215142949-d0b11bdaac8a/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY= +golang.org/x/sys v0.0.0-20190412213103-97732733099d/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs= +golang.org/x/sys v0.0.0-20190904154756-749cb33beabd/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs= +golang.org/x/sys v0.0.0-20191005200804-aed5e4c7ecf9/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs= +golang.org/x/sys v0.0.0-20191120155948-bd437916bb0e/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs= +golang.org/x/sys v0.0.0-20200323222414-85ca7c5b95cd/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs= +golang.org/x/sys v0.0.0-20200930185726-fdedc70b468f/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs= +golang.org/x/sys v0.0.0-20210112080510-489259a85091/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs= +golang.org/x/sys v0.0.0-20210316164454-77fc1eacc6aa/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs= +golang.org/x/sys v0.0.0-20220715151400-c0bba94af5f8/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= +golang.org/x/sys v0.28.0 h1:Fksou7UEQUWlKvIdsqzJmUmCX3cZuD2+P3XyyzwMhlA= +golang.org/x/sys v0.28.0/go.mod h1:/VUhepiaJMQUp4+oa/7Zr1D23ma6VTLIYjOOTFZPUcA= +golang.org/x/term v0.27.0 h1:WP60Sv1nlK1T6SupCHbXzSaN0b9wUmsPoRS9b61A23Q= +golang.org/x/term v0.27.0/go.mod h1:iMsnZpn0cago0GOrHO2+Y7u7JPn5AylBrcoWkElMTSM= +golang.org/x/text v0.3.0/go.mod h1:NqM8EUOU14njkJ3fqMW+pc6Ldnwhi/IjpwHt7yyuwOQ= +golang.org/x/text v0.3.3/go.mod h1:5Zoc/QRtKVWzQhOtBMvqHzDpF6irO9z98xDceosuGiQ= +golang.org/x/text v0.21.0 h1:zyQAAkrwaneQ066sspRyJaG9VNi/YJ1NfzcGB3hZ/qo= +golang.org/x/text v0.21.0/go.mod h1:4IBbMaMmOPCJ8SecivzSH54+73PCFmPWxNTLm+vZkEQ= +golang.org/x/tools v0.0.0-20180917221912-90fa682c2a6e/go.mod h1:n7NCudcB/nEzxVGmLbDWY5pfWTLqBcC2KZ6jyYvM4mQ= +golang.org/x/tools v0.0.0-20191119224855-298f0cb1881e/go.mod h1:b+2E5dAYhXwXZwtnZ6UAqBI28+e2cm9otk0dWdXHAEo= +golang.org/x/tools v0.0.0-20201224043029-2b0845dc783e/go.mod h1:emZCQorbCU4vsT4fOWvOPXz4eW1wZW4PmDk9uLelYpA= +golang.org/x/xerrors v0.0.0-20190717185122-a985d3407aa7/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0= +golang.org/x/xerrors v0.0.0-20191011141410-1b5146add898/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0= +golang.org/x/xerrors v0.0.0-20191204190536-9bdfabe68543/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0= +golang.org/x/xerrors v0.0.0-20200804184101-5ec99f83aff1/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0= +google.golang.org/protobuf v0.0.0-20200109180630-ec00e32a8dfd/go.mod h1:DFci5gLYBciE7Vtevhsrf46CRTquxDuWsQurQQe4oz8= +google.golang.org/protobuf v0.0.0-20200221191635-4d8936d0db64/go.mod h1:kwYJMbMJ01Woi6D6+Kah6886xMZcty6N08ah7+eCXa0= +google.golang.org/protobuf v0.0.0-20200228230310-ab0ca4ff8a60/go.mod h1:cfTl7dwQJ+fmap5saPgwCLgHXTUD7jkjRqWcaiX5VyM= +google.golang.org/protobuf v1.20.1-0.20200309200217-e05f789c0967/go.mod h1:A+miEFZTKqfCUM6K7xSMQL9OKL/b6hQv+e19PK+JZNE= +google.golang.org/protobuf v1.21.0/go.mod h1:47Nbq4nVaFHyn7ilMalzfO3qCViNmqZ2kzikPIcrTAo= +google.golang.org/protobuf v1.23.0/go.mod h1:EGpADcykh3NcUnDUJcl1+ZksZNG86OlYog2l/sGQquU= +google.golang.org/protobuf v1.26.0-rc.1/go.mod h1:jlhhOSvTdKEhbULTjvd4ARK9grFBp09yW+WbY/TyQbw= +google.golang.org/protobuf v1.26.0/go.mod h1:9q0QmTI4eRPtz6boOQmLYwt+qCgq0jsYwAQnmE0givc= +google.golang.org/protobuf v1.30.0 h1:kPPoIgf3TsEvrm0PFe15JQ+570QVxYzEvvHqChK+cng= +google.golang.org/protobuf v1.30.0/go.mod h1:HV8QOd/L58Z+nl8r43ehVNZIU/HEI6OcFqwMG9pJV4I= +gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0= +gopkg.in/check.v1 v1.0.0-20180628173108-788fd7840127/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0= +gopkg.in/check.v1 v1.0.0-20201130134442-10cb98267c6c h1:Hei/4ADfdWqJk1ZMxUNpqntNwaWcugrBjAiHlqqRiVk= +gopkg.in/check.v1 v1.0.0-20201130134442-10cb98267c6c/go.mod h1:JHkPIbrfpd72SG/EVd6muEfDQjcINNoR0C8j2r3qZ4Q= +gopkg.in/fsnotify.v1 v1.4.7/go.mod h1:Tz8NjZHkW78fSQdbUxIjBTcgA1z1m8ZHf0WmKUhAMys= +gopkg.in/tomb.v1 v1.0.0-20141024135613-dd632973f1e7 h1:uRGJdciOHaEIrze2W8Q3AKkepLTh2hOroT7a+7czfdQ= +gopkg.in/tomb.v1 v1.0.0-20141024135613-dd632973f1e7/go.mod h1:dt/ZhP58zS4L8KSrWDmTeBkI65Dw0HsyUHuEVlX15mw= +gopkg.in/yaml.v2 v2.2.1/go.mod h1:hI93XBmqTisBFMUTm0b8Fm+jr3Dg1NNxqwp+5A1VGuI= +gopkg.in/yaml.v2 v2.2.2/go.mod h1:hI93XBmqTisBFMUTm0b8Fm+jr3Dg1NNxqwp+5A1VGuI= +gopkg.in/yaml.v2 v2.2.4/go.mod h1:hI93XBmqTisBFMUTm0b8Fm+jr3Dg1NNxqwp+5A1VGuI= +gopkg.in/yaml.v2 v2.3.0/go.mod h1:hI93XBmqTisBFMUTm0b8Fm+jr3Dg1NNxqwp+5A1VGuI= +gopkg.in/yaml.v2 v2.4.0 h1:D8xgwECY7CYvx+Y2n4sBz93Jn9JRvxdiyyo8CTfuKaY= +gopkg.in/yaml.v2 v2.4.0/go.mod h1:RDklbk79AGWmwhnvt/jBztapEOGDOx6ZbXqjP6csGnQ= +gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM= +gopkg.in/yaml.v3 v3.0.0/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM= +gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA= +gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM= +nullprogram.com/x/optparse v1.0.0/go.mod h1:KdyPE+Igbe0jQUrVfMqDMeJQIJZEuyV7pjYmp6pbG50= diff --git a/pkg/adapters/eino/options.go b/pkg/adapters/eino/options.go new file mode 100644 index 000000000..842d0da9e --- /dev/null +++ b/pkg/adapters/eino/options.go @@ -0,0 +1,93 @@ +// Copyright 1999-2020 Alibaba Group Holding Ltd. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package eino + +import ( + "context" + + llmtokenratelimit "github.com/alibaba/sentinel-golang/core/llm_token_ratelimit" + "github.com/cloudwego/eino/schema" +) + +type Option func(*options) + +type options struct { + defaultResource string + resourceExtract func(context.Context) string + blockFallback func(context.Context) + requestInfosExtract func(context.Context) *llmtokenratelimit.RequestInfos + promptsExtract func([]*schema.Message) []string + usedTokenInfosExtract func(interface{}) *llmtokenratelimit.UsedTokenInfos +} + +func WithDefaultResource(resource string) Option { + return func(o *options) { + o.defaultResource = resource + } +} + +func WithResourceExtract(fn func(context.Context) string) Option { + return func(o *options) { + o.resourceExtract = fn + } +} + +func WithBlockFallback(fn func(context.Context)) Option { + return func(o *options) { + o.blockFallback = fn + } +} + +func WithRequestInfosExtract(fn func(context.Context) *llmtokenratelimit.RequestInfos) Option { + return func(o *options) { + o.requestInfosExtract = fn + } +} + +func WithPromptsExtract(fn func([]*schema.Message) []string) Option { + return func(o *options) { + o.promptsExtract = fn + } +} + +func WithUsedTokenInfosExtract(fn func(interface{}) *llmtokenratelimit.UsedTokenInfos) Option { + return func(o *options) { + o.usedTokenInfosExtract = fn + } +} + +func evaluateOptions(opts ...Option) *options { + optCopy := &options{ + defaultResource: llmtokenratelimit.DefaultResourcePattern, + promptsExtract: func(messages []*schema.Message) []string { + prompts := make([]string, 0, len(messages)) + for _, msg := range messages { + for _, part := range msg.MultiContent { + prompts = append(prompts, part.Text) + } + if len(msg.Content) != 0 { + prompts = append(prompts, msg.Content) + } + } + return prompts + }, + } + for _, opt := range opts { + if opt != nil { + opt(optCopy) + } + } + return optCopy +} diff --git a/pkg/adapters/eino/wrapper.go b/pkg/adapters/eino/wrapper.go new file mode 100644 index 000000000..a30828d6f --- /dev/null +++ b/pkg/adapters/eino/wrapper.go @@ -0,0 +1,104 @@ +// Copyright 1999-2020 Alibaba Group Holding Ltd. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package eino + +import ( + "context" + "fmt" + + sentinel "github.com/alibaba/sentinel-golang/api" + "github.com/alibaba/sentinel-golang/core/base" + llmtokenratelimit "github.com/alibaba/sentinel-golang/core/llm_token_ratelimit" + "github.com/cloudwego/eino/components/model" + "github.com/cloudwego/eino/schema" +) + +type LLMWrapper struct { + llm model.BaseChatModel + options *options +} + +func NewLLMWrapper(llm model.BaseChatModel, opts ...Option) *LLMWrapper { + return &LLMWrapper{ + llm: llm, + options: evaluateOptions(opts...), + } +} + +func (w *LLMWrapper) Generate(ctx context.Context, input []*schema.Message, opts ...model.Option) (*schema.Message, error) { + resource := w.options.defaultResource + if w.options.resourceExtract != nil { + resource = w.options.resourceExtract(ctx) + } + + prompts := []string{} + if w.options.promptsExtract != nil { + prompts = w.options.promptsExtract(input) + } + + reqInfos := llmtokenratelimit.GenerateRequestInfos( + llmtokenratelimit.WithPrompts(prompts), + ) + if w.options.requestInfosExtract != nil { + reqInfos = w.options.requestInfosExtract(ctx) + } + + // Check + entry, err := sentinel.Entry(resource, sentinel.WithTrafficType(base.Inbound), sentinel.WithArgs(reqInfos)) + + if err != nil { + // Block + if w.options.blockFallback != nil { + w.options.blockFallback(ctx) + } + return nil, err + } + // Pass + response, llmErr := w.llm.Generate(ctx, input, opts...) + if llmErr != nil { + return nil, llmErr + } + if err := w.validateResponse(response); err != nil { + return nil, err + } + + usedTokenInfos := &llmtokenratelimit.UsedTokenInfos{} + if w.options.usedTokenInfosExtract != nil { + usedTokenInfos = w.options.usedTokenInfosExtract(response.ResponseMeta.Usage) + } else { + // fallback to OpenAI extractor + infos, err := llmtokenratelimit.OpenAITokenExtractor(map[string]any{ + "prompt_tokens": response.ResponseMeta.Usage.PromptTokens, + "completion_tokens": response.ResponseMeta.Usage.CompletionTokens, + "total_tokens": response.ResponseMeta.Usage.TotalTokens, + }) + if err != nil { + return nil, err + } + usedTokenInfos = infos + } + + entry.SetPair(llmtokenratelimit.KeyUsedTokenInfos, usedTokenInfos) + entry.Exit() + + return response, nil +} + +func (w *LLMWrapper) validateResponse(response *schema.Message) error { + if response == nil || response.ResponseMeta == nil || response.ResponseMeta.Usage == nil { + return fmt.Errorf("llm response is nil or empty or missing Usage infos") + } + return nil +} diff --git a/pkg/adapters/kitex/go.mod b/pkg/adapters/kitex/go.mod index a3e93ce6a..27bf010a4 100644 --- a/pkg/adapters/kitex/go.mod +++ b/pkg/adapters/kitex/go.mod @@ -30,13 +30,17 @@ require ( github.com/cloudwego/netpoll v0.5.1 // indirect github.com/cloudwego/thriftgo v0.3.3 // indirect github.com/davecgh/go-spew v1.1.2-0.20180830191138-d8f796af33cc // indirect + github.com/dgryski/go-rendezvous v0.0.0-20200823014737-9f7001d12a5f // indirect + github.com/dlclark/regexp2 v1.10.0 // indirect github.com/fatih/structtag v1.2.0 // indirect github.com/go-ole/go-ole v1.2.4 // indirect + github.com/go-redis/redis/v8 v8.11.0 // indirect github.com/golang/protobuf v1.5.3 // indirect github.com/google/pprof v0.0.0-20220608213341-c488b8fa1db3 // indirect - github.com/google/uuid v1.1.2 // indirect + github.com/google/uuid v1.3.0 // indirect github.com/iancoleman/strcase v0.2.0 // indirect github.com/jhump/protoreflect v1.8.2 // indirect + github.com/jinzhu/copier v0.4.0 // indirect github.com/json-iterator/go v1.1.12 // indirect github.com/klauspost/cpuid/v2 v2.2.4 // indirect github.com/matttproud/golang_protobuf_extensions v1.0.4 // indirect @@ -45,12 +49,14 @@ require ( github.com/modern-go/reflect2 v1.0.2 // indirect github.com/oleiade/lane v1.0.1 // indirect github.com/pkg/errors v0.9.1 // indirect + github.com/pkoukk/tiktoken-go v0.1.7 // indirect github.com/pmezard/go-difflib v1.0.0 // indirect github.com/prometheus/client_golang v1.16.0 // indirect github.com/prometheus/client_model v0.3.0 // indirect github.com/prometheus/common v0.42.0 // indirect github.com/prometheus/procfs v0.10.1 // indirect github.com/shirou/gopsutil/v3 v3.21.6 // indirect + github.com/spaolacci/murmur3 v1.1.0 // indirect github.com/tidwall/gjson v1.9.3 // indirect github.com/tidwall/match v1.1.1 // indirect github.com/tidwall/pretty v1.2.0 // indirect diff --git a/pkg/adapters/kitex/go.sum b/pkg/adapters/kitex/go.sum index 2fb9c6039..985f02df5 100644 --- a/pkg/adapters/kitex/go.sum +++ b/pkg/adapters/kitex/go.sum @@ -39,6 +39,7 @@ github.com/bytedance/sonic v1.10.2 h1:GQebETVBxYB7JGWJtLBi07OVzWwt+8dWA00gEVW2ZF github.com/bytedance/sonic v1.10.2/go.mod h1:iZcSUejdk5aukTND/Eu/ivjQuEL0Cu9/rf50Hi0u/g4= github.com/census-instrumentation/opencensus-proto v0.2.1/go.mod h1:f6KPmirojxKA12rnyqOA5BBL4O983OfeGPqjHWSTneU= github.com/cespare/xxhash v1.1.0/go.mod h1:XrSqR1VqqWfGrhpAt58auRo0WTKS1nRRg3ghfAqPWnc= +github.com/cespare/xxhash/v2 v2.1.1/go.mod h1:VGX0DQ3Q6kWi7AoAeZDth3/j3BFtOZR5XLFGgcrjCOs= github.com/cespare/xxhash/v2 v2.2.0 h1:DC2CZ1Ep5Y4k3ZQ899DldepgrayRUGE6BBZ/cd9Cj44= github.com/cespare/xxhash/v2 v2.2.0/go.mod h1:VGX0DQ3Q6kWi7AoAeZDth3/j3BFtOZR5XLFGgcrjCOs= github.com/chenzhuoyu/base64x v0.0.0-20211019084208-fb5309c8db06/go.mod h1:DH46F32mSOjUmXrMHnKwZdA8wcEefY7UVqBKYGjpdQY= @@ -106,6 +107,9 @@ github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSs github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= github.com/davecgh/go-spew v1.1.2-0.20180830191138-d8f796af33cc h1:U9qPSI2PIWSS1VwoXQT9A3Wy9MM3WgvqSxFWenqJduM= github.com/davecgh/go-spew v1.1.2-0.20180830191138-d8f796af33cc/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= +github.com/dgryski/go-rendezvous v0.0.0-20200823014737-9f7001d12a5f h1:lO4WD4F/rVNCu3HqELle0jiPLLBs70cWOduZpkS1E78= +github.com/dgryski/go-rendezvous v0.0.0-20200823014737-9f7001d12a5f/go.mod h1:cuUVRXasLTGF7a8hSLbxyZXjz+1KgoB3wDUb6vlszIc= +github.com/dlclark/regexp2 v1.10.0 h1:+/GIL799phkJqYW+3YbOd8LCcbHzT0Pbo8zl70MHsq0= github.com/dlclark/regexp2 v1.10.0/go.mod h1:DHkYz0B9wPfa6wondMfaivmHpzrQ3v9q8cnmRbL6yW8= github.com/envoyproxy/go-control-plane v0.9.0/go.mod h1:YTl/9mNaCwkRvm6d1a2C3ymFceY/DCBVvsKhRF0iEA4= github.com/envoyproxy/go-control-plane v0.9.1-0.20191026205805-5f8ba28d4473/go.mod h1:YTl/9mNaCwkRvm6d1a2C3ymFceY/DCBVvsKhRF0iEA4= @@ -115,6 +119,9 @@ github.com/fatih/structtag v1.2.0 h1:/OdNE99OxoI/PqaW/SuSK9uxxT3f/tcSZgon/ssNSx4 github.com/fatih/structtag v1.2.0/go.mod h1:mBJUNpUnHmRKrKlQQlmCrh5PuhftFbNv8Ys4/aAZl94= github.com/fogleman/gg v1.2.1-0.20190220221249-0403632d5b90/go.mod h1:R/bRT+9gY/C5z7JzPU0zXsXHKM4/ayA+zqcVNZzPa1k= github.com/fogleman/gg v1.3.0/go.mod h1:R/bRT+9gY/C5z7JzPU0zXsXHKM4/ayA+zqcVNZzPa1k= +github.com/fsnotify/fsnotify v1.4.7/go.mod h1:jwhsz4b93w/PPRr/qN1Yymfu8t87LnFCMoQvtojpjFo= +github.com/fsnotify/fsnotify v1.4.9 h1:hsms1Qyu0jgnwNXIxa+/V/PDsU6CfLf6CNO8H7IWoS4= +github.com/fsnotify/fsnotify v1.4.9/go.mod h1:znqG4EE+3YCdAaPaxE2ZRY/06pZUdp0tY4IgpuI1SZQ= github.com/go-fonts/dejavu v0.1.0/go.mod h1:4Wt4I4OU2Nq9asgDCteaAaWZOV24E+0/Pwo0gppep4g= github.com/go-fonts/latin-modern v0.2.0/go.mod h1:rQVLdDMK+mK1xscDwsqM5J8U2jrRa3T0ecnM9pNujks= github.com/go-fonts/liberation v0.1.1/go.mod h1:K6qoJYypsmfVjWg8KOVDQhLc8UDgIK2HYqyqAO9z7GY= @@ -127,6 +134,8 @@ github.com/go-ole/go-ole v1.2.4 h1:nNBDSCOigTSiarFpYE9J/KtEA1IOW4CNeqT9TQDqCxI= github.com/go-ole/go-ole v1.2.4/go.mod h1:XCwSNxSkXRo4vlyPy93sltvi/qJq0jqQhjqQNIwKuxM= github.com/go-pdf/fpdf v0.5.0/go.mod h1:HzcnA+A23uwogo0tp9yU+l3V+KXhiESpt1PMayhOh5M= github.com/go-pdf/fpdf v0.6.0/go.mod h1:HzcnA+A23uwogo0tp9yU+l3V+KXhiESpt1PMayhOh5M= +github.com/go-redis/redis/v8 v8.11.0 h1:O1Td0mQ8UFChQ3N9zFQqo6kTU2cJ+/it88gDB+zg0wo= +github.com/go-redis/redis/v8 v8.11.0/go.mod h1:DLomh7y2e3ggQXQLd1YgmvIfecPJoFl7WU5SOQ/r06M= github.com/gogo/protobuf v1.3.2/go.mod h1:P1XiOD3dCwIKUDQYPy72D8LYyHL2YPYrpS2s69NZV8Q= github.com/golang/freetype v0.0.0-20170609003504-e2365dfdc4a0/go.mod h1:E/TSTwGwJL78qG/PmXZO1EjYhfJinVAhrmmHX6Z8B9k= github.com/golang/glog v0.0.0-20160126235308-23def4e6c14b/go.mod h1:SBH7ygxi8pfUlaOkMMuAQtPIUF8ecWP5IEl/CR7VP2Q= @@ -154,22 +163,27 @@ github.com/google/go-cmp v0.4.0/go.mod h1:v8dTdLbMG2kIc/vJvl+f65V22dbkXbowE6jgT/ github.com/google/go-cmp v0.5.0/go.mod h1:v8dTdLbMG2kIc/vJvl+f65V22dbkXbowE6jgT/gNBxE= github.com/google/go-cmp v0.5.4/go.mod h1:v8dTdLbMG2kIc/vJvl+f65V22dbkXbowE6jgT/gNBxE= github.com/google/go-cmp v0.5.5/go.mod h1:v8dTdLbMG2kIc/vJvl+f65V22dbkXbowE6jgT/gNBxE= +github.com/google/go-cmp v0.5.6/go.mod h1:v8dTdLbMG2kIc/vJvl+f65V22dbkXbowE6jgT/gNBxE= github.com/google/go-cmp v0.5.7/go.mod h1:n+brtR0CgQNWTVd5ZUFpTBC8YFBDLK/h/bpaJ8/DtOE= github.com/google/go-cmp v0.6.0 h1:ofyhxvXcZhMsU5ulbFiLKl/XBFqE1GSq7atu8tAmTRI= github.com/google/gofuzz v1.0.0/go.mod h1:dBl0BpW6vV/+mYPU4Po3pmUjxk6FQPldtuIdl/M65Eg= github.com/google/pprof v0.0.0-20220608213341-c488b8fa1db3 h1:mpL/HvfIgIejhVwAfxBQkwEjlhP5o0O9RAeTAjpwzxc= github.com/google/pprof v0.0.0-20220608213341-c488b8fa1db3/go.mod h1:gSuNB+gJaOiQKLEZ+q+PK9Mq3SOzhRcw2GsGS/FhYDk= github.com/google/renameio v0.1.0/go.mod h1:KWCgfxg9yswjAJkECMjeO8J8rahYeXnNhOm40UhjYkI= -github.com/google/uuid v1.1.2 h1:EVhdT+1Kseyi1/pUmXKaFxYsDNy9RQYkMWRH68J/W7Y= github.com/google/uuid v1.1.2/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo= +github.com/google/uuid v1.3.0 h1:t6JiXgmwXMjEs8VusXIJk2BXHsn+wx8BZdTaoZ5fu7I= +github.com/google/uuid v1.3.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo= github.com/gopherjs/gopherjs v0.0.0-20181017120253-0766667cb4d1 h1:EGx4pi6eqNxGaHF6qqu48+N2wcFQ5qg5FXgOdqsJ5d8= github.com/gopherjs/gopherjs v0.0.0-20181017120253-0766667cb4d1/go.mod h1:wJfORRmW1u3UXTncJ5qlYoELFm8eSnnEO6hX4iZ3EWY= github.com/gordonklaus/ineffassign v0.0.0-20200309095847-7953dde2c7bf/go.mod h1:cuNKsD1zp2v6XfE/orVX2QE1LC+i254ceGcVeDT3pTU= +github.com/hpcloud/tail v1.0.0/go.mod h1:ab1qPbhIpdTxEkNHXyeSf5vhxWSCs/tWer42PpOxQnU= github.com/iancoleman/strcase v0.2.0 h1:05I4QRnGpI0m37iZQRuskXh+w77mr6Z41lwQzuHLwW0= github.com/iancoleman/strcase v0.2.0/go.mod h1:iwCmte+B7n89clKwxIoIXy/HfoL7AsD47ZCWhYzw7ho= github.com/ianlancetaylor/demangle v0.0.0-20220319035150-800ac71e25c2/go.mod h1:aYm2/VgdVmcIU8iMfdMvDMsRAQjcfZSKFby6HOFvi/w= github.com/jhump/protoreflect v1.8.2 h1:k2xE7wcUomeqwY0LDCYA16y4WWfyTcMx5mKhk0d4ua0= github.com/jhump/protoreflect v1.8.2/go.mod h1:7GcYQDdMU/O/BBrl/cX6PNHpXh6cenjd8pneu5yW7Tg= +github.com/jinzhu/copier v0.4.0 h1:w3ciUoD19shMCRargcpm0cm91ytaBhDvuRpz1ODO/U8= +github.com/jinzhu/copier v0.4.0/go.mod h1:DfbEm0FYsaqBcKcFuvmOZb218JkPGtvSHsKg8S8hyyg= github.com/json-iterator/go v1.1.10/go.mod h1:KdQUCv79m/52Kvf8AW2vK1V8akMuk1QjK/uOdHXbAo4= github.com/json-iterator/go v1.1.11/go.mod h1:KdQUCv79m/52Kvf8AW2vK1V8akMuk1QjK/uOdHXbAo4= github.com/json-iterator/go v1.1.12 h1:PV8peI4a0ysnczrg+LtxykD8LfKY9ML6u2jnxaEnrnM= @@ -208,8 +222,18 @@ github.com/modern-go/reflect2 v1.0.2 h1:xBagoLtFs94CBntxluKeaWgTMpvLxC4ur3nMaC9G github.com/modern-go/reflect2 v1.0.2/go.mod h1:yWuevngMOJpCy52FWWMvUC8ws7m/LJsjYzDa0/r8luk= github.com/niemeyer/pretty v0.0.0-20200227124842-a10e7caefd8e/go.mod h1:zD1mROLANZcx1PVRCS0qkT7pwLkGfwJo4zjcN/Tysno= github.com/nishanths/predeclared v0.0.0-20200524104333-86fad755b4d3/go.mod h1:nt3d53pc1VYcphSCIaYAJtnPYnr3Zyn8fMq2wvPGPso= +github.com/nxadm/tail v1.4.4 h1:DQuhQpB1tVlglWS2hLQ5OV6B5r8aGxSrPc5Qo6uTN78= +github.com/nxadm/tail v1.4.4/go.mod h1:kenIhsEOeOJmVchQTgglprH7qJGnHDVpk1VPCcaMI8A= github.com/oleiade/lane v1.0.1 h1:hXofkn7GEOubzTwNpeL9MaNy8WxolCYb9cInAIeqShU= github.com/oleiade/lane v1.0.1/go.mod h1:IyTkraa4maLfjq/GmHR+Dxb4kCMtEGeb+qmhlrQ5Mk4= +github.com/onsi/ginkgo v1.6.0/go.mod h1:lLunBs/Ym6LB5Z9jYTR76FiuTmxDTDusOGeTQH+WWjE= +github.com/onsi/ginkgo v1.12.1/go.mod h1:zj2OWP4+oCPe1qIXoGWkgMRwljMUYCdkwsT2108oapk= +github.com/onsi/ginkgo v1.15.0 h1:1V1NfVQR87RtWAgp1lv9JZJ5Jap+XFGKPi00andXGi4= +github.com/onsi/ginkgo v1.15.0/go.mod h1:hF8qUzuuC8DJGygJH3726JnCZX4MYbRB8yFfISqnKUg= +github.com/onsi/gomega v1.7.1/go.mod h1:XdKZgCCFLUoM/7CFJVPcG8C1xQ1AJ0vpAezJrB7JYyY= +github.com/onsi/gomega v1.10.1/go.mod h1:iN09h71vgCQne3DLsj+A5owkum+a2tYe+TOCB1ybHNo= +github.com/onsi/gomega v1.10.5 h1:7n6FEkpFmfCoo2t+YYqXH0evK+a9ICQz0xcAy9dYcaQ= +github.com/onsi/gomega v1.10.5/go.mod h1:gza4q3jKQJijlu05nKWRCW/GavJumGt8aNRxWg7mt48= github.com/opentracing/opentracing-go v1.2.0/go.mod h1:GxEUsuufX4nBwe+T+Wl9TAgYrxe9dPLANfrWvHYVTgc= github.com/phpdave11/gofpdf v1.4.2/go.mod h1:zpO6xFn9yxo3YLyMvW8HcKWVdbNqgIfOOp2dXMnm1mY= github.com/phpdave11/gofpdi v1.0.12/go.mod h1:vBmVV0Do6hSBHC8uKUQ71JGW+ZGQq74llk/7bXwjDoI= @@ -217,6 +241,8 @@ github.com/phpdave11/gofpdi v1.0.13/go.mod h1:vBmVV0Do6hSBHC8uKUQ71JGW+ZGQq74llk github.com/pkg/errors v0.8.1/go.mod h1:bwawxfHBFNV+L2hUp1rHADufV3IMtnDRdf1r5NINEl0= github.com/pkg/errors v0.9.1 h1:FEBLx1zS214owpjy7qsBeixbURkuhQAwrK5UwLGTwt4= github.com/pkg/errors v0.9.1/go.mod h1:bwawxfHBFNV+L2hUp1rHADufV3IMtnDRdf1r5NINEl0= +github.com/pkoukk/tiktoken-go v0.1.7 h1:qOBHXX4PHtvIvmOtyg1EeKlwFRiMKAcoMp4Q+bLQDmw= +github.com/pkoukk/tiktoken-go v0.1.7/go.mod h1:9NiV+i9mJKGj1rYOT+njbv+ZwA/zJxYdewGl6qVatpg= github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM= github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4= github.com/prometheus/client_golang v1.16.0 h1:yk/hx9hDbrGHovbci4BY+pRMfSuuat626eFsHb7tmT8= @@ -240,6 +266,8 @@ github.com/smartystreets/assertions v0.0.0-20180927180507-b2de0cb4f26d/go.mod h1 github.com/smartystreets/goconvey v1.6.4 h1:fv0U8FUIMPNf1L9lnHLvLhgicrIVChEkdzIKYqbNC9s= github.com/smartystreets/goconvey v1.6.4/go.mod h1:syvi0/a8iFYH4r/RixwvyeAJjdLS9QV7WQ/tjFTllLA= github.com/spaolacci/murmur3 v0.0.0-20180118202830-f09979ecbc72/go.mod h1:JwIasOWyU6f++ZhiEuf87xNszmSA2myDM2Kzu9HwQUA= +github.com/spaolacci/murmur3 v1.1.0 h1:7c1g84S4BPRrfL5Xrdp6fOJ206sU9y293DDHaoy0bLI= +github.com/spaolacci/murmur3 v1.1.0/go.mod h1:JwIasOWyU6f++ZhiEuf87xNszmSA2myDM2Kzu9HwQUA= github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME= github.com/stretchr/objx v0.4.0/go.mod h1:YvHI0jy2hoMjB+UWwv71VJQ9isScKT/TqJzVSSt89Yw= github.com/stretchr/objx v0.5.0 h1:1zr/of2m5FGMsad5YfcqgdqdWrIhu+EBEJRhR1U7z/c= @@ -326,13 +354,16 @@ golang.org/x/mod v0.5.1/go.mod h1:5OXOZSfqPIIbmVBIIKWRFfZjPR0E5r58TLhUjH0a2Ro= golang.org/x/mod v0.6.0-dev.0.20220419223038-86c51ed26bb4/go.mod h1:jJ57K6gSWd91VN4djpZkiMVwK6gcyfeH4XE8wZrZaV4= golang.org/x/net v0.0.0-20180724234803-3673e40ba225/go.mod h1:mL1N/T3taQHkDXs73rZJwtUhF3w3ftmwwsq0BUmARs4= golang.org/x/net v0.0.0-20180826012351-8a410e7b638d/go.mod h1:mL1N/T3taQHkDXs73rZJwtUhF3w3ftmwwsq0BUmARs4= +golang.org/x/net v0.0.0-20180906233101-161cd47e91fd/go.mod h1:mL1N/T3taQHkDXs73rZJwtUhF3w3ftmwwsq0BUmARs4= golang.org/x/net v0.0.0-20190213061140-3a22650c66bd/go.mod h1:mL1N/T3taQHkDXs73rZJwtUhF3w3ftmwwsq0BUmARs4= golang.org/x/net v0.0.0-20190311183353-d8887717615a/go.mod h1:t9HGtf8HONx5eT2rtn7q6eTqICYqUVnKs3thJo3Qplg= golang.org/x/net v0.0.0-20190404232315-eb5bcb51f2a3/go.mod h1:t9HGtf8HONx5eT2rtn7q6eTqICYqUVnKs3thJo3Qplg= golang.org/x/net v0.0.0-20190620200207-3b0461eec859/go.mod h1:z5CRVTTTmAJ677TzLLGU+0bjPO0LkuOLi4/5GtJWs/s= golang.org/x/net v0.0.0-20200226121028-0de0cce0169b/go.mod h1:z5CRVTTTmAJ677TzLLGU+0bjPO0LkuOLi4/5GtJWs/s= +golang.org/x/net v0.0.0-20200520004742-59133d7f0dd7/go.mod h1:qpuaurCH72eLCgpAm/N6yyVIVM9cpaDIP3A8BGJEC5A= golang.org/x/net v0.0.0-20200625001655-4c5254603344/go.mod h1:/O7V0waA8r7cgGh81Ro3o1hOxt32SMVPicZroKQ2sZA= golang.org/x/net v0.0.0-20201021035429-f5854403a974/go.mod h1:sp8m0HH+o8qH0wwXwYZr8TS3Oi6o0r6Gce1SSxlDquU= +golang.org/x/net v0.0.0-20201202161906-c7110b5ffcbb/go.mod h1:sp8m0HH+o8qH0wwXwYZr8TS3Oi6o0r6Gce1SSxlDquU= golang.org/x/net v0.0.0-20210226172049-e18ecbb05110/go.mod h1:m0MpNAwzfU5UDzcl9v0D8zg8gWTRqZa9RBIspLL5mdg= golang.org/x/net v0.0.0-20210316092652-d523dce5a7f4/go.mod h1:RBQZq4jEuRlivfhVLdyRGr576XBO4/greRjx4P4O3yc= golang.org/x/net v0.0.0-20210405180319-a5a99cb37ef4/go.mod h1:p54w0d4576C0XHj96bSt6lcn1PtDYWL6XObtHCRCNQM= @@ -355,13 +386,18 @@ golang.org/x/sync v0.0.0-20220722155255-886fb9371eb4/go.mod h1:RxMgew5VJxzue5/jJ golang.org/x/sync v0.2.0 h1:PUR+T4wwASmuSTYdKjYHI5TD22Wy5ogLU5qZCOLxBrI= golang.org/x/sync v0.2.0/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM= golang.org/x/sys v0.0.0-20180830151530-49385e6e1522/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY= +golang.org/x/sys v0.0.0-20180909124046-d0be0721c37e/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY= golang.org/x/sys v0.0.0-20190215142949-d0b11bdaac8a/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY= golang.org/x/sys v0.0.0-20190312061237-fead79001313/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs= golang.org/x/sys v0.0.0-20190412213103-97732733099d/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs= +golang.org/x/sys v0.0.0-20190904154756-749cb33beabd/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs= +golang.org/x/sys v0.0.0-20191005200804-aed5e4c7ecf9/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs= +golang.org/x/sys v0.0.0-20191120155948-bd437916bb0e/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs= golang.org/x/sys v0.0.0-20200116001909-b77594299b42/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs= golang.org/x/sys v0.0.0-20200323222414-85ca7c5b95cd/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs= golang.org/x/sys v0.0.0-20200930185726-fdedc70b468f/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs= golang.org/x/sys v0.0.0-20201119102817-f84b799fce68/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs= +golang.org/x/sys v0.0.0-20210112080510-489259a85091/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs= golang.org/x/sys v0.0.0-20210119212857-b64e53b001e4/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs= golang.org/x/sys v0.0.0-20210304124612-50617c2ba197/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs= golang.org/x/sys v0.0.0-20210315160823-c6e025ad8005/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs= @@ -410,6 +446,7 @@ golang.org/x/tools v0.0.0-20200130002326-2f3ba24bd6e7/go.mod h1:TB2adYChydJhpapK golang.org/x/tools v0.0.0-20200522201501-cb1345f3a375/go.mod h1:EkVYQZoAsY45+roYkvgYkIh4xh/qjgUK9TdY2XT94GE= golang.org/x/tools v0.0.0-20200619180055-7c47624df98f/go.mod h1:EkVYQZoAsY45+roYkvgYkIh4xh/qjgUK9TdY2XT94GE= golang.org/x/tools v0.0.0-20200717024301-6ddee64345a6/go.mod h1:njjCfa9FT2d7l9Bc6FUM5FLjQPp3cFF28FI3qnDFljA= +golang.org/x/tools v0.0.0-20201224043029-2b0845dc783e/go.mod h1:emZCQorbCU4vsT4fOWvOPXz4eW1wZW4PmDk9uLelYpA= golang.org/x/tools v0.0.0-20210106214847-113979e3529a/go.mod h1:emZCQorbCU4vsT4fOWvOPXz4eW1wZW4PmDk9uLelYpA= golang.org/x/tools v0.1.0/go.mod h1:xkSsbof2nBLbhDlRMhhhyNLN/zl3eTqcnHD5viDpcZ0= golang.org/x/tools v0.1.1/go.mod h1:o0xws9oXOQQZyjljx8fwUC0k7L1pTE6eaCbjGeHmOkk= @@ -461,8 +498,12 @@ gopkg.in/check.v1 v1.0.0-20180628173108-788fd7840127/go.mod h1:Co6ibVJAznAaIkqp8 gopkg.in/check.v1 v1.0.0-20200227125254-8fa46927fb4f/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0= gopkg.in/check.v1 v1.0.0-20201130134442-10cb98267c6c h1:Hei/4ADfdWqJk1ZMxUNpqntNwaWcugrBjAiHlqqRiVk= gopkg.in/errgo.v2 v2.1.0/go.mod h1:hNsd1EY+bozCKY1Ytp96fpM3vjJbqLJn88ws8XvfDNI= +gopkg.in/fsnotify.v1 v1.4.7/go.mod h1:Tz8NjZHkW78fSQdbUxIjBTcgA1z1m8ZHf0WmKUhAMys= gopkg.in/natefinch/lumberjack.v2 v2.0.0/go.mod h1:l0ndWWf7gzL7RNwBG7wST/UCcT4T24xpD6X8LsfU/+k= +gopkg.in/tomb.v1 v1.0.0-20141024135613-dd632973f1e7 h1:uRGJdciOHaEIrze2W8Q3AKkepLTh2hOroT7a+7czfdQ= +gopkg.in/tomb.v1 v1.0.0-20141024135613-dd632973f1e7/go.mod h1:dt/ZhP58zS4L8KSrWDmTeBkI65Dw0HsyUHuEVlX15mw= gopkg.in/yaml.v2 v2.2.2/go.mod h1:hI93XBmqTisBFMUTm0b8Fm+jr3Dg1NNxqwp+5A1VGuI= +gopkg.in/yaml.v2 v2.2.4/go.mod h1:hI93XBmqTisBFMUTm0b8Fm+jr3Dg1NNxqwp+5A1VGuI= gopkg.in/yaml.v2 v2.3.0/go.mod h1:hI93XBmqTisBFMUTm0b8Fm+jr3Dg1NNxqwp+5A1VGuI= gopkg.in/yaml.v2 v2.4.0 h1:D8xgwECY7CYvx+Y2n4sBz93Jn9JRvxdiyyo8CTfuKaY= gopkg.in/yaml.v2 v2.4.0/go.mod h1:RDklbk79AGWmwhnvt/jBztapEOGDOx6ZbXqjP6csGnQ= diff --git a/pkg/adapters/kratos/go.mod b/pkg/adapters/kratos/go.mod index 143d7e64f..896ace005 100644 --- a/pkg/adapters/kratos/go.mod +++ b/pkg/adapters/kratos/go.mod @@ -12,14 +12,19 @@ require ( require ( github.com/beorn7/perks v1.0.1 // indirect github.com/cespare/xxhash/v2 v2.2.0 // indirect + github.com/dgryski/go-rendezvous v0.0.0-20200823014737-9f7001d12a5f // indirect + github.com/dlclark/regexp2 v1.10.0 // indirect github.com/go-ole/go-ole v1.2.6 // indirect github.com/go-playground/form/v4 v4.2.0 // indirect + github.com/go-redis/redis/v8 v8.11.0 // indirect github.com/golang/protobuf v1.5.4 // indirect github.com/google/uuid v1.4.0 // indirect + github.com/jinzhu/copier v0.4.0 // indirect github.com/kr/text v0.2.0 // indirect github.com/lufia/plan9stats v0.0.0-20230326075908-cb1d2100619a // indirect github.com/matttproud/golang_protobuf_extensions v1.0.4 // indirect github.com/pkg/errors v0.9.1 // indirect + github.com/pkoukk/tiktoken-go v0.1.7 // indirect github.com/power-devops/perfstat v0.0.0-20221212215047-62379fc7944b // indirect github.com/prometheus/client_golang v1.16.0 // indirect github.com/prometheus/client_model v0.3.0 // indirect @@ -27,6 +32,7 @@ require ( github.com/prometheus/procfs v0.10.1 // indirect github.com/shirou/gopsutil/v3 v3.23.6 // indirect github.com/shoenig/go-m1cpu v0.1.6 // indirect + github.com/spaolacci/murmur3 v1.1.0 // indirect github.com/tklauser/go-sysconf v0.3.11 // indirect github.com/tklauser/numcpus v0.6.1 // indirect github.com/yusufpapurcu/wmi v1.2.3 // indirect diff --git a/pkg/adapters/kratos/go.sum b/pkg/adapters/kratos/go.sum index b37fc6802..eab65cdce 100644 --- a/pkg/adapters/kratos/go.sum +++ b/pkg/adapters/kratos/go.sum @@ -1,11 +1,19 @@ github.com/beorn7/perks v1.0.1 h1:VlbKKnNfV8bJzeqoa4cOKqO6bYr3WgKZxO8Z16+hsOM= github.com/beorn7/perks v1.0.1/go.mod h1:G2ZrVWU2WbWT9wwq4/hrbKbnv/1ERSJQ0ibhJ6rlkpw= +github.com/cespare/xxhash/v2 v2.1.1/go.mod h1:VGX0DQ3Q6kWi7AoAeZDth3/j3BFtOZR5XLFGgcrjCOs= github.com/cespare/xxhash/v2 v2.2.0 h1:DC2CZ1Ep5Y4k3ZQ899DldepgrayRUGE6BBZ/cd9Cj44= github.com/cespare/xxhash/v2 v2.2.0/go.mod h1:VGX0DQ3Q6kWi7AoAeZDth3/j3BFtOZR5XLFGgcrjCOs= github.com/creack/pty v1.1.9/go.mod h1:oKZEueFk5CKHvIhNR5MUki03XCEU+Q6VDXinZuGJ33E= github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c= github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= +github.com/dgryski/go-rendezvous v0.0.0-20200823014737-9f7001d12a5f h1:lO4WD4F/rVNCu3HqELle0jiPLLBs70cWOduZpkS1E78= +github.com/dgryski/go-rendezvous v0.0.0-20200823014737-9f7001d12a5f/go.mod h1:cuUVRXasLTGF7a8hSLbxyZXjz+1KgoB3wDUb6vlszIc= +github.com/dlclark/regexp2 v1.10.0 h1:+/GIL799phkJqYW+3YbOd8LCcbHzT0Pbo8zl70MHsq0= +github.com/dlclark/regexp2 v1.10.0/go.mod h1:DHkYz0B9wPfa6wondMfaivmHpzrQ3v9q8cnmRbL6yW8= +github.com/fsnotify/fsnotify v1.4.7/go.mod h1:jwhsz4b93w/PPRr/qN1Yymfu8t87LnFCMoQvtojpjFo= +github.com/fsnotify/fsnotify v1.4.9/go.mod h1:znqG4EE+3YCdAaPaxE2ZRY/06pZUdp0tY4IgpuI1SZQ= +github.com/fsnotify/fsnotify v1.6.0 h1:n+5WquG0fcWoWp6xPWfHdbskMCQaFnG6PfBrh1Ky4HY= github.com/go-kratos/kratos/v2 v2.8.0 h1:qr27WRTRrI3o4jzJzNKf4XVVoMYIqnQD+4ws1C46yhM= github.com/go-kratos/kratos/v2 v2.8.0/go.mod h1:+Vfe3FzF0d+BfMdajA11jT0rAyJWublRE/seZQNZVxE= github.com/go-ole/go-ole v1.2.6 h1:/Fpf6oFPoeFik9ty7siob0G6Ke8QvQEuVcuChpwXzpY= @@ -14,15 +22,29 @@ github.com/go-playground/assert/v2 v2.0.1 h1:MsBgLAaY856+nPRTKrp3/OZK38U/wa0CcBY github.com/go-playground/assert/v2 v2.0.1/go.mod h1:VDjEfimB/XKnb+ZQfWdccd7VUvScMdVu0Titje2rxJ4= github.com/go-playground/form/v4 v4.2.0 h1:N1wh+Goz61e6w66vo8vJkQt+uwZSoLz50kZPJWR8eic= github.com/go-playground/form/v4 v4.2.0/go.mod h1:q1a2BY+AQUUzhl6xA/6hBetay6dEIhMHjgvJiGo6K7U= +github.com/go-redis/redis/v8 v8.11.0 h1:O1Td0mQ8UFChQ3N9zFQqo6kTU2cJ+/it88gDB+zg0wo= +github.com/go-redis/redis/v8 v8.11.0/go.mod h1:DLomh7y2e3ggQXQLd1YgmvIfecPJoFl7WU5SOQ/r06M= github.com/golang/protobuf v1.2.0/go.mod h1:6lQm79b+lXiMfvg/cZm0SGofjICqVBUtrP5yJMmIC1U= github.com/golang/protobuf v1.3.5/go.mod h1:6O5/vntMXwX2lRkT1hjjk0nAC1IDOTvTlVgjlRvqsdk= +github.com/golang/protobuf v1.4.0-rc.1/go.mod h1:ceaxUfeHdC40wWswd/P6IGgMaK3YpKi5j83Wpe3EHw8= +github.com/golang/protobuf v1.4.0-rc.1.0.20200221234624-67d41d38c208/go.mod h1:xKAWHe0F5eneWXFV3EuXVDTCmh+JuBKY0li0aMyXATA= +github.com/golang/protobuf v1.4.0-rc.2/go.mod h1:LlEzMj4AhA7rCAGe4KMBDvJI+AwstrUpVNzEA03Pprs= +github.com/golang/protobuf v1.4.0-rc.4.0.20200313231945-b860323f09d0/go.mod h1:WU3c8KckQ9AFe+yFwt9sWVRKCVIyN9cPHBJSNnbL67w= +github.com/golang/protobuf v1.4.0/go.mod h1:jodUvKwWbYaEsadDk5Fwe5c77LiNKVO9IDvqG2KuDX0= +github.com/golang/protobuf v1.4.2/go.mod h1:oDoupMAO8OvCJWAcko0GGGIgR6R6ocIYbsSw735rRwI= github.com/golang/protobuf v1.5.4 h1:i7eJL8qZTpSEXOPTxNKhASYpMn+8e5Q6AdndVa1dWek= github.com/golang/protobuf v1.5.4/go.mod h1:lnTiLA8Wa4RWRcIUkrtSVa5nRhsEGBg48fD6rSs7xps= +github.com/google/go-cmp v0.3.0/go.mod h1:8QqcDgzrUqlUb/G2PQTWiueGozuR1884gddMywk6iLU= +github.com/google/go-cmp v0.3.1/go.mod h1:8QqcDgzrUqlUb/G2PQTWiueGozuR1884gddMywk6iLU= +github.com/google/go-cmp v0.4.0/go.mod h1:v8dTdLbMG2kIc/vJvl+f65V22dbkXbowE6jgT/gNBxE= github.com/google/go-cmp v0.5.6/go.mod h1:v8dTdLbMG2kIc/vJvl+f65V22dbkXbowE6jgT/gNBxE= github.com/google/go-cmp v0.5.9/go.mod h1:17dUlkBOakJ0+DkrSSNjCkIjxS6bF9zb3elmeNGIjoY= github.com/google/go-cmp v0.6.0 h1:ofyhxvXcZhMsU5ulbFiLKl/XBFqE1GSq7atu8tAmTRI= github.com/google/uuid v1.4.0 h1:MtMxsa51/r9yyhkyLsVeVt0B+BGQZzpQiTQ4eHZ8bc4= github.com/google/uuid v1.4.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo= +github.com/hpcloud/tail v1.0.0/go.mod h1:ab1qPbhIpdTxEkNHXyeSf5vhxWSCs/tWer42PpOxQnU= +github.com/jinzhu/copier v0.4.0 h1:w3ciUoD19shMCRargcpm0cm91ytaBhDvuRpz1ODO/U8= +github.com/jinzhu/copier v0.4.0/go.mod h1:DfbEm0FYsaqBcKcFuvmOZb218JkPGtvSHsKg8S8hyyg= github.com/kr/pretty v0.3.1 h1:flRD4NNwYAUpkphVc1HcthR4KEIFJ65n8Mw5qdRn3LE= github.com/kr/text v0.2.0 h1:5Nx0Ya0ZqY2ygV366QzturHI13Jq95ApcVaJBhpS+AY= github.com/kr/text v0.2.0/go.mod h1:eLer722TekiGuMkidMxC/pM04lWEeraHUUmBw8l2grE= @@ -31,8 +53,20 @@ github.com/lufia/plan9stats v0.0.0-20230326075908-cb1d2100619a h1:N9zuLhTvBSRt0g github.com/lufia/plan9stats v0.0.0-20230326075908-cb1d2100619a/go.mod h1:JKx41uQRwqlTZabZc+kILPrO/3jlKnQ2Z8b7YiVw5cE= github.com/matttproud/golang_protobuf_extensions v1.0.4 h1:mmDVorXM7PCGKw94cs5zkfA9PSy5pEvNWRP0ET0TIVo= github.com/matttproud/golang_protobuf_extensions v1.0.4/go.mod h1:BSXmuO+STAnVfrANrmjBb36TMTDstsz7MSK+HVaYKv4= +github.com/nxadm/tail v1.4.4 h1:DQuhQpB1tVlglWS2hLQ5OV6B5r8aGxSrPc5Qo6uTN78= +github.com/nxadm/tail v1.4.4/go.mod h1:kenIhsEOeOJmVchQTgglprH7qJGnHDVpk1VPCcaMI8A= +github.com/onsi/ginkgo v1.6.0/go.mod h1:lLunBs/Ym6LB5Z9jYTR76FiuTmxDTDusOGeTQH+WWjE= +github.com/onsi/ginkgo v1.12.1/go.mod h1:zj2OWP4+oCPe1qIXoGWkgMRwljMUYCdkwsT2108oapk= +github.com/onsi/ginkgo v1.15.0 h1:1V1NfVQR87RtWAgp1lv9JZJ5Jap+XFGKPi00andXGi4= +github.com/onsi/ginkgo v1.15.0/go.mod h1:hF8qUzuuC8DJGygJH3726JnCZX4MYbRB8yFfISqnKUg= +github.com/onsi/gomega v1.7.1/go.mod h1:XdKZgCCFLUoM/7CFJVPcG8C1xQ1AJ0vpAezJrB7JYyY= +github.com/onsi/gomega v1.10.1/go.mod h1:iN09h71vgCQne3DLsj+A5owkum+a2tYe+TOCB1ybHNo= +github.com/onsi/gomega v1.10.5 h1:7n6FEkpFmfCoo2t+YYqXH0evK+a9ICQz0xcAy9dYcaQ= +github.com/onsi/gomega v1.10.5/go.mod h1:gza4q3jKQJijlu05nKWRCW/GavJumGt8aNRxWg7mt48= github.com/pkg/errors v0.9.1 h1:FEBLx1zS214owpjy7qsBeixbURkuhQAwrK5UwLGTwt4= github.com/pkg/errors v0.9.1/go.mod h1:bwawxfHBFNV+L2hUp1rHADufV3IMtnDRdf1r5NINEl0= +github.com/pkoukk/tiktoken-go v0.1.7 h1:qOBHXX4PHtvIvmOtyg1EeKlwFRiMKAcoMp4Q+bLQDmw= +github.com/pkoukk/tiktoken-go v0.1.7/go.mod h1:9NiV+i9mJKGj1rYOT+njbv+ZwA/zJxYdewGl6qVatpg= github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM= github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4= github.com/power-devops/perfstat v0.0.0-20210106213030-5aafc221ea8c/go.mod h1:OmDBASR4679mdNQnz2pUhc2G8CO2JrUAVFDRBDP/hJE= @@ -53,6 +87,8 @@ github.com/shoenig/go-m1cpu v0.1.6 h1:nxdKQNcEB6vzgA2E2bvzKIYRuNj7XNJ4S/aRSwKzFt github.com/shoenig/go-m1cpu v0.1.6/go.mod h1:1JJMcUBvfNwpq05QDQVAnx3gUHr9IYF7GNg9SUEw2VQ= github.com/shoenig/test v0.6.4 h1:kVTaSd7WLz5WZ2IaoM0RSzRsUD+m8wRR+5qvntpn4LU= github.com/shoenig/test v0.6.4/go.mod h1:byHiCGXqrVaflBLAMq/srcZIHynQPQgeyvkvXnjqq0k= +github.com/spaolacci/murmur3 v1.1.0 h1:7c1g84S4BPRrfL5Xrdp6fOJ206sU9y293DDHaoy0bLI= +github.com/spaolacci/murmur3 v1.1.0/go.mod h1:JwIasOWyU6f++ZhiEuf87xNszmSA2myDM2Kzu9HwQUA= github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME= github.com/stretchr/objx v0.4.0/go.mod h1:YvHI0jy2hoMjB+UWwv71VJQ9isScKT/TqJzVSSt89Yw= github.com/stretchr/objx v0.5.0 h1:1zr/of2m5FGMsad5YfcqgdqdWrIhu+EBEJRhR1U7z/c= @@ -66,27 +102,69 @@ github.com/tklauser/go-sysconf v0.3.11/go.mod h1:GqXfhXY3kiPa0nAXPDIQIWzJbMCB7Am github.com/tklauser/numcpus v0.6.0/go.mod h1:FEZLMke0lhOUG6w2JadTzp0a+Nl8PF/GFkQ5UVIcaL4= github.com/tklauser/numcpus v0.6.1 h1:ng9scYS7az0Bk4OZLvrNXNSAO2Pxr1XXRAPyjhIx+Fk= github.com/tklauser/numcpus v0.6.1/go.mod h1:1XfjsgE2zo8GVw7POkMbHENHzVg3GzmoZ9fESEdAacY= +github.com/yuin/goldmark v1.2.1/go.mod h1:3hX8gzYuyVAZsxl0MRgGTJEmQBFcNTphYh9decYSb74= github.com/yusufpapurcu/wmi v1.2.3 h1:E1ctvB7uKFMOJw3fdOW32DwGE9I7t++CRUEMKvFoFiw= github.com/yusufpapurcu/wmi v1.2.3/go.mod h1:SBZ9tNy3G9/m5Oi98Zks0QjeHVDvuK0qfxQmPyzfmi0= +golang.org/x/crypto v0.0.0-20190308221718-c2843e01d9a2/go.mod h1:djNgcEr1/C05ACkg1iLfiJU5Ep61QUkGW8qpdssI0+w= +golang.org/x/crypto v0.0.0-20191011191535-87dc89f01550/go.mod h1:yigFU9vqHzYiE8UmvKecakEJjdnWj3jj499lnFckfCI= +golang.org/x/crypto v0.0.0-20200622213623-75b288015ac9/go.mod h1:LzIPMQfyMNhhGPhUkYOs5KpL4U8rLKemX1yGLhDgUto= +golang.org/x/mod v0.3.0/go.mod h1:s0Qsj1ACt9ePp/hMypM3fl4fZqREWJwdYDEqhRiZZUA= +golang.org/x/net v0.0.0-20180906233101-161cd47e91fd/go.mod h1:mL1N/T3taQHkDXs73rZJwtUhF3w3ftmwwsq0BUmARs4= +golang.org/x/net v0.0.0-20190404232315-eb5bcb51f2a3/go.mod h1:t9HGtf8HONx5eT2rtn7q6eTqICYqUVnKs3thJo3Qplg= +golang.org/x/net v0.0.0-20190620200207-3b0461eec859/go.mod h1:z5CRVTTTmAJ677TzLLGU+0bjPO0LkuOLi4/5GtJWs/s= +golang.org/x/net v0.0.0-20200520004742-59133d7f0dd7/go.mod h1:qpuaurCH72eLCgpAm/N6yyVIVM9cpaDIP3A8BGJEC5A= +golang.org/x/net v0.0.0-20201021035429-f5854403a974/go.mod h1:sp8m0HH+o8qH0wwXwYZr8TS3Oi6o0r6Gce1SSxlDquU= +golang.org/x/net v0.0.0-20201202161906-c7110b5ffcbb/go.mod h1:sp8m0HH+o8qH0wwXwYZr8TS3Oi6o0r6Gce1SSxlDquU= golang.org/x/net v0.23.0 h1:7EYJ93RZ9vYSZAIb2x3lnuvqO5zneoD6IvWjuhfxjTs= +golang.org/x/sync v0.0.0-20180314180146-1d60e4601c6f/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM= golang.org/x/sync v0.0.0-20181221193216-37e7f081c4d4/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM= +golang.org/x/sync v0.0.0-20190423024810-112230192c58/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM= +golang.org/x/sync v0.0.0-20201020160332-67f06af15bc9/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM= +golang.org/x/sys v0.0.0-20180909124046-d0be0721c37e/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY= +golang.org/x/sys v0.0.0-20190215142949-d0b11bdaac8a/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY= +golang.org/x/sys v0.0.0-20190412213103-97732733099d/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs= +golang.org/x/sys v0.0.0-20190904154756-749cb33beabd/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs= golang.org/x/sys v0.0.0-20190916202348-b4ddaad3f8a3/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs= +golang.org/x/sys v0.0.0-20191005200804-aed5e4c7ecf9/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs= +golang.org/x/sys v0.0.0-20191120155948-bd437916bb0e/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs= +golang.org/x/sys v0.0.0-20200323222414-85ca7c5b95cd/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs= +golang.org/x/sys v0.0.0-20200930185726-fdedc70b468f/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs= golang.org/x/sys v0.0.0-20201204225414-ed752295db88/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs= +golang.org/x/sys v0.0.0-20210112080510-489259a85091/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs= golang.org/x/sys v0.2.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= golang.org/x/sys v0.8.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= golang.org/x/sys v0.9.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= golang.org/x/sys v0.21.0 h1:rF+pYz3DAGSQAxAu1CbC7catZg4ebC4UIeIhKxBZvws= golang.org/x/sys v0.21.0/go.mod h1:/VUhepiaJMQUp4+oa/7Zr1D23ma6VTLIYjOOTFZPUcA= +golang.org/x/text v0.3.0/go.mod h1:NqM8EUOU14njkJ3fqMW+pc6Ldnwhi/IjpwHt7yyuwOQ= +golang.org/x/text v0.3.3/go.mod h1:5Zoc/QRtKVWzQhOtBMvqHzDpF6irO9z98xDceosuGiQ= golang.org/x/text v0.14.0 h1:ScX5w1eTa3QqT8oi6+ziP7dTV1S2+ALU0bI+0zXKWiQ= +golang.org/x/tools v0.0.0-20180917221912-90fa682c2a6e/go.mod h1:n7NCudcB/nEzxVGmLbDWY5pfWTLqBcC2KZ6jyYvM4mQ= +golang.org/x/tools v0.0.0-20191119224855-298f0cb1881e/go.mod h1:b+2E5dAYhXwXZwtnZ6UAqBI28+e2cm9otk0dWdXHAEo= +golang.org/x/tools v0.0.0-20201224043029-2b0845dc783e/go.mod h1:emZCQorbCU4vsT4fOWvOPXz4eW1wZW4PmDk9uLelYpA= +golang.org/x/xerrors v0.0.0-20190717185122-a985d3407aa7/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0= +golang.org/x/xerrors v0.0.0-20191011141410-1b5146add898/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0= golang.org/x/xerrors v0.0.0-20191204190536-9bdfabe68543/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0= +golang.org/x/xerrors v0.0.0-20200804184101-5ec99f83aff1/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0= google.golang.org/genproto/googleapis/rpc v0.0.0-20240102182953-50ed04b92917 h1:6G8oQ016D88m1xAKljMlBOOGWDZkes4kMhgGFlf8WcQ= google.golang.org/genproto/googleapis/rpc v0.0.0-20240102182953-50ed04b92917/go.mod h1:xtjpI3tXFPP051KaWnhvxkiubL/6dJ18vLVf7q2pTOU= google.golang.org/grpc v1.61.1 h1:kLAiWrZs7YeDM6MumDe7m3y4aM6wacLzM1Y/wiLP9XY= google.golang.org/grpc v1.61.1/go.mod h1:VUbo7IFqmF1QtCAstipjG0GIoq49KvMe9+h1jFLBNJs= +google.golang.org/protobuf v0.0.0-20200109180630-ec00e32a8dfd/go.mod h1:DFci5gLYBciE7Vtevhsrf46CRTquxDuWsQurQQe4oz8= +google.golang.org/protobuf v0.0.0-20200221191635-4d8936d0db64/go.mod h1:kwYJMbMJ01Woi6D6+Kah6886xMZcty6N08ah7+eCXa0= +google.golang.org/protobuf v0.0.0-20200228230310-ab0ca4ff8a60/go.mod h1:cfTl7dwQJ+fmap5saPgwCLgHXTUD7jkjRqWcaiX5VyM= +google.golang.org/protobuf v1.20.1-0.20200309200217-e05f789c0967/go.mod h1:A+miEFZTKqfCUM6K7xSMQL9OKL/b6hQv+e19PK+JZNE= +google.golang.org/protobuf v1.21.0/go.mod h1:47Nbq4nVaFHyn7ilMalzfO3qCViNmqZ2kzikPIcrTAo= +google.golang.org/protobuf v1.23.0/go.mod h1:EGpADcykh3NcUnDUJcl1+ZksZNG86OlYog2l/sGQquU= google.golang.org/protobuf v1.33.0 h1:uNO2rsAINq/JlFpSdYEKIZ0uKD/R9cpdv0T+yoGwGmI= google.golang.org/protobuf v1.33.0/go.mod h1:c6P6GXX6sHbq/GpV6MGZEdwhWPcYBgnhAHhKbcUYpos= gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0= gopkg.in/check.v1 v1.0.0-20201130134442-10cb98267c6c h1:Hei/4ADfdWqJk1ZMxUNpqntNwaWcugrBjAiHlqqRiVk= +gopkg.in/fsnotify.v1 v1.4.7/go.mod h1:Tz8NjZHkW78fSQdbUxIjBTcgA1z1m8ZHf0WmKUhAMys= +gopkg.in/tomb.v1 v1.0.0-20141024135613-dd632973f1e7 h1:uRGJdciOHaEIrze2W8Q3AKkepLTh2hOroT7a+7czfdQ= +gopkg.in/tomb.v1 v1.0.0-20141024135613-dd632973f1e7/go.mod h1:dt/ZhP58zS4L8KSrWDmTeBkI65Dw0HsyUHuEVlX15mw= +gopkg.in/yaml.v2 v2.2.4/go.mod h1:hI93XBmqTisBFMUTm0b8Fm+jr3Dg1NNxqwp+5A1VGuI= +gopkg.in/yaml.v2 v2.3.0/go.mod h1:hI93XBmqTisBFMUTm0b8Fm+jr3Dg1NNxqwp+5A1VGuI= gopkg.in/yaml.v2 v2.4.0 h1:D8xgwECY7CYvx+Y2n4sBz93Jn9JRvxdiyyo8CTfuKaY= gopkg.in/yaml.v2 v2.4.0/go.mod h1:RDklbk79AGWmwhnvt/jBztapEOGDOx6ZbXqjP6csGnQ= gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM= diff --git a/pkg/adapters/langchaingo/go.mod b/pkg/adapters/langchaingo/go.mod new file mode 100644 index 000000000..9700d372c --- /dev/null +++ b/pkg/adapters/langchaingo/go.mod @@ -0,0 +1,41 @@ +module github.com/alibaba/sentinel-golang/pkg/adapters/langchaingo + +go 1.22.0 + +replace github.com/alibaba/sentinel-golang => ../../../ + +require ( + github.com/alibaba/sentinel-golang v1.0.4 + github.com/tmc/langchaingo v0.1.13 +) + +require ( + github.com/beorn7/perks v1.0.1 // indirect + github.com/cespare/xxhash/v2 v2.2.0 // indirect + github.com/dgryski/go-rendezvous v0.0.0-20200823014737-9f7001d12a5f // indirect + github.com/dlclark/regexp2 v1.10.0 // indirect + github.com/go-ole/go-ole v1.2.6 // indirect + github.com/go-redis/redis/v8 v8.11.5 // indirect + github.com/golang/protobuf v1.5.4 // indirect + github.com/google/uuid v1.6.0 // indirect + github.com/jinzhu/copier v0.4.0 // indirect + github.com/lufia/plan9stats v0.0.0-20211012122336-39d0f177ccd0 // indirect + github.com/matttproud/golang_protobuf_extensions v1.0.4 // indirect + github.com/pkg/errors v0.9.1 // indirect + github.com/pkoukk/tiktoken-go v0.1.7 // indirect + github.com/power-devops/perfstat v0.0.0-20210106213030-5aafc221ea8c // indirect + github.com/prometheus/client_golang v1.16.0 // indirect + github.com/prometheus/client_model v0.3.0 // indirect + github.com/prometheus/common v0.42.0 // indirect + github.com/prometheus/procfs v0.10.1 // indirect + github.com/shirou/gopsutil/v3 v3.23.12 // indirect + github.com/shoenig/go-m1cpu v0.1.6 // indirect + github.com/spaolacci/murmur3 v1.1.0 // indirect + github.com/stretchr/objx v0.5.2 // indirect + github.com/tklauser/go-sysconf v0.3.12 // indirect + github.com/tklauser/numcpus v0.6.1 // indirect + github.com/yusufpapurcu/wmi v1.2.3 // indirect + golang.org/x/sys v0.27.0 // indirect + google.golang.org/protobuf v1.34.1 // indirect + gopkg.in/yaml.v2 v2.4.0 // indirect +) diff --git a/pkg/adapters/langchaingo/go.sum b/pkg/adapters/langchaingo/go.sum new file mode 100644 index 000000000..0d22e9c59 --- /dev/null +++ b/pkg/adapters/langchaingo/go.sum @@ -0,0 +1,114 @@ +github.com/beorn7/perks v1.0.1 h1:VlbKKnNfV8bJzeqoa4cOKqO6bYr3WgKZxO8Z16+hsOM= +github.com/beorn7/perks v1.0.1/go.mod h1:G2ZrVWU2WbWT9wwq4/hrbKbnv/1ERSJQ0ibhJ6rlkpw= +github.com/cespare/xxhash/v2 v2.2.0 h1:DC2CZ1Ep5Y4k3ZQ899DldepgrayRUGE6BBZ/cd9Cj44= +github.com/cespare/xxhash/v2 v2.2.0/go.mod h1:VGX0DQ3Q6kWi7AoAeZDth3/j3BFtOZR5XLFGgcrjCOs= +github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= +github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c= +github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= +github.com/dgryski/go-rendezvous v0.0.0-20200823014737-9f7001d12a5f h1:lO4WD4F/rVNCu3HqELle0jiPLLBs70cWOduZpkS1E78= +github.com/dgryski/go-rendezvous v0.0.0-20200823014737-9f7001d12a5f/go.mod h1:cuUVRXasLTGF7a8hSLbxyZXjz+1KgoB3wDUb6vlszIc= +github.com/dlclark/regexp2 v1.10.0 h1:+/GIL799phkJqYW+3YbOd8LCcbHzT0Pbo8zl70MHsq0= +github.com/dlclark/regexp2 v1.10.0/go.mod h1:DHkYz0B9wPfa6wondMfaivmHpzrQ3v9q8cnmRbL6yW8= +github.com/fsnotify/fsnotify v1.4.9 h1:hsms1Qyu0jgnwNXIxa+/V/PDsU6CfLf6CNO8H7IWoS4= +github.com/fsnotify/fsnotify v1.4.9/go.mod h1:znqG4EE+3YCdAaPaxE2ZRY/06pZUdp0tY4IgpuI1SZQ= +github.com/go-ole/go-ole v1.2.6 h1:/Fpf6oFPoeFik9ty7siob0G6Ke8QvQEuVcuChpwXzpY= +github.com/go-ole/go-ole v1.2.6/go.mod h1:pprOEPIfldk/42T2oK7lQ4v4JSDwmV0As9GaiUsvbm0= +github.com/go-redis/redis/v8 v8.11.5 h1:AcZZR7igkdvfVmQTPnu9WE37LRrO/YrBH5zWyjDC0oI= +github.com/go-redis/redis/v8 v8.11.5/go.mod h1:gREzHqY1hg6oD9ngVRbLStwAWKhA0FEgq8Jd4h5lpwo= +github.com/golang/protobuf v1.2.0/go.mod h1:6lQm79b+lXiMfvg/cZm0SGofjICqVBUtrP5yJMmIC1U= +github.com/golang/protobuf v1.3.5/go.mod h1:6O5/vntMXwX2lRkT1hjjk0nAC1IDOTvTlVgjlRvqsdk= +github.com/golang/protobuf v1.5.4 h1:i7eJL8qZTpSEXOPTxNKhASYpMn+8e5Q6AdndVa1dWek= +github.com/golang/protobuf v1.5.4/go.mod h1:lnTiLA8Wa4RWRcIUkrtSVa5nRhsEGBg48fD6rSs7xps= +github.com/google/go-cmp v0.5.6/go.mod h1:v8dTdLbMG2kIc/vJvl+f65V22dbkXbowE6jgT/gNBxE= +github.com/google/go-cmp v0.5.9/go.mod h1:17dUlkBOakJ0+DkrSSNjCkIjxS6bF9zb3elmeNGIjoY= +github.com/google/go-cmp v0.6.0 h1:ofyhxvXcZhMsU5ulbFiLKl/XBFqE1GSq7atu8tAmTRI= +github.com/google/go-cmp v0.6.0/go.mod h1:17dUlkBOakJ0+DkrSSNjCkIjxS6bF9zb3elmeNGIjoY= +github.com/google/uuid v1.6.0 h1:NIvaJDMOsjHA8n1jAhLSgzrAzy1Hgr+hNrb57e+94F0= +github.com/google/uuid v1.6.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo= +github.com/jinzhu/copier v0.4.0 h1:w3ciUoD19shMCRargcpm0cm91ytaBhDvuRpz1ODO/U8= +github.com/jinzhu/copier v0.4.0/go.mod h1:DfbEm0FYsaqBcKcFuvmOZb218JkPGtvSHsKg8S8hyyg= +github.com/kr/pretty v0.3.1 h1:flRD4NNwYAUpkphVc1HcthR4KEIFJ65n8Mw5qdRn3LE= +github.com/kr/pretty v0.3.1/go.mod h1:hoEshYVHaxMs3cyo3Yncou5ZscifuDolrwPKZanG3xk= +github.com/kr/text v0.2.0 h1:5Nx0Ya0ZqY2ygV366QzturHI13Jq95ApcVaJBhpS+AY= +github.com/kr/text v0.2.0/go.mod h1:eLer722TekiGuMkidMxC/pM04lWEeraHUUmBw8l2grE= +github.com/lufia/plan9stats v0.0.0-20211012122336-39d0f177ccd0 h1:6E+4a0GO5zZEnZ81pIr0yLvtUWk2if982qA3F3QD6H4= +github.com/lufia/plan9stats v0.0.0-20211012122336-39d0f177ccd0/go.mod h1:zJYVVT2jmtg6P3p1VtQj7WsuWi/y4VnjVBn7F8KPB3I= +github.com/matttproud/golang_protobuf_extensions v1.0.4 h1:mmDVorXM7PCGKw94cs5zkfA9PSy5pEvNWRP0ET0TIVo= +github.com/matttproud/golang_protobuf_extensions v1.0.4/go.mod h1:BSXmuO+STAnVfrANrmjBb36TMTDstsz7MSK+HVaYKv4= +github.com/nxadm/tail v1.4.8 h1:nPr65rt6Y5JFSKQO7qToXr7pePgD6Gwiw05lkbyAQTE= +github.com/nxadm/tail v1.4.8/go.mod h1:+ncqLTQzXmGhMZNUePPaPqPvBxHAIsmXswZKocGu+AU= +github.com/onsi/ginkgo v1.16.5 h1:8xi0RTUf59SOSfEtZMvwTvXYMzG4gV23XVHOZiXNtnE= +github.com/onsi/ginkgo v1.16.5/go.mod h1:+E8gABHa3K6zRBolWtd+ROzc/U5bkGt0FwiG042wbpU= +github.com/onsi/gomega v1.18.1 h1:M1GfJqGRrBrrGGsbxzV5dqM2U2ApXefZCQpkukxYRLE= +github.com/onsi/gomega v1.18.1/go.mod h1:0q+aL8jAiMXy9hbwj2mr5GziHiwhAIQpFmmtT5hitRs= +github.com/pkg/errors v0.9.1 h1:FEBLx1zS214owpjy7qsBeixbURkuhQAwrK5UwLGTwt4= +github.com/pkg/errors v0.9.1/go.mod h1:bwawxfHBFNV+L2hUp1rHADufV3IMtnDRdf1r5NINEl0= +github.com/pkoukk/tiktoken-go v0.1.7 h1:qOBHXX4PHtvIvmOtyg1EeKlwFRiMKAcoMp4Q+bLQDmw= +github.com/pkoukk/tiktoken-go v0.1.7/go.mod h1:9NiV+i9mJKGj1rYOT+njbv+ZwA/zJxYdewGl6qVatpg= +github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM= +github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4= +github.com/power-devops/perfstat v0.0.0-20210106213030-5aafc221ea8c h1:ncq/mPwQF4JjgDlrVEn3C11VoGHZN7m8qihwgMEtzYw= +github.com/power-devops/perfstat v0.0.0-20210106213030-5aafc221ea8c/go.mod h1:OmDBASR4679mdNQnz2pUhc2G8CO2JrUAVFDRBDP/hJE= +github.com/prometheus/client_golang v1.16.0 h1:yk/hx9hDbrGHovbci4BY+pRMfSuuat626eFsHb7tmT8= +github.com/prometheus/client_golang v1.16.0/go.mod h1:Zsulrv/L9oM40tJ7T815tM89lFEugiJ9HzIqaAx4LKc= +github.com/prometheus/client_model v0.3.0 h1:UBgGFHqYdG/TPFD1B1ogZywDqEkwp3fBMvqdiQ7Xew4= +github.com/prometheus/client_model v0.3.0/go.mod h1:LDGWKZIo7rky3hgvBe+caln+Dr3dPggB5dvjtD7w9+w= +github.com/prometheus/common v0.42.0 h1:EKsfXEYo4JpWMHH5cg+KOUWeuJSov1Id8zGR8eeI1YM= +github.com/prometheus/common v0.42.0/go.mod h1:xBwqVerjNdUDjgODMpudtOMwlOwf2SaTr1yjz4b7Zbc= +github.com/prometheus/procfs v0.10.1 h1:kYK1Va/YMlutzCGazswoHKo//tZVlFpKYh+PymziUAg= +github.com/prometheus/procfs v0.10.1/go.mod h1:nwNm2aOCAYw8uTR/9bWRREkZFxAUcWzPHWJq+XBB/FM= +github.com/rogpeppe/go-internal v1.13.1 h1:KvO1DLK/DRN07sQ1LQKScxyZJuNnedQ5/wKSR38lUII= +github.com/rogpeppe/go-internal v1.13.1/go.mod h1:uMEvuHeurkdAXX61udpOXGD/AzZDWNMNyH2VO9fmH0o= +github.com/shirou/gopsutil/v3 v3.23.12 h1:z90NtUkp3bMtmICZKpC4+WaknU1eXtp5vtbQ11DgpE4= +github.com/shirou/gopsutil/v3 v3.23.12/go.mod h1:1FrWgea594Jp7qmjHUUPlJDTPgcsb9mGnXDxavtikzM= +github.com/shoenig/go-m1cpu v0.1.6 h1:nxdKQNcEB6vzgA2E2bvzKIYRuNj7XNJ4S/aRSwKzFtM= +github.com/shoenig/go-m1cpu v0.1.6/go.mod h1:1JJMcUBvfNwpq05QDQVAnx3gUHr9IYF7GNg9SUEw2VQ= +github.com/shoenig/test v0.6.4 h1:kVTaSd7WLz5WZ2IaoM0RSzRsUD+m8wRR+5qvntpn4LU= +github.com/shoenig/test v0.6.4/go.mod h1:byHiCGXqrVaflBLAMq/srcZIHynQPQgeyvkvXnjqq0k= +github.com/spaolacci/murmur3 v1.1.0 h1:7c1g84S4BPRrfL5Xrdp6fOJ206sU9y293DDHaoy0bLI= +github.com/spaolacci/murmur3 v1.1.0/go.mod h1:JwIasOWyU6f++ZhiEuf87xNszmSA2myDM2Kzu9HwQUA= +github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME= +github.com/stretchr/objx v0.4.0/go.mod h1:YvHI0jy2hoMjB+UWwv71VJQ9isScKT/TqJzVSSt89Yw= +github.com/stretchr/objx v0.5.0/go.mod h1:Yh+to48EsGEfYuaHDzXPcE3xhTkx73EhmCGUpEOglKo= +github.com/stretchr/objx v0.5.2 h1:xuMeJ0Sdp5ZMRXx/aWO6RZxdr3beISkG5/G/aIRr3pY= +github.com/stretchr/objx v0.5.2/go.mod h1:FRsXN1f5AsAjCGJKqEizvkpNtU+EGNCLh3NxZ/8L+MA= +github.com/stretchr/testify v1.7.1/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg= +github.com/stretchr/testify v1.8.0/go.mod h1:yNjHg4UonilssWZ8iaSj1OCr/vHnekPRkoO+kdMU+MU= +github.com/stretchr/testify v1.8.4/go.mod h1:sz/lmYIOXD/1dqDmKjjqLyZ2RngseejIcXlSw2iwfAo= +github.com/stretchr/testify v1.9.0 h1:HtqpIVDClZ4nwg75+f6Lvsy/wHu+3BoSGCbBAcpTsTg= +github.com/stretchr/testify v1.9.0/go.mod h1:r2ic/lqez/lEtzL7wO/rwa5dbSLXVDPFyf8C91i36aY= +github.com/tklauser/go-sysconf v0.3.12 h1:0QaGUFOdQaIVdPgfITYzaTegZvdCjmYO52cSFAEVmqU= +github.com/tklauser/go-sysconf v0.3.12/go.mod h1:Ho14jnntGE1fpdOqQEEaiKRpvIavV0hSfmBq8nJbHYI= +github.com/tklauser/numcpus v0.6.1 h1:ng9scYS7az0Bk4OZLvrNXNSAO2Pxr1XXRAPyjhIx+Fk= +github.com/tklauser/numcpus v0.6.1/go.mod h1:1XfjsgE2zo8GVw7POkMbHENHzVg3GzmoZ9fESEdAacY= +github.com/tmc/langchaingo v0.1.13 h1:rcpMWBIi2y3B90XxfE4Ao8dhCQPVDMaNPnN5cGB1CaA= +github.com/tmc/langchaingo v0.1.13/go.mod h1:vpQ5NOIhpzxDfTZK9B6tf2GM/MoaHewPWM5KXXGh7hg= +github.com/yusufpapurcu/wmi v1.2.3 h1:E1ctvB7uKFMOJw3fdOW32DwGE9I7t++CRUEMKvFoFiw= +github.com/yusufpapurcu/wmi v1.2.3/go.mod h1:SBZ9tNy3G9/m5Oi98Zks0QjeHVDvuK0qfxQmPyzfmi0= +golang.org/x/net v0.25.0 h1:d/OCCoBEUq33pjydKrGQhw7IlUPI2Oylr+8qLx49kac= +golang.org/x/net v0.25.0/go.mod h1:JkAGAh7GEvH74S6FOH42FLoXpXbE/aqXSrIQjXgsiwM= +golang.org/x/sync v0.0.0-20181221193216-37e7f081c4d4/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM= +golang.org/x/sys v0.0.0-20190916202348-b4ddaad3f8a3/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs= +golang.org/x/sys v0.0.0-20201204225414-ed752295db88/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs= +golang.org/x/sys v0.8.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= +golang.org/x/sys v0.11.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= +golang.org/x/sys v0.15.0/go.mod h1:/VUhepiaJMQUp4+oa/7Zr1D23ma6VTLIYjOOTFZPUcA= +golang.org/x/sys v0.27.0 h1:wBqf8DvsY9Y/2P8gAfPDEYNuS30J4lPHJxXSb/nJZ+s= +golang.org/x/sys v0.27.0/go.mod h1:/VUhepiaJMQUp4+oa/7Zr1D23ma6VTLIYjOOTFZPUcA= +golang.org/x/text v0.20.0 h1:gK/Kv2otX8gz+wn7Rmb3vT96ZwuoxnQlY+HlJVj7Qug= +golang.org/x/text v0.20.0/go.mod h1:D4IsuqiFMhST5bX19pQ9ikHC2GsaKyk/oF+pn3ducp4= +golang.org/x/xerrors v0.0.0-20191204190536-9bdfabe68543/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0= +google.golang.org/protobuf v1.34.1 h1:9ddQBjfCyZPOHPUiPxpYESBLc+T8P3E+Vo4IbKZgFWg= +google.golang.org/protobuf v1.34.1/go.mod h1:c6P6GXX6sHbq/GpV6MGZEdwhWPcYBgnhAHhKbcUYpos= +gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0= +gopkg.in/check.v1 v1.0.0-20201130134442-10cb98267c6c h1:Hei/4ADfdWqJk1ZMxUNpqntNwaWcugrBjAiHlqqRiVk= +gopkg.in/check.v1 v1.0.0-20201130134442-10cb98267c6c/go.mod h1:JHkPIbrfpd72SG/EVd6muEfDQjcINNoR0C8j2r3qZ4Q= +gopkg.in/tomb.v1 v1.0.0-20141024135613-dd632973f1e7 h1:uRGJdciOHaEIrze2W8Q3AKkepLTh2hOroT7a+7czfdQ= +gopkg.in/tomb.v1 v1.0.0-20141024135613-dd632973f1e7/go.mod h1:dt/ZhP58zS4L8KSrWDmTeBkI65Dw0HsyUHuEVlX15mw= +gopkg.in/yaml.v2 v2.4.0 h1:D8xgwECY7CYvx+Y2n4sBz93Jn9JRvxdiyyo8CTfuKaY= +gopkg.in/yaml.v2 v2.4.0/go.mod h1:RDklbk79AGWmwhnvt/jBztapEOGDOx6ZbXqjP6csGnQ= +gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM= +gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA= +gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM= +sigs.k8s.io/yaml v1.3.0 h1:a2VclLzOGrwOHDiV8EfBGhvjHvP46CtW5j6POvhYGGo= +sigs.k8s.io/yaml v1.3.0/go.mod h1:GeOyir5tyXNByN85N/dRIT9es5UQNerPYEKK56eTBm8= diff --git a/pkg/adapters/langchaingo/option.go b/pkg/adapters/langchaingo/option.go new file mode 100644 index 000000000..20e0b2571 --- /dev/null +++ b/pkg/adapters/langchaingo/option.go @@ -0,0 +1,92 @@ +// Copyright 1999-2020 Alibaba Group Holding Ltd. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package langchaingo + +import ( + "context" + + llmtokenratelimit "github.com/alibaba/sentinel-golang/core/llm_token_ratelimit" + "github.com/tmc/langchaingo/llms" +) + +type Option func(*options) + +type options struct { + defaultResource string + resourceExtract func(context.Context) string + blockFallback func(context.Context) + requestInfosExtract func(context.Context) *llmtokenratelimit.RequestInfos + promptsExtract func([]llms.MessageContent) []string + usedTokenInfosExtract func(interface{}) *llmtokenratelimit.UsedTokenInfos +} + +func WithDefaultResource(resource string) Option { + return func(o *options) { + o.defaultResource = resource + } +} + +func WithResourceExtract(fn func(context.Context) string) Option { + return func(o *options) { + o.resourceExtract = fn + } +} + +func WithBlockFallback(fn func(context.Context)) Option { + return func(o *options) { + o.blockFallback = fn + } +} + +func WithRequestInfosExtract(fn func(context.Context) *llmtokenratelimit.RequestInfos) Option { + return func(o *options) { + o.requestInfosExtract = fn + } +} + +func WithPromptsExtract(fn func([]llms.MessageContent) []string) Option { + return func(o *options) { + o.promptsExtract = fn + } +} + +func WithUsedTokenInfosExtract(fn func(interface{}) *llmtokenratelimit.UsedTokenInfos) Option { + return func(o *options) { + o.usedTokenInfosExtract = fn + } +} + +func evaluateOptions(opts ...Option) *options { + optCopy := &options{ + defaultResource: llmtokenratelimit.DefaultResourcePattern, + promptsExtract: func(messages []llms.MessageContent) []string { + prompts := make([]string, 0, len(messages)) + for _, msg := range messages { + for _, part := range msg.Parts { + if textPart, ok := part.(llms.TextContent); ok { + prompts = append(prompts, textPart.Text) + } + } + } + return prompts + }, + } + for _, opt := range opts { + if opt != nil { + opt(optCopy) + } + } + return optCopy +} diff --git a/pkg/adapters/langchaingo/wrapper.go b/pkg/adapters/langchaingo/wrapper.go new file mode 100644 index 000000000..1182a1358 --- /dev/null +++ b/pkg/adapters/langchaingo/wrapper.go @@ -0,0 +1,114 @@ +// Copyright 1999-2020 Alibaba Group Holding Ltd. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +package langchaingo + +import ( + "context" + "fmt" + + sentinel "github.com/alibaba/sentinel-golang/api" + "github.com/alibaba/sentinel-golang/core/base" + llmtokenratelimit "github.com/alibaba/sentinel-golang/core/llm_token_ratelimit" + "github.com/tmc/langchaingo/llms" +) + +type LLMWrapper struct { + llm llms.Model + options *options +} + +func NewLLMWrapper(llm llms.Model, opts ...Option) *LLMWrapper { + return &LLMWrapper{ + llm: llm, + options: evaluateOptions(opts...), + } +} + +func (w *LLMWrapper) GenerateContent(ctx context.Context, messages []llms.MessageContent, options ...llms.CallOption) (*llms.ContentResponse, error) { + resource := w.options.defaultResource + if w.options.resourceExtract != nil { + resource = w.options.resourceExtract(ctx) + } + + prompts := []string{} + if w.options.promptsExtract != nil { + prompts = w.options.promptsExtract(messages) + } + + reqInfos := llmtokenratelimit.GenerateRequestInfos( + llmtokenratelimit.WithPrompts(prompts), + ) + if w.options.requestInfosExtract != nil { + reqInfos = w.options.requestInfosExtract(ctx) + } + + // Check + entry, err := sentinel.Entry(resource, sentinel.WithTrafficType(base.Inbound), sentinel.WithArgs(reqInfos)) + + if err != nil { + // Block + if w.options.blockFallback != nil { + w.options.blockFallback(ctx) + } + return nil, err + } + // Pass + response, llmErr := w.llm.GenerateContent(ctx, messages, options...) + if llmErr != nil { + return nil, llmErr + } + if err := w.validateResponse(response); err != nil { + return nil, err + } + + usedTokenInfos := &llmtokenratelimit.UsedTokenInfos{} + if w.options.usedTokenInfosExtract != nil { + usedTokenInfos = w.options.usedTokenInfosExtract(response.Choices[0].GenerationInfo) + } else { + // fallback to OpenAI extractor + infos, err := llmtokenratelimit.OpenAITokenExtractor(map[string]any{ + "prompt_tokens": response.Choices[0].GenerationInfo["PromptTokens"], + "completion_tokens": response.Choices[0].GenerationInfo["CompletionTokens"], + "total_tokens": response.Choices[0].GenerationInfo["TotalTokens"], + }) + if err != nil { + return nil, err + } + usedTokenInfos = infos + } + + entry.SetPair(llmtokenratelimit.KeyUsedTokenInfos, usedTokenInfos) + entry.Exit() + + return response, nil +} + +func (w *LLMWrapper) validateResponse(response *llms.ContentResponse) error { + if response == nil || len(response.Choices) == 0 || response.Choices[0].GenerationInfo == nil { + return fmt.Errorf("llm response is nil or empty or missing GenerationInfo") + } + _, exists := response.Choices[0].GenerationInfo["PromptTokens"] + if !exists { + return fmt.Errorf("llm response missing PromptTokens info") + } + _, exists = response.Choices[0].GenerationInfo["CompletionTokens"] + if !exists { + return fmt.Errorf("llm response missing CompletionTokens info") + } + _, exists = response.Choices[0].GenerationInfo["TotalTokens"] + if !exists { + return fmt.Errorf("llm response missing TotalTokens info") + } + return nil +} diff --git a/pkg/adapters/micro/go.mod b/pkg/adapters/micro/go.mod index 5889c7f21..b688275c7 100644 --- a/pkg/adapters/micro/go.mod +++ b/pkg/adapters/micro/go.mod @@ -8,7 +8,7 @@ require ( github.com/alibaba/sentinel-golang v1.0.4 github.com/golang/protobuf v1.5.3 github.com/micro/go-micro/v2 v2.9.1 - github.com/stretchr/testify v1.8.0 + github.com/stretchr/testify v1.8.2 ) require ( @@ -23,19 +23,23 @@ require ( github.com/cpuguy83/go-md2man/v2 v2.0.0 // indirect github.com/davecgh/go-spew v1.1.1 // indirect github.com/dgrijalva/jwt-go v3.2.0+incompatible // indirect + github.com/dgryski/go-rendezvous v0.0.0-20200823014737-9f7001d12a5f // indirect + github.com/dlclark/regexp2 v1.10.0 // indirect github.com/emirpasic/gods v1.12.0 // indirect - github.com/fsnotify/fsnotify v1.4.7 // indirect + github.com/fsnotify/fsnotify v1.4.9 // indirect github.com/ghodss/yaml v1.0.0 // indirect github.com/go-git/gcfg v1.5.0 // indirect github.com/go-git/go-billy/v5 v5.0.0 // indirect github.com/go-git/go-git/v5 v5.1.0 // indirect github.com/go-ole/go-ole v1.2.4 // indirect + github.com/go-redis/redis/v8 v8.11.0 // indirect github.com/gogo/protobuf v1.2.1 // indirect - github.com/google/uuid v1.1.1 // indirect + github.com/google/uuid v1.3.0 // indirect github.com/hashicorp/hcl v1.0.0 // indirect github.com/hpcloud/tail v1.0.0 // indirect github.com/imdario/mergo v0.3.9 // indirect github.com/jbenet/go-context v0.0.0-20150711004518-d14ea06fba99 // indirect + github.com/jinzhu/copier v0.4.0 // indirect github.com/kevinburke/ssh_config v0.0.0-20190725054713-01f96b0aa0cd // indirect github.com/matttproud/golang_protobuf_extensions v1.0.4 // indirect github.com/micro/cli/v2 v2.1.2 // indirect @@ -49,6 +53,7 @@ require ( github.com/oxtoacart/bpool v0.0.0-20190530202638-03653db5a59c // indirect github.com/patrickmn/go-cache v2.1.0+incompatible // indirect github.com/pkg/errors v0.9.1 // indirect + github.com/pkoukk/tiktoken-go v0.1.7 // indirect github.com/pmezard/go-difflib v1.0.0 // indirect github.com/prometheus/client_golang v1.16.0 // indirect github.com/prometheus/client_model v0.3.0 // indirect @@ -58,6 +63,7 @@ require ( github.com/sergi/go-diff v1.1.0 // indirect github.com/shirou/gopsutil/v3 v3.21.6 // indirect github.com/shurcooL/sanitized_anchor_name v1.0.0 // indirect + github.com/spaolacci/murmur3 v1.1.0 // indirect github.com/tklauser/go-sysconf v0.3.6 // indirect github.com/tklauser/numcpus v0.2.2 // indirect github.com/xanzy/ssh-agent v0.2.1 // indirect diff --git a/pkg/adapters/micro/go.sum b/pkg/adapters/micro/go.sum index 314cb46a4..e92140565 100644 --- a/pkg/adapters/micro/go.sum +++ b/pkg/adapters/micro/go.sum @@ -65,6 +65,7 @@ github.com/caddyserver/certmagic v0.10.6/go.mod h1:Y8jcUBctgk/IhpAzlHKfimZNyXCkf github.com/cenkalti/backoff/v4 v4.0.0/go.mod h1:eEew/i+1Q6OrCDZh3WiXYv3+nJwBASZ8Bog/87DQnVg= github.com/census-instrumentation/opencensus-proto v0.2.0/go.mod h1:f6KPmirojxKA12rnyqOA5BBL4O983OfeGPqjHWSTneU= github.com/census-instrumentation/opencensus-proto v0.2.1/go.mod h1:f6KPmirojxKA12rnyqOA5BBL4O983OfeGPqjHWSTneU= +github.com/cespare/xxhash/v2 v2.1.1/go.mod h1:VGX0DQ3Q6kWi7AoAeZDth3/j3BFtOZR5XLFGgcrjCOs= github.com/cespare/xxhash/v2 v2.2.0 h1:DC2CZ1Ep5Y4k3ZQ899DldepgrayRUGE6BBZ/cd9Cj44= github.com/cespare/xxhash/v2 v2.2.0/go.mod h1:VGX0DQ3Q6kWi7AoAeZDth3/j3BFtOZR5XLFGgcrjCOs= github.com/cheekybits/genny v1.0.0/go.mod h1:+tQajlRqAUrPI7DOSpB0XAqZYtQakVtB7wXkRAgjxjQ= @@ -100,7 +101,11 @@ github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= github.com/dgrijalva/jwt-go v3.2.0+incompatible h1:7qlOGliEKZXTDg6OTjfoBKDXWrumCAMpl/TFQ4/5kLM= github.com/dgrijalva/jwt-go v3.2.0+incompatible/go.mod h1:E3ru+11k8xSBh+hMPgOLZmtrrCbhqsmaPHjLKYnJCaQ= +github.com/dgryski/go-rendezvous v0.0.0-20200823014737-9f7001d12a5f h1:lO4WD4F/rVNCu3HqELle0jiPLLBs70cWOduZpkS1E78= +github.com/dgryski/go-rendezvous v0.0.0-20200823014737-9f7001d12a5f/go.mod h1:cuUVRXasLTGF7a8hSLbxyZXjz+1KgoB3wDUb6vlszIc= github.com/dimchansky/utfbom v1.1.0/go.mod h1:rO41eb7gLfo8SF1jd9F8HplJm1Fewwi4mQvIirEdv+8= +github.com/dlclark/regexp2 v1.10.0 h1:+/GIL799phkJqYW+3YbOd8LCcbHzT0Pbo8zl70MHsq0= +github.com/dlclark/regexp2 v1.10.0/go.mod h1:DHkYz0B9wPfa6wondMfaivmHpzrQ3v9q8cnmRbL6yW8= github.com/dnaeon/go-vcr v0.0.0-20180814043457-aafff18a5cc2/go.mod h1:aBB1+wY4s93YsC3HHjMBMrwTj2R9FHDzUr9KyGc8n1E= github.com/dnsimple/dnsimple-go v0.30.0/go.mod h1:O5TJ0/U6r7AfT8niYNlmohpLbCSG+c71tQlGr9SeGrg= github.com/docker/distribution v2.7.1+incompatible/go.mod h1:J2gT2udsDAN96Uj4KfcMRqY0/ypR+oyYUYmja8H+y+w= @@ -120,8 +125,9 @@ github.com/exoscale/egoscale v0.18.1/go.mod h1:Z7OOdzzTOz1Q1PjQXumlz9Wn/CddH0zSY github.com/fatih/structs v1.1.0/go.mod h1:9NiDSp5zOcgEDl+j00MP/WkGVPOlPRLejGD8Ga6PJ7M= github.com/flynn/go-shlex v0.0.0-20150515145356-3f9db97f8568/go.mod h1:xEzjJPgXI435gkrCt3MPfRiAkVrwSbHsst4LCFVfpJc= github.com/forestgiant/sliceutil v0.0.0-20160425183142-94783f95db6c/go.mod h1:pFdJbAhRf7rh6YYMUdIQGyzne6zYL1tCUW8QV2B3UfY= -github.com/fsnotify/fsnotify v1.4.7 h1:IXs+QLmnXW2CcXuY+8Mzv/fWEsPGWxqefPtCP5CnV9I= github.com/fsnotify/fsnotify v1.4.7/go.mod h1:jwhsz4b93w/PPRr/qN1Yymfu8t87LnFCMoQvtojpjFo= +github.com/fsnotify/fsnotify v1.4.9 h1:hsms1Qyu0jgnwNXIxa+/V/PDsU6CfLf6CNO8H7IWoS4= +github.com/fsnotify/fsnotify v1.4.9/go.mod h1:znqG4EE+3YCdAaPaxE2ZRY/06pZUdp0tY4IgpuI1SZQ= github.com/fsouza/go-dockerclient v1.6.0/go.mod h1:YWwtNPuL4XTX1SKJQk86cWPmmqwx+4np9qfPbb+znGc= github.com/ghodss/yaml v1.0.0 h1:wQHKEahhL6wmXdzwWG11gIVCkOv05bNOh+Rxn0yngAk= github.com/ghodss/yaml v1.0.0/go.mod h1:4dBDuWmgqj2HViK6kFavaiC9ZROes6MMH2rRYeMEF04= @@ -145,6 +151,8 @@ github.com/go-logfmt/logfmt v0.3.0/go.mod h1:Qt1PoO58o5twSAckw1HlFXLmHsOX5/0LbT9 github.com/go-logfmt/logfmt v0.4.0/go.mod h1:3RMwSq7FuexP4Kalkev3ejPJsZTpXXBr9+V4qmtdjCk= github.com/go-ole/go-ole v1.2.4 h1:nNBDSCOigTSiarFpYE9J/KtEA1IOW4CNeqT9TQDqCxI= github.com/go-ole/go-ole v1.2.4/go.mod h1:XCwSNxSkXRo4vlyPy93sltvi/qJq0jqQhjqQNIwKuxM= +github.com/go-redis/redis/v8 v8.11.0 h1:O1Td0mQ8UFChQ3N9zFQqo6kTU2cJ+/it88gDB+zg0wo= +github.com/go-redis/redis/v8 v8.11.0/go.mod h1:DLomh7y2e3ggQXQLd1YgmvIfecPJoFl7WU5SOQ/r06M= github.com/go-stack/stack v1.8.0/go.mod h1:v0f6uXyyMGvRgIKkXu+yp6POWl0qKG85gN/melR3HDY= github.com/go-telegram-bot-api/telegram-bot-api v4.6.4+incompatible/go.mod h1:qf9acutJ8cwBUhm1bqgz6Bei9/C/c93FPDljKWwsOgM= github.com/gobwas/httphead v0.0.0-20180130184737-2c6c146eadee/go.mod h1:L0fX3K22YWvt/FAX9NnzrNzcI4wNYi9Yku4O0LKYflo= @@ -174,6 +182,7 @@ github.com/golang/protobuf v1.4.0-rc.1.0.20200221234624-67d41d38c208/go.mod h1:x github.com/golang/protobuf v1.4.0-rc.2/go.mod h1:LlEzMj4AhA7rCAGe4KMBDvJI+AwstrUpVNzEA03Pprs= github.com/golang/protobuf v1.4.0-rc.4.0.20200313231945-b860323f09d0/go.mod h1:WU3c8KckQ9AFe+yFwt9sWVRKCVIyN9cPHBJSNnbL67w= github.com/golang/protobuf v1.4.0/go.mod h1:jodUvKwWbYaEsadDk5Fwe5c77LiNKVO9IDvqG2KuDX0= +github.com/golang/protobuf v1.4.2/go.mod h1:oDoupMAO8OvCJWAcko0GGGIgR6R6ocIYbsSw735rRwI= github.com/golang/protobuf v1.5.0/go.mod h1:FsONVRAS9T7sI+LIUmWTfcYkHO4aIWwzhcaSAoJOfIk= github.com/golang/protobuf v1.5.3 h1:KhyjKVUg7Usr/dYsdSqoFveMYd5ko72D+zANwlG1mmg= github.com/golang/protobuf v1.5.3/go.mod h1:XVQd3VNwM+JqD3oG2Ue2ip4fOMUkwXdXDdiuN0vRsmY= @@ -186,6 +195,7 @@ github.com/google/go-cmp v0.3.0/go.mod h1:8QqcDgzrUqlUb/G2PQTWiueGozuR1884gddMyw github.com/google/go-cmp v0.3.1/go.mod h1:8QqcDgzrUqlUb/G2PQTWiueGozuR1884gddMywk6iLU= github.com/google/go-cmp v0.4.0/go.mod h1:v8dTdLbMG2kIc/vJvl+f65V22dbkXbowE6jgT/gNBxE= github.com/google/go-cmp v0.5.5/go.mod h1:v8dTdLbMG2kIc/vJvl+f65V22dbkXbowE6jgT/gNBxE= +github.com/google/go-cmp v0.5.6/go.mod h1:v8dTdLbMG2kIc/vJvl+f65V22dbkXbowE6jgT/gNBxE= github.com/google/go-cmp v0.6.0 h1:ofyhxvXcZhMsU5ulbFiLKl/XBFqE1GSq7atu8tAmTRI= github.com/google/go-querystring v1.0.0/go.mod h1:odCYkC5MyYFN7vkCjXpyrEuKhc/BUO6wN/zVPAxq5ck= github.com/google/gofuzz v1.0.0/go.mod h1:dBl0BpW6vV/+mYPU4Po3pmUjxk6FQPldtuIdl/M65Eg= @@ -193,8 +203,9 @@ github.com/google/martian v2.1.0+incompatible/go.mod h1:9I4somxYTbIHy5NJKHRl3wXi github.com/google/pprof v0.0.0-20181206194817-3ea8567a2e57/go.mod h1:zfwlbNMJ+OItoe0UupaVj+oy1omPYYDuagoSzA8v9mc= github.com/google/pprof v0.0.0-20190515194954-54271f7e092f/go.mod h1:zfwlbNMJ+OItoe0UupaVj+oy1omPYYDuagoSzA8v9mc= github.com/google/renameio v0.1.0/go.mod h1:KWCgfxg9yswjAJkECMjeO8J8rahYeXnNhOm40UhjYkI= -github.com/google/uuid v1.1.1 h1:Gkbcsh/GbpXz7lPftLA3P6TYMwjCLYm83jiFQZF/3gY= github.com/google/uuid v1.1.1/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo= +github.com/google/uuid v1.3.0 h1:t6JiXgmwXMjEs8VusXIJk2BXHsn+wx8BZdTaoZ5fu7I= +github.com/google/uuid v1.3.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo= github.com/googleapis/gax-go/v2 v2.0.4/go.mod h1:0Wqv26UfaUD9n4G6kQubkQ+KchISgw+vpHVxEJEs9eg= github.com/googleapis/gax-go/v2 v2.0.5/go.mod h1:DWXyrwAJ9X0FpwwEdw+IPEYBICEFu5mhpdKc/us6bOk= github.com/gophercloud/gophercloud v0.3.0/go.mod h1:vxM41WHh5uqHVBMZHzuwNOHh8XEoIEcSTewFxm1c5g8= @@ -230,6 +241,8 @@ github.com/imdario/mergo v0.3.9/go.mod h1:2EnlNZ0deacrJVfApfmtdGgDfMuh/nq6Ok1EcJ github.com/jbenet/go-context v0.0.0-20150711004518-d14ea06fba99 h1:BQSFePA1RWJOlocH6Fxy8MmwDt+yVQYULKfN0RoTN8A= github.com/jbenet/go-context v0.0.0-20150711004518-d14ea06fba99/go.mod h1:1lJo3i6rXxKeerYnT8Nvf0QmHCRC1n8sfWVwXF2Frvo= github.com/jessevdk/go-flags v1.4.0/go.mod h1:4FA24M0QyGHXBuZZK/XkWh8h0e1EYbRYJSGM75WSRxI= +github.com/jinzhu/copier v0.4.0 h1:w3ciUoD19shMCRargcpm0cm91ytaBhDvuRpz1ODO/U8= +github.com/jinzhu/copier v0.4.0/go.mod h1:DfbEm0FYsaqBcKcFuvmOZb218JkPGtvSHsKg8S8hyyg= github.com/jmespath/go-jmespath v0.0.0-20180206201540-c2b33e8439af/go.mod h1:Nht3zPeWKUH0NzdCt2Blrr5ys8VGpn0CEB0cQHVjt7k= github.com/jonboulle/clockwork v0.1.0 h1:VKV+ZcuP6l3yW9doeqz6ziZGgcynBVQO+obU0+0hcPo= github.com/jonboulle/clockwork v0.1.0/go.mod h1:Ii8DK3G1RaLaWxj9trq07+26W01tbo22gdxWY5EU2bo= @@ -313,10 +326,19 @@ github.com/nrdcg/auroradns v1.0.0/go.mod h1:6JPXKzIRzZzMqtTDgueIhTi6rFf1QvYE/Hzq github.com/nrdcg/dnspod-go v0.4.0/go.mod h1:vZSoFSFeQVm2gWLMkyX61LZ8HI3BaqtHZWgPTGKr6KQ= github.com/nrdcg/goinwx v0.6.1/go.mod h1:XPiut7enlbEdntAqalBIqcYcTEVhpv/dKWgDCX2SwKQ= github.com/nrdcg/namesilo v0.2.1/go.mod h1:lwMvfQTyYq+BbjJd30ylEG4GPSS6PII0Tia4rRpRiyw= +github.com/nxadm/tail v1.4.4 h1:DQuhQpB1tVlglWS2hLQ5OV6B5r8aGxSrPc5Qo6uTN78= +github.com/nxadm/tail v1.4.4/go.mod h1:kenIhsEOeOJmVchQTgglprH7qJGnHDVpk1VPCcaMI8A= github.com/olekukonko/tablewriter v0.0.1/go.mod h1:vsDQFd/mU46D+Z4whnwzcISnGGzXWMclvtLoiIKAKIo= github.com/onsi/ginkgo v1.6.0/go.mod h1:lLunBs/Ym6LB5Z9jYTR76FiuTmxDTDusOGeTQH+WWjE= github.com/onsi/ginkgo v1.7.0/go.mod h1:lLunBs/Ym6LB5Z9jYTR76FiuTmxDTDusOGeTQH+WWjE= +github.com/onsi/ginkgo v1.12.1/go.mod h1:zj2OWP4+oCPe1qIXoGWkgMRwljMUYCdkwsT2108oapk= +github.com/onsi/ginkgo v1.15.0 h1:1V1NfVQR87RtWAgp1lv9JZJ5Jap+XFGKPi00andXGi4= +github.com/onsi/ginkgo v1.15.0/go.mod h1:hF8qUzuuC8DJGygJH3726JnCZX4MYbRB8yFfISqnKUg= github.com/onsi/gomega v1.4.3/go.mod h1:ex+gbHU/CVuBBDIJjb2X0qEXbFg53c61hWP/1CpauHY= +github.com/onsi/gomega v1.7.1/go.mod h1:XdKZgCCFLUoM/7CFJVPcG8C1xQ1AJ0vpAezJrB7JYyY= +github.com/onsi/gomega v1.10.1/go.mod h1:iN09h71vgCQne3DLsj+A5owkum+a2tYe+TOCB1ybHNo= +github.com/onsi/gomega v1.10.5 h1:7n6FEkpFmfCoo2t+YYqXH0evK+a9ICQz0xcAy9dYcaQ= +github.com/onsi/gomega v1.10.5/go.mod h1:gza4q3jKQJijlu05nKWRCW/GavJumGt8aNRxWg7mt48= github.com/opencontainers/go-digest v0.0.0-20180430190053-c9281466c8b2/go.mod h1:cMLVZDEM3+U2I4VmLI6N8jQYUd2OVphdqWwCJHrFt2s= github.com/opencontainers/go-digest v1.0.0-rc1/go.mod h1:cMLVZDEM3+U2I4VmLI6N8jQYUd2OVphdqWwCJHrFt2s= github.com/opencontainers/image-spec v1.0.1/go.mod h1:BtxoFyWECRxE4U/7sNtV5W15zMzWCbyJoFRP3s7yZA0= @@ -337,6 +359,8 @@ github.com/pkg/errors v0.8.0/go.mod h1:bwawxfHBFNV+L2hUp1rHADufV3IMtnDRdf1r5NINE github.com/pkg/errors v0.8.1/go.mod h1:bwawxfHBFNV+L2hUp1rHADufV3IMtnDRdf1r5NINEl0= github.com/pkg/errors v0.9.1 h1:FEBLx1zS214owpjy7qsBeixbURkuhQAwrK5UwLGTwt4= github.com/pkg/errors v0.9.1/go.mod h1:bwawxfHBFNV+L2hUp1rHADufV3IMtnDRdf1r5NINEl0= +github.com/pkoukk/tiktoken-go v0.1.7 h1:qOBHXX4PHtvIvmOtyg1EeKlwFRiMKAcoMp4Q+bLQDmw= +github.com/pkoukk/tiktoken-go v0.1.7/go.mod h1:9NiV+i9mJKGj1rYOT+njbv+ZwA/zJxYdewGl6qVatpg= github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM= github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4= github.com/prometheus/client_golang v0.9.1/go.mod h1:7SWBe2y4D6OKWSNQJUaRYU/AaXPKyh/dDVn+NZz0KFw= @@ -387,17 +411,21 @@ github.com/smartystreets/assertions v0.0.0-20180927180507-b2de0cb4f26d/go.mod h1 github.com/smartystreets/goconvey v0.0.0-20190330032615-68dc04aab96a/go.mod h1:syvi0/a8iFYH4r/RixwvyeAJjdLS9QV7WQ/tjFTllLA= github.com/soheilhy/cmux v0.1.4 h1:0HKaf1o97UwFjHH9o5XsHUOF+tqmdA7KEzXLpiyaw0E= github.com/soheilhy/cmux v0.1.4/go.mod h1:IM3LyeVVIOuxMH7sFAkER9+bJ4dT7Ms6E4xg4kGIyLM= +github.com/spaolacci/murmur3 v1.1.0 h1:7c1g84S4BPRrfL5Xrdp6fOJ206sU9y293DDHaoy0bLI= +github.com/spaolacci/murmur3 v1.1.0/go.mod h1:JwIasOWyU6f++ZhiEuf87xNszmSA2myDM2Kzu9HwQUA= github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME= github.com/stretchr/objx v0.1.1/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME= -github.com/stretchr/objx v0.4.0 h1:M2gUjqZET1qApGOWNSnZ49BAIMX4F/1plDv3+l31EJ4= github.com/stretchr/objx v0.4.0/go.mod h1:YvHI0jy2hoMjB+UWwv71VJQ9isScKT/TqJzVSSt89Yw= +github.com/stretchr/objx v0.5.0 h1:1zr/of2m5FGMsad5YfcqgdqdWrIhu+EBEJRhR1U7z/c= +github.com/stretchr/objx v0.5.0/go.mod h1:Yh+to48EsGEfYuaHDzXPcE3xhTkx73EhmCGUpEOglKo= github.com/stretchr/testify v1.2.2/go.mod h1:a8OnRcib4nhh0OaRAV+Yts87kKdq0PP7pXfy6kDkUVs= github.com/stretchr/testify v1.3.0/go.mod h1:M5WIy9Dh21IEIfnGCwXGc5bZfKNJtfHm1UVUgZn+9EI= github.com/stretchr/testify v1.4.0/go.mod h1:j7eGeouHqKxXV5pUuKE4zz7dFj8WfuZ+81PSLYec5m4= github.com/stretchr/testify v1.6.1/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg= github.com/stretchr/testify v1.7.1/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg= -github.com/stretchr/testify v1.8.0 h1:pSgiaMZlXftHpm5L7V1+rVB+AZJydKsMxsQBIJw4PKk= github.com/stretchr/testify v1.8.0/go.mod h1:yNjHg4UonilssWZ8iaSj1OCr/vHnekPRkoO+kdMU+MU= +github.com/stretchr/testify v1.8.2 h1:+h33VjcLVPDHtOdpUCuF+7gSuG3yGIftsP1YvFihtJ8= +github.com/stretchr/testify v1.8.2/go.mod h1:w2LPCIKwWwSfY2zedu0+kehJoqGctiVI29o6fzry7u4= github.com/syndtr/gocapability v0.0.0-20170704070218-db04d3cc01c8/go.mod h1:hkRG7XYTFWNJGYcbNJQlaLq0fg1yr4J4t/NcTQtrfww= github.com/technoweenie/multipartstreamer v1.0.1/go.mod h1:jNVxdtShOxzAsukZwTSw6MDx5eUJoiEBsSvzDU9uzog= github.com/timewasted/linode v0.0.0-20160829202747-37e84520dcf7/go.mod h1:imsgLplxEC/etjIhdr3dNzV3JeT27LbVu5pYWm0JCBY= @@ -420,6 +448,7 @@ github.com/xeipuuv/gojsonschema v0.0.0-20180618132009-1d523034197f/go.mod h1:5yf github.com/xeipuuv/gojsonschema v1.1.0/go.mod h1:5yf86TLmAcydyeJq5YvxkGPE2fm/u4myDekKRoLuqhs= github.com/xiang90/probing v0.0.0-20190116061207-43a291ad63a2 h1:eY9dn8+vbi4tKz5Qo6v2eYzo7kUS51QINcR5jNpbZS8= github.com/xiang90/probing v0.0.0-20190116061207-43a291ad63a2/go.mod h1:UETIi67q53MR2AWcXfiuqkDkRtnGDLqkBTpCHuJHxtU= +github.com/yuin/goldmark v1.2.1/go.mod h1:3hX8gzYuyVAZsxl0MRgGTJEmQBFcNTphYh9decYSb74= go.etcd.io/bbolt v1.3.4/go.mod h1:G5EMThwa9y8QZGBClrRx5EY+Yw9kAhnjy3bSjsnlVTQ= go.opencensus.io v0.20.1/go.mod h1:6WKK9ahsWS3RSO+PY9ZHZUfv2irvY6gN279GOPZjmmk= go.opencensus.io v0.20.2/go.mod h1:6WKK9ahsWS3RSO+PY9ZHZUfv2irvY6gN279GOPZjmmk= @@ -456,6 +485,7 @@ golang.org/x/crypto v0.0.0-20191011191535-87dc89f01550/go.mod h1:yigFU9vqHzYiE8U golang.org/x/crypto v0.0.0-20200302210943-78000ba7a073/go.mod h1:LzIPMQfyMNhhGPhUkYOs5KpL4U8rLKemX1yGLhDgUto= golang.org/x/crypto v0.0.0-20200323165209-0ec3e9974c59/go.mod h1:LzIPMQfyMNhhGPhUkYOs5KpL4U8rLKemX1yGLhDgUto= golang.org/x/crypto v0.0.0-20200510223506-06a226fb4e37/go.mod h1:LzIPMQfyMNhhGPhUkYOs5KpL4U8rLKemX1yGLhDgUto= +golang.org/x/crypto v0.0.0-20200622213623-75b288015ac9/go.mod h1:LzIPMQfyMNhhGPhUkYOs5KpL4U8rLKemX1yGLhDgUto= golang.org/x/crypto v0.24.0 h1:mnl8DM0o513X8fdIkmyFE/5hTYxbwYOjDS/+rK6qpRI= golang.org/x/crypto v0.24.0/go.mod h1:Z1PMYSOR5nyMcyAVAIQSKCDwalqy85Aqn1x3Ws4L5DM= golang.org/x/exp v0.0.0-20190121172915-509febef88a4/go.mod h1:CJ0aWSM057203Lf6IL+f9T1iT9GByDxfZKAQTCR3kQA= @@ -479,6 +509,7 @@ golang.org/x/mobile v0.0.0-20190719004257-d2bd2a29d028/go.mod h1:E/iHnbuqvinMTCc golang.org/x/mod v0.0.0-20190513183733-4bf6d317e70e/go.mod h1:mXi4GBBbnImb6dmsKGUJ2LatrhH/nqhxcFungHvyanc= golang.org/x/mod v0.1.0/go.mod h1:0QHyrYULN0/3qlju5TqG8bIK38QM8yzMo5ekMj3DlcY= golang.org/x/mod v0.1.1-0.20191105210325-c90efee705ee/go.mod h1:QqPTAvyqsEbceGzBzNggFXnrqF1CaUcvgkdR5Ot7KZg= +golang.org/x/mod v0.3.0/go.mod h1:s0Qsj1ACt9ePp/hMypM3fl4fZqREWJwdYDEqhRiZZUA= golang.org/x/mod v0.18.0 h1:5+9lSbEzPSdWkH32vYPBwEpX8KwDbM52Ud9xBUvNlb0= golang.org/x/net v0.0.0-20180611182652-db08ff08e862/go.mod h1:mL1N/T3taQHkDXs73rZJwtUhF3w3ftmwwsq0BUmARs4= golang.org/x/net v0.0.0-20180724234803-3673e40ba225/go.mod h1:mL1N/T3taQHkDXs73rZJwtUhF3w3ftmwwsq0BUmARs4= @@ -502,7 +533,10 @@ golang.org/x/net v0.0.0-20190923162816-aa69164e4478/go.mod h1:z5CRVTTTmAJ677TzLL golang.org/x/net v0.0.0-20190930134127-c5a3c61f89f3/go.mod h1:z5CRVTTTmAJ677TzLLGU+0bjPO0LkuOLi4/5GtJWs/s= golang.org/x/net v0.0.0-20191027093000-83d349e8ac1a/go.mod h1:z5CRVTTTmAJ677TzLLGU+0bjPO0LkuOLi4/5GtJWs/s= golang.org/x/net v0.0.0-20200301022130-244492dfa37a/go.mod h1:z5CRVTTTmAJ677TzLLGU+0bjPO0LkuOLi4/5GtJWs/s= +golang.org/x/net v0.0.0-20200520004742-59133d7f0dd7/go.mod h1:qpuaurCH72eLCgpAm/N6yyVIVM9cpaDIP3A8BGJEC5A= golang.org/x/net v0.0.0-20200520182314-0ba52f642ac2/go.mod h1:qpuaurCH72eLCgpAm/N6yyVIVM9cpaDIP3A8BGJEC5A= +golang.org/x/net v0.0.0-20201021035429-f5854403a974/go.mod h1:sp8m0HH+o8qH0wwXwYZr8TS3Oi6o0r6Gce1SSxlDquU= +golang.org/x/net v0.0.0-20201202161906-c7110b5ffcbb/go.mod h1:sp8m0HH+o8qH0wwXwYZr8TS3Oi6o0r6Gce1SSxlDquU= golang.org/x/net v0.26.0 h1:soB7SVo0PWrY4vPW/+ay0jKDNScG2X9wFeYlXIvJsOQ= golang.org/x/net v0.26.0/go.mod h1:5YKkiSynbBIh3p6iOc/vibscux0x38BZDkn8sCUPxHE= golang.org/x/oauth2 v0.0.0-20180821212333-d2e6202438be/go.mod h1:N/0e6XlmueqKjAGxoOufVs8QHGRruUQn6yWY3a++T0U= @@ -514,6 +548,7 @@ golang.org/x/sync v0.0.0-20181221193216-37e7f081c4d4/go.mod h1:RxMgew5VJxzue5/jJ golang.org/x/sync v0.0.0-20190227155943-e225da77a7e6/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM= golang.org/x/sync v0.0.0-20190423024810-112230192c58/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM= golang.org/x/sync v0.0.0-20190911185100-cd5d95a43a6e/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM= +golang.org/x/sync v0.0.0-20201020160332-67f06af15bc9/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM= golang.org/x/sync v0.7.0 h1:YsImfSBoP9QPYL0xyKJPq0gcaJdG3rInoqxTWbfQu9M= golang.org/x/sys v0.0.0-20180622082034-63fc586f45fe/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY= golang.org/x/sys v0.0.0-20180830151530-49385e6e1522/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY= @@ -539,10 +574,14 @@ golang.org/x/sys v0.0.0-20190801041406-cbf593c0f2f3/go.mod h1:h1NjWce9XRLGQEsW7w golang.org/x/sys v0.0.0-20190904154756-749cb33beabd/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs= golang.org/x/sys v0.0.0-20190916202348-b4ddaad3f8a3/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs= golang.org/x/sys v0.0.0-20190924154521-2837fb4f24fe/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs= +golang.org/x/sys v0.0.0-20191005200804-aed5e4c7ecf9/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs= +golang.org/x/sys v0.0.0-20191120155948-bd437916bb0e/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs= golang.org/x/sys v0.0.0-20200202164722-d101bd2416d5/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs= golang.org/x/sys v0.0.0-20200302150141-5c8b2ff67527/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs= golang.org/x/sys v0.0.0-20200323222414-85ca7c5b95cd/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs= golang.org/x/sys v0.0.0-20200523222454-059865788121/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs= +golang.org/x/sys v0.0.0-20200930185726-fdedc70b468f/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs= +golang.org/x/sys v0.0.0-20210112080510-489259a85091/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs= golang.org/x/sys v0.0.0-20210316164454-77fc1eacc6aa/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs= golang.org/x/sys v0.21.0 h1:rF+pYz3DAGSQAxAu1CbC7catZg4ebC4UIeIhKxBZvws= golang.org/x/sys v0.21.0/go.mod h1:/VUhepiaJMQUp4+oa/7Zr1D23ma6VTLIYjOOTFZPUcA= @@ -550,6 +589,7 @@ golang.org/x/term v0.21.0 h1:WVXCp+/EBEHOj53Rvu+7KiT/iElMrO8ACK16SMZ3jaA= golang.org/x/text v0.3.0/go.mod h1:NqM8EUOU14njkJ3fqMW+pc6Ldnwhi/IjpwHt7yyuwOQ= golang.org/x/text v0.3.1-0.20180807135948-17ff2d5776d2/go.mod h1:NqM8EUOU14njkJ3fqMW+pc6Ldnwhi/IjpwHt7yyuwOQ= golang.org/x/text v0.3.2/go.mod h1:bEr9sfX3Q8Zfm5fL9x+3itogRgK3+ptLWKqgva+5dAk= +golang.org/x/text v0.3.3/go.mod h1:5Zoc/QRtKVWzQhOtBMvqHzDpF6irO9z98xDceosuGiQ= golang.org/x/text v0.16.0 h1:a94ExnEXNtEwYLGJSIUxnWoxoRz/ZcCsV63ROupILh4= golang.org/x/text v0.16.0/go.mod h1:GhwF1Be+LQoKShO3cGOHzqOgRrGaYc9AvblQOmPVHnI= golang.org/x/time v0.0.0-20181108054448-85acf8d2951c/go.mod h1:tRJNPiyCQ0inRvYxbN9jk5I+vvW/OXSQhTDSoE431IQ= @@ -577,14 +617,17 @@ golang.org/x/tools v0.0.0-20190911174233-4f2ddba30aff/go.mod h1:b+2E5dAYhXwXZwtn golang.org/x/tools v0.0.0-20191012152004-8de300cfc20a/go.mod h1:b+2E5dAYhXwXZwtnZ6UAqBI28+e2cm9otk0dWdXHAEo= golang.org/x/tools v0.0.0-20191029041327-9cc4af7d6b2c/go.mod h1:b+2E5dAYhXwXZwtnZ6UAqBI28+e2cm9otk0dWdXHAEo= golang.org/x/tools v0.0.0-20191029190741-b9c20aec41a5/go.mod h1:b+2E5dAYhXwXZwtnZ6UAqBI28+e2cm9otk0dWdXHAEo= +golang.org/x/tools v0.0.0-20191119224855-298f0cb1881e/go.mod h1:b+2E5dAYhXwXZwtnZ6UAqBI28+e2cm9otk0dWdXHAEo= golang.org/x/tools v0.0.0-20191125144606-a911d9008d1f/go.mod h1:b+2E5dAYhXwXZwtnZ6UAqBI28+e2cm9otk0dWdXHAEo= golang.org/x/tools v0.0.0-20191216052735-49a3e744a425/go.mod h1:TB2adYChydJhpapKDTa4BR/hXlZSLoq2Wpct/0txZ28= golang.org/x/tools v0.0.0-20191216173652-a0e659d51361/go.mod h1:TB2adYChydJhpapKDTa4BR/hXlZSLoq2Wpct/0txZ28= +golang.org/x/tools v0.0.0-20201224043029-2b0845dc783e/go.mod h1:emZCQorbCU4vsT4fOWvOPXz4eW1wZW4PmDk9uLelYpA= golang.org/x/tools v0.22.0 h1:gqSGLZqv+AI9lIQzniJ0nZDRG5GBPsSi+DRNHWNz6yA= golang.org/x/tools v0.22.0/go.mod h1:aCwcsjqvq7Yqt6TNyX7QMU2enbQ/Gt0bo6krSeEri+c= golang.org/x/xerrors v0.0.0-20190717185122-a985d3407aa7/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0= golang.org/x/xerrors v0.0.0-20191011141410-1b5146add898/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0= golang.org/x/xerrors v0.0.0-20191204190536-9bdfabe68543/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0= +golang.org/x/xerrors v0.0.0-20200804184101-5ec99f83aff1/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0= google.golang.org/api v0.3.1/go.mod h1:6wY9I6uQWHQ8EM57III9mq/AjF+i8G65rmVagqKMtkk= google.golang.org/api v0.4.0/go.mod h1:8k5glujaEP+g9n7WNsDg8QP6cUVNI86fCNMcbazEtwE= google.golang.org/api v0.7.0/go.mod h1:WtwebWUNSVBH/HAw79HIFXZNqEvBhG+Ra+ax0hx3E3M= @@ -621,6 +664,7 @@ google.golang.org/protobuf v0.0.0-20200228230310-ab0ca4ff8a60/go.mod h1:cfTl7dwQ google.golang.org/protobuf v1.20.1-0.20200309200217-e05f789c0967/go.mod h1:A+miEFZTKqfCUM6K7xSMQL9OKL/b6hQv+e19PK+JZNE= google.golang.org/protobuf v1.21.0/go.mod h1:47Nbq4nVaFHyn7ilMalzfO3qCViNmqZ2kzikPIcrTAo= google.golang.org/protobuf v1.22.0/go.mod h1:EGpADcykh3NcUnDUJcl1+ZksZNG86OlYog2l/sGQquU= +google.golang.org/protobuf v1.23.0/go.mod h1:EGpADcykh3NcUnDUJcl1+ZksZNG86OlYog2l/sGQquU= google.golang.org/protobuf v1.26.0-rc.1/go.mod h1:jlhhOSvTdKEhbULTjvd4ARK9grFBp09yW+WbY/TyQbw= google.golang.org/protobuf v1.26.0/go.mod h1:9q0QmTI4eRPtz6boOQmLYwt+qCgq0jsYwAQnmE0givc= google.golang.org/protobuf v1.30.0 h1:kPPoIgf3TsEvrm0PFe15JQ+570QVxYzEvvHqChK+cng= @@ -650,6 +694,7 @@ gopkg.in/yaml.v2 v2.0.0-20170812160011-eb3733d160e7/go.mod h1:JAlM8MvJe8wmxCU4Bl gopkg.in/yaml.v2 v2.2.1/go.mod h1:hI93XBmqTisBFMUTm0b8Fm+jr3Dg1NNxqwp+5A1VGuI= gopkg.in/yaml.v2 v2.2.2/go.mod h1:hI93XBmqTisBFMUTm0b8Fm+jr3Dg1NNxqwp+5A1VGuI= gopkg.in/yaml.v2 v2.2.4/go.mod h1:hI93XBmqTisBFMUTm0b8Fm+jr3Dg1NNxqwp+5A1VGuI= +gopkg.in/yaml.v2 v2.3.0/go.mod h1:hI93XBmqTisBFMUTm0b8Fm+jr3Dg1NNxqwp+5A1VGuI= gopkg.in/yaml.v2 v2.4.0 h1:D8xgwECY7CYvx+Y2n4sBz93Jn9JRvxdiyyo8CTfuKaY= gopkg.in/yaml.v2 v2.4.0/go.mod h1:RDklbk79AGWmwhnvt/jBztapEOGDOx6ZbXqjP6csGnQ= gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM= diff --git a/pkg/datasource/apollo/go.mod b/pkg/datasource/apollo/go.mod index a1ed5ce14..ee12696db 100644 --- a/pkg/datasource/apollo/go.mod +++ b/pkg/datasource/apollo/go.mod @@ -8,7 +8,7 @@ require ( github.com/alibaba/sentinel-golang v1.0.4 github.com/apolloconfig/agollo/v4 v4.0.9 github.com/pkg/errors v0.9.1 - github.com/stretchr/testify v1.8.0 + github.com/stretchr/testify v1.8.2 ) require ( @@ -16,27 +16,33 @@ require ( github.com/beorn7/perks v1.0.1 // indirect github.com/cespare/xxhash/v2 v2.2.0 // indirect github.com/davecgh/go-spew v1.1.1 // indirect + github.com/dgryski/go-rendezvous v0.0.0-20200823014737-9f7001d12a5f // indirect + github.com/dlclark/regexp2 v1.10.0 // indirect github.com/fsnotify/fsnotify v1.5.1 // indirect github.com/go-ole/go-ole v1.2.4 // indirect + github.com/go-redis/redis/v8 v8.11.0 // indirect github.com/golang/protobuf v1.5.3 // indirect - github.com/google/uuid v1.1.2 // indirect + github.com/google/uuid v1.3.0 // indirect github.com/hashicorp/hcl v1.0.0 // indirect + github.com/jinzhu/copier v0.4.0 // indirect github.com/magiconair/properties v1.8.5 // indirect github.com/matttproud/golang_protobuf_extensions v1.0.4 // indirect github.com/mitchellh/mapstructure v1.4.2 // indirect github.com/pelletier/go-toml v1.9.4 // indirect + github.com/pkoukk/tiktoken-go v0.1.7 // indirect github.com/pmezard/go-difflib v1.0.0 // indirect github.com/prometheus/client_golang v1.16.0 // indirect github.com/prometheus/client_model v0.3.0 // indirect github.com/prometheus/common v0.42.0 // indirect github.com/prometheus/procfs v0.10.1 // indirect github.com/shirou/gopsutil/v3 v3.21.6 // indirect + github.com/spaolacci/murmur3 v1.1.0 // indirect github.com/spf13/afero v1.6.0 // indirect github.com/spf13/cast v1.4.1 // indirect github.com/spf13/jwalterweatherman v1.1.0 // indirect github.com/spf13/pflag v1.0.5 // indirect github.com/spf13/viper v1.9.0 // indirect - github.com/stretchr/objx v0.4.0 // indirect + github.com/stretchr/objx v0.5.0 // indirect github.com/subosito/gotenv v1.2.0 // indirect github.com/tklauser/go-sysconf v0.3.6 // indirect github.com/tklauser/numcpus v0.2.2 // indirect diff --git a/pkg/datasource/apollo/go.sum b/pkg/datasource/apollo/go.sum index 171297755..201fdf338 100644 --- a/pkg/datasource/apollo/go.sum +++ b/pkg/datasource/apollo/go.sum @@ -61,6 +61,7 @@ github.com/bgentry/speakeasy v0.1.0/go.mod h1:+zsyZBPWlz7T6j88CTgSN5bM796AkVf0kB github.com/bketelsen/crypt v0.0.4/go.mod h1:aI6NrJ0pMGgvZKL1iVgXLnfIFJtfV+bKCoqOes/6LfM= github.com/census-instrumentation/opencensus-proto v0.2.1/go.mod h1:f6KPmirojxKA12rnyqOA5BBL4O983OfeGPqjHWSTneU= github.com/cespare/xxhash v1.1.0/go.mod h1:XrSqR1VqqWfGrhpAt58auRo0WTKS1nRRg3ghfAqPWnc= +github.com/cespare/xxhash/v2 v2.1.1/go.mod h1:VGX0DQ3Q6kWi7AoAeZDth3/j3BFtOZR5XLFGgcrjCOs= github.com/cespare/xxhash/v2 v2.2.0 h1:DC2CZ1Ep5Y4k3ZQ899DldepgrayRUGE6BBZ/cd9Cj44= github.com/cespare/xxhash/v2 v2.2.0/go.mod h1:VGX0DQ3Q6kWi7AoAeZDth3/j3BFtOZR5XLFGgcrjCOs= github.com/chzyer/logex v1.1.10/go.mod h1:+Ywpsq7O8HXn0nuIou7OrIPyXbp3wmkHB+jjWRnGsAI= @@ -76,6 +77,10 @@ github.com/coreos/go-systemd/v22 v22.3.2/go.mod h1:Y58oyj3AT4RCenI/lSvhwexgC+NSV github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c= github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= +github.com/dgryski/go-rendezvous v0.0.0-20200823014737-9f7001d12a5f h1:lO4WD4F/rVNCu3HqELle0jiPLLBs70cWOduZpkS1E78= +github.com/dgryski/go-rendezvous v0.0.0-20200823014737-9f7001d12a5f/go.mod h1:cuUVRXasLTGF7a8hSLbxyZXjz+1KgoB3wDUb6vlszIc= +github.com/dlclark/regexp2 v1.10.0 h1:+/GIL799phkJqYW+3YbOd8LCcbHzT0Pbo8zl70MHsq0= +github.com/dlclark/regexp2 v1.10.0/go.mod h1:DHkYz0B9wPfa6wondMfaivmHpzrQ3v9q8cnmRbL6yW8= github.com/envoyproxy/go-control-plane v0.9.0/go.mod h1:YTl/9mNaCwkRvm6d1a2C3ymFceY/DCBVvsKhRF0iEA4= github.com/envoyproxy/go-control-plane v0.9.1-0.20191026205805-5f8ba28d4473/go.mod h1:YTl/9mNaCwkRvm6d1a2C3ymFceY/DCBVvsKhRF0iEA4= github.com/envoyproxy/go-control-plane v0.9.4/go.mod h1:6rpuAdCZL397s3pYoYcLgu1mIlRU8Am5FuJP05cCM98= @@ -86,6 +91,7 @@ github.com/envoyproxy/go-control-plane v0.9.9-0.20210512163311-63b5d3c536b0/go.m github.com/envoyproxy/protoc-gen-validate v0.1.0/go.mod h1:iSmxcyjqTsJpI2R4NaDN7+kN2VEUnK/pcBlmesArF7c= github.com/fatih/color v1.7.0/go.mod h1:Zm6kSWBoL9eyXnKyktHP6abPY2pDugNf5KwzbycvMj4= github.com/fatih/color v1.9.0/go.mod h1:eQcE1qtQxscV5RaZvpXrrb8Drkc3/DdQ+uUYCNjL+zU= +github.com/fsnotify/fsnotify v1.4.7/go.mod h1:jwhsz4b93w/PPRr/qN1Yymfu8t87LnFCMoQvtojpjFo= github.com/fsnotify/fsnotify v1.4.9/go.mod h1:znqG4EE+3YCdAaPaxE2ZRY/06pZUdp0tY4IgpuI1SZQ= github.com/fsnotify/fsnotify v1.5.1 h1:mZcQUHVQUQWoPXXtuf9yuEXKudkV2sx1E06UadKWpgI= github.com/fsnotify/fsnotify v1.5.1/go.mod h1:T3375wBYaZdLLcVNkcVbzGHY7f1l/uK5T5Ai1i3InKU= @@ -95,6 +101,8 @@ github.com/go-gl/glfw/v3.3/glfw v0.0.0-20191125211704-12ad95a8df72/go.mod h1:tQ2 github.com/go-gl/glfw/v3.3/glfw v0.0.0-20200222043503-6f7a984d4dc4/go.mod h1:tQ2UAYgL5IevRw8kRxooKSPJfGvJ9fJQFa0TUsXzTg8= github.com/go-ole/go-ole v1.2.4 h1:nNBDSCOigTSiarFpYE9J/KtEA1IOW4CNeqT9TQDqCxI= github.com/go-ole/go-ole v1.2.4/go.mod h1:XCwSNxSkXRo4vlyPy93sltvi/qJq0jqQhjqQNIwKuxM= +github.com/go-redis/redis/v8 v8.11.0 h1:O1Td0mQ8UFChQ3N9zFQqo6kTU2cJ+/it88gDB+zg0wo= +github.com/go-redis/redis/v8 v8.11.0/go.mod h1:DLomh7y2e3ggQXQLd1YgmvIfecPJoFl7WU5SOQ/r06M= github.com/godbus/dbus/v5 v5.0.4/go.mod h1:xhWf0FNVPg57R7Z0UbKHbJfkEywrmjJnf7w5xrFpKfA= github.com/gogo/protobuf v1.3.2/go.mod h1:P1XiOD3dCwIKUDQYPy72D8LYyHL2YPYrpS2s69NZV8Q= github.com/golang/glog v0.0.0-20160126235308-23def4e6c14b/go.mod h1:SBH7ygxi8pfUlaOkMMuAQtPIUF8ecWP5IEl/CR7VP2Q= @@ -165,8 +173,9 @@ github.com/google/pprof v0.0.0-20210601050228-01bbb1931b22/go.mod h1:kpwsk12EmLe github.com/google/pprof v0.0.0-20210609004039-a478d1d731e9/go.mod h1:kpwsk12EmLew5upagYY7GY0pfYCcupk39gWOCRROcvE= github.com/google/pprof v0.0.0-20210720184732-4bb14d4b1be1/go.mod h1:kpwsk12EmLew5upagYY7GY0pfYCcupk39gWOCRROcvE= github.com/google/renameio v0.1.0/go.mod h1:KWCgfxg9yswjAJkECMjeO8J8rahYeXnNhOm40UhjYkI= -github.com/google/uuid v1.1.2 h1:EVhdT+1Kseyi1/pUmXKaFxYsDNy9RQYkMWRH68J/W7Y= github.com/google/uuid v1.1.2/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo= +github.com/google/uuid v1.3.0 h1:t6JiXgmwXMjEs8VusXIJk2BXHsn+wx8BZdTaoZ5fu7I= +github.com/google/uuid v1.3.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo= github.com/googleapis/gax-go/v2 v2.0.4/go.mod h1:0Wqv26UfaUD9n4G6kQubkQ+KchISgw+vpHVxEJEs9eg= github.com/googleapis/gax-go/v2 v2.0.5/go.mod h1:DWXyrwAJ9X0FpwwEdw+IPEYBICEFu5mhpdKc/us6bOk= github.com/googleapis/gax-go/v2 v2.1.0/go.mod h1:Q3nei7sK6ybPYH7twZdmQpAd1MKb7pfu6SK+H1/DsU0= @@ -201,8 +210,11 @@ github.com/hashicorp/memberlist v0.1.3/go.mod h1:ajVTdAv/9Im8oMAAj5G31PhhMCZJV2p github.com/hashicorp/memberlist v0.2.2/go.mod h1:MS2lj3INKhZjWNqd3N0m3J+Jxf3DAOnAH9VT3Sh9MUE= github.com/hashicorp/serf v0.8.2/go.mod h1:6hOLApaqBFA1NXqRQAsxw9QxuDEvNxSQRwA/JwenrHc= github.com/hashicorp/serf v0.9.5/go.mod h1:UWDWwZeL5cuWDJdl0C6wrvrUwEqtQ4ZKBKKENpqIUyk= +github.com/hpcloud/tail v1.0.0/go.mod h1:ab1qPbhIpdTxEkNHXyeSf5vhxWSCs/tWer42PpOxQnU= github.com/ianlancetaylor/demangle v0.0.0-20181102032728-5e5cf60278f6/go.mod h1:aSSvb/t6k1mPoxDqO4vJh6VOCGPwU4O0C2/Eqndh1Sc= github.com/ianlancetaylor/demangle v0.0.0-20200824232613-28f6c0f3b639/go.mod h1:aSSvb/t6k1mPoxDqO4vJh6VOCGPwU4O0C2/Eqndh1Sc= +github.com/jinzhu/copier v0.4.0 h1:w3ciUoD19shMCRargcpm0cm91ytaBhDvuRpz1ODO/U8= +github.com/jinzhu/copier v0.4.0/go.mod h1:DfbEm0FYsaqBcKcFuvmOZb218JkPGtvSHsKg8S8hyyg= github.com/json-iterator/go v1.1.11/go.mod h1:KdQUCv79m/52Kvf8AW2vK1V8akMuk1QjK/uOdHXbAo4= github.com/jstemmer/go-junit-report v0.0.0-20190106144839-af01ea7f8024/go.mod h1:6v2b51hI/fHJwM22ozAgKL4VKDeJcHhJFhtBdhmNjmU= github.com/jstemmer/go-junit-report v0.9.1/go.mod h1:Brl9GWCQeLvo8nXZwPNNblvFj/XSXhF0NWZEnDohbsk= @@ -245,6 +257,16 @@ github.com/mitchellh/mapstructure v1.4.2/go.mod h1:bFUtVrKA4DC2yAKiSyO/QUcy7e+RR github.com/modern-go/concurrent v0.0.0-20180228061459-e0a39a4cb421/go.mod h1:6dJC0mAP4ikYIbvyc7fijjWJddQyLn8Ig3JB5CqoB9Q= github.com/modern-go/reflect2 v0.0.0-20180701023420-4b7aa43c6742/go.mod h1:bx2lNnkwVCuqBIxFjflWJWanXIb3RllmbCylyMrvgv0= github.com/modern-go/reflect2 v1.0.1/go.mod h1:bx2lNnkwVCuqBIxFjflWJWanXIb3RllmbCylyMrvgv0= +github.com/nxadm/tail v1.4.4 h1:DQuhQpB1tVlglWS2hLQ5OV6B5r8aGxSrPc5Qo6uTN78= +github.com/nxadm/tail v1.4.4/go.mod h1:kenIhsEOeOJmVchQTgglprH7qJGnHDVpk1VPCcaMI8A= +github.com/onsi/ginkgo v1.6.0/go.mod h1:lLunBs/Ym6LB5Z9jYTR76FiuTmxDTDusOGeTQH+WWjE= +github.com/onsi/ginkgo v1.12.1/go.mod h1:zj2OWP4+oCPe1qIXoGWkgMRwljMUYCdkwsT2108oapk= +github.com/onsi/ginkgo v1.15.0 h1:1V1NfVQR87RtWAgp1lv9JZJ5Jap+XFGKPi00andXGi4= +github.com/onsi/ginkgo v1.15.0/go.mod h1:hF8qUzuuC8DJGygJH3726JnCZX4MYbRB8yFfISqnKUg= +github.com/onsi/gomega v1.7.1/go.mod h1:XdKZgCCFLUoM/7CFJVPcG8C1xQ1AJ0vpAezJrB7JYyY= +github.com/onsi/gomega v1.10.1/go.mod h1:iN09h71vgCQne3DLsj+A5owkum+a2tYe+TOCB1ybHNo= +github.com/onsi/gomega v1.10.5 h1:7n6FEkpFmfCoo2t+YYqXH0evK+a9ICQz0xcAy9dYcaQ= +github.com/onsi/gomega v1.10.5/go.mod h1:gza4q3jKQJijlu05nKWRCW/GavJumGt8aNRxWg7mt48= github.com/pascaldekloe/goe v0.0.0-20180627143212-57f6aae5913c/go.mod h1:lzWF7FIEvWOWxwDKqyGYQf6ZUaNfKdP144TG7ZOy1lc= github.com/pelletier/go-toml v1.9.3/go.mod h1:u1nR/EPcESfeI/szUZKdtJ0xRNbUoANCkoOuaOx1Y+c= github.com/pelletier/go-toml v1.9.4 h1:tjENF6MfZAg8e4ZmZTeWaWiT2vXtsoO6+iuOjFhECwM= @@ -253,6 +275,8 @@ github.com/pkg/errors v0.8.1/go.mod h1:bwawxfHBFNV+L2hUp1rHADufV3IMtnDRdf1r5NINE github.com/pkg/errors v0.9.1 h1:FEBLx1zS214owpjy7qsBeixbURkuhQAwrK5UwLGTwt4= github.com/pkg/errors v0.9.1/go.mod h1:bwawxfHBFNV+L2hUp1rHADufV3IMtnDRdf1r5NINEl0= github.com/pkg/sftp v1.10.1/go.mod h1:lYOWFsE0bwd1+KfKJaKeuokY15vzFx25BLbzYYoAxZI= +github.com/pkoukk/tiktoken-go v0.1.7 h1:qOBHXX4PHtvIvmOtyg1EeKlwFRiMKAcoMp4Q+bLQDmw= +github.com/pkoukk/tiktoken-go v0.1.7/go.mod h1:9NiV+i9mJKGj1rYOT+njbv+ZwA/zJxYdewGl6qVatpg= github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM= github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4= github.com/posener/complete v1.1.1/go.mod h1:em0nMJCgc9GFtwrmVmEMR/ZL6WyhyjMBndrE9hABlRI= @@ -277,6 +301,8 @@ github.com/shirou/gopsutil/v3 v3.21.6/go.mod h1:JfVbDpIBLVzT8oKbvMg9P3wEIMDDpVn+ github.com/smartystreets/assertions v0.0.0-20180927180507-b2de0cb4f26d/go.mod h1:OnSkiWE9lh6wB0YB77sQom3nweQdgAjqCqsofrRNTgc= github.com/smartystreets/goconvey v1.6.4/go.mod h1:syvi0/a8iFYH4r/RixwvyeAJjdLS9QV7WQ/tjFTllLA= github.com/spaolacci/murmur3 v0.0.0-20180118202830-f09979ecbc72/go.mod h1:JwIasOWyU6f++ZhiEuf87xNszmSA2myDM2Kzu9HwQUA= +github.com/spaolacci/murmur3 v1.1.0 h1:7c1g84S4BPRrfL5Xrdp6fOJ206sU9y293DDHaoy0bLI= +github.com/spaolacci/murmur3 v1.1.0/go.mod h1:JwIasOWyU6f++ZhiEuf87xNszmSA2myDM2Kzu9HwQUA= github.com/spf13/afero v1.6.0 h1:xoax2sJ2DT8S8xA2paPFjDCScCNeWsg75VG0DLRreiY= github.com/spf13/afero v1.6.0/go.mod h1:Ai8FlHk4v/PARR026UzYexafAt9roJ7LcLMAmO6Z93I= github.com/spf13/cast v1.3.1/go.mod h1:Qx5cxh0v+4UWYiBimWS+eyWzqEqokIECu5etghLkUJE= @@ -290,8 +316,9 @@ github.com/spf13/viper v1.8.1/go.mod h1:o0Pch8wJ9BVSWGQMbra6iw0oQ5oktSIBaujf1rJH github.com/spf13/viper v1.9.0 h1:yR6EXjTp0y0cLN8OZg1CRZmOBdI88UcGkhgyJhu6nZk= github.com/spf13/viper v1.9.0/go.mod h1:+i6ajR7OX2XaiBkrcZJFK21htRk7eDeLg7+O6bhUPP4= github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME= -github.com/stretchr/objx v0.4.0 h1:M2gUjqZET1qApGOWNSnZ49BAIMX4F/1plDv3+l31EJ4= github.com/stretchr/objx v0.4.0/go.mod h1:YvHI0jy2hoMjB+UWwv71VJQ9isScKT/TqJzVSSt89Yw= +github.com/stretchr/objx v0.5.0 h1:1zr/of2m5FGMsad5YfcqgdqdWrIhu+EBEJRhR1U7z/c= +github.com/stretchr/objx v0.5.0/go.mod h1:Yh+to48EsGEfYuaHDzXPcE3xhTkx73EhmCGUpEOglKo= github.com/stretchr/testify v1.2.2/go.mod h1:a8OnRcib4nhh0OaRAV+Yts87kKdq0PP7pXfy6kDkUVs= github.com/stretchr/testify v1.3.0/go.mod h1:M5WIy9Dh21IEIfnGCwXGc5bZfKNJtfHm1UVUgZn+9EI= github.com/stretchr/testify v1.4.0/go.mod h1:j7eGeouHqKxXV5pUuKE4zz7dFj8WfuZ+81PSLYec5m4= @@ -299,8 +326,9 @@ github.com/stretchr/testify v1.5.1/go.mod h1:5W2xD1RspED5o8YsWQXVCued0rvSQ+mT+I5 github.com/stretchr/testify v1.6.1/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg= github.com/stretchr/testify v1.7.0/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg= github.com/stretchr/testify v1.7.1/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg= -github.com/stretchr/testify v1.8.0 h1:pSgiaMZlXftHpm5L7V1+rVB+AZJydKsMxsQBIJw4PKk= github.com/stretchr/testify v1.8.0/go.mod h1:yNjHg4UonilssWZ8iaSj1OCr/vHnekPRkoO+kdMU+MU= +github.com/stretchr/testify v1.8.2 h1:+h33VjcLVPDHtOdpUCuF+7gSuG3yGIftsP1YvFihtJ8= +github.com/stretchr/testify v1.8.2/go.mod h1:w2LPCIKwWwSfY2zedu0+kehJoqGctiVI29o6fzry7u4= github.com/subosito/gotenv v1.2.0 h1:Slr1R9HxAlEKefgq5jn9U+DnETlIUa6HfgEzj0g5d7s= github.com/subosito/gotenv v1.2.0/go.mod h1:N0PQaV/YGNqwC0u51sEeR/aUtSLEXKX9iv69rRypqCw= github.com/tevid/gohamcrest v1.1.1 h1:ou+xSqlIw1xfGTg1uq1nif/htZ2S3EzRqLm2BP+tYU0= @@ -376,6 +404,7 @@ golang.org/x/mod v0.4.1/go.mod h1:s0Qsj1ACt9ePp/hMypM3fl4fZqREWJwdYDEqhRiZZUA= golang.org/x/mod v0.4.2/go.mod h1:s0Qsj1ACt9ePp/hMypM3fl4fZqREWJwdYDEqhRiZZUA= golang.org/x/net v0.0.0-20180724234803-3673e40ba225/go.mod h1:mL1N/T3taQHkDXs73rZJwtUhF3w3ftmwwsq0BUmARs4= golang.org/x/net v0.0.0-20180826012351-8a410e7b638d/go.mod h1:mL1N/T3taQHkDXs73rZJwtUhF3w3ftmwwsq0BUmARs4= +golang.org/x/net v0.0.0-20180906233101-161cd47e91fd/go.mod h1:mL1N/T3taQHkDXs73rZJwtUhF3w3ftmwwsq0BUmARs4= golang.org/x/net v0.0.0-20181023162649-9b4f9f5ad519/go.mod h1:mL1N/T3taQHkDXs73rZJwtUhF3w3ftmwwsq0BUmARs4= golang.org/x/net v0.0.0-20181201002055-351d144fa1fc/go.mod h1:mL1N/T3taQHkDXs73rZJwtUhF3w3ftmwwsq0BUmARs4= golang.org/x/net v0.0.0-20190108225652-1e06a53dbb7e/go.mod h1:mL1N/T3taQHkDXs73rZJwtUhF3w3ftmwwsq0BUmARs4= @@ -399,6 +428,7 @@ golang.org/x/net v0.0.0-20200324143707-d3edc9973b7e/go.mod h1:qpuaurCH72eLCgpAm/ golang.org/x/net v0.0.0-20200501053045-e0ff5e5a1de5/go.mod h1:qpuaurCH72eLCgpAm/N6yyVIVM9cpaDIP3A8BGJEC5A= golang.org/x/net v0.0.0-20200506145744-7e3656a0809f/go.mod h1:qpuaurCH72eLCgpAm/N6yyVIVM9cpaDIP3A8BGJEC5A= golang.org/x/net v0.0.0-20200513185701-a91f0712d120/go.mod h1:qpuaurCH72eLCgpAm/N6yyVIVM9cpaDIP3A8BGJEC5A= +golang.org/x/net v0.0.0-20200520004742-59133d7f0dd7/go.mod h1:qpuaurCH72eLCgpAm/N6yyVIVM9cpaDIP3A8BGJEC5A= golang.org/x/net v0.0.0-20200520182314-0ba52f642ac2/go.mod h1:qpuaurCH72eLCgpAm/N6yyVIVM9cpaDIP3A8BGJEC5A= golang.org/x/net v0.0.0-20200625001655-4c5254603344/go.mod h1:/O7V0waA8r7cgGh81Ro3o1hOxt32SMVPicZroKQ2sZA= golang.org/x/net v0.0.0-20200707034311-ab3426394381/go.mod h1:/O7V0waA8r7cgGh81Ro3o1hOxt32SMVPicZroKQ2sZA= @@ -406,12 +436,14 @@ golang.org/x/net v0.0.0-20200822124328-c89045814202/go.mod h1:/O7V0waA8r7cgGh81R golang.org/x/net v0.0.0-20201021035429-f5854403a974/go.mod h1:sp8m0HH+o8qH0wwXwYZr8TS3Oi6o0r6Gce1SSxlDquU= golang.org/x/net v0.0.0-20201031054903-ff519b6c9102/go.mod h1:sp8m0HH+o8qH0wwXwYZr8TS3Oi6o0r6Gce1SSxlDquU= golang.org/x/net v0.0.0-20201110031124-69a78807bb2b/go.mod h1:sp8m0HH+o8qH0wwXwYZr8TS3Oi6o0r6Gce1SSxlDquU= +golang.org/x/net v0.0.0-20201202161906-c7110b5ffcbb/go.mod h1:sp8m0HH+o8qH0wwXwYZr8TS3Oi6o0r6Gce1SSxlDquU= golang.org/x/net v0.0.0-20201209123823-ac852fbbde11/go.mod h1:m0MpNAwzfU5UDzcl9v0D8zg8gWTRqZa9RBIspLL5mdg= golang.org/x/net v0.0.0-20210119194325-5f4716e94777/go.mod h1:m0MpNAwzfU5UDzcl9v0D8zg8gWTRqZa9RBIspLL5mdg= golang.org/x/net v0.0.0-20210226172049-e18ecbb05110/go.mod h1:m0MpNAwzfU5UDzcl9v0D8zg8gWTRqZa9RBIspLL5mdg= golang.org/x/net v0.0.0-20210316092652-d523dce5a7f4/go.mod h1:RBQZq4jEuRlivfhVLdyRGr576XBO4/greRjx4P4O3yc= golang.org/x/net v0.0.0-20210405180319-a5a99cb37ef4/go.mod h1:p54w0d4576C0XHj96bSt6lcn1PtDYWL6XObtHCRCNQM= golang.org/x/net v0.0.0-20210503060351-7fd8e65b6420/go.mod h1:9nx3DQGgdP8bBQD5qxJ1jj9UTztislL4KSBs9R2vV5Y= +golang.org/x/net v0.7.0 h1:rJrUqqhjsgNp7KqAIc25s9pZnjU7TUcSY7HcVZjdn1g= golang.org/x/oauth2 v0.0.0-20180821212333-d2e6202438be/go.mod h1:N/0e6XlmueqKjAGxoOufVs8QHGRruUQn6yWY3a++T0U= golang.org/x/oauth2 v0.0.0-20190226205417-e64efc72b421/go.mod h1:gOpvHmFTYa4IltrdGE7lF6nIHvwfUNPOp7c8zoXwtLw= golang.org/x/oauth2 v0.0.0-20190604053449-0f29369cfe45/go.mod h1:gOpvHmFTYa4IltrdGE7lF6nIHvwfUNPOp7c8zoXwtLw= @@ -441,6 +473,7 @@ golang.org/x/sync v0.0.0-20201207232520-09787c993a3a/go.mod h1:RxMgew5VJxzue5/jJ golang.org/x/sync v0.0.0-20210220032951-036812b2e83c/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM= golang.org/x/sys v0.0.0-20180823144017-11551d06cbcc/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY= golang.org/x/sys v0.0.0-20180830151530-49385e6e1522/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY= +golang.org/x/sys v0.0.0-20180909124046-d0be0721c37e/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY= golang.org/x/sys v0.0.0-20181026203630-95b1ffbd15a5/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY= golang.org/x/sys v0.0.0-20190215142949-d0b11bdaac8a/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY= golang.org/x/sys v0.0.0-20190222072716-a9d3bda3a223/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY= @@ -451,12 +484,14 @@ golang.org/x/sys v0.0.0-20190507160741-ecd444e8653b/go.mod h1:h1NjWce9XRLGQEsW7w golang.org/x/sys v0.0.0-20190606165138-5da285871e9c/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs= golang.org/x/sys v0.0.0-20190624142023-c5567b49c5d0/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs= golang.org/x/sys v0.0.0-20190726091711-fc99dfbffb4e/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs= +golang.org/x/sys v0.0.0-20190904154756-749cb33beabd/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs= golang.org/x/sys v0.0.0-20190922100055-0a153f010e69/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs= golang.org/x/sys v0.0.0-20190924154521-2837fb4f24fe/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs= golang.org/x/sys v0.0.0-20191001151750-bb3f8db39f24/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs= golang.org/x/sys v0.0.0-20191005200804-aed5e4c7ecf9/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs= golang.org/x/sys v0.0.0-20191008105621-543471e840be/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs= golang.org/x/sys v0.0.0-20191026070338-33540a1f6037/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs= +golang.org/x/sys v0.0.0-20191120155948-bd437916bb0e/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs= golang.org/x/sys v0.0.0-20191204072324-ce4227a45e2e/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs= golang.org/x/sys v0.0.0-20191228213918-04cbcbbfeed8/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs= golang.org/x/sys v0.0.0-20200113162924-86b910548bc1/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs= @@ -479,6 +514,7 @@ golang.org/x/sys v0.0.0-20200930185726-fdedc70b468f/go.mod h1:h1NjWce9XRLGQEsW7w golang.org/x/sys v0.0.0-20201119102817-f84b799fce68/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs= golang.org/x/sys v0.0.0-20201201145000-ef89a241ccb3/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs= golang.org/x/sys v0.0.0-20210104204734-6f8348627aad/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs= +golang.org/x/sys v0.0.0-20210112080510-489259a85091/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs= golang.org/x/sys v0.0.0-20210119212857-b64e53b001e4/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs= golang.org/x/sys v0.0.0-20210220050731-9a76102bfb43/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs= golang.org/x/sys v0.0.0-20210305230114-8fe3ee5dd75b/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs= @@ -560,6 +596,7 @@ golang.org/x/tools v0.0.0-20200904185747-39188db58858/go.mod h1:Cj7w3i3Rnn0Xh82u golang.org/x/tools v0.0.0-20201110124207-079ba7bd75cd/go.mod h1:emZCQorbCU4vsT4fOWvOPXz4eW1wZW4PmDk9uLelYpA= golang.org/x/tools v0.0.0-20201201161351-ac6f37ff4c2a/go.mod h1:emZCQorbCU4vsT4fOWvOPXz4eW1wZW4PmDk9uLelYpA= golang.org/x/tools v0.0.0-20201208233053-a543418bbed2/go.mod h1:emZCQorbCU4vsT4fOWvOPXz4eW1wZW4PmDk9uLelYpA= +golang.org/x/tools v0.0.0-20201224043029-2b0845dc783e/go.mod h1:emZCQorbCU4vsT4fOWvOPXz4eW1wZW4PmDk9uLelYpA= golang.org/x/tools v0.0.0-20210105154028-b0ab187a4818/go.mod h1:emZCQorbCU4vsT4fOWvOPXz4eW1wZW4PmDk9uLelYpA= golang.org/x/tools v0.0.0-20210106214847-113979e3529a/go.mod h1:emZCQorbCU4vsT4fOWvOPXz4eW1wZW4PmDk9uLelYpA= golang.org/x/tools v0.1.0/go.mod h1:xkSsbof2nBLbhDlRMhhhyNLN/zl3eTqcnHD5viDpcZ0= @@ -705,13 +742,18 @@ gopkg.in/check.v1 v1.0.0-20180628173108-788fd7840127/go.mod h1:Co6ibVJAznAaIkqp8 gopkg.in/check.v1 v1.0.0-20190902080502-41f04d3bba15/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0= gopkg.in/check.v1 v1.0.0-20201130134442-10cb98267c6c h1:Hei/4ADfdWqJk1ZMxUNpqntNwaWcugrBjAiHlqqRiVk= gopkg.in/errgo.v2 v2.1.0/go.mod h1:hNsd1EY+bozCKY1Ytp96fpM3vjJbqLJn88ws8XvfDNI= +gopkg.in/fsnotify.v1 v1.4.7/go.mod h1:Tz8NjZHkW78fSQdbUxIjBTcgA1z1m8ZHf0WmKUhAMys= gopkg.in/ini.v1 v1.62.0/go.mod h1:pNLf8WUiyNEtQjuu5G5vTm06TEv9tsIgeAvK8hOrP4k= gopkg.in/ini.v1 v1.63.2/go.mod h1:pNLf8WUiyNEtQjuu5G5vTm06TEv9tsIgeAvK8hOrP4k= gopkg.in/ini.v1 v1.64.0 h1:Mj2zXEXcNb5joEiSA0zc3HZpTst/iyjNiR4CN8tDzOg= gopkg.in/ini.v1 v1.64.0/go.mod h1:pNLf8WUiyNEtQjuu5G5vTm06TEv9tsIgeAvK8hOrP4k= +gopkg.in/tomb.v1 v1.0.0-20141024135613-dd632973f1e7 h1:uRGJdciOHaEIrze2W8Q3AKkepLTh2hOroT7a+7czfdQ= +gopkg.in/tomb.v1 v1.0.0-20141024135613-dd632973f1e7/go.mod h1:dt/ZhP58zS4L8KSrWDmTeBkI65Dw0HsyUHuEVlX15mw= gopkg.in/yaml.v2 v2.2.2/go.mod h1:hI93XBmqTisBFMUTm0b8Fm+jr3Dg1NNxqwp+5A1VGuI= gopkg.in/yaml.v2 v2.2.3/go.mod h1:hI93XBmqTisBFMUTm0b8Fm+jr3Dg1NNxqwp+5A1VGuI= +gopkg.in/yaml.v2 v2.2.4/go.mod h1:hI93XBmqTisBFMUTm0b8Fm+jr3Dg1NNxqwp+5A1VGuI= gopkg.in/yaml.v2 v2.2.8/go.mod h1:hI93XBmqTisBFMUTm0b8Fm+jr3Dg1NNxqwp+5A1VGuI= +gopkg.in/yaml.v2 v2.3.0/go.mod h1:hI93XBmqTisBFMUTm0b8Fm+jr3Dg1NNxqwp+5A1VGuI= gopkg.in/yaml.v2 v2.4.0 h1:D8xgwECY7CYvx+Y2n4sBz93Jn9JRvxdiyyo8CTfuKaY= gopkg.in/yaml.v2 v2.4.0/go.mod h1:RDklbk79AGWmwhnvt/jBztapEOGDOx6ZbXqjP6csGnQ= gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM= diff --git a/tests/benchmark/llm_token_ratelimit/Dockerfile b/tests/benchmark/llm_token_ratelimit/Dockerfile new file mode 100644 index 000000000..76b9d68ce --- /dev/null +++ b/tests/benchmark/llm_token_ratelimit/Dockerfile @@ -0,0 +1,12 @@ +FROM ubuntu:22.04 + + +RUN sed -i 's/archive.ubuntu.com/mirrors.aliyun.com/g' /etc/apt/sources.list \ + && sed -i 's/security.ubuntu.com/mirrors.aliyun.com/g' /etc/apt/sources.list + +RUN apt-get update && \ + apt-get install -y vim net-tools ca-certificates less redis-tools && \ + apt-get clean && \ + rm -rf /var/lib/apt/lists/* + +WORKDIR /workspace diff --git a/tests/benchmark/llm_token_ratelimit/build-code.sh b/tests/benchmark/llm_token_ratelimit/build-code.sh new file mode 100755 index 000000000..4d70bccd3 --- /dev/null +++ b/tests/benchmark/llm_token_ratelimit/build-code.sh @@ -0,0 +1,22 @@ +#!/bin/zsh + +supported_modes=("single" "cluster") + +mode=$1 +if [[ ! " ${supported_modes[@]} " =~ " ${mode} " ]]; then + echo "Usage: $0 [single|cluster]" + exit 1 +fi + +if [[ ! -e "sentinel-$mode.yml" ]]; then + echo "sentinel-$mode.yml not found!" + exit 1 +fi + + +mkdir -p ./out + +# 编译go代码 +go build -o out/llm_token_ratelimit_benchmark main.go +# 复制 sentinel 配置文件到输出目录 +cp sentinel-$mode.yml out/sentinel.yml \ No newline at end of file diff --git a/tests/benchmark/llm_token_ratelimit/build-container.sh b/tests/benchmark/llm_token_ratelimit/build-container.sh new file mode 100755 index 000000000..43ed8f11f --- /dev/null +++ b/tests/benchmark/llm_token_ratelimit/build-container.sh @@ -0,0 +1,48 @@ +#!/bin/zsh + +supported_modes=("single" "cluster") + +mode=$1 +if [[ ! " ${supported_modes[@]} " =~ " ${mode} " ]]; then + echo "Usage: $0 [single|cluster]" + exit 1 +fi + +if [[ ! -e "docker-compose-$mode.yml" ]]; then + echo "docker-compose-$mode.yml not found!" + exit 1 +fi + +echo "Building application..." +./build-code.sh $mode + +echo "Building and starting containers..." +docker-compose -f docker-compose-$mode.yml down +docker-compose -f docker-compose-$mode.yml up --build -d + +echo "Waiting for Redis nodes to start..." +sleep 10 + +if [ "$mode" = "cluster" ]; then + echo "Checking Redis nodes status..." + for port in 7001 7002 7003; do + echo "Checking Redis on port $port..." + timeout 5 redis-cli -h 127.0.0.1 -p $port ping || echo "Redis on port $port not ready" + done + + echo "Initializing Redis cluster..." + # 使用容器内部执行集群初始化 + docker exec redis-node-1 redis-cli --cluster create \ + 172.20.0.11:6379 \ + 172.20.0.12:6379 \ + 172.20.0.13:6379 \ + --cluster-replicas 0 \ + --cluster-yes + + echo "Checking cluster status..." + docker exec redis-node-1 redis-cli cluster nodes +else + echo "Single node mode, no cluster initialization needed." +fi + +echo "Containers setup complete!" \ No newline at end of file diff --git a/tests/benchmark/llm_token_ratelimit/docker-compose-cluster.yml b/tests/benchmark/llm_token_ratelimit/docker-compose-cluster.yml new file mode 100644 index 000000000..5772fc7ab --- /dev/null +++ b/tests/benchmark/llm_token_ratelimit/docker-compose-cluster.yml @@ -0,0 +1,76 @@ +services: + # Redis 6.0 集群 + redis-node-1: + image: redis:6.2.7 + container_name: redis-node-1 + ports: + - "7001:6379" + - "17001:16379" + volumes: + - ./redis-cluster.conf:/usr/local/etc/redis/redis.conf + command: redis-server /usr/local/etc/redis/redis.conf + networks: + test-network: + ipv4_address: 172.20.0.11 + + redis-node-2: + image: redis:6.2.7 + container_name: redis-node-2 + ports: + - "7002:6379" + - "17002:16379" + volumes: + - ./redis-cluster.conf:/usr/local/etc/redis/redis.conf + command: redis-server /usr/local/etc/redis/redis.conf + networks: + test-network: + ipv4_address: 172.20.0.12 + + redis-node-3: + image: redis:6.2.7 + container_name: redis-node-3 + ports: + - "7003:6379" + - "17003:16379" + volumes: + - ./redis-cluster.conf:/usr/local/etc/redis/redis.conf + command: redis-server /usr/local/etc/redis/redis.conf + networks: + test-network: + ipv4_address: 172.20.0.13 + + # Sentinel Go LLM Token RateLimit 服务容器 + sentinel-go-llm-token-ratelimit: + build: + context: . + dockerfile: Dockerfile + container_name: sentinel-go-llm-token-ratelimit + command: ["tail", "-f", "/dev/null"] + ports: + - "9527:9527" + volumes: + - ./out:/workspace + networks: + - test-network + depends_on: + - redis-node-1 + - redis-node-2 + - redis-node-3 + environment: + - LLM_API_KEY=${LLM_API_KEY} + - LLM_BASE_URL=${LLM_BASE_URL} + - LLM_MODEL=${LLM_MODEL} + deploy: + resources: + limits: + cpus: '2' + memory: 4G + restart: on-failure + +# 定义网络和卷 +networks: + test-network: + driver: bridge + ipam: + config: + - subnet: 172.20.0.0/16 \ No newline at end of file diff --git a/tests/benchmark/llm_token_ratelimit/docker-compose-single.yml b/tests/benchmark/llm_token_ratelimit/docker-compose-single.yml new file mode 100644 index 000000000..9ed767f89 --- /dev/null +++ b/tests/benchmark/llm_token_ratelimit/docker-compose-single.yml @@ -0,0 +1,43 @@ +services: + # Redis 6.0 单节点 + redis-node-1: + image: redis:6.2.7 + container_name: redis-node-1 + ports: + - "7001:6379" + volumes: + - ./redis-single.conf:/usr/local/etc/redis/redis.conf + command: redis-server /usr/local/etc/redis/redis.conf + networks: + - test-network + + # Sentinel Go LLM Token RateLimit 服务容器 + sentinel-go-llm-token-ratelimit: + build: + context: . + dockerfile: Dockerfile + container_name: sentinel-go-llm-token-ratelimit + command: ["tail", "-f", "/dev/null"] + ports: + - "9527:9527" + volumes: + - ./out:/workspace + networks: + - test-network + depends_on: + - redis-node-1 + environment: + - LLM_API_KEY=${LLM_API_KEY} + - LLM_BASE_URL=${LLM_BASE_URL} + - LLM_MODEL=${LLM_MODEL} + deploy: + resources: + limits: + cpus: '2' + memory: 4G + restart: on-failure + +# 定义网络和卷 +networks: + test-network: + driver: bridge \ No newline at end of file diff --git a/tests/benchmark/llm_token_ratelimit/go.mod b/tests/benchmark/llm_token_ratelimit/go.mod new file mode 100644 index 000000000..49fec4c8f --- /dev/null +++ b/tests/benchmark/llm_token_ratelimit/go.mod @@ -0,0 +1,88 @@ +module llm_token_ratelimit + +go 1.22.0 + +replace github.com/alibaba/sentinel-golang => ../../../ + +require ( + github.com/alibaba/sentinel-golang v1.0.4 + github.com/cloudwego/eino v0.4.7 + github.com/cloudwego/eino-ext/components/model/openai v0.0.0-20250904121005-ad78ed3e5e49 + github.com/gin-gonic/gin v1.10.1 + github.com/tmc/langchaingo v0.1.13 +) + +require ( + github.com/bahlo/generic-list-go v0.2.0 // indirect + github.com/beorn7/perks v1.0.1 // indirect + github.com/buger/jsonparser v1.1.1 // indirect + github.com/bytedance/sonic v1.14.0 // indirect + github.com/bytedance/sonic/loader v0.3.0 // indirect + github.com/cespare/xxhash/v2 v2.2.0 // indirect + github.com/cloudwego/base64x v0.1.5 // indirect + github.com/cloudwego/eino-ext/libs/acl/openai v0.0.0-20250826113018-8c6f6358d4bb // indirect + github.com/dgryski/go-rendezvous v0.0.0-20200823014737-9f7001d12a5f // indirect + github.com/dlclark/regexp2 v1.10.0 // indirect + github.com/dustin/go-humanize v1.0.1 // indirect + github.com/eino-contrib/jsonschema v1.0.0 // indirect + github.com/evanphx/json-patch v0.5.2 // indirect + github.com/gabriel-vasile/mimetype v1.4.3 // indirect + github.com/getkin/kin-openapi v0.118.0 // indirect + github.com/gin-contrib/sse v0.1.0 // indirect + github.com/go-ole/go-ole v1.2.6 // indirect + github.com/go-openapi/jsonpointer v0.19.6 // indirect + github.com/go-openapi/swag v0.22.4 // indirect + github.com/go-playground/locales v0.14.1 // indirect + github.com/go-playground/universal-translator v0.18.1 // indirect + github.com/go-playground/validator/v10 v10.20.0 // indirect + github.com/go-redis/redis/v8 v8.11.0 // indirect + github.com/goccy/go-json v0.10.2 // indirect + github.com/golang/protobuf v1.5.4 // indirect + github.com/google/uuid v1.6.0 // indirect + github.com/goph/emperror v0.17.2 // indirect + github.com/invopop/yaml v0.1.0 // indirect + github.com/jinzhu/copier v0.4.0 // indirect + github.com/josharian/intern v1.0.0 // indirect + github.com/json-iterator/go v1.1.12 // indirect + github.com/klauspost/cpuid/v2 v2.2.7 // indirect + github.com/leodido/go-urn v1.4.0 // indirect + github.com/lufia/plan9stats v0.0.0-20211012122336-39d0f177ccd0 // indirect + github.com/mailru/easyjson v0.7.7 // indirect + github.com/mattn/go-isatty v0.0.20 // indirect + github.com/matttproud/golang_protobuf_extensions v1.0.4 // indirect + github.com/meguminnnnnnnnn/go-openai v0.0.0-20250821095446-07791bea23a0 // indirect + github.com/modern-go/concurrent v0.0.0-20180306012644-bacd9c7ef1dd // indirect + github.com/modern-go/reflect2 v1.0.2 // indirect + github.com/mohae/deepcopy v0.0.0-20170929034955-c48cc78d4826 // indirect + github.com/nikolalohinski/gonja v1.5.3 // indirect + github.com/pelletier/go-toml/v2 v2.2.2 // indirect + github.com/perimeterx/marshmallow v1.1.4 // indirect + github.com/pkg/errors v0.9.1 // indirect + github.com/pkoukk/tiktoken-go v0.1.7 // indirect + github.com/power-devops/perfstat v0.0.0-20210106213030-5aafc221ea8c // indirect + github.com/prometheus/client_golang v1.16.0 // indirect + github.com/prometheus/client_model v0.3.0 // indirect + github.com/prometheus/common v0.42.0 // indirect + github.com/prometheus/procfs v0.10.1 // indirect + github.com/shirou/gopsutil/v3 v3.23.12 // indirect + github.com/shoenig/go-m1cpu v0.1.6 // indirect + github.com/sirupsen/logrus v1.9.3 // indirect + github.com/slongfield/pyfmt v0.0.0-20220222012616-ea85ff4c361f // indirect + github.com/spaolacci/murmur3 v1.1.0 // indirect + github.com/tklauser/go-sysconf v0.3.12 // indirect + github.com/tklauser/numcpus v0.6.1 // indirect + github.com/twitchyliquid64/golang-asm v0.15.1 // indirect + github.com/ugorji/go/codec v1.2.12 // indirect + github.com/wk8/go-ordered-map/v2 v2.1.8 // indirect + github.com/yargevad/filepathx v1.0.0 // indirect + github.com/yusufpapurcu/wmi v1.2.3 // indirect + golang.org/x/arch v0.11.0 // indirect + golang.org/x/crypto v0.31.0 // indirect + golang.org/x/exp v0.0.0-20230713183714-613f0c0eb8a1 // indirect + golang.org/x/net v0.25.0 // indirect + golang.org/x/sys v0.28.0 // indirect + golang.org/x/text v0.21.0 // indirect + google.golang.org/protobuf v1.34.1 // indirect + gopkg.in/yaml.v2 v2.4.0 // indirect + gopkg.in/yaml.v3 v3.0.1 // indirect +) diff --git a/tests/benchmark/llm_token_ratelimit/go.sum b/tests/benchmark/llm_token_ratelimit/go.sum new file mode 100644 index 000000000..dad035217 --- /dev/null +++ b/tests/benchmark/llm_token_ratelimit/go.sum @@ -0,0 +1,337 @@ +github.com/airbrake/gobrake v3.6.1+incompatible/go.mod h1:wM4gu3Cn0W0K7GUuVWnlXZU11AGBXMILnrdOU8Kn00o= +github.com/bahlo/generic-list-go v0.2.0 h1:5sz/EEAK+ls5wF+NeqDpk5+iNdMDXrh3z3nPnH1Wvgk= +github.com/bahlo/generic-list-go v0.2.0/go.mod h1:2KvAjgMlE5NNynlg/5iLrrCCZ2+5xWbdbCW3pNTGyYg= +github.com/beorn7/perks v1.0.1 h1:VlbKKnNfV8bJzeqoa4cOKqO6bYr3WgKZxO8Z16+hsOM= +github.com/beorn7/perks v1.0.1/go.mod h1:G2ZrVWU2WbWT9wwq4/hrbKbnv/1ERSJQ0ibhJ6rlkpw= +github.com/bitly/go-simplejson v0.5.0/go.mod h1:cXHtHw4XUPsvGaxgjIAn8PhEWG9NfngEKAMDJEczWVA= +github.com/bmizerany/assert v0.0.0-20160611221934-b7ed37b82869/go.mod h1:Ekp36dRnpXw/yCqJaO+ZrUyxD+3VXMFFr56k5XYrpB4= +github.com/buger/jsonparser v1.1.1 h1:2PnMjfWD7wBILjqQbt530v576A/cAbQvEW9gGIpYMUs= +github.com/buger/jsonparser v1.1.1/go.mod h1:6RYKKt7H4d4+iWqouImQ9R2FZql3VbhNgx27UK13J/0= +github.com/bugsnag/bugsnag-go v1.4.0/go.mod h1:2oa8nejYd4cQ/b0hMIopN0lCRxU0bueqREvZLWFrtK8= +github.com/bugsnag/panicwrap v1.2.0/go.mod h1:D/8v3kj0zr8ZAKg1AQ6crr+5VwKN5eIywRkfhyM/+dE= +github.com/bytedance/mockey v1.2.14 h1:KZaFgPdiUwW+jOWFieo3Lr7INM1P+6adO3hxZhDswY8= +github.com/bytedance/mockey v1.2.14/go.mod h1:1BPHF9sol5R1ud/+0VEHGQq/+i2lN+GTsr3O2Q9IENY= +github.com/bytedance/sonic v1.14.0 h1:/OfKt8HFw0kh2rj8N0F6C/qPGRESq0BbaNZgcNXXzQQ= +github.com/bytedance/sonic v1.14.0/go.mod h1:WoEbx8WTcFJfzCe0hbmyTGrfjt8PzNEBdxlNUO24NhA= +github.com/bytedance/sonic/loader v0.1.1/go.mod h1:ncP89zfokxS5LZrJxl5z0UJcsk4M4yY2JpfqGeCtNLU= +github.com/bytedance/sonic/loader v0.3.0 h1:dskwH8edlzNMctoruo8FPTJDF3vLtDT0sXZwvZJyqeA= +github.com/bytedance/sonic/loader v0.3.0/go.mod h1:N8A3vUdtUebEY2/VQC0MyhYeKUFosQU6FxH2JmUe6VI= +github.com/certifi/gocertifi v0.0.0-20190105021004-abcd57078448/go.mod h1:GJKEexRPVJrBSOjoqN5VNOIKJ5Q3RViH6eu3puDRwx4= +github.com/cespare/xxhash/v2 v2.1.1/go.mod h1:VGX0DQ3Q6kWi7AoAeZDth3/j3BFtOZR5XLFGgcrjCOs= +github.com/cespare/xxhash/v2 v2.2.0 h1:DC2CZ1Ep5Y4k3ZQ899DldepgrayRUGE6BBZ/cd9Cj44= +github.com/cespare/xxhash/v2 v2.2.0/go.mod h1:VGX0DQ3Q6kWi7AoAeZDth3/j3BFtOZR5XLFGgcrjCOs= +github.com/cloudwego/base64x v0.1.5 h1:XPciSp1xaq2VCSt6lF0phncD4koWyULpl5bUxbfCyP4= +github.com/cloudwego/base64x v0.1.5/go.mod h1:0zlkT4Wn5C6NdauXdJRhSKRlJvmclQ1hhJgA0rcu/8w= +github.com/cloudwego/eino v0.4.7 h1:wwqsFWCuzCQuhw1dYKqHjGWULzjDjFfN9sTn/cezYV4= +github.com/cloudwego/eino v0.4.7/go.mod h1:1TDlOmwGSsbCJaWB92w9YLZi2FL0WRZoRcD4eMvqikg= +github.com/cloudwego/eino-ext/components/model/openai v0.0.0-20250904121005-ad78ed3e5e49 h1:Mb7GWb/YkjdyJz/b6XxsN5Wlgq1SHjngFGI47NAZa6Y= +github.com/cloudwego/eino-ext/components/model/openai v0.0.0-20250904121005-ad78ed3e5e49/go.mod h1:QQhCuQxuBAVWvu/YAZBhs/RsR76mUigw59Tl0kh04C8= +github.com/cloudwego/eino-ext/libs/acl/openai v0.0.0-20250826113018-8c6f6358d4bb h1:RMslzyijc3bi9EkqCulpS0hZupTl1y/wayR3+fVRN/c= +github.com/cloudwego/eino-ext/libs/acl/openai v0.0.0-20250826113018-8c6f6358d4bb/go.mod h1:fHn/6OqPPY1iLLx9wzz+MEVT5Dl9gwuZte1oLEnCoYw= +github.com/cloudwego/iasm v0.2.0/go.mod h1:8rXZaNYT2n95jn+zTI1sDr+IgcD2GVs0nlbbQPiEFhY= +github.com/creack/pty v1.1.9/go.mod h1:oKZEueFk5CKHvIhNR5MUki03XCEU+Q6VDXinZuGJ33E= +github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= +github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c= +github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= +github.com/dgryski/go-rendezvous v0.0.0-20200823014737-9f7001d12a5f h1:lO4WD4F/rVNCu3HqELle0jiPLLBs70cWOduZpkS1E78= +github.com/dgryski/go-rendezvous v0.0.0-20200823014737-9f7001d12a5f/go.mod h1:cuUVRXasLTGF7a8hSLbxyZXjz+1KgoB3wDUb6vlszIc= +github.com/dlclark/regexp2 v1.10.0 h1:+/GIL799phkJqYW+3YbOd8LCcbHzT0Pbo8zl70MHsq0= +github.com/dlclark/regexp2 v1.10.0/go.mod h1:DHkYz0B9wPfa6wondMfaivmHpzrQ3v9q8cnmRbL6yW8= +github.com/dustin/go-humanize v1.0.1 h1:GzkhY7T5VNhEkwH0PVJgjz+fX1rhBrR7pRT3mDkpeCY= +github.com/dustin/go-humanize v1.0.1/go.mod h1:Mu1zIs6XwVuF/gI1OepvI0qD18qycQx+mFykh5fBlto= +github.com/eino-contrib/jsonschema v1.0.0 h1:dXxbhGNZuI3+xNi8x3JT8AGyoXz6Pff6mRvmpjVl5Ww= +github.com/eino-contrib/jsonschema v1.0.0/go.mod h1:cpnX4SyKjWjGC7iN2EbhxaTdLqGjCi0e9DxpLYxddD4= +github.com/evanphx/json-patch v0.5.2 h1:xVCHIVMUu1wtM/VkR9jVZ45N3FhZfYMMYGorLCR8P3k= +github.com/evanphx/json-patch v0.5.2/go.mod h1:ZWS5hhDbVDyob71nXKNL0+PWn6ToqBHMikGIFbs31qQ= +github.com/fsnotify/fsnotify v1.4.7/go.mod h1:jwhsz4b93w/PPRr/qN1Yymfu8t87LnFCMoQvtojpjFo= +github.com/fsnotify/fsnotify v1.4.9 h1:hsms1Qyu0jgnwNXIxa+/V/PDsU6CfLf6CNO8H7IWoS4= +github.com/fsnotify/fsnotify v1.4.9/go.mod h1:znqG4EE+3YCdAaPaxE2ZRY/06pZUdp0tY4IgpuI1SZQ= +github.com/gabriel-vasile/mimetype v1.4.3 h1:in2uUcidCuFcDKtdcBxlR0rJ1+fsokWf+uqxgUFjbI0= +github.com/gabriel-vasile/mimetype v1.4.3/go.mod h1:d8uq/6HKRL6CGdk+aubisF/M5GcPfT7nKyLpA0lbSSk= +github.com/getkin/kin-openapi v0.118.0 h1:z43njxPmJ7TaPpMSCQb7PN0dEYno4tyBPQcrFdHoLuM= +github.com/getkin/kin-openapi v0.118.0/go.mod h1:l5e9PaFUo9fyLJCPGQeXI2ML8c3P8BHOEV2VaAVf/pc= +github.com/getsentry/raven-go v0.2.0/go.mod h1:KungGk8q33+aIAZUIVWZDr2OfAEBsO49PX4NzFV5kcQ= +github.com/gin-contrib/sse v0.1.0 h1:Y/yl/+YNO8GZSjAhjMsSuLt29uWRFHdHYUb5lYOV9qE= +github.com/gin-contrib/sse v0.1.0/go.mod h1:RHrZQHXnP2xjPF+u1gW/2HnVO7nvIa9PG3Gm+fLHvGI= +github.com/gin-gonic/gin v1.10.1 h1:T0ujvqyCSqRopADpgPgiTT63DUQVSfojyME59Ei63pQ= +github.com/gin-gonic/gin v1.10.1/go.mod h1:4PMNQiOhvDRa013RKVbsiNwoyezlm2rm0uX/T7kzp5Y= +github.com/go-check/check v0.0.0-20180628173108-788fd7840127 h1:0gkP6mzaMqkmpcJYCFOLkIBwI7xFExG03bbkOkCvUPI= +github.com/go-check/check v0.0.0-20180628173108-788fd7840127/go.mod h1:9ES+weclKsC9YodN5RgxqK/VD9HM9JsCSh7rNhMZE98= +github.com/go-ole/go-ole v1.2.6 h1:/Fpf6oFPoeFik9ty7siob0G6Ke8QvQEuVcuChpwXzpY= +github.com/go-ole/go-ole v1.2.6/go.mod h1:pprOEPIfldk/42T2oK7lQ4v4JSDwmV0As9GaiUsvbm0= +github.com/go-openapi/jsonpointer v0.19.5/go.mod h1:Pl9vOtqEWErmShwVjC8pYs9cog34VGT37dQOVbmoatg= +github.com/go-openapi/jsonpointer v0.19.6 h1:eCs3fxoIi3Wh6vtgmLTOjdhSpiqphQ+DaPn38N2ZdrE= +github.com/go-openapi/jsonpointer v0.19.6/go.mod h1:osyAmYz/mB/C3I+WsTTSgw1ONzaLJoLCyoi6/zppojs= +github.com/go-openapi/swag v0.19.5/go.mod h1:POnQmlKehdgb5mhVOsnJFsivZCEZ/vjK9gh66Z9tfKk= +github.com/go-openapi/swag v0.22.3/go.mod h1:UzaqsxGiab7freDnrUUra0MwWfN/q7tE4j+VcZ0yl14= +github.com/go-openapi/swag v0.22.4 h1:QLMzNJnMGPRNDCbySlcj1x01tzU8/9LTTL9hZZZogBU= +github.com/go-openapi/swag v0.22.4/go.mod h1:UzaqsxGiab7freDnrUUra0MwWfN/q7tE4j+VcZ0yl14= +github.com/go-playground/assert/v2 v2.2.0 h1:JvknZsQTYeFEAhQwI4qEt9cyV5ONwRHC+lYKSsYSR8s= +github.com/go-playground/assert/v2 v2.2.0/go.mod h1:VDjEfimB/XKnb+ZQfWdccd7VUvScMdVu0Titje2rxJ4= +github.com/go-playground/locales v0.14.1 h1:EWaQ/wswjilfKLTECiXz7Rh+3BjFhfDFKv/oXslEjJA= +github.com/go-playground/locales v0.14.1/go.mod h1:hxrqLVvrK65+Rwrd5Fc6F2O76J/NuW9t0sjnWqG1slY= +github.com/go-playground/universal-translator v0.18.1 h1:Bcnm0ZwsGyWbCzImXv+pAJnYK9S473LQFuzCbDbfSFY= +github.com/go-playground/universal-translator v0.18.1/go.mod h1:xekY+UJKNuX9WP91TpwSH2VMlDf28Uj24BCp08ZFTUY= +github.com/go-playground/validator/v10 v10.20.0 h1:K9ISHbSaI0lyB2eWMPJo+kOS/FBExVwjEviJTixqxL8= +github.com/go-playground/validator/v10 v10.20.0/go.mod h1:dbuPbCMFw/DrkbEynArYaCwl3amGuJotoKCe95atGMM= +github.com/go-redis/redis/v8 v8.11.0 h1:O1Td0mQ8UFChQ3N9zFQqo6kTU2cJ+/it88gDB+zg0wo= +github.com/go-redis/redis/v8 v8.11.0/go.mod h1:DLomh7y2e3ggQXQLd1YgmvIfecPJoFl7WU5SOQ/r06M= +github.com/go-test/deep v1.0.8 h1:TDsG77qcSprGbC6vTN8OuXp5g+J+b5Pcguhf7Zt61VM= +github.com/go-test/deep v1.0.8/go.mod h1:5C2ZWiW0ErCdrYzpqxLbTX7MG14M9iiw8DgHncVwcsE= +github.com/goccy/go-json v0.10.2 h1:CrxCmQqYDkv1z7lO7Wbh2HN93uovUHgrECaO5ZrCXAU= +github.com/goccy/go-json v0.10.2/go.mod h1:6MelG93GURQebXPDq3khkgXZkazVtN9CRI+MGFi0w8I= +github.com/gofrs/uuid v3.2.0+incompatible/go.mod h1:b2aQJv3Z4Fp6yNu3cdSllBxTCLRxnplIgP/c0N/04lM= +github.com/golang/protobuf v1.2.0/go.mod h1:6lQm79b+lXiMfvg/cZm0SGofjICqVBUtrP5yJMmIC1U= +github.com/golang/protobuf v1.3.5/go.mod h1:6O5/vntMXwX2lRkT1hjjk0nAC1IDOTvTlVgjlRvqsdk= +github.com/golang/protobuf v1.4.0-rc.1/go.mod h1:ceaxUfeHdC40wWswd/P6IGgMaK3YpKi5j83Wpe3EHw8= +github.com/golang/protobuf v1.4.0-rc.1.0.20200221234624-67d41d38c208/go.mod h1:xKAWHe0F5eneWXFV3EuXVDTCmh+JuBKY0li0aMyXATA= +github.com/golang/protobuf v1.4.0-rc.2/go.mod h1:LlEzMj4AhA7rCAGe4KMBDvJI+AwstrUpVNzEA03Pprs= +github.com/golang/protobuf v1.4.0-rc.4.0.20200313231945-b860323f09d0/go.mod h1:WU3c8KckQ9AFe+yFwt9sWVRKCVIyN9cPHBJSNnbL67w= +github.com/golang/protobuf v1.4.0/go.mod h1:jodUvKwWbYaEsadDk5Fwe5c77LiNKVO9IDvqG2KuDX0= +github.com/golang/protobuf v1.4.2/go.mod h1:oDoupMAO8OvCJWAcko0GGGIgR6R6ocIYbsSw735rRwI= +github.com/golang/protobuf v1.5.4 h1:i7eJL8qZTpSEXOPTxNKhASYpMn+8e5Q6AdndVa1dWek= +github.com/golang/protobuf v1.5.4/go.mod h1:lnTiLA8Wa4RWRcIUkrtSVa5nRhsEGBg48fD6rSs7xps= +github.com/google/go-cmp v0.3.0/go.mod h1:8QqcDgzrUqlUb/G2PQTWiueGozuR1884gddMywk6iLU= +github.com/google/go-cmp v0.3.1/go.mod h1:8QqcDgzrUqlUb/G2PQTWiueGozuR1884gddMywk6iLU= +github.com/google/go-cmp v0.4.0/go.mod h1:v8dTdLbMG2kIc/vJvl+f65V22dbkXbowE6jgT/gNBxE= +github.com/google/go-cmp v0.5.6/go.mod h1:v8dTdLbMG2kIc/vJvl+f65V22dbkXbowE6jgT/gNBxE= +github.com/google/go-cmp v0.5.9/go.mod h1:17dUlkBOakJ0+DkrSSNjCkIjxS6bF9zb3elmeNGIjoY= +github.com/google/go-cmp v0.6.0 h1:ofyhxvXcZhMsU5ulbFiLKl/XBFqE1GSq7atu8tAmTRI= +github.com/google/go-cmp v0.6.0/go.mod h1:17dUlkBOakJ0+DkrSSNjCkIjxS6bF9zb3elmeNGIjoY= +github.com/google/gofuzz v1.0.0/go.mod h1:dBl0BpW6vV/+mYPU4Po3pmUjxk6FQPldtuIdl/M65Eg= +github.com/google/uuid v1.6.0 h1:NIvaJDMOsjHA8n1jAhLSgzrAzy1Hgr+hNrb57e+94F0= +github.com/google/uuid v1.6.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo= +github.com/goph/emperror v0.17.2 h1:yLapQcmEsO0ipe9p5TaN22djm3OFV/TfM/fcYP0/J18= +github.com/goph/emperror v0.17.2/go.mod h1:+ZbQ+fUNO/6FNiUo0ujtMjhgad9Xa6fQL9KhH4LNHic= +github.com/gopherjs/gopherjs v1.17.2 h1:fQnZVsXk8uxXIStYb0N4bGk7jeyTalG/wsZjQ25dO0g= +github.com/gopherjs/gopherjs v1.17.2/go.mod h1:pRRIvn/QzFLrKfvEz3qUuEhtE/zLCWfreZ6J5gM2i+k= +github.com/gorilla/mux v1.8.0/go.mod h1:DVbg23sWSpFRCP0SfiEN6jmj59UnW/n46BH5rLB71So= +github.com/hpcloud/tail v1.0.0/go.mod h1:ab1qPbhIpdTxEkNHXyeSf5vhxWSCs/tWer42PpOxQnU= +github.com/invopop/yaml v0.1.0 h1:YW3WGUoJEXYfzWBjn00zIlrw7brGVD0fUKRYDPAPhrc= +github.com/invopop/yaml v0.1.0/go.mod h1:2XuRLgs/ouIrW3XNzuNj7J3Nvu/Dig5MXvbCEdiBN3Q= +github.com/jessevdk/go-flags v1.4.0/go.mod h1:4FA24M0QyGHXBuZZK/XkWh8h0e1EYbRYJSGM75WSRxI= +github.com/jinzhu/copier v0.4.0 h1:w3ciUoD19shMCRargcpm0cm91ytaBhDvuRpz1ODO/U8= +github.com/jinzhu/copier v0.4.0/go.mod h1:DfbEm0FYsaqBcKcFuvmOZb218JkPGtvSHsKg8S8hyyg= +github.com/josharian/intern v1.0.0 h1:vlS4z54oSdjm0bgjRigI+G1HpF+tI+9rE5LLzOg8HmY= +github.com/josharian/intern v1.0.0/go.mod h1:5DoeVV0s6jJacbCEi61lwdGj/aVlrQvzHFFd8Hwg//Y= +github.com/json-iterator/go v1.1.12 h1:PV8peI4a0ysnczrg+LtxykD8LfKY9ML6u2jnxaEnrnM= +github.com/json-iterator/go v1.1.12/go.mod h1:e30LSqwooZae/UwlEbR2852Gd8hjQvJoHmT4TnhNGBo= +github.com/jtolds/gls v4.20.0+incompatible h1:xdiiI2gbIgH/gLH7ADydsJ1uDOEzR8yvV7C0MuV77Wo= +github.com/jtolds/gls v4.20.0+incompatible/go.mod h1:QJZ7F/aHp+rZTRtaJ1ow/lLfFfVYBRgL+9YlvaHOwJU= +github.com/kardianos/osext v0.0.0-20190222173326-2bc1f35cddc0/go.mod h1:1NbS8ALrpOvjt0rHPNLyCIeMtbizbir8U//inJ+zuB8= +github.com/klauspost/cpuid/v2 v2.0.9/go.mod h1:FInQzS24/EEf25PyTYn52gqo7WaD8xa0213Md/qVLRg= +github.com/klauspost/cpuid/v2 v2.2.7 h1:ZWSB3igEs+d0qvnxR/ZBzXVmxkgt8DdzP6m9pfuVLDM= +github.com/klauspost/cpuid/v2 v2.2.7/go.mod h1:Lcz8mBdAVJIBVzewtcLocK12l3Y+JytZYpaMropDUws= +github.com/knz/go-libedit v1.10.1/go.mod h1:MZTVkCWyz0oBc7JOWP3wNAzd002ZbM/5hgShxwh4x8M= +github.com/konsorten/go-windows-terminal-sequences v1.0.1/go.mod h1:T0+1ngSBFLxvqU3pZ+m/2kptfBszLMUkC4ZK/EgS/cQ= +github.com/kr/pretty v0.1.0/go.mod h1:dAy3ld7l9f0ibDNOQOHHMYYIIbhfbHSm3C4ZsoJORNo= +github.com/kr/pretty v0.2.1/go.mod h1:ipq/a2n7PKx3OHsz4KJII5eveXtPO4qwEXGdVfWzfnI= +github.com/kr/pretty v0.3.1 h1:flRD4NNwYAUpkphVc1HcthR4KEIFJ65n8Mw5qdRn3LE= +github.com/kr/pretty v0.3.1/go.mod h1:hoEshYVHaxMs3cyo3Yncou5ZscifuDolrwPKZanG3xk= +github.com/kr/pty v1.1.1/go.mod h1:pFQYn66WHrOpPYNljwOMqo10TkYh1fy3cYio2l3bCsQ= +github.com/kr/text v0.1.0/go.mod h1:4Jbv+DJW3UT/LiOwJeYQe1efqtUx/iVham/4vfdArNI= +github.com/kr/text v0.2.0 h1:5Nx0Ya0ZqY2ygV366QzturHI13Jq95ApcVaJBhpS+AY= +github.com/kr/text v0.2.0/go.mod h1:eLer722TekiGuMkidMxC/pM04lWEeraHUUmBw8l2grE= +github.com/leodido/go-urn v1.4.0 h1:WT9HwE9SGECu3lg4d/dIA+jxlljEa1/ffXKmRjqdmIQ= +github.com/leodido/go-urn v1.4.0/go.mod h1:bvxc+MVxLKB4z00jd1z+Dvzr47oO32F/QSNjSBOlFxI= +github.com/lufia/plan9stats v0.0.0-20211012122336-39d0f177ccd0 h1:6E+4a0GO5zZEnZ81pIr0yLvtUWk2if982qA3F3QD6H4= +github.com/lufia/plan9stats v0.0.0-20211012122336-39d0f177ccd0/go.mod h1:zJYVVT2jmtg6P3p1VtQj7WsuWi/y4VnjVBn7F8KPB3I= +github.com/mailru/easyjson v0.0.0-20190614124828-94de47d64c63/go.mod h1:C1wdFJiN94OJF2b5HbByQZoLdCWB1Yqtg26g4irojpc= +github.com/mailru/easyjson v0.0.0-20190626092158-b2ccc519800e/go.mod h1:C1wdFJiN94OJF2b5HbByQZoLdCWB1Yqtg26g4irojpc= +github.com/mailru/easyjson v0.7.7 h1:UGYAvKxe3sBsEDzO8ZeWOSlIQfWFlxbzLZe7hwFURr0= +github.com/mailru/easyjson v0.7.7/go.mod h1:xzfreul335JAWq5oZzymOObrkdz5UnU4kGfJJLY9Nlc= +github.com/mattn/go-colorable v0.1.13 h1:fFA4WZxdEF4tXPZVKMLwD8oUnCTTo08duU7wxecdEvA= +github.com/mattn/go-colorable v0.1.13/go.mod h1:7S9/ev0klgBDR4GtXTXX8a3vIGJpMovkB8vQcUbaXHg= +github.com/mattn/go-isatty v0.0.20 h1:xfD0iDuEKnDkl03q4limB+vH+GxLEtL/jb4xVJSWWEY= +github.com/mattn/go-isatty v0.0.20/go.mod h1:W+V8PltTTMOvKvAeJH7IuucS94S2C6jfK/D7dTCTo3Y= +github.com/matttproud/golang_protobuf_extensions v1.0.4 h1:mmDVorXM7PCGKw94cs5zkfA9PSy5pEvNWRP0ET0TIVo= +github.com/matttproud/golang_protobuf_extensions v1.0.4/go.mod h1:BSXmuO+STAnVfrANrmjBb36TMTDstsz7MSK+HVaYKv4= +github.com/meguminnnnnnnnn/go-openai v0.0.0-20250821095446-07791bea23a0 h1:nIohpHs1ViKR0SVgW/cbBstHjmnqFZDM9RqgX9m9Xu8= +github.com/meguminnnnnnnnn/go-openai v0.0.0-20250821095446-07791bea23a0/go.mod h1:qs96ysDmxhE4BZoU45I43zcyfnaYxU3X+aRzLko/htY= +github.com/mgutz/ansi v0.0.0-20170206155736-9520e82c474b h1:j7+1HpAFS1zy5+Q4qx1fWh90gTKwiN4QCGoY9TWyyO4= +github.com/mgutz/ansi v0.0.0-20170206155736-9520e82c474b/go.mod h1:01TrycV0kFyexm33Z7vhZRXopbI8J3TDReVlkTgMUxE= +github.com/modern-go/concurrent v0.0.0-20180228061459-e0a39a4cb421/go.mod h1:6dJC0mAP4ikYIbvyc7fijjWJddQyLn8Ig3JB5CqoB9Q= +github.com/modern-go/concurrent v0.0.0-20180306012644-bacd9c7ef1dd h1:TRLaZ9cD/w8PVh93nsPXa1VrQ6jlwL5oN8l14QlcNfg= +github.com/modern-go/concurrent v0.0.0-20180306012644-bacd9c7ef1dd/go.mod h1:6dJC0mAP4ikYIbvyc7fijjWJddQyLn8Ig3JB5CqoB9Q= +github.com/modern-go/reflect2 v1.0.2 h1:xBagoLtFs94CBntxluKeaWgTMpvLxC4ur3nMaC9Gz0M= +github.com/modern-go/reflect2 v1.0.2/go.mod h1:yWuevngMOJpCy52FWWMvUC8ws7m/LJsjYzDa0/r8luk= +github.com/mohae/deepcopy v0.0.0-20170929034955-c48cc78d4826 h1:RWengNIwukTxcDr9M+97sNutRR1RKhG96O6jWumTTnw= +github.com/mohae/deepcopy v0.0.0-20170929034955-c48cc78d4826/go.mod h1:TaXosZuwdSHYgviHp1DAtfrULt5eUgsSMsZf+YrPgl8= +github.com/nikolalohinski/gonja v1.5.3 h1:GsA+EEaZDZPGJ8JtpeGN78jidhOlxeJROpqMT9fTj9c= +github.com/nikolalohinski/gonja v1.5.3/go.mod h1:RmjwxNiXAEqcq1HeK5SSMmqFJvKOfTfXhkJv6YBtPa4= +github.com/nxadm/tail v1.4.4 h1:DQuhQpB1tVlglWS2hLQ5OV6B5r8aGxSrPc5Qo6uTN78= +github.com/nxadm/tail v1.4.4/go.mod h1:kenIhsEOeOJmVchQTgglprH7qJGnHDVpk1VPCcaMI8A= +github.com/onsi/ginkgo v1.6.0/go.mod h1:lLunBs/Ym6LB5Z9jYTR76FiuTmxDTDusOGeTQH+WWjE= +github.com/onsi/ginkgo v1.8.0/go.mod h1:lLunBs/Ym6LB5Z9jYTR76FiuTmxDTDusOGeTQH+WWjE= +github.com/onsi/ginkgo v1.12.1/go.mod h1:zj2OWP4+oCPe1qIXoGWkgMRwljMUYCdkwsT2108oapk= +github.com/onsi/ginkgo v1.15.0/go.mod h1:hF8qUzuuC8DJGygJH3726JnCZX4MYbRB8yFfISqnKUg= +github.com/onsi/ginkgo v1.16.5 h1:8xi0RTUf59SOSfEtZMvwTvXYMzG4gV23XVHOZiXNtnE= +github.com/onsi/ginkgo v1.16.5/go.mod h1:+E8gABHa3K6zRBolWtd+ROzc/U5bkGt0FwiG042wbpU= +github.com/onsi/gomega v1.5.0/go.mod h1:ex+gbHU/CVuBBDIJjb2X0qEXbFg53c61hWP/1CpauHY= +github.com/onsi/gomega v1.7.1/go.mod h1:XdKZgCCFLUoM/7CFJVPcG8C1xQ1AJ0vpAezJrB7JYyY= +github.com/onsi/gomega v1.10.1/go.mod h1:iN09h71vgCQne3DLsj+A5owkum+a2tYe+TOCB1ybHNo= +github.com/onsi/gomega v1.10.5/go.mod h1:gza4q3jKQJijlu05nKWRCW/GavJumGt8aNRxWg7mt48= +github.com/onsi/gomega v1.27.3 h1:5VwIwnBY3vbBDOJrNtA4rVdiTZCsq9B5F12pvy1Drmk= +github.com/onsi/gomega v1.27.3/go.mod h1:5vG284IBtfDAmDyrK+eGyZmUgUlmi+Wngqo557cZ6Gw= +github.com/pelletier/go-toml/v2 v2.2.2 h1:aYUidT7k73Pcl9nb2gScu7NSrKCSHIDE89b3+6Wq+LM= +github.com/pelletier/go-toml/v2 v2.2.2/go.mod h1:1t835xjRzz80PqgE6HHgN2JOsmgYu/h4qDAS4n929Rs= +github.com/perimeterx/marshmallow v1.1.4 h1:pZLDH9RjlLGGorbXhcaQLhfuV0pFMNfPO55FuFkxqLw= +github.com/perimeterx/marshmallow v1.1.4/go.mod h1:dsXbUu8CRzfYP5a87xpp0xq9S3u0Vchtcl8we9tYaXw= +github.com/pkg/errors v0.8.0/go.mod h1:bwawxfHBFNV+L2hUp1rHADufV3IMtnDRdf1r5NINEl0= +github.com/pkg/errors v0.9.1 h1:FEBLx1zS214owpjy7qsBeixbURkuhQAwrK5UwLGTwt4= +github.com/pkg/errors v0.9.1/go.mod h1:bwawxfHBFNV+L2hUp1rHADufV3IMtnDRdf1r5NINEl0= +github.com/pkoukk/tiktoken-go v0.1.7 h1:qOBHXX4PHtvIvmOtyg1EeKlwFRiMKAcoMp4Q+bLQDmw= +github.com/pkoukk/tiktoken-go v0.1.7/go.mod h1:9NiV+i9mJKGj1rYOT+njbv+ZwA/zJxYdewGl6qVatpg= +github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM= +github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4= +github.com/power-devops/perfstat v0.0.0-20210106213030-5aafc221ea8c h1:ncq/mPwQF4JjgDlrVEn3C11VoGHZN7m8qihwgMEtzYw= +github.com/power-devops/perfstat v0.0.0-20210106213030-5aafc221ea8c/go.mod h1:OmDBASR4679mdNQnz2pUhc2G8CO2JrUAVFDRBDP/hJE= +github.com/prometheus/client_golang v1.16.0 h1:yk/hx9hDbrGHovbci4BY+pRMfSuuat626eFsHb7tmT8= +github.com/prometheus/client_golang v1.16.0/go.mod h1:Zsulrv/L9oM40tJ7T815tM89lFEugiJ9HzIqaAx4LKc= +github.com/prometheus/client_model v0.3.0 h1:UBgGFHqYdG/TPFD1B1ogZywDqEkwp3fBMvqdiQ7Xew4= +github.com/prometheus/client_model v0.3.0/go.mod h1:LDGWKZIo7rky3hgvBe+caln+Dr3dPggB5dvjtD7w9+w= +github.com/prometheus/common v0.42.0 h1:EKsfXEYo4JpWMHH5cg+KOUWeuJSov1Id8zGR8eeI1YM= +github.com/prometheus/common v0.42.0/go.mod h1:xBwqVerjNdUDjgODMpudtOMwlOwf2SaTr1yjz4b7Zbc= +github.com/prometheus/procfs v0.10.1 h1:kYK1Va/YMlutzCGazswoHKo//tZVlFpKYh+PymziUAg= +github.com/prometheus/procfs v0.10.1/go.mod h1:nwNm2aOCAYw8uTR/9bWRREkZFxAUcWzPHWJq+XBB/FM= +github.com/rogpeppe/go-internal v1.13.1 h1:KvO1DLK/DRN07sQ1LQKScxyZJuNnedQ5/wKSR38lUII= +github.com/rogpeppe/go-internal v1.13.1/go.mod h1:uMEvuHeurkdAXX61udpOXGD/AzZDWNMNyH2VO9fmH0o= +github.com/rollbar/rollbar-go v1.0.2/go.mod h1:AcFs5f0I+c71bpHlXNNDbOWJiKwjFDtISeXco0L5PKQ= +github.com/shirou/gopsutil/v3 v3.23.12 h1:z90NtUkp3bMtmICZKpC4+WaknU1eXtp5vtbQ11DgpE4= +github.com/shirou/gopsutil/v3 v3.23.12/go.mod h1:1FrWgea594Jp7qmjHUUPlJDTPgcsb9mGnXDxavtikzM= +github.com/shoenig/go-m1cpu v0.1.6 h1:nxdKQNcEB6vzgA2E2bvzKIYRuNj7XNJ4S/aRSwKzFtM= +github.com/shoenig/go-m1cpu v0.1.6/go.mod h1:1JJMcUBvfNwpq05QDQVAnx3gUHr9IYF7GNg9SUEw2VQ= +github.com/shoenig/test v0.6.4 h1:kVTaSd7WLz5WZ2IaoM0RSzRsUD+m8wRR+5qvntpn4LU= +github.com/shoenig/test v0.6.4/go.mod h1:byHiCGXqrVaflBLAMq/srcZIHynQPQgeyvkvXnjqq0k= +github.com/sirupsen/logrus v1.2.0/go.mod h1:LxeOpSwHxABJmUn/MG1IvRgCAasNZTLOkJPxbbu5VWo= +github.com/sirupsen/logrus v1.9.3 h1:dueUQJ1C2q9oE3F7wvmSGAaVtTmUizReu6fjN8uqzbQ= +github.com/sirupsen/logrus v1.9.3/go.mod h1:naHLuLoDiP4jHNo9R0sCBMtWGeIprob74mVsIT4qYEQ= +github.com/slongfield/pyfmt v0.0.0-20220222012616-ea85ff4c361f h1:Z2cODYsUxQPofhpYRMQVwWz4yUVpHF+vPi+eUdruUYI= +github.com/slongfield/pyfmt v0.0.0-20220222012616-ea85ff4c361f/go.mod h1:JqzWyvTuI2X4+9wOHmKSQCYxybB/8j6Ko43qVmXDuZg= +github.com/smarty/assertions v1.15.0 h1:cR//PqUBUiQRakZWqBiFFQ9wb8emQGDb0HeGdqGByCY= +github.com/smarty/assertions v1.15.0/go.mod h1:yABtdzeQs6l1brC900WlRNwj6ZR55d7B+E8C6HtKdec= +github.com/smartystreets/goconvey v1.8.1 h1:qGjIddxOk4grTu9JPOU31tVfq3cNdBlNa5sSznIX1xY= +github.com/smartystreets/goconvey v1.8.1/go.mod h1:+/u4qLyY6x1jReYOp7GOM2FSt8aP9CzCZL03bI28W60= +github.com/spaolacci/murmur3 v1.1.0 h1:7c1g84S4BPRrfL5Xrdp6fOJ206sU9y293DDHaoy0bLI= +github.com/spaolacci/murmur3 v1.1.0/go.mod h1:JwIasOWyU6f++ZhiEuf87xNszmSA2myDM2Kzu9HwQUA= +github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME= +github.com/stretchr/objx v0.1.1/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME= +github.com/stretchr/objx v0.4.0/go.mod h1:YvHI0jy2hoMjB+UWwv71VJQ9isScKT/TqJzVSSt89Yw= +github.com/stretchr/objx v0.5.0/go.mod h1:Yh+to48EsGEfYuaHDzXPcE3xhTkx73EhmCGUpEOglKo= +github.com/stretchr/objx v0.5.2 h1:xuMeJ0Sdp5ZMRXx/aWO6RZxdr3beISkG5/G/aIRr3pY= +github.com/stretchr/objx v0.5.2/go.mod h1:FRsXN1f5AsAjCGJKqEizvkpNtU+EGNCLh3NxZ/8L+MA= +github.com/stretchr/testify v1.2.2/go.mod h1:a8OnRcib4nhh0OaRAV+Yts87kKdq0PP7pXfy6kDkUVs= +github.com/stretchr/testify v1.3.0/go.mod h1:M5WIy9Dh21IEIfnGCwXGc5bZfKNJtfHm1UVUgZn+9EI= +github.com/stretchr/testify v1.7.0/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg= +github.com/stretchr/testify v1.7.1/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg= +github.com/stretchr/testify v1.8.0/go.mod h1:yNjHg4UonilssWZ8iaSj1OCr/vHnekPRkoO+kdMU+MU= +github.com/stretchr/testify v1.8.1/go.mod h1:w2LPCIKwWwSfY2zedu0+kehJoqGctiVI29o6fzry7u4= +github.com/stretchr/testify v1.8.4/go.mod h1:sz/lmYIOXD/1dqDmKjjqLyZ2RngseejIcXlSw2iwfAo= +github.com/stretchr/testify v1.9.0/go.mod h1:r2ic/lqez/lEtzL7wO/rwa5dbSLXVDPFyf8C91i36aY= +github.com/stretchr/testify v1.10.0 h1:Xv5erBjTwe/5IxqUQTdXv5kgmIvbHo3QQyRwhJsOfJA= +github.com/stretchr/testify v1.10.0/go.mod h1:r2ic/lqez/lEtzL7wO/rwa5dbSLXVDPFyf8C91i36aY= +github.com/tklauser/go-sysconf v0.3.12 h1:0QaGUFOdQaIVdPgfITYzaTegZvdCjmYO52cSFAEVmqU= +github.com/tklauser/go-sysconf v0.3.12/go.mod h1:Ho14jnntGE1fpdOqQEEaiKRpvIavV0hSfmBq8nJbHYI= +github.com/tklauser/numcpus v0.6.1 h1:ng9scYS7az0Bk4OZLvrNXNSAO2Pxr1XXRAPyjhIx+Fk= +github.com/tklauser/numcpus v0.6.1/go.mod h1:1XfjsgE2zo8GVw7POkMbHENHzVg3GzmoZ9fESEdAacY= +github.com/tmc/langchaingo v0.1.13 h1:rcpMWBIi2y3B90XxfE4Ao8dhCQPVDMaNPnN5cGB1CaA= +github.com/tmc/langchaingo v0.1.13/go.mod h1:vpQ5NOIhpzxDfTZK9B6tf2GM/MoaHewPWM5KXXGh7hg= +github.com/twitchyliquid64/golang-asm v0.15.1 h1:SU5vSMR7hnwNxj24w34ZyCi/FmDZTkS4MhqMhdFk5YI= +github.com/twitchyliquid64/golang-asm v0.15.1/go.mod h1:a1lVb/DtPvCB8fslRZhAngC2+aY1QWCk3Cedj/Gdt08= +github.com/ugorji/go v1.2.7/go.mod h1:nF9osbDWLy6bDVv/Rtoh6QgnvNDpmCalQV5urGCCS6M= +github.com/ugorji/go/codec v1.2.7/go.mod h1:WGN1fab3R1fzQlVQTkfxVtIBhWDRqOviHU95kRgeqEY= +github.com/ugorji/go/codec v1.2.12 h1:9LC83zGrHhuUA9l16C9AHXAqEV/2wBQ4nkvumAE65EE= +github.com/ugorji/go/codec v1.2.12/go.mod h1:UNopzCgEMSXjBc6AOMqYvWC1ktqTAfzJZUZgYf6w6lg= +github.com/wk8/go-ordered-map/v2 v2.1.8 h1:5h/BUHu93oj4gIdvHHHGsScSTMijfx5PeYkE/fJgbpc= +github.com/wk8/go-ordered-map/v2 v2.1.8/go.mod h1:5nJHM5DyteebpVlHnWMV0rPz6Zp7+xBAnxjb1X5vnTw= +github.com/x-cray/logrus-prefixed-formatter v0.5.2 h1:00txxvfBM9muc0jiLIEAkAcIMJzfthRT6usrui8uGmg= +github.com/x-cray/logrus-prefixed-formatter v0.5.2/go.mod h1:2duySbKsL6M18s5GU7VPsoEPHyzalCE06qoARUCeBBE= +github.com/yargevad/filepathx v1.0.0 h1:SYcT+N3tYGi+NvazubCNlvgIPbzAk7i7y2dwg3I5FYc= +github.com/yargevad/filepathx v1.0.0/go.mod h1:BprfX/gpYNJHJfc35GjRRpVcwWXS89gGulUIU5tK3tA= +github.com/yuin/goldmark v1.2.1/go.mod h1:3hX8gzYuyVAZsxl0MRgGTJEmQBFcNTphYh9decYSb74= +github.com/yusufpapurcu/wmi v1.2.3 h1:E1ctvB7uKFMOJw3fdOW32DwGE9I7t++CRUEMKvFoFiw= +github.com/yusufpapurcu/wmi v1.2.3/go.mod h1:SBZ9tNy3G9/m5Oi98Zks0QjeHVDvuK0qfxQmPyzfmi0= +golang.org/x/arch v0.11.0 h1:KXV8WWKCXm6tRpLirl2szsO5j/oOODwZf4hATmGVNs4= +golang.org/x/arch v0.11.0/go.mod h1:FEVrYAQjsQXMVJ1nsMoVVXPZg6p2JE2mx8psSWTDQys= +golang.org/x/crypto v0.0.0-20180904163835-0709b304e793/go.mod h1:6SG95UA2DQfeDnfUPMdvaQW0Q7yPrPDi9nlGo2tz2b4= +golang.org/x/crypto v0.0.0-20190308221718-c2843e01d9a2/go.mod h1:djNgcEr1/C05ACkg1iLfiJU5Ep61QUkGW8qpdssI0+w= +golang.org/x/crypto v0.0.0-20191011191535-87dc89f01550/go.mod h1:yigFU9vqHzYiE8UmvKecakEJjdnWj3jj499lnFckfCI= +golang.org/x/crypto v0.0.0-20200622213623-75b288015ac9/go.mod h1:LzIPMQfyMNhhGPhUkYOs5KpL4U8rLKemX1yGLhDgUto= +golang.org/x/crypto v0.31.0 h1:ihbySMvVjLAeSH1IbfcRTkD/iNscyz8rGzjF/E5hV6U= +golang.org/x/crypto v0.31.0/go.mod h1:kDsLvtWBEx7MV9tJOj9bnXsPbxwJQ6csT/x4KIN4Ssk= +golang.org/x/exp v0.0.0-20230713183714-613f0c0eb8a1 h1:MGwJjxBy0HJshjDNfLsYO8xppfqWlA5ZT9OhtUUhTNw= +golang.org/x/exp v0.0.0-20230713183714-613f0c0eb8a1/go.mod h1:FXUEEKJgO7OQYeo8N01OfiKP8RXMtf6e8aTskBGqWdc= +golang.org/x/mod v0.3.0/go.mod h1:s0Qsj1ACt9ePp/hMypM3fl4fZqREWJwdYDEqhRiZZUA= +golang.org/x/net v0.0.0-20180906233101-161cd47e91fd/go.mod h1:mL1N/T3taQHkDXs73rZJwtUhF3w3ftmwwsq0BUmARs4= +golang.org/x/net v0.0.0-20190404232315-eb5bcb51f2a3/go.mod h1:t9HGtf8HONx5eT2rtn7q6eTqICYqUVnKs3thJo3Qplg= +golang.org/x/net v0.0.0-20190620200207-3b0461eec859/go.mod h1:z5CRVTTTmAJ677TzLLGU+0bjPO0LkuOLi4/5GtJWs/s= +golang.org/x/net v0.0.0-20200520004742-59133d7f0dd7/go.mod h1:qpuaurCH72eLCgpAm/N6yyVIVM9cpaDIP3A8BGJEC5A= +golang.org/x/net v0.0.0-20201021035429-f5854403a974/go.mod h1:sp8m0HH+o8qH0wwXwYZr8TS3Oi6o0r6Gce1SSxlDquU= +golang.org/x/net v0.0.0-20201202161906-c7110b5ffcbb/go.mod h1:sp8m0HH+o8qH0wwXwYZr8TS3Oi6o0r6Gce1SSxlDquU= +golang.org/x/net v0.25.0 h1:d/OCCoBEUq33pjydKrGQhw7IlUPI2Oylr+8qLx49kac= +golang.org/x/net v0.25.0/go.mod h1:JkAGAh7GEvH74S6FOH42FLoXpXbE/aqXSrIQjXgsiwM= +golang.org/x/sync v0.0.0-20180314180146-1d60e4601c6f/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM= +golang.org/x/sync v0.0.0-20181221193216-37e7f081c4d4/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM= +golang.org/x/sync v0.0.0-20190423024810-112230192c58/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM= +golang.org/x/sync v0.0.0-20201020160332-67f06af15bc9/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM= +golang.org/x/sys v0.0.0-20180905080454-ebe1bf3edb33/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY= +golang.org/x/sys v0.0.0-20180909124046-d0be0721c37e/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY= +golang.org/x/sys v0.0.0-20190215142949-d0b11bdaac8a/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY= +golang.org/x/sys v0.0.0-20190412213103-97732733099d/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs= +golang.org/x/sys v0.0.0-20190904154756-749cb33beabd/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs= +golang.org/x/sys v0.0.0-20190916202348-b4ddaad3f8a3/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs= +golang.org/x/sys v0.0.0-20191005200804-aed5e4c7ecf9/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs= +golang.org/x/sys v0.0.0-20191120155948-bd437916bb0e/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs= +golang.org/x/sys v0.0.0-20200323222414-85ca7c5b95cd/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs= +golang.org/x/sys v0.0.0-20200930185726-fdedc70b468f/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs= +golang.org/x/sys v0.0.0-20201204225414-ed752295db88/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs= +golang.org/x/sys v0.0.0-20210112080510-489259a85091/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs= +golang.org/x/sys v0.0.0-20220715151400-c0bba94af5f8/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= +golang.org/x/sys v0.5.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= +golang.org/x/sys v0.6.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= +golang.org/x/sys v0.8.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= +golang.org/x/sys v0.11.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= +golang.org/x/sys v0.15.0/go.mod h1:/VUhepiaJMQUp4+oa/7Zr1D23ma6VTLIYjOOTFZPUcA= +golang.org/x/sys v0.28.0 h1:Fksou7UEQUWlKvIdsqzJmUmCX3cZuD2+P3XyyzwMhlA= +golang.org/x/sys v0.28.0/go.mod h1:/VUhepiaJMQUp4+oa/7Zr1D23ma6VTLIYjOOTFZPUcA= +golang.org/x/term v0.27.0 h1:WP60Sv1nlK1T6SupCHbXzSaN0b9wUmsPoRS9b61A23Q= +golang.org/x/term v0.27.0/go.mod h1:iMsnZpn0cago0GOrHO2+Y7u7JPn5AylBrcoWkElMTSM= +golang.org/x/text v0.3.0/go.mod h1:NqM8EUOU14njkJ3fqMW+pc6Ldnwhi/IjpwHt7yyuwOQ= +golang.org/x/text v0.3.3/go.mod h1:5Zoc/QRtKVWzQhOtBMvqHzDpF6irO9z98xDceosuGiQ= +golang.org/x/text v0.21.0 h1:zyQAAkrwaneQ066sspRyJaG9VNi/YJ1NfzcGB3hZ/qo= +golang.org/x/text v0.21.0/go.mod h1:4IBbMaMmOPCJ8SecivzSH54+73PCFmPWxNTLm+vZkEQ= +golang.org/x/tools v0.0.0-20180917221912-90fa682c2a6e/go.mod h1:n7NCudcB/nEzxVGmLbDWY5pfWTLqBcC2KZ6jyYvM4mQ= +golang.org/x/tools v0.0.0-20191119224855-298f0cb1881e/go.mod h1:b+2E5dAYhXwXZwtnZ6UAqBI28+e2cm9otk0dWdXHAEo= +golang.org/x/tools v0.0.0-20201224043029-2b0845dc783e/go.mod h1:emZCQorbCU4vsT4fOWvOPXz4eW1wZW4PmDk9uLelYpA= +golang.org/x/xerrors v0.0.0-20190717185122-a985d3407aa7/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0= +golang.org/x/xerrors v0.0.0-20191011141410-1b5146add898/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0= +golang.org/x/xerrors v0.0.0-20191204190536-9bdfabe68543/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0= +golang.org/x/xerrors v0.0.0-20200804184101-5ec99f83aff1/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0= +google.golang.org/protobuf v0.0.0-20200109180630-ec00e32a8dfd/go.mod h1:DFci5gLYBciE7Vtevhsrf46CRTquxDuWsQurQQe4oz8= +google.golang.org/protobuf v0.0.0-20200221191635-4d8936d0db64/go.mod h1:kwYJMbMJ01Woi6D6+Kah6886xMZcty6N08ah7+eCXa0= +google.golang.org/protobuf v0.0.0-20200228230310-ab0ca4ff8a60/go.mod h1:cfTl7dwQJ+fmap5saPgwCLgHXTUD7jkjRqWcaiX5VyM= +google.golang.org/protobuf v1.20.1-0.20200309200217-e05f789c0967/go.mod h1:A+miEFZTKqfCUM6K7xSMQL9OKL/b6hQv+e19PK+JZNE= +google.golang.org/protobuf v1.21.0/go.mod h1:47Nbq4nVaFHyn7ilMalzfO3qCViNmqZ2kzikPIcrTAo= +google.golang.org/protobuf v1.23.0/go.mod h1:EGpADcykh3NcUnDUJcl1+ZksZNG86OlYog2l/sGQquU= +google.golang.org/protobuf v1.34.1 h1:9ddQBjfCyZPOHPUiPxpYESBLc+T8P3E+Vo4IbKZgFWg= +google.golang.org/protobuf v1.34.1/go.mod h1:c6P6GXX6sHbq/GpV6MGZEdwhWPcYBgnhAHhKbcUYpos= +gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0= +gopkg.in/check.v1 v1.0.0-20180628173108-788fd7840127/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0= +gopkg.in/check.v1 v1.0.0-20201130134442-10cb98267c6c h1:Hei/4ADfdWqJk1ZMxUNpqntNwaWcugrBjAiHlqqRiVk= +gopkg.in/check.v1 v1.0.0-20201130134442-10cb98267c6c/go.mod h1:JHkPIbrfpd72SG/EVd6muEfDQjcINNoR0C8j2r3qZ4Q= +gopkg.in/fsnotify.v1 v1.4.7/go.mod h1:Tz8NjZHkW78fSQdbUxIjBTcgA1z1m8ZHf0WmKUhAMys= +gopkg.in/tomb.v1 v1.0.0-20141024135613-dd632973f1e7 h1:uRGJdciOHaEIrze2W8Q3AKkepLTh2hOroT7a+7czfdQ= +gopkg.in/tomb.v1 v1.0.0-20141024135613-dd632973f1e7/go.mod h1:dt/ZhP58zS4L8KSrWDmTeBkI65Dw0HsyUHuEVlX15mw= +gopkg.in/yaml.v2 v2.2.1/go.mod h1:hI93XBmqTisBFMUTm0b8Fm+jr3Dg1NNxqwp+5A1VGuI= +gopkg.in/yaml.v2 v2.2.2/go.mod h1:hI93XBmqTisBFMUTm0b8Fm+jr3Dg1NNxqwp+5A1VGuI= +gopkg.in/yaml.v2 v2.2.4/go.mod h1:hI93XBmqTisBFMUTm0b8Fm+jr3Dg1NNxqwp+5A1VGuI= +gopkg.in/yaml.v2 v2.3.0/go.mod h1:hI93XBmqTisBFMUTm0b8Fm+jr3Dg1NNxqwp+5A1VGuI= +gopkg.in/yaml.v2 v2.4.0 h1:D8xgwECY7CYvx+Y2n4sBz93Jn9JRvxdiyyo8CTfuKaY= +gopkg.in/yaml.v2 v2.4.0/go.mod h1:RDklbk79AGWmwhnvt/jBztapEOGDOx6ZbXqjP6csGnQ= +gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM= +gopkg.in/yaml.v3 v3.0.0/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM= +gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA= +gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM= +nullprogram.com/x/optparse v1.0.0/go.mod h1:KdyPE+Igbe0jQUrVfMqDMeJQIJZEuyV7pjYmp6pbG50= +sigs.k8s.io/yaml v1.3.0 h1:a2VclLzOGrwOHDiV8EfBGhvjHvP46CtW5j6POvhYGGo= +sigs.k8s.io/yaml v1.3.0/go.mod h1:GeOyir5tyXNByN85N/dRIT9es5UQNerPYEKK56eTBm8= diff --git a/tests/benchmark/llm_token_ratelimit/llm_client/llm_client.go b/tests/benchmark/llm_token_ratelimit/llm_client/llm_client.go new file mode 100644 index 000000000..4412ceee8 --- /dev/null +++ b/tests/benchmark/llm_token_ratelimit/llm_client/llm_client.go @@ -0,0 +1,227 @@ +// Copyright 1999-2020 Alibaba Group Holding Ltd. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package llm_client + +import ( + "context" + "fmt" + "os" + + langchaingo_llms "github.com/tmc/langchaingo/llms" + langchaingo_openai "github.com/tmc/langchaingo/llms/openai" + + eino_openai "github.com/cloudwego/eino-ext/components/model/openai" + eino_model "github.com/cloudwego/eino/components/model" + eino_schema "github.com/cloudwego/eino/schema" +) + +type LLMMessage struct { + Role string `json:"role"` + Content string `json:"content"` +} + +type LLMRequestInfos struct { + Provider LLMProvider `json:"provider"` + Messages []LLMMessage `json:"messages"` + Model string `json:"model"` +} + +type LLMResponseInfos struct { + Content string `json:"content"` + Usage map[string]any `json:"usage,omitempty"` +} + +type LLMProvider int32 + +const ( + LangChain LLMProvider = iota + Eino +) + +func (p LLMProvider) String() string { + switch p { + case LangChain: + return "langchain" + case Eino: + return "eino" + default: + return "unknown" + } +} + +func ParseLLMProvider(s string) (LLMProvider, error) { + switch s { + case "langchain": + return LangChain, nil + case "eino": + return Eino, nil + default: + return 0, fmt.Errorf("unknown LLM provider: %s", s) + } +} + +func (p LLMProvider) MarshalJSON() ([]byte, error) { + return []byte(`"` + p.String() + `"`), nil +} + +func (p *LLMProvider) UnmarshalJSON(data []byte) error { + // 去除引号 + s := string(data) + if len(s) >= 2 && s[0] == '"' && s[len(s)-1] == '"' { + s = s[1 : len(s)-1] + } + + provider, err := ParseLLMProvider(s) + if err != nil { + return err + } + *p = provider + return nil +} + +// ================================= LLMClient ==================================== + +type LLMClient interface { + GenerateContent(infos *LLMRequestInfos) (*LLMResponseInfos, error) + GetProvider() LLMProvider +} + +func NewLLMClient(infos *LLMRequestInfos) (LLMClient, error) { + switch infos.Provider { + case LangChain: + return NewLangChainClient(infos.Model) + case Eino: + return NewEinoClient(infos.Model) + default: + return nil, fmt.Errorf("unsupported provider: %v", infos.Provider) + } +} + +// ================================= LangChainClient ================================ + +type LangChainClient struct { + llm langchaingo_llms.Model +} + +func NewLangChainClient(model string) (LLMClient, error) { + apiKey := os.Getenv("LLM_API_KEY") + if apiKey == "" { + return nil, fmt.Errorf("LLM_API_KEY environment variable is not set") + } + + baseURL := os.Getenv("LLM_BASE_URL") + if baseURL == "" { + return nil, fmt.Errorf("LLM_BASE_URL environment variable is not set") + } + + llm, err := langchaingo_openai.New( + langchaingo_openai.WithToken(apiKey), + langchaingo_openai.WithBaseURL(baseURL), + langchaingo_openai.WithModel(model), + ) + if err != nil { + return nil, fmt.Errorf("failed to create LangChain LLM client: %w", err) + } + + return &LangChainClient{ + llm: llm, + }, nil +} + +func (c *LangChainClient) GenerateContent(infos *LLMRequestInfos) (*LLMResponseInfos, error) { + if infos == nil || infos.Messages == nil { + return nil, fmt.Errorf("invalid request infos") + } + content := make([]langchaingo_llms.MessageContent, len(infos.Messages)) + for i, msg := range infos.Messages { + content[i] = langchaingo_llms.TextParts(langchaingo_llms.ChatMessageType(msg.Role), msg.Content) + } + completion, err := c.llm.GenerateContent(context.Background(), content) + if err != nil { + return nil, fmt.Errorf("failed to generate content: %w", err) + } + return &LLMResponseInfos{ + Content: completion.Choices[0].Content, + Usage: map[string]any{ + "prompt_tokens": completion.Choices[0].GenerationInfo["PromptTokens"], + "completion_tokens": completion.Choices[0].GenerationInfo["CompletionTokens"], + "total_tokens": completion.Choices[0].GenerationInfo["TotalTokens"], + }, + }, nil +} + +func (c *LangChainClient) GetProvider() LLMProvider { + return LangChain +} + +// ================================= EinoClient ==================================== + +type EinoClient struct { + llm eino_model.BaseChatModel +} + +func NewEinoClient(model string) (LLMClient, error) { + apiKey := os.Getenv("LLM_API_KEY") + if apiKey == "" { + return nil, fmt.Errorf("LLM_API_KEY environment variable is not set") + } + + baseURL := os.Getenv("LLM_BASE_URL") + if baseURL == "" { + return nil, fmt.Errorf("LLM_BASE_URL environment variable is not set") + } + + llm, err := eino_openai.NewChatModel(context.Background(), &eino_openai.ChatModelConfig{ + BaseURL: baseURL, + APIKey: apiKey, + Model: model, + }) + if err != nil { + return nil, fmt.Errorf("failed to create Eino LLM client: %w", err) + } + + return &EinoClient{ + llm: llm, + }, nil +} + +func (c *EinoClient) GenerateContent(infos *LLMRequestInfos) (*LLMResponseInfos, error) { + if infos == nil || infos.Messages == nil { + return nil, fmt.Errorf("invalid request infos") + } + content := make([]*eino_schema.Message, len(infos.Messages)) + for i, msg := range infos.Messages { + content[i] = &eino_schema.Message{ + Role: eino_schema.RoleType(msg.Role), + Content: msg.Content, + } + } + completion, err := c.llm.Generate(context.Background(), content) + if err != nil { + return nil, fmt.Errorf("failed to generate content: %w", err) + } + return &LLMResponseInfos{ + Content: completion.Content, + Usage: map[string]any{ + "prompt_tokens": completion.ResponseMeta.Usage.PromptTokens, + "completion_tokens": completion.ResponseMeta.Usage.CompletionTokens, + "total_tokens": completion.ResponseMeta.Usage.TotalTokens, + }, + }, nil +} + +func (c *EinoClient) GetProvider() LLMProvider { + return Eino +} diff --git a/tests/benchmark/llm_token_ratelimit/main.go b/tests/benchmark/llm_token_ratelimit/main.go new file mode 100644 index 000000000..000b404a4 --- /dev/null +++ b/tests/benchmark/llm_token_ratelimit/main.go @@ -0,0 +1,32 @@ +// Copyright 1999-2020 Alibaba Group Holding Ltd. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package main + +import ( + "llm_token_ratelimit/ratelimit" + "llm_token_ratelimit/server" + + "github.com/gin-gonic/gin" +) + +func StartSerivce() { + gin.SetMode(gin.ReleaseMode) + ratelimit.InitSentinel() + server.StartServer("0.0.0.0", 9527) +} + +func main() { + StartSerivce() +} diff --git a/tests/benchmark/llm_token_ratelimit/ratelimit/options.go b/tests/benchmark/llm_token_ratelimit/ratelimit/options.go new file mode 100644 index 000000000..c0ecc8d5d --- /dev/null +++ b/tests/benchmark/llm_token_ratelimit/ratelimit/options.go @@ -0,0 +1,62 @@ +// Copyright 1999-2020 Alibaba Group Holding Ltd. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package ratelimit + +import ( + llmtokenratelimit "github.com/alibaba/sentinel-golang/core/llm_token_ratelimit" + "github.com/gin-gonic/gin" +) + +type Option func(*options) + +type options struct { + resourceExtract func(*gin.Context) string + blockFallback func(*gin.Context) + requestInfosExtract func(*gin.Context) *llmtokenratelimit.RequestInfos + promptsExtract func(*gin.Context) []string +} + +func WithBlockFallback(fn func(ctx *gin.Context)) Option { + return func(opts *options) { + opts.blockFallback = fn + } +} + +func WithResourceExtractor(fn func(*gin.Context) string) Option { + return func(opts *options) { + opts.resourceExtract = fn + } +} + +func WithRequestInfosExtractor(fn func(*gin.Context) *llmtokenratelimit.RequestInfos) Option { + return func(opts *options) { + opts.requestInfosExtract = fn + } +} + +func WithPromptsExtractor(fn func(*gin.Context) []string) Option { + return func(opts *options) { + opts.promptsExtract = fn + } +} + +func evaluateOptions(opts []Option) *options { + optCopy := &options{} + for _, opt := range opts { + opt(optCopy) + } + + return optCopy +} diff --git a/tests/benchmark/llm_token_ratelimit/ratelimit/ratelimit.go b/tests/benchmark/llm_token_ratelimit/ratelimit/ratelimit.go new file mode 100644 index 000000000..a6d47d61d --- /dev/null +++ b/tests/benchmark/llm_token_ratelimit/ratelimit/ratelimit.go @@ -0,0 +1,98 @@ +// Copyright 1999-2020 Alibaba Group Holding Ltd. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package ratelimit + +import ( + sentinel "github.com/alibaba/sentinel-golang/api" + "github.com/alibaba/sentinel-golang/core/base" + llmtokenratelimit "github.com/alibaba/sentinel-golang/core/llm_token_ratelimit" + "github.com/gin-gonic/gin" +) + +func InitSentinel() { + if err := sentinel.InitDefault(); err != nil { + panic(err) + } +} + +func SentinelMiddleware(opts ...Option) gin.HandlerFunc { + options := evaluateOptions(opts) + return func(c *gin.Context) { + resource := c.Request.Method + ":" + c.FullPath() + + if options.resourceExtract != nil { + resource = options.resourceExtract(c) + } + + prompts := []string{} + if options.promptsExtract != nil { + prompts = options.promptsExtract(c) + } + + reqInfos := llmtokenratelimit.GenerateRequestInfos( + llmtokenratelimit.WithHeader(c.Request.Header), + llmtokenratelimit.WithPrompts(prompts), + ) + + if options.requestInfosExtract != nil { + reqInfos = options.requestInfosExtract(c) + } + + // Check + entry, err := sentinel.Entry(resource, sentinel.WithTrafficType(base.Inbound), sentinel.WithArgs(reqInfos)) + if err != nil { + // Block + if options.blockFallback != nil { + options.blockFallback(c) + } else { + responseHeader, ok := err.TriggeredValue().(*llmtokenratelimit.ResponseHeader) + if !ok || responseHeader == nil { + c.AbortWithStatusJSON(500, gin.H{ + "error": "internal server error. invalid response header.", + }) + return + } + setResponseHeaders(c, responseHeader) + c.AbortWithStatusJSON(int(responseHeader.ErrorCode), gin.H{ + "error": responseHeader.ErrorMessage, + }) + } + return + } + // Set response headers + responseHeader, ok := entry.Context().GetPair(llmtokenratelimit.KeyResponseHeaders).(*llmtokenratelimit.ResponseHeader) + if ok && responseHeader != nil { + setResponseHeaders(c, responseHeader) + } + // Pass or Disabled + c.Next() + // Update used token info + usedTokenInfos, exists := c.Get(llmtokenratelimit.KeyUsedTokenInfos) + if exists && usedTokenInfos != nil { + entry.SetPair(llmtokenratelimit.KeyUsedTokenInfos, usedTokenInfos) + } + entry.Exit() // Must be executed immediately after the SetPair function + } +} + +func setResponseHeaders(c *gin.Context, header *llmtokenratelimit.ResponseHeader) { + if c == nil || header == nil { + return + } + + for key, value := range header.GetAll() { + c.Header(key, value) + } +} diff --git a/tests/benchmark/llm_token_ratelimit/redis-cluster.conf b/tests/benchmark/llm_token_ratelimit/redis-cluster.conf new file mode 100644 index 000000000..64538f398 --- /dev/null +++ b/tests/benchmark/llm_token_ratelimit/redis-cluster.conf @@ -0,0 +1,9 @@ +port 6379 +bind 0.0.0.0 +protected-mode no +daemonize no +appendonly yes +cluster-enabled yes +cluster-config-file nodes.conf +cluster-node-timeout 5000 +cluster-announce-bus-port 16379 diff --git a/tests/benchmark/llm_token_ratelimit/redis-single.conf b/tests/benchmark/llm_token_ratelimit/redis-single.conf new file mode 100644 index 000000000..069884f6c --- /dev/null +++ b/tests/benchmark/llm_token_ratelimit/redis-single.conf @@ -0,0 +1,6 @@ +port 6379 +bind 0.0.0.0 +protected-mode no +daemonize no +appendonly yes +cluster-enabled no diff --git a/tests/benchmark/llm_token_ratelimit/sentinel-cluster.yml b/tests/benchmark/llm_token_ratelimit/sentinel-cluster.yml new file mode 100644 index 000000000..9f8c38310 --- /dev/null +++ b/tests/benchmark/llm_token_ratelimit/sentinel-cluster.yml @@ -0,0 +1,48 @@ +version: "v1" +sentinel: + app: + name: sentinel-go-demo + log: + metric: + maxFileCount: 7 + llmTokenRatelimit: + enabled: true + rules: + - resource: "POST:/v1/chat/completion/fixed_window" + strategy: "fixed-window" + specificItems: + - identifier: + type: "header" + value: ".*" + keyItems: + - key: ".*" + token: + number: 0 + countStrategy: "total-tokens" + time: + unit: "second" + value: 60 + - resource: "POST:/v1/chat/completion/peta" + strategy: "peta" + specificItems: + - identifier: + type: "header" + value: ".*" + keyItems: + - key: ".*" + token: + number: 0 + countStrategy: "input-tokens" + time: + unit: "second" + value: 60 + errorCode: 429 + errorMessage: "Too Many Requests" + redis: + addrs: + - name: "redis-node-1" + port: 6379 + - name: "redis-node-2" + port: 6379 + - name: "redis-node-3" + port: 6379 \ No newline at end of file diff --git a/tests/benchmark/llm_token_ratelimit/sentinel-single.yml b/tests/benchmark/llm_token_ratelimit/sentinel-single.yml new file mode 100644 index 000000000..acbb90efb --- /dev/null +++ b/tests/benchmark/llm_token_ratelimit/sentinel-single.yml @@ -0,0 +1,44 @@ +version: "v1" +sentinel: + app: + name: sentinel-go-demo + log: + metric: + maxFileCount: 7 + llmTokenRatelimit: + enabled: true + rules: + - resource: "POST:/v1/chat/completion/fixed_window" + strategy: "fixed-window" + specificItems: + - identifier: + type: "header" + value: ".*" + keyItems: + - key: ".*" + token: + number: 0 + countStrategy: "total-tokens" + time: + unit: "second" + value: 60 + - resource: "POST:/v1/chat/completion/peta" + strategy: "peta" + specificItems: + - identifier: + type: "header" + value: ".*" + keyItems: + - key: ".*" + token: + number: 0 + countStrategy: "input-tokens" + time: + unit: "second" + value: 60 + errorCode: 429 + errorMessage: "Too Many Requests" + redis: + addrs: + - name: "redis-node-1" + port: 6379 \ No newline at end of file diff --git a/tests/benchmark/llm_token_ratelimit/server/route.go b/tests/benchmark/llm_token_ratelimit/server/route.go new file mode 100644 index 000000000..b771f14c7 --- /dev/null +++ b/tests/benchmark/llm_token_ratelimit/server/route.go @@ -0,0 +1,94 @@ +// Copyright 1999-2020 Alibaba Group Holding Ltd. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package server + +import ( + "llm_token_ratelimit/llm_client" + "net/http" + "time" + + llmtokenratelimit "github.com/alibaba/sentinel-golang/core/llm_token_ratelimit" + "github.com/gin-gonic/gin" +) + +func (s *Server) chatCompletion(c *gin.Context) { + infos, err := bindJSONFromCache[llm_client.LLMRequestInfos](c) + if err != nil { + c.JSON(http.StatusBadRequest, gin.H{ + "error": gin.H{ + "message": "failed to process LLM request", + "details": err.Error(), + "type": "api_error", + }, + }) + return + } + client, err := llm_client.NewLLMClient(infos) + if err != nil { + c.JSON(http.StatusInternalServerError, gin.H{ + "error": gin.H{ + "message": "failed to create LLM client", + "details": err.Error(), + "type": "client_error", + }, + }) + return + } + response, err := client.GenerateContent(infos) + if err != nil { + c.JSON(http.StatusInternalServerError, gin.H{ + "error": gin.H{ + "message": "failed to generate content", + "details": err.Error(), + "type": "generation_error", + }, + }) + return + } + + if response == nil { + c.JSON(http.StatusInternalServerError, gin.H{ + "error": gin.H{ + "message": "received nil response from LLM", + "type": "response_error", + }, + }) + return + } + + usedTokenInfos, err := llmtokenratelimit.OpenAITokenExtractor(response.Usage) + if err != nil { + c.JSON(http.StatusInternalServerError, gin.H{ + "error": gin.H{ + "message": "failed to extract token usage info", + "details": err.Error(), + "type": "extraction_error", + }, + }) + return + } + c.Set(llmtokenratelimit.KeyUsedTokenInfos, usedTokenInfos) + + c.JSON(http.StatusOK, gin.H{ + "message": "success", + "timestamp": time.Now().Unix(), + "choices": response.Content, + "usage": gin.H{ + "input_tokens": usedTokenInfos.InputTokens, + "output_tokens": usedTokenInfos.OutputTokens, + "total_tokens": usedTokenInfos.TotalTokens, + }, + }) +} diff --git a/tests/benchmark/llm_token_ratelimit/server/server.go b/tests/benchmark/llm_token_ratelimit/server/server.go new file mode 100644 index 000000000..490633084 --- /dev/null +++ b/tests/benchmark/llm_token_ratelimit/server/server.go @@ -0,0 +1,156 @@ +// Copyright 1999-2020 Alibaba Group Holding Ltd. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package server + +import ( + "bytes" + "fmt" + "io" + "llm_token_ratelimit/llm_client" + "llm_token_ratelimit/ratelimit" + "net/http" + + "github.com/gin-gonic/gin" +) + +const ( + KeyRawBody = "rawBody" +) + +func corsMiddleware() gin.HandlerFunc { + return func(c *gin.Context) { + c.Header("Access-Control-Allow-Origin", "*") + c.Header("Access-Control-Allow-Methods", "GET, POST, PUT, DELETE, OPTIONS") + c.Header("Access-Control-Allow-Headers", "Content-Type, Authorization") + + if c.Request.Method == "OPTIONS" { + c.AbortWithStatus(http.StatusNoContent) + return + } + + c.Next() + } +} + +func cacheBodyMiddleware() gin.HandlerFunc { + return func(c *gin.Context) { + if c.Request.Body == nil { + c.Next() + return + } + + bodyBytes, err := io.ReadAll(c.Request.Body) + if err != nil { + c.JSON(http.StatusBadRequest, gin.H{ + "error": "failed to read request body", + }) + c.Abort() + return + } + + c.Request.Body.Close() + c.Request.Body = io.NopCloser(bytes.NewBuffer(bodyBytes)) + + c.Set(KeyRawBody, bodyBytes) + + c.Next() + } +} + +func bindJSONFromCache[T any](c *gin.Context) (*T, error) { + rawBody, exists := c.Get(KeyRawBody) + if !exists { + return nil, fmt.Errorf("raw body not found in context") + } + + bodyBytes, ok := rawBody.([]byte) + if !ok { + return nil, fmt.Errorf("invalid raw body type") + } + + tempBody := io.NopCloser(bytes.NewBuffer(bodyBytes)) + originalBody := c.Request.Body + c.Request.Body = tempBody + + var result T + err := c.ShouldBindJSON(&result) + + c.Request.Body = originalBody + + if err != nil { + return nil, err + } + + return &result, nil +} + +type Server struct { + engine *gin.Engine + ip string + port uint16 +} + +func NewServer(ip string, port uint16) *Server { + engine := gin.New() + + engine.Use(gin.Logger()) + engine.Use(gin.Recovery()) + engine.Use(corsMiddleware()) + engine.Use(cacheBodyMiddleware()) + engine.Use(ratelimit.SentinelMiddleware( + ratelimit.WithPromptsExtractor(func(c *gin.Context) []string { + infos, err := bindJSONFromCache[llm_client.LLMRequestInfos](c) + if err != nil { + c.JSON(http.StatusBadRequest, gin.H{ + "error": gin.H{ + "message": "failed to process LLM request", + "details": err.Error(), + "type": "api_error", + }, + }) + return nil + } + prompts := make([]string, 0, len(infos.Messages)) + for _, msg := range infos.Messages { + prompts = append(prompts, msg.Content) + } + return prompts + }), + )) + + return &Server{ + engine: engine, + ip: ip, + port: port, + } +} + +func (s *Server) setupRoutes() { + s.engine.POST("/v1/chat/completion/fixed_window", s.chatCompletion) + s.engine.POST("/v1/chat/completion/peta", s.chatCompletion) +} + +func (s *Server) Start() error { + s.setupRoutes() + addr := fmt.Sprintf("%s:%d", s.ip, s.port) + return s.engine.Run(addr) +} + +func StartServer(ip string, port uint16) { + server := NewServer(ip, port) + if err := server.Start(); err != nil { + panic(err) + } +} diff --git a/util/regex_match.go b/util/regex_match.go new file mode 100644 index 000000000..c105d637d --- /dev/null +++ b/util/regex_match.go @@ -0,0 +1,71 @@ +package util + +import ( + "regexp" + "strings" + "sync" +) + +var ( + regexCache = make(map[string]*regexp.Regexp) + cacheMu sync.RWMutex +) + +const ( + RegexBeginPattern = "^" + RegexEndPattern = "$" +) + +func RegexMatch(pattern, text string) bool { + if pattern == "" && text == "" { + return true + } + if pattern == "" || text == "" { + return false + } + + exactPattern := ensureExactRegexMatch(pattern) + + regex, err := getCompiledRegex(exactPattern) + if err != nil { + return false + } + + return regex.MatchString(text) +} + +func ensureExactRegexMatch(pattern string) string { + if strings.HasPrefix(pattern, RegexBeginPattern) && strings.HasSuffix(pattern, RegexEndPattern) { + return pattern + } + + if strings.HasPrefix(pattern, RegexBeginPattern) && !strings.HasSuffix(pattern, RegexEndPattern) { + return pattern + RegexEndPattern + } + + if !strings.HasPrefix(pattern, RegexBeginPattern) && strings.HasSuffix(pattern, RegexEndPattern) { + return RegexBeginPattern + pattern + } + + return RegexBeginPattern + pattern + RegexEndPattern +} + +func getCompiledRegex(pattern string) (*regexp.Regexp, error) { + cacheMu.RLock() + if regex, exists := regexCache[pattern]; exists { + cacheMu.RUnlock() + return regex, nil + } + cacheMu.RUnlock() + + regex, err := regexp.Compile(pattern) + if err != nil { + return nil, err + } + + cacheMu.Lock() + regexCache[pattern] = regex + cacheMu.Unlock() + + return regex, nil +} diff --git a/util/regex_match_test.go b/util/regex_match_test.go new file mode 100644 index 000000000..39b8b7df3 --- /dev/null +++ b/util/regex_match_test.go @@ -0,0 +1,337 @@ +// Copyright 1999-2020 Alibaba Group Holding Ltd. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package util + +import ( + "regexp" + "testing" +) + +func TestRegexMatch(t *testing.T) { + tests := []struct { + name string + pattern string + text string + want bool + }{ + {"empty pattern", "", "test", false}, + {"empty text", "test", "", false}, + {"both empty", "", "", true}, + {"asterisk pattern empty text", ".*", "", false}, + + {"asterisk pattern", ".*", "anything", true}, + {"asterisk pattern special chars", ".*", "!@#$%^&*()", true}, + + {"exact match", "hello", "hello", true}, + {"no match", "hello", "world", false}, + {"case sensitive", "Hello", "hello", false}, + {"partial match should fail", "hello", "hello world", false}, + {"partial match should fail 2", "world", "hello world", false}, + + {"dot wildcard", "h.llo", "hello", true}, + {"dot wildcard", "h.llo", "hallo", true}, + {"dot wildcard no match", "h.llo", "hllo", false}, + {"dot wildcard partial should fail", "h.llo", "say hello", false}, + {"star quantifier", "hel*o", "helo", true}, + {"star quantifier", "hel*o", "hello", true}, + {"star quantifier multiple", "hel*o", "helllo", true}, + {"star quantifier partial should fail", "hel*o", "say hello", false}, + {"plus quantifier", "hel+o", "hello", true}, + {"plus quantifier no match", "hel+o", "heo", false}, + {"plus quantifier partial should fail", "hel+o", "say hello", false}, + + {"digit class", "[0-9]+", "12345", true}, + {"digit class no match", "[0-9]+", "abc", false}, + {"digit class partial should fail", "[0-9]+", "abc123", false}, + {"letter class", "[a-z]+", "hello", true}, + {"letter class no match", "[a-z]+", "HELLO", false}, + {"letter class partial should fail", "[a-z]+", "hello123", false}, + {"mixed class", "[a-zA-Z0-9]+", "Hello123", true}, + {"mixed class partial should fail", "[a-zA-Z0-9]+", "Hello123!", false}, + + {"start anchor exact", "^hello$", "hello", true}, + {"start anchor partial should fail", "^hello$", "hello world", false}, + {"start anchor only", "^hello", "hello", true}, + {"start anchor only partial should fail", "^hello", "hello world", false}, + {"end anchor only", "world$", "world", true}, + {"end anchor only partial should fail", "world$", "hello world", false}, + + {"email pattern valid", `^[a-zA-Z0-9._%+-]+@[a-zA-Z0-9.-]+\.[a-zA-Z]{2,}$`, "test@example.com", true}, + {"email pattern invalid", `^[a-zA-Z0-9._%+-]+@[a-zA-Z0-9.-]+\.[a-zA-Z]{2,}$`, "invalid-email", false}, + {"email pattern without anchors", `[a-zA-Z0-9._%+-]+@[a-zA-Z0-9.-]+\.[a-zA-Z]{2,}`, "test@example.com", true}, + {"phone pattern valid", `^\+?[1-9]\d{1,14}$`, "+1234567890", true}, + {"phone pattern invalid", `^\+?[1-9]\d{1,14}$`, "abc123", false}, + + {"group match", "(hello|world)", "hello", true}, + {"group match alt", "(hello|world)", "world", true}, + {"group no match", "(hello|world)", "test", false}, + {"group partial should fail", "(hello|world)", "say hello", false}, + + {"escaped dot", `hello\.world`, "hello.world", true}, + {"escaped dot no match", `hello\.world`, "helloXworld", false}, + {"escaped dot partial should fail", `hello\.world`, "say hello.world", false}, + {"escaped plus", `test\+`, "test+", true}, + {"escaped star", `test\*`, "test*", true}, + + {"invalid pattern unclosed paren", "(hello", "hello", false}, + {"invalid pattern unclosed bracket", "[hello", "hello", false}, + {"invalid pattern bad escape", `\`, "test", false}, + {"invalid pattern bad quantifier", "*hello", "hello", true}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + got := RegexMatch(tt.pattern, tt.text) + if got != tt.want { + t.Errorf("RegexMatch(%q, %q) = %v; want %v", tt.pattern, tt.text, got, tt.want) + } + }) + } +} + +func TestEnsureExactRegexMatch(t *testing.T) { + tests := []struct { + name string + pattern string + expected string + }{ + {"no anchors", "hello", "^hello$"}, + {"both anchors", "^hello$", "^hello$"}, + {"start anchor only", "^hello", "^hello$"}, + {"end anchor only", "hello$", "^hello$"}, + {"complex pattern", "[a-z]+", "^[a-z]+$"}, + {"already anchored complex", "^[a-z]+$", "^[a-z]+$"}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + got := ensureExactRegexMatch(tt.pattern) + if got != tt.expected { + t.Errorf("ensureExactRegexMatch(%q) = %q; want %q", tt.pattern, got, tt.expected) + } + }) + } +} + +func TestGetCompiledRegex(t *testing.T) { + cacheMu.Lock() + regexCache = make(map[string]*regexp.Regexp) + cacheMu.Unlock() + + tests := []struct { + name string + pattern string + wantErr bool + }{ + {"valid pattern", "^hello$", false}, + {"valid complex pattern", `^[a-z]+$`, false}, + {"invalid pattern", "(hello", true}, + {"invalid pattern bracket", "[hello", true}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + regex, err := getCompiledRegex(tt.pattern) + + if tt.wantErr { + if err == nil { + t.Errorf("getCompiledRegex(%q) expected error, got nil", tt.pattern) + } + if regex != nil { + t.Errorf("getCompiledRegex(%q) expected nil regex when error, got %v", tt.pattern, regex) + } + } else { + if err != nil { + t.Errorf("getCompiledRegex(%q) unexpected error: %v", tt.pattern, err) + } + if regex == nil { + t.Errorf("getCompiledRegex(%q) expected regex, got nil", tt.pattern) + } + } + }) + } +} + +func TestRegexCaching(t *testing.T) { + cacheMu.Lock() + regexCache = make(map[string]*regexp.Regexp) + cacheMu.Unlock() + + pattern := "^test_pattern_for_caching$" + + regex1, err1 := getCompiledRegex(pattern) + if err1 != nil { + t.Fatalf("First call failed: %v", err1) + } + + cacheMu.RLock() + cachedRegex, exists := regexCache[pattern] + cacheMu.RUnlock() + + if !exists { + t.Error("Pattern should be cached after first call") + } + + if cachedRegex != regex1 { + t.Error("Cached regex should be the same instance as returned") + } + + regex2, err2 := getCompiledRegex(pattern) + if err2 != nil { + t.Fatalf("Second call failed: %v", err2) + } + + if regex1 != regex2 { + t.Error("Second call should return the same cached instance") + } +} + +func TestRegexMatchConcurrency(t *testing.T) { + pattern := "concurrent_test_[0-9]+" + text := "concurrent_test_123" + + cacheMu.Lock() + regexCache = make(map[string]*regexp.Regexp) + cacheMu.Unlock() + + const numGoroutines = 100 + results := make(chan bool, numGoroutines) + + for i := 0; i < numGoroutines; i++ { + go func() { + result := RegexMatch(pattern, text) + results <- result + }() + } + + for i := 0; i < numGoroutines; i++ { + result := <-results + if !result { + t.Error("All concurrent calls should return true") + } + } + + cacheMu.RLock() + if len(regexCache) != 1 { + t.Errorf("Expected 1 cached regex, got %d", len(regexCache)) + } + cacheMu.RUnlock() +} + +func BenchmarkRegexMatch(b *testing.B) { + // 清空缓存 + cacheMu.Lock() + regexCache = make(map[string]*regexp.Regexp) + cacheMu.Unlock() + + pattern := `[a-zA-Z0-9._%+-]+@[a-zA-Z0-9.-]+\.[a-zA-Z]{2,}` + text := "test@example.com" + + b.ResetTimer() + for i := 0; i < b.N; i++ { + RegexMatch(pattern, text) + } +} + +func BenchmarkRegexMatchCached(b *testing.B) { + pattern := `[a-zA-Z0-9._%+-]+@[a-zA-Z0-9.-]+\.[a-zA-Z]{2,}` + text := "test@example.com" + + RegexMatch(pattern, text) + + b.ResetTimer() + for i := 0; i < b.N; i++ { + RegexMatch(pattern, text) + } +} + +func BenchmarkRegexMatchWithoutCache(b *testing.B) { + pattern := `^[a-zA-Z0-9._%+-]+@[a-zA-Z0-9.-]+\.[a-zA-Z]{2,}$` + text := "test@example.com" + + b.ResetTimer() + for i := 0; i < b.N; i++ { + regex, err := regexp.Compile(pattern) + if err != nil { + b.Fatal(err) + } + regex.MatchString(text) + } +} + +func TestRegexMatchUnicode(t *testing.T) { + tests := []struct { + name string + pattern string + text string + want bool + }{ + {"chinese exact match", "你好", "你好", true}, + {"chinese no match - partial", "你好", "你好世界", false}, + {"chinese no match - different", "你好", "世界", false}, + + {"emoji exact match", "😀", "😀", true}, + {"emoji no match - in sentence", "😀", "Hello 😀 World", false}, + {"emoji no match", "😀", "Hello World", false}, + + {"mixed unicode exact", "[\\p{Han}]+", "中文测试", true}, + {"mixed unicode no match", "[\\p{Han}]+", "English test", false}, + {"mixed unicode partial should fail", "[\\p{Han}]+", "English 中文 test", false}, + + {"chinese with punctuation", "你好！", "你好！", true}, + {"chinese with punctuation partial", "你好！", "你好！世界", false}, + {"japanese hiragana", "こんにちは", "こんにちは", true}, + {"japanese hiragana partial", "こんにちは", "こんにちは世界", false}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + got := RegexMatch(tt.pattern, tt.text) + if got != tt.want { + t.Errorf("RegexMatch(%q, %q) = %v; want %v", tt.pattern, tt.text, got, tt.want) + } + }) + } +} + +func TestRegexMatchSpecialCases(t *testing.T) { + tests := []struct { + name string + pattern string + text string + want bool + }{ + {"whitespace exact", "\\s+", " ", true}, + {"whitespace partial should fail", "\\s+", "hello world", false}, + + {"number exact", "\\d+", "12345", true}, + {"number partial should fail", "\\d+", "abc12345xyz", false}, + + {"word boundary", "\\bhello\\b", "hello", true}, + {"word boundary partial should fail", "\\bhello\\b", "say hello world", false}, + + {"line start and end", "^test$", "test", true}, + {"line start and end multiline should fail", "^test$", "test\nmore", false}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + got := RegexMatch(tt.pattern, tt.text) + if got != tt.want { + t.Errorf("RegexMatch(%q, %q) = %v; want %v", tt.pattern, tt.text, got, tt.want) + } + }) + } +}