Update app.py

app.py

@@ -1,4 +1,5 @@
 # app.py
+# app.py
 from __future__ import annotations
 import os
 import traceback
@@ -8,6 +9,11 @@ from typing import List, Tuple, Dict, Any
 import gradio as gr
 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
+
 # ---- Local modules
 from settings import (
     HEALTHCARE_SETTINGS, GENERAL_CONVERSATION_PROMPT, USE_SCENARIO_ENGINE, DEBUG_PLAN,
@@ -75,11 +81,12 @@ def ping_cohere() -> str:
 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.
-    - Scenario mode: planner -> deterministic executor -> LLM narrative (Cohere) -> safety-net narrative if needed.
-    - General mode: direct to Cohere/open fallback with a light system prompt.
+    - 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.
     """
     try:
-        # Safety
+        # 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)
@@ -87,70 +94,74 @@ def handle(user_msg: str, history_messages: List[Dict[str, str]], files: list) -
             new_hist = _append_msg(new_hist, "assistant", reply)
             return new_hist, ""
 
-        # Normalize files into paths (Gradio can return temp file objects or paths)
         file_paths: List[str] = [getattr(f, "name", None) or f for f in (files or [])]
 
-        # Register CSVs for deterministic analysis
-        registry = DataRegistry()
-        for p in file_paths:
+        # --- NEW LOGIC: Activate data agent if files are uploaded ---
+        if file_paths:
+            try:
+                # For this example, we'll load the first CSV file.
+                # This can be extended to handle multiple DataFrames.
+                df = pd.read_csv(file_paths[0])
+
+                # Initialize the Cohere Chat LLM for the agent
+                llm = ChatCohere(model=COHERE_MODEL_PRIMARY, temperature=0)
+
+                # Create the pandas DataFrame agent, powered by Cohere
+                agent = create_pandas_dataframe_agent(
+                    llm,
+                    df,
+                    agent_type=AgentType.OPENAI_FUNCTIONS, # Recommended for reliability
+                    verbose=True # Set to False in production
+                )
+
+                # Run the agent with the user's scenario text. The agent will
+                # write and execute code to answer the query based on the dataframe.
+                reply = agent.run(safe_in)
+                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 now primarily handles scenarios where no data files are provided,
+            # relying on the original deterministic analysis logic.
+            registry = DataRegistry() # This part might be simplified if files always trigger the agent
+            rag = RAGIndex()
             try:
-                if p:
-                    registry.add_path(p)
+                ing = extract_text_from_files(file_paths) # For text extraction from markdown/txt
+                rag.add(ing.get("chunks", []))
             except Exception as e:
-                log_event("ingest_error", None, {"file": p, "err": str(e)})
-
-        # Lightweight RAG: ingest any text/markdown for grounding (embeddings via Cohere)
-        rag = RAGIndex()
-        try:
-            ing = extract_text_from_files(file_paths)
-            rag.add(ing.get("chunks", []))
-        except Exception as e:
-            log_event("rag_ingest_error", None, {"err": str(e)})
-
-        # Decide mode
-        if is_healthcare_scenario(safe_in, bool(file_paths)) and USE_SCENARIO_ENGINE:
-            # 1) Deterministic dataset exposure
+                log_event("rag_ingest_error", None, {"err": str(e)})
+
             analyzer = HealthcareAnalyzer(registry)
-            datasets = analyzer.comprehensive_analysis(safe_in)  # dict[str, DataFrame]
+            datasets = analyzer.comprehensive_analysis(safe_in)
             catalog = _dataset_catalog(datasets)
-
-            # 2) Plan (Cohere-first; auto safety-net if LLM parse fails)
             plan = parse_to_plan(safe_in, catalog)
-
-            # 3) Execute plan deterministically
             structured_md = ScenarioEngine.execute_plan(plan, datasets)
-
-            # 4) Narrative (Cohere-first), grounded with RAG hits
             rag_hits = [txt for txt, _ in rag.retrieve(safe_in, k=6)]
             narrative = generate_narrative(safe_in, structured_md, rag_hits)
 
-            # 5) Safety-net narrative if LLM narrative absent/failed
             if not narrative or "Unable to generate narrative" in narrative:
-                # Provide generic hints only (dynamic, not hard-coded to any schema)
                 narrative = build_narrative(
-                    scenario_text=safe_in,
-                    datasets=datasets,
-                    structured_tables=None,
+                    scenario_text=safe_in, datasets=datasets, structured_tables=None,
                     metric_hints=["surgery_median", "consult_median", "wait", "median", "p90", "90th"],
                     group_hints=["facility", "specialty", "zone", "hospital", "city", "region"],
                     min_sample=5
                 )
 
-            debug_note = ""
-            if DEBUG_PLAN and getattr(plan, "notes", None):
-                debug_note = f"\n\n> **Planner note:** {getattr(plan, 'notes', '')}"
-
-            reply = _sanitize_text(
-                f"{structured_md}\n\n# Narrative & Recommendations\n\n{narrative}{debug_note}"
-            )
+            debug_note = f"\n\n> **Planner note:** {getattr(plan, 'notes', '')}" if DEBUG_PLAN and getattr(plan, "notes", None) else ""
+            reply = _sanitize_text(f"{structured_md}\n\n# Narrative & Recommendations\n\n{narrative}{debug_note}")
 
         else:
-            # General conversation mode (Cohere-first; open-weights fallback)
+            # General conversation mode (no files, not a structured scenario)
             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 to chat history
+        # 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, ""
@@ -159,7 +170,7 @@ def handle(user_msg: str, history_messages: List[Dict[str, str]], files: list) -
         tb = traceback.format_exc()
         log_event("app_error", None, {"err": str(e), "tb": tb})
         new_hist = _append_msg(history_messages, "user", user_msg)
-        new_hist = _append_msg(new_hist, "assistant", f"Error: {e}\n\n{tb}")
+        new_hist = _append_msg(new_hist, "assistant", f"A critical error occurred: {e}\n\n{tb}")
         return new_hist, ""
 
 
@@ -168,15 +179,15 @@ with gr.Blocks(analytics_enabled=False) as demo:
     gr.Markdown("## Canadian Healthcare AI • Cohere API • Scenario-Agnostic • Deterministic Analytics")
 
     with gr.Row():
-        # Use messages format to avoid deprecation warnings and enable role-based history
-        chat = gr.Chatbot(type="messages", height=520)
+        chat = gr.Chatbot(label="Chat History", type="messages", height=520)
         files = gr.Files(
+            label="Upload Data Files (CSV recommended)",
             file_count="multiple",
             type="filepath",
             file_types=HEALTHCARE_SETTINGS["supported_file_types"]
         )
 
-    msg = gr.Textbox(placeholder="Paste any scenario (Background / Situation / Tasks / Deliverables) or just chat.")
+    msg = gr.Textbox(label="Prompt", placeholder="Paste any scenario (Background / Situation / Tasks / Deliverables) or just chat.")
     with gr.Row():
         send = gr.Button("Send")
         clear = gr.Button("Clear")
@@ -189,14 +200,18 @@ with gr.Blocks(analytics_enabled=False) as demo:
 
     send.click(_on_send, inputs=[msg, chat, files], outputs=[chat, msg])
     msg.submit(_on_send, inputs=[msg, chat, files], outputs=[chat, msg])
-    clear.click(lambda: ([], ""), outputs=[chat, msg])
+    clear.click(lambda: ([], "", None), outputs=[chat, msg, files])
     ping_btn.click(lambda: ping_cohere(), outputs=[ping_out])
 
 if __name__ == "__main__":
+    # Ensure you have your COHERE_API_KEY set as an environment variable
+    if not os.getenv("COHERE_API_KEY"):
+        print("🔴 COHERE_API_KEY environment variable not set. Application may not function correctly.")
+
     log_event("startup", None, {
         "cohere_key_present": bool(os.getenv("COHERE_API_KEY")),
         "cohere_model": COHERE_MODEL_PRIMARY,
         "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")))