Update src/streamlit_app.py

src/streamlit_app.py

@@ -9,8 +9,8 @@ import plotly.graph_objects as go
 from datetime import datetime
 
 # ── Path setup so `src` is importable when running from src/ or project root ──
-_here = Path(__file__).resolve().parent  # src/
-_root = _here.parent  # project root
+_here = Path(__file__).resolve().parent          # src/
+_root = _here.parent                             # project root
 for _p in [str(_here), str(_root)]:
     if _p not in sys.path:
         sys.path.insert(0, _p)
@@ -27,8 +27,7 @@ st.set_page_config(
 )
 
 # ── Custom CSS ────────────────────────────────────────────────────────────────
-st.markdown(
-    """
+st.markdown("""
 <style>
 @import url('https://fonts.googleapis.com/css2?family=JetBrains+Mono:wght@400;600&family=Syne:wght@400;600;700;800&display=swap');
 
@@ -260,9 +259,7 @@ html, body, [class*="css"] {
     font-weight: 600; margin-bottom: 8px; font-family: var(--mono);
 }
 </style>
-""",
-    unsafe_allow_html=True,
-)
+""", unsafe_allow_html=True)
 
 # ── Plotly dark theme template ────────────────────────────────────────────────
 PLOTLY_TEMPLATE = dict(
@@ -272,16 +269,8 @@ PLOTLY_TEMPLATE = dict(
         font=dict(color="#e8eaf0", family="JetBrains Mono, monospace", size=11),
         xaxis=dict(gridcolor="#2a2f42", linecolor="#2a2f42", zerolinecolor="#2a2f42"),
         yaxis=dict(gridcolor="#2a2f42", linecolor="#2a2f42", zerolinecolor="#2a2f42"),
-        colorway=[
-            "#4f8ef7",
-            "#7c5cfc",
-            "#22d3a5",
-            "#f75f5f",
-            "#f7a24f",
-            "#e879f9",
-            "#38bdf8",
-            "#fb923c",
-        ],
+        colorway=["#4f8ef7", "#7c5cfc", "#22d3a5", "#f75f5f", "#f7a24f",
+                  "#e879f9", "#38bdf8", "#fb923c"],
         legend=dict(bgcolor="#1c2030", bordercolor="#2a2f42", borderwidth=1),
         margin=dict(l=40, r=20, t=40, b=40),
     )
@@ -292,20 +281,14 @@ PLOTLY_TEMPLATE = dict(
 
 CHART_ALIASES = {
     # ── ChartsType enum values (uppercase from Pydantic/Enum .value) ──
-    "pie": "pie",
-    "bar": "bar",
+    "pie":  "pie",
+    "bar":  "bar",
     "line": "line",
     # ── Common string variants (lowercase) ──
-    "bar_chart": "bar",
-    "vertical_bar": "bar",
-    "column": "bar",
-    "grouped_bar": "bar",
-    "stacked_bar": "bar",
-    "line_chart": "line",
-    "time_series": "line",
-    "trend": "line",
-    "donut": "pie",
-    "doughnut": "pie",
+    "bar_chart": "bar", "vertical_bar": "bar", "column": "bar",
+    "grouped_bar": "bar", "stacked_bar": "bar",
+    "line_chart": "line", "time_series": "line", "trend": "line",
+    "donut": "pie", "doughnut": "pie",
 }
 
 
@@ -339,9 +322,7 @@ def _guess_chart_type(df: pd.DataFrame) -> str:
     return "bar"
 
 
-def render_chart(
-    df: pd.DataFrame, chart_type_raw: str | None = None, key_prefix: str = "chart"
-):
+def render_chart(df: pd.DataFrame, chart_type_raw: str | None = None, key_prefix: str = "chart"):
     """
     Render a Plotly chart with user-controlled column selectors.
     The user picks x, y, and optional color columns via st.selectbox.
@@ -353,19 +334,19 @@ def render_chart(
         return
 
     chart_type = _normalise_chart_type(chart_type_raw) or _guess_chart_type(df)
-    cols = list(df.columns)
+    cols    = list(df.columns)
     numeric = df.select_dtypes(include="number").columns.tolist()
-    cat = df.select_dtypes(exclude="number").columns.tolist()
+    cat     = df.select_dtypes(exclude="number").columns.tolist()
 
     # ── Smart defaults: pre-select the most sensible column per role ─────────────────────
-    default_x = cat[0] if cat else cols[0]
+    default_x = cat[0]     if cat     else cols[0]
     default_y = numeric[0] if numeric else (cols[1] if len(cols) > 1 else cols[0])
 
     # ── Column selector UI ─────────────────────────────────────────────────────────────────
     st.markdown(
         '<div style="font-size:0.7rem;color:var(--muted);font-family:var(--mono);'
         'text-transform:uppercase;letter-spacing:0.08em;margin-bottom:8px;">'
-        "⚙️  Configure columns</div>",
+        '⚙️  Configure columns</div>',
         unsafe_allow_html=True,
     )
 
@@ -427,60 +408,36 @@ def render_chart(
     try:
         if chart_type == "bar":
             if color_col:
-                fig = px.bar(
-                    df,
-                    x=x_col,
-                    y=y_col,
-                    color=color_col,
-                    barmode="group",
-                    template=PLOTLY_TEMPLATE,
-                )
+                fig = px.bar(df, x=x_col, y=y_col, color=color_col,
+                             barmode="group", template=PLOTLY_TEMPLATE)
             else:
                 fig = px.bar(df, x=x_col, y=y_col, template=PLOTLY_TEMPLATE)
 
         elif chart_type == "line":
             if color_col:
-                fig = px.line(
-                    df,
-                    x=x_col,
-                    y=y_col,
-                    color=color_col,
-                    markers=True,
-                    template=PLOTLY_TEMPLATE,
-                )
+                fig = px.line(df, x=x_col, y=y_col, color=color_col,
+                              markers=True, template=PLOTLY_TEMPLATE)
             else:
-                fig = px.line(
-                    df, x=x_col, y=y_col, markers=True, template=PLOTLY_TEMPLATE
-                )
+                fig = px.line(df, x=x_col, y=y_col, markers=True,
+                              template=PLOTLY_TEMPLATE)
 
         elif chart_type == "pie":
-            fig = px.pie(
-                df, names=x_col, values=y_col, hole=0.35, template=PLOTLY_TEMPLATE
-            )
+            fig = px.pie(df, names=x_col, values=y_col,
+                         hole=0.35, template=PLOTLY_TEMPLATE)
             fig.update_traces(textinfo="percent+label")
 
         else:  # unrecognised / None → heuristic fallback bar
-            fig = px.bar(
-                df,
-                x=x_col,
-                y=y_col,
-                template=PLOTLY_TEMPLATE,
-                title=f"Chart type '{chart_type_raw}' not recognized — showing bar chart",
-            )
+            fig = px.bar(df, x=x_col, y=y_col, template=PLOTLY_TEMPLATE,
+                         title=f"Chart type '{chart_type_raw}' not recognized — showing bar chart")
 
     except Exception as e:
-        st.warning(
-            f"⚠ Could not render `{chart_type}` chart: {e}. Check the column types selected above."
-        )
+        st.warning(f"⚠ Could not render `{chart_type}` chart: {e}. Check the column types selected above.")
         return
 
     if fig:
-        fig.update_layout(
-            paper_bgcolor="#141720", plot_bgcolor="#0d0f14", font_color="#e8eaf0"
-        )
+        fig.update_layout(paper_bgcolor="#141720", plot_bgcolor="#0d0f14", font_color="#e8eaf0")
         st.plotly_chart(fig, use_container_width=True, config={"displayModeBar": False})
 
-
 def render_crosstab(df: pd.DataFrame):
     """
     Auto-build a crosstab-style pivot summary.
@@ -496,51 +453,41 @@ def render_crosstab(df: pd.DataFrame):
         return
 
     numeric = df.select_dtypes(include="number").columns.tolist()
-    cat = df.select_dtypes(exclude="number").columns.tolist()
+    cat     = df.select_dtypes(exclude="number").columns.tolist()
 
     try:
         if len(cat) >= 2 and len(numeric) >= 1:
             pivot = df.pivot_table(
-                index=cat[0],
-                columns=cat[1],
-                values=numeric[0],
-                aggfunc="sum",
-                fill_value=0,
-            )
-            st.markdown(
-                f'<div class="table-label">📐 Crosstab — {cat[0]} × {cat[1]} (sum of {numeric[0]})</div>',
-                unsafe_allow_html=True,
+                index=cat[0], columns=cat[1],
+                values=numeric[0], aggfunc="sum", fill_value=0,
             )
+            st.markdown(f'<div class="table-label">📐 Crosstab — {cat[0]} × {cat[1]} (sum of {numeric[0]})</div>',
+                        unsafe_allow_html=True)
             st.dataframe(pivot, use_container_width=True)
 
         elif len(cat) == 1 and len(numeric) >= 1:
-            summary = df.groupby(cat[0])[numeric].agg(["sum", "mean", "count"])
+            summary = (
+                df.groupby(cat[0])[numeric]
+                  .agg(["sum", "mean", "count"])
+            )
             summary.columns = [f"{v}_{f}" for v, f in summary.columns]
             summary = summary.reset_index()
-            st.markdown(
-                f'<div class="table-label">📐 Summary — grouped by {cat[0]}</div>',
-                unsafe_allow_html=True,
-            )
+            st.markdown(f'<div class="table-label">📐 Summary — grouped by {cat[0]}</div>',
+                        unsafe_allow_html=True)
             st.dataframe(summary, use_container_width=True, hide_index=True)
 
         elif len(numeric) >= 2:
             corr = df[numeric].corr().round(3)
-            st.markdown(
-                '<div class="table-label">📐 Correlation Matrix</div>',
-                unsafe_allow_html=True,
-            )
-            st.dataframe(
-                corr.style.background_gradient(cmap="Blues", axis=None),
-                use_container_width=True,
-            )
+            st.markdown('<div class="table-label">📐 Correlation Matrix</div>',
+                        unsafe_allow_html=True)
+            st.dataframe(corr.style.background_gradient(cmap="Blues", axis=None),
+                         use_container_width=True)
 
         else:
             desc = df.describe(include="all").T.reset_index()
             desc.rename(columns={"index": "column"}, inplace=True)
-            st.markdown(
-                '<div class="table-label">📐 Statistical Summary</div>',
-                unsafe_allow_html=True,
-            )
+            st.markdown('<div class="table-label">📐 Statistical Summary</div>',
+                        unsafe_allow_html=True)
             st.dataframe(desc, use_container_width=True, hide_index=True)
 
     except Exception as e:
@@ -553,7 +500,6 @@ def render_crosstab(df: pd.DataFrame):
 def get_controller():
     return DataExtractorController()
 
-
 controller = get_controller()
 
 # ── Session state ────────────────────────────────────────────────────────────
@@ -564,7 +510,6 @@ if "total_queries" not in st.session_state:
 if "successful_queries" not in st.session_state:
     st.session_state.successful_queries = 0
 
-
 # ── Helpers ──────────────────────────────────────────────────────────────────
 def build_message_history() -> list[Message]:
     return [
@@ -572,28 +517,23 @@ def build_message_history() -> list[Message]:
         for msg in st.session_state.chat_history
     ]
 
-
 def call_controller(user_query: str):
     uq = UserQuery(user_query=user_query)
     history = build_message_history()
     response = asyncio.run(controller.extrcat(user_query=uq, message_history=history))
     return response
 
-
 def render_message(msg):
     is_user = msg["role"] == "user"
     role_class = "user" if is_user else "bot"
     avatar = "U" if is_user else "AI"
     ts = msg.get("ts", "")
 
-    st.markdown(
-        f"""
+    st.markdown(f"""
     <div class="msg-wrap {role_class}">
         <div class="avatar {role_class}">{avatar}</div>
         <div style="max-width:72%">
-    """,
-        unsafe_allow_html=True,
-    )
+    """, unsafe_allow_html=True)
 
     with st.container():
         if "status" in msg:
@@ -601,36 +541,62 @@ def render_message(msg):
             label = "✓ success" if msg["status"] == "success" else "✗ error"
             st.markdown(
                 f'<span class="badge {badge_class}">{label}</span>',
-                unsafe_allow_html=True,
+                unsafe_allow_html=True
             )
 
         st.markdown(msg["content"])
 
         if msg.get("sql"):
-            st.markdown(
-                '<div class="sql-label">⚡ Generated SQL</div>', unsafe_allow_html=True
-            )
+            st.markdown('<div class="sql-label">⚡ Generated SQL</div>', unsafe_allow_html=True)
             st.code(msg["sql"], language="sql")
 
     st.markdown(f'<div class="ts">{ts}</div>', unsafe_allow_html=True)
     st.markdown("</div></div>", unsafe_allow_html=True)
 
-    # ── Data table ──
-    if msg.get("data") and len(msg["data"]) > 0:
-        st.markdown(
-            '<div class="table-label">📊 Query Results</div>', unsafe_allow_html=True
-        )
-        df = pd.DataFrame(msg["data"])
-        st.dataframe(df, use_container_width=True, hide_index=True)
-    elif (
-        msg.get("status") == "success"
-        and "data" in msg
-        and len(msg.get("data", [])) == 0
-    ):
+    # ── Multi-view data panel ──────────────────────────────────────────────
+    data = msg.get("data", [])
+    chart_hint = msg.get("best_suitable_chart")   # e.g. "bar", "line", "pie"
+
+    if data and len(data) > 0 and msg.get("status") == "success":
+        df = pd.DataFrame(data)
+
+        tab_table, tab_crosstab, tab_chart = st.tabs([
+            "📋  Table",
+            "📐  Crosstab / Summary",
+            "📊  Chart",
+        ])
+
+        with tab_table:
+            st.markdown('<div class="table-label">Query Results</div>', unsafe_allow_html=True)
+            st.dataframe(df, use_container_width=True, hide_index=True)
+
+        with tab_crosstab:
+            render_crosstab(df)
+
+        with tab_chart:
+            if chart_hint:
+                icon_map = {"BAR": "📊", "PIE": "🥧", "LINE": "📈"}
+                icon = icon_map.get(str(chart_hint).upper(), "📊")
+                st.markdown(
+                    f'<div class="chart-label">{icon} {chart_hint}'
+                    f' <span style="color:var(--muted);font-weight:400">(AI suggested)</span></div>',
+                    unsafe_allow_html=True,
+                )
+            else:
+                st.markdown(
+                    '<div class="chart-label">📊 Auto-detected chart type</div>',
+                    unsafe_allow_html=True,
+                )
+            # key_prefix is unique per message (uses timestamp) so each
+            # chart's column selectors have independent Streamlit widget keys.
+            chart_key = f"chart_{msg.get('ts', 'x').replace(':', '_')}"
+            render_chart(df, chart_hint, key_prefix=chart_key)
+
+    elif msg.get("status") == "success" and not data:
         st.markdown(
             '<div style="color:#6b7280;font-size:0.8rem;margin-top:8px;'
             'font-family:monospace">⚠ Query returned 0 rows.</div>',
-            unsafe_allow_html=True,
+            unsafe_allow_html=True
         )
 
 
@@ -669,30 +635,24 @@ with st.sidebar:
 
 
 # ── Main layout ───────────────────────────────────────────────────────────────
-st.markdown(
-    """
+st.markdown("""
 <div style="margin-bottom: 1.5rem;">
   <div class="page-title">Firerms Data Extractor Chatbot</div>
   <div class="page-sub">Natural language → SQL → Results</div>
 </div>
-""",
-    unsafe_allow_html=True,
-)
+""", unsafe_allow_html=True)
 
 # ── Chat area ─────────────────────────────────────────────────────────────────
 chat_container = st.container()
 with chat_container:
     if not st.session_state.chat_history:
-        st.markdown(
-            """
+        st.markdown("""
         <div class="empty-state">
             <div class="empty-icon">🔍</div>
             <div class="empty-title">Ask anything about your data</div>
             <div class="empty-hint">Type a natural language question and the AI will generate SQL and return results.</div>
         </div>
-        """,
-            unsafe_allow_html=True,
-        )
+        """, unsafe_allow_html=True)
     else:
         for msg in st.session_state.chat_history:
             render_message(msg)
@@ -706,28 +666,24 @@ if not prompt and prefill:
 if prompt:
     ts_now = datetime.now().strftime("%H:%M:%S")
 
-    st.session_state.chat_history.append(
-        {
-            "role": "user",
-            "content": prompt,
-            "ts": ts_now,
-        }
-    )
+    st.session_state.chat_history.append({
+        "role": "user",
+        "content": prompt,
+        "ts": ts_now,
+    })
     st.session_state.total_queries += 1
 
     with st.spinner("Generating SQL and fetching results…"):
         try:
             result = call_controller(prompt)
-            status = result.status
-            sql = result.sql_query
-            data = result.data or []
+            status    = result.status
+            sql       = result.sql_query
+            data      = result.data or []
 
             # ── Extract best_suitable_chart from result.output (SQLQueryExtractor) ──
             # result.output.best_suitable_chart is a ChartsType enum → use .value for the string
             try:
-                best_chart = (
-                    result.output.best_suitable_chart.value
-                )  # e.g. "PIE", "BAR", "LINE"
+                best_chart = result.output.best_suitable_chart.value  # e.g. "PIE", "BAR", "LINE"
             except Exception:
                 best_chart = None
 
@@ -738,29 +694,25 @@ if prompt:
                 else f"Query returned status: `{status}`."
             )
 
-            st.session_state.chat_history.append(
-                {
-                    "role": "assistant",
-                    "content": content,
-                    "sql": sql,
-                    "data": data,
-                    "status": status,
-                    "best_suitable_chart": best_chart,
-                    "ts": datetime.now().strftime("%H:%M:%S"),
-                }
-            )
+            st.session_state.chat_history.append({
+                "role": "assistant",
+                "content": content,
+                "sql": sql,
+                "data": data,
+                "status": status,
+                "best_suitable_chart": best_chart,
+                "ts": datetime.now().strftime("%H:%M:%S"),
+            })
 
             if status == "success":
                 st.session_state.successful_queries += 1
 
         except Exception as e:
-            st.session_state.chat_history.append(
-                {
-                    "role": "assistant",
-                    "content": f"❌ Error: {str(e)}",
-                    "status": "error",
-                    "ts": datetime.now().strftime("%H:%M:%S"),
-                }
-            )
-
-    st.rerun()
+            st.session_state.chat_history.append({
+                "role": "assistant",
+                "content": f"❌ Error: {str(e)}",
+                "status": "error",
+                "ts": datetime.now().strftime("%H:%M:%S"),
+            })
+
+    st.rerun()