weave/pipelines/orm/data_science_dev/congress/bill.py

"""Bill models for legislation, official actions, text versions, and topic tags."""

from __future__ import annotations

from datetime import date, datetime
from enum import StrEnum
from typing import TYPE_CHECKING

from sqlalchemy import DateTime, Enum, ForeignKey, Index, UniqueConstraint
from sqlalchemy.orm import Mapped, mapped_column, relationship

from pipelines.orm.data_science_dev.base import DataScienceDevTableBase

if TYPE_CHECKING:
    from pipelines.orm.data_science_dev.congress.context import VoteMeasureLink


class BillTopicPosition(StrEnum):
    """Whether a yes vote on a bill is for or against a topic."""

    FOR = "for"
    AGAINST = "against"


def _enum_column(enum_cls: type[StrEnum], *, name: str) -> Enum:
    """Build a portable SQLAlchemy enum column for StrEnum values."""

    return Enum(
        enum_cls,
        values_callable=lambda enum_type: [member.value for member in enum_type],
        native_enum=False,
        name=name,
    )


class Bill(DataScienceDevTableBase):
    """Legislation with congress number, type, titles, status, and sponsor."""

    __tablename__ = "bill"
    __table_args__ = (
        UniqueConstraint(
            "congress", "bill_type", "number", name="uq_bill_congress_type_number"
        ),
        Index("ix_bill_congress", "congress"),
    )

    congress: Mapped[int]
    bill_type: Mapped[str]
    number: Mapped[int]

    title: Mapped[str | None]
    title_short: Mapped[str | None]
    official_title: Mapped[str | None]

    status: Mapped[str | None]
    status_at: Mapped[date | None]

    sponsor_bioguide_id: Mapped[str | None]

    subjects_top_term: Mapped[str | None]
    score_processed_at: Mapped[datetime | None] = mapped_column(DateTime(timezone=True))

    bill_texts: Mapped[list[BillText]] = relationship(
        "BillText",
        back_populates="bill",
        cascade="all, delete-orphan",
    )
    topics: Mapped[list[BillTopic]] = relationship(
        "BillTopic",
        back_populates="bill",
        cascade="all, delete-orphan",
    )
    bill_actions: Mapped[list[BillAction]] = relationship(
        "BillAction",
        back_populates="bill",
        cascade="all, delete-orphan",
    )
    outgoing_bill_relations: Mapped[list[BillRelation]] = relationship(
        "BillRelation",
        foreign_keys="BillRelation.bill_id",
        back_populates="bill",
        cascade="all, delete-orphan",
    )
    incoming_bill_relations: Mapped[list[BillRelation]] = relationship(
        "BillRelation",
        foreign_keys="BillRelation.related_bill_id",
        back_populates="related_bill",
        cascade="all, delete-orphan",
    )
    vote_measure_links: Mapped[list[VoteMeasureLink]] = relationship(
        "VoteMeasureLink",
        back_populates="measure",
        cascade="all, delete-orphan",
    )


class BillText(DataScienceDevTableBase):
    """Stores different text versions of a bill (introduced, enrolled, etc.)."""

    __tablename__ = "bill_text"
    __table_args__ = (
        UniqueConstraint(
            "bill_id", "version_code", name="uq_bill_text_bill_id_version_code"
        ),
    )

    bill_id: Mapped[int] = mapped_column(ForeignKey("main.bill.id", ondelete="CASCADE"))
    primary_summary_id: Mapped[int | None] = mapped_column(
        ForeignKey("main.bill_text_summary.id", ondelete="SET NULL")
    )
    version_code: Mapped[str]
    version_name: Mapped[str | None]
    text_content: Mapped[str | None]
    date: Mapped[date | None]
    source_datetime_raw: Mapped[str | None]
    text_url_xml: Mapped[str | None]
    text_url_pdf: Mapped[str | None]
    text_url_html: Mapped[str | None]
    source_artifact_id: Mapped[int | None] = mapped_column(
        ForeignKey("main.source_artifact.id", ondelete="SET NULL")
    )

    bill: Mapped[Bill] = relationship("Bill", back_populates="bill_texts")
    summaries: Mapped[list[BillTextSummary]] = relationship(
        "BillTextSummary",
        back_populates="bill_text",
        cascade="all, delete-orphan",
        foreign_keys="BillTextSummary.bill_text_id",
        order_by=lambda: (
            BillTextSummary.created.desc(),
            BillTextSummary.id.desc(),
        ),
    )
    primary_summary: Mapped[BillTextSummary | None] = relationship(
        "BillTextSummary",
        foreign_keys=[primary_summary_id],
        post_update=True,
    )

    def latest_summary(self) -> BillTextSummary | None:
        """Return the newest summary row for this bill text."""
        return self.summaries[0] if self.summaries else None

    def default_summary(self) -> BillTextSummary | None:
        """Return the primary summary when set, otherwise the newest summary."""
        return self.primary_summary or self.latest_summary()


