Update app_web.py

app_web.py

@@ -173,17 +173,34 @@ def render_chat():
         cat_cfg = cfg.get("query_categorization", {})
         max_history = cat_cfg.get("max_history", 6)
 
+        # ── Check if this is a clarification response ───────────────────
+        unresolved = st.session_state.unresolved_category
+        if unresolved is not None:
+            # This prompt is the user's clarification answer
+            # Include the clarification question for context
+            unresolved_question = st.session_state.get("unresolved_category_question", "")
+            if unresolved_question:
+                combined_cat = f"{unresolved} (Clarification: Q: {unresolved_question} A: {prompt})"
+            else:
+                combined_cat = f"{unresolved} — {prompt}"
+            st.session_state.unresolved_category_question = None
+            original_query_cat = unresolved
+            st.session_state.unresolved_category = None
+        else:
+            combined_cat = prompt
+            original_query_cat = prompt
+            st.session_state.resolution_rounds = 0
+
         prior_messages = st.session_state.messages[:-1]
         history = [
             {"role": m["role"], "content": m["content"]}
             for m in prior_messages[-max_history:]
         ]
         try:
-            qcat_result = categorize_query(prompt, cfg, history)
-        except Exception as e:
-            print(e)
+            qcat_result = categorize_query(combined_cat, cfg, history)
+        except Exception:
             qcat_result = {"category": "general", "action": "continue"}
-        print(qcat_result.get("category"))
+
         if qcat_result.get("category") == "contact_and_info" and qcat_result.get("action") == "contact":
             with st.chat_message("assistant"):
                 st.markdown("For additional assistance, please contact the WEAI Helpdesk at IFPRI-WEAI@cgiar.org.\
@@ -213,8 +230,11 @@ def render_chat():
             return
 
         if qcat_result.get("category") == "stata":
-            if "awaiting_selection" not in st.session_state:
-                st.session_state.awaiting_selection = True
+            # Ask clarification — store original query and question for context
+            st.session_state.unresolved_category = original_query_cat
+            st.session_state.messages.append(
+                {"role": "assistant", "content": clarification_msg}
+            )
 
             file_path = ""
             if qcat_result.get("action") == "do":
@@ -241,35 +261,29 @@ def render_chat():
 
             if qcat_result.get("action") == "error":
                 options = {"explain the error":"explain", 
-                           "suggest fix(es) for the error":"suggestfix"}                
+                           "suggest fix(es) for the error":"suggestfix"}        
+                        
+            st.session_state.unresolved_category_question = f"Please let me know what you need help with: {'or '.join(options.keys())}. List all options that apply."
+            st.session_state.resolution_rounds += 1
+            resolution_msg = f"**Before I search, could you clarify?** Please let me know what you need help with: {'or '.join(options.keys())}. List all options that apply."
+            st.session_state.messages.append({"role": "user", "content": resolution_msg})
             
             with st.chat_message("assistant"):
-                with st.status("Waiting for input...", expanded=True) as status:
-                    selection = st.multiselect("Please let me know what you need help with:", 
-                                            options.keys())
-                    nxt_btn = st.button("Continue")
-                    
-                    if nxt_btn:
-                        status.update(label=f"Seeing how I can {'and '.join(selection)}")
-                        if selection:
-                            choice_text = f"Help me {'and '.join(selection)}"
-                            st.session_state.messages.append({"role": "user", "content": choice_text})
-                            selected_options = [options.get(s) for s in selection]
-                            opts = {
-                                    "rewrite":    True if "rewrite" in selected_options else False,
-                                    "explain":    True if "explain" in selected_options else False,
-                                    "suggestfix": True if "suggestfix" in selected_options else False,
-                                    # "capture":    True if "capture" in selected_options else False,
-                                    # "verbose":    True if "verbose" in selected_options else False,
-                                    "lines":      None#lines if "lines" in selected_options else None
-                                }
-                            result = explain_(qcat_result.get("action"), prompt, cfg, file_path, opts)
-                            st.session_state.messages.append({"role": "assistant", "content": result})
-                            st.markdown(result)
-                            return
-                        else:
-                            st.warning("Please select at least one option you need help with.")
-
+                selection = qcat_result.get("sub-action", ["explain"])
+                selected_options = [options.get(s) for s in selection]
+                opts = {
+                        "rewrite":    True if "rewrite" in selected_options else False,
+                        "explain":    True if "explain" in selected_options else False,
+                        "suggestfix": True if "suggestfix" in selected_options else False,
+                        # "capture":    True if "capture" in selected_options else False,
+                        # "verbose":    True if "verbose" in selected_options else False,
+                        "lines":      None#lines if "lines" in selected_options else None
+                    }
+                result = explain_(qcat_result.get("action"), prompt, cfg, file_path, opts)
+                st.session_state.messages.append({"role": "assistant", "content": result})
+                st.markdown(result)
+                return
+            
         # ── Check if this is a clarification response ───────────────────
         pending = st.session_state.pending_clarification
         if pending is not None: