app.py

"""
n8n Workflow Generator - Gradio Web Interface
Deploy this to Hugging Face Spaces
"""

import gradio as gr
from transformers import AutoModelForCausalLM, AutoTokenizer
from peft import PeftModel
import torch
import json
import re

# ==============================================================================
# CONFIGURATION
# ==============================================================================

MODEL_REPO = "Nishan30/n8n-workflow-generator"  # Update with your HF repo
BASE_MODEL = "Qwen/Qwen2.5-Coder-1.5B-Instruct"

SYSTEM_PROMPT = """You are an expert n8n workflow generator. Given a user's request, you generate clean, functional TypeScript code using the @n8n-generator/core DSL.

Your output should:
- Only contain the code, no explanations
- Use the Workflow class from @n8n-generator/core
- Use workflow.add() to create nodes
- Use .to() or workflow.connect() for connections
- Be ready to compile directly to n8n JSON

Example:
User: "Create a webhook that sends data to Slack"
Assistant:
```typescript
const workflow = new Workflow('Webhook to Slack');
const webhook = workflow.add('n8n-nodes-base.webhook', { path: 'data' });
const slack = workflow.add('n8n-nodes-base.slack', { channel: '#general' });
webhook.to(slack);
```"""

# ==============================================================================
# MODEL LOADING
# ==============================================================================

@gr.cache
def load_model():
    """Load model once and cache it"""
    print("Loading model...")

    # Load base model
    base_model = AutoModelForCausalLM.from_pretrained(
        BASE_MODEL,
        torch_dtype=torch.float16,
        device_map="auto",
        trust_remote_code=True
    )

    # Load LoRA adapter
    model = PeftModel.from_pretrained(base_model, MODEL_REPO)
    tokenizer = AutoTokenizer.from_pretrained(MODEL_REPO)

    print("Model loaded successfully!")
    return model, tokenizer

# Load model at startup
model, tokenizer = load_model()

# ==============================================================================
# CODE GENERATION
# ==============================================================================

def generate_workflow(prompt, temperature=0.3, max_tokens=512):
    """Generate n8n workflow code from prompt"""

    if not prompt.strip():
        return "Please enter a workflow description.", None, None

    # Format messages
    messages = [
        {"role": "system", "content": SYSTEM_PROMPT},
        {"role": "user", "content": prompt}
    ]

    # Apply chat template
    text = tokenizer.apply_chat_template(
        messages,
        tokenize=False,
        add_generation_prompt=True
    )

    # Tokenize
    inputs = tokenizer(text, return_tensors="pt").to(model.device)

    # Generate
    with torch.no_grad():
        outputs = model.generate(
            **inputs,
            max_new_tokens=max_tokens,
            temperature=temperature,
            do_sample=True if temperature > 0 else False,
            top_p=0.9,
            repetition_penalty=1.1
        )

    # Decode
    generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)

    # Extract code from response
    code = extract_code(generated_text)

    # Convert to n8n JSON
    n8n_json = convert_to_n8n_json(code)

    # Create visualization
    visualization = create_visualization(n8n_json)

    return code, json.dumps(n8n_json, indent=2), visualization

def extract_code(text):
    """Extract TypeScript code from generated text"""

    # Try to find code block
    code_match = re.search(r'```(?:typescript|ts)?\n(.*?)```', text, re.DOTALL)
    if code_match:
        return code_match.group(1).strip()

    # If no code block, look for code after assistant response
    if "assistant" in text.lower():
        parts = text.split("assistant", 1)
        if len(parts) > 1:
            # Remove any markdown code blocks
            code = parts[1].strip()
            code = re.sub(r'```(?:typescript|ts)?\n', '', code)
            code = re.sub(r'```', '', code)
            return code.strip()

    return text.strip()

# ==============================================================================
# N8N JSON CONVERSION
# ==============================================================================

