scripts/ingest_ebm.py

from __future__ import annotations

import argparse
import json
import sys
from pathlib import Path

ROOT = Path(__file__).resolve().parent.parent
sys.path.insert(0, str(ROOT))

from src.chunking import dataframe_to_search_corpus
from src.parser import filter_df_by_fachgruppe, parse_ebm_xml_to_dataframe


def main() -> None:
    parser = argparse.ArgumentParser(description="Parse the EBM XML into structured artifacts.")
    parser.add_argument("--xml", required=True, help="Path to the official EBM XML.")
    parser.add_argument("--output", required=True, help="Directory for generated artifacts.")
    parser.add_argument("--fachgruppe-filter", action="store_true", help="Filter to Fachgruppe 001 only (for full EBM downloads).")
    args = parser.parse_args()

    xml_path = Path(args.xml)
    output_dir = Path(args.output)
    output_dir.mkdir(parents=True, exist_ok=True)

    df = parse_ebm_xml_to_dataframe(str(xml_path))
    
    # Apply Fachgruppe filter only if requested
    if args.fachgruppe_filter:
        print("Applying Fachgruppe 001 filter...")
        df = filter_df_by_fachgruppe(df)
        if df.empty:
            raise ValueError(
                "No Fachgruppe 001 documents found in the provided XML. "
                "Please check the XML file or remove the --fachgruppe-filter flag."
            )
    
    if df.empty:
        raise ValueError(
            "No documents found in the provided XML. "
            "Please provide a valid KBV EBM XML file."
        )
    
    print(f"Processing {len(df)} documents...")
    df.to_parquet(output_dir / "ebm.parquet", index=False)
    df.to_json(output_dir / "ebm.jsonl", orient="records", lines=True, force_ascii=False)

    corpus = dataframe_to_search_corpus(df)
    (output_dir / "ebm_documents.jsonl").write_text(
        "\n".join(json.dumps(item, ensure_ascii=False) for item in corpus),
        encoding="utf-8",
    )


if __name__ == "__main__":
    main()