Update app.py

app.py

@@ -10,6 +10,8 @@ import torch.nn as nn
 import torch.optim as optim
 from torch.utils.data import DataLoader, TensorDataset
 import base64
+import glob
+import time
 
 # Page Configuration
 st.set_page_config(
@@ -79,6 +81,10 @@ def sanitize_label(label):
     """Remove invalid characters for Mermaid labels."""
     return re.sub(r'[^\w\s-]', '', label).replace(' ', '_')
 
+def sanitize_filename(label):
+    """Make a valid filename from a label."""
+    return re.sub(r'[^\w\s-]', '', label).replace(' ', '_')
+
 def parse_outline_to_mermaid(outline_text, search_agent):
     """Convert tree outline to Mermaid syntax with clickable nodes."""
     lines = outline_text.strip().split('\n')
@@ -137,6 +143,45 @@ def get_download_link(file_path, mime_type="text/plain"):
     b64 = base64.b64encode(data).decode()
     return f'<a href="data:{mime_type};base64,{b64}" download="{file_path}">Download {file_path}</a>'
 
+def save_tree_to_file(tree_text, parent_node, new_node):
+    """Save tree to a markdown file with name based on nodes."""
+    root_node = tree_text.strip().split('\n')[0].split('.')[1].strip() if tree_text.strip() else "Knowledge_Tree"
+    filename = f"{sanitize_filename(root_node)}_{sanitize_filename(parent_node)}_{sanitize_filename(new_node)}_{int(time.time())}.md"
+    
+    mermaid_code = parse_outline_to_mermaid(tree_text, "🔮Google")  # Default search engine for saved trees
+    export_md = f"# Knowledge Tree: {root_node}\n\n## Outline\n{tree_text}\n\n## Mermaid Diagram\n```mermaid\n{mermaid_code}\n```"
+    
+    with open(filename, "w") as f:
+        f.write(export_md)
+    return filename
+
+def load_trees_from_files():
+    """Load all saved tree markdown files."""
+    tree_files = glob.glob("*.md")
+    trees_dict = {}
+    
+    for file in tree_files:
+        if file != "README.md" and file != "knowledge_tree.md":  # Skip project README and temp export
+            try:
+                with open(file, 'r') as f:
+                    content = f.read()
+                    # Extract the tree name from the first line
+                    match = re.search(r'# Knowledge Tree: (.*)', content)
+                    if match:
+                        tree_name = match.group(1)
+                    else:
+                        tree_name = os.path.splitext(file)[0]
+                    
+                    # Extract the outline section
+                    outline_match = re.search(r'## Outline\n(.*?)(?=\n## |$)', content, re.DOTALL)
+                    if outline_match:
+                        tree_outline = outline_match.group(1).strip()
+                        trees_dict[f"{tree_name} ({file})"] = tree_outline
+            except Exception as e:
+                print(f"Error loading {file}: {e}")
+    
+    return trees_dict
+
 # Search Agents (Highest resolution social network default: X)
 search_urls = {
     "📚📖ArXiv": lambda k: f"/?q={quote(k)}",
@@ -150,10 +195,22 @@ search_urls = {
 # Main App
 st.title("🌳 AI Knowledge Tree Builder 🌱")
 
+# Sidebar with saved trees
+st.sidebar.title("Saved Trees")
+saved_trees = load_trees_from_files()
+selected_saved_tree = st.sidebar.selectbox("Select a saved tree", ["None"] + list(saved_trees.keys()))
+
 # Select Project Type
 project_type = st.selectbox("Select Project Type", ["Code Project", "Papers Project", "AI Project"])
+
+# Initialize or load tree
 if 'current_tree' not in st.session_state:
-    st.session_state['current_tree'] = trees.get("ML Engineering", project_seeds[project_type])
+    if selected_saved_tree != "None" and selected_saved_tree in saved_trees:
+        st.session_state['current_tree'] = saved_trees[selected_saved_tree]
+    else:
+        st.session_state['current_tree'] = trees.get("ML Engineering", project_seeds[project_type])
+elif selected_saved_tree != "None" and selected_saved_tree in saved_trees:
+    st.session_state['current_tree'] = saved_trees[selected_saved_tree]
 
 # Select Search Agent for Node Links
 search_agent = st.selectbox("Select Search Agent for Node Links", list(search_urls.keys()), index=5)  # Default to X
@@ -163,9 +220,14 @@ new_node = st.text_input("Add New Node")
 parent_node = st.text_input("Parent Node")
 if st.button("Grow Tree 🌱") and new_node and parent_node:
     st.session_state['current_tree'] = grow_tree(st.session_state['current_tree'], new_node, parent_node)
+    
+    # Save to a new file with the node names
+    saved_file = save_tree_to_file(st.session_state['current_tree'], parent_node, new_node)
+    st.success(f"Added '{new_node}' under '{parent_node}' and saved to {saved_file}!")
+    
+    # Also update the temporary current_tree.md for compatibility
     with open("current_tree.md", "w") as f:
         f.write(st.session_state['current_tree'])
-    st.success(f"Added '{new_node}' under '{parent_node}'!")
 
 # Display Mermaid Diagram
 st.markdown("### Knowledge Tree Visualization")