Update main.py

main.py

@@ -191,11 +191,13 @@ 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.
     """
+
     # Guards
     if not GEMINI_API_KEY:
         return None
@@ -216,21 +218,33 @@ def _try_wms_tool(user_text: str) -> dict | None:
             timeout=25,
             verify=GEMINI_SSL_VERIFY
         )
+        # Raise if HTTP not 2xx so we handle cleanly
+        resp.raise_for_status()
         data = resp.json()
+
         candidates = data.get("candidates", [])
-        part = candidates[0].get("content", {}).get("parts", [{}])[0] if candidates else {}
-        if "functionCall" not in part:
+        if not candidates:
+            return None
+
+        # Defensive extraction of the first part
+        content = candidates[0].get("content", {})
+        parts = content.get("parts", [])
+        part = parts[0] if parts else {}
+
+        # Gemini returns a function call in part["functionCall"] when a tool is selected
+        fn_call = part.get("functionCall")
+        if not fn_call or "name" not in fn_call:
             return None
 
-        fn = part["functionCall"]
-        tool_name = fn["name"]
+        tool_name = fn_call["name"]
+        args = fn_call.get("args", {}) or {}
+
         cfg = TOOL_REGISTRY.get(tool_name)
         if not cfg:
+            # Tool not registered → bail
             return None
 
-        # -------- Prepare arguments and warehouse --------
-        args = fn.get("args", {}) or {}
-        # Respect explicit warehouse override if present, else use session default
+        # -------- Resolve warehouse --------
         wh = (
             args.pop("warehouse_id", None)
             or args.pop("warehouseId", None)
@@ -238,16 +252,60 @@ def _try_wms_tool(user_text: str) -> dict | None:
             or _WMS_WAREHOUSE_ID
         )
 
-        # -------- Call the underlying function with correct signature --------
-        # Holds & Location use kwargs (e.g., lodnum/subnum/dtlnum OR stoloc); Item uses single ID
+        # -------- Invoke the tool callable correctly --------
+        tool_fn = cfg.get("function")
+        if not callable(tool_fn):
+            # Misconfigured registry
+            return {
+                "bot_response": f"⚠️ Tool '{tool_name}' is not callable. Check TOOL_REGISTRY['{tool_name}']['function'].",
+                "status": "PARTIAL",
+                "context_found": False,
+                "ask_resolved": False,
+                "suggest_incident": True,
+                "followup": "Should I raise a ServiceNow ticket?",
+                "top_hits": [],
+                "sources": [],
+                "debug": {"intent": "wms_config_error", "tool": tool_name},
+                "source": "ERROR",
+            }
+
+        # Holds & Location use kwargs (e.g., lodnum/subnum/dtlnum OR stoloc);
+        # Item uses single ID parameter (e.g., item_id) plus warehouse.
         if tool_name in ("get_inventory_holds", "get_location_status"):
-            raw = cfg"function"   # <<< invoke function
+            # Pass warehouse as kwarg if present; rest comes from args
+            call_kwargs = {"warehouse_id": wh} if wh is not None else {}
+            call_kwargs.update(args)
+            raw = tool_fn(**call_kwargs)
         else:
-            target = args.get(cfg["id_param"])
-            raw = cfg"function"   # <<< invoke function
+            id_param = cfg.get("id_param")
+            target = args.get(id_param)
+            if id_param and target is None:
+                return {
+                    "bot_response": f"Couldn’t find required parameter '{id_param}' in function args.",
+                    "status": "PARTIAL",
+                    "context_found": False,
+                    "ask_resolved": False,
+                    "suggest_incident": False,
+                    "followup": f"Can you provide the {id_param}?",
+                    "top_hits": [],
+                    "sources": [],
+                    "debug": {"intent": "wms_missing_id_param", "tool": tool_name},
+                    "source": "CLIENT",
+                }
+
+            call_kwargs = {"warehouse_id": wh} if wh is not None else {}
+            # Include the ID param, plus any other Gemini-parsed args
+            if id_param:
+                call_kwargs[id_param] = target
+            # Keep other non-ID args (filters, etc.)
+            for k, v in args.items():
+                if k != id_param:
+                    call_kwargs[k] = v
+
+            raw = tool_fn(**call_kwargs)
 
-        # -------- Error handling --------
-        if raw.get("error"):
+        # -------- Error handling from tool layer --------
+        if isinstance(raw, dict) and raw.get("error"):
             return {
                 "bot_response": f"⚠️ {raw['error']}",
                 "status": "PARTIAL",
@@ -258,10 +316,12 @@ 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
+                "source": "ERROR",
             }
 
-        data_list = raw.get(cfg.get("response_key", "data"), [])
+        # -------- Extract data list (default key = 'data') --------
+        response_key = cfg.get("response_key", "data")
+        data_list = (raw.get(response_key, []) if isinstance(raw, dict) else []) or []
         if not data_list:
             return {
                 "bot_response": "I searched WMS but couldn't find matching records.",
@@ -273,7 +333,7 @@ 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
+                "source": "LIVE_API",
             }
 
         # -------- TABLE: Holds / Location --------
@@ -288,7 +348,13 @@ def _try_wms_tool(user_text: str) -> dict | None:
 
             if tool_name == "get_inventory_holds":
                 # Only holds have quantity in most schemas
-                total_qty = sum(int(it.get("untqty", 0) or 0) for it in data_list)
+                def _to_int(x):
+                    try:
+                        return int(x or 0)
+                    except Exception:
+                        return 0
+
+                total_qty = sum(_to_int(it.get("untqty", 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 ""
@@ -304,16 +370,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 />
+                "data": data_list,
                 "show_export": True if tool_name == "get_inventory_holds" else False,
             }
 
-        # -------- ITEM: summary + single-row table (consistent with "data => table") --------
+        # -------- ITEM: summary + single-row table --------
         flat = flatten_json(data_list[0])
+
         requested = args.get("fields", [])
         if requested:
             filtered_map = {k: v for k, v in cfg["mapping"].items() if k in requested or k == "Item"}
@@ -337,18 +402,30 @@ 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
+            "data": [data_list[0]],
             "show_export": False,
         }
 
+    except requests.HTTPError as e:
+        # Surface Gemini API issues more clearly
+        return {
+            "bot_response": f"Gemini API error: {e}",
+            "status": "PARTIAL",
+            "context_found": False,
+            "ask_resolved": False,
+            "suggest_incident": True,
+            "followup": "Should I raise a ServiceNow ticket?",
+            "top_hits": [],
+            "sources": [],
+            "debug": {"intent": "gemini_http_error"},
+            "source": "ERROR",
+        }
     except Exception as e:
         print("[WMS] tool call error:", e)
         return None
 
-
 def _detect_error_families(msg: str) -> list:
     low = (msg or "").lower()
     low_norm = re.sub(r"[^\w\s]", " ", low)