Update app.py

app.py

@@ -1,132 +1,101 @@
 import gradio as gr
 import graphviz
 import os
+from huggingface_hub import InferenceClient
 from PIL import Image, ImageDraw, ImageFont
 
-# --- Setup: Create directories and placeholder image ---
+# --- Constants and Configuration ---
+# Using a powerful open-source instruction-tuned model from Mistral AI
+MODEL_ID = "mistralai/Mistral-7B-Instruct-v0.2"
 
-# Create a directory for output images if it doesn't exist
-if not os.path.exists("outputs"):
-    os.makedirs("outputs")
+# System prompt to instruct the AI model
+# This is the most crucial part for getting reliable, structured output.
+SYSTEM_PROMPT = """You are an expert at creating flowcharts. A user will provide a text prompt, and your task is to generate the flowchart description in the Graphviz DOT language.
 
-def create_placeholder_image(path="placeholder.png"):
-    """Creates a placeholder image to display on app startup."""
+Your response MUST be ONLY the Graphviz DOT language source code for a directed graph (digraph).
+- The graph should be top-to-bottom (`rankdir=TB`).
+- Use rounded boxes for process steps (`shape=box, style="rounded,filled", fillcolor="#EAEAFB"`).
+- Use diamonds for decision points (`shape=diamond, fillcolor="#F9EED5"`).
+- Use ellipses for start and end nodes (`shape=ellipse, fillcolor="#D5D6F9"`).
+- Do not include any explanations, comments, or markdown formatting like ```dot ... ```. The entire response must be valid DOT code."""
+
+# --- Helper Functions ---
+def create_placeholder_image(text="Flowchart will be generated here", size=(600, 800), path="placeholder.png"):
+    """Creates a placeholder or error image with text."""
     try:
-        img = Image.new('RGB', (600, 800), color=(255, 255, 255))
+        img = Image.new('RGB', size, color=(255, 255, 255))
         draw = ImageDraw.Draw(img)
         try:
-            # Use a common font, fallback to default
             font = ImageFont.truetype("DejaVuSans.ttf", 24)
         except IOError:
             font = ImageFont.load_default()
-        text = "Flowchart will be generated here"
-        # Get text size using textbbox
+            
         bbox = draw.textbbox((0, 0), text, font=font)
         text_width = bbox[2] - bbox[0]
         text_height = bbox[3] - bbox[1]
-        position = ((600 - text_width) / 2, (800 - text_height) / 2)
+        position = ((size[0] - text_width) / 2, (size[1] - text_height) / 2)
         draw.text(position, text, fill=(200, 200, 200), font=font)
         img.save(path)
         return path
     except Exception:
         return None
 
-# --- Core Logic: Flowchart Generation ---
-
-def generate_flowchart(prompt: str) -> str:
+# --- Core AI and Rendering Logic ---
+def generate_flowchart(prompt: str, hf_token: str):
     """
-    Generates a flowchart using Graphviz based on the input prompt.
+    Calls the Hugging Face Inference API to generate DOT code and then renders it.
     Returns the file path of the generated PNG image.
     """
-    # A simple keyword-based logic to generate different flowcharts
-    if "coffee" in prompt.lower():
-        dot_string = """
-        digraph G {
-            graph [pad="0.5", nodesep="0.5", ranksep="0.8", fontname="helvetica"];
-            node [shape=box, style="rounded,filled", fillcolor="#EAEAFB", fontname="helvetica"];
-            edge [color="#666666"];
-            rankdir=TB;
-
-            start [label="Start", shape=ellipse, fillcolor="#D5D6F9"];
-            have_beans [label="Have coffee beans?", shape=diamond, fillcolor="#F9EED5"];
-            buy_beans [label="Buy coffee beans"];
-            grind_beans [label="Grind beans"];
-            boil_water [label="Boil water"];
-            add_grounds [label="Add grounds to filter"];
-            pour_water [label="Pour water over grounds"];
-            wait [label="Wait for brewing"];
-            pour_cup [label="Pour coffee into cup"];
-            add_extras [label="Add milk/sugar?", shape=diamond, fillcolor="#F9EED5"];
-            add_milk_sugar [label="Add milk and/or sugar"];
-            drink [label="Drink coffee"];
-            end_node [label="Enjoy your coffee!", shape=ellipse, fillcolor="#D5D6F9"];
-
-            start -> have_beans;
-            have_beans -> grind_beans [label="  Yes"];
-            have_beans -> buy_beans [label="No  "];
-            buy_beans -> grind_beans;
-            grind_beans -> boil_water;
-            boil_water -> add_grounds;
-            add_grounds -> pour_water;
-            pour_water -> wait;
-            wait -> pour_cup;
-            pour_cup -> add_extras;
-            add_extras -> add_milk_sugar [label="  Yes"];
-            add_extras -> drink [label="No  "];
-            add_milk_sugar -> drink;
-            drink -> end_node;
-        }
-        """
-    else:
-        # A default flowchart for any other prompt
-        cleaned_prompt = prompt.replace("'", "").replace("\"", "")[:60]
-        dot_string = f"""
-        digraph G {{
-            graph [fontname="helvetica"];
-            node [shape=box, style="rounded,filled", fillcolor="#EAEAFB", fontname="helvetica"];
-            rankdir=TB;
-            
-            start [label="Start", shape=ellipse, fillcolor="#D5D6F9"];
-            process [label="Process user request:\\n'{cleaned_prompt}... '"];
-            end_node [label="End", shape=ellipse, fillcolor="#D5D6F9"];
-            
-            start -> process -> end_node;
-        }}
-        """
+    if not hf_token:
+        return create_placeholder_image("Error: Hugging Face API Token is not set.\nPlease add it to your Space's secrets."), None
+
+    if not prompt:
+        return create_placeholder_image("Please enter a prompt to generate a flowchart."), None
 
     try:
