app.py

import gradio as gr
import requests
import spacy
import re
import time
import os  # Import the 'os' module
import base64  # Import base64

# Try to load the spaCy model, download if it's not available
try:
    nlp = spacy.load("en_core_web_sm")
except OSError:
    print("Downloading en_core_web_sm model...")
    import spacy.cli
    spacy.cli.download("en_core_web_sm")
    nlp = spacy.load("en_core_web_sm")


def text_to_diagram(text_description):
    """Convert process description into Mermaid.js flowchart."""
    doc = nlp(text_description.lower())
    steps = []
    conditions = []
    loops = []
    complex_flow = []
    sentences = [sent.text.strip() for sent in doc.sents]

    for sentence in sentences:
        sentence = sentence.replace("2, conditional branch", "").replace("3 loop start", "").replace("4 complex flow", "").strip()

        if "then" in sentence:
            tokens = sentence.split("then")
            for token in tokens:
                token = token.strip()
                if token:
                    step = token.strip().replace(",", "").replace(".", "")
                    if step and step not in ["conditional branch", "loop start", "complex flow"]:
                        if "flowchart that shows" in step:
                            continue
                        steps.append(step)

        cond_match = re.search(r"if ([a-z\s]+), (?:go to|proceed to) ([a-z\s]+); otherwise, (?:go to|proceed to) ([a-z\s]+)", sentence)
        if not cond_match:
            cond_match = re.search(r"if ([a-z\s]+), ([a-z\s]+); otherwise, ([a-z\s]+)", sentence)
        if cond_match:
            condition, true_action, false_action = cond_match.groups()
            conditions.append((condition.strip(), true_action.strip(), false_action.strip()))

        loop_match = re.search(r"(?:do|perform) (?:action )?([a-z]), repeat ([a-z]) until ([a-z]) is true", sentence)
        if not loop_match:
            loop_match = re.search(r"loop: ([a-z\s]+), repeat ([a-z\s]+) until ([a-z\s]+)", sentence)
        if loop_match:
            action, repeat_action, condition = loop_match.groups()
            loops.append((action.strip(), condition.strip()))

        if "flowchart that shows" in sentence:
            flow_part = sentence.split("flowchart that shows")[1]
            flow_steps = flow_part.split("then")
            for step in flow_steps:
                step = step.replace("and finally", "").replace(".", "").replace(",", "").strip()
                if step:
                    complex_flow.append(step)
        elif "then" in sentence and not steps:
            tokens = sentence.split("then")
            for token in tokens:
                token = token.strip()
                if token and token not in steps:
                    step = token.replace(",", "").replace(".", "").strip()
                    if step:
                        complex_flow.append(step)

    if not complex_flow:
        complex_flow = steps[1:] if steps else []
        steps = steps[:1] if steps else []

    diagram_code = "graph TD\n"
    node_counter = 0
    node_map = {}

    def get_node_id(label):
        nonlocal node_counter
        if label not in node_map:
            node_map[label] = chr(65 + node_counter)
            node_counter += 1
        return node_map[label]

    if steps:
        for i in range(len(steps)):
            step = steps[i]
            node_id = get_node_id(step)
            diagram_code += f"{node_id}[{step.capitalize()}]\n"
            if i > 0:
                prev_node_id = get_node_id(steps[i - 1])
                diagram_code += f"{prev_node_id} --> {node_id}\n"
    else:
        start_node = get_node_id("start")
        diagram_code += f"{start_node}[Start]\n"

    last_node = steps[-1] if steps else "start"
    last_node_id = get_node_id(last_node)
    for condition, true_action, false_action in conditions:
        cond_node = get_node_id(f"cond_{condition}")
        true_node = get_node_id(true_action)
        false_node = get_node_id(false_action)
        diagram_code += f"{last_node_id} --> {cond_node}{{{condition.capitalize()}?}}\n"
        diagram_code += f"{cond_node} -->|Yes| {true_node}[{true_action.capitalize()}]\n"
        diagram_code += f"{cond_node} -->|No| {false_node}[{false_action.capitalize()}]\n"
        last_node = true_action
        last_node_id = true_node
        false_node_id = false_node

    for action, condition in loops:
        loop_node = get_node_id(f"loop_{action}")
        cond_node = get_node_id(f"cond_{condition}")
        finish_node = get_node_id(f"finish_{action}")
        diagram_code += f"{last_node_id} --> {loop_node}[Loop: Perform {action.capitalize()}]\n"
        diagram_code += f"{loop_node} --> {cond_node}{{{condition.capitalize()} is True?}}\n"
        diagram_code += f"{cond_node} -->|No| {loop_node}\n"
        diagram_code += f"{cond_node} -->|Yes| {finish_node}[Finish Loop]\n"
        last_node_id = finish_node

    for i in range(len(complex_flow)):
        step = complex_flow[i]
        node_id = get_node_id(f"flow_{i}_{step}")
        diagram_code += f"{node_id}[{step.capitalize()}]\n"
        if i == 0:
            diagram_code += f"{last_node_id} --> {node_id}\n"
        else:
            prev_node_id = get_node_id(f"flow_{i - 1}_{complex_flow[i - 1]}")
            diagram_code += f"{prev_node_id} --> {node_id}\n"
        last_node_id = node_id

    end_node = get_node_id("end")
    diagram_code += f"{end_node}[End Process]\n"
    diagram_code += f"{last_node_id} --> {end_node}\n"

    if conditions:
        diagram_code += f"{false_node_id} --> {end_node}\n"

    try:
        img_url = render_mermaid_to_url(diagram_code)
        response = requests.get(img_url, timeout=10)
        response.raise_for_status()  # Raise HTTPError for bad responses (4xx or 5xx)
        image_data = response.content
        temp_img_path = f"temp_diagram_{int(time.time())}.png"
        with open(temp_img_path, "wb") as f:
            f.write(image_data)
        return diagram_code, temp_img_path
    except requests.exceptions.RequestException as e:
        print(f"Error rendering diagram (network): {e}")
        return diagram_code, None
    except Exception as e:
        print(f"Error rendering diagram: {e}")
        return diagram_code, None