def convert_to_n8n_json(typescript_code):
    """Convert TypeScript DSL to n8n JSON format"""

    nodes = []
    connections = {}
    workflow_name = "Generated Workflow"

    # Extract workflow name
    name_match = re.search(r"new Workflow\(['\"](.*?)['\"]\)", typescript_code)
    if name_match:
        workflow_name = name_match.group(1)

    # Extract node definitions
    node_pattern = r'const\s+(\w+)\s*=\s*workflow\.add\([\'"]([^\'\"]+)[\'"](?:,\s*({[^}]+}))?\)'
    node_matches = re.finditer(node_pattern, typescript_code)

    node_map = {}  # variable name -> node id
    position_y = 250
    position_x = 300

    for i, match in enumerate(node_matches):
        var_name = match.group(1)
        node_type = match.group(2)
        params_str = match.group(3) if match.group(3) else "{}"

        # Parse parameters (basic JSON parsing)
        try:
            parameters = json.loads(params_str)
        except:
            parameters = {}

        node_id = str(i)
        node_map[var_name] = node_id

        nodes.append({
            "id": node_id,
            "name": var_name,
            "type": node_type,
            "typeVersion": 1,
            "position": [position_x, position_y],
            "parameters": parameters
        })

        position_x += 300

    # Extract connections
    connection_pattern = r'(\w+)\.to\((\w+)\)'
    connection_matches = re.finditer(connection_pattern, typescript_code)

    for match in connection_matches:
        source_var = match.group(1)
        target_var = match.group(2)

        if source_var in node_map and target_var in node_map:
            source_id = node_map[source_var]
            target_id = node_map[target_var]

            # Find source node name
            source_node = next((n for n in nodes if n["id"] == source_id), None)
            if source_node:
                source_name = source_node["name"]

                if source_name not in connections:
                    connections[source_name] = {"main": [[]] }

                connections[source_name]["main"][0].append({
                    "node": target_var,
                    "type": "main",
                    "index": 0
                })

    return {
        "name": workflow_name,
        "nodes": nodes,
        "connections": connections,
        "active": False,
        "settings": {}
    }

# ==============================================================================
# VISUALIZATION
# ==============================================================================

def create_visualization(n8n_json):
    """Create HTML visualization of the workflow"""

    nodes = n8n_json.get("nodes", [])
    connections = n8n_json.get("connections", {})

    if not nodes:
        return "<div style='padding:20px;text-align:center;color:#666;'>No nodes found in workflow</div>"

    html = """
    <div style="font-family: Arial, sans-serif; padding: 20px; background: #f5f5f5; border-radius: 8px;">
        <h3 style="margin-top:0; color: #ff6d5a;">📊 Workflow Visualization</h3>
        <div style="display: flex; flex-direction: column; gap: 15px;">
    """

    # Display nodes
    for i, node in enumerate(nodes):
        node_name = node.get("name", f"Node{i}")
        node_type = node.get("type", "unknown").split(".")[-1]
        params = node.get("parameters", {})

        # Count outgoing connections
        outgoing = 0
        for source, conns in connections.items():
            if source == node_name:
                outgoing = len(conns.get("main", [[]])[0])

        # Node card
        html += f"""
        <div style="background: white; padding: 15px; border-radius: 8px; border-left: 4px solid #ff6d5a; box-shadow: 0 2px 4px rgba(0,0,0,0.1);">
            <div style="display: flex; justify-content: space-between; align-items: center;">
                <div>
                    <div style="font-weight: bold; font-size: 16px; color: #333;">{node_name}</div>
                    <div style="color: #666; font-size: 14px; margin-top: 4px;">
                        <code style="background: #f0f0f0; padding: 2px 6px; border-radius: 3px;">{node_type}</code>
                    </div>
                </div>
                <div style="text-align: right; color: #999; font-size: 12px;">
                    Node #{i+1}
                </div>
            </div>
        """

        # Show key parameters
        if params:
            html += "<div style='margin-top: 10px; font-size: 13px; color: #555;'>"
            html += "<strong>Parameters:</strong><br>"
            for key, value in list(params.items())[:3]:  # Show first 3 params
                value_str = str(value)[:50]
                html += f"&nbsp;&nbsp;• {key}: <code style='background:#f9f9f9;padding:1px 4px;'>{value_str}</code><br>"
            html += "</div>"

        # Show connections
        if outgoing > 0:
            html += f"<div style='margin-top: 8px; color: #4CAF50; font-size: 12px;'>→ {outgoing} connection(s)</div>"

        html += "</div>"

        # Show arrow between nodes
        if i < len(nodes) - 1:
            html += "<div style='text-align: center; color: #999; font-size: 20px;'>↓</div>"

    html += """
        </div>
        <div style="margin-top: 15px; padding: 10px; background: #e3f2fd; border-radius: 4px; font-size: 12px; color: #1976d2;">
            💡 <strong>Tip:</strong> Copy the n8n JSON and import it directly into your n8n instance!
        </div>
    </div>
    """

    return html

