diff --git a/tab_timeline.py b/tab_timeline.py index 5bfc86c..b62c570 100644 --- a/tab_timeline.py +++ b/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 - st.subheader("🕰️ Version History") - try: - graph_dot = htree.generate_horizontal_graph() - st.graphviz_chart(graph_dot, use_container_width=True) - except Exception as e: - st.error(f"Graph Error: {e}") + # 2. CONVERT TREE TO AGRAPH FORMAT + nodes = [] + edges = [] + + # Sort for consistent processing + sorted_nodes = sorted(htree.nodes.values(), key=lambda x: x["timestamp"]) + + 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 - col_sel, col_act = st.columns([3, 1]) + # 5. DETERMINE INSPECTION TARGET + # Priority: Clicked Node > Dropdown > Head - all_nodes = list(htree.nodes.values()) - all_nodes.sort(key=lambda x: x["timestamp"], reverse=True) + # Helper list for dropdown + all_nodes_list = list(htree.nodes.values()) + all_nodes_list.sort(key=lambda x: x["timestamp"], reverse=True) - def fmt_node(n): - return f"{n.get('note', 'Step')} ({n['id']})" + target_node_id = None + + # 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: - # Auto-select HEAD if possible - current_idx = 0 - 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: + # 6. INSPECTOR UI + if target_node_id and target_node_id in htree.nodes: + selected_node = htree.nodes[target_node_id] 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]) - c1.markdown(f"**{k}**") - c2.markdown(f"🔴 Current: `{str(v_now)[:50]}`") - c3.markdown(f"🟢 Selected: `{str(v_then)[:50]}`") + dc1, dc2, dc3 = st.columns([1, 2, 2]) + dc1.markdown(f"**{k}**") + dc2.markdown(f"🔴 `{str(v_now)[:30]}`") + dc3.markdown(f"🟢 `{str(v_then)[:30]}`") - # --- B. RENAME TOOL (NEW) --- - c_ren1, c_ren2 = st.columns([3, 1]) - new_note = c_ren1.text_input("Rename Node Label", value=selected_node.get("note", "")) - if c_ren2.button("Update Label"): - 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 + # --- B. ACTIONS --- + with c_h2: + st.write("") + if st.button("⏪ Restore Version", type="primary", use_container_width=True): 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 - node_label = f"{selected_node.get('note', 'Step')} ({selected_node['id'][:4]})" - st.session_state.restored_indicator = node_label + label = f"{selected_node.get('note')} ({target_node_id[:4]})" + st.session_state.restored_indicator = label - st.toast(f"Restored to {selected_node['id']}!", icon="🔄") + st.toast(f"Restored {target_node_id}!", icon="🔄") st.rerun() - # 5. RAW DATA & DELETE - with st.expander("Advanced Options (Raw JSON & Delete)"): - st.json(node_data, expanded=False) - if st.button("🗑️ Delete Node"): - if selected_node['id'] in htree.nodes: - del htree.nodes[selected_node['id']] - for b, tip in list(htree.branches.items()): - if tip == selected_node['id']: - del htree.branches[b] - - data["history_tree"] = htree.to_dict() - save_json(file_path, data) - st.rerun() + # --- C. RENAME --- + rn_col1, rn_col2 = st.columns([3, 1]) + new_label = rn_col1.text_input("Rename Label", value=selected_node.get("note", "")) + if rn_col2.button("Update"): + selected_node["note"] = new_label + data["history_tree"] = htree.to_dict() + save_json(file_path, data) + st.rerun()