Update app.py

app.py

@@ -1,6 +1,5 @@
 import json
 import textwrap
-import math
 from typing import Dict, Any, List, Tuple
 
 import gradio as gr
@@ -8,21 +7,31 @@ import requests
 import matplotlib.pyplot as plt
 from matplotlib.figure import Figure
 
-# ============================================================
-#  LLM CALLER — GPT-4.1 ONLY
-# ============================================================
+
+# -----------------------------
+# LLM CALL HELPERS
+# -----------------------------
 
 def call_chat_completion(
     api_key: str,
     base_url: str,
+    model: str,
     system_prompt: str,
     user_prompt: str,
-    model: str = "gpt-4.1",
-    max_completion_tokens: int = 2000,
+    max_completion_tokens: int = 1800,
 ) -> 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("Missing API key.")
+        raise ValueError("API key is required.")
+
+    if not base_url:
+        base_url = "https://api.openai.com"
 
     url = base_url.rstrip("/") + "/v1/chat/completions"
 
@@ -31,7 +40,8 @@ def call_chat_completion(
         "Content-Type": "application/json",
     }
 
-    payload = {
+    # Primary payload using max_completion_tokens
+    new_payload = {
         "model": model,
         "messages": [
             {"role": "system", "content": system_prompt},
@@ -40,338 +50,538 @@ def call_chat_completion(
         "max_completion_tokens": max_completion_tokens,
     }
 
-    resp = requests.post(url, headers=headers, json=payload, timeout=60)
+    resp = requests.post(url, headers=headers, json=new_payload, timeout=60)
 
-    # Fallback for providers requiring max_tokens
+    # If provider doesn't support `max_completion_tokens`, try legacy `max_tokens`
     if resp.status_code == 400 and "max_completion_tokens" in resp.text:
-        payload["max_tokens"] = max_completion_tokens
-        payload.pop("max_completion_tokens", None)
-        resp = requests.post(url, headers=headers, json=payload, timeout=60)
+        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)
 
     if resp.status_code != 200:
         raise RuntimeError(
-            f"LLM API Error {resp.status_code}:\n{resp.text[:400]}"
+            f"LLM API error: {resp.status_code} - {resp.text[:400]}"
         )
 
     data = resp.json()
     try:
         return data["choices"][0]["message"]["content"]
-    except:
-        raise RuntimeError(f"Malformed response:\n\n{json.dumps(data, indent=2)}")
+    except Exception as e:
+        raise RuntimeError(f"Unexpected LLM response format: {e}\n\n{data}")
 
 
-# ============================================================
-#  SOP PROMPT + JSON PARSER
-# ============================================================
+# -----------------------------
+# SOP GENERATION LOGIC
+# -----------------------------
 
 SOP_SYSTEM_PROMPT = """
-You are an expert process engineer. Produce SOPs as JSON using:
+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:
 
 {
-  "title": "",
-  "purpose": "",
-  "scope": "",
-  "definitions": [],
-  "roles": [{"name": "", "responsibilities": []}],
-  "prerequisites": [],
+  "title": "string",
+  "purpose": "string",
+  "scope": "string",
+  "definitions": ["string", ...],
+  "roles": [
+    {
+      "name": "string",
+      "responsibilities": ["string", ...]
+    }
+  ],
+  "prerequisites": ["string", ...],
   "steps": [
     {
       "step_number": 1,
-      "title": "",
-      "description": "",
-      "owner_role": "",
-      "inputs": [],
-      "outputs": []
+      "title": "string",
+      "description": "string",
+      "owner_role": "string",
+      "inputs": ["string", ...],
+      "outputs": ["string", ...]
     }
   ],
-  "escalation": [],
-  "metrics": [],
-  "risks": [],
-  "versioning": {"version": "1.0","owner": "","last_updated": ""}
+  "escalation": ["string", ...],
+  "metrics": ["string", ...],
+  "risks": ["string", ...],
+  "versioning": {
+    "version": "1.0",
+    "owner": "string",
+    "last_updated": "string"
+  }
 }
-
-Return ONLY JSON.
 """
 
-def build_user_prompt(title, desc, industry, tone, detail):
+
+def build_user_prompt(
+    sop_title: str,
+    description: str,
+    industry: str,
+    tone: str,
+    detail_level: str,
+) -> str:
     return f"""
-SOP Title: {title}
-Context: {desc}
-Industry: {industry}
-Tone: {tone}
-Detail Level: {detail}
-Audience: mid-career professionals.
+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.
 """
 
 
-def parse_sop_json(raw: str) -> Dict[str, Any]:
-    txt = raw.strip()
+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
     if txt.startswith("```"):
-        txt = txt.split("```")[1]
+        parts = txt.split("```")
+        txt = next((p for p in parts if "{" in p), parts[-1])
 
+    # Extract JSON between first '{' and last '}'
     first = txt.find("{")
     last = txt.rfind("}")
-    return json.loads(txt[first:last+1])
+    if first != -1 and last != -1:
+        txt = txt[first:last + 1]
+
+    return json.loads(txt)
 
 
 def sop_to_markdown(sop: Dict[str, Any]) -> str:
+    """Format JSON SOP into Markdown."""
 
-    def bullet(items):
+    def bullet(items: List[str]) -> str:
         if not items:
-            return "_None provided._"
+            return "_None specified._"
         return "\n".join(f"- {i}" for i in items)
 
-    md = []
-    md.append(f"# {sop.get('title','Untitled SOP')}\n")
+    md: List[str] = []
+
+    md.append(f"# {sop.get('title', 'Standard Operating Procedure')}\n")
 
-    md.append("## 1. Purpose\n" + sop.get("purpose","N/A"))
-    md.append("## 2. Scope\n" + sop.get("scope","N/A"))
+    md.append("## 1. Purpose")
+    md.append(sop.get("purpose", "N/A"))
 
-    md.append("## 3. Definitions\n" + bullet(sop.get("definitions", [])))
+    md.append("\n## 2. Scope")
+    md.append(sop.get("scope", "N/A"))
 
-    md.append("## 4. Roles & Responsibilities")
+    md.append("\n## 3. Definitions")
+    md.append(bullet(sop.get("definitions", [])))
+
+    md.append("\n## 4. Roles & Responsibilities")
     for r in sop.get("roles", []):
-        md.append(f"### {r.get('name','Role')}")
+        name = r.get("name", "Role")
+        md.append(f"### {name}")
         md.append(bullet(r.get("responsibilities", [])))
 
-    md.append("## 5. Prerequisites\n" + bullet(sop.get("prerequisites", [])))
+    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("## 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## 7. Escalation")
+    md.append(bullet(sop.get("escalation", [])))
 
-    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")))
+    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", [])))
 
     v = sop.get("versioning", {})
-    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')}")
+    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')}")
 
     return "\n\n".join(md)
 
 
-# ============================================================
-#  🆕 PERFECTED DIAGRAM — AUTO-SIZE CARDS
-# ============================================================
+# -----------------------------
+# INFOGRAPHIC / DATA VISUAL
+# -----------------------------
 
 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=(6,2))
-        ax.text(0.5,0.5,"No steps to visualize.",ha="center",va="center")
+        fig, ax = plt.subplots(figsize=(7, 2))
+        ax.text(
+            0.5,
+            0.5,
+            "No steps available to visualize.",
+            ha="center",
+            va="center",
+            fontsize=12,
+        )
         ax.axis("off")
+        fig.tight_layout()
         return fig
 
-    # Dynamically compute figure height based on text amount
-    total_height = 0
-    block_heights = []
+    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))
 