# ==============================================================================
# GRADIO INTERFACE
# ==============================================================================

def create_ui():
    """Create Gradio interface"""

    with gr.Blocks(title="n8n Workflow Generator", theme=gr.themes.Soft()) as demo:

        gr.Markdown("""
        # 🚀 n8n Workflow Generator

        Generate n8n workflows using natural language! Powered by fine-tuned **Qwen2.5-Coder-1.5B**.

        ### How to use:
        1. Describe your workflow in plain English
        2. Click "Generate Workflow"
        3. Copy the generated code or n8n JSON
        4. Import into your n8n instance
        """)

        with gr.Row():
            with gr.Column(scale=1):
                prompt_input = gr.Textbox(
                    label="Workflow Description",
                    placeholder="Example: Create a webhook that receives data, filters active users, and sends to Slack",
                    lines=3
                )

                with gr.Row():
                    temperature = gr.Slider(
                        minimum=0.0,
                        maximum=1.0,
                        value=0.3,
                        step=0.1,
                        label="Temperature (creativity)",
                        info="Lower = more consistent, Higher = more creative"
                    )
                    max_tokens = gr.Slider(
                        minimum=128,
                        maximum=1024,
                        value=512,
                        step=128,
                        label="Max tokens",
                        info="Maximum length of generated code"
                    )

                generate_btn = gr.Button("🎯 Generate Workflow", variant="primary", size="lg")

                gr.Markdown("""
                ### 📝 Example Prompts:
                - *Create a webhook that sends data to Slack*
                - *Schedule that runs daily and backs up database to Google Drive*
                - *Webhook receives form data, validates email, saves to Airtable*
                - *Monitor RSS feed and post new items to Twitter*
                """)

            with gr.Column(scale=1):
                visualization_output = gr.HTML(label="Visual Workflow")

        with gr.Row():
            with gr.Column():
                code_output = gr.Code(
                    label="Generated TypeScript Code",
                    language="typescript",
                    lines=15
                )

            with gr.Column():
                json_output = gr.Code(
                    label="n8n JSON (import this into n8n)",
                    language="json",
                    lines=15
                )

        # Examples
        gr.Examples(
            examples=[
                ["Create a webhook that sends data to Slack"],
                ["Build a workflow that fetches GitHub issues and sends daily summary email"],
                ["Webhook receives order, if amount > $1000 send to priority queue, else standard processing"],
                ["Schedule that runs every Monday, fetches data from API, transforms it, and updates Google Sheets"],
                ["Monitor RSS feeds, remove duplicates, and post to Twitter"],
            ],
            inputs=prompt_input
        )

        # Event handler
        generate_btn.click(
            fn=generate_workflow,
            inputs=[prompt_input, temperature, max_tokens],
            outputs=[code_output, json_output, visualization_output]
        )

        gr.Markdown("""
        ---
        ### ℹ️ About

        This model achieved **92.4% accuracy** on diverse n8n workflow generation tasks.

        **Model:** Fine-tuned Qwen2.5-Coder-1.5B with LoRA
        **Training:** 247 curated workflow examples
        **Performance:** Production-ready quality

        [🤗 Model Card](https://huggingface.co/{}) | [📊 GitHub](https://github.com/yourusername/n8n-generator)
        """.format(MODEL_REPO))

    return demo

# ==============================================================================
# LAUNCH
# ==============================================================================

if __name__ == "__main__":
    demo = create_ui()
    demo.launch(
        server_name="0.0.0.0",
        server_port=7860,
        share=False
    )