graph_rag/graph_builder.py

import asyncio
import logging
import pickle
import re
from pathlib import Path
from typing import List, Optional

import networkx as nx

from .external_sources import fetch_agris, fetch_agricola
from .kg_data import (
    CLIMATE_CONDITIONS,
    CROPS,
    DISEASES,
    PESTICIDES,
    PESTS,
    RELATIONSHIPS,
    SOIL_TYPES,
)

_enrich_logger = logging.getLogger(__name__ + ".enrich")


class AgroKGBuilder:
    """Build and persist the TerraMind agricultural knowledge graph."""

    def __init__(self):
        self.G = nx.MultiDiGraph()
        self.node_index = {}

    def build(self) -> nx.MultiDiGraph:
        self._add_crop_nodes()
        self._add_pest_nodes()
        self._add_disease_nodes()
        self._add_pesticide_nodes()
        self._add_soil_nodes()
        self._add_climate_nodes()
        self._add_all_edges()
        self._validate()
        return self.G

    def _add_crop_nodes(self):
        for crop_id, data in CROPS.items():
            self.G.add_node(crop_id, node_type="crop", **data)
            self.node_index[data["name_en"].lower()] = crop_id
            self.node_index[data["name_hi"]] = crop_id
            self.node_index[crop_id] = crop_id
            if "scientific_name" in data:
                self.node_index[data["scientific_name"].lower()] = crop_id

    def _add_pest_nodes(self):
        for pest_id, data in PESTS.items():
            self.G.add_node(pest_id, node_type="pest", **data)
            self.node_index[data["name_en"].lower()] = pest_id
            self.node_index[data["name_hi"]] = pest_id
            self.node_index[pest_id] = pest_id
            self.node_index[data["scientific_name"].lower()] = pest_id

    def _add_disease_nodes(self):
        for disease_id, data in DISEASES.items():
            self.G.add_node(disease_id, node_type="disease", **data)
            self.node_index[data["name_en"].lower()] = disease_id
            self.node_index[data["name_hi"]] = disease_id
            self.node_index[disease_id] = disease_id

    def _add_pesticide_nodes(self):
        for pesticide_id, data in PESTICIDES.items():
            self.G.add_node(pesticide_id, node_type="pesticide", **data)
            self.node_index[data["name_en"].lower()] = pesticide_id
            self.node_index[data["name_hi"]] = pesticide_id
            self.node_index[pesticide_id] = pesticide_id

    def _add_soil_nodes(self):
        for soil_id, data in SOIL_TYPES.items():
            self.G.add_node(soil_id, node_type="soil_type", **data)
            self.node_index[data["name_en"].lower()] = soil_id
            self.node_index[data["name_hi"]] = soil_id
            self.node_index[soil_id] = soil_id

    def _add_climate_nodes(self):
        for climate_id, data in CLIMATE_CONDITIONS.items():
            self.G.add_node(climate_id, node_type="climate", **data)
            self.node_index[data["name_en"].lower()] = climate_id
            self.node_index[climate_id] = climate_id

    def _add_all_edges(self):
        rels = RELATIONSHIPS

        for r in rels["crop_pest"]:
            self.G.add_edge(
                r["crop"],
                r["pest"],
                relation="SUSCEPTIBLE_TO",
                **{k: v for k, v in r.items() if k not in ["crop", "pest"]},
            )

        for r in rels["crop_disease"]:
            self.G.add_edge(
                r["crop"],
                r["disease"],
                relation="VULNERABLE_TO",
                **{k: v for k, v in r.items() if k not in ["crop", "disease"]},
            )

        for r in rels["disease_treatment"]:
            self.G.add_edge(
                r["disease"],
                r["pesticide"],
                relation="TREATED_BY",
                **{k: v for k, v in r.items() if k not in ["disease", "pesticide"]},
            )

        for r in rels["pest_treatment"]:
            self.G.add_edge(
                r["pest"],
                r["pesticide"],
                relation="CONTROLLED_BY",
                **{k: v for k, v in r.items() if k not in ["pest", "pesticide"]},
            )

        for r in rels["pesticide_soil_conflict"]:
            self.G.add_edge(
                r["pesticide"],
                r["soil"],
                relation="CONFLICTS_WITH",
                **{k: v for k, v in r.items() if k not in ["pesticide", "soil"]},
            )

        for r in rels["pest_climate"]:
            self.G.add_edge(
                r["pest"],
                r["climate"],
                relation="PEAKS_DURING",
                **{k: v for k, v in r.items() if k not in ["pest", "climate"]},
            )

        for r in rels["tank_mix_conflicts"]:
            for a, b in [(r["pesticide_a"], r["pesticide_b"]), (r["pesticide_b"], r["pesticide_a"])]:
                self.G.add_edge(
                    a,
                    b,
                    relation="INCOMPATIBLE_WITH",
                    reason=r["reason"],
                    severity=r["severity"],
                )

        for r in rels["disease_climate"]:
            self.G.add_edge(
                r["disease"],
                r["climate"],
                relation="FAVORED_BY",
                **{k: v for k, v in r.items() if k not in ["disease", "climate"]},
            )

    def _validate(self):
        import logging

        log = logging.getLogger("AgroKGBuilder")

        for node, data in self.G.nodes(data=True):
            if data.get("node_type") == "disease":
                treatments = [
                    v
                    for _, v, d in self.G.out_edges(node, data=True)
                    if d.get("relation") == "TREATED_BY"
                ]
                if not treatments:
                    log.warning("Disease with no treatment: %s", node)

        log.info(
            "KG built: %s nodes, %s edges",
            self.G.number_of_nodes(),
            self.G.number_of_edges(),
        )

    def resolve_node(self, name: str) -> Optional[str]:
        if not name:
            return None
        if name in self.node_index:
            return self.node_index[name]
        if name.lower() in self.node_index:
            return self.node_index[name.lower()]

        keys = list(self.node_index.keys())

        try:
            from rapidfuzz import fuzz, process

            match = process.extractOne(
                name.lower(),
                keys,
                scorer=fuzz.WRatio,
                score_cutoff=80,
            )
            if match:
                return self.node_index[match[0]]
        except Exception:
            from difflib import SequenceMatcher

            best_key = None
            best_score = 0.0
            target = name.lower()
            for key in keys:
                score = SequenceMatcher(None, target, key).ratio()
                if score > best_score:
                    best_score = score
                    best_key = key

            if best_key and best_score >= 0.8:
                return self.node_index[best_key]
        return None

    def save(self, path: str = "graph_rag/agrokg.pkl"):
        Path(path).parent.mkdir(parents=True, exist_ok=True)
        with open(path, "wb") as f:
            pickle.dump({"graph": self.G, "node_index": self.node_index}, f)

    @classmethod
    def load(cls, path: str = "graph_rag/agrokg.pkl"):
        obj = cls()
        try:
            with open(path, "rb") as f:
                data = pickle.load(f)
            obj.G = data["graph"]
            obj.node_index = data["node_index"]
        except FileNotFoundError:
            obj.build()
            obj.save(path)
        return obj


