Update agents/brain.py

agents/brain.py

@@ -1,88 +1,28 @@
-import pandas as pd
-from .planner import detect_intent
-from .reasoner import (
+from agents.planner import detect_intent
+from agents.reasoner import (
     run_slotting_analysis,
-    run_picking_analysis,
-    build_operational_report
+    run_picking_optimization
 )
-from tools.heatmap import generate_heatmap
-
 
 class AutoWarehouseAgent:
 
-    def __init__(self):
-        self.memory = {}
-
-    def run(self, user_query, slotting_df, picking_df):
-        """
-        Autonomous agent interface.
-        Always returns 4 outputs:
-        - report text
-        - picking route image path
-        - heatmap image path
-        - slotting dataframe
-        """
-
-        intent = detect_intent(user_query)
-        self.memory["slotting"] = slotting_df
-        self.memory["picking"] = picking_df
-
-        if intent == "slotting_analysis":
-            return self._do_slotting()
-
-        elif intent == "picking_route":
-            return self._do_picking()
-
-        elif intent == "full_warehouse_assessment":
-            return self._do_full_assessment()
-
-        return (
-            "I can help with slotting, picking route optimization, or full warehouse assessment.",
-            None, None, pd.DataFrame()
-        )
-
-    # ---------------- INTERNAL ---------------- #
-
-    def _do_slotting(self):
-        slot_df = self.memory["slotting"]
-        slot_output, insights = run_slotting_analysis(slot_df)
-
-        return (
-            f"### Slotting Optimization Report\n{insights}",
-            None,
-            None,
-            slot_output
-        )
-
-    def _do_picking(self):
-        picking_df = self.memory["picking"]
-        route_img, picking_summary = run_picking_analysis(picking_df)
+    def run(self, message, slotting_df, picking_df):
 
-        return (
-            "### Picking Route Optimization Completed",
-            route_img,
-            None,
-            picking_summary
-        )
+        task = detect_intent(message)
 
-    def _do_full_assessment(self):
-        slot_df = self.memory["slotting"]
-        picking_df = self.memory["picking"]
+        if task == "slotting":
+            explanation, slotting_plan, _ = run_slotting_analysis(message, slotting_df)
+            return explanation, None, None, slotting_plan
 
-        slot_output, slot_insights = run_slotting_analysis(slot_df)
-        route_img, picking_summary = run_picking_analysis(picking_df)
-        heatmap_img = generate_heatmap(slot_output)
+        if task == "picking":
+            explanation, route_img = run_picking_optimization(message, picking_df)
+            return explanation, route_img, None, None
 
-        report = build_operational_report(
-            slotting_table=slot_output,
-            picking_table=picking_summary,
-            insights=slot_insights,
-            heatmap_path=heatmap_img
-        )
+        if task == "report":
+            # Combine both
+            exp1, slot_plan, _ = run_slotting_analysis(message, slotting_df)
+            exp2, route_img = run_picking_optimization(message, picking_df)
+            full_exp = exp1 + "\n\n" + exp2
+            return full_exp, route_img, None, slot_plan
 
-        return (
-            report,
-            route_img,
-            heatmap_img,
-            slot_output
-        )
+        return "Sorry, I could not understand your request.", None, None, None