app.py

# app.py
from __future__ import annotations
import os
import traceback
import regex as re2
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_cohere 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,
    COHERE_MODEL_PRIMARY, COHERE_TIMEOUT_S, USE_OPEN_FALLBACKS
)
from audit_log import log_event
from privacy import safety_filter, refusal_reply
from data_registry import DataRegistry
from upload_ingest import extract_text_from_files
from healthcare_analysis import HealthcareAnalyzer
from scenario_planner import parse_to_plan
from scenario_engine import ScenarioEngine
from rag import RAGIndex
from llm_router import generate_narrative, cohere_chat, open_fallback_chat, _co_client, cohere_embed
from narrative_safetynet import build_narrative


# ---------------- Utilities ----------------
def _sanitize_text(s: str) -> str:
    if not isinstance(s, str):
        return s
    return re2.sub(r'[\p{C}--[\n\t]]+', '', s)

# --- NEW: The "Intake Analyst" AI ---
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.
    """
    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)
    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()
    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:
        safe_in, blocked_in, reason_in = safety_filter(user_msg, mode="input")
        if blocked_in:
            reply = refusal_reply(reason_in)
            new_hist = _append_msg(history_messages, "user", user_msg)
            new_hist = _append_msg(new_hist, "assistant", reply)
            return new_hist, ""

        file_paths: List[str] = [getattr(f, "name", None) or f for f in (files or [])]

        if file_paths:
            try:
                dataframes = [pd.read_csv(p) for p in file_paths if p.endswith('.csv')]
                if not dataframes:
                    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, ""

                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  # <-- THE FINAL FIX IS HERE
                )

                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:
            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)

        new_hist = _append_msg(history_messages, "user", user_msg)
        new_hist = _append_msg(new_hist, "assistant", reply)
        return new_hist, ""

    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", reply)
        return new_hist, ""

# ---------------- UI ----------------
with gr.Blocks(analytics_enabled=False) as demo:
    gr.Markdown("## Universal AI Data Analyst")

    with gr.Row():
        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=[".csv"]
        )

    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")
        ping_btn = gr.Button("Ping Cohere")
    ping_out = gr.Markdown()

    def _on_send(m, h, f):
        h2, _ = handle(m, h, f or [])
        return h2, ""

    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: ([], "", None), outputs=[chat, msg, files])
    ping_btn.click(lambda: ping_cohere(), outputs=[ping_out])

if __name__ == "__main__":
    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")))