Update agents/brain.py

agents/brain.py

@@ -1,7 +1,9 @@
 from agents.planner import detect_intent
 from agents.reasoner import (
     run_slotting_analysis,
-    run_picking_optimization
+    run_picking_optimization,
+    run_demand_forecast,
+    run_replenishment_analysis
 )
 import pandas as pd
 
@@ -20,7 +22,7 @@ class AutoWarehouseAgent:
 
         task = detect_intent(message.lower().strip())
 
-        # Debug prints
+        # Debug logs
         print(f"🧠 DETECTED TASK: {task}")
         print("📦 Slotting DF:", slotting_df)
         print("🚚 Picking DF:", picking_df)
@@ -50,7 +52,7 @@ class AutoWarehouseAgent:
             }
 
         # ------------------------------------------------------
-        # 3️⃣ FULL REPORT (both slotting + picking)
+        # 3️⃣ FULL REPORT (slotting + picking)
         # ------------------------------------------------------
         if task == "report":
             exp1, slot_plan = run_slotting_analysis(message, slotting_df)
@@ -71,6 +73,30 @@ class AutoWarehouseAgent:
                 "slotting_table": slot_plan
             }
 
+        # ------------------------------------------------------
+        # 4️⃣ DEMAND FORECASTING
+        # ------------------------------------------------------
+        if task == "forecast":
+            explanation, forecast_plot, forecast_table = run_demand_forecast(message, slotting_df)
+
+            return {
+                "report": explanation,
+                "route_image": forecast_plot,
+                "slotting_table": forecast_table
+            }
+
+        # ------------------------------------------------------
+        # 5️⃣ REPLENISHMENT ANALYSIS (NEW)
+        # ------------------------------------------------------
+        if task == "replenishment":
+            explanation, repl_table = run_replenishment_analysis(message, slotting_df)
+
+            return {
+                "report": explanation,
+                "route_image": None,
+                "slotting_table": repl_table
+            }
+
         # ------------------------------------------------------
         # ❌ FALLBACK
         # ------------------------------------------------------
@@ -79,9 +105,11 @@ class AutoWarehouseAgent:
                 "### ❓ Unable to Understand Request\n"
                 "I couldn't determine whether your request is about:\n"
                 "- Slotting optimization\n"
-                "- Picking route planning\n"
+                "- Picking optimization\n"
+                "- Demand forecasting\n"
+                "- Replenishment needs\n"
                 "- Full warehouse report\n\n"
-                "Try rewriting your query."
+                "Try rephrasing your query."
             ),
             "route_image": None,
             "slotting_table": pd.DataFrame()