src/prompt_injection_framework/io/data_loader.py

from __future__ import annotations

import html
import unicodedata
from pathlib import Path
from typing import Any, Iterator

from prompt_injection_framework.utils.io import read_jsonl

REQUIRED_CORE_FIELDS = (
    "case_id",
    "source_name",
    "source_type",
    "attack_family",
    "attack_surface",
    "prompt",
    "success_criteria",
)


def read_jsonl_records(path: Path) -> Iterator[dict[str, Any]]:
    for record in read_jsonl(path):
        validate_core_fields(record)
        yield record


def normalize_text(value: str | None, *, decode_html: bool = False) -> str:
    if value is None:
        return ""

    text = unicodedata.normalize("NFKC", value)
    if decode_html:
        text = html.unescape(text)
    text = text.replace("\u200b", "").replace("\u200c", "").replace("\u200d", "")
    return " ".join(text.lower().split())


def validate_core_fields(case: dict[str, Any]) -> None:
    missing = [field for field in REQUIRED_CORE_FIELDS if not str(case.get(field, "")).strip()]
    if missing:
        raise ValueError(f"Missing required attack-case fields: {', '.join(missing)}")