| """Long-document batching and safety tests."""
|
|
|
| from __future__ import annotations
|
|
|
| import os
|
|
|
| os.environ.setdefault("LANGUAGE_TOOL_ENABLED", "false")
|
| os.environ.setdefault("LANGUAGE_TOOL_URL", "")
|
| os.environ.setdefault("GRAMMAR_FIX_OUTPUT", "false")
|
| os.environ.setdefault("ENGINE_USE_MINILM_SAFETY", "false")
|
|
|
| from app.engine.normalize import detect_blocks, normalize_text
|
| from app.engine.orchestrator import rewrite_document
|
| from app.engine.safety import check_safety
|
| from app.engine.segment import iter_paragraph_batches, word_count
|
|
|
|
|
| def _make_long_doc(target_words: int = 10500) -> str:
|
| templates = [
|
| "Ram went to school yesterday happily.",
|
| "The committee approved the budget today carefully.",
|
| "Workers cleaned the factory floor this morning quickly.",
|
| "Students completed the assignment last night diligently.",
|
| "Although rain fell heavily, traffic moved slowly through town.",
|
| "# Section Heading\n\n",
|
| "- Keep this list item unchanged forever.\n",
|
| "She sent the report to https://files.example.org/a.pdf yesterday.\n\n",
|
| ]
|
| parts: list[str] = []
|
| while word_count("\n\n".join(parts)) < target_words:
|
| for t in templates:
|
| parts.append(t.strip())
|
| if word_count("\n\n".join(parts)) >= target_words:
|
| break
|
| return "\n\n".join(parts)
|
|
|
|
|
| def test_long_document_batching():
|
| doc = _make_long_doc(10500)
|
| assert word_count(doc) >= 10000
|
| result = rewrite_document(doc, batch_paras=15)
|
| assert result.stats.batches >= 2
|
| assert result.input_words >= 10000
|
| assert result.stats.sentences > 50
|
|
|
| assert result.text.count("\n\n") >= 10
|
| assert result.stats.seconds < 180
|
|
|
|
|
| def test_bibliography_passthrough():
|
| src = (
|
| "Ram went to school yesterday happily.\n\n"
|
| "References\n\n"
|
| "Smith, J. (2020). A paper about schools. Journal of Education.\n"
|
| "Doe, A. (2019). Another citation here with many words about learning."
|
| )
|
| result = rewrite_document(src)
|
| assert "Smith, J." in result.text
|
| assert "References" in result.text
|
|
|
|
|
| def test_bibliography_ends_at_following_heading():
|
| src = (
|
| "References\n\n"
|
| "Smith, J. (2020). A paper about schools.\n\n"
|
| "# Appendix\n\n"
|
| "Ram went to school yesterday happily."
|
| )
|
| result = rewrite_document(src)
|
| appendix_sentence = next(
|
| record for record in result.sentences if record.original.startswith("Ram ")
|
| )
|
| assert appendix_sentence.status == "rewritten"
|
| assert "Yesterday, Ram happily went to school." in result.text
|
|
|
|
|
| def test_safety_rejects_entity_loss():
|
| original = "Alice met Bob yesterday at the park."
|
| bad = "Someone met a friend yesterday at the park."
|
| safety = check_safety(original, bad, use_minilm=False)
|
| assert not safety.ok
|
| assert any(r.startswith("entity:") or r == "negation" for r in safety.reasons) or (
|
| "entity" in " ".join(safety.reasons)
|
| )
|
|
|
|
|
| def test_safety_rejects_negation_flip():
|
| original = "Ram did not go to school yesterday."
|
| bad = "Ram did go to school yesterday."
|
| safety = check_safety(original, bad, use_minilm=False)
|
| assert not safety.ok
|
| assert "negation" in safety.reasons or "polarity" in safety.reasons
|
|
|
|
|
| def test_normalize_and_blocks():
|
| raw = "Hello world.\n\n# Title\n\n- item one\n\n```\ncode\n```"
|
| norm = normalize_text(raw)
|
| blocks = detect_blocks(norm)
|
| kinds = {b.kind for b in blocks}
|
| assert "heading" in kinds or "paragraph" in kinds
|
| batches = list(iter_paragraph_batches(blocks, batch_paras=2))
|
| assert batches
|
|
|
