Update agents/brain.py

agents/brain.py

@@ -3,7 +3,10 @@ from agents.reasoner import (
     run_slotting_analysis,
     run_picking_optimization,
     run_demand_forecast,
-    run_replenishment_analysis
+    run_replenishment_analysis,
+    run_rebalancing_analysis,
+    run_workforce_optimization,
+    run_dock_scheduling
 )
 import pandas as pd
 
@@ -12,6 +15,7 @@ class AutoWarehouseAgent:
 
     def run(self, message, slotting_df, picking_df):
         """
+        Central warehouse AI engine.
         Returns a structured dictionary:
         {
             "report": markdown text,
@@ -28,31 +32,29 @@ class AutoWarehouseAgent:
         print("🚚 Picking DF:", picking_df)
 
         # ------------------------------------------------------
-        # 1️⃣ SLOTTING
+        # 1️⃣ SLOTING OPTIMIZATION
         # ------------------------------------------------------
         if task == "slotting":
             explanation, slot_plan = run_slotting_analysis(message, slotting_df)
-
             return {
-                "report": f"### 📦 Slotting Optimization Report\n\n{explanation}",
+                "report": explanation,
                 "route_image": None,
                 "slotting_table": slot_plan
             }
 
         # ------------------------------------------------------
-        # 2️⃣ PICKING
+        # 2️⃣ PICKING ROUTE OPTIMIZATION
         # ------------------------------------------------------
         if task == "picking":
             explanation, route_img = run_picking_optimization(message, picking_df)
-
             return {
-                "report": f"### 🚚 Picking Route Optimization\n\n{explanation}",
+                "report": explanation,
                 "route_image": route_img,
                 "slotting_table": pd.DataFrame()
             }
 
         # ------------------------------------------------------
-        # 3️⃣ FULL REPORT (slotting + picking)
+        # 3️⃣ FULL REPORT — Slotting + Picking
         # ------------------------------------------------------
         if task == "report":
             exp1, slot_plan = run_slotting_analysis(message, slotting_df)
@@ -60,11 +62,9 @@ class AutoWarehouseAgent:
 
             combined_report = (
                 "### 🧾 Full Warehouse Intelligence Report\n\n"
-                "#### 📦 Slotting Optimization\n"
-                + exp1 +
+                "#### 📦 Slotting Optimization\n" + exp1 +
                 "\n\n---\n\n"
-                "#### 🚚 Picking Optimization\n"
-                + exp2
+                "#### 🚚 Picking Optimization\n" + exp2
             )
 
             return {
@@ -78,7 +78,6 @@ class AutoWarehouseAgent:
         # ------------------------------------------------------
         if task == "forecast":
             explanation, forecast_plot, forecast_table = run_demand_forecast(message, slotting_df)
-
             return {
                 "report": explanation,
                 "route_image": forecast_plot,
@@ -86,11 +85,10 @@ class AutoWarehouseAgent:
             }
 
         # ------------------------------------------------------
-        # 5️⃣ REPLENISHMENT ANALYSIS (NEW)
+        # 5️⃣ REPLENISHMENT ANALYSIS
         # ------------------------------------------------------
         if task == "replenishment":
             explanation, repl_table = run_replenishment_analysis(message, slotting_df)
-
             return {
                 "report": explanation,
                 "route_image": None,
@@ -98,18 +96,53 @@ class AutoWarehouseAgent:
             }
 
         # ------------------------------------------------------
-        # ❌ FALLBACK
+        # 6️⃣ INVENTORY REBALANCING
+        # ------------------------------------------------------
+        if task == "rebalancing":
+            explanation, move_table = run_rebalancing_analysis(message, slotting_df)
+            return {
+                "report": explanation,
+                "route_image": None,
+                "slotting_table": move_table
+            }
+
+        # ------------------------------------------------------
+        # 7️⃣ WORKFORCE OPTIMIZATION
+        # ------------------------------------------------------
+        if task == "workforce":
+            explanation, workforce_table = run_workforce_optimization(message, slotting_df)
+            return {
+                "report": explanation,
+                "route_image": None,
+                "slotting_table": workforce_table
+            }
+
+        # ------------------------------------------------------
+        # 8️⃣ DOCK SCHEDULING OPTIMIZATION
+        # ------------------------------------------------------
+        if task == "dock":
+            explanation, dock_table = run_dock_scheduling(message, slotting_df)
+            return {
+                "report": explanation,
+                "route_image": None,
+                "slotting_table": dock_table
+            }
+
+        # ------------------------------------------------------
+        # ❌ FALLBACK HANDLER
         # ------------------------------------------------------
         return {
             "report": (
-                "### ❓ Unable to Understand Request\n"
-                "I couldn't determine whether your request is about:\n"
+                "### ❓ Unable to Understand Your Request\n"
+                "Please try queries related to:\n"
                 "- Slotting optimization\n"
                 "- Picking optimization\n"
-                "- Demand forecasting\n"
-                "- Replenishment needs\n"
-                "- Full warehouse report\n\n"
-                "Try rephrasing your query."
+                "- Forecasting\n"
+                "- Replenishment\n"
+                "- Inventory rebalancing\n"
+                "- Workforce planning\n"
+                "- Dock scheduling\n"
+                "- Full warehouse report\n"
             ),
             "route_image": None,
             "slotting_table": pd.DataFrame()