Update app.py

app.py

@@ -8,9 +8,9 @@ import matplotlib.pyplot as plt
 from matplotlib.figure import Figure
 
 
-# -----------------------------
-# LLM CALL HELPERS
-# -----------------------------
+# ============================================================
+#  LLM CALLER (GPT-4.1 BY DEFAULT)
+# ============================================================
 
 def call_chat_completion(
     api_key: str,
@@ -18,14 +18,14 @@ def call_chat_completion(
     model: str,
     system_prompt: str,
     user_prompt: str,
-    max_completion_tokens: int = 1800,
+    max_completion_tokens: int = 2000,
 ) -> str:
     """
-    OpenAI-compatible ChatCompletion caller.
+    OpenAI-compatible chat completion call.
 
-    - 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).
+    - Uses new-style `max_completion_tokens` (for GPT-4.1, GPT-4o, etc.)
+    - Falls back to `max_tokens` if the provider doesn't support it.
+    - No temperature / top_p to avoid incompatibility with some models.
     """
     if not api_key:
         raise ValueError("API key is required.")
@@ -40,8 +40,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,19 +49,13 @@ 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 still expecting `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.pop("max_completion_tokens", None)
+        payload["max_tokens"] = max_completion_tokens
+        resp = requests.post(url, headers=headers, json=payload, timeout=60)
 
     if resp.status_code != 200:
         raise RuntimeError(
@@ -73,19 +66,15 @@ def call_chat_completion(
     try:
         return data["choices"][0]["message"]["content"]
     except Exception as e:
-        raise RuntimeError(f"Unexpected LLM response format: {e}\n\n{data}")
+        raise RuntimeError(f"Unexpected LLM response format: {e}\n\n{json.dumps(data, indent=2)}") from e
 
 
-# -----------------------------
-# SOP GENERATION LOGIC
-# -----------------------------
+# ============================================================
+#  SOP PROMPT + PARSING
+# ============================================================
 
 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 strictly as JSON with this schema:
 
 {
   "title": "string",
@@ -118,8 +107,9 @@ You MUST respond strictly as JSON using this schema:
     "last_updated": "string"
   }
 }
-"""
 
+Return ONLY JSON. No explanation or commentary.
+"""
 
 def build_user_prompt(
     sop_title: str,
@@ -129,43 +119,42 @@ def build_user_prompt(
     detail_level: str,
 ) -> str:
     return f"""
-Process Title: {sop_title or "Untitled SOP"}
+SOP 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.
-"""
+Audience: mid-career professionals who need clarity and accountability.
+""".strip()
 
 
 def parse_sop_json(raw_text: str) -> Dict[str, Any]:
-    """Clean model output and extract JSON."""
+    """Extract JSON from LLM output, stripping code fences if present."""
     txt = raw_text.strip()
 
-    # Strip markdown fences if present
     if txt.startswith("```"):
         parts = txt.split("```")
-        txt = next((p for p in parts if "{" in p), parts[-1])
+        # choose the first part that looks like JSON
+        txt = next((p for p in parts if "{" in p and "}" in p), parts[-1])
 
-    # Extract JSON between first '{' and last '}'
     first = txt.find("{")
     last = txt.rfind("}")
-    if first != -1 and last != -1:
-        txt = txt[first:last + 1]
+    if first == -1 or last == -1:
+        raise ValueError("No JSON object detected in model output.")
+    txt = txt[first:last + 1]
 
     return json.loads(txt)
 
 
 def sop_to_markdown(sop: Dict[str, Any]) -> str:
-    """Format JSON SOP into Markdown."""
+    """Render SOP JSON → readable Markdown document."""
 
-    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 = []
 
     md.append(f"# {sop.get('title', 'Standard Operating Procedure')}\n")
 
@@ -179,10 +168,9 @@ def sop_to_markdown(sop: Dict[str, Any]) -> str:
     md.append(bullet(sop.get("definitions", [])))
 
     md.append("\n## 4. Roles & Responsibilities")
-    for r in sop.get("roles", []):
-        name = r.get("name", "Role")
-        md.append(f"### {name}")
-        md.append(bullet(r.get("responsibilities", [])))
+    for role in sop.get("roles", []):
+        md.append(f"### {role.get('name', 'Role')}")
+        md.append(bullet(role.get("responsibilities", [])))
 
     md.append("\n## 5. Prerequisites")
     md.append(bullet(sop.get("prerequisites", [])))
@@ -200,10 +188,10 @@ def sop_to_markdown(sop: Dict[str, Any]) -> str:
     md.append("\n## 7. Escalation")
     md.append(bullet(sop.get("escalation", [])))
 
-    md.append("\n## 8. Metrics & Success Criteria")
+    md.append("\n## 8. Metrics")
     md.append(bullet(sop.get("metrics", [])))
 
-    md.append("\n## 9. Risks & Controls")
+    md.append("\n## 9. Risks")
     md.append(bullet(sop.get("risks", [])))
 
     v = sop.get("versioning", {})
@@ -215,182 +203,166 @@ def sop_to_markdown(sop: Dict[str, Any]) -> str:
     return "\n\n".join(md)
 
 
-# -----------------------------
-# INFOGRAPHIC / DATA VISUAL
-# -----------------------------
+# ============================================================
+#  IMPROVED DIAGRAM — AUTO-SIZED CARDS, NO OVERFLOW
+# ============================================================
 
 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
+    Draw each step as a stacked card with:
+    - dynamic height based on description length
+    - number block on the left
+    - title + owner + wrapped description inside card
     """
 
     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,
-        )
+        ax.text(0.5, 0.5, "No steps available 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))
-
-    # Coordinate system: y from 0 (bottom) to n (top)
+    # First pass: determine required height for each card
+    card_heights = []
+    total_height = 0.0
+
+    for step in steps:
+        desc_lines = textwrap.wrap(step.get("description", ""), width=70)
+        # base height (title + owner) + 0.3 per line of description
+        base = 1.0  # title + owner + padding
+        per_line = 0.32
+        h = base + per_line * max(len(desc_lines), 1)
+        h += 0.3  # bottom padding
+        card_heights.append(h)
+        total_height += h
+
+    # Add spacing between cards
+    spacing = 0.4
+    total_height += spacing * (len(steps) + 1)
+
+    fig_height = min(20, max(5, total_height))
+    fig, ax = plt.subplots(figsize=(10, fig_height))
     ax.set_xlim(0, 1)
-    ax.set_ylim(0, n)
+    ax.set_ylim(0, total_height)
 
-    card_top_margin = 0.25
-    card_bottom_margin = 0.25
-    card_height = 1 - (card_top_margin + card_bottom_margin)
+    y = total_height - spacing  # start from top
 
-    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 step, h in zip(steps, card_heights):
+        y_bottom = y - h
+        y_top = y
 
-        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", "")
-
-        # 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
+        # Card boundaries
+        x0 = 0.05
+        x1 = 0.95
 
+        # Draw outer card
         ax.add_patch(
             plt.Rectangle(
-                (card_x0, row_bottom),
-                card_width,
-                card_height,
+                (x0, y_bottom),
+                x1 - x0,
+                h,
                 fill=False,
-                linewidth=1.6,
+                linewidth=1.8,
             )
         )
 
-        # 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 block
+        num_block_w = 0.08
         ax.add_patch(
             plt.Rectangle(
-                (num_block_x0, num_block_y0),
-                num_block_width,
-                num_block_height,
+                (x0, y_bottom),
+                num_block_w,
+                h,
                 fill=False,
-                linewidth=1.4,
+                linewidth=1.6,
             )
         )
 
+        # Step number text in the center of the number block
         ax.text(
-            num_block_x0 + num_block_width / 2,
-            center_y,
-            str(step_number),
+            x0 + num_block_w / 2,
+            y_bottom + h / 2,
+            str(step.get("step_number", "?")),
             ha="center",
             va="center",
-            fontsize=12,
+            fontsize=13,
             fontweight="bold",
         )
 
-        # Text block (title, owner, description)
-        text_x0 = num_block_x0 + num_block_width + 0.02
+        # Text area start
+        text_x = x0 + num_block_w + 0.02
 
         # Title
         ax.text(
-            text_x0,
-            row_top - 0.08,
-            title,
+            text_x,
+            y_top - 0.25,
+            step.get("title", ""),
             ha="left",
             va="top",
             fontsize=12,
             fontweight="bold",
         )
 
-        # Owner line (optional)
+        # Owner
+        owner = step.get("owner_role", "")
         if owner:
+            owner_y = y_top - 0.55
             ax.text(
-                text_x0,
-                row_top - 0.28,
+                text_x,
+                owner_y,
                 f"Owner: {owner}",
                 ha="left",
                 va="top",
                 fontsize=10,
                 style="italic",
             )
-            desc_y = row_top - 0.48
         else:
-            desc_y = row_top - 0.3
+            owner_y = y_top - 0.5
 
         # Description (wrapped)
-        ax.text(
-            text_x0,
-            desc_y,
-            desc_wrapped,
-            ha="left",
-            va="top",
-            fontsize=9,
-        )
+        desc_lines = textwrap.wrap(step.get("description", ""), width=70)
+        desc_y = owner_y - 0.4
+        for line in desc_lines:
+            ax.text(
+                text_x,
+                desc_y,
+                line,
+                ha="left",
+                va="top",
+                fontsize=9,
+            )
+            desc_y -= 0.3  # vertical spacing per line
+
+        y = y_bottom - spacing  # move down for next card
 
     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."
-        ),
+    "Volunteer Onboarding": {
+        "title": "Volunteer Onboarding",
+        "description": "Onboard new volunteers including application review, background checks, orientation, training, and site placement.",
         "industry": "Nonprofit / Youth Development",
     },
     "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": "Design a remote onboarding SOP for hybrid employees including IT setup, HR paperwork, and culture onboarding.",
+        "industry": "HR / General",
     },
-    "IT Outage Incident Response": {
+    "IT Outage 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."
-        ),
+        "description": "Major outage response SOP including detection, triage, escalation, communication, restoration, and post-mortem.",
         "industry": "IT / Operations",
     },
 }
 
-
 def load_sample(sample_name: str) -> Tuple[str, str, str]:
     if not sample_name or sample_name not in SAMPLE_SOPS:
         return "", "", "General"
@@ -398,33 +370,32 @@ def load_sample(sample_name: str) -> Tuple[str, str, str]:
     return s["title"], s["description"], s["industry"]
 
 
-# -----------------------------
-# MAIN HANDLER (CALLED BY UI)
-# -----------------------------
+# ============================================================
+#  MAIN HANDLER FOR GRADIO
+# ============================================================
 
 def generate_sop_ui(
     api_key_state: str,
     api_key_input: str,
     base_url: str,
-    model: str,
+    model_name: 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.",
+            "⚠️ Please enter your API key in the left panel.",
             "",
             create_sop_steps_figure({"steps": []}),
             api_key_state,
         )
 
-    if not model:
-        model = "gpt-4.1-mini"
+    model = model_name or "gpt-4.1"
 
     user_prompt = build_user_prompt(sop_title, description, industry, tone, detail_level)
 
@@ -435,7 +406,7 @@ def generate_sop_ui(
             model=model,
             system_prompt=SOP_SYSTEM_PROMPT,
             user_prompt=user_prompt,
-            max_completion_tokens=1800,
+            max_completion_tokens=2000,
         )
 
         sop = parse_sop_json(raw)
@@ -443,8 +414,7 @@ def generate_sop_ui(
         fig = create_sop_steps_figure(sop)
         json_out = json.dumps(sop, indent=2, ensure_ascii=False)
 
-        # Save key into session state
-        return md, json_out, fig, api_key
+        return md, json_out, fig, api_key  # persist key in session state
 
     except Exception as e:
         return (
@@ -455,65 +425,65 @@ def generate_sop_ui(
         )
 
 
-# -----------------------------
-# 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
+Generate clean, professional Standard Operating Procedures (SOPs) from a short description,  
+plus an auto-generated visual diagram of the steps.
 
-> Your API key stays in this browser session and is not logged to disk.
+Powered by your own API key (GPT-4.1 by default).
 """
     )
 
     api_key_state = gr.State("")
 
     with gr.Row():
