Update app.py

app.py

@@ -1,3 +1,4 @@
+# space/app.py
 import os
 import json
 import gradio as gr
@@ -7,10 +8,11 @@ from tools.sql_tool import SQLTool
 from tools.predict_tool import PredictTool
 from tools.explain_tool import ExplainTool
 from tools.report_tool import ReportTool
+from tools.ts_preprocess import build_timeseries
+from tools.ts_forecast_tool import TimeseriesForecastTool
+
 from utils.tracing import Tracer
 from utils.config import AppConfig
-from tools.ts_forecast_tool import TimeseriesForecastTool
-from tools.ts_preprocess import build_timeseries
 
 # Optional tiny CPU LLM for planning (can be disabled by not setting ORCHESTRATOR_MODEL)
 llm = None
@@ -31,17 +33,12 @@ sql_tool = SQLTool(cfg, tracer)
 predict_tool = PredictTool(cfg, tracer)
 explain_tool = ExplainTool(cfg, tracer)
 report_tool = ReportTool(cfg, tracer)
-
-ts_tool = TimeseriesForecastTool(cfg, tracer,
-    context_length=512, forecast_length=96,
-    target_cols=["portfolio_value"],          # set after preprocessing
-    control_cols=["rate_deposit", "rate_asset"]  # optional exogenous
-)
+ts_tool = TimeseriesForecastTool(cfg, tracer)  # Granite wrapper
 
 SYSTEM_PROMPT = (
     "You are an analytical assistant for tabular data. "
     "Decide which tools to call in order: "
-    "1) SQL (retrieve) 2) Predict (score) 3) Explain (SHAP) 4) Report (document). "
+    "1) SQL (retrieve) 2) Predict (score) 3) Explain (SHAP) 4) Report (document) 5) Forecast (Granite TTM). "
     "Always disclose the steps taken."
 )
 
@@ -49,7 +46,7 @@ def plan_actions(message: str):
     if llm is not None:
         prompt = (
             f"{SYSTEM_PROMPT}\nUser: {message}\n"
-            "Return JSON with fields: steps (array subset of ['sql','predict','explain','report']), rationale."
+            "Return JSON with fields: steps (array subset of ['sql','predict','explain','report','forecast']), rationale."
         )
         try:
             out = llm(prompt)[0]["generated_text"]
@@ -66,6 +63,7 @@ def plan_actions(message: str):
     if any(k in m for k in ["predict", "score", "risk", "propensity", "probability"]): steps.append("predict")
     if any(k in m for k in ["why", "explain", "shap", "feature", "attribution"]): steps.append("explain")
     if any(k in m for k in ["report", "download", "pdf", "summary"]): steps.append("report")
+    if any(k in m for k in ["forecast", "next", "horizon", "granite"]): steps.append("forecast")
     if not steps: steps = ["sql"]
     return {"steps": steps, "rationale": "Rule-based plan."}
 
@@ -78,34 +76,43 @@ def run_agent(message: str, hitl_decision: str = "Approve", reviewer_note: str =
     predict_df = None
     explain_imgs = {}
     artifacts = {}
+    ts_forecast_df = None
 
     if "sql" in plan["steps"]:
         sql_df = sql_tool.run(message)
-        # 👉 Build the model-friendly time series (monthly)
+        artifacts["sql_rows"] = int(len(sql_df)) if isinstance(sql_df, pd.DataFrame) else 0
+
+    if "predict" in plan["steps"]:
+        predict_df = predict_tool.run(sql_df)
+
+    ts_df = None
+    if sql_df is not None:
         try:
             ts_df = build_timeseries(sql_df)
         except Exception:
             ts_df = None
-        
-    if "sql" in plan["steps"]:
-        sql_df = sql_tool.run(message)
-        artifacts["sql_rows"] = int(len(sql_df)) if isinstance(sql_df, pd.DataFrame) else 0
 
-    if "predict" in plan["steps"]:
-       predict_df = predict_tool.run(sql_df)
+    if "forecast" in plan["steps"] and ts_df is not None:
+        # Expect 'portfolio_value' after preprocessing
+        # Use the combined series — e.g., sum over instruments by timestamp
+        agg = ts_df.groupby("timestamp", as_index=True)["portfolio_value"].sum().sort_index()
+        try:
+            ts_forecast_df = ts_tool.zeroshot_forecast(agg, horizon=96)
+        except Exception as e:
+            # Surface a readable error in response
+            ts_forecast_df = pd.DataFrame({"error": [str(e)]})
 
     if "explain" in plan["steps"]:
         explain_imgs = explain_tool.run(predict_df or sql_df)
 
-    if "forecast" in plan["steps"] and ts_df is not None:
-        fc = ts_tool.zeroshot_forecast(ts_df[["timestamp","portfolio_value","interest_rate"]].dropna())
-
     report_link = None
     if "report" in plan["steps"]:
+        # Add forecast preview if available
+        forecast_preview = ts_forecast_df.head(50) if isinstance(ts_forecast_df, pd.DataFrame) else None
         report_link = report_tool.render_and_save(
             user_query=message,
             sql_preview=sql_df.head(50) if isinstance(sql_df, pd.DataFrame) else None,
-            predict_preview=predict_df.head(50) if isinstance(predict_df, pd.DataFrame) else None,
+            predict_preview=predict_df.head(50) if isinstance(predict_df, pd.DataFrame) else forecast_preview,
             explain_images=explain_imgs,
             plan=plan,
         )
@@ -118,13 +125,18 @@ def run_agent(message: str, hitl_decision: str = "Approve", reviewer_note: str =
         "plan": plan,
     })
 
+    # Compose response
     response = f"**Plan:** {plan['steps']}\n**Rationale:** {plan['rationale']}\n"
     if isinstance(sql_df, pd.DataFrame): response += f"\n**SQL rows:** {len(sql_df)}"
     if isinstance(predict_df, pd.DataFrame): response += f"\n**Predictions rows:** {len(predict_df)}"
+    if isinstance(ts_forecast_df, pd.DataFrame) and "forecast" in ts_forecast_df.columns:
+        response += f"\n**Forecast horizon:** {len(ts_forecast_df)}"
     if report_link: response += f"\n**Report:** {report_link}"
     if tracer.trace_url: response += f"\n**Trace:** {tracer.trace_url}"
 
-    preview_df = predict_df if isinstance(predict_df, pd.DataFrame) and len(predict_df) else sql_df
+    # Prefer to show forecast if present, else predictions, else raw query
+    preview_df = ts_forecast_df if isinstance(ts_forecast_df, pd.DataFrame) and not ts_forecast_df.empty else \
+                 (predict_df if isinstance(predict_df, pd.DataFrame) and not predict_df.empty else sql_df)
     return response, (preview_df if isinstance(preview_df, pd.DataFrame) else pd.DataFrame())
 
 with gr.Blocks() as demo: