add unit tests

Author: rbs333

redisvl/extensions/threshold_optimizer/base.py

@@ -2,7 +2,7 @@
 from enum import Enum
 from typing import Any, Callable, Dict, List, TypeVar
 
-from redisvl.extensions.threshold_optimizer.utils import validate_test_dict
+from redisvl.extensions.threshold_optimizer.utils import _validate_test_dict
 
 
 class EvalMetric(Enum):
@@ -47,7 +47,7 @@ def __init__(
             eval_fn: Function to evaluate performance
             opt_fn: Function to perform optimization
         """
-        self.test_data = validate_test_dict(test_dict)
+        self.test_data = _validate_test_dict(test_dict)
         self.optimizable = optimizable
         self.eval_metric = EvalMetric(eval_metric)
         self.opt_fn = opt_fn

redisvl/extensions/threshold_optimizer/cache.py

@@ -9,7 +9,10 @@
     EvalMetric,
 )
 from redisvl.extensions.threshold_optimizer.schema import TestData
-from redisvl.extensions.threshold_optimizer.utils import NULL_RESPONSE_KEY, format_qrels
+from redisvl.extensions.threshold_optimizer.utils import (
+    NULL_RESPONSE_KEY,
+    _format_qrels,
+)
 from redisvl.query import RangeQuery
 
 
@@ -63,7 +66,7 @@ def _grid_search_opt_cache(
         res = cache.index.query(query)
         td.response = res
 
-    qrels = format_qrels(test_data)
+    qrels = _format_qrels(test_data)
 
     for threshold in thresholds:
         score = _eval_cache(test_data, threshold, qrels, eval_metric.value)

redisvl/extensions/threshold_optimizer/router.py

@@ -10,7 +10,10 @@
     EvalMetric,
 )
 from redisvl.extensions.threshold_optimizer.schema import TestData
-from redisvl.extensions.threshold_optimizer.utils import NULL_RESPONSE_KEY, format_qrels
+from redisvl.extensions.threshold_optimizer.utils import (
+    NULL_RESPONSE_KEY,
+    _format_qrels,
+)
 
 
 def _generate_run_router(test_data: List[TestData], router: SemanticRouter) -> Run:
@@ -98,5 +101,5 @@ def __init__(
 
     def optimize(self, **kwargs: Any):
         """Optimize thresholds using the provided optimization function for router case."""
-        qrels = format_qrels(self.test_data)
+        qrels = _format_qrels(self.test_data)
         self.opt_fn(self.optimizable, self.test_data, qrels, self.eval_metric, **kwargs)

redisvl/extensions/threshold_optimizer/schema.py

@@ -1,4 +1,4 @@
-from typing import List
+from typing import List, Optional
 
 from pydantic import BaseModel, Field
 from ulid import ULID
@@ -7,5 +7,5 @@
 class TestData(BaseModel):
     q_id: str = Field(default_factory=lambda: str(ULID()))
     query: str
-    query_match: str | None
+    query_match: Optional[str]
     response: List[dict] = []

redisvl/extensions/threshold_optimizer/utils.py

@@ -7,7 +7,8 @@
 NULL_RESPONSE_KEY = "no_match"
 
 
-def format_qrels(test_data: List[TestData]) -> Qrels:
+def _format_qrels(test_data: List[TestData]) -> Qrels:
+    """Utility function for creating qrels for evaluation with ranx"""
     qrels_dict = {}
 
     for td in test_data:
@@ -20,5 +21,6 @@ def format_qrels(test_data: List[TestData]) -> Qrels:
     return Qrels(qrels_dict)
 
 
-def validate_test_dict(test_dict: List[dict]) -> List[TestData]:
+def _validate_test_dict(test_dict: List[dict]) -> List[TestData]:
+    """Convert/validate test_dict for use in optimizer"""
     return [TestData(**d) for d in test_dict]

tests/unit/test_threshold_optimizer_utility.py

@@ -0,0 +1,74 @@
+import pytest
+from ranx import evaluate
+
+from redisvl.extensions.threshold_optimizer.cache import _generate_run_cache
+from redisvl.extensions.threshold_optimizer.schema import TestData
+from redisvl.extensions.threshold_optimizer.utils import (
+    NULL_RESPONSE_KEY,
+    _format_qrels,
+)
+
+# Note: these tests are not intended to test ranx but to test that our data formatting for the package is correct
+
+
+def test_known_precision_case():
+    """
+    Test case with known precision value.
+
+    Setup:
+    - 2 queries
+    - Query 1 expects doc1, gets doc1 and doc2 (precision 0.5)
+    - Query 2 expects doc3, gets doc3 (precision 1.0)
+    Expected overall precision: 0.75
+    """
+    # Setup test data
+    test_data = [
+        TestData(
+            query="test query 1",
+            query_match="doc1",
+            response=[
+                {"id": "doc1", "vector_distance": 0.2},
+                {"id": "doc2", "vector_distance": 0.3},
+            ],
+        ),
+        TestData(
+            query="test query 2",
+            query_match="doc3",
+            response=[
+                {"id": "doc3", "vector_distance": 0.2},
+                {"id": "doc4", "vector_distance": 0.8},
+            ],
+        ),
+    ]
+
+    # Create qrels (ground truth)
+    qrels = _format_qrels(test_data)
+
+    threshold = 0.4
+    run = _generate_run_cache(test_data, threshold)
+
+    # Calculate precision using ranx
+    precision = evaluate(qrels, run, "precision")
+    assert precision == 0.75  # (0.5 + 1.0) / 2
+
+
+def test_known_precision_with_no_matches():
+    """Test case where some queries have no matches."""
+    test_data = [
+        TestData(
+            query="test query 2",
+            query_match="",  # Expecting no match
+            response=[],
+        ),
+    ]
+
+    # Create qrels
+    qrels = _format_qrels(test_data)
+
+    # Generate run with threshold that excludes all docs for first query
+    threshold = 0.3
+    run = _generate_run_cache(test_data, threshold)
+
+    # Calculate precision
+    precision = evaluate(qrels, run, "precision")
+    assert precision == 1.0  # (0.0 + 1.0) / 2