59 lines
1.8 KiB
Python
59 lines
1.8 KiB
Python
from sqlalchemy.dialects import postgresql
|
|
|
|
from pipelines.jobs.summarize_bills import (
|
|
create_select_bill_texts_for_summarization,
|
|
store_bill_summary_result,
|
|
)
|
|
from pipelines.orm.data_science_dev.congress import BillText, BillTextSummary
|
|
|
|
|
|
class FakeSession:
|
|
def __init__(self) -> None:
|
|
self.added: list[object] = []
|
|
|
|
def add(self, value: object) -> None:
|
|
self.added.append(value)
|
|
|
|
|
|
def _compile_sql(statement: object) -> str:
|
|
return str(
|
|
statement.compile(
|
|
dialect=postgresql.dialect(),
|
|
compile_kwargs={"literal_binds": True},
|
|
)
|
|
)
|
|
|
|
|
|
def test_store_bill_summary_result_creates_summary_row() -> None:
|
|
session = FakeSession()
|
|
bill_text = BillText(id=10, bill_id=5, version_code="ih")
|
|
|
|
summary_row = store_bill_summary_result(
|
|
session=session,
|
|
bill_text=bill_text,
|
|
summary="A summary",
|
|
model="gpt-5.4-mini",
|
|
)
|
|
|
|
assert session.added == [summary_row]
|
|
assert isinstance(summary_row, BillTextSummary)
|
|
assert summary_row.bill_text is bill_text
|
|
assert summary_row.summary == "A summary"
|
|
assert summary_row.summarization_model == "gpt-5.4-mini"
|
|
assert summary_row.summarization_system_prompt_version == "v1.2"
|
|
assert summary_row.summarization_user_prompt_version == "v1"
|
|
|
|
|
|
def test_create_select_bill_texts_for_summarization_excludes_existing_summaries() -> None:
|
|
sql = _compile_sql(create_select_bill_texts_for_summarization(force=False))
|
|
|
|
assert "bill_text_summary" in sql
|
|
assert "NOT (EXISTS" in sql or "NOT EXISTS" in sql
|
|
assert "bill_text.summary" not in sql
|
|
|
|
|
|
def test_create_select_bill_texts_for_summarization_force_skips_summary_filter() -> None:
|
|
sql = _compile_sql(create_select_bill_texts_for_summarization(force=True))
|
|
|
|
assert "bill_text_summary" not in sql
|