Update app.py

app.py

@@ -1,278 +1,278 @@
-import os
-import math
-import streamlit as st
-from typing import Dict, Optional
-from groq import Groq
-import cairosvg
-import re
-
-# --------------------------------------------------------------------
-# Streamlit page configuration
-# --------------------------------------------------------------------
-st.set_page_config(
-    layout="wide",
-    page_title="AI Mind Map Generator",
-    initial_sidebar_state="expanded"
-)
-
-# --------------------------------------------------------------------
-# Helper Functions for Mermaid Parsing and SVG Generation
-# --------------------------------------------------------------------
-def parse_mermaid_to_svg(mermaid_code, layout="flowchart"):
-    """
-    Parses Mermaid code to extract nodes and edges and generates SVG elements based on the chosen layout.
-
-    :param mermaid_code: Mermaid graph syntax as a string.
-    :param layout: The desired layout type (e.g., "wireframe" or "flowchart").
-    :return: SVG content as a string.
-    """
-    nodes = {}
-    edges = []
-
-    # Extract nodes
-    node_pattern = re.compile(r'(\w+)\[(.*?)\]')
-    for match in node_pattern.finditer(mermaid_code):
-        node_id, label = match.groups()
-        nodes[node_id] = label
-
-    # Extract edges
-    edge_pattern = re.compile(r'(\w+)\s*-->\s*(\w+)')
-    for match in edge_pattern.finditer(mermaid_code):
-        source, target = match.groups()
-        edges.append((source, target))
-
-    # Initialize SVG content
-    svg_content = '''
-    <svg viewBox="0 0 1200 800" xmlns="http://www.w3.org/2000/svg">
-        <!-- Background -->
-        <rect width="1200" height="800" fill="#ffffff"/>
-    '''
-
-    # Layout logic
-    if layout == "flowchart":
-        columns = 4  # Default columns for flowchart
-        spacing_x = 250
-        spacing_y = 150
-        start_x = 100
-        start_y = 100
-
-        node_positions = {}
-
-        for i, (node_id, label) in enumerate(nodes.items()):
-            x = start_x + (i % columns) * spacing_x
-            y = start_y + (i // columns) * spacing_y
-            node_positions[node_id] = (x, y)
-
-            svg_content += f'''
-            <g>
-                <rect x="{x - 75}" y="{y - 25}" width="150" height="50" fill="#4CAF50" rx="10" ry="10"/>
-                <text x="{x}" y="{y + 5}" text-anchor="middle" fill="white" font-family="Arial" font-size="14">{label}</text>
-            </g>
-            '''
-
-        # Draw edges
-        for source, target in edges:
-            if source in node_positions and target in node_positions:
-                x1, y1 = node_positions[source]
-                x2, y2 = node_positions[target]
-                svg_content += f'''
-                <line x1="{x1}" y1="{y1}" x2="{x2}" y2="{y2}" stroke="#000000" stroke-width="2" marker-end="url(#arrowhead)"/>
-                '''
-    elif layout == "wireframe":
-        center_x, center_y = 600, 400
-        radius = max(150, min(300, 250 + 10 * len(nodes)))
-        angle_step = 2 * math.pi / max(1, len(nodes))  # Avoid division by zero
-
-        node_positions = {}
-        for i, (node_id, label) in enumerate(nodes.items()):
-            angle = i * angle_step
-            x = center_x + radius * math.cos(angle)
-            y = center_y + radius * math.sin(angle)
-            node_positions[node_id] = (x, y)
-
-            svg_content += f'''
-            <g>
-                <rect x="{x - 75}" y="{y - 25}" width="150" height="50" fill="#f39c12" rx="10" ry="10"/>
-                <text x="{x}" y="{y + 5}" text-anchor="middle" fill="black" font-family="Arial" font-size="14">{label}</text>
-            </g>
-            '''
-
-        # Draw edges
-        for source, target in edges:
-            if source in node_positions and target in node_positions:
-                x1, y1 = node_positions[source]
-                x2, y2 = node_positions[target]
-                svg_content += f'''
-                <line x1="{x1}" y1="{y1}" x2="{x2}" y2="{y2}" stroke="#d35400" stroke-width="2" marker-end="url(#arrowhead)"/>
-                '''
-
-    # Add arrowhead for edges
-    svg_content += '''
-        <defs>
-            <marker id="arrowhead" markerWidth="10" markerHeight="7" refX="10" refY="3.5" orient="auto">
-                <polygon points="0 0, 10 3.5, 0 7" fill="black"/>
-            </marker>
-        </defs>
-    '''
-
-    svg_content += '</svg>'
-    return svg_content
-
-def generate_and_save_final_image(mermaid_code, layout="flowchart"):
-    """
-    Generates the final image based on Mermaid code and saves it as a PNG file.
-
-    :param mermaid_code: Mermaid graph syntax as a string.
-    :param layout: The desired layout type (e.g., "wireframe" or "flowchart").
-    :return: Tuple (success, output_file or error message).
-    """
-    try:
-        # Parse Mermaid to SVG
-        svg_content = parse_mermaid_to_svg(mermaid_code, layout)
-        
-        # Ensure the output directory exists
-        output_directory = "images"
-        os.makedirs(output_directory, exist_ok=True)
-        
-        # Define output file path
-        output_file = os.path.join(output_directory, f"final_image_{layout}.png")
-        
-        # Convert SVG to PNG and save
-        cairosvg.svg2png(bytestring=svg_content.encode("utf-8"), write_to=output_file)
-        return True, output_file
-    except Exception as e:
-        return False, str(e)
-
-# --------------------------------------------------------------------
-# Supported Models
-# --------------------------------------------------------------------
-SUPPORTED_MODELS: Dict[str, str] = {
-    "Llama 3 8B": "llama3-8b-8192",
-    "Llama 3.2 1B (Preview)": "llama-3.2-1b-preview",
-    "Llama 3 70B": "llama3-70b-8192",
-    "Mixtral 8x7B": "mixtral-8x7b-32768",
-    "Gemma 2 9B": "gemma2-9b-it",
-    "Llama 3.2 11B Vision (Preview)": "llama-3.2-11b-vision-preview",
-    "Llama 3.2 11B Text (Preview)": "llama-3.2-11b-text-preview",
-    "Llama 3.1 8B Instant (Text-Only Workloads)": "llama-3.1-8b-instant",
-    "Llama 3.2 90B Vision (Preview)": "llama-3.2-90b-vision-preview",
-    "Llama 3.1 70B Versatile": "llama-3.1-70b-versatile",
-    "Llama 3.1 8B Instant": "llama-3.1-8b-instant",
-    "Llama 3.2 11B Vision (Preview)": "llama-3.2-11b-vision-preview",
-    "Llama 3.2 1B (Preview)": "llama-3.2-1b-preview",
-    "Llama 3.2 3B (Preview)": "llama-3.2-3b-preview",
-    "Llama 3.2 90B Vision (Preview)": "llama-3.2-90b-vision-preview",
-    "Llama 3.3 70B SpecDec": "llama-3.3-70b-specdec",
-    "Llama 3.3 70B Versatile": "llama-3.3-70b-versatile",
-}
-
-MAX_TOKENS: int = 1500
-
-# --------------------------------------------------------------------
-# Initialize Groq client with API key
-# --------------------------------------------------------------------
-@st.cache_resource
-def get_groq_client() -> Optional[Groq]:
-    groq_api_key = os.getenv("GROQ_API_KEY")
-    if not groq_api_key:
-        st.error("GROQ_API_KEY not found in environment variables. Please set it and restart the app.")
-        return None
-    return Groq(api_key=groq_api_key)
-
-client = get_groq_client()
-
-# --------------------------------------------------------------------
-# SIDEBAR
-# --------------------------------------------------------------------
-st.sidebar.image("icon.png", width=300)
-st.sidebar.title("Model Configuration")
-
-selected_model = st.sidebar.selectbox("Choose an AI Model", list(SUPPORTED_MODELS.keys()))
-
-st.sidebar.subheader("Temperature")
-temperature = st.sidebar.slider(
-    "Set temperature for generation variability:", 
-    min_value=0.0, 
-    max_value=1.0, 
-    value=0.7
-)
-
-# Add layout selection
-st.sidebar.subheader("Layout Configuration")
-layout = st.sidebar.radio(
-    "Select the layout for the mind map:",
-    options=["flowchart", "wireframe"]
-)
-
-# --------------------------------------------------------------------
-# MAIN CONTENT
-# --------------------------------------------------------------------
-st.title("AI Mind Map Generator")
-st.markdown(
-    """
-    Enter your concepts or a short description below, then click **Generate Mind Map**. 
-    The Groq LLM will produce Mermaid diagram code, which we'll display below.
-    """
-)
-
-# Text area for user input
-mind_map_prompt = st.text_area(
-    "Describe your mind map focus:",
-    placeholder="e.g. 'Attention and Intention in personal development'"
-)
-
-if st.button("Generate Mind Map"):
-    if not mind_map_prompt.strip():
-        st.warning("Please provide a description or concept for the mind map.")
-    elif client:
-        with st.spinner("Generating your mind map..."):
-            prompt = f"""
-            You are an AI that generates a Mind Map in Mermaid format. 
-            The user wants a mind map about: {mind_map_prompt}.
-            Please output ONLY the Mermaid diagram, nothing else.
-            """
-
-            try:
-                response = client.chat.completions.create(
-                    model=SUPPORTED_MODELS[selected_model],
-                    messages=[
-                        {"role": "system", "content": "You are an AI that generates mind maps in Mermaid code."},
-                        {"role": "user", "content": prompt},
-                    ],
-                    temperature=temperature,
-                    max_tokens=MAX_TOKENS,
-                )
-
-                mermaid_code = response.choices[0].message.content.strip()
-
-                st.subheader("Generated Mind Map")
-                st.markdown(
-                    f"""
-                    ```mermaid
-                    {mermaid_code}
-                    ```
-                    """,
-                    unsafe_allow_html=True
-                )
-
-                st.download_button(
-                    label="Download Mermaid Code",
-                    data=mermaid_code,
-                    file_name="mind_map_mermaid.txt",
-                    mime="text/plain"
-                )
-
-                # Generate and display the final image based on layout
-                success, result = generate_and_save_final_image(mermaid_code, layout)
-                if success:
-                    st.image(result, caption=f"Generated Mind Map ({layout.capitalize()} Layout)", use_column_width=True)
-                else:
-                    st.error(f"Failed to generate image: {result}")
-
-            except Exception as e:
-                st.error(f"Error generating mind map: {e}")
-    else:
-        st.error("Groq client not initialized. Make sure you have set your GROQ_API_KEY environment variable.")
-
-st.info("Built by dw — This app uses Groq LLM to generate Mermaid-based mind maps.")
-
+import os
+import math
+import streamlit as st
+from typing import Dict, Optional
+from groq import Groq
+import cairosvg
+import re
+
+# --------------------------------------------------------------------
+# Streamlit page configuration
+# --------------------------------------------------------------------
+st.set_page_config(
+    layout="wide",
+    page_title="AI Mind Map Generator",
+    initial_sidebar_state="expanded"
+)
+
+# --------------------------------------------------------------------
+# Helper Functions for Mermaid Parsing and SVG Generation
+# --------------------------------------------------------------------
+def parse_mermaid_to_svg(mermaid_code, layout="flowchart"):
+    """
+    Parses Mermaid code to extract nodes and edges and generates SVG elements based on the chosen layout.
+
+    :param mermaid_code: Mermaid graph syntax as a string.
+    :param layout: The desired layout type (e.g., "wireframe" or "flowchart").
+    :return: SVG content as a string.
+    """
+    nodes = {}
+    edges = []
+
+    # Extract nodes
+    node_pattern = re.compile(r'(\w+)\[(.*?)\]')
+    for match in node_pattern.finditer(mermaid_code):
+        node_id, label = match.groups()
+        nodes[node_id] = label
+
+    # Extract edges
+    edge_pattern = re.compile(r'(\w+)\s*-->\s*(\w+)')
+    for match in edge_pattern.finditer(mermaid_code):
+        source, target = match.groups()
+        edges.append((source, target))
+
+    # Initialize SVG content
+    svg_content = '''
+    <svg viewBox="0 0 1200 800" xmlns="http://www.w3.org/2000/svg">
+        <!-- Background -->
+        <rect width="1200" height="800" fill="#ffffff"/>
+    '''
+
+    # Layout logic
+    if layout == "flowchart":
+        columns = 4  # Default columns for flowchart
+        spacing_x = 250
+        spacing_y = 150
+        start_x = 100
+        start_y = 100
+
+        node_positions = {}
+
+        for i, (node_id, label) in enumerate(nodes.items()):
+            x = start_x + (i % columns) * spacing_x
+            y = start_y + (i // columns) * spacing_y
+            node_positions[node_id] = (x, y)
+
+            svg_content += f'''
+            <g>
+                <rect x="{x - 75}" y="{y - 25}" width="150" height="50" fill="#4CAF50" rx="10" ry="10"/>
+                <text x="{x}" y="{y + 5}" text-anchor="middle" fill="white" font-family="Arial" font-size="14">{label}</text>
+            </g>
+            '''
+
+        # Draw edges
+        for source, target in edges:
+            if source in node_positions and target in node_positions:
+                x1, y1 = node_positions[source]
+                x2, y2 = node_positions[target]
+                svg_content += f'''
+                <line x1="{x1}" y1="{y1}" x2="{x2}" y2="{y2}" stroke="#000000" stroke-width="2" marker-end="url(#arrowhead)"/>
+                '''
+    elif layout == "wireframe":
+        center_x, center_y = 600, 400
+        radius = max(150, min(300, 250 + 10 * len(nodes)))
+        angle_step = 2 * math.pi / max(1, len(nodes))  # Avoid division by zero
+
+        node_positions = {}
+        for i, (node_id, label) in enumerate(nodes.items()):
+            angle = i * angle_step
+            x = center_x + radius * math.cos(angle)
+            y = center_y + radius * math.sin(angle)
+            node_positions[node_id] = (x, y)
+
+            svg_content += f'''
+            <g>
+                <rect x="{x - 75}" y="{y - 25}" width="150" height="50" fill="#f39c12" rx="10" ry="10"/>
+                <text x="{x}" y="{y + 5}" text-anchor="middle" fill="black" font-family="Arial" font-size="14">{label}</text>
+            </g>
+            '''
+
+        # Draw edges
+        for source, target in edges:
+            if source in node_positions and target in node_positions:
+                x1, y1 = node_positions[source]
+                x2, y2 = node_positions[target]
+                svg_content += f'''
+                <line x1="{x1}" y1="{y1}" x2="{x2}" y2="{y2}" stroke="#d35400" stroke-width="2" marker-end="url(#arrowhead)"/>
+                '''
+
+    # Add arrowhead for edges
+    svg_content += '''
+        <defs>
+            <marker id="arrowhead" markerWidth="10" markerHeight="7" refX="10" refY="3.5" orient="auto">
+                <polygon points="0 0, 10 3.5, 0 7" fill="black"/>
+            </marker>
+        </defs>
+    '''
+
+    svg_content += '</svg>'
+    return svg_content
+
+def generate_and_save_final_image(mermaid_code, layout="flowchart"):
+    """
+    Generates the final image based on Mermaid code and saves it as a PNG file.
+
+    :param mermaid_code: Mermaid graph syntax as a string.
+    :param layout: The desired layout type (e.g., "wireframe" or "flowchart").
+    :return: Tuple (success, output_file or error message).
+    """
+    try:
+        # Parse Mermaid to SVG
+        svg_content = parse_mermaid_to_svg(mermaid_code, layout)
+        
+        # Ensure the output directory exists
+        output_directory = "images"
+        os.makedirs(output_directory, exist_ok=True)
+        
+        # Define output file path
+        output_file = os.path.join(output_directory, f"final_image_{layout}.png")
+        
+        # Convert SVG to PNG and save
+        cairosvg.svg2png(bytestring=svg_content.encode("utf-8"), write_to=output_file)
+        return True, output_file
+    except Exception as e:
+        return False, str(e)
+
+# --------------------------------------------------------------------
+# Supported Models
+# --------------------------------------------------------------------
+SUPPORTED_MODELS: Dict[str, str] = {
+    "Llama 3 8B": "llama3-8b-8192",
+    "Llama 3.2 1B (Preview)": "llama-3.2-1b-preview",
+    "Llama 3 70B": "llama3-70b-8192",
+    "Mixtral 8x7B": "mixtral-8x7b-32768",
+    "Gemma 2 9B": "gemma2-9b-it",
+    "Llama 3.2 11B Vision (Preview)": "llama-3.2-11b-vision-preview",
+    "Llama 3.2 11B Text (Preview)": "llama-3.2-11b-text-preview",
+    "Llama 3.1 8B Instant (Text-Only Workloads)": "llama-3.1-8b-instant",
+    "Llama 3.2 90B Vision (Preview)": "llama-3.2-90b-vision-preview",
+    "Llama 3.1 70B Versatile": "llama-3.1-70b-versatile",
+    "Llama 3.1 8B Instant": "llama-3.1-8b-instant",
+    "Llama 3.2 11B Vision (Preview)": "llama-3.2-11b-vision-preview",
+    "Llama 3.2 1B (Preview)": "llama-3.2-1b-preview",
+    "Llama 3.2 3B (Preview)": "llama-3.2-3b-preview",
+    "Llama 3.2 90B Vision (Preview)": "llama-3.2-90b-vision-preview",
+    "Llama 3.3 70B SpecDec": "llama-3.3-70b-specdec",
+    "Llama 3.3 70B Versatile": "llama-3.3-70b-versatile",
+}
+
+MAX_TOKENS: int = 1500
+
+# --------------------------------------------------------------------
+# Initialize Groq client with API key
+# --------------------------------------------------------------------
+@st.cache_resource
+def get_groq_client() -> Optional[Groq]:
+    groq_api_key = os.getenv("GROQ_API_KEY")
+    if not groq_api_key:
+        st.error("GROQ_API_KEY not found in environment variables. Please set it and restart the app.")
+        return None
+    return Groq(api_key=groq_api_key)
+
+client = get_groq_client()
+
+# --------------------------------------------------------------------
+# SIDEBAR
+# --------------------------------------------------------------------
+st.sidebar.image("icon.png", width=300)
+st.sidebar.title("Model Configuration")
+
+selected_model = st.sidebar.selectbox("Choose an AI Model", list(SUPPORTED_MODELS.keys()))
+
+st.sidebar.subheader("Temperature")
+temperature = st.sidebar.slider(
+    "Set temperature for generation variability:", 
+    min_value=0.0, 
+    max_value=1.0, 
+    value=0.7
+)
+
+# Add layout selection
+st.sidebar.subheader("Layout Configuration")
+layout = st.sidebar.radio(
+    "Select the layout for the mind map:",
+    options=["flowchart", "wireframe"]
+)
+
+# --------------------------------------------------------------------
+# MAIN CONTENT
+# --------------------------------------------------------------------
+st.title("AI Mind Map Generator")
+st.markdown(
+    """
+    Enter your concepts or a short description below, then click **Generate Mind Map**. 
+    The Groq LLM will produce Mermaid diagram code, which we'll display below.
+    """
+)
+
+# Text area for user input
+mind_map_prompt = st.text_area(
+    "Describe your mind map focus:",
+    placeholder="e.g. 'Attention and Intention in personal development'"
+)
+
+if st.button("Generate Mind Map"):
+    if not mind_map_prompt.strip():
+        st.warning("Please provide a description or concept for the mind map.")
+    elif client:
+        with st.spinner("Generating your mind map..."):
+            prompt = f"""
+            You are an AI that generates a Mind Map in Mermaid format. 
+            The user wants a mind map about: {mind_map_prompt}.
+            Please output ONLY the Mermaid diagram, nothing else.
+            """
+
+            try:
+                response = client.chat.completions.create(
+                    model=SUPPORTED_MODELS[selected_model],
+                    messages=[
+                        {"role": "system", "content": "You are an AI that generates mind maps in Mermaid code."},
+                        {"role": "user", "content": prompt},
+                    ],
+                    temperature=temperature,
+                    max_tokens=MAX_TOKENS,
+                )
+
+                mermaid_code = response.choices[0].message.content.strip()
+
+                st.subheader("Generated Mind Map")
+                st.markdown(
+                    f"""
+                    ```mermaid
+                    {mermaid_code}
+                    ```
+                    """,
+                    unsafe_allow_html=True
+                )
+
+                st.download_button(
+                    label="Download Mermaid Code",
+                    data=mermaid_code,
+                    file_name="mind_map_mermaid.txt",
+                    mime="text/plain"
+                )
+
+                # Generate and display the final image based on layout
+                success, result = generate_and_save_final_image(mermaid_code, layout)
+                if success:
+                    st.image(result, caption=f"Generated Mind Map ({layout.capitalize()} Layout)", use_container_width=True)
+                else:
+                    st.error(f"Failed to generate image: {result}")
+
+            except Exception as e:
+                st.error(f"Error generating mind map: {e}")
+    else:
+        st.error("Groq client not initialized. Make sure you have set your GROQ_API_KEY environment variable.")
+
+st.info("Built by dw — This app uses Groq LLM to generate Mermaid-based mind maps.")
+