BizIntel_AI / tools /plot_generator.py
mgbam's picture
Update tools/plot_generator.py
ccdbd61 verified
# tools/plot_generator.py
# ------------------------------------------------------------
# Creates an interactive line‑and‑marker trend chart for any
# (date_col, value_col) pair and saves a hi‑res PNG copy.
import os
import tempfile
from typing import Tuple, Union
import pandas as pd
import plotly.graph_objects as go
# Alias for typing β€” every helper returns a go.Figure
Plot = go.Figure
def plot_metric_tool(
file_path: str,
date_col: str,
value_col: str,
output_dir: str = "/tmp",
title: str | None = None,
line_width: int = 2,
marker_size: int = 6,
) -> Union[Tuple[Plot, str], str]:
"""
Build a (date, metric) trend chart.
Returns
-------
(fig, png_path) on success
error string on failure (string starts with '❌')
"""
# ── 1. Load CSV or Excel ──────────────────────────────────
ext = os.path.splitext(file_path)[1].lower()
try:
df = (
pd.read_excel(file_path)
if ext in (".xls", ".xlsx")
else pd.read_csv(file_path)
)
except Exception as exc:
return f"❌ Failed to load file: {exc}"
# ── 2. Validate columns ───────────────────────────────────
missing = [c for c in (date_col, value_col) if c not in df.columns]
if missing:
return f"❌ Missing column(s): {', '.join(missing)}"
# ── 3. Parse & clean ──────────────────────────────────────
df[date_col] = pd.to_datetime(df[date_col], errors="coerce")
df[value_col] = pd.to_numeric(df[value_col], errors="coerce")
df = df.dropna(subset=[date_col, value_col])
if df.empty:
return f"❌ No valid data after cleaning '{date_col}' / '{value_col}'."
# Aggregate duplicate timestamps, sort by date
df = (
df[[date_col, value_col]]
.groupby(date_col, as_index=True)
.mean()
.sort_index()
)
# ── 4. Build Plotly figure ────────────────────────────────
fig = go.Figure(
go.Scatter(
x=df.index,
y=df[value_col],
mode="lines+markers",
line=dict(width=line_width),
marker=dict(size=marker_size),
name=value_col,
)
)
fig.update_layout(
title=title or f"{value_col} Trend",
xaxis_title=date_col,
yaxis_title=value_col,
template="plotly_dark",
hovermode="x unified",
)
# ── 5. Save static PNG copy ───────────────────────────────
os.makedirs(output_dir, exist_ok=True)
tmp = tempfile.NamedTemporaryFile(
prefix="trend_", suffix=".png", dir=output_dir, delete=False
)
png_path = tmp.name
tmp.close()
try:
fig.write_image(png_path, scale=2)
except Exception as exc:
return f"❌ Failed saving image: {exc}"
return fig, png_path