Update app.py

app.py

@@ -1,286 +1,67 @@
-import ast
-import logging
 import os
 import re
-from datetime import datetime
-
 import gradio as gr
 from groq import Groq
 
-# ---------------------------------------------------------------------------
-# Config
-# ---------------------------------------------------------------------------
 GROQ_API_KEY = os.environ.get("GROQ_API_KEY")
-MODEL = "llama-3.3-70b-versatile"
-TEMPERATURE = 0.1
-MAX_GOAL_CHARS = 4000
-
-logging.basicConfig(level=logging.INFO)
-logger = logging.getLogger(__name__)
-
-# Netmiko device_type -> NAPALM driver (when user picks NAPALM)
-NAPALM_DRIVER_MAP = {
-    "cisco_ios": "ios",
-    "cisco_nxos": "nxos",
-    "arista_eos": "eos",
-    "juniper_junos": "junos",
-}
-
-SYSTEM_PROMPT = """You are a Senior Network Automation Architect.
-Your task is to generate production-ready Python automation scripts.
-
-Rules:
-- Output ONLY raw Python code. No markdown fences, no ```python, no prose before or after.
-- Use Python 3.10+ syntax.
-- Never hardcode real passwords, API keys, or production IPs. Use placeholders and environment variables.
-- Include try/except for connection failures, sensible timeouts, and basic logging where appropriate.
-- Add brief inline comments only for non-obvious steps.
-- Include if __name__ == "__main__": with argparse or clear placeholders for hosts, credentials, and paths.
-- For Netmiko: use ConnectHandler with device_type exactly as provided in the user message.
-- For NAPALM: use the napalm driver name provided in the user message.
-- For Nornir: show inventory YAML structure or a minimal InitNornir example with placeholders.
-- Scripts must be reviewable in a lab before production use."""
-
-CUSTOM_CSS = """
-.gradio-container { font-family: 'Inter', system-ui, sans-serif; }
-.header-box {
-    text-align: center;
-    padding: 20px;
-    background: linear-gradient(90deg, #1e3a8a, #3b82f6);
-    color: white;
-    border-radius: 10px;
-    margin-bottom: 16px;
-}
-.header-box h1 { margin: 0; font-size: 28px; font-weight: bold; color: white !important; }
-.header-box p { margin: 5px 0 0 0; opacity: 0.9; }
-.disclaimer {
-    font-size: 0.85rem;
-    color: #b45309;
-    background: #fffbeb;
-    border: 1px solid #fcd34d;
-    border-radius: 8px;
-    padding: 10px 14px;
-    margin-bottom: 12px;
-}
-"""
-
-
-def strip_code_fences(text: str) -> str:
-    """Remove ```python ... ``` wrappers if the model ignores instructions."""
-    if not text:
-        return ""
-    text = text.strip()
-    pattern = r"^```(?:python)?\s*\n?(.*?)\n?```\s*$"
-    match = re.match(pattern, text, re.DOTALL | re.IGNORECASE)
-    return match.group(1).strip() if match else text
-
-
-def validate_syntax(code: str) -> str:
-    """Return status line for UI (does not block output)."""
-    try:
-        ast.parse(code)
-        return "Syntax check: OK"
-    except SyntaxError as e:
-        return f"Syntax check: possible issue — {e.msg} (line {e.lineno})"
-
-
-def build_user_prompt(user_goal: str, library: str, device_type: str) -> str:
-    napalm_driver = NAPALM_DRIVER_MAP.get(device_type, device_type)
-    return f"""Library: {library}
-Netmiko device_type (if using Netmiko): {device_type}
-NAPALM driver (if using NAPALM): {napalm_driver}
-
-Task:
-{user_goal.strip()}
-
-Requirements:
-- Match idioms and APIs for {library}.
-- Use placeholders for inventory (hostnames, usernames, paths).
-- If the task involves many devices, use concurrent execution safely with clear limits.
-- End with runnable structure and documented placeholders."""
-
 
