lexsi_ds/agent/reference_resolver.py

"""Reference resolver — ground anaphora / positional / predicate refs
against session state before the planner sees them.

Used by the multi-turn loop after `classify_turn` to convert phrases
like "show me 20 instead", "the second one", "those", "loans over
30000" into explicit args the planner can pass straight to a tool.
"""

from __future__ import annotations

import logging
import re
from dataclasses import dataclass, field
from typing import TYPE_CHECKING, Any, Literal

import pandas as pd

if TYPE_CHECKING:
    from lexsi_ds.agent.session import SessionContext

log = logging.getLogger(__name__)


ResolutionKind = Literal["slot", "positional", "predicate", "drilldown", "none"]


@dataclass
class ResolvedReference:
    kind: ResolutionKind = "none"
    source_slot: str = ""
    resolved_args: dict[str, Any] = field(default_factory=dict)
    confidence: float = 0.0
    notes: str = ""


_DRILLDOWN_RE = re.compile(
    r"\bwhy is (?:loan |row |id |entity |case )?#?(\d+)\b", re.IGNORECASE,
)
_TOP_K_RE = re.compile(r"\btop\s+(\d+)\b", re.IGNORECASE)
_POSITIONAL_RE = re.compile(
    r"\bthe\s+(first|second|third|fourth|fifth|last)\b", re.IGNORECASE,
)
_PREDICATE_RE = re.compile(
    r"\b(higher than|larger than|more than|greater than|over|above|>)\s+"
    r"\$?(\d+(?:[.,]\d{3})*(?:\.\d+)?)\b",
    re.IGNORECASE,
)
_THOSE_RE = re.compile(r"\b(those|them|they|these)\b", re.IGNORECASE)

_POSITION_MAP = {
    "first": 0, "second": 1, "third": 2,
    "fourth": 3, "fifth": 4, "last": -1,
}


def resolve(user_msg: str, session: "SessionContext") -> ResolvedReference:
    """Find the highest-priority grounded reference in `user_msg`."""
    msg = user_msg or ""

    # 1. Drilldown — "why is loan 5314 here?"
    m = _DRILLDOWN_RE.search(msg)
    if m:
        return ResolvedReference(
            kind="drilldown",
            source_slot="last_predictions",
            resolved_args={"row_id": int(m.group(1))},
            confidence=0.85,
            notes="drilldown on numeric id",
        )

    # 2. Top-K override — "top 20"
    m = _TOP_K_RE.search(msg)
    if m:
        return ResolvedReference(
            kind="positional",
            resolved_args={"top_k": int(m.group(1))},
            confidence=0.9,
            notes="top-k override",
        )

    # 3. Positional — "the second one", "the last entry"
    m = _POSITIONAL_RE.search(msg)
    if m:
        return ResolvedReference(
            kind="positional",
            resolved_args={"index": _POSITION_MAP[m.group(1).lower()]},
            confidence=0.8,
            notes=f"positional: {m.group(1)}",
        )

    # 4. Predicate filter — "over 30000"
    m = _PREDICATE_RE.search(msg)
    if m:
        value = m.group(2).replace(",", "")
        cols = _columns_in_last_table(session)
        column_hint = cols[0] if cols else "amount"
        return ResolvedReference(
            kind="predicate",
            source_slot="last_predictions" if session.cache.get("last_predictions") is not None
            else "last_sql_result",
            resolved_args={"filter": f"> {value}", "column_hint": column_hint},
            confidence=0.7,
            notes=f"predicate over {value}",
        )

    # 5. Anaphora — "those", "them" → previous predictions if present
    if _THOSE_RE.search(msg):
        if session.cache.get("last_predictions") is not None:
            return ResolvedReference(
                kind="slot",
                source_slot="last_predictions",
                resolved_args={},
                confidence=0.6,
                notes="anaphora to last predictions",
            )

    return ResolvedReference(kind="none", confidence=1.0, notes="no reference detected")


def _columns_in_last_table(session: "SessionContext") -> list[str]:
    """Return column names from the freshest cached table, if any."""
    df = session.cache.get("last_predictions")
    if isinstance(df, pd.DataFrame) and not df.empty:
        return list(df.columns)
    df = session.cache.get("last_sql_result")
    if isinstance(df, pd.DataFrame) and not df.empty:
        return list(df.columns)
    return []