-        graph = graphviz.Source(dot_string)
-        # Save the rendered graph to a file, overwriting previous ones
+        # 1. Call the LLM to get the DOT language code
+        client = InferenceClient(model=MODEL_ID, token=hf_token)
+        response = client.chat_completion(
+            messages=[
+                {"role": "system", "content": SYSTEM_PROMPT},
+                {"role": "user", "content": prompt},
+            ],
+            max_tokens=1500,
+            temperature=0.7,
+        )
+        dot_code = response.choices[0].message.content.strip()
+        
+        # Sometimes the model still adds markdown, let's strip it just in case
+        if dot_code.startswith("```dot"):
+            dot_code = dot_code[len("```dot"):].strip()
+        if dot_code.startswith("```"):
+            dot_code = dot_code[len("```"):].strip()
+        if dot_code.endswith("```"):
+            dot_code = dot_code[:-len("```")].strip()
+
+        # 2. Render the DOT code using Graphviz
+        graph = graphviz.Source(dot_code)
         output_path = graph.render(os.path.join("outputs", "flowchart"), format='png', cleanup=True)
-        return output_path
-    except Exception as e:
-        print(f"Error rendering graphviz: {e}")
-        return create_placeholder_image("error.png") # Return an error image if something goes wrong
-
-
-def handle_generate_click(prompt: str):
-    """Function to be called on button click. Updates the UI."""
-    if not prompt:
-        # Handle empty prompt case
-        placeholder_path = create_placeholder_image()
-        return placeholder_path, gr.update(visible=False)
         
-    image_path = generate_flowchart(prompt)
-    
-    # Return the image path to update the gr.Image component
-    # and a gr.update object for the gr.DownloadButton
-    return image_path, gr.update(value=image_path, visible=True)
+        return output_path, gr.update(value=output_path, visible=True)
 
+    except Exception as e:
+        print(f"An error occurred: {e}")
+        error_message = f"An error occurred.\nThis could be due to an invalid API token,\nan issue with the AI model, or invalid generated DOT code.\n\nDetails: {str(e)}"
+        return create_placeholder_image(error_message), gr.update(visible=False)
 
-# --- Gradio UI ---
 
-# Custom CSS to hide the footer for a cleaner look like the screenshot
+# --- Gradio UI ---
 css = """
 footer {display: none !important}
 .gradio-container {background-color: #f8f9fa}
+#status_display {text-align: center; color: #888;}
 """
 
 with gr.Blocks(theme=gr.themes.Soft(), css=css) as demo:
+    # Get the Hugging Face token from environment variables
+    hf_token = os.environ.get("HF_TOKEN")
+
     gr.Markdown("# AI Flowchart Generator")
     gr.Markdown(
         "Our AI Flowchart Generator allows you to create detailed flowcharts instantly. Whether you need a free "
@@ -135,29 +104,18 @@ with gr.Blocks(theme=gr.themes.Soft(), css=css) as demo:
         "AI solution available online."
     )
 
-    with gr.Group(): # Using gr.Group to create the white card effect
+    with gr.Group():
         with gr.Row(equal_height=False):
             with gr.Column(scale=1):
-                lang_dropdown = gr.Dropdown(
-                    label="Language",
-                    choices=["English"],
-                    value="English"
-                )
-                chart_type_dropdown = gr.Dropdown(
-                    choices=["Flowchart"],
-                    value="Flowchart",
-                    label=" " # Using space for an empty label like in the screenshot
-                )
                 prompt_input = gr.Textbox(
                     lines=10,
-                    placeholder="Generate a flowchart for making coffee",
-                    value="Generate a flowchart for making coffee",
-                    label=" " # Empty label
+                    placeholder="e.g., Explain the process of making a cup of tea",
+                    label="Enter your process description here"
                 )
                 with gr.Row():
-                    # These buttons are for UI purposes and are not connected to logic
-                    templates_btn = gr.Button("📄 Templates")
                     generate_btn = gr.Button("✨ Generate", variant="primary")
+                
+                status_display = gr.Markdown("", elem_id="status_display")
 
             with gr.Column(scale=1):
                 output_image = gr.Image(
@@ -171,22 +129,35 @@ with gr.Blocks(theme=gr.themes.Soft(), css=css) as demo:
                 download_btn = gr.DownloadButton(
                     "⬇️ Download",
                     variant="primary",
-                    visible=False, # Initially hidden
+                    visible=False,
                 )
 
     # --- Event Listeners ---
+    def on_generate_click(prompt):
+        # Provide user feedback that generation is in progress
+        yield (
+            gr.update(interactive=False), # Disable button
+            gr.update(visible=False),     # Hide download button
+            create_placeholder_image("🧠 Generating with AI... Please wait."),
+            "Generating... this can take up to 30 seconds."
+        )
+        # Call the main generation function
+        img_path, download_btn_update = generate_flowchart(prompt, hf_token)
+        # Update UI with the result
+        yield (
+            gr.update(interactive=True), # Re-enable button
+            download_btn_update,         # Show/update download button
+            img_path,
+            "" # Clear status message
+        )
+        
     generate_btn.click(
-        fn=handle_generate_click,
-        inputs=[prompt_input],
-        outputs=[output_image, download_btn]
-    )
-    
-    # Pre-load the default example when the app starts
-    demo.load(
-        fn=handle_generate_click,
+        fn=on_generate_click,
         inputs=[prompt_input],
-        outputs=[output_image, download_btn]
+        outputs=[generate_btn, download_btn, output_image, status_display]
     )
 
 if __name__ == "__main__":
+    if not os.path.exists("outputs"):
+        os.makedirs("outputs")
     demo.launch()