add ml_rerank metadata to snapshot response

Author: SentienceDEV

sentience/models.py

@@ -151,6 +151,34 @@ class GridInfo(BaseModel):
     viewport_coverage: float = 0.0  # Ratio of grid area to viewport area (0.0-1.0)
 
 
+class MlRerankTags(BaseModel):
+    """ML rerank tag configuration used for candidate text"""
+
+    repeated: bool
+    sponsored_ish: bool
+    non_sponsored: bool
+    pos: bool
+    occ: bool
+    vocc: bool
+    short: bool
+    action_ish: bool
+    nav_ish: bool
+
+
+class MlRerankInfo(BaseModel):
+    """ML rerank metadata for a snapshot response"""
+
+    enabled: bool
+    applied: bool
+    reason: str | None = None
+    candidate_count: int | None = None
+    top_probability: float | None = None
+    min_confidence: float | None = None
+    is_high_confidence: bool | None = None
+    tags: MlRerankTags | None = None
+    error: str | None = None
+
+
 class Snapshot(BaseModel):
     """Snapshot response from extension"""
 
@@ -170,6 +198,8 @@ class Snapshot(BaseModel):
     # Modal detection fields (from gateway)
     modal_detected: bool | None = None  # True if a modal/overlay grid was detected
     modal_grids: list[GridInfo] | None = None  # Array of GridInfo for detected modal grids
+    # ML rerank metadata (optional)
+    ml_rerank: MlRerankInfo | None = None
 
     def save(self, filepath: str) -> None:
         """Save snapshot as JSON file"""

tests/test_snapshot.py

@@ -196,3 +196,39 @@ def test_element_ml_fields_optional():
     assert element_partial.heuristic_index is None
     assert element_partial.ml_probability == 0.87
     assert element_partial.ml_score is None
+
+
+def test_snapshot_ml_rerank_metadata_optional():
+    """Test snapshot ML rerank metadata model"""
+    from sentience.models import MlRerankInfo, MlRerankTags, Snapshot
+
+    snap = Snapshot(
+        status="success",
+        url="https://example.com",
+        elements=[],
+        ml_rerank=MlRerankInfo(
+            enabled=True,
+            applied=False,
+            reason="low_confidence",
+            candidate_count=25,
+            top_probability=0.42,
+            min_confidence=0.6,
+            is_high_confidence=False,
+            tags=MlRerankTags(
+                repeated=True,
+                sponsored_ish=True,
+                non_sponsored=False,
+                pos=True,
+                occ=True,
+                vocc=False,
+                short=True,
+                action_ish=False,
+                nav_ish=False,
+            ),
+        ),
+    )
+
+    assert snap.ml_rerank is not None
+    assert snap.ml_rerank.enabled is True
+    assert snap.ml_rerank.is_high_confidence is False
+