pdf-quiz-generator/backend/scripts/article_pipeline.py

"""Retrieval and import for articles written by hand rather than by an API model.

Splits the job in two, because only one half belongs to a machine:

    fetch   — bge-m3 embeds the topic, Milvus returns the passages, and the
              references are derived from their metadata. No writing model is
              involved at any point.
    import  — takes a finished article as JSON and stores it as a draft.

Whoever writes the prose in between reads the passages and writes original text
from them. The references are built here from what retrieval actually returned,
so they cannot be invented by whoever is writing, which is the same property the
generated route had.

    docker compose exec backend python -m scripts.article_pipeline fetch "Croup" --out /app/uploads/_work
    docker compose exec backend python -m scripts.article_pipeline topics --limit 40
    docker compose exec backend python -m scripts.article_pipeline import /app/uploads/_work/croup.article.json
"""
import argparse
import json
import pathlib
import re
import sys
import uuid

from sqlalchemy import text as sa_text

from app.database import SessionLocal
from app.models.article import Article
from app.services import article_service, clinical_library
from app.services.article_writer import (
    MIN_SOURCE_CHARS, PASSAGES, SHELF, slugify, _unique_slugs,
)

VARIANTS = ("short", "long", "clinical")

# A bare condition name is a thin query. "Rickets" alone retrieved five passages
# about Rickettsia — an embedding has little to go on in one word, and the
# nearest neighbours of a short string are whatever looks like it. Saying what
# kind of thing is wanted removes the collision entirely.
QUERY_SHAPE = "{topic} in children: definition, causes, clinical features, diagnosis and management"

# Category names that are a shelf rather than a condition. They retrieve chapter
# headings and whatever happens to sit near them, and an article called
# "Pediatric Nephrology" is a department, not something to revise.
UMBRELLA = re.compile(
    r"^(pediatric|paediatric)\b|\b(medicine|surgery|disorder|disorders|care|health|"
    r"nephrology|neurology|cardiology|oncology|dermatology|psychiatry|radiology|"
    r"pulmonology|endocrinology|gastroenterology|rheumatology|urology|"
    r"hematology|immunology|genetics|orthopedics|ophthalmology)$",
    re.I)


def cmd_topics(args):
    """Conditions that still have no article, biggest first."""
    db = SessionLocal()
    try:
        # One condition, one article. The same name is a leaf under several
        # disciplines — "Hemolytic Uremic Syndrome" sits under Infectious
        # Disease, Nephrology and Emergency Medicine — and writing it three
        # times would be three articles nobody asked for, plus a collision in
        # the importer, which keys on the name.
        rows = db.execute(sa_text("""
            SELECT DISTINCT ON (lower(c.name)) c.id, c.name, SUM(COUNT(q.id)) OVER (
                       PARTITION BY lower(c.name)) AS uses
            FROM question_categories c
            JOIN questions q ON q.question_category_id = c.id
            WHERE NOT EXISTS (SELECT 1 FROM question_categories k WHERE k.parent_id = c.id)
            GROUP BY c.id, c.name
            ORDER BY lower(c.name), COUNT(q.id) DESC, c.id
        """)).fetchall()
        rows = sorted(rows, key=lambda r: (-int(r[2]), r[1]))
        have = {row[0] for row in db.query(Article.slug).all()}
        todo = [(cid, name, uses) for cid, name, uses in rows
                if slugify(name) not in have and not UMBRELLA.search(name.strip())]
        for category_id, name, uses in todo[:args.limit]:
            print(f"{category_id}\t{uses}\t{name}")
        print(f"\n# {len(todo)} topics without an article", file=sys.stderr)
    finally:
        db.close()
    return 0


def cmd_fetch(args):
    """Everything needed to write one article, and nothing that writes it."""
    passages = clinical_library.search(QUERY_SHAPE.format(topic=args.topic),
                                       limit=PASSAGES, folder_contains=SHELF)
    chars = sum(len(p["text"]) for p in passages)
    payload = {
        "topic": args.topic,
        "slug": slugify(args.topic),
        "category_id": args.category_id,
        "passage_count": len(passages),
        "source_chars": chars,
        "enough_material": len(passages) >= 3 and chars >= MIN_SOURCE_CHARS,
        "references": clinical_library.references_from(passages),
        "passages": [{
            "source": f"{p['source'].get('title', 'source')}"
                      f"{', p. ' + str(p['source']['page']) if p['source'].get('page') else ''}",
            "text": p["text"],
        } for p in passages],
    }
    out = pathlib.Path(args.out)
    out.mkdir(parents=True, exist_ok=True)
    target = out / f"{payload['slug']}.sources.json"
    target.write_text(json.dumps(payload, indent=1))
    print(target)
    return 0


def _sections(blocks, variant):
    out = []
    for block in blocks or []:
        if not isinstance(block, dict):
            continue
        title = str(block.get("title") or "").strip()
        content = str(block.get("content") or "").strip()
        if not title or not content:
            continue
        out.append({"id": uuid.uuid4().hex, "slug": slugify(title)[:60] or f"s{len(out) + 1}",
                    "title": title[:300], "content": content,
                    "parent_id": None, "variant": variant})
    return out


def cmd_import(args):
    """Store a finished article as a draft, refusing anything half-written."""
    data = json.loads(pathlib.Path(args.path).read_text())
    topic = (data.get("topic") or "").strip()
    if not topic:
        print("  No topic in that file.")
        return 1

    sections = _unique_slugs([s for v in VARIANTS for s in _sections(data.get(v), v)])
    missing = [v for v in VARIANTS if not any(s["variant"] == v for s in sections)]
    if missing:
        # A view the reader is offered and finds empty is worse than one that was
        # never promised, so an incomplete article is refused rather than stored.
        print(f"  {topic}: missing {', '.join(missing)} — not imported.")
        return 1

    db = SessionLocal()
    try:
        slug = data.get("slug") or slugify(topic)
        if db.query(Article.id).filter(Article.slug == slug).first():
            print(f"  {topic}: already exists.")
            return 0
        article = Article(
            slug=slug, title=topic[:300],
            summary=(data.get("summary") or "").strip()[:2000] or None,
            sections=sections, category_id=data.get("category_id"),
            status="draft",
            references_json=data.get("references") or [],
            generated_by=data.get("written_by") or "claude",
        )
        db.add(article)
        db.flush()
        article_service.record_slug(db, article)
        db.commit()
        print(f"  {topic}: imported as draft #{article.id} "
              f"({len(sections)} sections, {len(article.references_json)} references)")
    finally:
        db.close()
    return 0


def main():
    parser = argparse.ArgumentParser()
    sub = parser.add_subparsers(dest="cmd", required=True)

    topics = sub.add_parser("topics", help="conditions still without an article")
    topics.add_argument("--limit", type=int, default=50)
    topics.set_defaults(func=cmd_topics)

    fetch = sub.add_parser("fetch", help="retrieve the source passages for one topic")
    fetch.add_argument("topic")
    fetch.add_argument("--category-id", type=int, default=None)
    fetch.add_argument("--out", default="/app/uploads/_work")
    fetch.set_defaults(func=cmd_fetch)

    imp = sub.add_parser("import", help="store a written article as a draft")
    imp.add_argument("path")
    imp.set_defaults(func=cmd_import)

    args = parser.parse_args()
    return args.func(args)


if __name__ == "__main__":
    sys.exit(main())