core/query_data.py

"""Validation and sanitization helpers for sample-query retrieval data."""

from __future__ import annotations

import copy
import re
import unicodedata
from collections import defaultdict
from typing import Iterable


_SPACE_RE = re.compile(r"\s+")
_TRAILING_PUNCT_RE = re.compile(r"[\s?.!,~。？！]+$")


def normalize_query_text(text: str) -> str:
    value = unicodedata.normalize("NFKC", str(text or "")).strip().lower()
    value = _TRAILING_PUNCT_RE.sub("", value)
    return _SPACE_RE.sub(" ", value)


def conflicting_query_texts(rows: Iterable[dict]) -> set[str]:
    """Return normalized texts assigned to more than one menu."""
    menu_ids_by_query: dict[str, set[str]] = defaultdict(set)
    for row in rows:
        query = normalize_query_text(row.get("query", ""))
        menu_id = str(row.get("menu_id", "")).strip()
        if query and menu_id:
            menu_ids_by_query[query].add(menu_id)
    return {
        query
        for query, menu_ids in menu_ids_by_query.items()
        if len(menu_ids) > 1
    }


def sanitize_query_rows(
    rows: Iterable[dict],
    *,
    excluded_queries: set[str] | None = None,
) -> tuple[list[dict], dict]:
    """Drop conflicting, held-out, empty, and duplicate query rows."""
    source_rows = list(rows)
    rows_by_query: dict[str, list[dict]] = defaultdict(list)
    for row in source_rows:
        normalized = normalize_query_text(row.get("query", ""))
        if normalized:
            rows_by_query[normalized].append(row)
    shared_conflicts = {
        query
        for query, query_rows in rows_by_query.items()
        if len({str(row.get("menu_id", "")).strip() for row in query_rows}) > 1
        and all(row.get("allow_shared_query", False) for row in query_rows)
    }
    conflicts = conflicting_query_texts(source_rows) - shared_conflicts
    excluded = {
        normalize_query_text(query)
        for query in (excluded_queries or set())
        if normalize_query_text(query)
    }

    output = []
    seen_pairs = set()
    stats = {
        "input": len(source_rows),
        "empty": 0,
        "conflicting": 0,
        "held_out": 0,
        "duplicate": 0,
        "output": 0,
        "conflicting_texts": len(conflicts),
        "shared_conflicting_texts": len(shared_conflicts),
    }

    for row in source_rows:
        query = normalize_query_text(row.get("query", ""))
        menu_id = str(row.get("menu_id", "")).strip()
        if not query or not menu_id:
            stats["empty"] += 1
            continue
        if query in conflicts:
            stats["conflicting"] += 1
            continue
        if query in excluded:
            stats["held_out"] += 1
            continue
        pair = (query, menu_id)
        if pair in seen_pairs:
            stats["duplicate"] += 1
            continue
        seen_pairs.add(pair)
        clean = dict(row)
        clean["query"] = str(row.get("query", "")).strip()
        clean["menu_id"] = menu_id
        output.append(clean)

    stats["output"] = len(output)
    return output, stats


def sanitize_description_queries(
    descriptions: Iterable[dict],
) -> tuple[list[dict], dict]:
    """Remove exact sample-query conflicts without modifying source objects."""
    source = list(descriptions)
    rows = []
    for desc in source:
        for query in desc.get("sample_queries", []) or []:
            rows.append({"query": query, "menu_id": desc.get("menu_id", "")})
    conflicts = conflicting_query_texts(rows)

    output = []
    removed = 0
    duplicates = 0
    for desc in source:
        clean = copy.deepcopy(desc)
        kept = []
        seen = set()
        for query in desc.get("sample_queries", []) or []:
            normalized = normalize_query_text(query)
            if not normalized or normalized in conflicts:
                removed += 1
                continue
            if normalized in seen:
                duplicates += 1
                continue
            seen.add(normalized)
            kept.append(str(query).strip())
        clean["sample_queries"] = kept
        output.append(clean)

    return output, {
        "menus": len(source),
        "input_queries": len(rows),
        "conflicting_texts": len(conflicts),
        "removed_conflicting": removed,
        "removed_duplicate": duplicates,
        "output_queries": len(rows) - removed - duplicates,
    }