hyungi_document_server/scripts/presegment_attended.py

"""presegment PR3 — HOLD 거대문서 유인 분할 CLI (plan ds-presegment-mapreduce-2).

deep_summary 워커가 HOLD(payload.presegment.awaiting_split=true) 로 보류한
hybrid/whole tier 거대문서를, 사람이(유인 클로드 세션) 경계를 완성해 재개시키는 도구.

사용법 (fastapi 컨테이너 안에서 실행):
  docker compose exec fastapi python /app/scripts/presegment_attended.py list
  docker compose exec fastapi python /app/scripts/presegment_attended.py export --doc 44443 --out /app/logs/preseg_44443
  docker compose exec fastapi python /app/scripts/presegment_attended.py apply --doc 44443 --boundaries /app/logs/preseg_44443/boundaries_template.json --dry-run
  docker compose exec fastapi python /app/scripts/presegment_attended.py apply --doc 44443 --boundaries /app/logs/preseg_44443/boundaries_template.json

워크플로우:
  1. list   — awaiting_split 문서 확인.
  2. export — 문서 통계·hier 개요·자동 pack 유닛 제안·초과 섹션 본문 덤프·boundaries
              템플릿 JSON 생성. 유인 클로드 세션은 이 파일들만 읽고 TODO 스팬을
              CAP 이하 경계들로 분할해 템플릿을 완성한다.
  3. apply  — 완성된 boundaries 검증(단조·비중첩·범위·커버리지 90%+·유닛 캡) 후
              payload.presegment.units_override 기록 + awaiting_split 해제 +
              deferred_until 제거(즉시 재개). 워커가 다음 사이클에 map-reduce 재개.

stdout 규약: 사람이 읽는 요약 행 + '{' 로 시작하는 기계 파싱용 JSON 라인(1건 1라인).
사람용 행은 절대 '{' 로 시작하지 않는다.
"""

from __future__ import annotations

import argparse
import asyncio
import json
import os
import re
import sys
from datetime import datetime, timezone
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parent.parent / "app"))

from sqlalchemy import text as sql_text  # noqa: E402
from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker, create_async_engine  # noqa: E402

from services.hier_decomp.builder import build_hier_tree  # noqa: E402
from services.summarize_units import (  # noqa: E402
    CAP_TOKENS,
    OVERRIDE_MIN_COVERAGE_PCT,
    TRIGGER_TOKENS,
    choose_override_source,
    estimate_tokens,
    extract_leaves,
    leaf_spans,
    plan_summarize_units,
    validate_override_boundaries,
)

# 초과 섹션 본문 덤프 분할 단위 (유인 세션 컨텍스트 보호)
DUMP_CHUNK_CHARS = 200_000


def _jsonl(obj: dict) -> None:
    """기계 파싱용 JSON 라인 — 반드시 '{' 로 시작하는 단독 라인."""
    print(json.dumps(obj, ensure_ascii=False, default=str))


def _session_factory():
    database_url = os.getenv(
        "DATABASE_URL",
        "postgresql+asyncpg://pkm:pkm@postgres:5432/pkm",
    )
    engine = create_async_engine(database_url)
    return engine, async_sessionmaker(engine, class_=AsyncSession, expire_on_commit=False)


# ─── list ────────────────────────────────────────────────────────────────────

LIST_SQL = """
SELECT q.id AS queue_id, q.document_id, q.status, q.attempts,
       q.payload::text AS payload_text,
       LEFT(COALESCE(d.title, '(제목 없음)'), 80) AS title
  FROM processing_queue q
  JOIN documents d ON d.id = q.document_id
 WHERE q.stage = 'deep_summary'
   AND q.status IN ('pending', 'processing', 'failed')
   AND (q.payload -> 'presegment' ->> 'awaiting_split') = 'true'
 ORDER BY q.id
"""


