Update main.py

main.py

@@ -202,6 +202,92 @@ def _ensure_numbering(text: str) -> str:
         out.append(f"{marker} {seg}")
     return "\n".join(out)
 
+# --- Next-step helpers (generic; SOP-agnostic) ---
+
+def _norm_text(s: str) -> str:
+    import re
+    s = (s or "").lower()
+    s = re.sub(r"[^\w\s]", " ", s)
+    s = re.sub(r"\s+", " ", s).strip()
+    return s
+
+def _split_sop_into_steps(numbered_text: str) -> list:
+    """
+    Split a numbered/bulleted SOP block (already passed through _ensure_numbering)
+    into atomic steps. Returns a list of raw step strings (order preserved).
+    Safe for circled digits, '1.' styles, and bullets.
+    """
+    lines = [ln.strip() for ln in (numbered_text or "").splitlines() if ln.strip()]
+    steps = []
+    for ln in lines:
+        # Strip circled/number/bullet marker
+        cleaned = ln
+        cleaned = re.sub(r"^\s*(?:[\u2460-\u2473]|\d+[.)]|[-*•])\s*", "", cleaned)
+        if cleaned:
+            steps.append(cleaned)
+    return steps
+
+def _soft_match_score(a: str, b: str) -> float:
+    # Simple Jaccard-like score on tokens for fuzzy matching
+    ta = set(_norm_text(a).split())
+    tb = set(_norm_text(b).split())
+    if not ta or not tb:
+        return 0.0
+    inter = len(ta & tb)
+    union = len(ta | tb)
+    return inter / union if union else 0.0
+
+def _detect_next_intent(user_query: str) -> bool:
+    q = _norm_text(user_query)
+    # Conservative rules to avoid false triggers
+    keys = [
+        "after", "after this", "what next", "whats next", "next step",
+        "then what", "following step", "continue", "subsequent", "proceed"
+    ]
+    return any(k in q for k in keys)
+
+def _resolve_next_steps(user_query: str, numbered_text: str, max_next: int = 6, min_score: float = 0.35):
+    """
+    If 'what's next' intent is detected and we can reliably match the user's
+    referenced line to a SOP step, return ONLY the subsequent steps.
+    Else return None to preserve current behavior.
+    """
+    if not _detect_next_intent(user_query):
+        return None
+
+    steps = _split_sop_into_steps(numbered_text)
+    if not steps:
+        return None
+
+    q = user_query or ""
+    best_idx, best_score = -1, -1.0
+    for idx, step in enumerate(steps):
+        # Exact substring match gets max score; else use soft match
+        score = 1.0 if _norm_text(step) in _norm_text(q) else _soft_match_score(q, step)
+        if score > best_score:
+            best_score, best_idx = score, idx
+
+    if best_idx < 0 or best_score < min_score:
+        return None  # fallback to full SOP
+
+    start = best_idx + 1
+    if start >= len(steps):
+        return []  # user is at final step
+
+    end = min(start + max_next, len(steps))
+    return steps[start:end]
+
+def _format_steps_as_numbered(steps: list) -> str:
+    """
+    Render a small list of steps with circled numbers for visual continuity.
+    """
+    circled = {1:"\u2460",2:"\u2461",3:"\u2462",4:"\u2463",5:"\u2464",6:"\u2465",7:"\u2466",8:"\u2467",9:"\u2468",10:"\u2469",
+               11:"\u246a",12:"\u246b",13:"\u246c",14:"\u246d",15:"\u246e",16:"\u246f",17:"\u2470",18:"\u2471",19:"\u2472",20:"\u2473"}
+    out = []
+    for i, s in enumerate(steps, start=1):
+        out.append(f"{circled.get(i, str(i))} {s}")
+    return "\n".join(out)
+
 def _filter_error_lines_by_query(text: str, query: str, max_lines: int = 1) -> str:
     """
     Pick the most relevant 'Common Errors & Resolution' bullet(s) for the user's message.
@@ -906,8 +992,18 @@ async def chat_with_ai(input_data: ChatInput):
                     sec = (top_meta or {}).get("section")
                     if sec: full_steps = get_section_text(best_doc, sec)
                 if full_steps:
-                    context = _ensure_numbering(full_steps)
-
+                    
+                   numbered = _ensure_numbering(full_steps)  # keep your existing formatting
+        # NEW: return only subsequent steps when user asks "what's next"
+                   next_only = _resolve_next_steps(input_data.user_message, numbered, max_next=6, min_score=0.35)
+                   if next_only is not None:
+                      if len(next_only) == 0:
+                         context = "You are at the final step of this SOP. No further steps."
+                      else:
+                          context = _format_steps_as_numbered(next_only)
+                   else:
+                       context=numbered
+                    
             elif detected_intent == "errors":
                 full_errors = get_best_errors_section_text(best_doc)
                 #assist_followup = None  # collect a helpful follow-up for generic cases