Update app.py

app.py

@@ -34,55 +34,43 @@ from narrative_safetynet import build_narrative
 def _sanitize_text(s: str) -> str:
     if not isinstance(s, str):
         return s
-    # remove non-printing/control chars except newlines & tabs
     return re2.sub(r'[\p{C}--[\n\t]]+', '', s)
 
-
-def _dataset_catalog(results: Dict[str, Any]) -> Dict[str, List[str]]:
-    """Simple catalog of dataset columns for the planner prompt; dynamic & scenario-agnostic."""
-    cat: Dict[str, List[str]] = {}
-    for k, v in results.items():
-        if isinstance(v, pd.DataFrame):
-            cat[k] = v.columns.tolist()
-    return cat
-
-
-def is_healthcare_scenario(text: str, has_files: bool) -> bool:
+# --- NEW: The "Intake Analyst" AI ---
+def _create_enhanced_prompt(user_scenario: str) -> str:
     """
-    Dynamic detection: require uploaded files AND either structured scenario sections
-    or healthcare keywords (configured in settings).
+    Uses an LLM to pre-process the user's messy prompt into a structured brief
+    for the data analysis agent.
     """
-    t = (text or "").lower()
-    kws = HEALTHCARE_SETTINGS["healthcare_keywords"]
-    structured = any(s in t for s in ["background", "situation", "tasks", "deliverables"])
-    return has_files and (structured or any(k in t for k in kws))
+    # This prompt instructs the first LLM to act as a project manager.
+    prompt_for_planner = f"""
+You are an expert data analysis project manager. Your task is to read the user's unstructured scenario below and create a clear, structured brief for a data analysis AI.
 
+From the user's text, extract the following:
+1.  **Primary Objective:** A one-sentence summary of the user's main goal.
+2.  **Key Tasks:** A numbered list of the specific questions the user wants answered.
+3.  **Expert Guidelines & Assumptions:** A bulleted list of EVERY specific number, metric, calculation method, or assumption mentioned in the text. This is critical for high-quality analysis.
+4.  **Required Output Format:** A description of how the user wants the final answer to be structured.
 
-def _append_msg(history_messages: List[Dict[str, str]], role: str, content: str) -> List[Dict[str, str]]:
-    return (history_messages or []) + [{"role": role, "content": content}]
-
-
-def ping_cohere() -> str:
-    """Lightweight health check against Cohere (embeddings call)."""
-    try:
-        cli = _co_client()
-        if not cli:
-            return "Cohere client not initialized. Is COHERE_API_KEY set?"
-        vecs = cohere_embed(["hello", "world"])
-        if vecs and len(vecs) == 2:
-            return f"Cohere OK ✅ (model={COHERE_MODEL_PRIMARY}, timeout={COHERE_TIMEOUT_S}s)"
-        return "Cohere reachable, but embeddings returned no vectors."
-    except Exception as e:
-        return f"Cohere ping failed: {e}"
+Present this as a clean brief. Then, include the user's original text at the end.
 
+--- USER'S SCENARIO ---
+{user_scenario}
+"""
+    
+    # Use the existing cohere_chat function to get the structured brief
+    structured_brief = cohere_chat(prompt_for_planner)
+    
+    # If the LLM call fails, just use the original message
+    if not structured_brief:
+        return user_scenario
+        
+    return structured_brief
 
 # ---------------- Core handler ----------------
 def handle(user_msg: str, history_messages: List[Dict[str, str]], files: list) -> Tuple[List[Dict[str, str]], str]:
     """
-    One entrypoint for both healthcare scenarios and general conversation.
-    - NEW: If files are uploaded, a data-aware agent is used to perform analysis.
-    - Scenario mode (no files): planner -> deterministic executor -> LLM narrative (Cohere).
-    - General mode: direct to Cohere with a light system prompt.
+    Core logic handler with the new two-step AI process.
     """
     try:
         # Safety filter for user input
@@ -95,10 +83,9 @@ def handle(user_msg: str, history_messages: List[Dict[str, str]], files: list) -
 
         file_paths: List[str] = [getattr(f, "name", None) or f for f in (files or [])]
 
-        # --- NEW LOGIC: Activate data agent if files are uploaded ---
         if file_paths:
             try:
-                # --- FINAL UPGRADE: Load ALL uploaded CSVs into a list of DataFrames ---
+                # Load ALL uploaded CSVs into a list of DataFrames
                 dataframes = [pd.read_csv(p) for p in file_paths if p.endswith('.csv')]
                 if not dataframes:
                     return _append_msg(history_messages, "assistant", "Please upload at least one CSV file."), ""
@@ -106,6 +93,10 @@ def handle(user_msg: str, history_messages: List[Dict[str, str]], files: list) -
                 # Initialize the Cohere Chat LLM for the agent
                 llm = ChatCohere(model=COHERE_MODEL_PRIMARY, temperature=0)
 
+                # STEP 1: The "Intake Analyst" AI creates a structured brief.
+                enhanced_prompt = _create_enhanced_prompt(safe_in)
+
+                # This UNIVERSAL prefix contains only behavioral rules.
                 AGENT_PREFIX = """
 You are a data analysis agent. You have access to one or more pandas dataframes.
 You MUST respond in one of two formats.
@@ -116,17 +107,17 @@ Action: python_repl_ast
 Action Input: The Python code to run.
 
 FORMAT 2: To give the final answer. Your response must be a single block of text with ONLY these two sections:
-Thought: I can now answer the user's query.
-Final Answer: The complete answer.
+Thought: I can now answer the user's query based on the analysis.
+Final Answer: The complete answer, structured as the user requested.
 
 CRITICAL RULE: NEVER combine `Action` and `Final Answer` in the same response. Choose one format.
-Begin by analyzing the user's query and provide your first thought and action using FORMAT 1.
+Begin by analyzing the structured brief provided.
 """
 
-                # Create the pandas DataFrame agent, now giving it the LIST of dataframes
+                # STEP 2: The "Data Scientist" AI (Agent) executes the clean brief.
                 agent = create_pandas_dataframe_agent(
                     llm,
-                    dataframes, # <-- PASSING THE LIST OF DATAFRAMES
+                    dataframes,
                     agent_type=AgentType.ZERO_SHOT_REACT_DESCRIPTION,
                     verbose=True,
                     allow_dangerous_code=True,
@@ -134,24 +125,15 @@ Begin by analyzing the user's query and provide your first thought and action us
                     prefix=AGENT_PREFIX
                 )
 
-                # Run the agent with the user's scenario text.
-                reply = agent.run(safe_in)
+                reply = agent.run(enhanced_prompt)
                 reply = _sanitize_text(reply)
 
             except Exception as e:
                 tb = traceback.format_exc()
                 log_event("agent_error", None, {"err": str(e), "tb": tb})
                 reply = f"An error occurred while analyzing the data: {e}"
-
-        # --- ORIGINAL LOGIC: Fallback for scenarios without files or general chat ---
-        elif is_healthcare_scenario(safe_in, bool(file_paths)) and USE_SCENARIO_ENGINE:
-            # This block remains for scenarios without data files
-            registry = DataRegistry()
-            rag = RAGIndex()
-            # ... (rest of the original logic)
-
         else:
-            # General conversation mode
+            # Fallback to general conversation if no files are uploaded
             prompt = f"{GENERAL_CONVERSATION_PROMPT}\n\nUser: {safe_in}\nAssistant:"
             reply = cohere_chat(prompt) or open_fallback_chat(prompt) or "How can I help further?"
             reply = _sanitize_text(reply)
@@ -171,7 +153,7 @@ Begin by analyzing the user's query and provide your first thought and action us
 
 # ---------------- UI ----------------
 with gr.Blocks(analytics_enabled=False) as demo:
-    gr.Markdown("## Canadian Healthcare AI • Cohere API • Scenario-Agnostic • Deterministic Analytics")
+    gr.Markdown("## Universal AI Data Analyst")
 
     with gr.Row():
         chat = gr.Chatbot(label="Chat History", type="messages", height=520)
@@ -179,10 +161,10 @@ with gr.Blocks(analytics_enabled=False) as demo:
             label="Upload Data Files (CSV recommended)",
             file_count="multiple",
             type="filepath",
-            file_types=HEALTHCARE_SETTINGS["supported_file_types"]
+            file_types=[".csv"]
         )
 
-    msg = gr.Textbox(label="Prompt", placeholder="Paste any scenario (Background / Situation / Tasks / Deliverables) or just chat.")
+    msg = gr.Textbox(label="Prompt", placeholder="Paste your scenario, tasks, and any specific instructions here.")
     with gr.Row():
         send = gr.Button("Send")
         clear = gr.Button("Clear")