async def cmd_list() -> int:
    engine, factory = _session_factory()
    try:
        async with factory() as session:
            rows = (await session.execute(sql_text(LIST_SQL))).mappings().all()
    finally:
        await engine.dispose()

    print(f"awaiting_split 보류 문서 {len(rows)}건")
    for r in rows:
        payload = json.loads(r["payload_text"] or "{}")
        preseg = payload.get("presegment") or {}
        oversized = preseg.get("oversized_sections") or []
        print(
            f"  doc {r['document_id']} [{r['title']}] queue={r['queue_id']} status={r['status']} "
            f"tier={preseg.get('tier')} over%={preseg.get('over_pct')} "
            f"tokens={preseg.get('total_est_tokens'):,} units={preseg.get('units')}"
            if isinstance(preseg.get("total_est_tokens"), int)
            else f"  doc {r['document_id']} [{r['title']}] queue={r['queue_id']} status={r['status']}"
        )
        print(f"    초과 섹션 {len(oversized)}건: {', '.join(str(t) for t in oversized[:3] if t)}")
        print(
            f"    보류 알람={preseg.get('alerted_at') or '-'} / "
            f"재개 예정={payload.get('deferred_until') or '(즉시)'}"
            + (f" / 거부 사유={preseg.get('override_rejected')}" if preseg.get("override_rejected") else "")
        )
        _jsonl({
            "cmd": "list",
            "queue_id": r["queue_id"],
            "doc_id": r["document_id"],
            "title": r["title"],
            "status": r["status"],
            "tier": preseg.get("tier"),
            "over_pct": preseg.get("over_pct"),
            "total_est_tokens": preseg.get("total_est_tokens"),
            "units": preseg.get("units"),
            "oversized_sections": oversized,
            "alerted_at": preseg.get("alerted_at"),
            "deferred_until": payload.get("deferred_until"),
            "override_rejected": preseg.get("override_rejected"),
        })
    return 0


# ─── export ──────────────────────────────────────────────────────────────────

def _safe_name(title: str | None, fallback: str) -> str:
    t = re.sub(r"[^0-9A-Za-z가-힣._-]+", "_", (title or fallback)).strip("_")
    return (t or fallback)[:60]


def _build_outline(text: str) -> str:
    """hier_decomp builder 재사용 — window-split 억제(요약 계획과 동일 환경) 개요."""
    nodes = build_hier_tree(text, leaf_target_max=sys.maxsize, leaf_hard_max=sys.maxsize)
    lines = []
    for n in nodes:
        indent = "  " * max(n.level, 0)
        title = n.section_title or "(preamble)"
        tok = estimate_tokens(n.text)
        mark = " [CAP 초과]" if n.is_leaf and tok > CAP_TOKENS else ""
        lines.append(f"{indent}- L{n.level} {title} — {tok:,} tok{mark}")
    return "\n".join(lines)