-    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
+    # Coordinate system: y from 0 (bottom) to n (top)
+    ax.set_xlim(0, 1)
+    ax.set_ylim(0, n)
 
-    fig_height = min(18, max(5, total_height))
-    fig, ax = plt.subplots(figsize=(10, fig_height))
+    card_top_margin = 0.25
+    card_bottom_margin = 0.25
+    card_height = 1 - (card_top_margin + card_bottom_margin)
 
-    y = total_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
 
-    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]
+        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", "")
 
-        x0, x1 = 0.05, 0.95
+        # 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
 
         ax.add_patch(
             plt.Rectangle(
-                (x0, y - block_h),
-                x1 - x0,
-                block_h,
+                (card_x0, row_bottom),
+                card_width,
+                card_height,
                 fill=False,
-                linewidth=1.7
+                linewidth=1.6,
             )
         )
 
-        # Number box
-        nbw = 0.08
+        # 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
+
         ax.add_patch(
             plt.Rectangle(
-                (x0, y - block_h),
-                nbw,
-                block_h,
+                (num_block_x0, num_block_y0),
+                num_block_width,
+                num_block_height,
                 fill=False,
-                linewidth=1.5
+                linewidth=1.4,
             )
         )
 
         ax.text(
-            x0 + nbw/2,
-            y - block_h/2,
-            str(st["step_number"]),
-            ha="center", va="center",
-            fontsize=13, fontweight="bold"
+            num_block_x0 + num_block_width / 2,
+            center_y,
+            str(step_number),
+            ha="center",
+            va="center",
+            fontsize=12,
+            fontweight="bold",
         )
 
-        text_x = x0 + nbw + 0.02
+        # Text block (title, owner, description)
+        text_x0 = num_block_x0 + num_block_width + 0.02
 
         # Title
-        ax.text(text_x, y - 0.2,
-                title,
-                fontsize=12,
-                fontweight="bold",
-                ha="left", va="top")
-
-        # Owner
-        ax.text(text_x, y - 0.45,
+        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,
                 f"Owner: {owner}",
+                ha="left",
+                va="top",
                 fontsize=10,
                 style="italic",
-                ha="left", va="top")
+            )
+            desc_y = row_top - 0.48
+        else:
+            desc_y = row_top - 0.3
 
         # Description (wrapped)
