@@ -52,7 +52,7 @@ var _ = Describe("Security Checks", func() {
5252 },
5353 }
5454 router .PIIChecker = pii .NewPolicyChecker (cfg , cfg .ModelConfig )
55- router .Classifier = classification .NewClassifier (cfg , router .Classifier .CategoryMapping , router .Classifier .PIIMapping , nil , router . Classifier . ModelTTFT )
55+ router .Classifier = classification .NewClassifier (cfg , router .Classifier .CategoryMapping , router .Classifier .PIIMapping , nil )
5656 })
5757
5858 It ("should allow requests with no PII" , func () {
@@ -97,7 +97,7 @@ var _ = Describe("Security Checks", func() {
9797 piiMapping , err := classification .LoadPIIMapping (cfg .Classifier .PIIModel .PIIMappingPath )
9898 Expect (err ).NotTo (HaveOccurred ())
9999
100- router .Classifier = classification .NewClassifier (cfg , router .Classifier .CategoryMapping , piiMapping , nil , router . Classifier . ModelTTFT )
100+ router .Classifier = classification .NewClassifier (cfg , router .Classifier .CategoryMapping , piiMapping , nil )
101101 })
102102
103103 Describe ("ClassifyPII method" , func () {
@@ -339,7 +339,7 @@ var _ = Describe("Security Checks", func() {
339339 piiMapping , err := classification .LoadPIIMapping (cfg .Classifier .PIIModel .PIIMappingPath )
340340 Expect (err ).NotTo (HaveOccurred ())
341341
342- router .Classifier = classification .NewClassifier (cfg , router .Classifier .CategoryMapping , piiMapping , nil , router . Classifier . ModelTTFT )
342+ router .Classifier = classification .NewClassifier (cfg , router .Classifier .CategoryMapping , piiMapping , nil )
343343 })
344344
345345 Describe ("Error handling and edge cases" , func () {
@@ -524,7 +524,7 @@ var _ = Describe("Security Checks", func() {
524524 IdxToLabel : map [string ]string {"0" : "benign" , "1" : "jailbreak" },
525525 }
526526
527- router .Classifier = classification .NewClassifier (cfg , router .Classifier .CategoryMapping , router .Classifier .PIIMapping , jailbreakMapping , router . Classifier . ModelTTFT )
527+ router .Classifier = classification .NewClassifier (cfg , router .Classifier .CategoryMapping , router .Classifier .PIIMapping , jailbreakMapping )
528528 })
529529
530530 It ("should process potential jailbreak attempts" , func () {
0 commit comments