enhanced app.py file

app.py

@@ -1,8 +1,10 @@
 # space/app.py
 import os
 import json
+import logging
 import gradio as gr
 import pandas as pd
+from typing import Optional, Tuple
 
 from tools.sql_tool import SQLTool
 from tools.predict_tool import PredictTool
@@ -14,26 +16,48 @@ from tools.ts_forecast_tool import TimeseriesForecastTool
 from utils.tracing import Tracer
 from utils.config import AppConfig
 
-# Optional tiny CPU LLM for planning (can be disabled by not setting ORCHESTRATOR_MODEL)
+# Configure logging
+logging.basicConfig(
+    level=logging.INFO,
+    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
+)
+logger = logging.getLogger(__name__)
+
+# Constants
+MAX_RESPONSE_LENGTH = 10000
+MAX_FORECAST_HORIZON = 365
+DEFAULT_FORECAST_HORIZON = 96
+
+# Optional LLM for planning
 llm = None
 LLM_ID = os.getenv("ORCHESTRATOR_MODEL")
 if LLM_ID:
     try:
         from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
+        logger.info(f"Loading orchestrator model: {LLM_ID}")
         _tok = AutoTokenizer.from_pretrained(LLM_ID)
         _mdl = AutoModelForCausalLM.from_pretrained(LLM_ID)
         llm = pipeline("text-generation", model=_mdl, tokenizer=_tok, max_new_tokens=512)
-    except Exception:
+        logger.info("Orchestrator model loaded successfully")
+    except Exception as e:
+        logger.warning(f"Failed to load orchestrator model: {e}. Using fallback planner.")
         llm = None
 
