Skip to content

Commit dc5751d

Browse files
gidimGideon Mendelsclaudeandriidudar
authored
[NA] [FE] Maintain state of annotation view when moving between traces (#3703)
* Add persistent tree state and scroll position for annotation queue JSON viewer Implement full state persistence across trace navigation in annotation queues. When users expand/collapse nodes and scroll to specific positions in the JSON tree viewer, these states now persist when navigating to the next/previous trace. Key changes: - Created AnnotationTreeStateContext to manage tree expansion and scroll state across trace navigation - Modified JsonKeyValueTable to support controlled mode where parent component manages expansion state - Updated SyntaxHighlighter and MarkdownHighlighter to pass through controlled state props - Enhanced TraceDataViewer with scroll position tracking and restoration - Wrapped AnnotationView with AnnotationTreeStateProvider to enable state sharing across input/output/metadata sections Benefits: - Improved annotation workflow efficiency by maintaining visual context - Separate state management for input, output, and metadata sections - Scroll position restoration ensures users see the same part of the tree - Only applies to annotation queues, does not affect other parts of the app 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude <[email protected]> * Add key-based tree expansion matching for robust state persistence Enhance the JSON tree viewer to use key-based paths instead of index-based paths for tracking expansion state. This makes the expansion state robust across different tree structures when navigating between traces. How it works: - Each node in the tree now has a keyPath (e.g., "metadata.user.name") - When user expands/collapses nodes, we track which key paths are expanded - When data changes (new trace), we map stored key paths to new structure - If a key exists in the new trace, it stays expanded at its new position Benefits: - Expansion state persists even if keys appear in different order - Works across traces with slightly different structures - Example: If "metadata.session" is expanded in trace 1, it stays expanded in trace 2 even if other keys are added/removed Technical details: - Added keyPath tracking to JsonRowData interface - Created buildPathMaps() to build bidirectional index↔key mappings - Added useEffect hooks to track and restore expansion state in controlled mode - Only applies to controlled mode (annotation queues) 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude <[email protected]> * Update .gitignore to exclude Python venv and test scripts * Revision 3: Remove unnecessary metadata tree state tracking Metadata does not have prettifyConfig and only displays as flat text (YAML/JSON) via CodeMirrorHighlighter. It never uses JsonKeyValueTable or expandable tree components, so tree expansion state tracking is unnecessary and has been removed. Changes: - Removed metadata from AnnotationTreeState interface - Removed metadata scroll position tracking - Removed metadata tree expansion handlers - Added explanatory comments about why metadata is excluded * Refactor TraceDataViewer to streamline scroll position management - Removed unnecessary scroll refs and handlers, simplifying the component. - Introduced new scroll position handlers for input and output sections. - Updated SyntaxHighlighter components to support scroll position tracking and restoration. - Enhanced overall code readability and maintainability. This refactor improves the user experience by ensuring scroll positions are accurately managed and restored during trace navigation. * fix scroll restoration for codemirror view --------- Co-authored-by: Gideon Mendels <[email protected]> Co-authored-by: Claude <[email protected]> Co-authored-by: andrii.dudar <[email protected]>
1 parent 2120ccd commit dc5751d

File tree

9 files changed

+625
-104
lines changed

9 files changed

+625
-104
lines changed

‎apps/opik-frontend/.gitignore‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3,3 +3,10 @@ node_modules/
33
/playwright-report/
44
/blob-report/
55
/playwright/.cache/
6+
# Python virtual environments
7+
.venv/
8+
venv/
9+
env/
10+
11+
# Test scripts
12+
test_nested_json_traces.py

‎apps/opik-frontend/src/components/pages/SMEFlowPage/AnnotationView/AnnotationTreeStateContext.tsx‎

Lines changed: 141 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,141 @@
1+
import React, {
2+
createContext,
3+
useContext,
4+
useState,
5+
useCallback,
6+
ReactNode,
7+
} from "react";
8+
import { ExpandedState } from "@tanstack/react-table";
9+
10+
/**
11+
* State for a single section (input/output) in the annotation viewer
12+
* Note: expanded state uses index-based paths but JsonKeyValueTable will map these
13+
* to key-based paths internally for robustness across different tree structures
14+
*/
15+
interface SectionState {
16+
expanded: ExpandedState;
17+
scrollTop: number;
18+
}
19+
20+
/**
21+
* Complete tree state for all sections in the annotation viewer
22+
* Note: metadata is excluded because it doesn't have prettifyConfig and only displays
23+
* as flat text (YAML/JSON) via CodeMirrorHighlighter, never as an expandable tree
24+
*/
25+
interface AnnotationTreeState {
26+
input: SectionState;
27+
output: SectionState;
28+
}
29+
30+
interface AnnotationTreeStateContextValue {
31+
state: AnnotationTreeState;
32+
updateExpanded: (
33+
section: keyof AnnotationTreeState,
34+
updaterOrValue: ExpandedState | ((old: ExpandedState) => ExpandedState),
35+
) => void;
36+
updateScrollTop: (
37+
section: keyof AnnotationTreeState,
38+
scrollTop: number,
39+
) => void;
40+
getScrollTop: (section: keyof AnnotationTreeState) => number;
41+
}
42+
43+
const AnnotationTreeStateContext = createContext<
44+
AnnotationTreeStateContextValue | undefined
45+
>(undefined);
46+
47+
const createDefaultSectionState = (): SectionState => ({
48+
expanded: {},
49+
scrollTop: 0,
50+
});
51+
52+
const createDefaultState = (): AnnotationTreeState => ({
53+
input: createDefaultSectionState(),
54+
output: createDefaultSectionState(),
55+
});
56+
57+
interface AnnotationTreeStateProviderProps {
58+
children: ReactNode;
59+
}
60+
61+
/**
62+
* Provider that maintains tree expansion state and scroll positions across
63+
* trace navigation in annotation queues.
64+
*
65+
* This allows users to expand/collapse nodes and scroll to specific positions,
66+
* and those states persist when navigating to the next/previous trace.
67+
*/
68+
export const AnnotationTreeStateProvider: React.FC<
69+
AnnotationTreeStateProviderProps
70+
> = ({ children }) => {
71+
const [state, setState] = useState<AnnotationTreeState>(createDefaultState);
72+
73+
const updateExpanded = useCallback(
74+
(
75+
section: keyof AnnotationTreeState,
76+
updaterOrValue: ExpandedState | ((old: ExpandedState) => ExpandedState),
77+
) => {
78+
setState((prev) => {
79+
const newExpanded =
80+
typeof updaterOrValue === "function"
81+
? updaterOrValue(prev[section].expanded)
82+
: updaterOrValue;
83+
84+
return {
85+
...prev,
86+
[section]: {
87+
...prev[section],
88+
expanded: newExpanded,
89+
},
90+
};
91+
});
92+
},
93+
[],
94+
);
95+
96+
const updateScrollTop = useCallback(
97+
(section: keyof AnnotationTreeState, scrollTop: number) => {
98+
setState((prev) => ({
99+
...prev,
100+
[section]: {
101+
...prev[section],
102+
scrollTop,
103+
},
104+
}));
105+
},
106+
[],
107+
);
108+
109+
const getScrollTop = useCallback(
110+
(section: keyof AnnotationTreeState) => {
111+
return state[section].scrollTop;
112+
},
113+
[state],
114+
);
115+
116+
const value: AnnotationTreeStateContextValue = {
117+
state,
118+
updateExpanded,
119+
updateScrollTop,
120+
getScrollTop,
121+
};
122+
123+
return (
124+
<AnnotationTreeStateContext.Provider value={value}>
125+
{children}
126+
</AnnotationTreeStateContext.Provider>
127+
);
128+
};
129+
130+
/**
131+
* Hook to access annotation tree state context
132+
*/
133+
export const useAnnotationTreeState = () => {
134+
const context = useContext(AnnotationTreeStateContext);
135+
if (!context) {
136+
throw new Error(
137+
"useAnnotationTreeState must be used within AnnotationTreeStateProvider",
138+
);
139+
}
140+
return context;
141+
};

‎apps/opik-frontend/src/components/pages/SMEFlowPage/AnnotationView/AnnotationView.tsx‎

Lines changed: 84 additions & 79 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,7 @@ import { useSMEFlow } from "../SMEFlowContext";
1515
import { ANNOTATION_QUEUE_SCOPE } from "@/types/annotation-queues";
1616
import ThreadDataViewer from "./ThreadDataViewer";
1717
import { SME_ACTION, SME_HOTKEYS } from "../hotkeys";
18+
import { AnnotationTreeStateProvider } from "./AnnotationTreeStateContext";
1819

1920
interface AnnotationViewProps {
2021
header: React.ReactNode;
@@ -73,90 +74,94 @@ const AnnotationView: React.FunctionComponent<AnnotationViewProps> = ({
7374
const isThread = annotationQueue?.scope === ANNOTATION_QUEUE_SCOPE.THREAD;
7475

7576
return (
76-
<SMEFlowLayout
77-
header={header}
78-
footer={
79-
<>
80-
<ReturnToAnnotationQueueButton />
81-
<div className="flex items-center gap-2">
82-
<div className="comet-body-s flex items-center text-light-slate">
83-
{currentIndex + 1} of {queueItems.length}
84-
</div>
85-
<TooltipWrapper
86-
content="Previous item"
87-
hotkeys={[SME_HOTKEYS[SME_ACTION.PREVIOUS].display]}
88-
>
89-
<Button
90-
variant="outline"
91-
onClick={handlePrevious}
92-
disabled={isFirstItem}
77+
<AnnotationTreeStateProvider>
78+
<SMEFlowLayout
79+
header={header}
80+
footer={
81+
<>
82+
<ReturnToAnnotationQueueButton />
83+
<div className="flex items-center gap-2">
84+
<div className="comet-body-s flex items-center text-light-slate">
85+
{currentIndex + 1} of {queueItems.length}
86+
</div>
87+
<TooltipWrapper
88+
content="Previous item"
89+
hotkeys={[SME_HOTKEYS[SME_ACTION.PREVIOUS].display]}
9390
>
94-
<ChevronLeft className="mr-2 size-4" />
95-
Previous
96-
<HotkeyDisplay
97-
hotkey={SME_HOTKEYS[SME_ACTION.PREVIOUS].display}
91+
<Button
9892
variant="outline"
99-
size="sm"
100-
className="ml-2"
101-
/>
102-
</Button>
103-
</TooltipWrapper>
104-
<TooltipWrapper
105-
content="Next item"
106-
hotkeys={[SME_HOTKEYS[SME_ACTION.NEXT].display]}
107-
>
108-
<Button
109-
variant="outline"
110-
onClick={handleNext}
111-
disabled={isLastItem}
93+
onClick={handlePrevious}
94+
disabled={isFirstItem}
95+
>
96+
<ChevronLeft className="mr-2 size-4" />
97+
Previous
98+
<HotkeyDisplay
99+
hotkey={SME_HOTKEYS[SME_ACTION.PREVIOUS].display}
100+
variant="outline"
101+
size="sm"
102+
className="ml-2"
103+
/>
104+
</Button>
105+
</TooltipWrapper>
106+
<TooltipWrapper
107+
content="Next item"
108+
hotkeys={[SME_HOTKEYS[SME_ACTION.NEXT].display]}
112109
>
113-
<HotkeyDisplay
114-
hotkey={SME_HOTKEYS[SME_ACTION.NEXT].display}
110+
<Button
115111
variant="outline"
116-
size="sm"
117-
className="mr-2"
118-
/>
119-
Next
120-
<ChevronRight className="ml-2 size-4" />
121-
</Button>
122-
</TooltipWrapper>
123-
<TooltipWrapper
124-
content="Submit and continue"
125-
hotkeys={[SME_HOTKEYS[SME_ACTION.DONE].display]}
126-
>
127-
<Button
128-
onClick={handleSubmit}
129-
disabled={!validationState.canSubmit}
112+
onClick={handleNext}
113+
disabled={isLastItem}
114+
>
115+
<HotkeyDisplay
116+
hotkey={SME_HOTKEYS[SME_ACTION.NEXT].display}
117+
variant="outline"
118+
size="sm"
119+
className="mr-2"
120+
/>
121+
Next
122+
<ChevronRight className="ml-2 size-4" />
123+
</Button>
124+
</TooltipWrapper>
125+
<TooltipWrapper
126+
content="Submit and continue"
127+
hotkeys={[SME_HOTKEYS[SME_ACTION.DONE].display]}
130128
>
131-
{isLastUnprocessedItem ? "Submit & Complete" : "Submit + Next"}
132-
<HotkeyDisplay
133-
hotkey={SME_HOTKEYS[SME_ACTION.DONE].display}
134-
size="sm"
135-
className="ml-2"
136-
/>
137-
</Button>
138-
</TooltipWrapper>
139-
</div>
140-
</>
141-
}
142-
>
143-
<div className="flex h-full flex-col">
144-
{validationState.errors.length > 0 && (
145-
<div className="mb-4">
146-
<ValidationAlert errors={validationState.errors} />
147-
</div>
148-
)}
149-
<Card className="flex h-full flex-row items-stretch p-6">
150-
<div className="flex-[2] overflow-y-auto">
151-
{isThread ? <ThreadDataViewer /> : <TraceDataViewer />}
152-
</div>
153-
<Separator orientation="vertical" className="mx-3" />
154-
<div className="flex-[1] overflow-y-auto">
155-
<CommentAndScoreViewer />
156-
</div>
157-
</Card>
158-
</div>
159-
</SMEFlowLayout>
129+
<Button
130+
onClick={handleSubmit}
131+
disabled={!validationState.canSubmit}
132+
>
133+
{isLastUnprocessedItem
134+
? "Submit & Complete"
135+
: "Submit + Next"}
136+
<HotkeyDisplay
137+
hotkey={SME_HOTKEYS[SME_ACTION.DONE].display}
138+
size="sm"
139+
className="ml-2"
140+
/>
141+
</Button>
142+
</TooltipWrapper>
143+
</div>
144+
</>
145+
}
146+
>
147+
<div className="flex h-full flex-col">
148+
{validationState.errors.length > 0 && (
149+
<div className="mb-4">
150+
<ValidationAlert errors={validationState.errors} />
151+
</div>
152+
)}
153+
<Card className="flex h-full flex-row items-stretch p-6">
154+
<div className="flex-[2] overflow-y-auto">
155+
{isThread ? <ThreadDataViewer /> : <TraceDataViewer />}
156+
</div>
157+
<Separator orientation="vertical" className="mx-3" />
158+
<div className="flex-[1] overflow-y-auto">
159+
<CommentAndScoreViewer />
160+
</div>
161+
</Card>
162+
</div>
163+
</SMEFlowLayout>
164+
</AnnotationTreeStateProvider>
160165
);
161166
};
162167

0 commit comments

Comments
 (0)