Update main.py

main.py

@@ -189,12 +189,14 @@ ERROR_FAMILY_SYNS = {
     ),
 }
 
+
+
 def _try_wms_tool(user_text: str) -> dict | None:
     """
     Uses Gemini function calling + TOOL_REGISTRY to decide if a WMS tool should be called.
     Returns YOUR existing response dict shape. If no tool match or any error -> returns None.
     """
-    # Safe guards: only attempt if Gemini key exists and WMS login works
+    # Guards
     if not GEMINI_API_KEY:
         return None
     if not _ensure_wms_session():
@@ -207,7 +209,13 @@ def _try_wms_tool(user_text: str) -> dict | None:
     }
 
     try:
-        resp = requests.post(GEMINI_URL, headers=headers, json=payload, timeout=25, verify=GEMINI_SSL_VERIFY)
+        resp = requests.post(
+            GEMINI_URL,
+            headers=headers,
+            json=payload,
+            timeout=25,
+            verify=GEMINI_SSL_VERIFY
+        )
         data = resp.json()
         candidates = data.get("candidates", [])
         part = candidates[0].get("content", {}).get("parts", [{}])[0] if candidates else {}
@@ -220,20 +228,25 @@ def _try_wms_tool(user_text: str) -> dict | None:
         if not cfg:
             return None
 
-        # Prepare arguments and warehouse
+        # -------- Prepare arguments and warehouse --------
         args = fn.get("args", {}) or {}
-        for k in ("wh_id", "warehouseId", "warehouse_id"):  # WMS variations
-            args.pop(k, None)
-        wh = _WMS_WAREHOUSE_ID
+        # Respect explicit warehouse override if present, else use session default
+        wh = (
+            args.pop("warehouse_id", None)
+            or args.pop("warehouseId", None)
+            or args.pop("wh_id", None)
+            or _WMS_WAREHOUSE_ID
+        )
 
-        # Call underlying function per signature
+        # -------- Call the underlying function with correct signature --------
+        # Holds & Location use kwargs (e.g., lodnum/subnum/dtlnum OR stoloc); Item uses single ID
         if tool_name in ("get_inventory_holds", "get_location_status"):
-            raw = cfg["function"]
+            raw = cfg"function"   # <<< invoke function
         else:
             target = args.get(cfg["id_param"])
-            raw = cfg["function"]
+            raw = cfg"function"   # <<< invoke function
 
-        # Error path -> keep your response shape
+        # -------- Error handling --------
         if raw.get("error"):
             return {
                 "bot_response": f"⚠️ {raw['error']}",
@@ -245,6 +258,7 @@ def _try_wms_tool(user_text: str) -> dict | None:
                 "top_hits": [],
                 "sources": [],
                 "debug": {"intent": "wms_error", "tool": tool_name},
+                "source": "ERROR",  # optional for frontend branches
             }
 
         data_list = raw.get(cfg.get("response_key", "data"), [])
@@ -259,22 +273,26 @@ def _try_wms_tool(user_text: str) -> dict | None:
                 "top_hits": [],
                 "sources": [],
                 "debug": {"intent": "wms_empty", "tool": tool_name},
+                "source": "LIVE_API",  # still a live call, just no rows
             }
 
-        # Holds / Location → concise lines
+        # -------- TABLE: Holds / Location --------
         if tool_name in ("get_inventory_holds", "get_location_status"):
-            rows, total_qty = [], 0
+            rows_for_text = []
+            total_qty = 0
+
             for item in data_list:
                 flat = flatten_json(item)
                 row = extract_fields(flat, cfg["mapping"], cfg.get("formatters", {}))
-                rows.append("• " + "; ".join(f"{k}: {v}" for k, v in row.items()))
-                total_qty += int(item.get("untqty", 0))
+                rows_for_text.append("• " + "; ".join(f"{k}: {v}" for k, v in row.items()))
 
             if tool_name == "get_inventory_holds":
-                msg = f"Found {len(rows)} hold records (Total Qty: {total_qty}).\n" + "\n".join(rows[:10])
+                # Only holds have quantity in most schemas
+                total_qty = sum(int(it.get("untqty", 0) or 0) for it in data_list)
+                msg = f"Found {len(rows_for_text)} hold records (Total Qty: {total_qty}).\n" + "\n".join(rows_for_text[:10])
             else:
                 target_loc = args.get("stoloc") or ""
-                msg = f"Location '{target_loc}' status:\n" + ("\n".join(rows[:1]) if rows else "No status returned")
+                msg = f"Location '{target_loc}' status:\n" + ("\n".join(rows_for_text[:1]) if rows_for_text else "No status returned")
 
             return {
                 "bot_response": msg,
@@ -286,9 +304,15 @@ def _try_wms_tool(user_text: str) -> dict | None:
                 "top_hits": [],
                 "sources": [],
                 "debug": {"intent": tool_name, "warehouse_id": wh},
+
+                # Front-end signals → render table exactly like your friend's code
+                "source": "LIVE_API",
+                "type": "table",
+                "data": data_list,  # rows for <TableMessage />
+                "show_export": True if tool_name == "get_inventory_holds" else False,
             }
 
-        # Item → short summary using mapping & requested fields
+        # -------- ITEM: summary + single-row table (consistent with "data => table") --------
         flat = flatten_json(data_list[0])
         requested = args.get("fields", [])
         if requested:
@@ -313,6 +337,11 @@ def _try_wms_tool(user_text: str) -> dict | None:
             "top_hits": [],
             "sources": [],
             "debug": {"intent": "get_item_data", "warehouse_id": wh},
+
+            "source": "LIVE_API",
+            "type": "table",
+            "data": [data_list[0]],  # single-row table for item
+            "show_export": False,
         }
 
     except Exception as e: