agents/preprocessor.py

"""Preprocessor agent (M6): the LLM decides a column mapping; code applies it (golden rule #6).

Input per file = filename + header + first 5 rows. The LLM returns a `ColumnMapping` (canonical
field -> source column); `core.parse.apply_mapping` transforms the full table — the LLM never
rewrites data rows. With no key (or low confidence / a mapping that won't apply) we fall back to the
default clusterProfiler mapping, and only flag `needs_confirmation` when even that can't resolve the
required fields — so the deterministic core stays usable without an LLM.
"""
from __future__ import annotations

import json
from pathlib import Path
from typing import Optional, Union

import pandas as pd

from agents.llm import Provider
from core.parse import apply_mapping, default_mapping, read_table
from core.schema import ColumnMapping, GseaResult, MappingReport

FileInput = Union[str, Path]

SYSTEM = (
    "You map the columns of a GSEA results table to a canonical schema. Return ONLY a JSON object "
    "with keys: pathway, nes, padj, pval, size, leading_edge, collection, suggested_name, confidence. "
    "pathway/nes/padj/pval/size/leading_edge must each be an EXACT column name from the provided "
    "header, or null if absent. 'pathway' is the gene-set/term name; 'nes' is the signed normalized "
    "enrichment score; 'padj' is the adjusted p-value (FDR). For clusterProfiler output the "
    "leading-edge GENES are in 'core_enrichment', NOT in a column literally named 'leading_edge' "
    "(that one holds a stats string). 'confidence' is your certainty from 0 to 1."
)


def _preview(df: pd.DataFrame, name: str, n: int = 5) -> str:
    return json.dumps(
        {"filename": name, "columns": list(df.columns), "first_rows": df.head(n).to_dict(orient="records")},
        default=str,
    )[:4000]


def _parse_json(text: str) -> dict:
    try:
        return json.loads(text)
    except json.JSONDecodeError:
        i, j = text.find("{"), text.rfind("}")
        if i != -1 and j > i:
            return json.loads(text[i : j + 1])
        raise


def _llm_mapping(df: pd.DataFrame, name: str, llm: Provider) -> Optional[ColumnMapping]:
    user = f"Header and sample rows:\n{_preview(df, name)}\nReturn the JSON mapping."
    resp = llm.chat(
        [{"role": "system", "content": SYSTEM}, {"role": "user", "content": user}],
        json_mode=True,
    )
    try:
        return ColumnMapping.model_validate(_parse_json(resp.text))
    except Exception:
        return None


def preprocess_one(
    path: FileInput, llm: Optional[Provider] = None, confidence_floor: float = 0.5
) -> tuple[GseaResult, MappingReport]:
    path = Path(path)
    name = path.stem
    df = read_table(path)

    mapped = _llm_mapping(df, name, llm) if llm is not None else None
    if mapped is not None and mapped.confidence >= confidence_floor:
        mapping, collection = mapped.to_mapping(), mapped.collection
        result_name, source, confidence = (mapped.suggested_name or name), "llm", mapped.confidence
    else:
        mapping, collection = default_mapping(), None
        result_name, source = name, "default"
        confidence = mapped.confidence if mapped is not None else 1.0

    try:
        result = apply_mapping(df, mapping=mapping, name=result_name, collection=collection)
        return result, MappingReport(
            name=result.name, mapping=mapping, collection=result.collection,
            confidence=confidence, source=source, needs_confirmation=False,
        )
    except ValueError as err:
        # An LLM mapping that won't apply: retry with the deterministic default before giving up.
        if source != "default":
            try:
                result = apply_mapping(df, mapping=default_mapping(), name=name)
                return result, MappingReport(
                    name=result.name, mapping=default_mapping(), collection=result.collection,
                    confidence=confidence, source="default", needs_confirmation=False,
                    message=f"LLM mapping did not apply ({err}); used the default mapping.",
                )
            except ValueError as err2:
                err = err2
        # Even the default can't resolve required fields → surface for manual confirmation.
        return GseaResult(name=name, rows=[]), MappingReport(
            name=name, mapping=mapping, collection=None, confidence=confidence,
            source=source, needs_confirmation=True,
            message=f"{err}. Columns seen: {list(df.columns)}",
        )


def preprocess(
    files: list[FileInput], llm: Optional[Provider] = None, confidence_floor: float = 0.5
) -> list[tuple[GseaResult, MappingReport]]:
    """Map and parse each uploaded file → (GseaResult, MappingReport)."""
    return [preprocess_one(f, llm=llm, confidence_floor=confidence_floor) for f in files]