ArchCoder/llm-excel-plotter-agent
0
1import logging2import os3import time4import uuid5 6import matplotlib7matplotlib.use("Agg")8import matplotlib.pyplot as plt9import pandas as pd10import plotly.graph_objects as go11 12logger = logging.getLogger(__name__)13 14_PLOTLY_LAYOUT = dict(15 font=dict(family="Inter, system-ui, sans-serif", size=13),16 plot_bgcolor="#0f1117",17 paper_bgcolor="#0f1117",18 font_color="#e2e8f0",19 margin=dict(l=60, r=30, t=60, b=60),20 legend=dict(bgcolor="rgba(0,0,0,0)", borderwidth=0),21 xaxis=dict(gridcolor="#1e2d3d", linecolor="#2d3748", zerolinecolor="#2d3748"),22 yaxis=dict(gridcolor="#1e2d3d", linecolor="#2d3748", zerolinecolor="#2d3748"),23 colorway=["#4f8cff", "#34d399", "#f59e0b", "#ef4444", "#a78bfa", "#06b6d4"],24)25 26 27class ChartGenerator:28 def __init__(self, data=None):29 logger.info("Initializing ChartGenerator")30 if data is not None and not (isinstance(data, pd.DataFrame) and data.empty):31 self.data = data32 else:33 default_csv = os.path.join(34 os.path.dirname(__file__), "data", "sample_data.csv"35 )36 self.data = pd.read_csv(default_csv) if os.path.exists(default_csv) else pd.DataFrame()37 38 # -----------------------------------------------------------------------39 # Public40 # -----------------------------------------------------------------------41 42 def generate_chart(self, plot_args: dict) -> dict:43 """Return {"chart_path": str, "chart_spec": dict}."""44 t0 = time.time()45 logger.info(f"Generating chart: {plot_args}")46 47 x_col = plot_args["x"]48 y_cols = plot_args["y"]49 chart_type = plot_args.get("chart_type", "line")50 color = plot_args.get("color", None)51 52 self._validate_columns(x_col, y_cols)53 54 chart_path = self._save_matplotlib(x_col, y_cols, chart_type, color)55 chart_spec = self._build_plotly_spec(x_col, y_cols, chart_type, color)56 57 logger.info(f"Chart ready in {time.time() - t0:.2f}s")58 return {"chart_path": chart_path, "chart_spec": chart_spec}59 60 # -----------------------------------------------------------------------61 # Validation62 # -----------------------------------------------------------------------63 64 def _validate_columns(self, x_col: str, y_cols: list):65 missing = [c for c in [x_col] + y_cols if c not in self.data.columns]66 if missing:67 raise ValueError(68 f"Columns not found in data: {missing}. "69 f"Available: {list(self.data.columns)}"70 )71 72 # -----------------------------------------------------------------------73 # Matplotlib (static PNG)74 # -----------------------------------------------------------------------75 76 def _save_matplotlib(self, x_col, y_cols, chart_type, color) -> str:77 plt.clf()78 plt.close("all")79 fig, ax = plt.subplots(figsize=(10, 6))80 fig.patch.set_facecolor("#0f1117")81 ax.set_facecolor("#0f1117")82 83 palette = ["#4f8cff", "#34d399", "#f59e0b", "#ef4444", "#a78bfa"]84 x = self.data[x_col]85 86 for i, y_col in enumerate(y_cols):87 c = color or palette[i % len(palette)]88 y = self.data[y_col]89 if chart_type == "bar":90 ax.bar(x, y, label=y_col, color=c, alpha=0.85)91 elif chart_type == "scatter":92 ax.scatter(x, y, label=y_col, color=c, alpha=0.8)93 elif chart_type == "area":94 ax.fill_between(x, y, label=y_col, color=c, alpha=0.4)95 ax.plot(x, y, color=c)96 elif chart_type == "histogram":97 ax.hist(y, label=y_col, color=c, alpha=0.8, bins="auto", edgecolor="#1e2d3d")98 elif chart_type == "box":99 ax.boxplot(100 [self.data[y_col].dropna().values for y_col in y_cols],101 labels=y_cols,102 patch_artist=True,103 boxprops=dict(facecolor=c, color="#e2e8f0"),104 medianprops=dict(color="#f59e0b", linewidth=2),105 )106 break107 elif chart_type == "pie":108 ax.pie(109 y, labels=x, autopct="%1.1f%%",110 colors=palette, startangle=90,111 wedgeprops=dict(edgecolor="#0f1117"),112 )113 ax.set_aspect("equal")114 break115 else:116 ax.plot(x, y, label=y_col, color=c, marker="o", linewidth=2)117 118 for spine in ax.spines.values():119 spine.set_edgecolor("#2d3748")120 ax.tick_params(colors="#94a3b8")121 ax.xaxis.label.set_color("#94a3b8")122 ax.yaxis.label.set_color("#94a3b8")123 ax.set_xlabel(x_col, fontsize=11)124 ax.set_ylabel(" / ".join(y_cols), fontsize=11)125 ax.set_title(f"{chart_type.title()} \u2014 {', '.join(y_cols)} vs {x_col}",126 color="#e2e8f0", fontsize=13, pad=12)127 ax.grid(True, alpha=0.15, color="#1e2d3d")128 if chart_type not in ("pie", "histogram"):129 ax.legend(facecolor="#161b27", edgecolor="#2d3748", labelcolor="#e2e8f0")130 if chart_type not in ("pie", "histogram", "box") and len(x) > 5:131 plt.xticks(rotation=45, ha="right")132 133 output_dir = os.path.join(os.path.dirname(__file__), "static", "images")134 os.makedirs(output_dir, exist_ok=True)135 filename = f"chart_{uuid.uuid4().hex[:12]}.png"136 full_path = os.path.join(output_dir, filename)137 plt.savefig(full_path, dpi=150, bbox_inches="tight", facecolor=fig.get_facecolor())138 plt.close(fig)139 logger.info(f"Saved PNG: {full_path} ({os.path.getsize(full_path)} bytes)")140 return os.path.join("static", "images", filename)141 142 # -----------------------------------------------------------------------143 # Plotly (interactive JSON spec for frontend)144 # -----------------------------------------------------------------------145 146 def _build_plotly_spec(self, x_col, y_cols, chart_type, color) -> dict:147 palette = ["#4f8cff", "#34d399", "#f59e0b", "#ef4444", "#a78bfa"]148 x = self.data[x_col].tolist()149 traces = []150 151 for i, y_col in enumerate(y_cols):152 c = color or palette[i % len(palette)]153 y = self.data[y_col].tolist()154 155 if chart_type == "bar":156 traces.append(go.Bar(x=x, y=y, name=y_col, marker_color=c, opacity=0.85).to_plotly_json())157 elif chart_type == "scatter":158 traces.append(go.Scatter(x=x, y=y, name=y_col, mode="markers",159 marker=dict(color=c, size=8, opacity=0.8)).to_plotly_json())160 elif chart_type == "area":161 traces.append(go.Scatter(x=x, y=y, name=y_col, mode="lines",162 fill="tozeroy", line=dict(color=c)).to_plotly_json())163 elif chart_type == "histogram":164 traces.append(go.Histogram(x=y, name=y_col, marker_color=c, opacity=0.8).to_plotly_json())165 elif chart_type == "box":166 traces.append(go.Box(y=y, name=y_col, marker_color=c,167 line_color="#e2e8f0", fillcolor=c).to_plotly_json())168 elif chart_type == "pie":169 traces.append(go.Pie(labels=x, values=y, name=y_col,170 marker=dict(colors=palette)).to_plotly_json())171 break172 else: # line173 traces.append(go.Scatter(x=x, y=y, name=y_col, mode="lines+markers",174 line=dict(color=c, width=2),175 marker=dict(size=6)).to_plotly_json())176 177 layout = {**_PLOTLY_LAYOUT}178 layout["title"] = {179 "text": f"{chart_type.title()} \u2014 {', '.join(y_cols)} vs {x_col}",180 "font": {"size": 15, "color": "#e2e8f0"},181 }182 layout["xaxis"] = {**_PLOTLY_LAYOUT["xaxis"], "title": x_col}183 layout["yaxis"] = {**_PLOTLY_LAYOUT["yaxis"], "title": " / ".join(y_cols)}184 185 return {"data": traces, "layout": layout}