agents/data_extraction.py

"""
Data extraction agent: builds the graph node that extracts metrics data
based on user queries and filters.
"""
from __future__ import annotations

from typing import Callable

import pandas as pd

from tools import data_extractor, filter_extractor
from tools.state import GraphState

_DEFAULT_FILTER_COLUMNS = [
    "Region", "Period", "Cluster", "Product", "Calculation_Type", "TA Market", "Class", "Level",
]


def _tool_call(tool_or_fn, **kwargs):
    """Invoke a tool or callable, handling both raw functions and tool wrappers."""
    callable_obj = getattr(tool_or_fn, "func", tool_or_fn)
    return callable_obj(**kwargs)


def _candidate_values(df: pd.DataFrame, max_values: int = 100) -> dict[str, list[str]]:
    values: dict[str, list[str]] = {}
    for column in df.columns:
        unique_values = df[column].dropna().astype(str).str.strip().unique().tolist()
        values[column] = sorted([v for v in unique_values if v])[:max_values]
    return values


def _validate_filters(filters: dict, columns: list[str]) -> dict:
    allowed = set(columns)
    validated = {}
    for key, value in (filters or {}).items():
        if key not in allowed:
            continue
        if isinstance(value, list):
            cleaned = [str(v).strip() for v in value if str(v).strip()]
            if cleaned:
                validated[key] = cleaned
    return validated


def build_data_extraction_node(
    metrics_df: pd.DataFrame | None = None,
    llm_invoke: Callable[[str], object] | None = None,
    db_query_fn: Callable[[dict], pd.DataFrame] | None = None,
    column_values: dict[str, list[str]] | None = None,
):
    """
    Factory that returns a data extraction node for the LangGraph.

    SQL mode (preferred):
        Pass ``db_query_fn`` and ``column_values``. Filters are extracted from the
        user message, then a targeted SQL query is executed.

    In-memory mode (legacy):
        Pass ``metrics_df`` (the full metrics DataFrame). Filters are applied in
        memory via pandas.
    """
    if db_query_fn is not None:
        available_columns = list(column_values.keys()) if column_values else _DEFAULT_FILTER_COLUMNS
        value_map: dict[str, list[str]] = column_values or {}
    else:
        if metrics_df is None:
            raise ValueError("Either metrics_df or db_query_fn must be provided.")
        available_columns = metrics_df.columns.tolist()
        value_map = _candidate_values(metrics_df)

    def data_extraction_node(state: GraphState) -> dict:
        user_query = state.get("user_query", "")
        prior_filters = _validate_filters(state.get("filters", {}), available_columns)
        conversation_history = state.get("conversation_history", []) or []

        raw_filters = _tool_call(
            filter_extractor,
            user_query=user_query,
            available_columns=available_columns,
            column_values=value_map,
            llm=llm_invoke,
            prior_filters=prior_filters,
            conversation_history=conversation_history,
        )

        filters = _validate_filters(raw_filters, available_columns)

        if db_query_fn is not None:
            fetched_df = db_query_fn(filters)
            rows = fetched_df.to_dict(orient="records") if not fetched_df.empty else []
        else:
            rows = _tool_call(data_extractor, filters=filters, dataframe=metrics_df)

        error_message = state.get("error_message", "")
        if not rows:
            error_message = "I couldn't find any data matching your criteria."

        return {
            "filters": filters,
            "extracted_data": rows,
            "error_message": error_message,
        }

    return data_extraction_node