apm-wo-analysis/lambdas/classifier/handler.py

"""S3-triggered classifier Lambda: parse export -> two-axis classify -> Parquet.

Triggered on ``s3:ObjectCreated`` under the ``raw/`` prefix regardless of how the
file arrives (direct upload or local drop-folder). For each non-blank-comment row
it strips HTML, runs the two-axis classifier, derives ``is_escalation`` /
``is_action`` / ``mismatch``, writes a per-WO snapshot to ``analytics/dt=YYYY-MM-DD/``
as Parquet (registering the Glue partition), and emits a small ``summary.json``
for the slack-post Lambda to read cheaply.

The slack-post invocation is wired in Phase 4 (the function does not exist yet).
Runtime: Python 3.12, ARM64, 512 MB, 120 s. pandas/pyarrow/awswrangler are
bundled from requirements.txt; boto3 is provided by the Lambda runtime.
"""

from __future__ import annotations

import csv
import json
import os
import tempfile
from collections import Counter
from datetime import datetime, timezone
from urllib.parse import unquote_plus

import awswrangler as wr
import boto3
import openpyxl
import pandas as pd

import classify as clf

GLUE_DATABASE = "apm_wo_analysis"
GLUE_TABLE = "apm_wo_snapshots"
ANALYTICS_PREFIX = "analytics"

# Map snapshot field -> substring matched (case-insensitively) against the export
# header, tolerating minor header drift in the 13-column APM export.
COLUMN_MATCHERS = {
    "wo_number": "wo number",
    "wo_description": "wo description",
    "equipment_code": "equipment",
    "site": "organization",
    "due_date": "due date",
    "department": "department",
    "wo_status": "wo status",
    "hold_reason": "hold reason",
    "last_comment": "last comment",
    "last_comment_by": "last comment by",
    "last_comment_date": "last comment date",
    "contractor": "contractor",
    "contractor_description": "contractor description",
}

_s3 = boto3.client("s3")


def _haiku_enabled() -> bool:
    return os.environ.get("APM_HAIKU_FALLBACK", "on").strip().lower() in (
        "1",
        "on",
        "true",
        "yes",
    )


def _resolve_classifier():
    """Use the Haiku-fallback classifier when enabled, else the deterministic one."""
    if _haiku_enabled() and hasattr(clf, "classify_with_haiku"):
        return clf.classify_with_haiku
    return clf.classify


def _norm(value) -> str:
    if value is None:
        return ""
    if isinstance(value, datetime):
        return value.isoformat()
    return str(value).strip()


def _resolve_columns(header: list[str]) -> dict[str, int]:
    """Map each snapshot field to its column index in the export header."""
    lowered = [(_norm(h).lower(), i) for i, h in enumerate(header)]
    resolved: dict[str, int] = {}
    for field, needle in COLUMN_MATCHERS.items():
        # "last comment" is a prefix of "last comment by"/"date"; prefer exact-ish.
        exact = [i for h, i in lowered if h == needle]
        contains = [i for h, i in lowered if needle in h]
        match = exact or contains
        if match:
            resolved[field] = match[0]
    return resolved


def _read_rows(path: str, key: str) -> tuple[list[str], list[list]]:
    """Return (header, data_rows) from an xlsx or csv export."""
    if key.lower().endswith(".csv"):
        with open(path, newline="", encoding="utf-8-sig") as fh:
            reader = list(csv.reader(fh))
        return reader[0], reader[1:]
    wb = openpyxl.load_workbook(path, read_only=True, data_only=True)
    ws = wb.active
    rows = list(ws.iter_rows(values_only=True))
    wb.close()
    return list(rows[0]), [list(r) for r in rows[1:]]


