diff --git a/src/tracing/types.ts b/src/tracing/types.ts index 79448cf..a942c92 100644 --- a/src/tracing/types.ts +++ b/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; diff --git a/src/types.ts b/src/types.ts index 0d5b292..4727338 100644 --- a/src/types.ts +++ b/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 { diff --git a/src/utils/snapshot-event-builder.ts b/src/utils/snapshot-event-builder.ts index 019ef10..ec80751 100644 --- a/src/utils/snapshot-event-builder.ts +++ b/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, diff --git a/src/utils/trace-event-builder.ts b/src/utils/trace-event-builder.ts index c4a0790..ea1998f 100644 --- a/src/utils/trace-event-builder.ts +++ b/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, diff --git a/tests/snapshot.test.ts b/tests/snapshot.test.ts index c5bfe2e..9c86ae0 100644 --- a/tests/snapshot.test.ts +++ b/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); + }); +});