test: add more test cases and refactor SelectBestModelForCategory/SelectBestModelFromList/InitializeJailbreakClassifier for testability (#101)

* test: add more test cases and refactor SelectBestModelForCategory/SelectBestModelFromList for testability

Signed-off-by: Alex Wang <[email protected]>

* test: refactor InitializeJailbreakClassifier and add test cases

Signed-off-by: Alex Wang <[email protected]>

* refactor: remove load aware algorithm

Signed-off-by: Alex Wang <[email protected]>

* refactor: remove quality

Signed-off-by: Alex Wang <[email protected]>

---------

Signed-off-by: Alex Wang <[email protected]>
Co-authored-by: Huamin Chen <[email protected]>

Author: aeft

config/config.yaml

@@ -112,7 +112,6 @@ classifier:
     threshold: 0.7
     use_cpu: true
     pii_mapping_path: "models/pii_classifier_modernbert-base_presidio_token_model/pii_type_mapping.json"
-  load_aware: false
 categories:
 - name: business
   use_reasoning: false

deploy/kubernetes/config.yaml

@@ -78,7 +78,6 @@ classifier:
     threshold: 0.7
     use_cpu: true
     pii_mapping_path: "models/pii_classifier_modernbert-base_presidio_token_model/pii_type_mapping.json"
-  load_aware: false
 categories:
 - name: business
   model_scores:

src/semantic-router/pkg/config/config.go

@@ -34,7 +34,6 @@ type RouterConfig struct {
 			UseCPU         bool    `yaml:"use_cpu"`
 			PIIMappingPath string  `yaml:"pii_mapping_path"`
 		} `yaml:"pii_model"`
-		LoadAware bool `yaml:"load_aware"`
 	} `yaml:"classifier"`
 
 	// Categories for routing queries

src/semantic-router/pkg/config/config_test.go

@@ -60,7 +60,6 @@ classifier:
     use_cpu: true
     use_modernbert: false
     pii_mapping_path: "/path/to/pii.json"
-  load_aware: true
 
 categories:
   - name: "general"
@@ -138,7 +137,6 @@ tools:
 				// Verify classifier config
 				Expect(cfg.Classifier.CategoryModel.ModelID).To(Equal("test-category-model"))
 				Expect(cfg.Classifier.CategoryModel.UseModernBERT).To(BeTrue())
-				Expect(cfg.Classifier.LoadAware).To(BeTrue())
 
 				// Verify categories
 				Expect(cfg.Categories).To(HaveLen(1))

src/semantic-router/pkg/extproc/request_handler.go