async def cmd_export(doc_id: int, out_dir: str) -> int:
    engine, factory = _session_factory()
    try:
        async with factory() as session:
            row = (await session.execute(
                sql_text(
                    "SELECT id, title, md_content, extracted_text FROM documents WHERE id = :d"
                ),
                {"d": doc_id},
            )).mappings().first()
    finally:
        await engine.dispose()

    if not row:
        print(f"[error] 문서 id={doc_id} 없음")
        _jsonl({"cmd": "export", "ok": False, "doc_id": doc_id, "error": "document_not_found"})
        return 1

    source, text = choose_override_source(row["md_content"], row["extracted_text"])
    if not text.strip():
        print(f"[error] 문서 id={doc_id} 본문 비어있음 (md_content/extracted_text 둘 다)")
        _jsonl({"cmd": "export", "ok": False, "doc_id": doc_id, "error": "empty_text"})
        return 1

    plan = plan_summarize_units(text)
    leaves = extract_leaves(text)
    spans = leaf_spans(text, leaves)

    out = Path(out_dir)
    out.mkdir(parents=True, exist_ok=True)
    files: list[str] = []
    now_iso = datetime.now(timezone.utc).isoformat(timespec="seconds")

    # ① 통계 + hier 개요
    oversized_units = [u for u in plan.units if u.over_cap]
    overview = [
        f"# presegment export — doc {doc_id}",
        "",
        f"- 제목: {row['title'] or '(제목 없음)'}",
        f"- source: {source} (len={len(text):,}자, 오프셋 기준 텍스트)",
        f"- 추정 토큰: {plan.total_est_tokens:,} (trigger={TRIGGER_TOKENS:,} / cap={CAP_TOKENS:,})",
        f"- plan: mode={plan.mode} tier={plan.tier} over%={plan.over_pct}",
        f"- 유닛: 자동 pack {len(plan.units) - len(oversized_units)}개 + CAP 초과 {len(oversized_units)}개",
        f"- 생성: {now_iso}",
        "",
        "## 유닛 제안 (summarize_units greedy-pack)",
        "",
    ]
    for u in plan.units:
        if not u.leaf_indexes:
            continue
        s = spans[u.leaf_indexes[0]][0]
        e = spans[u.leaf_indexes[-1]][1]
        titles = " · ".join(t for t in u.section_titles if t) or "(무제 구간)"
        flag = " ★CAP 초과 — 분할 필요" if u.over_cap else ""
        overview.append(f"- 유닛 {u.index}: [{s}, {e}) {u.est_tokens:,} tok — {titles[:120]}{flag}")
    overview += ["", "## hier 개요", "", _build_outline(text), ""]
    (out / "overview.md").write_text("\n".join(overview), encoding="utf-8")
    files.append("overview.md")

    # ③ 초과 섹션 본문 덤프 (섹션당 파일 · 200K자 단위 분할 · 파일명에 절대 스팬)
    boundaries: list[dict] = []
    for u in plan.units:
        if not u.leaf_indexes:
            continue
        s = spans[u.leaf_indexes[0]][0]
        e = spans[u.leaf_indexes[-1]][1]
        title = next((t for t in u.section_titles if t), None)
        if not u.over_cap:
            # ④ 자동 pack 유닛은 템플릿에 채워둔다
            boundaries.append({"start": s, "end": e, "title": title or f"유닛 {u.index}"})
            continue
        boundaries.append({
            "start": s, "end": e, "title": title or f"유닛 {u.index}",
            "todo": (
                f"CAP 초과({u.est_tokens:,} tok > {CAP_TOKENS:,}) — 이 스팬을 cap 이하 "
                "경계 여러 개로 교체하고 todo 키를 제거할 것"
            ),
        })
        seg = text[s:e]
        base = f"oversized_{u.index:03d}_{_safe_name(title, f'unit{u.index}')}"
        for k in range(0, len(seg), DUMP_CHUNK_CHARS):
            cs, ce = s + k, s + min(k + DUMP_CHUNK_CHARS, len(seg))
            fname = f"{base}.{cs}_{ce}.md"
            # 본문 원문 그대로 (헤더 미부착 — 파일 내 로컬 오프셋 + 파일명 cs 로 절대 오프셋 계산)
            (out / fname).write_text(text[cs:ce], encoding="utf-8")
            files.append(fname)

    # ④ boundaries 템플릿 JSON
    template = {
        "doc_id": doc_id,
        "source": source,
        "source_len": len(text),
        "cap_tokens": CAP_TOKENS,
        "generated_at": now_iso,
        "boundaries": boundaries,
    }
    (out / "boundaries_template.json").write_text(
        json.dumps(template, ensure_ascii=False, indent=2), encoding="utf-8"
    )
    files.append("boundaries_template.json")

    # 유인 세션용 작업 안내
    readme = f"""# doc {doc_id} 유인 분할 안내

1. overview.md 로 구조 파악 (유닛 제안 + hier 개요).
2. oversized_*.md 본문을 읽고 의미 경계를 정한다.
   - 파일명 `..._<cs>_<ce>.md` 의 cs = 파일 첫 문자의 절대 오프셋.
   - 절대 오프셋 = cs + 파일 내 로컬 오프셋.
3. boundaries_template.json 의 todo 항목을 cap({CAP_TOKENS:,} tok) 이하 경계 여러 개로
   교체하고 todo 키를 제거한다. 나머지 자동 pack 항목은 그대로 둬도 된다.
   - 토큰 추정: 한글 0.529 tok/자 · 기타 0.217 tok/자 (services/summarize_units.py).
   - 규칙: start 단조증가 · 비중첩 · 전체 커버리지 {OVERRIDE_MIN_COVERAGE_PCT:.0f}%+ · 유닛당 cap 이하.
4. 검증/적용:
   python /app/scripts/presegment_attended.py apply --doc {doc_id} --boundaries <FILE> --dry-run
   python /app/scripts/presegment_attended.py apply --doc {doc_id} --boundaries <FILE>
"""
    (out / "README.md").write_text(readme, encoding="utf-8")
    files.append("README.md")

    print(f"doc {doc_id} [{row['title'] or '(제목 없음)'}] export 완료 → {out}")
    print(
        f"  source={source} len={len(text):,}자 tokens={plan.total_est_tokens:,} "
        f"tier={plan.tier} over%={plan.over_pct}"
    )
    print(f"  자동 pack 유닛 {len(plan.units) - len(oversized_units)}개 / TODO(초과) {len(oversized_units)}개")
    print(f"  파일 {len(files)}개: {', '.join(files[:6])}{' ...' if len(files) > 6 else ''}")
    _jsonl({
        "cmd": "export", "ok": True, "doc_id": doc_id, "out": str(out),
        "source": source, "source_len": len(text),
        "total_est_tokens": plan.total_est_tokens, "tier": plan.tier,
        "over_pct": plan.over_pct,
        "units_auto": len(plan.units) - len(oversized_units),
        "units_todo": len(oversized_units),
        "files": files,
    })
    return 0