def _build_snapshot(header: list[str], data: list[list]) -> tuple[pd.DataFrame, int]:
    """Classify each non-blank-comment row into a snapshot DataFrame."""
    cols = _resolve_columns(header)
    if "last_comment" not in cols or "wo_status" not in cols:
        raise ValueError(f"Export missing required columns; resolved={list(cols)}")

    classifier = _resolve_classifier()
    records: list[dict] = []
    blank = 0
    for row in data:

        def cell(field: str):
            idx = cols.get(field)
            return row[idx] if idx is not None and idx < len(row) else None

        raw_comment = cell("last_comment")
        comment = clf.strip_html(raw_comment)
        if not comment:
            blank += 1  # blank-comment rows excluded from the classified total
            continue

        status = _norm(cell("wo_status"))
        hold = _norm(cell("hold_reason"))
        category, mismatch = classifier(status, hold, comment)

        records.append(
            {
                "wo_number": _norm(cell("wo_number")),
                "wo_description": _norm(cell("wo_description")),
                "equipment_code": _norm(cell("equipment_code")),
                "site": _norm(cell("site")),
                "due_date": _norm(cell("due_date")),
                "department": _norm(cell("department")),
                "wo_status": status,
                "hold_reason": hold,
                "last_comment": comment,
                "last_comment_by": _norm(cell("last_comment_by")),
                "last_comment_date": _norm(cell("last_comment_date")),
                "contractor": _norm(cell("contractor")),
                "contractor_description": _norm(cell("contractor_description")),
                "category": category,
                "is_escalation": category in clf.ESCALATION_CATEGORIES,
                "is_action": category in clf.ACTION_NEEDED_CATEGORIES,
                "mismatch": mismatch or "",
            }
        )
    return pd.DataFrame.from_records(records), blank


def _build_summary(df: pd.DataFrame, dt: str, key: str, blank: int) -> dict:
    category_counts = Counter(df["category"])
    site_counts = Counter(s for s in df["site"] if s)
    mismatches = [
        {"wo_number": r.wo_number, "category": r.category, "mismatch": r.mismatch}
        for r in df.itertuples()
        if r.mismatch
    ]
    return {
        "dt": dt,
        "source_key": key,
        "classified_total": int(len(df)),
        "blank_comment_rows": int(blank),
        "category_counts": dict(category_counts),
        "escalation_total": int(df["is_escalation"].sum()),
        "third_escalation_count": int(category_counts.get("3rd Escalation", 0)),
        "action_needed": int(df["is_action"].sum()),
        "routine": int((~df["is_action"]).sum()),
        "top_sites": [{"site": s, "count": n} for s, n in site_counts.most_common(10)],
        "mismatches": mismatches,
        "generated_at": datetime.now(timezone.utc).isoformat(),
    }


def handler(event, context):
    """Classify each export dropped under raw/ into a daily Parquet snapshot."""
    record = event["Records"][0]
    bucket = record["s3"]["bucket"]["name"]
    key = unquote_plus(record["s3"]["object"]["key"])

    if not key.startswith("raw/") or not key.lower().endswith((".xlsx", ".csv")):
        print(f"Skipping non-export object: s3://{bucket}/{key}")
        return {"skipped": key}

    dt = datetime.now(timezone.utc).strftime("%Y-%m-%d")
    print(f"Classifying s3://{bucket}/{key} into dt={dt}")

    with tempfile.NamedTemporaryFile(suffix=os.path.splitext(key)[1]) as tmp:
        _s3.download_fileobj(bucket, key, tmp)
        tmp.flush()
        header, data = _read_rows(tmp.name, key)

    df, blank = _build_snapshot(header, data)
    if df.empty:
        print("No classifiable rows (all comments blank); nothing written.")
        return {"classified": 0, "blank": blank}

    df["dt"] = dt
    wr.s3.to_parquet(
        df=df,
        path=f"s3://{bucket}/{ANALYTICS_PREFIX}/",
        dataset=True,
        partition_cols=["dt"],
        mode="overwrite_partitions",
        database=GLUE_DATABASE,
        table=GLUE_TABLE,
    )

    summary = _build_summary(df, dt, key, blank)
    _s3.put_object(
        Bucket=bucket,
        Key=f"{ANALYTICS_PREFIX}/dt={dt}/summary.json",
        Body=json.dumps(summary, indent=2).encode("utf-8"),
        ContentType="application/json",
    )

    # Phase 4: async-invoke the slack-post Lambda here once it exists.
    print(
        f"Wrote {len(df)} rows, {summary['escalation_total']} escalations "
        f"({summary['third_escalation_count']} 3rd), {len(summary['mismatches'])} mismatches."
    )
    return {
        "classified": int(len(df)),
        "dt": dt,
        "summary_key": f"dt={dt}/summary.json",
    }