Update app.py

app.py

@@ -9,87 +9,167 @@ import gradio as gr
 import pandas as pd
 from datetime import datetime
 
-# --- (All your previous backend imports remain the same) ---
+# --- BACKEND IMPORTS ---
 from langchain.agents.agent_types import AgentType
 from langchain_cohere import ChatCohere
 from langchain_experimental.agents.agent_toolkits import create_pandas_dataframe_agent
+
+# --- LOCAL MODULE IMPORTS ---
+# (Assuming these files exist in your project)
 from settings import (
     HEALTHCARE_SETTINGS, GENERAL_CONVERSATION_PROMPT, USE_SCENARIO_ENGINE, DEBUG_PLAN,
     COHERE_MODEL_PRIMARY, COHERE_TIMEOUT_S, USE_OPEN_FALLBACKS
 )
 from audit_log import log_event
 from privacy import safety_filter, refusal_reply
-# ... (and so on for the rest of your backend functions)
-
-# --- (The entire backend logic from the previous version should be pasted here) ---
-# This includes:
-# _sanitize_text, _create_enhanced_prompt, is_healthcare_scenario,
-# _append_msg, ping_cohere, and the main handle() function.
-# For brevity, I am omitting them here, but they are ESSENTIAL.
-# Please copy them from the previous version you have.
-# The ONLY CHANGE is that we will now call handle() from a new UI wrapper function.
-
-# (For demonstration, I will paste a minimal version of the backend functions here)
-def _sanitize_text(s: str) -> str: return s
-def _create_enhanced_prompt(s: str) -> str: return s
-def _append_msg(h, r, c): return (h or []) + [{"role": r, "content": c}]
-
-# This is your perfected backend engine. It does not need to be changed.
-def handle(user_msg: str, history_messages: List[Dict[str, str]], files: list) -> Tuple[List[Dict[str, str]], str]:
-    # --- (Paste your entire, working handle() function here) ---
-    # For now, I'll use a placeholder that returns a success message
-    # In your real app, this will be your full LangChain agent logic.
-    file_names = [os.path.basename(f) for f in files]
-    response_text = f"### Analysis Complete\n**Prompt:** {user_msg}\n**Files Used:** {', '.join(file_names)}\n\nThis is where the structured output from the AI agent would appear."
-    new_hist = _append_msg(history_messages, "user", user_msg)
-    new_hist = _append_msg(new_hist, "assistant", response_text)
-    return new_hist, ""
-
-
-# ---------------- THE NEW UI ----------------
+from llm_router import cohere_chat, _co_client, cohere_embed
+
+# --- BACKEND UTILITY FUNCTIONS ---
+
+def _sanitize_text(s: str) -> str:
+    if not isinstance(s, str):
+        return s
+    return re2.sub(r'[\p{C}--[\n\t]]+', '', s)
+
+def _create_enhanced_prompt(user_scenario: str) -> str:
+    """Uses an LLM to pre-process the user's messy prompt into a structured brief."""
+    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 ALL the specific questions the user wants answered.
+3.  **Expert Guidelines & Assumptions:** A bulleted list of any specific numbers, metrics, or calculation methods mentioned.
+4.  **Required Output Format:** A description of how the user wants the final answer structured.
+
+CRITICAL INSTRUCTION: Tell the data analyst that it MUST answer ALL of the key tasks before providing its final answer.
+
+--- USER'S SCENARIO ---
+{user_scenario}
+"""
+    structured_brief = cohere_chat(prompt_for_planner)
+    return structured_brief if structured_brief else user_scenario
+
+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."""
+    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}"
+
+# --- THE CORE ANALYSIS ENGINE ---
+
+def handle(user_msg: str, files: list) -> str:
+    """
+    This is the powerful backend engine. It takes the user's query and files
+    and returns only the final AI-generated text response.
+    """
+    try:
+        safe_in, blocked_in, reason_in = safety_filter(user_msg, mode="input")
+        if blocked_in:
+            return refusal_reply(reason_in)
+
+        file_paths: List[str] = [getattr(f, "name", None) or f for f in (files or [])]
+
+        if file_paths:
+            dataframes = [pd.read_csv(p) for p in file_paths if p.endswith('.csv')]
+            if not dataframes:
+                return "Please upload at least one CSV file."
+
+            llm = ChatCohere(model=COHERE_MODEL_PRIMARY, temperature=0)
+            enhanced_prompt = _create_enhanced_prompt(safe_in)
+            
+            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.
+
+FORMAT 1: To perform a task. Your response must be a single block of text with ONLY these three sections:
+Thought: Your step-by-step reasoning.
+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 have now answered all the user's questions and can provide the final report.
+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.
+"""
+
+            agent = create_pandas_dataframe_agent(
+                llm,
+                dataframes,
+                agent_type=AgentType.ZERO_SHOT_REACT_DESCRIPTION,
+                verbose=True,
+                allow_dangerous_code=True,
+                prefix=AGENT_PREFIX,
+                max_iterations=50
+            )
+
+            result = agent.invoke({"input": enhanced_prompt})
+            reply = _sanitize_text(result.get("output", "No output generated."))
+            return reply
+        else:
+            # General conversation mode if no files are uploaded
+            prompt = f"{GENERAL_CONVERSATION_PROMPT}\n\nUser: {safe_in}\nAssistant:"
+            reply = cohere_chat(prompt) or "How can I help further?"
+            return _sanitize_text(reply)
+
+    except Exception as e:
+        tb = traceback.format_exc()
+        log_event("app_error", None, {"err": str(e), "tb": tb})
+        return f"A critical error occurred: {e}"
+
+# ---------------- THE NEW PROFESSIONAL UI ----------------
 with gr.Blocks(theme="soft", css="style.css") as demo:
     # State to store the history of all assessments in this session
     assessment_history = gr.State([])
 
-    gr.Markdown("# Universal AI Data Analyst", elem_classes="h1")
+    gr.Markdown("# Universal AI Data Analyst")
 
     with gr.Row(variant="panel"):
         # --- LEFT COLUMN: CONTROLS ---
         with gr.Column(scale=1):
-            gr.Markdown("## New Assessment", elem_classes="h2")
-
-            files = gr.Files(
+            gr.Markdown("## New Assessment")
+            files_input = gr.Files(
                 label="Upload Data Files (CSV recommended)",
                 file_count="multiple",
                 type="filepath",
                 file_types=[".csv"]
             )
-            msg = gr.Textbox(
+            prompt_input = gr.Textbox(
                 label="Prompt",
                 placeholder="Paste your scenario, tasks, and any specific instructions here.",
-                lines=10
+                lines=15
             )
             with gr.Row():
-                send_btn = gr.Button("▶️ Run Analysis", variant="primary")
-                clear_btn = gr.Button("🗑️ Clear", variant="secondary")
+                send_btn = gr.Button("▶️ Run Analysis", variant="primary", scale=2)
+                clear_btn = gr.Button("🗑️ Clear")
+            
+            ping_btn = gr.Button("Ping Cohere")
+            ping_out = gr.Markdown()
 
         # --- RIGHT COLUMN: RESULTS & HISTORY ---
         with gr.Column(scale=2):
             with gr.Tabs():
                 # --- TAB 1: CURRENT ASSESSMENT ---
                 with gr.TabItem("Current Assessment", id=0):
-                    chat_history = gr.Chatbot(
-                        label="Chat History",
+                    chat_history_output = gr.Chatbot(
+                        label="Analysis Output",
                         bubble_full_width=True,
-                        height=500
+                        height=600
                     )
-                    ping_btn = gr.Button("Ping Cohere")
-                    ping_out = gr.Markdown()
-
-
                 # --- TAB 2: ASSESSMENT HISTORY ---
                 with gr.TabItem("Assessment History", id=1):
-                    gr.Markdown("## Review Past Assessments", elem_classes="h2")
+                    gr.Markdown("## Review Past Assessments")
                     history_dropdown = gr.Dropdown(
                         label="Select an assessment to review",
                         choices=[]
@@ -99,70 +179,60 @@ with gr.Blocks(theme="soft", css="style.css") as demo:
                     )
 
     # --- UI LOGIC ---
-
-    # Function to run when "Run Analysis" is clicked
-    def run_analysis(prompt, files, chat, history_state):
+    def run_analysis_wrapper(prompt, files, chat_history_list, history_state_list):
         if not prompt or not files:
             gr.Warning("Please provide both a prompt and at least one data file.")
-            return chat, history_state, gr.update()
+            return chat_history_list, history_state_list, gr.update()
+
+        # 1. Append the user's message to the chat
+        chat_with_user_msg = _append_msg(chat_history_list, "user", prompt)
+        
+        # 2. Call the powerful backend engine to get the AI response
+        ai_response_text = handle(prompt, files)
 
-        # Call your powerful backend function
-        final_chat, _ = handle(prompt, chat, files)
+        # 3. Append the AI's response to the chat
+        final_chat = _append_msg(chat_with_user_msg, "assistant", ai_response_text)
         
-        # Save the completed assessment to our history state
+        # 4. Save the completed assessment to our history state
         timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
         file_names = [os.path.basename(f) for f in files]
         
         new_assessment = {
-            "id": timestamp,
-            "prompt": prompt,
-            "files": file_names,
-            "response": final_chat[-1]['content'] # Get the AI's final response
+            "id": timestamp, "prompt": prompt, "files": file_names,
+            "response": ai_response_text
         }
+        updated_history = history_state_list + [new_assessment]
         
-        updated_history = history_state + [new_assessment]
-        
-        # Create user-friendly labels for the dropdown
+        # 5. Create user-friendly labels for the history dropdown
         history_labels = [f"{item['id']} - {item['prompt'][:40]}..." for item in updated_history]
         
         return final_chat, updated_history, gr.update(choices=history_labels)
 
-    # Function to run when a history item is selected from the dropdown
-    def view_history(selection, history_state):
-        if not selection or not history_state:
-            return ""
-            
-        # Find the selected assessment
-        # The selection string is "TIMESTAMP - PROMPT...", so we match by the timestamp
+    def view_history(selection, history_state_list):
+        if not selection or not history_state_list: return ""
         selected_id = selection.split(" - ")[0]
-        selected_assessment = next((item for item in history_state if item["id"] == selected_id), None)
+        selected_assessment = next((item for item in history_state_list if item["id"] == selected_id), None)
 
         if selected_assessment:
-            # Format the past assessment for beautiful display in Markdown
-            display_text = f"""
-            ### Assessment from: {selected_assessment['id']}
-            
-            **Files Used:**
-            - {'- '.join(selected_assessment['files'])}
-            
-            ---
-            
-            **Original Prompt:**
-            > {selected_assessment['prompt']}
-            
-            ---
-            
-            **AI Generated Response:**
-            {selected_assessment['response']}
-            """
-            return display_text
+            file_list_md = "\n- ".join(selected_assessment['files'])
+            return f"""
+### Assessment from: {selected_assessment['id']}
+**Files Used:**
+- {file_list_md}
+---
+**Original Prompt:**
+> {selected_assessment['prompt']}
+---
+**AI Generated Response:**
+{selected_assessment['response']}
+"""
         return "Could not find the selected assessment."
 
     # Wire up the components
     send_btn.click(
-        run_analysis,
-        inputs=[msg, files, chat_history, assessment_history],
-        outputs=[chat_history, assessment_history, history_dropdown]
+        run_analysis_wrapper,
+        inputs=[prompt_input, files_input, chat_history_output, assessment_history],
+        outputs=[chat_history_output, assessment_history, history_dropdown]
     )
     
     history_dropdown.change(
@@ -171,12 +241,11 @@ with gr.Blocks(theme="soft", css="style.css") as demo:
         outputs=[history_display]
     )
 
-    clear_btn.click(lambda: (None, None, None), outputs=[msg, files, chat_history])
+    clear_btn.click(lambda: (None, None, [], []), outputs=[prompt_input, files_input, chat_history_output, assessment_history])
     ping_btn.click(lambda: ping_cohere(), outputs=[ping_out])
 
-
 if __name__ == "__main__":
-    # --- (Your startup logic remains the same) ---
     if not os.getenv("COHERE_API_KEY"):
         print("🔴 COHERE_API_KEY environment variable not set. Application may not function correctly.")
+
     demo.launch(server_name="0.0.0.0", server_port=int(os.getenv("PORT", "7860")))