-cfg = AppConfig.from_env()
-tracer = Tracer.from_env()
-
-sql_tool = SQLTool(cfg, tracer)
-predict_tool = PredictTool(cfg, tracer)
-explain_tool = ExplainTool(cfg, tracer)
-report_tool = ReportTool(cfg, tracer)
-ts_tool = TimeseriesForecastTool(cfg, tracer)  # Granite wrapper
+# Initialize configuration and tools
+try:
+    cfg = AppConfig.from_env()
+    tracer = Tracer.from_env()
+    
+    sql_tool = SQLTool(cfg, tracer)
+    predict_tool = PredictTool(cfg, tracer)
+    explain_tool = ExplainTool(cfg, tracer)
+    report_tool = ReportTool(cfg, tracer)
+    ts_tool = TimeseriesForecastTool(cfg, tracer)
+    
+    logger.info("All tools initialized successfully")
+except Exception as e:
+    logger.error(f"Failed to initialize application: {e}")
+    raise
 
 SYSTEM_PROMPT = (
     "You are an analytical assistant for tabular data. "
@@ -42,7 +66,40 @@ SYSTEM_PROMPT = (
     "Always disclose the steps taken."
 )
 
-def plan_actions(message: str):
+
+def validate_message(message: str) -> Tuple[bool, str]:
+    """Validate user input message."""
+    if not message or not message.strip():
+        return False, "Please enter a valid question."
+    
+    if len(message) > MAX_RESPONSE_LENGTH:
+        return False, f"Message too long. Please limit to {MAX_RESPONSE_LENGTH} characters."
+    
+    # Basic SQL injection pattern detection
+    suspicious_patterns = [
+        r';\s*drop\s+table',
+        r';\s*delete\s+from',
+        r';\s*truncate',
+        r'union\s+select.*from',
+        r'exec\s*\(',
+        r'execute\s*\('
+    ]
+    
+    import re
+    msg_lower = message.lower()
+    for pattern in suspicious_patterns:
+        if re.search(pattern, msg_lower):
+            logger.warning(f"Suspicious SQL pattern detected: {pattern}")
+            return False, "Query contains potentially unsafe patterns. Please rephrase."
+    
+    return True, ""
+
+
+def plan_actions(message: str) -> dict:
+    """
+    Determine which tools to execute based on the user message.
+    Uses LLM if available, otherwise falls back to heuristics.
+    """
     if llm is not None:
         prompt = (
             f"{SYSTEM_PROMPT}\nUser: {message}\n"
@@ -52,104 +109,335 @@ def plan_actions(message: str):
             out = llm(prompt)[0]["generated_text"]
             last = out.split("\n")[-1].strip()
             obj = json.loads(last) if last.startswith("{") else json.loads(out[out.rfind("{"):])
+            
             if isinstance(obj, dict) and "steps" in obj:
-                return obj
-        except Exception:
-            pass
-    # Fallback heuristic:
+                # Validate steps
+                valid_steps = {'sql', 'predict', 'explain', 'report', 'forecast'}
+                obj["steps"] = [s for s in obj["steps"] if s in valid_steps]
+                if obj["steps"]:
+                    logger.info(f"LLM plan: {obj['steps']}")
+                    return obj
+        except json.JSONDecodeError as e:
+            logger.warning(f"Failed to parse LLM output as JSON: {e}")
+        except Exception as e:
+            logger.warning(f"LLM planning failed: {e}")
+    
+    # Fallback heuristic planning
     m = message.lower()
     steps = []
-    if any(k in m for k in ["show", "average", "count", "trend", "top", "sql", "query", "kpi"]): steps.append("sql")
-    if any(k in m for k in ["predict", "score", "risk", "propensity", "probability"]): steps.append("predict")
-    if any(k in m for k in ["why", "explain", "shap", "feature", "attribution"]): steps.append("explain")
-    if any(k in m for k in ["report", "download", "pdf", "summary"]): steps.append("report")
-    if any(k in m for k in ["forecast", "next", "horizon", "granite"]): steps.append("forecast")
-    if not steps: steps = ["sql"]
-    return {"steps": steps, "rationale": "Rule-based plan."}
-
-def run_agent(message: str, hitl_decision: str = "Approve", reviewer_note: str = ""):
-    tracer.trace_event("user_message", {"message": message})
-    plan = plan_actions(message)
-    tracer.trace_event("plan", plan)
-
-    sql_df = None
-    predict_df = None
-    explain_imgs = {}
-    artifacts = {}
-    ts_forecast_df = None
-
-    if "sql" in plan["steps"]:
-        sql_df = sql_tool.run(message)
-        artifacts["sql_rows"] = int(len(sql_df)) if isinstance(sql_df, pd.DataFrame) else 0
-
-    if "predict" in plan["steps"]:
-        predict_df = predict_tool.run(sql_df)
-
-    ts_df = None
-    if sql_df is not None:
-        try:
-            ts_df = build_timeseries(sql_df)
-        except Exception:
-            ts_df = None
-
-    if "forecast" in plan["steps"] and ts_df is not None:
-        # Expect 'portfolio_value' after preprocessing
-        # Use the combined series — e.g., sum over instruments by timestamp
-        agg = ts_df.groupby("timestamp", as_index=True)["portfolio_value"].sum().sort_index()
+    
+    # SQL keywords
+    if any(k in m for k in ["show", "average", "count", "trend", "top", "sql", "query", 
+                             "kpi", "data", "retrieve", "fetch", "list", "view"]):
+        steps.append("sql")
+    
+    # Prediction keywords
+    if any(k in m for k in ["predict", "score", "risk", "propensity", "probability", 
+                             "classification", "regression"]):
+        steps.append("predict")
+        if "sql" not in steps:
+            steps.insert(0, "sql")  # Need data first
+    
+    # Explanation keywords
+    if any(k in m for k in ["why", "explain", "shap", "feature", "attribution", 
+                             "importance", "interpret"]):
+        steps.append("explain")
+        if "predict" not in steps:
+            steps.insert(0, "predict")
+        if "sql" not in steps:
+            steps.insert(0, "sql")
+    
+    # Report keywords
+    if any(k in m for k in ["report", "download", "pdf", "summary", "document", "export"]):
+        steps.append("report")
+    
+    # Forecast keywords
+    if any(k in m for k in ["forecast", "next", "horizon", "granite", "predict future", 
+                             "time series", "timeseries"]):
+        steps.append("forecast")
+        if "sql" not in steps:
+            steps.insert(0, "sql")
+    
+    # Default to SQL if no steps identified
+    if not steps:
+        steps = ["sql"]
+    
+    rationale = f"Rule-based plan based on keywords: {', '.join(steps)}"
+    logger.info(f"Heuristic plan: {steps}")
+    return {"steps": steps, "rationale": rationale}
+
+
+def run_agent(
+    message: str, 
+    hitl_decision: str = "Approve", 
+    reviewer_note: str = ""
+) -> Tuple[str, pd.DataFrame]:
+    """
+    Main agent execution function.
+    
+    Args:
+        message: User query
+        hitl_decision: Human-in-the-loop decision
+        reviewer_note: Optional review notes
+    
+    Returns:
+        Tuple of (response_text, preview_dataframe)
+    """
+    try:
+        # Validate input
+        is_valid, error_msg = validate_message(message)
+        if not is_valid:
+            logger.warning(f"Invalid message: {error_msg}")
+            return f"❌ **Error:** {error_msg}", pd.DataFrame()
+        
+        tracer.trace_event("user_message", {"message": message[:500]})  # Limit traced message length
+        
+        # Plan actions
         try:
-            ts_forecast_df = ts_tool.zeroshot_forecast(agg, horizon=96)
+            plan = plan_actions(message)
+            tracer.trace_event("plan", plan)
         except Exception as e:
-            # Surface a readable error in response
-            ts_forecast_df = pd.DataFrame({"error": [str(e)]})
-
-    if "explain" in plan["steps"]:
-        explain_imgs = explain_tool.run(predict_df or sql_df)
-
-    report_link = None
-    if "report" in plan["steps"]:
-        # Add forecast preview if available
-        forecast_preview = ts_forecast_df.head(50) if isinstance(ts_forecast_df, pd.DataFrame) else None
-        report_link = report_tool.render_and_save(
-            user_query=message,
-            sql_preview=sql_df.head(50) if isinstance(sql_df, pd.DataFrame) else None,
-            predict_preview=predict_df.head(50) if isinstance(predict_df, pd.DataFrame) else forecast_preview,
-            explain_images=explain_imgs,
-            plan=plan,
-        )
+            logger.error(f"Planning failed: {e}")
+            return f"❌ **Planning Error:** Unable to create execution plan. {str(e)}", pd.DataFrame()
+        
+        # Initialize result containers
+        sql_df = None
+        predict_df = None
+        explain_imgs = {}
+        artifacts = {}
+        ts_forecast_df = None
+        errors = []
+        
+        # Execute SQL step
+        if "sql" in plan["steps"]:
+            try:
+                sql_df = sql_tool.run(message)
+                if isinstance(sql_df, pd.DataFrame):
+                    artifacts["sql_rows"] = len(sql_df)
+                    logger.info(f"SQL returned {len(sql_df)} rows")
+                else:
+                    errors.append("SQL query returned no data")
+            except Exception as e:
+                error_msg = f"SQL execution failed: {str(e)}"
+                logger.error(error_msg)
+                errors.append(error_msg)
+        
+        # Execute prediction step
+        if "predict" in plan["steps"]:
+            try:
+                if sql_df is not None and not sql_df.empty:
+                    predict_df = predict_tool.run(sql_df)
+                    if isinstance(predict_df, pd.DataFrame):
+                        artifacts["predict_rows"] = len(predict_df)
+                        logger.info(f"Predictions generated for {len(predict_df)} rows")
+                else:
+                    errors.append("Prediction skipped: no data available")
+            except Exception as e:
+                error_msg = f"Prediction failed: {str(e)}"
+                logger.error(error_msg)
+                errors.append(error_msg)
+        
+        # Build time series if possible
+        ts_df = None
+        if sql_df is not None and not sql_df.empty:
+            try:
+                ts_df = build_timeseries(sql_df)
+                logger.info(f"Time series built with {len(ts_df)} records")
+            except Exception as e:
+                logger.info(f"Time series preprocessing skipped: {e}")
+                # Not always an error - data might not be suitable for TS
+        
+        # Execute forecast step
+        if "forecast" in plan["steps"]:
+            if ts_df is not None and not ts_df.empty:
+                try:
+                    # Aggregate portfolio value by timestamp
+                    agg = ts_df.groupby("timestamp", as_index=True)["portfolio_value"].sum().sort_index()
+                    
+                    if len(agg) < 2:
+                        errors.append("Insufficient time series data for forecasting (need at least 2 points)")
+                    else:
+                        # Validate horizon
+                        horizon = min(DEFAULT_FORECAST_HORIZON, MAX_FORECAST_HORIZON)
+                        ts_forecast_df = ts_tool.zeroshot_forecast(agg, horizon=horizon)
+                        
+                        if isinstance(ts_forecast_df, pd.DataFrame):
+                            if "error" in ts_forecast_df.columns:
+                                errors.append(f"Forecast error: {ts_forecast_df['error'].iloc[0]}")
+                                ts_forecast_df = None
+                            else:
+                                artifacts["forecast_horizon"] = len(ts_forecast_df)
+                                logger.info(f"Forecast generated for {len(ts_forecast_df)} periods")
+                except Exception as e:
+                    error_msg = f"Forecasting failed: {str(e)}"
+                    logger.error(error_msg)
+                    errors.append(error_msg)
+            else:
+                errors.append("Forecast skipped: no suitable time series data")
+        
+        # Execute explanation step
+        if "explain" in plan["steps"]:
+            try:
+                explain_data = predict_df if predict_df is not None else sql_df
+                if explain_data is not None and not explain_data.empty:
+                    explain_imgs = explain_tool.run(explain_data)
+                    artifacts["explain_charts"] = len(explain_imgs)
+                    logger.info(f"Generated {len(explain_imgs)} explanation charts")
+                else:
+                    errors.append("Explanation skipped: no data available")
+            except Exception as e:
+                error_msg = f"Explanation failed: {str(e)}"
+                logger.error(error_msg)
+                errors.append(error_msg)
+        
+        # Execute report generation
+        report_link = None
+        if "report" in plan["steps"]:
+            try:
+                forecast_preview = ts_forecast_df.head(50) if isinstance(ts_forecast_df, pd.DataFrame) else None
+                report_link = report_tool.render_and_save(
+                    user_query=message,
+                    sql_preview=sql_df.head(50) if isinstance(sql_df, pd.DataFrame) else None,
+                    predict_preview=predict_df.head(50) if isinstance(predict_df, pd.DataFrame) else forecast_preview,
+                    explain_images=explain_imgs,
+                    plan=plan,
+                )
+                logger.info(f"Report generated: {report_link}")
+            except Exception as e:
+                error_msg = f"Report generation failed: {str(e)}"
+                logger.error(error_msg)
+                errors.append(error_msg)
+        
+        # Log human-in-the-loop decision
+        tracer.trace_event("hitl", {
+            "message": message[:500],
+            "decision": hitl_decision,
+            "reviewer_note": reviewer_note[:500] if reviewer_note else "",
+            "artifacts": artifacts,
+            "plan": plan,
+            "errors": errors,
+        })
+        
+        # Compose response
+        response = f"**Plan:** {', '.join(plan['steps'])}\n\n**Rationale:** {plan['rationale']}\n\n"
+        
+        # Add artifacts info
+        if artifacts:
+            response += "**Results:**\n"
+            if "sql_rows" in artifacts:
+                response += f"- SQL query returned {artifacts['sql_rows']} rows\n"
+            if "predict_rows" in artifacts:
+                response += f"- Generated predictions for {artifacts['predict_rows']} rows\n"
+            if "forecast_horizon" in artifacts:
+                response += f"- Forecast generated for {artifacts['forecast_horizon']} periods\n"
+            if "explain_charts" in artifacts:
+                response += f"- Created {artifacts['explain_charts']} explanation charts\n"
+            response += "\n"
+        
+        # Add report link
+        if report_link:
+            response += f"📄 **Report:** {report_link}\n\n"
+        
+        # Add trace URL
+        if tracer.trace_url:
+            response += f"🔍 **Trace:** {tracer.trace_url}\n\n"
+        
+        # Add errors if any
+        if errors:
+            response += "**⚠️ Warnings/Errors:**\n"
+            for err in errors:
+                response += f"- {err}\n"
+        
+        # Determine preview dataframe
+        if isinstance(ts_forecast_df, pd.DataFrame) and not ts_forecast_df.empty:
+            preview_df = ts_forecast_df.head(100)
+        elif isinstance(predict_df, pd.DataFrame) and not predict_df.empty:
+            preview_df = predict_df.head(100)
+        elif isinstance(sql_df, pd.DataFrame) and not sql_df.empty:
+            preview_df = sql_df.head(100)
+        else:
+            preview_df = pd.DataFrame({"message": ["No data to display"]})
+        
+        return response, preview_df
+        
+    except Exception as e:
+        error_msg = f"Unexpected error in agent execution: {str(e)}"
+        logger.exception(error_msg)
+        tracer.trace_event("error", {"message": error_msg})
+        return f"❌ **Critical Error:** {error_msg}", pd.DataFrame()
 
-    tracer.trace_event("hitl", {
-        "message": message,
-        "decision": hitl_decision,
-        "reviewer_note": reviewer_note,
-        "artifacts": artifacts,
-        "plan": plan,
-    })
-
-    # Compose response
-    response = f"**Plan:** {plan['steps']}\n**Rationale:** {plan['rationale']}\n"
-    if isinstance(sql_df, pd.DataFrame): response += f"\n**SQL rows:** {len(sql_df)}"
-    if isinstance(predict_df, pd.DataFrame): response += f"\n**Predictions rows:** {len(predict_df)}"
-    if isinstance(ts_forecast_df, pd.DataFrame) and "forecast" in ts_forecast_df.columns:
-        response += f"\n**Forecast horizon:** {len(ts_forecast_df)}"
-    if report_link: response += f"\n**Report:** {report_link}"
-    if tracer.trace_url: response += f"\n**Trace:** {tracer.trace_url}"
-
-    # Prefer to show forecast if present, else predictions, else raw query
-    preview_df = ts_forecast_df if isinstance(ts_forecast_df, pd.DataFrame) and not ts_forecast_df.empty else \
-                 (predict_df if isinstance(predict_df, pd.DataFrame) and not predict_df.empty else sql_df)
-    return response, (preview_df if isinstance(preview_df, pd.DataFrame) else pd.DataFrame())
-
-with gr.Blocks() as demo:
-    gr.Markdown("# Tabular Agentic XAI (Free-Tier)")
+
+# Gradio Interface
+with gr.Blocks(title="Tabular Agentic XAI") as demo:
+    gr.Markdown("""
+    # 🤖 Tabular Agentic XAI (Enterprise Edition)
+    
+    An intelligent assistant for analyzing tabular data with ML predictions, explanations, and time-series forecasting.
+    
+    **Capabilities:**
+    - 📊 SQL queries and data retrieval
+    - 🎯 ML predictions with confidence scores
+    - 🔍 SHAP-based model explanations
+    - 📈 Time-series forecasting with Granite TTM
+    - 📄 Automated report generation
+    """)
+    
     with gr.Row():
-        msg = gr.Textbox(label="Ask your question")
+        msg = gr.Textbox(
+            label="Ask your question",
+            placeholder="e.g., Show me the top 10 customers by revenue, predict churn risk, forecast next quarter...",
+            lines=3
+        )
+    
     with gr.Row():
-        hitl = gr.Radio(["Approve", "Needs Changes"], value="Approve", label="Human Review")
-        note = gr.Textbox(label="Reviewer note (optional)")
-    out_md = gr.Markdown()
-    out_df = gr.Dataframe(interactive=False)
-    ask = gr.Button("Run")
-    ask.click(run_agent, inputs=[msg, hitl, note], outputs=[out_md, out_df])
+        hitl = gr.Radio(
+            ["Approve", "Needs Changes"], 
+            value="Approve", 
+            label="Human Review",
+            info="Review the planned actions before execution"
+        )
+        note = gr.Textbox(
+            label="Reviewer note (optional)",
+            placeholder="Add any review comments...",
+            lines=2
+        )
+    
+    out_md = gr.Markdown(label="Response")
+    out_df = gr.Dataframe(
+        interactive=False,
+        label="Data Preview (max 100 rows)",
+        wrap=True
+    )
+    
+    with gr.Row():
+        ask = gr.Button("🚀 Run Analysis", variant="primary")
+        clear = gr.Button("🔄 Clear")
+    
+    ask.click(
+        run_agent, 
+        inputs=[msg, hitl, note], 
+        outputs=[out_md, out_df]
+    )
+    
+    clear.click(
+        lambda: ("", "Approve", "", "", pd.DataFrame()),
+        outputs=[msg, hitl, note, out_md, out_df]
+    )
+    
+    gr.Markdown("""
+    ---
+    **Tips:**
+    - Be specific in your queries for better results
+    - Use natural language - the system will interpret your intent
+    - Review the execution plan before approving
+    - Check the trace link for detailed execution logs
+    """)
+
 
 if __name__ == "__main__":
-    demo.launch()
+    logger.info("Starting Gradio application...")
+    demo.launch(
+        server_name="0.0.0.0",
+        server_port=7860,
+        show_error=True
+    )