@@ -370,9 +370,6 @@ func (r *OpenAIRouter) handleModelRouting(openAIRequest *openai.ChatCompletionNe
 				effortForMetrics := r.getReasoningEffort(categoryName)
 				metrics.RecordReasoningDecision(categoryName, matchedModel, useReasoning, effortForMetrics)
 
-				// Track the model load for the selected model
-				r.Classifier.IncrementModelLoad(matchedModel)
-
 				// Track the model routing change
 				metrics.RecordModelRouting(originalModel, matchedModel)
 

src/semantic-router/pkg/extproc/response_handler.go

@@ -52,7 +52,6 @@ func (r *OpenAIRouter) handleResponseBody(v *ext_proc.ProcessingRequest_Response
 			float64(completionTokens),
 		)
 		metrics.RecordModelCompletionLatency(ctx.RequestModel, completionLatency.Seconds())
-		r.Classifier.DecrementModelLoad(ctx.RequestModel)
 
 		// Compute and record cost if pricing is configured
 		if r.Config != nil {

src/semantic-router/pkg/extproc/router.go

@@ -131,8 +131,7 @@ func NewOpenAIRouter(configPath string) (*OpenAIRouter, error) {
 
 	// Create utility components
 	piiChecker := pii.NewPolicyChecker(cfg, cfg.ModelConfig)
-	modelTTFT := make(map[string]float64) // Empty TTFT map since load balancing is disabled
-	classifier := classification.NewClassifier(cfg, categoryMapping, piiMapping, jailbreakMapping, modelTTFT)
+	classifier := classification.NewClassifier(cfg, categoryMapping, piiMapping, jailbreakMapping)
 
 	// Create global classification service for API access
 	services.NewClassificationService(classifier, cfg)

src/semantic-router/pkg/extproc/security_test.go

@@ -52,7 +52,7 @@ var _ = Describe("Security Checks", func() {
 				},
 			}
 			router.PIIChecker = pii.NewPolicyChecker(cfg, cfg.ModelConfig)
-			router.Classifier = classification.NewClassifier(cfg, router.Classifier.CategoryMapping, router.Classifier.PIIMapping, nil, router.Classifier.ModelTTFT)
+			router.Classifier = classification.NewClassifier(cfg, router.Classifier.CategoryMapping, router.Classifier.PIIMapping, nil)
 		})
 
 		It("should allow requests with no PII", func() {
@@ -97,7 +97,7 @@ var _ = Describe("Security Checks", func() {
 			piiMapping, err := classification.LoadPIIMapping(cfg.Classifier.PIIModel.PIIMappingPath)
 			Expect(err).NotTo(HaveOccurred())
 
-			router.Classifier = classification.NewClassifier(cfg, router.Classifier.CategoryMapping, piiMapping, nil, router.Classifier.ModelTTFT)
+			router.Classifier = classification.NewClassifier(cfg, router.Classifier.CategoryMapping, piiMapping, nil)
 		})
 
 		Describe("ClassifyPII method", func() {
@@ -339,7 +339,7 @@ var _ = Describe("Security Checks", func() {
 			piiMapping, err := classification.LoadPIIMapping(cfg.Classifier.PIIModel.PIIMappingPath)
 			Expect(err).NotTo(HaveOccurred())
 
-			router.Classifier = classification.NewClassifier(cfg, router.Classifier.CategoryMapping, piiMapping, nil, router.Classifier.ModelTTFT)
+			router.Classifier = classification.NewClassifier(cfg, router.Classifier.CategoryMapping, piiMapping, nil)
 		})
 
 		Describe("Error handling and edge cases", func() {
@@ -524,7 +524,7 @@ var _ = Describe("Security Checks", func() {
 				IdxToLabel: map[string]string{"0": "benign", "1": "jailbreak"},
 			}
 
-			router.Classifier = classification.NewClassifier(cfg, router.Classifier.CategoryMapping, router.Classifier.PIIMapping, jailbreakMapping, router.Classifier.ModelTTFT)
+			router.Classifier = classification.NewClassifier(cfg, router.Classifier.CategoryMapping, router.Classifier.PIIMapping, jailbreakMapping)
 		})
 
 		It("should process potential jailbreak attempts", func() {

src/semantic-router/pkg/extproc/test_utils_test.go

@@ -95,7 +95,6 @@ func CreateTestConfig() *config.RouterConfig {
 				UseCPU         bool    `yaml:"use_cpu"`
 				PIIMappingPath string  `yaml:"pii_mapping_path"`
 			} `yaml:"pii_model"`
-			LoadAware bool `yaml:"load_aware"`
 		}{
 			CategoryModel: struct {
 				ModelID             string  `yaml:"model_id"`
@@ -119,7 +118,6 @@ func CreateTestConfig() *config.RouterConfig {
 				UseCPU:         true,
 				PIIMappingPath: "../../../../models/pii_classifier_modernbert-base_presidio_token_model/pii_type_mapping.json",
 			},
-			LoadAware: true,
 		},
 		Categories: []config.Category{
 			{
@@ -220,11 +218,7 @@ func CreateTestRouter(cfg *config.RouterConfig) (*extproc.OpenAIRouter, error) {
 	toolsDatabase := tools.NewToolsDatabase(toolsOptions)
 
 	// Create classifier
-	modelTTFT := map[string]float64{
-		"model-a": 2.5,
-		"model-b": 1.8,
-	}
-	classifier := classification.NewClassifier(cfg, categoryMapping, piiMapping, nil, modelTTFT)
+	classifier := classification.NewClassifier(cfg, categoryMapping, piiMapping, nil)
 
 	// Create PII checker
 	piiChecker := pii.NewPolicyChecker(cfg, cfg.ModelConfig)