Update app.py

app.py

@@ -1,5 +1,6 @@
 import json
 import textwrap
+import math
 from typing import Dict, Any, List, Tuple
 
 import gradio as gr
@@ -7,31 +8,21 @@ import requests
 import matplotlib.pyplot as plt
 from matplotlib.figure import Figure
 
-
-# -----------------------------
-# LLM CALL HELPERS
-# -----------------------------
+# ============================================================
+#  LLM CALLER — GPT-4.1 ONLY
+# ============================================================
 
 def call_chat_completion(
     api_key: str,
     base_url: str,
-    model: str,
     system_prompt: str,
     user_prompt: str,
-    max_completion_tokens: int = 1800,
+    model: str = "gpt-4.1",
+    max_completion_tokens: int = 2000,
 ) -> str:
-    """
-    OpenAI-compatible ChatCompletion caller.
 
-    - Uses `max_completion_tokens` (new OpenAI spec).
-    - Falls back to `max_tokens` for providers that still expect it.
-    - No temperature param (some models only allow default).
-    """
     if not api_key:
-        raise ValueError("API key is required.")
-
-    if not base_url:
-        base_url = "https://api.openai.com"
+        raise ValueError("Missing API key.")
 
     url = base_url.rstrip("/") + "/v1/chat/completions"
 
@@ -40,8 +31,7 @@ def call_chat_completion(
         "Content-Type": "application/json",
     }
 
