Skip to content

Commit 5fbd817

Browse files
ae86zhizhiautopear
ae86zhizhi
authored and
autopear
committed
refactor: Define environment variables as constants for better maintainability
- Add constants for remote tokenizer environment variables in pkg/constants/kv_event_sync.go: - EnvUseRemoteTokenizer for AIBRIX_USE_REMOTE_TOKENIZER - EnvPrefixCacheTokenizerType for AIBRIX_PREFIX_CACHE_TOKENIZER_TYPE - EnvRemoteTokenizerEndpoint for AIBRIX_REMOTE_TOKENIZER_ENDPOINT - Replace all hardcoded environment variable strings with constants throughout: - pkg/kvevent/manager.go - pkg/cache/kv_event_manager_zmq.go - Test files for consistent usage This change improves code maintainability by preventing typos and providing a single source of truth for environment variable names, following the established pattern in the AIBrix codebase. Signed-off-by: ZHENYU <[email protected]>
1 parent 87dabd5 commit 5fbd817

File tree

5 files changed

+39
-27
lines changed

5 files changed

+39
-27
lines changed

‎pkg/cache/kv_event_manager_validation_test.go

Lines changed: 12 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -43,36 +43,36 @@ func TestKVEventManagerValidation(t *testing.T) {
4343
name: "enabled without tokenizer - error",
4444
envVars: map[string]string{
4545
constants.EnvKVEventSyncEnabled: "true",
46-
"AIBRIX_USE_REMOTE_TOKENIZER": "false",
46+
constants.EnvUseRemoteTokenizer: "false",
4747
},
4848
wantError: true,
4949
},
5050
{
5151
name: "enabled without tokenizer type - error",
5252
envVars: map[string]string{
53-
constants.EnvKVEventSyncEnabled: "true",
54-
"AIBRIX_USE_REMOTE_TOKENIZER": "true",
55-
"AIBRIX_PREFIX_CACHE_TOKENIZER_TYPE": "local",
53+
constants.EnvKVEventSyncEnabled: "true",
54+
constants.EnvUseRemoteTokenizer: "true",
55+
constants.EnvPrefixCacheTokenizerType: "local",
5656
},
5757
wantError: true,
5858
},
5959
{
6060
name: "enabled without endpoint - error",
6161
envVars: map[string]string{
62-
constants.EnvKVEventSyncEnabled: "true",
63-
"AIBRIX_USE_REMOTE_TOKENIZER": "true",
64-
"AIBRIX_PREFIX_CACHE_TOKENIZER_TYPE": "remote",
65-
"AIBRIX_REMOTE_TOKENIZER_ENDPOINT": "",
62+
constants.EnvKVEventSyncEnabled: "true",
63+
constants.EnvUseRemoteTokenizer: "true",
64+
constants.EnvPrefixCacheTokenizerType: "remote",
65+
constants.EnvRemoteTokenizerEndpoint: "",
6666
},
6767
wantError: true,
6868
},
6969
{
7070
name: "enabled with all config - no error",
7171
envVars: map[string]string{
72-
constants.EnvKVEventSyncEnabled: "true",
73-
"AIBRIX_USE_REMOTE_TOKENIZER": "true",
74-
"AIBRIX_PREFIX_CACHE_TOKENIZER_TYPE": "remote",
75-
"AIBRIX_REMOTE_TOKENIZER_ENDPOINT": "http://localhost:8080",
72+
constants.EnvKVEventSyncEnabled: "true",
73+
constants.EnvUseRemoteTokenizer: "true",
74+
constants.EnvPrefixCacheTokenizerType: "remote",
75+
constants.EnvRemoteTokenizerEndpoint: "http://localhost:8080",
7676
},
7777
wantError: false,
7878
},

‎pkg/cache/kv_event_manager_zmq.go

Lines changed: 7 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -54,26 +54,26 @@ func validateKVEventConfiguration() error {
5454
}
5555

5656
// If enabled, check requirements
57-
remoteTokenValue := utils.LoadEnv("AIBRIX_USE_REMOTE_TOKENIZER", "false")
57+
remoteTokenValue := utils.LoadEnv(constants.EnvUseRemoteTokenizer, "false")
5858
remoteTokenizerEnabled, err := strconv.ParseBool(remoteTokenValue)
5959
if err != nil {
60-
return fmt.Errorf("invalid boolean value for AIBRIX_USE_REMOTE_TOKENIZER: %q", remoteTokenValue)
60+
return fmt.Errorf("invalid boolean value for %s: %q", constants.EnvUseRemoteTokenizer, remoteTokenValue)
6161
}
6262

6363
if !remoteTokenizerEnabled {
64-
return fmt.Errorf("KV event sync requires remote tokenizer (set AIBRIX_USE_REMOTE_TOKENIZER=true)")
64+
return fmt.Errorf("KV event sync requires remote tokenizer (set %s=true)", constants.EnvUseRemoteTokenizer)
6565
}
6666

6767
// Check tokenizer type
68-
tokenizerType := utils.LoadEnv("AIBRIX_PREFIX_CACHE_TOKENIZER_TYPE", "")
68+
tokenizerType := utils.LoadEnv(constants.EnvPrefixCacheTokenizerType, "")
6969
if tokenizerType != "remote" {
70-
return fmt.Errorf("KV event sync requires AIBRIX_PREFIX_CACHE_TOKENIZER_TYPE=remote (got %q)", tokenizerType)
70+
return fmt.Errorf("KV event sync requires %s=remote (got %q)", constants.EnvPrefixCacheTokenizerType, tokenizerType)
7171
}
7272