+        # LEFT COLUMN — API + Samples
         with gr.Column(scale=1):
             gr.Markdown("### Step 1 — API & Model Settings")
 
             api_key_input = gr.Textbox(
                 label="LLM API Key",
-                placeholder="Enter your API key (OpenAI or compatible provider)",
+                placeholder="Enter your OpenAI (or compatible) API key",
                 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",
+                placeholder="e.g. https://api.openai.com or custom OpenAI-compatible 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.",
+                value="gpt-4.1",
+                placeholder="e.g. gpt-4.1, gpt-4o, etc.",
             )
 
-            gr.Markdown("### Step 2 — Try a Sample Scenario")
+            gr.Markdown("### Load a Sample SOP")
+
             sample_dropdown = gr.Dropdown(
-                label="Sample SOPs",
+                label="Sample scenarios",
                 choices=list(SAMPLE_SOPS.keys()),
                 value=None,
-                info="Optional: load a predefined example.",
+                info="Optional: load a ready-made example to test the tool.",
             )
+
             load_button = gr.Button("Load Sample into Form")
 
+        # RIGHT COLUMN — SOP Description
         with gr.Column(scale=2):
-            gr.Markdown("### Step 3 — Describe Your SOP")
+            gr.Markdown("### Step 2 — Describe the SOP")
 
             sop_title = gr.Textbox(
                 label="SOP Title",
-                placeholder="e.g. Volunteer Onboarding Workflow, IT Outage Response",
+                placeholder="e.g. Volunteer Onboarding Workflow",
             )
 
             description = gr.Textbox(
@@ -542,13 +512,13 @@ Perfect for mid-career professionals who need clarity, structure, and ownership
 
             generate_button = gr.Button("🚀 Generate SOP", variant="primary")
 
-    gr.Markdown("### Step 4 — Results")
+    gr.Markdown("### Step 3 — Generated SOP")
 
     with gr.Row():
         with gr.Column(scale=3):
             sop_output = gr.Markdown(
-                label="Generated SOP",
-                value="Your SOP will appear here.",
+                label="SOP (Markdown)",
+                value="Your SOP will appear here after generation.",
             )
         with gr.Column(scale=2):
             sop_json_output = gr.Code(
@@ -556,10 +526,10 @@ Perfect for mid-career professionals who need clarity, structure, and ownership
                 language="json",
             )
 
-    gr.Markdown("### Visual Flow of Steps")
+    gr.Markdown("### Step 4 — Visual Workflow Diagram")
     sop_figure = gr.Plot(label="SOP Steps Diagram")
 
-    # Wire up events
+    # Wire up actions
     load_button.click(
         fn=load_sample,
         inputs=[sample_dropdown],