From 7b1fb5476be1f3998189cc8826836256b5c97f44 Mon Sep 17 00:00:00 2001
From: rcholic <ivytony@gmail.com>
Date: Tue, 30 Dec 2025 16:21:02 -0800
Subject: [PATCH] support for reranker

---
 package.json           |  2 +-
 src/types.ts           |  6 ++++
 tests/snapshot.test.ts | 66 ++++++++++++++++++++++++++++++++++++++++++
 3 files changed, 73 insertions(+), 1 deletion(-)

diff --git a/package.json b/package.json
index b9101733..ab0aa74e 100644
--- a/package.json
+++ b/package.json
@@ -1,6 +1,6 @@
 {
   "name": "sentienceapi",
-  "version": "0.90.12",
+  "version": "0.90.14",
   "description": "TypeScript SDK for Sentience AI Agent Browser Automation",
   "main": "dist/index.js",
   "types": "dist/index.d.ts",
diff --git a/src/types.ts b/src/types.ts
index dab36c80..3313cebf 100644
--- a/src/types.ts
+++ b/src/types.ts
@@ -30,6 +30,12 @@ export interface Element {
   in_viewport: boolean;
   is_occluded: boolean;
   z_index: number;
+
+  // ML reranking metadata (optional - can be absent or null)
+  rerank_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)
 }
 
 export interface Snapshot {
diff --git a/tests/snapshot.test.ts b/tests/snapshot.test.ts
index c5fe8191..2101bbb4 100644
--- a/tests/snapshot.test.ts
+++ b/tests/snapshot.test.ts
@@ -75,3 +75,69 @@ describe('Snapshot', () => {
   }, 60000); // 60 seconds - browser startup can be slow
 });
 
+describe('Element ML Fields', () => {
+  it('should accept elements without ML reranking fields', () => {
+    const element = {
+      id: 1,
+      role: 'button',
+      text: 'Click me',
+      importance: 100,
+      bbox: { x: 10, y: 20, width: 100, height: 50 },
+      visual_cues: { is_primary: true, background_color_name: 'blue', is_clickable: true },
+      in_viewport: true,
+      is_occluded: false,
+      z_index: 0,
+    };
+
+    expect(element.id).toBe(1);
+    expect(element).not.toHaveProperty('rerank_index');
+    expect(element).not.toHaveProperty('heuristic_index');
+    expect(element).not.toHaveProperty('ml_probability');
+    expect(element).not.toHaveProperty('ml_score');
+  });
+
+  it('should accept elements with ML reranking fields', () => {
+    const element = {
+      id: 2,
+      role: 'link',
+      text: 'Learn more',
+      importance: 80,
+      bbox: { x: 15, y: 25, width: 120, height: 40 },
+      visual_cues: { is_primary: false, background_color_name: 'white', is_clickable: true },
+      in_viewport: true,
+      is_occluded: false,
+      z_index: 1,
+      rerank_index: 0,
+      heuristic_index: 5,
+      ml_probability: 0.95,
+      ml_score: 2.34,
+    };
+
+    expect(element.rerank_index).toBe(0);
+    expect(element.heuristic_index).toBe(5);
+    expect(element.ml_probability).toBe(0.95);
+    expect(element.ml_score).toBe(2.34);
+  });
+
+  it('should accept elements with partial ML fields', () => {
+    const element = {
+      id: 3,
+      role: 'textbox',
+      text: null,
+      importance: 60,
+      bbox: { x: 20, y: 30, width: 200, height: 30 },
+      visual_cues: { is_primary: false, background_color_name: null, is_clickable: true },
+      in_viewport: true,
+      is_occluded: false,
+      z_index: 0,
+      rerank_index: 1,
+      ml_probability: 0.87,
+    };
+
+    expect(element.rerank_index).toBe(1);
+    expect(element).not.toHaveProperty('heuristic_index');
+    expect(element.ml_probability).toBe(0.87);
+    expect(element).not.toHaveProperty('ml_score');
+  });
+});
+