-    # Primary payload using max_completion_tokens
-    new_payload = {
+    payload = {
         "model": model,
         "messages": [
             {"role": "system", "content": system_prompt},
@@ -50,538 +40,338 @@ def call_chat_completion(
         "max_completion_tokens": max_completion_tokens,
     }
 
-    resp = requests.post(url, headers=headers, json=new_payload, timeout=60)
+    resp = requests.post(url, headers=headers, json=payload, timeout=60)
 
-    # If provider doesn't support `max_completion_tokens`, try legacy `max_tokens`
+    # Fallback for providers requiring max_tokens
     if resp.status_code == 400 and "max_completion_tokens" in resp.text:
-        legacy_payload = {
-            "model": model,
-            "messages": [
-                {"role": "system", "content": system_prompt},
-                {"role": "user", "content": user_prompt},
-            ],
-            "max_tokens": max_completion_tokens,
-        }
-        resp = requests.post(url, headers=headers, json=legacy_payload, timeout=60)
+        payload["max_tokens"] = max_completion_tokens
+        payload.pop("max_completion_tokens", None)
+        resp = requests.post(url, headers=headers, json=payload, timeout=60)
 
     if resp.status_code != 200:
         raise RuntimeError(
-            f"LLM API error: {resp.status_code} - {resp.text[:400]}"
+            f"LLM API Error {resp.status_code}:\n{resp.text[:400]}"
         )
 
     data = resp.json()
     try:
         return data["choices"][0]["message"]["content"]
-    except Exception as e:
-        raise RuntimeError(f"Unexpected LLM response format: {e}\n\n{data}")
+    except:
+        raise RuntimeError(f"Malformed response:\n\n{json.dumps(data, indent=2)}")
 
 
-# -----------------------------
-# SOP GENERATION LOGIC
-# -----------------------------
+# ============================================================
+#  SOP PROMPT + JSON PARSER
+# ============================================================
 
 SOP_SYSTEM_PROMPT = """
-You are an expert operations consultant and technical writer.
-
-You generate clear, professional, implementation-ready Standard Operating Procedures (SOPs).
-
-You MUST respond strictly as JSON using this schema:
+You are an expert process engineer. Produce SOPs as JSON using:
 
 {
-  "title": "string",
-  "purpose": "string",
-  "scope": "string",
-  "definitions": ["string", ...],
-  "roles": [
-    {
-      "name": "string",
-      "responsibilities": ["string", ...]
-    }
-  ],
-  "prerequisites": ["string", ...],
+  "title": "",
+  "purpose": "",
+  "scope": "",
+  "definitions": [],
+  "roles": [{"name": "", "responsibilities": []}],
+  "prerequisites": [],
   "steps": [
     {
       "step_number": 1,
-      "title": "string",
-      "description": "string",
-      "owner_role": "string",
-      "inputs": ["string", ...],
-      "outputs": ["string", ...]
+      "title": "",
+      "description": "",
+      "owner_role": "",
+      "inputs": [],
+      "outputs": []
     }
   ],
-  "escalation": ["string", ...],
-  "metrics": ["string", ...],
-  "risks": ["string", ...],
-  "versioning": {
-    "version": "1.0",
-    "owner": "string",
-    "last_updated": "string"
-  }
+  "escalation": [],
+  "metrics": [],
+  "risks": [],
+  "versioning": {"version": "1.0","owner": "","last_updated": ""}
 }
-"""
 
+Return ONLY JSON.
+"""
 
-def build_user_prompt(
-    sop_title: str,
-    description: str,
-    industry: str,
-    tone: str,
-    detail_level: str,
-) -> str:
+def build_user_prompt(title, desc, industry, tone, detail):
     return f"""
-Process Title: {sop_title or "Untitled SOP"}
-Context: {description or "N/A"}
-Industry: {industry or "General"}
-Tone: {tone or "Professional"}
-Detail Level: {detail_level or "Standard"}
-
-Audience: Mid-career professionals.
+SOP Title: {title}
+Context: {desc}
+Industry: {industry}
+Tone: {tone}
+Detail Level: {detail}
+Audience: mid-career professionals.
 """
 
 
-def parse_sop_json(raw_text: str) -> Dict[str, Any]:
-    """Clean model output and extract JSON."""
-    txt = raw_text.strip()
-
-    # Strip markdown fences if present
+def parse_sop_json(raw: str) -> Dict[str, Any]:
+    txt = raw.strip()
     if txt.startswith("```"):
-        parts = txt.split("```")
-        txt = next((p for p in parts if "{" in p), parts[-1])
+        txt = txt.split("```")[1]
 
-    # Extract JSON between first '{' and last '}'
     first = txt.find("{")
     last = txt.rfind("}")
-    if first != -1 and last != -1:
-        txt = txt[first:last + 1]
-
-    return json.loads(txt)
+    return json.loads(txt[first:last+1])
 
 
 def sop_to_markdown(sop: Dict[str, Any]) -> str:
-    """Format JSON SOP into Markdown."""
 
-    def bullet(items: List[str]) -> str:
+    def bullet(items):
         if not items:
-            return "_None specified._"
+            return "_None provided._"
         return "\n".join(f"- {i}" for i in items)
 
-    md: List[str] = []
-
-    md.append(f"# {sop.get('title', 'Standard Operating Procedure')}\n")
+    md = []
+    md.append(f"# {sop.get('title','Untitled SOP')}\n")
 
-    md.append("## 1. Purpose")
-    md.append(sop.get("purpose", "N/A"))
+    md.append("## 1. Purpose\n" + sop.get("purpose","N/A"))
+    md.append("## 2. Scope\n" + sop.get("scope","N/A"))
 
-    md.append("\n## 2. Scope")
-    md.append(sop.get("scope", "N/A"))
+    md.append("## 3. Definitions\n" + bullet(sop.get("definitions", [])))
 
-    md.append("\n## 3. Definitions")
-    md.append(bullet(sop.get("definitions", [])))
-
-    md.append("\n## 4. Roles & Responsibilities")
+    md.append("## 4. Roles & Responsibilities")
     for r in sop.get("roles", []):
-        name = r.get("name", "Role")
-        md.append(f"### {name}")
+        md.append(f"### {r.get('name','Role')}")
         md.append(bullet(r.get("responsibilities", [])))
 
-    md.append("\n## 5. Prerequisites")
-    md.append(bullet(sop.get("prerequisites", [])))
-
-    md.append("\n## 6. Procedure (Step-by-Step)")
-    for step in sop.get("steps", []):
-        md.append(f"### Step {step.get('step_number', '?')}: {step.get('title', 'Step')}")
-        md.append(f"**Owner:** {step.get('owner_role', 'N/A')}")
-        md.append(step.get("description", ""))
-        md.append("**Inputs:**")
-        md.append(bullet(step.get("inputs", [])))
-        md.append("**Outputs:**")
-        md.append(bullet(step.get("outputs", [])))
+    md.append("## 5. Prerequisites\n" + bullet(sop.get("prerequisites", [])))
 
-    md.append("\n## 7. Escalation")
-    md.append(bullet(sop.get("escalation", [])))
+    md.append("## 6. Procedure")
+    for st in sop.get("steps", []):
+        md.append(f"### Step {st['step_number']}: {st['title']}")
+        md.append(f"**Owner:** {st['owner_role']}")
+        md.append(st["description"])
+        md.append("**Inputs:**\n" + bullet(st["inputs"]))
+        md.append("**Outputs:**\n" + bullet(st["outputs"]))
 
-    md.append("\n## 8. Metrics & Success Criteria")
-    md.append(bullet(sop.get("metrics", [])))
-
-    md.append("\n## 9. Risks & Controls")
-    md.append(bullet(sop.get("risks", [])))
+    md.append("## 7. Escalation\n" + bullet(sop.get("escalation", [])))
+    md.append("## 8. Metrics\n" + bullet(sop.get("metrics")))
+    md.append("## 9. Risks\n" + bullet(sop.get("risks")))
 
     v = sop.get("versioning", {})
-    md.append("\n## 10. Version Control")
-    md.append(f"- Version: {v.get('version', '1.0')}")
-    md.append(f"- Owner: {v.get('owner', 'N/A')}")
-    md.append(f"- Last Updated: {v.get('last_updated', 'N/A')}")
+    md.append("## 10. Version Control")
+    md.append(f"- Version: {v.get('version','1.0')}")
+    md.append(f"- Owner: {v.get('owner','N/A')}")
+    md.append(f"- Last Updated: {v.get('last_updated','N/A')}")
 
     return "\n\n".join(md)
 
 
-# -----------------------------
-# INFOGRAPHIC / DATA VISUAL
-# -----------------------------
+# ============================================================
+#  🆕 PERFECTED DIAGRAM — AUTO-SIZE CARDS
+# ============================================================
 
 def create_sop_steps_figure(sop: Dict[str, Any]) -> Figure:
-    """
-    Create a clearer, more readable infographic-style figure
-    showing the SOP steps as stacked cards.
-
-    - Large, legible fonts
-    - Number block on the left
-    - Wrapped description text
-    """
-
     steps = sop.get("steps", [])
-
-    # Empty state
     if not steps:
-        fig, ax = plt.subplots(figsize=(7, 2))
-        ax.text(
-            0.5,
-            0.5,
-            "No steps available to visualize.",
-            ha="center",
-            va="center",
-            fontsize=12,
-        )
+        fig, ax = plt.subplots(figsize=(6,2))
+        ax.text(0.5,0.5,"No steps to visualize.",ha="center",va="center")
         ax.axis("off")
-        fig.tight_layout()
         return fig
 
-    n = len(steps)
-
-    # Figure height scales with number of steps (capped)
-    fig_height = min(14, max(4, 1.6 * n))
-    fig, ax = plt.subplots(figsize=(9, fig_height))
+    # Dynamically compute figure height based on text amount
+    total_height = 0
+    block_heights = []
 
-    # Coordinate system: y from 0 (bottom) to n (top)
-    ax.set_xlim(0, 1)
-    ax.set_ylim(0, n)
+    for st in steps:
+        desc_lines = textwrap.wrap(st["description"], width=65)
+        num_lines = 2 + len(desc_lines)  # title + owner + description lines
+        block_h = 0.35 * num_lines
+        block_heights.append(block_h)
+        total_height += block_h + 0.3  # spacing
 
-    card_top_margin = 0.25
-    card_bottom_margin = 0.25
-    card_height = 1 - (card_top_margin + card_bottom_margin)
+    fig_height = min(18, max(5, total_height))
+    fig, ax = plt.subplots(figsize=(10, fig_height))
 
-    for idx, step in enumerate(steps):
-        # y coordinate from top down
-        row_top = n - idx - card_top_margin
-        row_bottom = row_top - card_height
-        center_y = (row_top + row_bottom) / 2
+    y = total_height
 
-        step_number = step.get("step_number", idx + 1)
-        title = step.get("title", f"Step {step_number}")
-        owner = step.get("owner_role", "")
-        desc = step.get("description", "")
+    for idx, st in enumerate(steps):
+        title = st["title"]
+        owner = st["owner_role"]
+        desc_lines = textwrap.wrap(st["description"], width=70)
+        block_h = block_heights[idx]
 
-        # Wrap description into multiple lines
-        desc_wrapped = textwrap.fill(desc, width=80)
-
-        # Card rectangle (full width)
-        card_x0 = 0.03
-        card_x1 = 0.97
-        card_width = card_x1 - card_x0
+        x0, x1 = 0.05, 0.95
 
         ax.add_patch(
             plt.Rectangle(
-                (card_x0, row_bottom),
-                card_width,
-                card_height,
+                (x0, y - block_h),
+                x1 - x0,
+                block_h,
                 fill=False,
-                linewidth=1.6,
+                linewidth=1.7
             )
         )
 
-        # Number block on the left
-        num_block_width = 0.08
-        num_block_x0 = card_x0
-        num_block_y0 = row_bottom
-        num_block_height = card_height
-
+        # Number box
+        nbw = 0.08
         ax.add_patch(
             plt.Rectangle(
-                (num_block_x0, num_block_y0),
-                num_block_width,
-                num_block_height,
+                (x0, y - block_h),
+                nbw,
+                block_h,
                 fill=False,
-                linewidth=1.4,
+                linewidth=1.5
             )
         )
 
         ax.text(
-            num_block_x0 + num_block_width / 2,
-            center_y,
-            str(step_number),
-            ha="center",
-            va="center",
-            fontsize=12,
-            fontweight="bold",
+            x0 + nbw/2,
+            y - block_h/2,
+            str(st["step_number"]),
+            ha="center", va="center",
+            fontsize=13, fontweight="bold"
         )
 
-        # Text block (title, owner, description)
-        text_x0 = num_block_x0 + num_block_width + 0.02
+        text_x = x0 + nbw + 0.02
 
         # Title
-        ax.text(
-            text_x0,
-            row_top - 0.08,
-            title,
-            ha="left",
-            va="top",
-            fontsize=12,
-            fontweight="bold",
-        )
-
-        # Owner line (optional)
-        if owner:
-            ax.text(
-                text_x0,
-                row_top - 0.28,
+        ax.text(text_x, y - 0.2,
+                title,
+                fontsize=12,
+                fontweight="bold",
+                ha="left", va="top")
+
+        # Owner
+        ax.text(text_x, y - 0.45,
                 f"Owner: {owner}",
-                ha="left",
-                va="top",
                 fontsize=10,
                 style="italic",
-            )
-            desc_y = row_top - 0.48
-        else:
-            desc_y = row_top - 0.3
+                ha="left", va="top")
 
         # Description (wrapped)
-        ax.text(
-            text_x0,
-            desc_y,
-            desc_wrapped,
-            ha="left",
-            va="top",
-            fontsize=9,
-        )
+        text_y = y - 0.75
+        for line in desc_lines:
+            ax.text(text_x, text_y, line, fontsize=9, ha="left", va="top")
+            text_y -= 0.28
+
+        y -= (block_h + 0.3)
 
     ax.axis("off")
     fig.tight_layout()
     return fig
 
 
-# -----------------------------
-# SAMPLE PRESETS
-# -----------------------------
+# ============================================================
+#  SAMPLE SCENARIOS
+# ============================================================
 
-SAMPLE_SOPS: Dict[str, Dict[str, str]] = {
-    "Volunteer Onboarding Workflow": {
-        "title": "Volunteer Onboarding Workflow",
-        "description": (
-            "Create a clear SOP for onboarding new volunteers at a youth-serving "
-            "nonprofit. Include background checks, orientation, training, and site placement."
-        ),
-        "industry": "Nonprofit / Youth Development",
+SAMPLES = {
+    "Volunteer Onboarding": {
+        "title": "Volunteer Onboarding",
+        "description": "Create SOP for onboarding volunteers: background checks, orientation, training, placement.",
+        "industry": "Nonprofit"
     },
     "Remote Employee Onboarding": {
         "title": "Remote Employee Onboarding",
-        "description": (
-            "Design a remote onboarding SOP for new employees in a hybrid org, "
-            "covering IT setup, HR paperwork, culture onboarding, and 30-60-90 day milestones."
-        ),
-        "industry": "General / HR",
+        "description": "SOP for remote hires including IT setup, HR docs, culture onboarding.",
+        "industry": "HR"
     },
-    "IT Outage Incident Response": {
-        "title": "IT Outage Incident Response",
-        "description": (
-            "Create an SOP for responding to major IT outages affecting multiple sites, "
-            "including triage, communication, escalation, and post-mortem."
-        ),
-        "industry": "IT / Operations",
+    "IT Outage Response": {
+        "title": "IT Outage Response",
+        "description": "Major outage response: detection, triage, escalation, comms, restoration, post-mortem.",
+        "industry": "IT"
     },
 }
 
-
-def load_sample(sample_name: str) -> Tuple[str, str, str]:
-    if not sample_name or sample_name not in SAMPLE_SOPS:
+def load_sample(name):
+    if name not in SAMPLES:
         return "", "", "General"
-    s = SAMPLE_SOPS[sample_name]
+    s = SAMPLES[name]
     return s["title"], s["description"], s["industry"]
 
 
-# -----------------------------
-# MAIN HANDLER (CALLED BY UI)
-# -----------------------------
+# ============================================================
+#  MAIN GENERATOR
+# ============================================================
+
+def generate_sop(
+    api_key_state,
+    api_key_input,
+    base_url,
+    model,
+    title,
+    desc,
+    industry,
+    tone,
+    detail
+):
 
-def generate_sop_ui(
-    api_key_state: str,
-    api_key_input: str,
-    base_url: str,
-    model: str,
-    sop_title: str,
-    description: str,
-    industry: str,
-    tone: str,
-    detail_level: str,
-) -> Tuple[str, str, Figure, str]:
-    """Gradio event handler: generate SOP + JSON + figure."""
     api_key = api_key_input or api_key_state
     if not api_key:
-        return (
-            "⚠️ Please enter an API key.",
-            "",
-            create_sop_steps_figure({"steps": []}),
-            api_key_state,
-        )
-
-    if not model:
-        model = "gpt-4.1-mini"
-
-    user_prompt = build_user_prompt(sop_title, description, industry, tone, detail_level)
+        return ("⚠️ Enter an API key.",
+                "",
+                create_sop_steps_figure({"steps": []}),
+                api_key_state)
 
     try:
+        user_prompt = build_user_prompt(title, desc, industry, tone, detail)
+
         raw = call_chat_completion(
             api_key=api_key,
             base_url=base_url,
-            model=model,
             system_prompt=SOP_SYSTEM_PROMPT,
             user_prompt=user_prompt,
-            max_completion_tokens=1800,
+            model="gpt-4.1",         # 🔥 Forced stable model
+            max_completion_tokens=2000
         )
 
         sop = parse_sop_json(raw)
         md = sop_to_markdown(sop)
         fig = create_sop_steps_figure(sop)
-        json_out = json.dumps(sop, indent=2, ensure_ascii=False)
+        json_out = json.dumps(sop, indent=2)
 
-        # Save key into session state
         return md, json_out, fig, api_key
 
     except Exception as e:
-        return (
-            f"❌ Error generating SOP:\n\n{e}",
-            "",
-            create_sop_steps_figure({"steps": []}),
-            api_key_state,
-        )
+        return (f"❌ Error generating SOP:\n{e}",
+                "",
+                create_sop_steps_figure({"steps": []}),
+                api_key_state)
 
 
-# -----------------------------
-# GRADIO UI
-# -----------------------------
+# ============================================================
+#  GRADIO UI
+# ============================================================
 
 with gr.Blocks(title="ZEN Simple SOP Builder") as demo:
-    gr.Markdown(
-        """
-# 🧭 ZEN Simple SOP Builder
 
-Generate clean, professional Standard Operating Procedures (SOPs) from a short description.  
-Perfect for mid-career professionals who need clarity, structure, and ownership — fast.
-
-1. Configure your API settings  
-2. Describe the process you want to document  
-3. Generate a full SOP + visual flow of the steps
-
-> Your API key stays in this browser session and is not logged to disk.
-"""
-    )
+    gr.Markdown("""
+# 🧭 ZEN Simple SOP Builder
+Generate clean SOPs + auto diagrams using **GPT-4.1**.
+""")
 
     api_key_state = gr.State("")
 
     with gr.Row():
         with gr.Column(scale=1):
-            gr.Markdown("### Step 1 — API & Model Settings")
+            api_input = gr.Textbox("API Key", type="password")
+            base_url = gr.Textbox("Base URL", value="https://api.openai.com")
+            model_name = gr.Textbox("Model (GPT-4.1 only)", value="gpt-4.1")
 
-            api_key_input = gr.Textbox(
-                label="LLM API Key",
-                placeholder="Enter your API key (OpenAI or compatible provider)",
-                type="password",
-            )
-
-            base_url = gr.Textbox(
-                label="Base URL",
-                value="https://api.openai.com",
-                placeholder="e.g. https://api.openai.com or your custom endpoint",
-            )
-
-            model_name = gr.Textbox(
-                label="Model Name",
-                value="gpt-4.1-mini",
-                placeholder="e.g. gpt-4.1, gpt-4o, deepseek-chat, mistral-large, etc.",
-            )
-
-            gr.Markdown("### Step 2 — Try a Sample Scenario")
-            sample_dropdown = gr.Dropdown(
-                label="Sample SOPs",
-                choices=list(SAMPLE_SOPS.keys()),
-                value=None,
-                info="Optional: load a predefined example.",
-            )
-            load_button = gr.Button("Load Sample into Form")
+            sample = gr.Dropdown("Sample SOP", choices=list(SAMPLES.keys()))
+            load_btn = gr.Button("Load Sample")
 
         with gr.Column(scale=2):
-            gr.Markdown("### Step 3 — Describe Your SOP")
-
-            sop_title = gr.Textbox(
-                label="SOP Title",
-                placeholder="e.g. Volunteer Onboarding Workflow, IT Outage Response",
-            )
-
-            description = gr.Textbox(
-                label="Describe the process / context",
-                placeholder="What should this SOP cover? Who is it for? Any constraints?",
-                lines=6,
-            )
-
-            industry = gr.Textbox(
-                label="Industry / Domain",
-                value="General",
-                placeholder="e.g. Nonprofit, HR, Education, Healthcare, IT",
-            )
-
-            tone = gr.Dropdown(
-                label="Tone",
-                choices=["Professional", "Executive", "Supportive", "Direct", "Compliance-focused"],
-                value="Professional",
-            )
-
-            detail_level = gr.Dropdown(
-                label="Detail Level",
-                choices=["Standard", "High detail", "Checklist-style", "Overview only"],
-                value="Standard",
-            )
+            title = gr.Textbox("SOP Title")
+            desc = gr.Textbox("Description", lines=5)
+            industry = gr.Textbox("Industry", value="General")
+            tone = gr.Dropdown("Tone", ["Professional","Executive","Supportive"], value="Professional")
+            detail = gr.Dropdown("Detail Level", ["Standard","High detail","Checklist"], value="Standard")
 
-            generate_button = gr.Button("🚀 Generate SOP", variant="primary")
+            gen_btn = gr.Button("🚀 Generate SOP", variant="primary")
 
-    gr.Markdown("### Step 4 — Results")
+    sop_md = gr.Markdown()
+    sop_json = gr.Code(language="json")
+    sop_fig = gr.Plot()
 
-    with gr.Row():
-        with gr.Column(scale=3):
-            sop_output = gr.Markdown(
-                label="Generated SOP",
-                value="Your SOP will appear here.",
-            )
-        with gr.Column(scale=2):
-            sop_json_output = gr.Code(
-                label="Raw SOP JSON (for automation / export)",
-                language="json",
-            )
-
-    gr.Markdown("### Visual Flow of Steps")
-    sop_figure = gr.Plot(label="SOP Steps Diagram")
+    load_btn.click(load_sample, sample, [title, desc, industry])
 
-    # Wire up events
-    load_button.click(
-        fn=load_sample,
-        inputs=[sample_dropdown],
-        outputs=[sop_title, description, industry],
+    gen_btn.click(
+        generate_sop,
+        [api_key_state, api_input, base_url, model_name, title, desc, industry, tone, detail],
+        [sop_md, sop_json, sop_fig, api_key_state],
     )
 
-    generate_button.click(
-        fn=generate_sop_ui,
-        inputs=[
-            api_key_state,
-            api_key_input,
-            base_url,
-            model_name,
-            sop_title,
-            description,
-            industry,
-            tone,
-            detail_level,
-        ],
-        outputs=[sop_output, sop_json_output, sop_figure, api_key_state],
-    )
-
-
 if __name__ == "__main__":
     demo.launch()