Update tab_timeline.py

tab_timeline.py

@@ -1,13 +1,14 @@
 import streamlit as st
 import json
-import graphviz
 from history_tree import HistoryTree
 from utils import save_json
+# NEW IMPORTS
+from streamlit_agraph import agraph, Node, Edge, Config
 
 def render_timeline_tab(data, file_path):
     tree_data = data.get("history_tree", {})
     if not tree_data:
-        st.info("No history timeline exists for this file yet. Make some changes in the Editor first!")
+        st.info("No history timeline exists. Make some changes in the Editor first!")
         return
 
     htree = HistoryTree(tree_data)
@@ -16,110 +17,148 @@ def render_timeline_tab(data, file_path):
     if 'restored_indicator' in st.session_state and st.session_state.restored_indicator:
         st.info(f"📍 Editing Restored Version: **{st.session_state.restored_indicator}**")
 
-    # 2. COMPACT VISUALIZER
+    # 2. CONVERT TREE TO AGRAPH FORMAT
-    st.subheader("🕰️ Version History")
+    nodes = []
-    try:
+    edges = []
-        graph_dot = htree.generate_horizontal_graph()
+    
-        st.graphviz_chart(graph_dot, use_container_width=True)
+    # Sort for consistent processing
-    except Exception as e:
+    sorted_nodes = sorted(htree.nodes.values(), key=lambda x: x["timestamp"])
-        st.error(f"Graph Error: {e}")
+    
+    for n in sorted_nodes:
+        nid = n["id"]
+        note = n.get('note', 'Step')
+        short_note = (note[:15] + '..') if len(note) > 15 else note
+        
+        # Colors
+        color = "#ffffff"     # Default White
+        border = "#cccccc"
+        shape = "box"
+        
+        # Highlight HEAD (Current)
+        if nid == htree.head_id:
+            color = "#fff6cd" # Yellow
+            border = "#eebb00"
+        
+        # Highlight Branch Tips
+        if nid in htree.branches.values():
+            if color == "#ffffff": 
+                color = "#e6ffe6" # Green
+                border = "#99cc99"
+
+        # Create Node
+        nodes.append(Node(
+            id=nid,
+            label=f"{short_note}\n{nid[:4]}",
+            size=25,
+            shape=shape,
+            color=color,
+            borderWidth=1,
+            borderColor=border,
+            font={'color': 'black', 'face': 'Arial', 'size': 12}
+        ))
+        
+        # Create Edge
+        if n["parent"] and n["parent"] in htree.nodes:
+            edges.append(Edge(
+                source=n["parent"],
+                target=nid,
+                color="#aaaaaa",
+                type="STRAIGHT" # Keeps timeline looking clean
+            ))
+
+    # 3. CONFIGURE GRAPH VISUALS
+    config = Config(
+        width="100%",
+        height=300,  # Fixed height! Solves "Way too big" issue.
+        directed=True, 
+        physics=False, # Disable physics for a rigid timeline structure
+        hierarchical=True, # Force Tree Layout
+        direction="LR", # Left to Right
+        sortMethod="directed",
+        levelSeparation=150, # Space between time steps
+        nodeSpacing=100
+    )
+
+    st.subheader("🕰️ Interactive Timeline")
+    st.caption("Scroll to Zoom • Drag to Pan • Click to Inspect")
+    
+    # 4. RENDER & CAPTURE CLICK
+    # agraph returns the 'id' of the clicked node, or None
+    clicked_node_id = agraph(nodes=nodes, edges=edges, config=config)
 
     st.markdown("---")
 
-    # 3. INSPECTOR AREA
+    # 5. DETERMINE INSPECTION TARGET
-    col_sel, col_act = st.columns([3, 1])
+    # Priority: Clicked Node > Dropdown > Head
     
-    all_nodes = list(htree.nodes.values())
+    # Helper list for dropdown
-    all_nodes.sort(key=lambda x: x["timestamp"], reverse=True) 
+    all_nodes_list = list(htree.nodes.values())
+    all_nodes_list.sort(key=lambda x: x["timestamp"], reverse=True)
     
-    def fmt_node(n):
+    target_node_id = None
-        return f"{n.get('note', 'Step')} ({n['id']})"
+    
+    # If user clicked graph, use that
+    if clicked_node_id:
+        target_node_id = clicked_node_id
+    else:
+        # Fallback to current HEAD
+        target_node_id = htree.head_id
 
-    with col_sel:
+    # 6. INSPECTOR UI
-        # Auto-select HEAD if possible
+    if target_node_id and target_node_id in htree.nodes:
-        current_idx = 0
+        selected_node = htree.nodes[target_node_id]
-        for i, n in enumerate(all_nodes):
-            if n["id"] == htree.head_id:
-                current_idx = i
-                break
-                
-        selected_node = st.selectbox(
-            "Inspect Node:", 
-            all_nodes, 
-            format_func=fmt_node,
-            index=current_idx
-        )
-
-    # 4. SMART INSPECTOR
-    if selected_node:
         node_data = selected_node["data"]