# ── Module-level enrichment function ─────────────────────────────────────────

# Shared builder instance used by the enrichment function.
_shared_builder: Optional[AgroKGBuilder] = None


def _get_shared_builder() -> AgroKGBuilder:
    """Lazily initialize and return the shared AgroKGBuilder instance."""
    global _shared_builder
    if _shared_builder is None:
        _shared_builder = AgroKGBuilder.load()
    return _shared_builder


async def enrich_graph_for_disease(crop: str, disease: str) -> List[str]:
    """
    Enrich the knowledge graph with external research data for a given
    crop-disease pair.  Calls both AGRIS and AGRICOLA fetchers, then
    upserts the returned document strings as new nodes in the graph.

    Parameters
    ----------
    crop : str
        The crop name (e.g. ``"Tomato"``).
    disease : str
        The disease label (e.g. ``"Late_blight"``).

    Returns
    -------
    list[str]
        Combined list of all retrieved text strings (may be empty).
    """
    _enrich_logger.info("Enriching graph for crop=%s, disease=%s", crop, disease)

    agris_results, agricola_results = await asyncio.gather(
        fetch_agris(crop, disease),
        fetch_agricola(crop, disease),
    )

    all_texts: List[str] = agris_results + agricola_results

    if not all_texts:
        _enrich_logger.warning(
            "Both AGRIS and AGRICOLA returned empty results for "
            "crop='%s', disease='%s'. Report will rely solely on "
            "local Knowledge Graph data.",
            crop, disease,
        )
        return all_texts

    # Upsert enrichment nodes into the shared graph
    builder = _get_shared_builder()
    disease_clean = disease.replace("___", " ").replace("__", " ").replace("_", " ")

    for idx, text in enumerate(all_texts):
        source = "agris" if idx < len(agris_results) else "agricola"
        node_id = re.sub(r"[^a-zA-Z0-9]", "_", f"enrich_{crop}_{disease}_{source}_{idx}").lower()

        if node_id not in builder.G.nodes:
            builder.G.add_node(
                node_id,
                node_type="enrichment",
                source=source,
                crop=crop,
                disease=disease_clean,
                text=text[:500],
                name_en=f"{crop} {disease_clean} ({source} #{idx + 1})",
            )
            builder.node_index[node_id] = node_id
            _enrich_logger.debug("Added enrichment node: %s", node_id)

    _enrich_logger.info(
        "Graph enrichment complete: %d new texts added (%d AGRIS, %d AGRICOLA)",
        len(all_texts), len(agris_results), len(agricola_results),
    )
    return all_texts