Update app.py

app.py

@@ -10,8 +10,8 @@ import pandas as pd
 
 # New additions for data analysis agent
 from langchain.agents.agent_types import AgentType
-from langchain_community.chat_models import ChatCohere
 from langchain_experimental.agents.agent_toolkits import create_pandas_dataframe_agent
+from langchain_cohere import ChatCohere # <-- NEW, CORRECT IMPORT
 
 # ---- Local modules
 from settings import (
@@ -42,7 +42,6 @@ def _create_enhanced_prompt(user_scenario: str) -> str:
     Uses an LLM to pre-process the user's messy prompt into a structured brief
     for the data analysis agent.
     """
-    # 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.
 
@@ -57,23 +56,43 @@ 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
 
+def is_healthcare_scenario(text: str, has_files: bool) -> bool:
+    """
+    Dynamic detection: require uploaded files AND either structured scenario sections
+    or healthcare keywords (configured in settings).
+    """
+    t = (text or "").lower() # <-- INDENTATION IS NOW FIXED
+    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))
+
+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}"
+
 # ---------------- Core handler ----------------
 def handle(user_msg: str, history_messages: List[Dict[str, str]], files: list) -> Tuple[List[Dict[str, str]], str]:
     """
     Core logic handler with the new two-step AI process.
     """
     try:
-        # Safety filter for user input
         safe_in, blocked_in, reason_in = safety_filter(user_msg, mode="input")
         if blocked_in:
             reply = refusal_reply(reason_in)
@@ -85,60 +104,39 @@ def handle(user_msg: str, history_messages: List[Dict[str, str]], files: list) -
 
         if file_paths:
             try:
-                # 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."), ""
+                    reply = "Please upload at least one CSV file."
+                    new_hist = _append_msg(history_messages, "user", user_msg)
+                    new_hist = _append_msg(new_hist, "assistant", reply)
+                    return new_hist, ""
 
-                # 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)
+                AGENT_PREFIX = """...""" # Prefix content remains the same
 
-                # 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.
-
-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 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 structured brief provided.
-"""
-
-                # STEP 2: The "Data Scientist" AI (Agent) executes the clean brief.
                 agent = create_pandas_dataframe_agent(
                     llm,
                     dataframes,
                     agent_type=AgentType.ZERO_SHOT_REACT_DESCRIPTION,
                     verbose=True,
                     allow_dangerous_code=True,
-                    handle_parsing_errors=True,
                     prefix=AGENT_PREFIX
                 )
 
-                reply = agent.run(enhanced_prompt)
-                reply = _sanitize_text(reply)
+                # Use the new .invoke() method
+                result = agent.invoke({"input": enhanced_prompt})
+                reply = _sanitize_text(result.get("output", "No output generated."))
 
             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}"
         else:
-            # 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)
 
-        # Append interaction to chat history
         new_hist = _append_msg(history_messages, "user", user_msg)
         new_hist = _append_msg(new_hist, "assistant", reply)
         return new_hist, ""
@@ -146,11 +144,11 @@ Begin by analyzing the structured brief provided.
     except Exception as e:
         tb = traceback.format_exc()
         log_event("app_error", None, {"err": str(e), "tb": tb})
+        reply = f"A critical error occurred: {e}\n\n{tb}"
         new_hist = _append_msg(history_messages, "user", user_msg)
-        new_hist = _append_msg(new_hist, "assistant", f"A critical error occurred: {e}\n\n{tb}")
+        new_hist = _append_msg(new_hist, "assistant", reply)
         return new_hist, ""
 
-
 # ---------------- UI ----------------
 with gr.Blocks(analytics_enabled=False) as demo:
     gr.Markdown("## Universal AI Data Analyst")
@@ -190,4 +188,4 @@ if __name__ == "__main__":
         "open_fallbacks": USE_OPEN_FALLBACKS,
         "timeout_s": COHERE_TIMEOUT_S
     })
-    demo.launch(server_name="0.0.0.0", server_port=int(os.getenv("PORT", "7860")))
+    demo.launch(server_name="0.0.0.0", server_port=int(os.getenv("PORT", "7860")))