def render_mermaid_to_url(mermaid_code):
    """Render Mermaid code to an image URL using Mermaid.ink."""
    try:
        encoded_code = base64.urlsafe_b64encode(mermaid_code.encode()).decode()
        return f"https://mermaid.ink/img/{encoded_code}"
    except Exception as e:
        print(f"Error encoding Mermaid code: {e}")
        return None

def gradio_interface(text_input):
    """Process user input and return diagram output for Gradio."""
    try:
        diagram_code, img_path = text_to_diagram(text_input)
        print("Generated Mermaid Code:")
        print(diagram_code)
        return diagram_code, img_path if img_path else None
    except Exception as e:
        print(f"Error in Gradio interface: {e}")
        return f"Error: {str(e)}", None

iface = gr.Interface(
    fn=gradio_interface,
    inputs=gr.Textbox(lines=5, placeholder="Describe your process (e.g., 'User logs in, then enters details...')"),
    outputs=[
        gr.Textbox(label="Generated Mermaid Code"),
        gr.Image(label="Diagram Visualization", type="filepath"),
    ],
    title="Text-to-Diagram Converter (Napkin AI Style)",
    description="Convert process descriptions into structured flowcharts like Napkin AI.",
    examples=[
        ["Developer writes code, then tests code. If tests pass, proceed to deploy; otherwise, debug code. Then release software to production."],
        ["Customer browses products, then adds to cart, then proceeds to checkout. If payment is verified, confirm order; otherwise, request payment again. Loop: Validate payment, repeat until successful. Then ship order and notify customer."],
        ["HR posts job opening, then reviews applications, then conducts interviews. If candidate is selected, offer job; otherwise, reject candidate. Then onboard new employee."],
        ["Student registers for exam, then studies material. If prepared, proceed to take exam; otherwise, study more. If exam is passed, receive certificate; otherwise, retake exam. Then celebrate success."],
        ["Customer places order, then kitchen prepares food. Loop: Check food quality, repeat until perfect. Then serve food and collect payment."],
        ["Traveler searches for flights, then selects flight, then books ticket. If booking is confirmed, proceed to payment; otherwise, search again. If payment is successful, receive itinerary; otherwise, cancel booking. Then plan trip."],
        ["User signs up, then verifies email, then logs in."],
    ],
    allow_flagging="never",
)

if __name__ == "__main__":
    iface.launch(share=True)