cascade_risk / scripts /02_03_quick_cascade.py
Lucasoppem's picture
Sync from GitHub main (part 2)
36f9d47 verified
Raw
History Blame Contribute Delete
4.77 kB
"""Quick pipeline: Generate cascade chains directly from event metadata using LLM.
Skips GDELT/scraping (slow, rate-limited) and uses LLM + expert knowledge
to generate cascade chains from event metadata alone. GDELT-based pipeline
(02_collect_news.py + 03_extract_cascades.py) can be used later to augment
with real news data.
"""
import json
import logging
from datetime import date
from pathlib import Path
from src.data.cascade_extractor import save_cascade_chain, _extract_json, _infer_severity
from src.llm import create_llm_client
from src.llm.client import load_config
from src.models.schemas import CascadeChain, CascadeNode, FloodEvent
logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
logger = logging.getLogger(__name__)
def generate_cascade_from_metadata(
event: FloodEvent, llm_client, config: dict
) -> CascadeChain | None:
"""Generate a cascade chain using LLM + expert knowledge from event metadata."""
variables = {
"event_id": event.event_id,
"country": event.country,
"iso": event.iso,
"location": event.location or event.country,
"start_date": str(event.start_date),
"origin": event.origin or "Unknown",
"total_deaths": event.total_deaths or "Unknown",
"total_affected": event.total_affected or "Unknown",
"total_damage_k_usd": event.total_damage_k_usd or "Unknown",
}
response = llm_client.call_with_config(
prompt_key="generate_cascade",
knowledge_key="extraction",
variables=variables,
config=config,
)
try:
json_str = _extract_json(response)
data = json.loads(json_str)
nodes = []
for nd in data.get("cascade_events", []):
nodes.append(CascadeNode(
id=nd.get("id", f"E{len(nodes)+1}"),
description=nd.get("description", ""),
domain=nd.get("domain", "unknown"),
severity=nd.get("severity", "medium"),
time_offset_hours=nd.get("time_offset_hours"),
mechanism=nd.get("mechanism", ""),
parent_ids=nd.get("parent_ids", []),
))
return CascadeChain(
event_id=event.event_id,
trigger_summary=data.get(
"trigger_summary",
f"Flood in {event.location or event.country} on {event.start_date}"
),
trigger_country=event.country,
trigger_iso=event.iso,
trigger_date=event.start_date,
trigger_severity=_infer_severity(event),
cascade_events=nodes,
source_articles=[],
extraction_date=date.today(),
)
except (json.JSONDecodeError, KeyError, ValueError) as e:
logger.error(f"Failed to parse LLM response for {event.event_id}: {e}")
logger.error(f"Response: {response[:500]}")
return None
def main():
config = load_config()
paths = config["paths"]
# Load ALL events (both train and test — we want cascade chains for all)
all_events_data = json.loads(Path(paths["events_catalog"]).read_text())
all_events = [FloodEvent(**e) for e in all_events_data]
logger.info(f"Loaded {len(all_events)} events total")
llm_client = create_llm_client(config)
chains_index = []
for i, event in enumerate(all_events):
logger.info(
f"\n[{i+1}/{len(all_events)}] Generating cascade chain for: "
f"{event.event_id} ({event.country}, {event.start_date})"
)
chain = generate_cascade_from_metadata(event, llm_client, config)
if chain:
path = save_cascade_chain(chain, config)
logger.info(f" Saved: {path} ({len(chain.cascade_events)} cascade nodes)")
chains_index.append({
"event_id": chain.event_id,
"trigger_summary": chain.trigger_summary,
"country": chain.trigger_country,
"date": str(chain.trigger_date),
"num_cascade_nodes": len(chain.cascade_events),
"domains": list({n.domain for n in chain.cascade_events}),
"file": str(path),
})
else:
logger.warning(f" Failed to generate chain for {event.event_id}")
# Save index
index_path = Path(paths["cascade_index"])
index_path.parent.mkdir(parents=True, exist_ok=True)
index_path.write_text(json.dumps(chains_index, indent=2, ensure_ascii=False))
logger.info(f"\nSaved cascade chains index: {index_path}")
logger.info(f"Total chains generated: {len(chains_index)}/{len(all_events)}")
logger.info("Done. Run scripts/04_build_vectordb.py next.")
if __name__ == "__main__":
main()