DiplomovaPraca/scripts/lightrag_ingest/build_section_docs.py

"""Build short LightRAG documents from selected ADC leaflet sections.

Instead of the full leaflet, each document contains only the parts relevant
for the agent: indications, contraindications, warnings and interactions.
Products are deduplicated by their active substance set (see
ingest_adc_structured.dedupe_key), keeping the product with the richest text.

The output is a JSON list in the same shape that ingest_adc_structured.py
expects (`urls`, `product`, `graph_hints`, `lightrag_text`), so it can be
uploaded with:

    python scripts/lightrag_ingest/ingest_adc_structured.py \
        --input outputs/lightrag_sections/adc_section_docs.json \
        --progress outputs/lightrag_sections/progress.json --dedupe none --limit 100000
"""

from __future__ import annotations

import argparse
import json
import sys
from collections import Counter
from pathlib import Path
from typing import Any

sys.path.insert(0, str(Path(__file__).resolve().parent))
from ingest_adc_structured import DEFAULT_INPUT, dedupe_key, load_json, write_json  # noqa: E402


DEFAULT_OUTPUT = Path("outputs/lightrag_sections/adc_section_docs.json")

# (heading in the document, where to read it in `pil`, fallback key in `graph_hints`)
SECTIONS = (
    ("Indikácie (na čo sa liek používa)", ("sections", "what_is_it"), "indications_text"),
    ("Kontraindikácie (kedy liek neužívať)", ("subsections", "contraindications"), "contraindications_text"),
    ("Upozornenia a opatrenia", ("subsections", "warnings"), "warnings_text"),
    ("Interakcie s inými liekmi", ("subsections", "interactions"), "interactions_text"),
)

if hasattr(sys.stdout, "reconfigure"):
    sys.stdout.reconfigure(encoding="utf-8")


def section_text(record: dict[str, Any], location: tuple[str, str], fallback: str) -> str:
    pil = record.get("pil") or {}
    group = pil.get(location[0]) if isinstance(pil, dict) else None
    text = group.get(location[1]) if isinstance(group, dict) else None
    if not text:
        text = (record.get("graph_hints") or {}).get(fallback)
    return str(text or "").strip()


def build_document(record: dict[str, Any]) -> tuple[str, int]:
    """Return the document text and how many clinical sections it contains."""
    product = record.get("product") or {}
    hints = record.get("graph_hints") or {}
    urls = record.get("urls") or {}

    substances = hints.get("active_substances") or []
    if isinstance(substances, str):
        substances = [substances]
    classes = [
        str(item.get("name"))
        for item in hints.get("classification_codes") or []
        if isinstance(item, dict) and item.get("name")
    ]

    lines = [f"Liek: {hints.get('drug') or product.get('name') or ''}"]
    if substances:
        lines.append(f"Účinná látka: {', '.join(str(s) for s in substances)}")
    if hints.get("dosage_form"):
        lines.append(f"Lieková forma: {hints['dosage_form']}")
    if classes:
        lines.append(f"Skupina: {' > '.join(classes)}")
    if urls.get("pil"):
        lines.append(f"Zdroj: {urls['pil']}")

    found = 0
    for heading, location, fallback in SECTIONS:
        text = section_text(record, location, fallback)
        if text:
            found += 1
            lines.append(f"\n## {heading}\n{text}")
    return "\n".join(lines), found


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    parser.add_argument("--input", type=Path, default=DEFAULT_INPUT)
    parser.add_argument("--output", type=Path, default=DEFAULT_OUTPUT)
    parser.add_argument("--dedupe", choices=("none", "text", "substance"), default="substance")
    args = parser.parse_args()

    print(f"Loading {args.input} ...")
    records = load_json(args.input)
    stats: Counter = Counter(input=len(records))

    groups: dict[str, tuple[int, dict[str, Any]]] = {}
    for record in records:
        text, found = build_document(record)
        # Keep only products with at least one of contraindications / warnings / interactions.
        if found == 0 or not any(
            section_text(record, location, fallback) for _, location, fallback in SECTIONS[1:]
        ):
            stats["no_clinical_sections"] += 1
            continue
        doc = {
            "urls": record.get("urls") or {},
            "product": {"name": (record.get("product") or {}).get("name")},
            "graph_hints": {
                key: (record.get("graph_hints") or {}).get(key)
                for key in ("drug", "active_substances", "contraindications_text",
                            "warnings_text", "interactions_text", "side_effects_text")
            },
            "lightrag_text": text,
        }
        key = dedupe_key(record, args.dedupe)
        if key in groups:
            stats["duplicate"] += 1
            if len(text) <= groups[key][0]:
                continue
        groups[key] = (len(text), doc)

    docs = [doc for _, doc in groups.values()]
    write_json(args.output, docs)

    total_chars = sum(len(doc["lightrag_text"]) for doc in docs)
    stats["documents"] = len(docs)
    print(f"Stats: {dict(stats)}")
    print(f"Average document length: {total_chars // max(len(docs), 1)} chars, total {total_chars / 1e6:.1f} M chars")
    print(f"Saved: {args.output}")


if __name__ == "__main__":
    main()