Update tools/forecaster.py

tools/forecaster.py

@@ -6,39 +6,61 @@ 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.
-    Saves PNG to /tmp via our safe write monkey-patch, returns forecast table as text.
+    Forecast next 3 periods for any numeric metric, and return a textual table.
+    Saves a date-indexed PNG under /tmp via our safe monkey-patch.
     """
-    # 1) Load
+    # 1) Load & parse
     df = pd.read_csv(file_path)
-
-    # 2) Parse dates
     try:
         df[date_col] = pd.to_datetime(df[date_col])
     except Exception:
         return f"❌ Could not parse '{date_col}' as dates."
-
-    # 3) Coerce metric to numeric, drop invalid
+    
+    # 2) Coerce metric to numeric & drop invalid
     df[value_col] = pd.to_numeric(df[value_col], errors="coerce")
-    series = df.set_index(date_col)[value_col].dropna()
-
-    if series.empty:
-        return f"❌ Column '{value_col}' has no valid numeric data after coercion."
-
+    df = df.dropna(subset=[date_col, value_col])
+    
+    if df.empty:
+        return f"❌ No valid data for '{value_col}'."
+    
+    # 3) Sort and set index, infer frequency
+    df = df.sort_values(date_col)
+    df.set_index(date_col, inplace=True)
+    freq = pd.infer_freq(df.index)
+    if freq is None:
+        # fallback to daily if pandas can't infer
+        freq = "D"
+    df = df.asfreq(freq)
+    
     # 4) Fit ARIMA
     try:
-        model     = ARIMA(series, order=(1, 1, 1))
+        model     = ARIMA(df[value_col], order=(1, 1, 1))
         model_fit = model.fit()
     except Exception as e:
         return f"❌ ARIMA fitting failed: {e}"
-
-    # 5) Forecast & plot
-    forecast = model_fit.forecast(steps=3)
+    
+    # 5) Get a proper date-indexed forecast
+    fc_res   = model_fit.get_forecast(steps=3)
+    forecast = fc_res.predicted_mean  # a pd.Series with a DatetimeIndex
+    
+    # 6) Plot historical + forecast
     fig = go.Figure()
-    fig.add_scatter(x=series.index, y=series, mode="lines", name=value_col)
-    fig.add_scatter(x=forecast.index, y=forecast, mode="lines", name="Forecast")
-    fig.update_layout(title=f"{value_col} Forecast", template="plotly_dark")
-    fig.write_image("forecast_plot.png")  # goes into /tmp thanks to our monkey-patch
-
-    # 6) Return textual table
-    return forecast.to_frame(name="Forecast").to_string()
+    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=str(date_col),
+        yaxis_title=str(value_col),
+        template="plotly_dark"
+    )
+    fig.write_image("forecast_plot.png")
+    
+    # 7) Return the forecast table as text
+    tbl = forecast.to_frame(name="Forecast")
+    return tbl.to_string()