210 lines
6.9 KiB
Python
210 lines
6.9 KiB
Python
"""llm.source — comment + corpus Doc iterators."""
|
|
|
|
from unittest.mock import patch
|
|
|
|
import pytest
|
|
|
|
from bib.item import Item, Rule
|
|
from bib.store import Store
|
|
from llm.source import (
|
|
_default_root,
|
|
comment_key_map,
|
|
iter_comment_docs,
|
|
iter_corpus_docs,
|
|
iter_rule_docs,
|
|
)
|
|
|
|
DOCKET = "CMS-2019-0111"
|
|
CID = f"{DOCKET}-0042"
|
|
|
|
COMBINED = f"""---
|
|
comment_id: {CID}
|
|
docket_id: {DOCKET}
|
|
---
|
|
|
|
We object to the E/M consolidation.
|
|
"""
|
|
|
|
|
|
@pytest.fixture
|
|
def store(tmp_path):
|
|
s = Store(":memory:", storage_dir=tmp_path / "storage")
|
|
key = s.create(
|
|
Item(
|
|
item_type="report",
|
|
title="A comment",
|
|
url=f"https://www.regulations.gov/comment/{CID}",
|
|
abstract="Inline abstract body.",
|
|
)
|
|
)
|
|
for tag in (f"docket:{DOCKET}", "doctype:comment", "year:2019"):
|
|
s.add_tag(key, tag)
|
|
s._comment_key = key # test convenience
|
|
return s
|
|
|
|
|
|
@pytest.fixture
|
|
def root(tmp_path):
|
|
d = tmp_path / DOCKET / CID
|
|
d.mkdir(parents=True)
|
|
(d / "combined.md").write_text(COMBINED)
|
|
return tmp_path
|
|
|
|
|
|
class TestCommentDocs:
|
|
def test_extracted_comment_uses_combined_body(self, store, root):
|
|
docs = list(iter_comment_docs(store, docket=DOCKET, root=root))
|
|
assert len(docs) == 1
|
|
assert docs[0].key == store._comment_key
|
|
assert "E/M consolidation" in docs[0].text
|
|
assert docs[0].metadata == {
|
|
"docket": DOCKET,
|
|
"comment_id": CID,
|
|
"doctype": "comment",
|
|
"year": "2019",
|
|
}
|
|
|
|
def test_unextracted_comment_falls_back_to_abstract(self, store, tmp_path):
|
|
docs = list(iter_comment_docs(store, docket=DOCKET, root=tmp_path))
|
|
assert docs[0].text == "Inline abstract body."
|
|
|
|
def test_docket_filter_excludes_others(self, store, root):
|
|
assert list(iter_comment_docs(store, docket="CMS-2021-0119", root=root)) == []
|
|
|
|
|
|
class TestCommentKeyMap:
|
|
def test_maps_comment_id_to_key_and_year(self, store):
|
|
mapping = comment_key_map(store, DOCKET)
|
|
assert mapping[CID] == (store._comment_key, "2019")
|
|
|
|
|
|
class TestCorpusDocs:
|
|
def test_non_comment_item_with_abstract(self, store):
|
|
key = store.create(
|
|
Item(item_type="rule", title="Final rule", abstract="Rule text.")
|
|
)
|
|
store.add_tag(key, "year:2020")
|
|
docs = list(iter_corpus_docs(store))
|
|
assert [d.key for d in docs] == [key]
|
|
assert docs[0].text == "Rule text."
|
|
assert docs[0].metadata["doctype"] == "rule"
|
|
|
|
def test_comments_excluded_from_corpus(self, store):
|
|
assert all(d.metadata["doctype"] != "comment" for d in iter_corpus_docs(store))
|
|
|
|
def test_attachment_text_included(self, store, tmp_path):
|
|
key = store.create(Item(item_type="rule", title="Rule with attachment"))
|
|
store.add_tag(key, "year:2021")
|
|
att = tmp_path / "letter.txt"
|
|
att.write_text("Attachment body text " * 10) # > 50 chars => status "ok"
|
|
store.attach_file(key, att)
|
|
docs = {d.key: d for d in iter_corpus_docs(store)}
|
|
assert "Attachment body text" in docs[key].text
|
|
|
|
def test_empty_text_item_skipped(self, store):
|
|
key = store.create(Item(item_type="rule", title="No content", abstract=" "))
|
|
store.add_tag(key, "year:2022")
|
|
assert key not in {d.key for d in iter_corpus_docs(store)}
|
|
|
|
|
|
class TestRuleDocs:
|
|
@pytest.fixture
|
|
def rule_key(self, store):
|
|
key = store.create(
|
|
Rule(
|
|
title="CY2026 PFS Proposed Rule",
|
|
document_number="2025-13271",
|
|
)
|
|
)
|
|
for tag in ("cms-rule:CMS-1832-P", "year:2026"):
|
|
store.add_tag(key, tag)
|
|
return key
|
|
|
|
def test_txt_attachment_yields_doc_with_metadata(self, store, rule_key, tmp_path):
|
|
txt = tmp_path / "2025-13271.txt"
|
|
txt.write_text(
|
|
"<html><head><title>x</title></head><body><pre>\n"
|
|
"The Secretary proposes to amend 42 CFR part 414.\n"
|
|
"</pre></body></html>"
|
|
)
|
|
store.attach_file(rule_key, txt)
|
|
|
|
docs = list(iter_rule_docs(store))
|
|
|
|
assert len(docs) == 1
|
|
doc = docs[0]
|
|
assert doc.key == rule_key
|
|
assert "The Secretary proposes to amend 42 CFR part 414." in doc.text
|
|
assert "<html>" not in doc.text
|
|
assert "<pre>" not in doc.text
|
|
assert doc.metadata == {
|
|
"doctype": "rule",
|
|
"cms_rule_id": "CMS-1832-P",
|
|
"fr_document_number": "2025-13271",
|
|
"year": "2026",
|
|
"item_key": rule_key,
|
|
}
|
|
|
|
def test_inline_html_and_entities_stripped(self, store, rule_key, tmp_path):
|
|
txt = tmp_path / "2025-13271.txt"
|
|
txt.write_text(
|
|
"<html><head><title>x</title></head><body><pre>\n"
|
|
"[[Page 12345]]\n"
|
|
'Contact us at <a href="/cdn-cgi/l/email-protection#x">'
|
|
"someone</a> regarding CMS & PFS.\n"
|
|
"A finding was significant (p<0.05 and >2 cm).\n"
|
|
"</pre></body></html>"
|
|
)
|
|
store.attach_file(rule_key, txt)
|
|
|
|
docs = list(iter_rule_docs(store))
|
|
|
|
assert len(docs) == 1
|
|
text = docs[0].text
|
|
assert "<a href=" not in text
|
|
assert "</a>" not in text
|
|
assert "someone" in text
|
|
assert "CMS & PFS" in text
|
|
assert "[[Page 12345]]" in text
|
|
assert "p<0.05 and >2 cm" in text
|
|
|
|
def test_pdf_only_falls_back_to_extract_attachment(self, store, rule_key, tmp_path):
|
|
pdf = tmp_path / "rule.pdf"
|
|
pdf.write_bytes(b"%PDF-1.4 fake")
|
|
store.attach_file(rule_key, pdf)
|
|
|
|
with patch("rex.comments.combine.extract_attachment") as mock_extract:
|
|
from rex.comments.extract import ExtractResult
|
|
|
|
mock_extract.return_value = ExtractResult(
|
|
text="Extracted PDF body.", status="ok", chars=20
|
|
)
|
|
docs = list(iter_rule_docs(store))
|
|
|
|
assert len(docs) == 1
|
|
assert docs[0].text == "Extracted PDF body."
|
|
|
|
def test_keys_filters_to_named_items(self, store, rule_key, tmp_path):
|
|
other_key = store.create(Rule(title="Other rule", document_number="2025-00001"))
|
|
store.add_tag(other_key, "cms-rule:CMS-9999-P")
|
|
for key in (rule_key, other_key):
|
|
txt = tmp_path / f"{key}.txt"
|
|
txt.write_text("<pre>Some rule text.</pre>")
|
|
store.attach_file(key, txt)
|
|
|
|
docs = list(iter_rule_docs(store, keys=(rule_key,)))
|
|
|
|
assert [d.key for d in docs] == [rule_key]
|
|
|
|
def test_non_rule_items_not_yielded(self, store):
|
|
assert all(d.metadata["doctype"] == "rule" for d in iter_rule_docs(store))
|
|
# the "report" comment item from the `store` fixture is never yielded
|
|
assert store._comment_key not in {d.key for d in iter_rule_docs(store)}
|
|
|
|
|
|
class TestDefaultRoot:
|
|
def test_default_root_under_state_comments(self):
|
|
from conf import ROOT
|
|
|
|
assert _default_root() == ROOT / ".state" / "comments"
|