show charts

src/controllers/_data_extractor.py

@@ -18,6 +18,7 @@ class Response(BaseModel):
     sql_query: str
     output: SQLQueryExtractor
 
+
 class DataExtractorController:
     def __init__(self):
         self.router = APIRouter()
@@ -36,7 +37,9 @@ class DataExtractorController:
                 response, sql_query, output = await service.extract(
                     user_query=user_query.user_query, message_history=message_history
                 )
-                return Response(status="success", data=response, sql_query=sql_query, output=output)
+                return Response(
+                    status="success", data=response, sql_query=sql_query, output=output
+                )
         except HTTPException as e:
             logger.error(e)
             raise e

src/schemas/_pydantic_agent.py

@@ -118,12 +118,13 @@ class OrderConditions(BaseModel):
             "Use 'ASC' for: oldest, lowest, alphabetical, earliest."
         ),
     )
-    
-    
+
+
 class ChartsType(Enum):
     """
     Enumeration of supported chart types for visualization.
     """
+
     PIE: str = "PIE"
     BAR: str = "BAR"
     LINE: str = "LINE"
@@ -221,7 +222,7 @@ class SQLQueryExtractor(BaseModel):
             "This must be set True if provided tool return true else set False. "
         ),
     )
-    
+
     best_suitable_chart: ChartsType = Field(
         ...,
         description=(
@@ -233,7 +234,6 @@ class SQLQueryExtractor(BaseModel):
         ),
     )
 
-
     @model_validator(mode="after")
     def validate_sql_query(self):
         if not self.is_sql_query_verified_using_provided_tool:

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,7 +27,8 @@ 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');
 
@@ -259,7 +260,9 @@ 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(
@@ -269,8 +272,16 @@ 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),
     )
@@ -281,14 +292,20 @@ 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",
 }
 
 
@@ -322,7 +339,9 @@ 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.
@@ -334,19 +353,19 @@ def render_chart(df: pd.DataFrame, chart_type_raw: str | None = None, key_prefix
         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,
     )
 
@@ -408,36 +427,60 @@ def render_chart(df: pd.DataFrame, chart_type_raw: str | None = None, key_prefix
     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.
@@ -453,41 +496,51 @@ 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,
+                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.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:
@@ -500,6 +553,7 @@ def render_crosstab(df: pd.DataFrame):
 def get_controller():
     return DataExtractorController()
 
+
 controller = get_controller()
 
 # ── Session state ────────────────────────────────────────────────────────────
@@ -510,6 +564,7 @@ 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 [
@@ -517,23 +572,28 @@ 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:
@@ -541,13 +601,15 @@ 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)
@@ -555,14 +617,20 @@ def render_message(msg):
 
     # ── Data table ──
     if msg.get("data") and len(msg["data"]) > 0:
-        st.markdown('<div class="table-label">📊 Query Results</div>', unsafe_allow_html=True)
+        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:
+    elif (
+        msg.get("status") == "success"
+        and "data" in msg
+        and len(msg.get("data", [])) == 0
+    ):
         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,
         )
 
 
@@ -601,24 +669,30 @@ 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)
@@ -632,24 +706,28 @@ 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
 
@@ -660,25 +738,29 @@ 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()