app.py

# app.py
from __future__ import annotations
import os
import traceback
import regex as re2
from typing import List, Dict, Any

import gradio as gr
import pandas as pd
from datetime import datetime

# --- BACKEND IMPORTS ---
from langchain_cohere import ChatCohere
# --- THE FIXED IMPORT IS HERE ---
from langchain_experimental.utilities.python import PythonREPL 

# --- LOCAL MODULE IMPORTS ---
from settings import (
    HEALTHCARE_SETTINGS, GENERAL_CONVERSATION_PROMPT,
    COHERE_MODEL_PRIMARY, COHERE_TIMEOUT_S, USE_OPEN_FALLBACKS
)
from audit_log import log_event
from privacy import safety_filter, refusal_reply
from llm_router import cohere_chat, _co_client, cohere_embed

# --- UTILITY FUNCTIONS ---

def load_markdown_text(filepath: str) -> str:
    """Safely loads text content from a markdown file."""
    try:
        with open(filepath, 'r', encoding='utf-8') as f:
            return f.read()
    except FileNotFoundError:
        return f"**Error:** The document `{os.path.basename(filepath)}` was not found."

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

def _create_python_script(user_scenario: str, schema_context: str) -> str:
    """Uses an LLM to act as an "AI Coder", writing a complete Python script."""
    prompt_for_coder = f"""
You are an expert Python data scientist. Your sole job is to write a single, complete, and executable Python script to answer the user's request.
You have access to a list of pandas dataframes loaded into a variable named `dfs`. The first dataframe is `dfs[0]`, the second is `dfs[1]`, and so on.

CRITICAL CONTEXT: Before writing any code, you MUST first understand the data you have been given. Here is the schema for each dataframe:
--- DATA SCHEMA ---
{schema_context}
--- END SCHEMA ---

CRITICAL RULE: You MUST use the exact column names provided in the DATA SCHEMA. Column names are case-sensitive. Pay close attention to capitalization (e.g., 'Zone' vs 'zone'). A KeyError will cause a failure.

Based on the user's scenario below, write a single Python script that performs the entire analysis.

RULES FOR YOUR SCRIPT:
1.  **Use the DataFrames:** Your script MUST use the `dfs` list to access the data.
2.  **Print Your Findings:** Use the `print()` function at each step of your analysis to output the results as a formatted report.
3.  **No Placeholders:** Do not use placeholder data.
4.  **Self-Contained:** The script must be entirely self-contained.

--- USER'S SCENARIO ---
{user_scenario}

--- PYTHON SCRIPT ---
Now, write the complete Python script to be executed. The script should start with `import pandas as pd` and contain all the logic.
```python
"""
    generated_text = cohere_chat(prompt_for_coder)
    match = re2.search(r"```python\n(.*?)```", generated_text, re2.DOTALL)
    if match:
        return match.group(1).strip()
    else:
        return "print('Error: The AI failed to generate a valid Python script.')"

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."
        vecs = cohere_embed(["hello", "world"])
        return f"Cohere OK ✅ (model={COHERE_MODEL_PRIMARY}, timeout={COHERE_TIMEOUT_S}s)" if vecs else "Cohere reachable."
    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 using the "Coder" pattern."""
    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 = []
            schema_parts = []
            for i, p in enumerate(file_paths):
                if p.endswith('.csv'):
                    try:
                        df = pd.read_csv(p)
                    except UnicodeDecodeError:
                        df = pd.read_csv(p, encoding='latin1')
                    dataframes.append(df)
                    schema_parts.append(f"DataFrame `dfs[{i}]` (from file `{os.path.basename(p)}`):\n{df.head().to_markdown()}\n")
            
            if not dataframes: return "Please upload at least one CSV file."

            schema_context = "\n".join(schema_parts)
            analysis_script = _create_python_script(safe_in, schema_context)

            # Initialize the Python Executor
            python_repl = PythonREPL()
            # Pass the dataframes into the execution environment
            local_vars = {"dfs": dataframes}
            
            try:
                # Execute the AI-generated script
                res = python_repl.run(command=analysis_script, locals=local_vars)
                return _sanitize_text(res)
            except Exception as e:
                # If execution fails, return the error and the script for debugging
                return f"An error occurred executing the script: {e}\n\nGenerated Script:\n```python\n{analysis_script}\n```"
        else:
            prompt = f"{GENERAL_CONVERSATION_PROMPT}\n\nUser: {safe_in}\nAssistant:"
            return _sanitize_text(cohere_chat(prompt) or "How can I help further?")

    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}"

# --- PRE-LOAD LEGAL DOCUMENTS ---
PRIVACY_POLICY_TEXT = load_markdown_text("privacy_policy.md")
TERMS_OF_SERVICE_TEXT = load_markdown_text("terms_of_service.md")

# ---------------- THE PROFESSIONAL UI WITH INTEGRATED LEGAL DOCS ----------------
with gr.Blocks(theme="soft", css="style.css") as demo:
    assessment_history = gr.State([])
    
    with gr.Group(visible=False) as privacy_modal:
        with gr.Blocks():
            gr.Markdown(PRIVACY_POLICY_TEXT)
            close_privacy_btn = gr.Button("Close")
            
    with gr.Group(visible=False) as terms_modal:
        with gr.Blocks():
            gr.Markdown(TERMS_OF_SERVICE_TEXT)
            close_terms_btn = gr.Button("Close")

    gr.Markdown("# Universal AI Data Analyst")
    with gr.Row(variant="panel"):
        with gr.Column(scale=1):
            gr.Markdown("## New Assessment")
            files_input = gr.Files(label="Upload Data Files (.csv)", file_count="multiple", type="filepath", file_types=[".csv"])
            prompt_input = gr.Textbox(label="Prompt", placeholder="Paste your scenario here.", lines=15)
            with gr.Row():
                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()
        with gr.Column(scale=2):
            with gr.Tabs():
                with gr.TabItem("Current Assessment", id=0):
                    chat_history_output = gr.Chatbot(label="Analysis Output", type="messages", height=600)
                with gr.TabItem("Assessment History", id=1):
                    gr.Markdown("## Review Past Assessments")
                    history_dropdown = gr.Dropdown(label="Select an assessment to review", choices=[])
                    history_display = gr.Markdown(label="Selected Assessment Details")
    with gr.Row(): gr.Markdown("---")
    with gr.Row():
        privacy_link = gr.Button("Privacy Policy", variant="link")
        terms_link = gr.Button("Terms of Service", variant="link")

    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.")
            yield chat_history_list, history_state_list, gr.update()
            return

        chat_with_user_msg = _append_msg(chat_history_list, "user", prompt)
        thinking_message = _append_msg(chat_with_user_msg, "assistant", "```\n🧠 Generating and executing analysis script... This may take a moment.\n```")
        yield thinking_message, history_state_list, gr.update()
        
        ai_response_text = handle(prompt, files)
        
        final_chat = _append_msg(chat_with_user_msg, "assistant", ai_response_text)
        timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
        file_names = [os.path.basename(f.name if hasattr(f, 'name') else f) for f in files]
        new_assessment = {"id": timestamp, "prompt": prompt, "files": file_names, "response": ai_response_text}
        updated_history = history_state_list + [new_assessment]
        history_labels = [f"{item['id']} - {item['prompt'][:40]}..." for item in updated_history]
        yield final_chat, updated_history, gr.update(choices=history_labels)

    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_list if item["id"] == selected_id), None)
        if selected_assessment:
            file_list_md = "\n- ".join(selected_assessment['files'])
            return f"""### Assessment from: {selected_assessment['id']}\n**Files Used:**\n- {file_list_md}\n---\n**Original Prompt:**\n> {selected_assessment['prompt']}\n---\n**AI Generated Response:**\n{selected_assessment['response']}"""
        return "Could not find the selected assessment."

    send_btn.click(
        run_analysis_wrapper,
        inputs=[prompt_input, files_input, chat_history_output, assessment_history],
        outputs=[chat_history_output, assessment_history, history_dropdown]
    )
    history_dropdown.change(
        view_history,
        inputs=[history_dropdown, assessment_history],
        outputs=[history_display]
    )
    clear_btn.click(lambda: (None, None, [], []), outputs=[prompt_input, files_input, chat_history_output, assessment_history])
    ping_btn.click(ping_cohere, outputs=[ping_out])
    privacy_link.click(lambda: gr.update(visible=True), outputs=[privacy_modal])
    close_privacy_btn.click(lambda: gr.update(visible=False), outputs=[privacy_modal])
    terms_link.click(lambda: gr.update(visible=True), outputs=[terms_modal])
    close_terms_btn.click(lambda: gr.update(visible=False), outputs=[terms_modal])

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