Skip to content

Commit 43ee299

Browse files
TataKKKLTataKKKL
authored
Merge pull request #69 from PathOnAI/add-d3-simple-search
bfs/dfs v1 working
2 parents 9911362 + 6bc8e3f commit 43ee299

File tree

1 file changed

+285
-45
lines changed

1 file changed

+285
-45
lines changed

‎visual-tree-search-app/components/SimpleSearchVisual.tsx‎

Lines changed: 285 additions & 45 deletions
Original file line numberDiff line numberDiff line change
@@ -1,17 +1,22 @@
1-
import React, { useEffect, useRef } from 'react';
1+
import React, { useEffect, useRef, useState } from 'react';
22
import * as d3 from 'd3';
33
import { useTheme } from 'next-themes';
44

5+
interface TreeNode {
6+
id: number;
7+
parent_id: number | null;
8+
action: string;
9+
description: string | null;
10+
depth?: number;
11+
is_terminal?: boolean;
12+
value?: number;
13+
visits?: number;
14+
feedback?: string;
15+
reward?: number;
16+
}
17+
518
interface Message {
6-
content: string | {
7-
type: string;
8-
tree?: Array<{
9-
id: number;
10-
parent_id: number | null;
11-
action: string;
12-
description: string | null;
13-
}>;
14-
};
19+
content: string;
1520
type: 'incoming' | 'outgoing';
1621
timestamp: string;
1722
}
@@ -23,54 +28,286 @@ interface SimpleSearchVisualProps {
2328
const SimpleSearchVisual: React.FC<SimpleSearchVisualProps> = ({ messages }) => {
2429
const svgRef = useRef<SVGSVGElement>(null);
2530
const containerRef = useRef<HTMLDivElement>(null);
31+
const tooltipRef = useRef<HTMLDivElement | null>(null);
2632
const { theme } = useTheme();
33+
const [selectedNodeId, setSelectedNodeId] = useState<number | null>(null);
34+
const [treeNodes, setTreeNodes] = useState<TreeNode[]>([]);
35+
const [containerWidth, setContainerWidth] = useState<number>(0);
36+
37+
// Set up resize observer to make the visualization responsive
38+
useEffect(() => {
39+
if (!containerRef.current) return;
40+
41+
const resizeObserver = new ResizeObserver(entries => {
42+
for (const entry of entries) {
43+
if (entry.target === containerRef.current) {
44+
const newWidth = entry.contentRect.width;
45+
if (newWidth > 0) {
46+
setContainerWidth(newWidth);
47+
}
48+
}
49+
}
50+
});
51+
52+
resizeObserver.observe(containerRef.current);
53+
54+
return () => {
55+
resizeObserver.disconnect();
56+
};
57+
}, []);
2758

28-
// Simple placeholder visualization
59+
// Cleanup tooltip on unmount
2960
useEffect(() => {
30-
if (!svgRef.current || !messages.length) return;
61+
return () => {
62+
if (tooltipRef.current) {
63+
tooltipRef.current.remove();
64+
}
65+
};
66+
}, []);
67+
68+
// Process messages to extract tree data
69+
useEffect(() => {
70+
if (!messages.length) return;
71+
72+
let updatedTreeNodes: TreeNode[] = [...treeNodes];
73+
let newSelectedNodeId = selectedNodeId;
74+
let hasChanges = false;
75+
76+
messages.forEach(msg => {
77+
try {
78+
const data = JSON.parse(msg.content);
79+
80+
// Handle node selection updates
81+
if (data.type === 'node_selected' && data.node_id !== undefined) {
82+
newSelectedNodeId = data.node_id;
83+
hasChanges = true;
84+
}
85+
86+
// Handle tree structure updates
87+
if (data.type === 'tree_update_node_expansion' && Array.isArray(data.tree)) {
88+
updatedTreeNodes = data.tree;
89+
hasChanges = true;
90+
}
91+
92+
// Handle node evaluation updates
93+
if (data.type === 'tree_update_node_evaluation' && Array.isArray(data.tree)) {
94+
updatedTreeNodes = data.tree;
95+
hasChanges = true;
96+
}
97+
} catch {
98+
// Skip messages that can't be parsed
99+
}
100+
});
101+
102+
if (hasChanges) {
103+
setTreeNodes(updatedTreeNodes);
104+
setSelectedNodeId(newSelectedNodeId);
105+
}
106+
}, [messages, treeNodes, selectedNodeId]);
107+
108+
// Render the tree visualization
109+
useEffect(() => {
110+
if (!svgRef.current || !treeNodes.length) return;
31111

32112
// Clear previous content
33113
const svg = d3.select(svgRef.current);
34114
svg.selectAll("*").remove();
35115

36-
// Create a simple placeholder visualization
116+
// Create or update tooltip
117+
const createTooltip = () => {
118+
// Remove existing tooltip if any
119+
if (tooltipRef.current) {
120+
tooltipRef.current.remove();
121+
}
122+
123+
// Create tooltip
124+
tooltipRef.current = d3.select("body")
125+
.append("div")
126+
.attr("class", "tooltip")
127+
.style("opacity", 0)
128+
.style("position", "absolute")
129+
.style("background-color", theme === 'dark' ? "#1F2937" : "#FFFFFF")
130+
.style("border", `1px solid ${theme === 'dark' ? "#374151" : "#E5E7EB"}`)
131+
.style("border-radius", "4px")
132+
.style("padding", "12px 16px")
133+
.style("font-size", "14px")
134+
.style("color", theme === 'dark' ? "#FFFFFF" : "#111827")
135+
.style("pointer-events", "none")
136+
.style("z-index", "1000")
137+
.style("max-width", "400px")
138+
.style("box-shadow", "0 4px 6px -1px rgba(0, 0, 0, 0.1), 0 2px 4px -1px rgba(0, 0, 0, 0.06)")
139+
.node() as HTMLDivElement;
140+
};
141+
142+
createTooltip();
143+
144+
// Create hierarchical data structure
145+
const stratify = d3.stratify<TreeNode>()
146+
.id(d => d.id.toString())
147+
.parentId(d => d.parent_id !== null ? d.parent_id.toString() : null);
148+
149+
// Try to create hierarchy, fall back to simple layout if it fails
150+
let root;
151+
try {
152+
root = stratify(treeNodes);
153+
} catch (error) {
154+
console.error("Failed to create hierarchy:", error);
155+
return;
156+
}
157+
158+
// Set up dimensions
37159
const width = 400;
38160
const height = 700;
39-
const margin = { top: 20, right: 20, bottom: 20, left: 20 };
161+
const margin = { top: 40, right: 30, bottom: 40, left: 140 };
162+
const innerWidth = width - margin.left - margin.right;
163+
const innerHeight = height - margin.top - margin.bottom;
40164

41-
// Filter tree_update messages
42-
const treeUpdates = messages.filter(msg => {
43-
try {
44-
const data = typeof msg.content === 'string' ? JSON.parse(msg.content) : msg.content;
45-
return data.type === 'tree_update';
46-
} catch {
47-
return false;
48-
}
49-
});
165+
// Create tree layout - horizontal tree
166+
const treeLayout = d3.tree<TreeNode>()
167+
.size([innerHeight, innerWidth]);
168+
169+
// Apply the tree layout
170+
treeLayout(root);
50171

51-
// Draw placeholder circles for each tree update
172+
// Create container group
52173
const g = svg.append("g")
53174
.attr("transform", `translate(${margin.left},${margin.top})`);
54175

55-
g.selectAll("circle")
56-
.data(treeUpdates)
176+
// Add links
177+
g.selectAll(".link")
178+
.data(root.links())
57179
.enter()
58-
.append("circle")
59-
.attr("cx", (d, i) => (i % 3) * 100 + 50)
60-
.attr("cy", (d, i) => Math.floor(i / 3) * 100 + 50)
61-
.attr("r", 20)
62-
.attr("fill", theme === 'dark' ? "#4B5563" : "#9CA3AF")
63-
.attr("stroke", theme === 'dark' ? "#374151" : "#E5E7EB");
64-
65-
// Add placeholder text
66-
g.append("text")
67-
.attr("x", width / 2)
68-
.attr("y", height / 2)
69-
.attr("text-anchor", "middle")
70-
.attr("fill", theme === 'dark' ? "#FFFFFF" : "#111827")
71-
.text(`Tree Updates: ${treeUpdates.length}`);
72-
73-
}, [messages, theme]);
180+
.append("path")
181+
.attr("class", "link")
182+
.attr("d", d => {
183+
const sourceY = d.source.y ?? 0; // Default to 0 if undefined
184+
const sourceX = d.source.x ?? 0; // Default to 0 if undefined
185+
const targetY = d.target.y ?? 0; // Default to 0 if undefined
186+
const targetX = d.target.x ?? 0; // Default to 0 if undefined
187+
return `M${sourceY},${sourceX}C${(sourceY + targetY) / 2},${sourceX} ${(sourceY + targetY) / 2},${targetX} ${targetY},${targetX}`;
188+
})
189+
.attr("fill", "none")
190+
.attr("stroke", theme === 'dark' ? "#9CA3AF" : "#6B7280")
191+
.attr("stroke-width", 1.5)
192+
.attr("stroke-opacity", 0.7);
193+
194+
// Create node groups
195+
const nodes = g.selectAll(".node")
196+
.data(root.descendants())
197+
.enter()
198+
.append("g")
199+
.attr("class", d => `node ${d.children ? "node--internal" : "node--leaf"}`)
200+
.attr("transform", d => `translate(${d.y},${d.x})`);
201+
202+
// Add node circles
203+
nodes.append("circle")
204+
.attr("r", 12)
205+
.attr("fill", d => {
206+
// Selected node (blue)
207+
if (d.data.id === selectedNodeId) {
208+
return theme === 'dark' ? "#3B82F6" : "#60A5FA";
209+
}
210+
211+
// Root node (gray)
212+
if (d.data.parent_id === null) {
213+
return theme === 'dark' ? "#6B7280" : "#D1D5DB";
214+
}
215+
216+
// Action node (default)
217+
return theme === 'dark' ? "#4B5563" : "#E5E7EB";
218+
})
219+
.attr("stroke", d => d.data.id === selectedNodeId
220+
? theme === 'dark' ? "#93C5FD" : "#2563EB"
221+
: theme === 'dark' ? "#374151" : "#D1D5DB")
222+
.attr("stroke-width", d => d.data.id === selectedNodeId ? 3 : 2);
223+
224+
// Add node labels directly on the node circles
225+
nodes.append("text")
226+
.attr("dy", ".35em")
227+
.attr("x", d => d.children ? -18 : 18)
228+
.attr("text-anchor", d => d.children ? "end" : "start")
229+
.text(d => {
230+
// For root node
231+
if (d.data.parent_id === null) return "ROOT";
232+
233+
// Extract action name from action string
234+
if (d.data.action) {
235+
const actionMatch = d.data.action.match(/^([a-zA-Z_]+)\(/);
236+
return actionMatch ? actionMatch[1] : "action";
237+
}
238+
239+
return d.data.id.toString().slice(-4);
240+
})
241+
.attr("font-size", "14px")
242+
.attr("font-weight", "500")
243+
.attr("fill", d => {
244+
if (d.data.id === selectedNodeId) {
245+
return theme === 'dark' ? "#93C5FD" : "#2563EB";
246+
}
247+
return theme === 'dark' ? "#FFFFFF" : "#111827";
248+
});
249+
250+
// Add reward values near nodes
251+
nodes.append("text")
252+
.attr("dy", "1.5em")
253+
.attr("x", d => d.children ? -18 : 18)
254+
.attr("text-anchor", d => d.children ? "end" : "start")
255+
.text(d => {
256+
if (typeof d.data.reward === 'number') {
257+
return `R: ${d.data.reward.toFixed(2)}`;
258+
}
259+
return "";
260+
})
261+
.attr("font-size", "12px")
262+
.attr("fill", theme === 'dark' ? "#E5E7EB" : "#4B5563");
263+
264+
// Add tooltip interactions
265+
nodes
266+
.on("mouseover", function(event, d) {
267+
if (tooltipRef.current) {
268+
let content = `<div><strong>Node ID:</strong> ${d.data.id}</div>`;
269+
if (d.data.action) content += `<div><strong>Action:</strong> ${d.data.action}</div>`;
270+
if (d.data.description) content += `<div><strong>Description:</strong> ${d.data.description}</div>`;
271+
if (typeof d.data.value === 'number') content += `<div><strong>Value:</strong> ${d.data.value.toFixed(2)}</div>`;
272+
if (typeof d.data.reward === 'number') content += `<div><strong>Reward:</strong> ${d.data.reward.toFixed(2)}</div>`;
273+
if (typeof d.data.visits === 'number') content += `<div><strong>Visits:</strong> ${d.data.visits}</div>`;
274+
if (d.data.feedback) content += `<div><strong>Feedback:</strong> ${d.data.feedback}</div>`;
275+
276+
const tooltip = d3.select(tooltipRef.current);
277+
tooltip.transition()
278+
.duration(200)
279+
.style("opacity", .9);
280+
tooltip.html(content)
281+
.style("left", (event.pageX + 15) + "px")
282+
.style("top", (event.pageY - 30) + "px");
283+
}
284+
})
285+
.on("mousemove", function(event) {
286+
if (tooltipRef.current) {
287+
d3.select(tooltipRef.current)
288+
.style("left", (event.pageX + 15) + "px")
289+
.style("top", (event.pageY - 30) + "px");
290+
}
291+
})
292+
.on("mouseout", function() {
293+
if (tooltipRef.current) {
294+
d3.select(tooltipRef.current)
295+
.transition()
296+
.duration(500)
297+
.style("opacity", 0);
298+
}
299+
});
300+
301+
// Add zoom behavior
302+
const zoom = d3.zoom<SVGSVGElement, unknown>()
303+
.scaleExtent([0.3, 3])
304+
.on("zoom", (event) => {
305+
g.attr("transform", event.transform);
306+
});
307+
308+
svg.call(zoom);
309+
310+
}, [treeNodes, selectedNodeId, theme, containerWidth]);
74311

75312
return (
76313
<div className="w-[30%] bg-white dark:bg-slate-800 rounded-r-lg overflow-hidden">
@@ -82,10 +319,13 @@ const SimpleSearchVisual: React.FC<SimpleSearchVisualProps> = ({ messages }) =>
82319
Tree Visualization
83320
</h2>
84321
</div>
85-
<div ref={containerRef} className="h-[calc(100%-48px)] w-full overflow-auto bg-gradient-to-r from-sky-50 to-white dark:from-slate-900 dark:to-slate-800">
322+
<div
323+
ref={containerRef}
324+
className="h-[calc(100%-48px)] w-full overflow-auto bg-gradient-to-r from-sky-50 to-white dark:from-slate-900 dark:to-slate-800"
325+
>
86326
<svg
87327
ref={svgRef}
88-
width="400"
328+
width="100%"
89329
height="700"
90330
className="overflow-visible"
91331
></svg>
@@ -94,4 +334,4 @@ const SimpleSearchVisual: React.FC<SimpleSearchVisualProps> = ({ messages }) =>
94334
);
95335
};
96336

97-
export default SimpleSearchVisual;
337+
export default SimpleSearchVisual;

0 commit comments

Comments
 (0)