class BillTextSummary(DataScienceDevTableBase):
    """Stores one generated summary for a bill text version."""

    __tablename__ = "bill_text_summary"
    __table_args__ = (
        Index("ix_bill_text_summary_bill_text_id", "bill_text_id"),
        Index(
            "ix_bill_text_summary_bill_text_id_created",
            "bill_text_id",
            "created",
        ),
    )

    bill_text_id: Mapped[int] = mapped_column(
        ForeignKey("main.bill_text.id", ondelete="CASCADE")
    )
    summary: Mapped[str]
    summarization_model: Mapped[str | None]
    summarization_user_prompt_version: Mapped[str | None]
    summarization_system_prompt_version: Mapped[str | None]

    bill_text: Mapped[BillText] = relationship(
        "BillText",
        back_populates="summaries",
        foreign_keys=[bill_text_id],
    )


class BillAction(DataScienceDevTableBase):
    """Official action row from Bill Status XML."""

    __tablename__ = "bill_action"
    __table_args__ = (
        UniqueConstraint("bill_id", "sequence", name="uq_bill_action_bill_id_sequence"),
    )

    bill_id: Mapped[int] = mapped_column(ForeignKey("main.bill.id", ondelete="CASCADE"))
    sequence: Mapped[int]
    action_date: Mapped[date]
    action_time: Mapped[str | None]
    action_text: Mapped[str]
    action_type: Mapped[str | None]
    action_code: Mapped[str | None]
    source_system_code: Mapped[str | None]
    source_system_name: Mapped[str | None]
    source_artifact_id: Mapped[int | None] = mapped_column(
        ForeignKey("main.source_artifact.id", ondelete="SET NULL")
    )

    bill: Mapped[Bill] = relationship("Bill", back_populates="bill_actions")
    recorded_votes: Mapped[list[BillActionRecordedVote]] = relationship(
        "BillActionRecordedVote",
        back_populates="bill_action",
        cascade="all, delete-orphan",
    )


class BillActionRecordedVote(DataScienceDevTableBase):
    """Recorded vote nested under one official bill action."""

    __tablename__ = "bill_action_recorded_vote"
    __table_args__ = (
        UniqueConstraint(
            "bill_action_id",
            "congress",
            "chamber",
            "session_number",
            "roll_number",
            name="uq_bill_action_recorded_vote_match_key",
        ),
        Index(
            "ix_bill_action_recorded_vote_match_tuple",
            "congress",
            "chamber",
            "session_number",
            "roll_number",
        ),
    )

    bill_action_id: Mapped[int] = mapped_column(
        ForeignKey("main.bill_action.id", ondelete="CASCADE")
    )
    congress: Mapped[int]
    chamber: Mapped[str]
    session_number: Mapped[int]
    roll_number: Mapped[int]
    vote_datetime: Mapped[datetime | None] = mapped_column(DateTime(timezone=True))
    vote_url: Mapped[str | None]

    bill_action: Mapped[BillAction] = relationship(
        "BillAction",
        back_populates="recorded_votes",
    )


class BillRelation(DataScienceDevTableBase):
    """Relationship between one bill/resolution and another."""

    __tablename__ = "bill_relation"
    __table_args__ = (
        Index("ix_bill_relation_bill_id", "bill_id"),
        Index("ix_bill_relation_related_bill_id", "related_bill_id"),
    )

    bill_id: Mapped[int] = mapped_column(ForeignKey("main.bill.id", ondelete="CASCADE"))
    related_bill_id: Mapped[int] = mapped_column(
        ForeignKey("main.bill.id", ondelete="CASCADE")
    )
    relationship_type: Mapped[str]
    identified_by: Mapped[str | None]
    latest_action_date: Mapped[date | None]
    latest_action_text: Mapped[str | None]

    bill: Mapped[Bill] = relationship(
        "Bill",
        foreign_keys=[bill_id],
        back_populates="outgoing_bill_relations",
    )
    related_bill: Mapped[Bill] = relationship(
        "Bill",
        foreign_keys=[related_bill_id],
        back_populates="incoming_bill_relations",
    )


class BillTopic(DataScienceDevTableBase):
    """One bill stance on one topic used to score roll-call votes."""

    __tablename__ = "bill_topic"
    __table_args__ = (
        UniqueConstraint(
            "bill_id",
            "topic",
            "support_position",
            name="uq_bill_topic_bill_id_topic_support_position",
        ),
        Index("ix_bill_topic_topic", "topic"),
    )

    bill_id: Mapped[int] = mapped_column(ForeignKey("main.bill.id", ondelete="CASCADE"))
    topic: Mapped[str]
    support_position: Mapped[BillTopicPosition] = mapped_column(
        _enum_column(BillTopicPosition, name="bill_topic_position")
    )

    bill: Mapped[Bill] = relationship("Bill", back_populates="topics")