agent/tools.py

"""LangChain tool definitions for EO processing."""
import json
from pathlib import Path
from typing import Any, Dict
from langchain_core.tools import tool
from agent.schema import SampleSpec


def load_sample_registry() -> Dict[str, SampleSpec]:
    """Load sample registry from config/samples_index.json."""
    registry_path = Path("config/samples_index.json")
    with open(registry_path) as f:
        samples = json.load(f)
    return {s["title"]: SampleSpec(**s) for s in samples}


@tool
def run_ndvi(sample_title: str) -> Dict[str, Any]:
    """
    Compute NDVI (Normalized Difference Vegetation Index) for a Sentinel-2 or Landsat scene.

    REQUIRES: Sample containing RED and NIR bands (Band4 and Band5).

    Args:
        sample_title: title of the sample from the registry

    Returns:
        Dictionary with raster_path, preview_png, histogram_png, metrics, extras
    """
    from eo.ndvi import run
    registry = load_sample_registry()
    sample = registry[sample_title]
    result = run(sample)
    return result.model_dump()


@tool
def run_lst_landsat(sample_title: str) -> Dict[str, Any]:
    """
    Compute Land Surface Temperature for a Landsat 8/9 scene and summarize hotspots/coldspots.

    REQUIRES: Sample providing the thermal band, QA_PIXEL mask, and metadata with Kelvin scaling constants.

    Args:
        sample_title: title of the sample from the registry

    Returns:
        Dictionary with raster_path, preview_png, histogram_png, metrics, and extras describing the coolest/hottest detected patches
    """
    from eo.lst import run
    registry = load_sample_registry()
    sample = registry[sample_title]
    result = run(sample)
    return result.model_dump()


@tool
def run_dnbr_s2pair(sample_title: str) -> Dict[str, Any]:
    """
    Compute dNBR (Differenced Normalized Burn Ratio) for a Sentinel-2 pre/post pair.

    REQUIRES: Sample with task='DNBR' containing two scenes (pre, post) with NIR and SWIR2 bands.

    Args:
        sample_title: title of the sample from the registry

    Returns:
        Dictionary with raster_path, preview_png, histogram_png, metrics, extras (including severity class counts)
    """
    from eo.dnbr import run
    registry = load_sample_registry()
    sample = registry[sample_title]
    result = run(sample)
    return result.model_dump()


# Tool list for router
TOOLS = [run_ndvi, run_lst_landsat, run_dnbr_s2pair]