# ─── apply ───────────────────────────────────────────────────────────────────

QUEUE_ROW_SQL = """
SELECT id, status, attempts, payload::text AS payload_text
  FROM processing_queue
 WHERE document_id = :d AND stage = 'deep_summary'
   AND status IN ('pending', 'processing', 'failed')
 ORDER BY id DESC
 LIMIT 1
"""

APPLY_UPDATE_SQL = """
UPDATE processing_queue
   SET payload = CAST(:payload AS JSONB),
       status = 'pending',
       attempts = 0,
       error_message = NULL
 WHERE id = :qid
"""


async def cmd_apply(doc_id: int, boundaries_file: str, dry_run: bool) -> int:
    raw = json.loads(Path(boundaries_file).read_text(encoding="utf-8"))
    if isinstance(raw, dict):
        boundaries = raw.get("boundaries") or []
        declared_source = raw.get("source")
        declared_len = raw.get("source_len")
        if raw.get("doc_id") not in (None, doc_id):
            print(f"[error] boundaries 파일 doc_id={raw.get('doc_id')} != --doc {doc_id}")
            _jsonl({"cmd": "apply", "ok": False, "doc_id": doc_id, "error": "doc_id_mismatch"})
            return 1
    else:
        boundaries, declared_source, declared_len = raw, None, None

    engine, factory = _session_factory()
    try:
        async with factory() as session:
            doc = (await session.execute(
                sql_text(
                    "SELECT id, title, md_content, extracted_text FROM documents WHERE id = :d"
                ),
                {"d": doc_id},
            )).mappings().first()
            if not doc:
                print(f"[error] 문서 id={doc_id} 없음")
                _jsonl({"cmd": "apply", "ok": False, "doc_id": doc_id, "error": "document_not_found"})
                return 1

            qrow = (await session.execute(sql_text(QUEUE_ROW_SQL), {"d": doc_id})).mappings().first()
            if not qrow:
                print(f"[error] doc {doc_id} 의 활성 deep_summary 큐 행 없음 (pending/processing/failed)")
                _jsonl({"cmd": "apply", "ok": False, "doc_id": doc_id, "error": "queue_row_not_found"})
                return 1
            if qrow["status"] == "processing":
                print(f"[error] queue {qrow['id']} 가 processing 중 — 워커 완료/보류 후 재시도")
                _jsonl({"cmd": "apply", "ok": False, "doc_id": doc_id, "error": "queue_processing"})
                return 1

            # 오프셋 기준 텍스트 — export 와 동일 규칙 (파일에 source 선언 시 그 선언 우선)
            if declared_source in ("md_content", "extracted_text"):
                source = declared_source
                text = (doc["md_content"] if source == "md_content" else doc["extracted_text"]) or ""
            else:
                source, text = choose_override_source(doc["md_content"], doc["extracted_text"])
            if declared_len is not None and declared_len != len(text):
                print(
                    f"[error] source_len 불일치 — 파일={declared_len:,} vs 현재 {source}={len(text):,}"
                    " (본문 재생성됨 — export 부터 재실행)"
                )
                _jsonl({"cmd": "apply", "ok": False, "doc_id": doc_id, "error": "source_len_mismatch"})
                return 1

            check = validate_override_boundaries(text, boundaries)
            for w in check.warnings:
                print(f"  [warn] {w}")
            if not check.ok:
                print(f"[error] 경계 검증 실패 — {len(check.errors)}건:")
                for e in check.errors:
                    print(f"  - {e}")
                _jsonl({
                    "cmd": "apply", "ok": False, "doc_id": doc_id,
                    "error": "validation_failed", "errors": check.errors,
                    "warnings": check.warnings, "coverage_pct": check.coverage_pct,
                })
                return 1

            print(
                f"doc {doc_id} [{doc['title'] or '(제목 없음)'}] 경계 검증 통과 — "
                f"유닛 {len(check.boundaries)}개 / 커버리지 {check.coverage_pct}% / "
                f"최대 유닛 {max(check.unit_tokens):,} tok (cap {CAP_TOKENS:,})"
            )
            for i, ((s, e, t), tok) in enumerate(zip(check.boundaries, check.unit_tokens)):
                print(f"  유닛 {i}: [{s}, {e}) {tok:,} tok — {t or '(무제)'}")

            payload = json.loads(qrow["payload_text"] or "{}")
            preseg = dict(payload.get("presegment") or {})
            preseg["units_override"] = {
                "source": source,
                "source_len": len(text),
                "boundaries": [[s, e, t] for s, e, t in check.boundaries],
                "applied_at": datetime.now(timezone.utc).isoformat(timespec="seconds"),
            }
            preseg["awaiting_split"] = False
            # 알람 dedupe 리셋(다음 이벤트는 신선하게 발화) + 이전 거부/맵 잔재 제거
            for k in ("alerted_at", "override_rejected", "override_rejected_at", "map_results"):
                preseg.pop(k, None)
            payload["presegment"] = preseg
            payload.pop("deferred_until", None)  # 즉시 재개

            if dry_run:
                print(f"  [dry-run] queue {qrow['id']} 미변경 — 위 경계로 적용 가능")
                _jsonl({
                    "cmd": "apply", "ok": True, "dry_run": True, "doc_id": doc_id,
                    "queue_id": qrow["id"], "units": len(check.boundaries),
                    "coverage_pct": check.coverage_pct, "unit_tokens": check.unit_tokens,
                })
                return 0

            await session.execute(
                sql_text(APPLY_UPDATE_SQL),
                {"payload": json.dumps(payload, ensure_ascii=False), "qid": qrow["id"]},
            )
            await session.commit()
            print(
                f"  적용 완료 — queue {qrow['id']} status=pending, deferred_until 제거 "
                f"(다음 queue_consumer 사이클에 재개)"
            )
            _jsonl({
                "cmd": "apply", "ok": True, "dry_run": False, "doc_id": doc_id,
                "queue_id": qrow["id"], "units": len(check.boundaries),
                "coverage_pct": check.coverage_pct, "unit_tokens": check.unit_tokens,
            })
            return 0
    finally:
        await engine.dispose()