-        
-        # --- A. DIFF VIEWER (NEW) ---
-        # Calculate differences between CURRENT Data and SELECTED Node
-        diffs = []
-        
-        # Compare keys
-        all_keys = set(data.keys()) | set(node_data.keys())
-        ignore_keys = {"history_tree", "prompt_history", "batch_data", "ui_reset_token"}
-        
-        for k in all_keys:
-            if k in ignore_keys: continue
-            
-            val_now = data.get(k, "N/A")
-            val_then = node_data.get(k, "N/A")
-            
-            if str(val_now) != str(val_then):
-                diffs.append((k, val_now, val_then))
 
-        with st.expander(f"🔍 Delta Inspector (Differences from Current)", expanded=True):
+        # Header
+        c_h1, c_h2 = st.columns([3, 1])
+        c_h1.markdown(f"### 🔎 Inspecting: {selected_node.get('note', 'Step')}")
+        c_h1.caption(f"ID: {target_node_id}")
+
+        # --- A. DIFF VIEWER ---
+        with st.expander(f"Compare with Current State", expanded=True):
+            diffs = []
+            all_keys = set(data.keys()) | set(node_data.keys())
+            ignore_keys = {"history_tree", "prompt_history", "batch_data", "ui_reset_token"}
+            
+            for k in all_keys:
+                if k in ignore_keys: continue
+                val_now = data.get(k, "N/A")
+                val_then = node_data.get(k, "N/A")
+                if str(val_now) != str(val_then):
+                    diffs.append((k, val_now, val_then))
+
             if not diffs:
-                st.caption("✅ This node is identical to your current state.")
+                st.caption("✅ Identical to current state")
             else:
                 for k, v_now, v_then in diffs:
-                    c1, c2, c3 = st.columns([1, 2, 2])
+                    dc1, dc2, dc3 = st.columns([1, 2, 2])
-                    c1.markdown(f"**{k}**")
+                    dc1.markdown(f"**{k}**")
-                    c2.markdown(f"🔴 Current: `{str(v_now)[:50]}`")
+                    dc2.markdown(f"🔴 `{str(v_now)[:30]}`")
-                    c3.markdown(f"🟢 Selected: `{str(v_then)[:50]}`")
+                    dc3.markdown(f"🟢 `{str(v_then)[:30]}`")
 
-        # --- B. RENAME TOOL (NEW) ---
+        # --- B. ACTIONS ---
-        c_ren1, c_ren2 = st.columns([3, 1])
+        with c_h2:
-        new_note = c_ren1.text_input("Rename Node Label", value=selected_node.get("note", ""))
+            st.write("")
-        if c_ren2.button("Update Label"):
+            if st.button("⏪ Restore Version", type="primary", use_container_width=True):
-            selected_node["note"] = new_note
-            data["history_tree"] = htree.to_dict()
-            save_json(file_path, data)
-            st.rerun()
-
-        # --- C. RESTORE BUTTON ---
-        with col_act:
-            st.write(""); st.write("")
-            if st.button("⏪ Restore", type="primary", use_container_width=True):
-                # Restore Logic
                 data.update(node_data)
-                htree.head_id = selected_node['id']
+                htree.head_id = target_node_id
                 
                 data["history_tree"] = htree.to_dict()
                 save_json(file_path, data)
                 
                 st.session_state.ui_reset_token += 1
                 
-                # Set Indicator
+                label = f"{selected_node.get('note')} ({target_node_id[:4]})"
-                node_label = f"{selected_node.get('note', 'Step')} ({selected_node['id'][:4]})"
+                st.session_state.restored_indicator = label
-                st.session_state.restored_indicator = node_label
                 
-                st.toast(f"Restored to {selected_node['id']}!", icon="🔄")
+                st.toast(f"Restored {target_node_id}!", icon="🔄")
                 st.rerun()
 
-    # 5. RAW DATA & DELETE
+        # --- C. RENAME ---
-    with st.expander("Advanced Options (Raw JSON & Delete)"):
+        rn_col1, rn_col2 = st.columns([3, 1])
-        st.json(node_data, expanded=False)
+        new_label = rn_col1.text_input("Rename Label", value=selected_node.get("note", ""))
-        if st.button("🗑️ Delete Node"):
+        if rn_col2.button("Update"):
-            if selected_node['id'] in htree.nodes:
+            selected_node["note"] = new_label
-                del htree.nodes[selected_node['id']]
+            data["history_tree"] = htree.to_dict()
-                for b, tip in list(htree.branches.items()):
+            save_json(file_path, data)
-                    if tip == selected_node['id']:
+            st.rerun()
-                        del htree.branches[b] 
-                
-                data["history_tree"] = htree.to_dict()
-                save_json(file_path, data)
-                st.rerun()