Ethanfel
b02bf124fb
- save_json() now writes to a temp file then uses os.replace() for atomic writes - Replace hardcoded "batch_data", "history_tree", "prompt_history", "sequence_number" strings with constants (KEY_BATCH_DATA, etc.) across all modules - Add 29 unit tests for history_tree, utils, and json_loader - Add type hints to public functions in utils.py, json_loader.py, history_tree.py - Remove ALLOWED_BASE_DIR restriction that blocked navigating outside app CWD - Fix path text input not updating on navigation by using session state key - Add unpin button (❌) for removing pinned folders Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
130 lines
4.7 KiB
Python
130 lines
4.7 KiB
Python
import time
|
|
import uuid
|
|
from typing import Any
|
|
|
|
KEY_PROMPT_HISTORY = "prompt_history"
|
|
|
|
|
|
class HistoryTree:
|
|
def __init__(self, raw_data: dict[str, Any]) -> None:
|
|
self.nodes: dict[str, dict[str, Any]] = raw_data.get("nodes", {})
|
|
self.branches: dict[str, str | None] = raw_data.get("branches", {"main": None})
|
|
self.head_id: str | None = raw_data.get("head_id", None)
|
|
|
|
if KEY_PROMPT_HISTORY in raw_data and isinstance(raw_data[KEY_PROMPT_HISTORY], list) and not self.nodes:
|
|
self._migrate_legacy(raw_data[KEY_PROMPT_HISTORY])
|
|
|
|
def _migrate_legacy(self, old_list: list[dict[str, Any]]) -> None:
|
|
parent = None
|
|
for item in reversed(old_list):
|
|
node_id = str(uuid.uuid4())[:8]
|
|
self.nodes[node_id] = {
|
|
"id": node_id, "parent": parent, "timestamp": time.time(),
|
|
"data": item, "note": item.get("note", "Legacy Import")
|
|
}
|
|
parent = node_id
|
|
self.branches["main"] = parent
|
|
self.head_id = parent
|
|
|
|
def commit(self, data: dict[str, Any], note: str = "Snapshot") -> str:
|
|
new_id = str(uuid.uuid4())[:8]
|
|
|
|
# Cycle detection: walk parent chain from head to verify no cycle
|
|
if self.head_id:
|
|
visited = set()
|
|
current = self.head_id
|
|
while current:
|
|
if current in visited:
|
|
raise ValueError(f"Cycle detected in history tree at node {current}")
|
|
visited.add(current)
|
|
node = self.nodes.get(current)
|
|
current = node["parent"] if node else None
|
|
|
|
active_branch = None
|
|
for b_name, tip_id in self.branches.items():
|
|
if tip_id == self.head_id:
|
|
active_branch = b_name
|
|
break
|
|
|
|
if not active_branch:
|
|
base_name = "branch"
|
|
count = 1
|
|
while f"{base_name}_{count}" in self.branches: count += 1
|
|
active_branch = f"{base_name}_{count}"
|
|
|
|
self.nodes[new_id] = {
|
|
"id": new_id, "parent": self.head_id, "timestamp": time.time(),
|
|
"data": data, "note": note
|
|
}
|
|
self.branches[active_branch] = new_id
|
|
self.head_id = new_id
|
|
return new_id
|
|
|
|
def checkout(self, node_id: str) -> dict[str, Any] | None:
|
|
if node_id in self.nodes:
|
|
self.head_id = node_id
|
|
return self.nodes[node_id]["data"]
|
|
return None
|
|
|
|
def to_dict(self) -> dict[str, Any]:
|
|
return {"nodes": self.nodes, "branches": self.branches, "head_id": self.head_id}
|
|
|
|
# --- UPDATED GRAPH GENERATOR ---
|
|
def generate_graph(self, direction: str = "LR") -> str:
|
|
"""
|
|
Generates Graphviz source.
|
|
direction: "LR" (Horizontal) or "TB" (Vertical)
|
|
"""
|
|
dot = [
|
|
'digraph History {',
|
|
f' rankdir={direction};', # Dynamic Direction
|
|
' bgcolor="white";',
|
|
' splines=ortho;',
|
|
|
|
# TIGHT SPACING
|
|
' nodesep=0.2;',
|
|
' ranksep=0.3;',
|
|
|
|
# GLOBAL STYLES
|
|
' node [shape=plain, fontname="Arial"];',
|
|
' edge [color="#888888", arrowsize=0.6, penwidth=1.0];'
|
|
]
|
|
|
|
sorted_nodes = sorted(self.nodes.values(), key=lambda x: x["timestamp"])
|
|
|
|
for n in sorted_nodes:
|
|
nid = n["id"]
|
|
full_note = n.get('note', 'Step')
|
|
|
|
display_note = (full_note[:15] + '..') if len(full_note) > 15 else full_note
|
|
|
|
# COLORS
|
|
bg_color = "#f9f9f9"
|
|
border_color = "#999999"
|
|
border_width = "1"
|
|
|
|
if nid == self.head_id:
|
|
bg_color = "#fff6cd" # Yellow for Current
|
|
border_color = "#eebb00"
|
|
border_width = "2"
|
|
elif nid in self.branches.values():
|
|
bg_color = "#e6ffe6" # Green for Tips
|
|
border_color = "#66aa66"
|
|
|
|
# HTML LABEL
|
|
label = (
|
|
f'<<TABLE BORDER="{border_width}" CELLBORDER="0" CELLSPACING="0" CELLPADDING="4" BGCOLOR="{bg_color}" COLOR="{border_color}">'
|
|
f'<TR><TD><B><FONT POINT-SIZE="10">{display_note}</FONT></B></TD></TR>'
|
|
f'<TR><TD><FONT POINT-SIZE="8" COLOR="#555555">{nid[:4]}</FONT></TD></TR>'
|
|
f'</TABLE>>'
|
|
)
|
|
|
|
safe_tooltip = full_note.replace('"', "'")
|
|
dot.append(f' "{nid}" [label={label}, tooltip="{safe_tooltip}"];')
|
|
|
|
if n["parent"] and n["parent"] in self.nodes:
|
|
dot.append(f' "{n["parent"]}" -> "{nid}";')
|
|
|
|
dot.append("}")
|
|
return "\n".join(dot)
|