7373
// Check remote tokenizer endpoint
74-
endpoint := utils.LoadEnv("AIBRIX_REMOTE_TOKENIZER_ENDPOINT", "")
74+
endpoint := utils.LoadEnv(constants.EnvRemoteTokenizerEndpoint, "")
7575
if endpoint == "" {
76-
return fmt.Errorf("KV event sync requires AIBRIX_REMOTE_TOKENIZER_ENDPOINT to be set")
76+
return fmt.Errorf("KV event sync requires %s to be set", constants.EnvRemoteTokenizerEndpoint)
7777
}
7878

7979
return nil

‎pkg/constants/kv_event_sync.go

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -48,6 +48,18 @@ const (
4848
// EnvPrefixCacheMetricsEnabled enables prefix cache metrics
4949
// Added as part of KV Event Sync to control metrics registration
5050
EnvPrefixCacheMetricsEnabled = "AIBRIX_PREFIX_CACHE_METRICS_ENABLED"
51+
52+
// EnvUseRemoteTokenizer enables remote tokenizer usage
53+
// When true, uses remote tokenizer service instead of local tokenization
54+
EnvUseRemoteTokenizer = "AIBRIX_USE_REMOTE_TOKENIZER"
55+
56+
// EnvPrefixCacheTokenizerType specifies the tokenizer type for prefix cache
57+
// Options: "character", "tiktoken", "remote"
58+
EnvPrefixCacheTokenizerType = "AIBRIX_PREFIX_CACHE_TOKENIZER_TYPE"
59+
60+
// EnvRemoteTokenizerEndpoint specifies the remote tokenizer service endpoint
61+
// Format: "http://service:port" - required when using remote tokenizer
62+
EnvRemoteTokenizerEndpoint = "AIBRIX_REMOTE_TOKENIZER_ENDPOINT"
5163
)
5264

5365
// Helper functions for KV Event Sync labels

‎pkg/kvevent/manager.go

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -307,10 +307,10 @@ func validateConfiguration() bool {
307307
}
308308

309309
// Check remote tokenizer
310-
remoteTokenValue := utils.LoadEnv("AIBRIX_USE_REMOTE_TOKENIZER", "false")
310+
remoteTokenValue := utils.LoadEnv(constants.EnvUseRemoteTokenizer, "false")
311311
remoteTokenizerEnabled, err := strconv.ParseBool(remoteTokenValue)
312312
if err != nil {
313-
klog.Warningf("Invalid boolean value for AIBRIX_USE_REMOTE_TOKENIZER: %q. Defaulting to false.", remoteTokenValue)
313+
klog.Warningf("Invalid boolean value for %s: %q. Defaulting to false.", constants.EnvUseRemoteTokenizer, remoteTokenValue)
314314
remoteTokenizerEnabled = false
315315
}
316316

‎pkg/kvevent/manager_comprehensive_test.go

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -75,15 +75,15 @@ func TestManagerConfiguration(t *testing.T) {
7575
name: "sync enabled with remote tokenizer",
7676
envVars: map[string]string{
7777
constants.EnvKVEventSyncEnabled: "true",
78-
"AIBRIX_USE_REMOTE_TOKENIZER": "true",
78+
constants.EnvUseRemoteTokenizer: "true",
7979
},
8080
expectEnabled: true,
8181
},
8282
{
8383
name: "sync enabled without remote tokenizer",
8484
envVars: map[string]string{
8585
constants.EnvKVEventSyncEnabled: "true",
86-
"AIBRIX_USE_REMOTE_TOKENIZER": "false",
86+
constants.EnvUseRemoteTokenizer: "false",
8787
},
8888
expectEnabled: false,
8989
},
@@ -117,7 +117,7 @@ func TestManagerConfiguration(t *testing.T) {
117117
func TestManagerPodLifecycle(t *testing.T) {
118118
// Enable KV sync
119119
t.Setenv(constants.EnvKVEventSyncEnabled, "true")
120-
t.Setenv("AIBRIX_USE_REMOTE_TOKENIZER", "true")
120+
t.Setenv(constants.EnvUseRemoteTokenizer, "true")
121121

122122
// Track method calls
123123
syncIndexer := &mockSyncIndexerWithCallbacks{}
@@ -188,7 +188,7 @@ func TestManagerPodLifecycle(t *testing.T) {
188188
// Test sync indexer initialization retry
189189
func TestManagerSyncIndexerRetry(t *testing.T) {
190190
t.Setenv(constants.EnvKVEventSyncEnabled, "true")
191-
t.Setenv("AIBRIX_USE_REMOTE_TOKENIZER", "true")
191+
t.Setenv(constants.EnvUseRemoteTokenizer, "true")
192192

193193
podProvider := &mockPodProvider{pods: make(map[string]*kvevent.PodInfo)}
194194

@@ -226,7 +226,7 @@ func TestManagerSyncIndexerRetry(t *testing.T) {
226226
func TestManagerPodSubscriptionCriteria(t *testing.T) {
227227
// Enable KV sync
228228
t.Setenv(constants.EnvKVEventSyncEnabled, "true")
229-
t.Setenv("AIBRIX_USE_REMOTE_TOKENIZER", "true")
229+
t.Setenv(constants.EnvUseRemoteTokenizer, "true")
230230

231231
tests := []struct {
232232
name string
@@ -354,7 +354,7 @@ func TestManagerPodSubscriptionCriteria(t *testing.T) {
354354
func TestManagerConcurrentOperations(t *testing.T) {
355355
// Enable KV sync
356356
t.Setenv(constants.EnvKVEventSyncEnabled, "true")
357-
t.Setenv("AIBRIX_USE_REMOTE_TOKENIZER", "true")
357+
t.Setenv(constants.EnvUseRemoteTokenizer, "true")
358358

359359
podProvider := &mockPodProvider{
360360
pods: make(map[string]*kvevent.PodInfo),

0 commit comments

Comments
 (0)