-def generate_network_script(
-    user_goal: str,
-    library: str,
-    device_type: str,
-    progress: gr.Progress = gr.Progress(),
-) -> tuple[str, str, str]:
-    """
-    Returns: (code, status_message, download_filename_suggestion)
-    """
-    progress(0, desc="Validating input…")
+def strip_code_fences(text):
+    text = (text or "").strip()
+    m = re.match(r"^```(?:python)?\s*\n?(.*?)\n?```\s*$", text, re.DOTALL | re.IGNORECASE)
+    return m.group(1).strip() if m else text
 
+def generate_network_script(user_goal, library, device_type):
     if not GROQ_API_KEY:
-        err = (
-            "# Error\n"
-            "GROQ_API_KEY is missing. Add it in Space Settings → "
-            "Variables and secrets."
-        )
-        return err, "Missing API key", ""
-
+        return "# Error\nGROQ_API_KEY is missing. Add it in Space Settings → Variables and secrets."
     goal = (user_goal or "").strip()
     if not goal:
-        return "# Error\nPlease describe the task.", "Empty task", ""
-    if len(goal) > MAX_GOAL_CHARS:
-        return (
-            f"# Error\nTask description exceeds {MAX_GOAL_CHARS} characters.",
-            "Input too long",
-            "",
-        )
+        return "# Error\nPlease describe the task."
 
-    progress(0.2, desc="Calling Groq…")
-    client = Groq(api_key=GROQ_API_KEY)
-    user_prompt = build_user_prompt(goal, library, device_type)
+    system_prompt = """You are a Senior Network Automation Architect.
+Output ONLY raw Python code. No markdown fences. No extra text."""
+    user_prompt = f"Library: {library}\nDevice: {device_type}\nTask: {goal}"
 
     try:
