Merge pull request #166 from SentienceAPI/new_fields

new fields with ml_rerank metadata

Author: rcholic

src/tracing/types.ts

@@ -39,7 +39,7 @@ export interface TraceElement {
   in_viewport?: boolean;
   is_occluded?: boolean;
   z_index?: number;
-  rerank_index?: number;
+  fused_rank_index?: number;
   heuristic_index?: number;
   ml_probability?: number;
   ml_score?: number;

src/types.ts

@@ -33,7 +33,7 @@ export interface Element {
   z_index: number;
 
   // ML reranking metadata (optional - can be absent or null)
-  rerank_index?: number; // 0-based, The rank after ML reranking
+  fused_rank_index?: number; // 0-based, The rank after ML reranking
   heuristic_index?: number; // 0-based, Where it would have been without ML
   ml_probability?: number; // Confidence score from ONNX model (0.0 - 1.0)
   ml_score?: number; // Raw logit score (optional, for debugging)
@@ -143,6 +143,30 @@ export interface GridInfo {
   viewport_coverage?: number;
 }
 
+export interface MlRerankTags {
+  repeated: boolean;
+  sponsored_ish: boolean;
+  non_sponsored: boolean;
+  pos: boolean;
+  occ: boolean;
+  vocc: boolean;
+  short: boolean;
+  action_ish: boolean;
+  nav_ish: boolean;
+}
+
+export interface MlRerankInfo {
+  enabled: boolean;
+  applied: boolean;
+  reason?: string | null;
+  candidate_count?: number | null;
+  top_probability?: number | null;
+  min_confidence?: number | null;
+  is_high_confidence?: boolean | null;
+  tags?: MlRerankTags | null;
+  error?: string | null;
+}
+
 export interface Snapshot {
   status: 'success' | 'error';
   timestamp?: string;
@@ -162,6 +186,8 @@ export interface Snapshot {
   modal_detected?: boolean;
   /** Array of GridInfo for detected modal grids */
   modal_grids?: GridInfo[];
+  /** ML rerank metadata (optional) */
+  ml_rerank?: MlRerankInfo;
 }
 
 export interface StepHookContext {

src/utils/snapshot-event-builder.ts

@@ -47,7 +47,7 @@ export class SnapshotEventBuilder {
         in_viewport: el.in_viewport,
         is_occluded: el.is_occluded,
         z_index: el.z_index,
-        rerank_index: el.rerank_index,
+        fused_rank_index: el.fused_rank_index,
         heuristic_index: el.heuristic_index,
         ml_probability: el.ml_probability,
         ml_score: el.ml_score,

src/utils/trace-event-builder.ts

@@ -219,7 +219,7 @@ export class TraceEventBuilder {
         in_viewport: el.in_viewport,
         is_occluded: el.is_occluded,
         z_index: el.z_index,
-        rerank_index: el.rerank_index,
+        fused_rank_index: el.fused_rank_index,
         heuristic_index: el.heuristic_index,
         ml_probability: el.ml_probability,
         ml_score: el.ml_score,
@@ -429,7 +429,7 @@ export class TraceEventBuilder {
           in_viewport: el.in_viewport,
           is_occluded: el.is_occluded,
           z_index: el.z_index,
-          rerank_index: el.rerank_index,
+          fused_rank_index: el.fused_rank_index,
           heuristic_index: el.heuristic_index,
           ml_probability: el.ml_probability,
           ml_score: el.ml_score,

tests/snapshot.test.ts

@@ -95,7 +95,7 @@ describe('Element ML Fields', () => {
     };
 
     expect(element.id).toBe(1);
-    expect(element).not.toHaveProperty('rerank_index');
+    expect(element).not.toHaveProperty('fused_rank_index');
     expect(element).not.toHaveProperty('heuristic_index');
     expect(element).not.toHaveProperty('ml_probability');
     expect(element).not.toHaveProperty('ml_score');
@@ -112,13 +112,13 @@ describe('Element ML Fields', () => {
       in_viewport: true,
       is_occluded: false,
       z_index: 1,
-      rerank_index: 0,
+      fused_rank_index: 0,
       heuristic_index: 5,
       ml_probability: 0.95,
       ml_score: 2.34,
     };
 
-    expect(element.rerank_index).toBe(0);
+    expect(element.fused_rank_index).toBe(0);
     expect(element.heuristic_index).toBe(5);
     expect(element.ml_probability).toBe(0.95);
     expect(element.ml_score).toBe(2.34);
@@ -135,13 +135,46 @@ describe('Element ML Fields', () => {
       in_viewport: true,
       is_occluded: false,
       z_index: 0,
-      rerank_index: 1,
+      fused_rank_index: 1,
       ml_probability: 0.87,
     };
 
-    expect(element.rerank_index).toBe(1);
+    expect(element.fused_rank_index).toBe(1);
     expect(element).not.toHaveProperty('heuristic_index');
     expect(element.ml_probability).toBe(0.87);
     expect(element).not.toHaveProperty('ml_score');
   });
 });
+
+describe('Snapshot ML Rerank Metadata', () => {
+  it('should accept snapshot with ml_rerank metadata', () => {
+    const snap = {
+      status: 'success' as const,
+      url: 'https://example.com',
+      elements: [],
+      ml_rerank: {
+        enabled: true,
+        applied: false,
+        reason: 'low_confidence',
+        candidate_count: 25,
+        top_probability: 0.42,
+        min_confidence: 0.6,
+        is_high_confidence: false,
+        tags: {
+          repeated: true,
+          sponsored_ish: true,
+          non_sponsored: false,
+          pos: true,
+          occ: true,
+          vocc: false,
+          short: true,
+          action_ish: false,
+          nav_ish: false,
+        },
+      },
+    };
+
+    expect(snap.ml_rerank.enabled).toBe(true);
+    expect(snap.ml_rerank.is_high_confidence).toBe(false);
+  });
+});