query.py

"""
Query pre-processing: normalization, spelling corrections, alias expansion,
product/source hints, and intent detection.

Product aliases and spelling fixes are loaded from knowledge/aliases.yaml so
that adding a new product requires only editing the YAML file.
"""
import re
import logging
from pathlib import Path

logger = logging.getLogger(__name__)

# ---------------------------------------------------------------------------
# Alias configuration loader
# ---------------------------------------------------------------------------

_ALIASES_PATH = Path(__file__).resolve().parent.parent / "knowledge" / "aliases.yaml"
_aliases_config: dict | None = None


def _load_aliases() -> dict:
    global _aliases_config
    if _aliases_config is not None:
        return _aliases_config
    try:
        import yaml
        with open(_ALIASES_PATH, "r", encoding="utf-8") as f:
            _aliases_config = yaml.safe_load(f) or {}
    except ImportError:
        logger.warning("PyYAML not installed; falling back to built-in alias config.")
        _aliases_config = _builtin_aliases()
    except FileNotFoundError:
        logger.warning("aliases.yaml not found; using built-in aliases.")
        _aliases_config = _builtin_aliases()
    except Exception as e:
        logger.warning("Failed to load aliases.yaml (%s); using built-in aliases.", e)
        _aliases_config = _builtin_aliases()
    return _aliases_config


def _builtin_aliases() -> dict:
    """Fallback hardcoded config matching the original main.py logic."""
    return {
        "spelling_fixes": {
            "handysrping": "handyspring",
            "key chain": "keychain",
            "blu tooth": "bluetooth",
        },
        "products": {
            "handyspring": {
                "canonical": "handyspring",
                "aliases": ["handyspring", "drink water reminder", "handy spring"],
                "source_hints": ["handyspring", "drink water reminder"],
            },
            "coolfire": {
                "canonical": "coolfire",
                "aliases": ["coolfire", "vibration watch", "vibrating watch"],
                "source_hints": ["coolfire"],
            },
            "keychain": {
                "canonical": "keychain",
                "aliases": ["keychain", "key chain", "talking keychain", "talking key chain"],
                "source_hints": [
                    "radio control talking key chain",
                    "digital talking keychain",
                    "key chain",
                ],
            },
        },
        "compound_triggers": {
            "five_senses_bt_watch": {
                "require_all": ["bluetooth", "talking", "watch"],
                "source_hints": [
                    "five senses bluetooth talking watch",
                    "five senses bluetooth talking",
                ],
            }
        },
    }


# ---------------------------------------------------------------------------
# Core normalization helpers
# ---------------------------------------------------------------------------

def normalize_for_match(text: str) -> str:
    """Lower-case, strip punctuation, collapse whitespace."""
    lowered = (text or "").lower()
    lowered = re.sub(r"[^a-z0-9]+", " ", lowered)
    return " ".join(lowered.split())


# ---------------------------------------------------------------------------
# Intent detection
# ---------------------------------------------------------------------------

def is_procedural_question(question: str) -> bool:
    """Return True when the query is asking how to do something."""
    q = (question or "").lower()
    hints = ["how to", "set", "setup", "set up", "step", "time", "alarm", "connect", "pair"]
    return any(h in q for h in hints)


def is_image_intent(question: str) -> bool:
    """Return True when the user appears to want an image shown."""
    q = (question or "").lower()
    hard_keywords = ["image", "photo", "picture", "圖片", "相片", "img", "jpg", "png"]
    if any(k in q for k in hard_keywords):
        return True
    soft_triggers = ["show me", "show", "睇"]
    if any(k in q for k in soft_triggers) and not is_procedural_question(q):
        return True
    return False


# ---------------------------------------------------------------------------
# Source / product hint inference
# ---------------------------------------------------------------------------

def infer_source_hints(question: str) -> list[str]:
    """
    Return a list of substring hints that should appear in the source filename
    when a specific product is mentioned in the query.
    """
    cfg = _load_aliases()
    q = normalize_for_match(question)
    q_compact = q.replace(" ", "")
    hints: list[str] = []

    for _product_key, product in cfg.get("products", {}).items():
        canonical = product.get("canonical", "")
        aliases = [normalize_for_match(a) for a in product.get("aliases", [])]
        match = (
            canonical in q
            or canonical in q_compact
            or any(a in q for a in aliases)
        )
        if match:
            hints.extend(product.get("source_hints", []))

    for _trigger_key, trigger in cfg.get("compound_triggers", {}).items():
        if all(token in q for token in trigger.get("require_all", [])):
            hints.extend(trigger.get("source_hints", []))

    return hints


def source_matches_hints(source: str, hints: list[str]) -> bool:
    s = normalize_for_match(source)
    if not hints:
        return False
    return any(h in s for h in hints)


def path_matches_hints(image_path: str, hints: list[str]) -> bool:
    import os
    if not hints:
        return False
    normalized = normalize_for_match(os.path.basename(image_path))
    compact = normalized.replace(" ", "")
    for hint in hints:
        h = normalize_for_match(hint)
        if h in normalized:
            return True
        if h.replace(" ", "") in compact:
            return True
    return False


# ---------------------------------------------------------------------------
# Query rewriting (alias expansion + spelling fixes)
# ---------------------------------------------------------------------------

def rewrite_query_for_retrieval(question: str) -> str:
    """
    Normalise spelling, then expand recognised product terms into aliases
    so both vector and BM25 searches can match the right documents.
    """
    original = (question or "").strip()
    if not original:
        return original

    cfg = _load_aliases()
    q = normalize_for_match(original)

    # Apply global spelling fixes
    for wrong, right in cfg.get("spelling_fixes", {}).items():
        q = q.replace(normalize_for_match(wrong), normalize_for_match(right))

    # Apply per-product spelling fixes
    for _product_key, product in cfg.get("products", {}).items():
        for wrong, right in product.get("spelling_fixes", {}).items():
            q = q.replace(normalize_for_match(wrong), normalize_for_match(right))

    # Expand recognised product terms into all aliases
    expanded_terms = [q]
    q_compact = q.replace(" ", "")
    for _product_key, product in cfg.get("products", {}).items():
        canonical = product.get("canonical", "")
        aliases = [normalize_for_match(a) for a in product.get("aliases", [])]
        hit = canonical in q or canonical in q_compact or any(a in q for a in aliases)
        if hit:
            expanded_terms.extend(aliases)

    return " ".join(dict.fromkeys(" ".join(expanded_terms).split()))