-        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.text(
+            text_x0,
+            desc_y,
+            desc_wrapped,
+            ha="left",
+            va="top",
+            fontsize=9,
+        )
 
     ax.axis("off")
     fig.tight_layout()
     return fig
 
 
-# ============================================================
-#  SAMPLE SCENARIOS
-# ============================================================
+# -----------------------------
+# SAMPLE PRESETS
+# -----------------------------
 
-SAMPLES = {
-    "Volunteer Onboarding": {
-        "title": "Volunteer Onboarding",
-        "description": "Create SOP for onboarding volunteers: background checks, orientation, training, placement.",
-        "industry": "Nonprofit"
+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",
     },
     "Remote Employee Onboarding": {
         "title": "Remote Employee Onboarding",
-        "description": "SOP for remote hires including IT setup, HR docs, culture onboarding.",
-        "industry": "HR"
+        "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",
     },
-    "IT Outage Response": {
-        "title": "IT Outage Response",
-        "description": "Major outage response: detection, triage, escalation, comms, restoration, post-mortem.",
-        "industry": "IT"
+    "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",
     },
 }
 
-def load_sample(name):
-    if name not in SAMPLES:
+
+def load_sample(sample_name: str) -> Tuple[str, str, str]:
+    if not sample_name or sample_name not in SAMPLE_SOPS:
         return "", "", "General"
-    s = SAMPLES[name]
+    s = SAMPLE_SOPS[sample_name]
     return s["title"], s["description"], s["industry"]
 
 
-# ============================================================
-#  MAIN GENERATOR
-# ============================================================
-
-def generate_sop(
-    api_key_state,
-    api_key_input,
-    base_url,
-    model,
-    title,
-    desc,
-    industry,
-    tone,
-    detail
-):
+# -----------------------------
+# MAIN HANDLER (CALLED BY UI)
+# -----------------------------
 
+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 ("⚠️ Enter an API key.",
-                "",
-                create_sop_steps_figure({"steps": []}),
-                api_key_state)
+        return (
+            "⚠️ Please enter an API key.",
+            "",
+            create_sop_steps_figure({"steps": []}),
+            api_key_state,
+        )
 
-    try:
-        user_prompt = build_user_prompt(title, desc, industry, tone, detail)
+    if not model:
+        model = "gpt-4.1-mini"
+
+    user_prompt = build_user_prompt(sop_title, description, industry, tone, detail_level)
 
+    try:
         raw = call_chat_completion(
             api_key=api_key,
             base_url=base_url,
+            model=model,
             system_prompt=SOP_SYSTEM_PROMPT,
             user_prompt=user_prompt,
-            model="gpt-4.1",         # 🔥 Forced stable model
-            max_completion_tokens=2000
+            max_completion_tokens=1800,
         )
 
         sop = parse_sop_json(raw)
         md = sop_to_markdown(sop)
         fig = create_sop_steps_figure(sop)
-        json_out = json.dumps(sop, indent=2)
+        json_out = json.dumps(sop, indent=2, ensure_ascii=False)
 
+        # Save key into session state
         return md, json_out, fig, api_key
 
     except Exception as e:
-        return (f"❌ Error generating SOP:\n{e}",
-                "",
-                create_sop_steps_figure({"steps": []}),
-                api_key_state)
+        return (
+            f"❌ Error generating SOP:\n\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("""
+    gr.Markdown(
+        """
 # 🧭 ZEN Simple SOP Builder
-Generate clean SOPs + auto diagrams using **GPT-4.1**.
-""")
+
+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.
+"""
+    )
 
     api_key_state = gr.State("")
 
     with gr.Row():
         with gr.Column(scale=1):
-            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")
+            gr.Markdown("### Step 1 — API & Model Settings")
 
-            sample = gr.Dropdown("Sample SOP", choices=list(SAMPLES.keys()))
-            load_btn = gr.Button("Load Sample")
+            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")
 
         with gr.Column(scale=2):
-            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")
+            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",
+            )
 
-            gen_btn = gr.Button("🚀 Generate SOP", variant="primary")
+            generate_button = gr.Button("🚀 Generate SOP", variant="primary")
 
-    sop_md = gr.Markdown()
-    sop_json = gr.Code(language="json")
-    sop_fig = gr.Plot()
+    gr.Markdown("### Step 4 — Results")
 
-    load_btn.click(load_sample, sample, [title, desc, industry])
+    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")
 
-    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],
+    # Wire up events
+    load_button.click(
+        fn=load_sample,
+        inputs=[sample_dropdown],
+        outputs=[sop_title, description, industry],
     )
 
+    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()