# ─── main ────────────────────────────────────────────────────────────────────

def main() -> int:
    parser = argparse.ArgumentParser(description="presegment 유인 분할 CLI (PR3)")
    sub = parser.add_subparsers(dest="cmd", required=True)

    sub.add_parser("list", help="awaiting_split 보류 문서 목록")

    p_export = sub.add_parser("export", help="유인 분할 작업 패키지 덤프")
    p_export.add_argument("--doc", type=int, required=True)
    p_export.add_argument("--out", default=None, help="출력 디렉토리 (기본 ./preseg_export_<doc>)")

    p_apply = sub.add_parser("apply", help="완성된 boundaries 검증·적용(재개)")
    p_apply.add_argument("--doc", type=int, required=True)
    p_apply.add_argument("--boundaries", required=True, help="boundaries JSON 파일 경로")
    p_apply.add_argument("--dry-run", action="store_true", help="검증만 하고 DB 미변경")

    args = parser.parse_args()
    if args.cmd == "list":
        return asyncio.run(cmd_list())
    if args.cmd == "export":
        out = args.out or f"./preseg_export_{args.doc}"
        return asyncio.run(cmd_export(args.doc, out))
    if args.cmd == "apply":
        return asyncio.run(cmd_apply(args.doc, args.boundaries, args.dry_run))
    return 2


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