-        chat_completion = client.chat.completions.create(
+        client = Groq(api_key=GROQ_API_KEY)
+        r = client.chat.completions.create(
+            model="llama-3.3-70b-versatile",
             messages=[
-                {"role": "system", "content": SYSTEM_PROMPT},
+                {"role": "system", "content": system_prompt},
                 {"role": "user", "content": user_prompt},
             ],
-            model=MODEL,
-            temperature=TEMPERATURE,
-        )
-        raw = chat_completion.choices[0].message.content or ""
-        code = strip_code_fences(raw)
-        progress(0.9, desc="Checking syntax…")
-        status = validate_syntax(code)
-        progress(1.0, desc="Done")
-        filename = (
-            f"autonet_{library.lower()}_{device_type}_"
-            f"{datetime.utcnow().strftime('%Y%m%d_%H%M%S')}.py"
-        )
-        return code, status, filename
-
-    except Exception:
-        logger.exception("Groq generation failed")
-        return (
-            "# Error\n"
-            "Generation failed. Check Space logs or try again in a moment.",
-            "Generation failed",
-            "",
+            temperature=0.1,
         )
-
-
-def on_generate(user_goal, library, device_type, progress=gr.Progress()):
-    code, status, filename = generate_network_script(
-        user_goal, library, device_type, progress
-    )
-    return code, status, code, filename
-
-
-# Example prompts for first-time users
-EXAMPLES = [
-    [
-        "Log into a list of Cisco switches via SSH, backup running-config, "
-        "and save each file as {hostname}_{date}_running.cfg in ./backups/",
-        "Netmiko",
-        "cisco_ios",
-    ],
-    [
-        "Connect to three Arista switches, collect interface status and "
-        "print any interface that is admin down but operationally up.",
-        "NAPALM",
-        "arista_eos",
-    ],
-    [
-        "Using an inventory file, run 'show version' on all Juniper devices "
-        "and save output to logs/{host}.txt",
-        "Nornir",
-        "juniper_junos",
-    ],
-]
-
-HEADER_HTML = """
-<div class="header-box">
-    <h1>⚙️ AutoNet Architect</h1>
-    <p>AI-Powered Python Script Generator for Network Engineers</p>
-</div>
-"""
-
-DISCLAIMER_HTML = """
-<div class="disclaimer">
-    <strong>Review before you run.</strong> Generated code may be wrong or unsafe.
-    Test in a lab first. Never paste production passwords into this form.
-    You are responsible for anything executed on your network.
-</div>
+        return strip_code_fences(r.choices[0].message.content)
+    except Exception as e:
+        return f"# Error\n{str(e)}"
+
+custom_css = """
+.gradio-container { font-family: 'Inter', sans-serif; }
+.header-box { text-align: center; padding: 20px; background: linear-gradient(90deg, #1e3a8a, #3b82f6);
+  color: white; border-radius: 10px; margin-bottom: 20px; }
+.header-box h1 { margin: 0; font-size: 28px; color: white !important; }
 """
 
-
-with gr.Blocks(
-    theme=gr.themes.Soft(primary_hue="blue", secondary_hue="indigo"),
-    css=CUSTOM_CSS,
-    title="AutoNet Architect",
-) as app:
-
-    gr.HTML(HEADER_HTML)
-    gr.HTML(DISCLAIMER_HTML)
-
+with gr.Blocks(theme=gr.themes.Soft(primary_hue="blue", secondary_hue="indigo"), css=custom_css) as app:
+    gr.HTML('<div class="header-box"><h1>⚙️ AutoNet Architect</h1><p>AI script generator</p></div>')
     with gr.Row():
         with gr.Column(scale=1):
-            library = gr.Dropdown(
-                choices=["Netmiko", "NAPALM", "Nornir"],
-                value="Netmiko",
-                label="1. Select Python Library",
-                info="Which automation framework do you want to use?",
-            )
+            library = gr.Dropdown(["Netmiko", "NAPALM", "Nornir"], value="Netmiko", label="Library")
             device_type = gr.Dropdown(
-                choices=["cisco_ios", "cisco_nxos", "arista_eos", "juniper_junos"],
+                ["cisco_ios", "cisco_nxos", "arista_eos", "juniper_junos"],
                 value="cisco_ios",
-                label="2. Select Device OS",
-                info="Netmiko device_type; NAPALM driver is mapped automatically.",
-            )
-            user_goal = gr.Textbox(
-                lines=6,
-                label="3. Describe the Task",
-                placeholder=(
-                    "e.g., Log into 50 switches, backup running-config, and save "
-                    "them as text files with hostname and date."
-                ),
+                label="Device OS",
             )
-
-            gr.Examples(
-                examples=EXAMPLES,
-                inputs=[user_goal, library, device_type],
-                label="Example tasks",
-            )
-
-            generate_btn = gr.Button(
-                "✨ Generate Code",
-                variant="primary",
-                size="lg",
-            )
-
+            user_goal = gr.Textbox(lines=5, label="Describe the task")
+            generate_btn = gr.Button("✨ Generate Code", variant="primary")
         with gr.Column(scale=2):
-            status_md = gr.Markdown("*Ready.*")
-            output_code = gr.Code(
-                label="Production-Ready Python Script",
-                language="python",
-                lines=28,
-                interactive=True,
-            )
-            with gr.Row():
-                download_btn = gr.DownloadButton(
-                    "⬇️ Download .py",
-                    variant="secondary",
-                )
-
-    # Hidden state for download filename content (same as code)
-    download_file = gr.File(visible=False)
+            output_code = gr.Code(label="Python Script", language="python", interactive=True)
 
     generate_btn.click(
-        fn=on_generate,
+        fn=generate_network_script,
         inputs=[user_goal, library, device_type],
-        outputs=[output_code, status_md, download_btn, download_btn],
-        show_progress="full",
+        outputs=output_code,
     )
 
-
-# Hugging Face Spaces
 if __name__ == "__main__":
-    app.queue(default_concurrency_limit=2)
-    app.launch(
-        server_name="0.0.0.0",
-        server_port=int(os.environ.get("PORT", 7860)),
-    )
+    app.launch(server_name="0.0.0.0", server_port=int(os.environ.get("PORT", 7860)))