src/inventory_db.py

"""
inventory_db.py
----------------
Very small "NL -> structured query" layer over the synthetic inventory and
orders tables. It uses the intent classifier's output plus simple regex
slot extraction (SKU codes, order IDs, zone names) to filter the
in-memory DataFrames -- a lightweight stand-in for the kind of
WMS/WCS query interface a production assistant would call as a tool.
"""

import re

import pandas as pd

SKU_RE = re.compile(r"SKU-\d{3,4}", re.IGNORECASE)
ORDER_RE = re.compile(r"#\d{4,6}")
ZONE_RE = re.compile(r"zone [a-d]", re.IGNORECASE)


def extract_sku(text: str):
    m = SKU_RE.search(text)
    return m.group(0).upper() if m else None


def extract_order_id(text: str):
    m = ORDER_RE.search(text)
    return m.group(0) if m else None


def extract_zone(text: str):
    m = ZONE_RE.search(text)
    return m.group(0).title() if m else None


def query_inventory(inventory_df: pd.DataFrame, text: str) -> pd.DataFrame:
    sku = extract_sku(text)
    zone = extract_zone(text)
    df = inventory_df.copy()
    if sku:
        df = df[df["sku"].str.upper() == sku]
    if zone:
        df = df[df["zone"].str.lower() == zone.lower()]
    if df.empty and not sku and not zone:
        # no specific filters recognised -> show low-stock items as a useful default
        df = inventory_df[inventory_df["on_hand_units"] <= inventory_df["reorder_point"]]
    return df.reset_index(drop=True)


def query_orders(orders_df: pd.DataFrame, text: str) -> pd.DataFrame:
    order_id = extract_order_id(text)
    zone = extract_zone(text)
    df = orders_df.copy()
    if order_id:
        df = df[df["order_id"] == order_id]
    elif zone:
        df = df[df["zone"].str.lower() == zone.lower()]
    elif "delay" in text.lower():
        df = df[df["status"] == "Delayed"]
    return df.reset_index(drop=True)