feat: add comprehensive unit tests for entropy-based routing. Tests c… (#112)

* feat: add comprehensive unit tests for entropy-based routing. Tests cover all 5 key scenarios and probability distribution validation.

Signed-off-by: Huamin Chen <[email protected]>

* review feedback

Signed-off-by: Huamin Chen <[email protected]>

* review feedback

Signed-off-by: Huamin Chen <[email protected]>

* review feedback

Signed-off-by: Huamin Chen <[email protected]>

* review feedback

Signed-off-by: Huamin Chen <[email protected]>

---------

Signed-off-by: Huamin Chen <[email protected]>

Author: rootfs

candle-binding/semantic-router.go

@@ -78,22 +78,42 @@ typedef struct {
     float confidence;
 } ClassificationResult;
 
+// Classification result with full probability distribution structure
+typedef struct {
+    int class;
+    float confidence;
+    float* probabilities;
+    int num_classes;
+} ClassificationResultWithProbs;
+
 // ModernBERT Classification result structure
 typedef struct {
     int class;
     float confidence;
 } ModernBertClassificationResult;
 
+// ModernBERT Classification result with full probability distribution structure
+typedef struct {
+    int class;
+    float confidence;
+    float* probabilities;
+    int num_classes;
+} ModernBertClassificationResultWithProbs;
+
 extern SimilarityResult find_most_similar(const char* query, const char** candidates, int num_candidates, int max_length);
 extern EmbeddingResult get_text_embedding(const char* text, int max_length);
 extern TokenizationResult tokenize_text(const char* text, int max_length);
 extern void free_cstring(char* s);
 extern void free_embedding(float* data, int length);
 extern void free_tokenization_result(TokenizationResult result);
 extern ClassificationResult classify_text(const char* text);
+extern ClassificationResultWithProbs classify_text_with_probabilities(const char* text);
+extern void free_probabilities(float* probabilities, int num_classes);
 extern ClassificationResult classify_pii_text(const char* text);
 extern ClassificationResult classify_jailbreak_text(const char* text);
 extern ModernBertClassificationResult classify_modernbert_text(const char* text);
+extern ModernBertClassificationResultWithProbs classify_modernbert_text_with_probabilities(const char* text);
+extern void free_modernbert_probabilities(float* probabilities, int num_classes);
 extern ModernBertClassificationResult classify_modernbert_pii_text(const char* text);
 extern ModernBertClassificationResult classify_modernbert_jailbreak_text(const char* text);
 */
@@ -137,6 +157,14 @@ type ClassResult struct {
 	Confidence float32 // Confidence score
 }
 
+// ClassResultWithProbs represents the result of a text classification with full probability distribution
+type ClassResultWithProbs struct {
+	Class         int       // Class index
+	Confidence    float32   // Confidence score
+	Probabilities []float32 // Full probability distribution
+	NumClasses    int       // Number of classes
+}
+
 // TokenEntity represents a single detected entity in token classification
 type TokenEntity struct {
 	EntityType string  // Type of entity (e.g., "PERSON", "EMAIL", "PHONE")
@@ -452,6 +480,36 @@ func ClassifyText(text string) (ClassResult, error) {
 	}, nil
 }
 
+// ClassifyTextWithProbabilities classifies the provided text and returns the predicted class, confidence, and full probability distribution
+func ClassifyTextWithProbabilities(text string) (ClassResultWithProbs, error) {
+	cText := C.CString(text)
+	defer C.free(unsafe.Pointer(cText))
+
+	result := C.classify_text_with_probabilities(cText)
+
+	if result.class < 0 {
+		return ClassResultWithProbs{}, fmt.Errorf("failed to classify text with probabilities")
+	}
+
+	// Convert C array to Go slice
+	probabilities := make([]float32, int(result.num_classes))
+	if result.probabilities != nil && result.num_classes > 0 {
+		probsSlice := (*[1 << 30]C.float)(unsafe.Pointer(result.probabilities))[:result.num_classes:result.num_classes]
+		for i, prob := range probsSlice {
+			probabilities[i] = float32(prob)
+		}
+		// Free the C-allocated memory
+		C.free_probabilities(result.probabilities, result.num_classes)
+	}
+
+	return ClassResultWithProbs{
+		Class:         int(result.class),
+		Confidence:    float32(result.confidence),
+		Probabilities: probabilities,
+		NumClasses:    int(result.num_classes),
+	}, nil
+}
+
 // ClassifyPIIText classifies the provided text for PII detection and returns the predicted class and confidence
 func ClassifyPIIText(text string) (ClassResult, error) {
 	cText := C.CString(text)
@@ -599,6 +657,36 @@ func ClassifyModernBertText(text string) (ClassResult, error) {
 	}, nil
 }
 
+// ClassifyModernBertTextWithProbabilities classifies the provided text using ModernBERT and returns the predicted class, confidence, and full probability distribution
+func ClassifyModernBertTextWithProbabilities(text string) (ClassResultWithProbs, error) {
+	cText := C.CString(text)
+	defer C.free(unsafe.Pointer(cText))
+
+	result := C.classify_modernbert_text_with_probabilities(cText)
+
+	if result.class < 0 {
+		return ClassResultWithProbs{}, fmt.Errorf("failed to classify text with probabilities using ModernBERT")
+	}
+
+	// Convert C array to Go slice
+	probabilities := make([]float32, int(result.num_classes))
+	if result.probabilities != nil && result.num_classes > 0 {
+		probsSlice := (*[1 << 30]C.float)(unsafe.Pointer(result.probabilities))[:result.num_classes:result.num_classes]
+		for i, prob := range probsSlice {
+			probabilities[i] = float32(prob)
+		}
+		// Free the C-allocated memory
+		C.free_modernbert_probabilities(result.probabilities, result.num_classes)
+	}
+
+	return ClassResultWithProbs{
+		Class:         int(result.class),
+		Confidence:    float32(result.confidence),
+		Probabilities: probabilities,
+		NumClasses:    int(result.num_classes),
+	}, nil
+}
+
 // ClassifyModernBertPIIText classifies the provided text for PII detection using ModernBERT and returns the predicted class and confidence
 func ClassifyModernBertPIIText(text string) (ClassResult, error) {
 	cText := C.CString(text)