Update main.py

main.py

@@ -15,6 +15,15 @@ from pydantic import BaseModel
 from dotenv import load_dotenv
 from difflib import SequenceMatcher
 
+# --- WMS API integration imports ---
+from public_api.public_api_auth import login
+from public_api.public_api_item import get_item
+from public_api.public_api_inventory import get_inventory_holds
+from public_api.public_api_location import get_location_status
+from public_api.utils import flatten_json, extract_fields
+from public_api.field_mapping import TOOL_REGISTRY, ALL_TOOLS
+
+
 # KB services
 from services.kb_creation import (
     collection,
@@ -43,6 +52,25 @@ GEMINI_URL = (
 )
 os.environ["POSTHOG_DISABLED"] = "true"
 
+# --- API WMS session cache ---
+_WMS_SESSION_DATA = None
+_WMS_WAREHOUSE_ID = None
+
+def _ensure_wms_session() -> bool:
+    """Login once to WMS and cache session_data + default warehouse."""
+    global _WMS_SESSION_DATA, _WMS_WAREHOUSE_ID
+    if _WMS_SESSION_DATA and _WMS_WAREHOUSE_ID:
+        return True
+    try:
+        session, warehouse_id, ok = login()
+        if ok:
+            _WMS_SESSION_DATA = {"wms_auth": session, "user_warehouse_id": warehouse_id}
+            _WMS_WAREHOUSE_ID = warehouse_id
+            return True
+    except Exception as e:
+        print("[WMS] login failed:", e)
+    return False
+
 # ---------------------------------------------------------------------
 # Minimal server-side cache (used to populate short description if frontend didn’t pass last_issue)
 # ---------------------------------------------------------------------
@@ -161,6 +189,136 @@ 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
+    if not GEMINI_API_KEY:
+        return None
+    if not _ensure_wms_session():
+        return None
+
+    headers = {"Content-Type": "application/json"}
+    payload = {
+        "contents": [{"parts": [{"text": user_text}]}],
+        "tools": [{"function_declarations": ALL_TOOLS}],
+    }
+
+    try:
+        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 {}
+        if "functionCall" not in part:
+            return None
+
+        fn = part["functionCall"]
+        tool_name = fn["name"]
+        cfg = TOOL_REGISTRY.get(tool_name)
+        if not cfg:
+            return None
+
+        # 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
+
+        # Call underlying function per signature
+        if tool_name in ("get_inventory_holds", "get_location_status"):
+            raw = cfg"function"
+        else:
+            target = args.get(cfg["id_param"])
+            raw = cfg"function"
+
+        # Error path -> keep your response shape
+        if 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", "tool": tool_name},
+            }
+
+        data_list = raw.get(cfg.get("response_key", "data"), [])
+        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", "tool": tool_name},
+            }
+
+        # Holds / Location → concise lines
+        if tool_name in ("get_inventory_holds", "get_location_status"):
+            rows, 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))
+
+            if tool_name == "get_inventory_holds":
+                msg = f"Found {len(rows)} hold records (Total Qty: {total_qty}).\n" + "\n".join(rows[: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")
+
+            return {
+                "bot_response": msg,
+                "status": "OK",
+                "context_found": True,
+                "ask_resolved": False,
+                "suggest_incident": False,
+                "followup": None,
+                "top_hits": [],
+                "sources": [],
+                "debug": {"intent": tool_name, "warehouse_id": wh},
+            }
+
+        # Item → short summary using mapping & requested fields
+        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"}
+            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", "warehouse_id": wh},
+        }
+
+    except Exception as e:
+        print("[WMS] tool call error:", e)
+        return None
+
 
 def _detect_error_families(msg: str) -> list:
     low = (msg or "").lower()
@@ -782,6 +940,12 @@ async def chat_with_ai(input_data: ChatInput):
                 }
             except Exception as e:
                 raise HTTPException(status_code=500, detail=safe_str(e))
+                
+        # --- Try WMS API tools before KB search ---
+        res = _try_wms_tool(input_data.user_message)
+        if res is not None:
+           return res
+
 
         # ------------------ Hybrid KB search ------------------
         kb_results = hybrid_search_knowledge_base(input_data.user_message, top_k=10, alpha=0.6, beta=0.4)