tools/data_tools.py

from __future__ import annotations

from typing import Any

import pandas as pd

from langchain_core.tools import tool


def _normalize(values: list[Any]) -> set[str]:
    return {str(v).strip().lower() for v in values if str(v).strip()}


@tool
def data_extractor(filters: dict, dataframe: pd.DataFrame) -> list[dict]:
    """Extract dataframe rows using validated filters."""
    if dataframe is None or dataframe.empty:
        return []

    result = dataframe.copy()
    valid_columns = set(result.columns)

    for key, values in (filters or {}).items():
        if key not in valid_columns:
            continue
        if not isinstance(values, list) or not values:
            continue
        wanted = _normalize(values)
        series = result[key].fillna("").astype(str).str.strip().str.lower()
        result = result[series.isin(wanted)]
        if result.empty:
            return []

    return result.to_dict(orient="records")