Update main.py

main.py

@@ -218,26 +218,229 @@ 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.
+    WMS tool orchestrator:
+    1) Fast-path: regex-detects item/location/holds queries and calls WMS immediately.
+    2) Fallback: Gemini function-calling (functionDeclarations) if fast-path doesn't trigger.
+    Returns your existing bot response dict shape; None when no tool applies.
     """
-    # Guards
-    if not GEMINI_API_KEY:
-        return None
+    # --- Guards / session ---
     if not _ensure_wms_session():
         return None
 
-    # Prepare default warehouse from cached session
     session_data = _WMS_SESSION_DATA
     default_wh = session_data.get("user_warehouse_id") if session_data else None
 
+    # ------------------------------------------------------------------
+    # 1) FAST-PATH (no Gemini needed): item / location / holds
+    # ------------------------------------------------------------------
+    import re
+    text = (user_text or "").strip()
+
+    # a) ITEM: "Get item number 100001 ..." / "Show item 100001 ..."
+    m_item = re.search(r"\bitem\s+(?:number|id)?\s*[: ]?([A-Za-z0-9\-]+)", text, flags=re.IGNORECASE)
+    if m_item:
+        item_num = m_item.group(1)
+
+        # Detect requested field labels (matches labels in your ITEM_MAPPING)
+        try:
+            item_labels = set(TOOL_REGISTRY["get_item_data"]["mapping"].keys())
+        except Exception:
+            item_labels = set()
+        requested = [lbl for lbl in item_labels if re.search(rf"\b{re.escape(lbl)}\b", text, flags=re.IGNORECASE)]
+
+        try:
+            raw = get_item(session_data, default_wh, item_num)  # <-- your real function
+            if isinstance(raw, dict) and raw.get("error"):
+                return {
+                    "bot_response": f"⚠️ {raw['error']}",
+                    "status": "PARTIAL",
+                    "context_found": False,
+                    "ask_resolved": False,
+                    "suggest_incident": True,
+                    "followup": "Should I raise a ServiceNow ticket?",
+                    "top_hits": [],
+                    "sources": [],
+                    "debug": {"intent": "wms_error_fastpath", "tool": "get_item_data"},
+                    "source": "ERROR",
+                }
+
+            data_list = raw.get("data", []) or []
+            if not data_list:
+                return {
+                    "bot_response": "I searched WMS but couldn't find matching records.",
+                    "status": "PARTIAL",
+                    "context_found": False,
+                    "ask_resolved": False,
+                    "suggest_incident": True,
+                    "followup": "Do you want me to raise a ticket or try a different ID?",
+                    "top_hits": [],
+                    "sources": [],
+                    "debug": {"intent": "wms_empty_fastpath", "tool": "get_item_data"},
+                    "source": "LIVE_API",
+                }
+
+            # Build summary table
+            cfg = TOOL_REGISTRY["get_item_data"]
+            flat = flatten_json(data_list[0])
+            if requested:
+                filtered_map = {k: v for k, v in cfg["mapping"].items() if k in requested or k == "Item"}
+                order = ["Item"] + [f for f in requested if f != "Item"]
+            else:
+                summary_keys = ["Item", "Description", "Item type", "Warehouse ID"]
+                filtered_map = {k: v for k, v in cfg["mapping"].items() if k in summary_keys}
+                order = summary_keys
+
+            cleaned = extract_fields(flat, filtered_map, cfg.get("formatters", {}), requested_order=order)
+            lines = [f"{k}: {cleaned.get(k, 'N/A')}" for k in order]
+            msg = "Details:\n" + "\n".join(lines)
+
+            return {
+                "bot_response": msg,
+                "status": "OK",
+                "context_found": True,
+                "ask_resolved": False,
+                "suggest_incident": False,
+                "followup": None,
+                "top_hits": [],
+                "sources": [],
+                "debug": {"intent": "get_item_data_fastpath", "warehouse_id": default_wh},
+                "source": "LIVE_API",
+                "type": "table",
+                "data": [data_list[0]],
+                "show_export": False,
+            }
+        except Exception as e:
+            print("[WMS fastpath:item] error:", e)
+            # fall through to Gemini
+
+    # b) LOCATION: "Check location A1-01-01" / "Status of bin A1-01-01"
+    m_loc = re.search(r"\b(?:check|status|loc(?:ation)?)\s+(?:bin\s+)?([A-Za-z0-9\-]+)", text, flags=re.IGNORECASE)
+    if m_loc:
+        stoloc = m_loc.group(1)
+        try:
+            raw = get_location_status(session_data, default_wh, stoloc)
+            if isinstance(raw, dict) and raw.get("error"):
+                return {
+                    "bot_response": f"⚠️ {raw['error']}",
+                    "status": "PARTIAL",
+                    "context_found": False,
+                    "ask_resolved": False,
+                    "suggest_incident": True,
+                    "followup": "Should I raise a ServiceNow ticket?",
+                    "top_hits": [],
+                    "sources": [],
+                    "debug": {"intent": "wms_error_fastpath", "tool": "get_location_status"},
+                    "source": "ERROR",
+                }
+
+            data_list = raw.get("data", []) or []
+            cfg = TOOL_REGISTRY["get_location_status"]
+            rows_for_text = []
+            for item in data_list:
+                flat = flatten_json(item)
+                row = extract_fields(flat, cfg["mapping"], cfg.get("formatters", {}))
+                rows_for_text.append("• " + "; ".join(f"{k}: {v}" for k, v in row.items()))
+            msg = f"Location '{stoloc}' status:\n" + ("\n".join(rows_for_text[:1]) if rows_for_text else "No status returned")
+
+            return {
+                "bot_response": msg,
+                "status": "OK",
+                "context_found": True,
+                "ask_resolved": False,
+                "suggest_incident": False,
+                "followup": None,
+                "top_hits": [],
+                "sources": [],
+                "debug": {"intent": "get_location_status_fastpath", "warehouse_id": default_wh},
+                "source": "LIVE_API",
+                "type": "table",
+                "data": data_list,
+                "show_export": False,
+            }
+        except Exception as e:
+            print("[WMS fastpath:location] error:", e)
+            # fall through
+
+    # c) HOLDS: "Show holds for LPN 123456" / "Show all inventory holds"
+    m_lpn = re.search(r"\b(?:lpn|lodnum)\s*[: ]?([A-Za-z0-9\-]+)", text, flags=re.IGNORECASE)
+    want_all_holds = re.search(r"\b(all\s+inventory\s+holds|list\s+holds|show\s+holds)\b", text, flags=re.IGNORECASE)
+    if m_lpn or want_all_holds:
+        kwargs = {}
+        if m_lpn:
+            kwargs["lodnum"] = m_lpn.group(1)
+        try:
+            raw = get_inventory_holds(session_data, default_wh, **kwargs)
+            if isinstance(raw, dict) and raw.get("error"):
+                return {
+                    "bot_response": f"⚠️ {raw['error']}",
+                    "status": "PARTIAL",
+                    "context_found": False,
+                    "ask_resolved": False,
+                    "suggest_incident": True,
+                    "followup": "Should I raise a ServiceNow ticket?",
+                    "top_hits": [],
+                    "sources": [],
+                    "debug": {"intent": "wms_error_fastpath", "tool": "get_inventory_holds"},
+                    "source": "ERROR",
+                }
+
+            data_list = raw.get("data", []) or []
+            if not data_list:
+                return {
+                    "bot_response": "I searched WMS but couldn't find matching records.",
+                    "status": "PARTIAL",
+                    "context_found": False,
+                    "ask_resolved": False,
+                    "suggest_incident": True,
+                    "followup": "Do you want me to raise a ticket or try a different ID?",
+                    "top_hits": [],
+                    "sources": [],
+                    "debug": {"intent": "wms_empty_fastpath", "tool": "get_inventory_holds"},
+                    "source": "LIVE_API",
+                }
+
+            cfg = TOOL_REGISTRY["get_inventory_holds"]
+            rows_for_text = []
+            def _to_int(x):
+                try: return int(x or 0)
+                except Exception: return 0
+            total_qty = 0
+            for item in data_list:
+                flat = flatten_json(item)
+                row = extract_fields(flat, cfg["mapping"], cfg.get("formatters", {}))
+                rows_for_text.append("• " + "; ".join(f"{k}: {v}" for k, v in row.items()))
+                total_qty += _to_int(item.get("untqty", 0))
+
+            msg = f"Found {len(rows_for_text)} hold records (Total Qty: {total_qty}).\n" + "\n".join(rows_for_text[:10])
+
+            return {
+                "bot_response": msg,
+                "status": "OK",
+                "context_found": True,
+                "ask_resolved": False,
+                "suggest_incident": False,
+                "followup": None,
+                "top_hits": [],
+                "sources": [],
+                "debug": {"intent": "get_inventory_holds_fastpath", "warehouse_id": default_wh},
+                "source": "LIVE_API",
+                "type": "table",
+                "data": data_list,
+                "show_export": True,
+            }
+        except Exception as e:
+            print("[WMS fastpath:holds] error:", e)
+            # fall through
+
+    # ------------------------------------------------------------------
+    # 2) Gemini function-calling fallback (if fast-path didn't match)
+    # ------------------------------------------------------------------
     payload = {
         "contents": [{"role": "user", "parts": [{"text": user_text}]}],
-        "tools": [{"functionDeclarations": ALL_TOOLS}],  # from public_api.field_mapping
-        "toolConfig": {"functionCallingConfig": {"mode": "ANY"}},  # nudge function calling
+        "tools": [{"functionDeclarations": ALL_TOOLS}],
+        "toolConfig": {"functionCallingConfig": {"mode": "ANY"}},
     }
 
     try:
@@ -263,7 +466,6 @@ def _try_wms_tool(user_text: str) -> dict | None:
         if not cfg:
             return None
 
-        # Resolve warehouse override, else use session default
         wh = (
             args.pop("warehouse_id", None)
             or args.pop("warehouseId", None)
@@ -271,7 +473,6 @@ def _try_wms_tool(user_text: str) -> dict | None:
             or default_wh
         )
 
-        # Get callable
         tool_fn = cfg.get("function")
         if not callable(tool_fn):
             return {
@@ -287,12 +488,9 @@ def _try_wms_tool(user_text: str) -> dict | None:
                 "source": "ERROR",
             }
 
-        # ---- Correct invocation according to your real function signatures ----
-        # get_inventory_holds(session_data, warehouse_id, **kwargs)
-        # get_location_status(session_data, warehouse_id, stoloc)
-        # get_item(session_data, warehouse_id, item_number, **kwargs)
+        # ---- Correct invocation for your real signatures ----
         if tool_name == "get_inventory_holds":
-            call_kwargs = {**args}  # lodnum / subnum / dtlnum
+            call_kwargs = {**args}
             raw = tool_fn(session_data, wh, **call_kwargs)
 
         elif tool_name == "get_location_status":
@@ -318,7 +516,7 @@ def _try_wms_tool(user_text: str) -> dict | None:
             extra_kwargs = {k: v for k, v in args.items() if k != id_param}
             raw = tool_fn(session_data, wh, target, **extra_kwargs)
 
-        # ---- Tool-layer error handling ----
+        # ---- Common handling ----
         if isinstance(raw, dict) and raw.get("error"):
             return {
                 "bot_response": f"⚠️ {raw['error']}",
@@ -333,7 +531,6 @@ def _try_wms_tool(user_text: str) -> dict | None:
                 "source": "ERROR",
             }
 
-        # ---- Extract rows from response ----
         response_key = cfg.get("response_key", "data")
         data_list = (raw.get(response_key, []) if isinstance(raw, dict) else []) or []
         if not data_list:
@@ -350,7 +547,6 @@ def _try_wms_tool(user_text: str) -> dict | None:
                 "source": "LIVE_API",
             }
 
-        # ---- Table rendering (holds + location) ----
         if tool_name in ("get_inventory_holds", "get_location_status"):
             rows_for_text, total_qty = [], 0
             for item in data_list:
@@ -360,10 +556,8 @@ def _try_wms_tool(user_text: str) -> dict | None:
 
             if tool_name == "get_inventory_holds":
                 def _to_int(x):
-                    try:
-                        return int(x or 0)
-                    except Exception:
-                        return 0
+                    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:
@@ -386,7 +580,7 @@ def _try_wms_tool(user_text: str) -> dict | None:
                 "show_export": True if tool_name == "get_inventory_holds" else False,
             }
 
-        # ---- Item summary + one-row table ----
+        # Item summary + one-row table
         flat = flatten_json(data_list[0])
         requested = args.get("fields", [])
         if requested:
@@ -434,7 +628,6 @@ def _try_wms_tool(user_text: str) -> dict | None:
         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)