Update tools/forecaster.py

tools/forecaster.py

@@ -6,55 +6,70 @@ import plotly.graph_objects as go
 
 def forecast_metric_tool(file_path: str, date_col: str, value_col: str):
     """
-    Forecast next 3 periods for any numeric metric, saving
-    the PNG under /tmp and returning the forecast table as text.
+    Forecast the next 3 periods for any numeric metric.
+    - Saves a date‐indexed Plotly PNG under /tmp via the safe write monkey‐patch.
+    - Returns a text table of the forecast.
     """
-    # 1) Load & parse dates
+
+    # 0) Read full CSV
     df = pd.read_csv(file_path)
+
+    # 1) Check that both columns actually exist
+    if date_col not in df.columns:
+        return f"❌ Date column '{date_col}' not found in your data."
+    if value_col not in df.columns:
+        return f"❌ Metric column '{value_col}' not found in your data."
+
+    # 2) Parse dates
     try:
         df[date_col] = pd.to_datetime(df[date_col])
     except Exception:
         return f"❌ Could not parse '{date_col}' as dates."
 
-    # 2) Coerce metric to numeric & drop invalid rows
+    # 3) Coerce metric to numeric & drop invalid rows
     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 for '{value_col}'."
+        return f"❌ After coercion, no valid data remains for '{value_col}'."
 
-    # 3) Sort by date, set index, then collapse any duplicate timestamps
+    # 4) Sort & index by date, collapse duplicates
     df = df.sort_values(date_col).set_index(date_col)
-    # If you have multiple rows for the same timestamp, take their mean
     df = df[[value_col]].groupby(level=0).mean()
 
-    # 4) Infer frequency (e.g. 'D', 'M', etc.) and reindex
+    # 5) Infer a frequency and re‐index
     freq = pd.infer_freq(df.index)
     if freq is None:
         freq = "D"  # fallback to daily
     df = df.asfreq(freq)
 
-    # 5) Fit ARIMA
+    # 6) Fit ARIMA (1,1,1)
     try:
         model     = ARIMA(df[value_col], order=(1, 1, 1))
         model_fit = model.fit()
     except Exception as e:
         return f"❌ ARIMA fitting failed: {e}"
 
-    # 6) Forecast with a proper DatetimeIndex
+    # 7) Produce a proper date‐indexed forecast
     fc_res   = model_fit.get_forecast(steps=3)
-    forecast = fc_res.predicted_mean  # pd.Series indexed by future dates
+    forecast = fc_res.predicted_mean
 
-    # 7) Plot history + forecast
+    # 8) Plot history + forecast
     fig = go.Figure()
-    fig.add_scatter(x=df.index,       y=df[value_col], mode="lines", name=value_col)
-    fig.add_scatter(x=forecast.index, y=forecast,     mode="lines+markers", name="Forecast")
+    fig.add_scatter(
+        x=df.index, y=df[value_col],
+        mode="lines", name=value_col
+    )
+    fig.add_scatter(
+        x=forecast.index, y=forecast,
+        mode="lines+markers", name="Forecast"
+    )
     fig.update_layout(
         title=f"{value_col} Forecast",
         xaxis_title=date_col,
         yaxis_title=value_col,
         template="plotly_dark",
     )
-    fig.write_image("forecast_plot.png")  # safely lands in /tmp via monkey-patch
+    fig.write_image("forecast_plot.png")  # lands in /tmp via our monkey‐patch
 
-    # 8) Return the forecast table as plain text
+    # 9) Return the forecast as a text table
     return forecast.to_frame(name="Forecast").to_string()