Hands-on · ~30 min

Persistence: The Contract Data Model

Extend the schema with suppliers, agreements, extractions tables. Upsert extraction results idempotently per (document_id, model_version). Create agreements from MSA extractions.

What we're building

Suppliers, agreements, extractions tables in Postgres.
A function that takes an MSA extraction and writes it to the DB.
Supplier creation / lookup based on extracted party names.
Idempotency: re-extracting with the same model_version updates rather than duplicates.

SQLAlchemy models

Extend src/ci/models.py.

src/ci/models.py — additions

from __future__ import annotations
import uuid
from datetime import date, datetime
from sqlalchemy import String, Integer, DateTime, Date, Float, ForeignKey, UniqueConstraint, JSON
from sqlalchemy.orm import Mapped, mapped_column, relationship


class Supplier(Base):
    __tablename__ = "suppliers"
    __table_args__ = (UniqueConstraint("customer_id", "canonical_name"),)

    supplier_id: Mapped[str] = mapped_column(String, primary_key=True)
    customer_id: Mapped[str] = mapped_column(String, nullable=False, index=True)
    canonical_name: Mapped[str] = mapped_column(String, nullable=False)
    parent_supplier_id: Mapped[str | None] = mapped_column(ForeignKey("suppliers.supplier_id"), nullable=True)
    country: Mapped[str | None] = mapped_column(String, nullable=True)
    tax_id: Mapped[str | None] = mapped_column(String, nullable=True)
    industry: Mapped[str | None] = mapped_column(String, nullable=True)
    first_seen_at: Mapped[datetime] = mapped_column(DateTime, default=datetime.utcnow, nullable=False)


class Agreement(Base):
    __tablename__ = "agreements"

    agreement_id: Mapped[str] = mapped_column(String, primary_key=True)
    customer_id: Mapped[str] = mapped_column(String, nullable=False, index=True)
    supplier_id: Mapped[str] = mapped_column(ForeignKey("suppliers.supplier_id"), nullable=False)
    agreement_type: Mapped[str] = mapped_column(String, nullable=False)
    effective_date: Mapped[date | None] = mapped_column(Date, nullable=True)
    current_term_end: Mapped[date | None] = mapped_column(Date, nullable=True)
    parent_agreement_id: Mapped[str | None] = mapped_column(ForeignKey("agreements.agreement_id"), nullable=True)
    status: Mapped[str] = mapped_column(String, default="active", nullable=False)
    status_changed_at: Mapped[datetime] = mapped_column(DateTime, default=datetime.utcnow, nullable=False)


class Extraction(Base):
    __tablename__ = "extractions"
    __table_args__ = (UniqueConstraint("document_id", "model_version"),)

    extraction_id: Mapped[str] = mapped_column(String, primary_key=True)
    document_id: Mapped[str] = mapped_column(ForeignKey("documents.document_id"), nullable=False, index=True)
    customer_id: Mapped[str] = mapped_column(String, nullable=False, index=True)
    model_version: Mapped[str] = mapped_column(String, nullable=False)
    pipeline_version: Mapped[str] = mapped_column(String, nullable=False)
    fields: Mapped[dict] = mapped_column(JSON, nullable=False)
    qa: Mapped[dict | None] = mapped_column(JSON, nullable=True)
    extracted_at: Mapped[datetime] = mapped_column(DateTime, default=datetime.utcnow, nullable=False)


class DocumentAgreementLink(Base):
    __tablename__ = "document_agreement_links"

    document_id: Mapped[str] = mapped_column(ForeignKey("documents.document_id"), primary_key=True)
    agreement_id: Mapped[str] = mapped_column(ForeignKey("agreements.agreement_id"), primary_key=True)
    role: Mapped[str] = mapped_column(String, nullable=False)

Migration

alembic revision --autogenerate -m "suppliers, agreements, extractions"
alembic upgrade head

docker compose exec postgres psql -U ci -d contract_intel \
  -c "\dt"

Verify all five tables exist: documents, suppliers, agreements, extractions, document_agreement_links.

Persist extractions

src/ci/persistence.py

from __future__ import annotations
import uuid
from datetime import datetime

from ci.config import settings
from ci.db import db_session
from ci.models import Extraction
from ci.extraction.schemas import MSAExtraction


def persist_extraction(
    document_id: str,
    extraction: MSAExtraction,
    text_source: str,  # 'native' or 'ocr'
) -> str:
    """Upsert the extraction record. Idempotent per (document_id, model_version)."""
    extraction_id = f"ext_{document_id}_{settings.extraction_model}"

    with db_session() as s:
        existing = s.get(Extraction, extraction_id)
        fields_json = extraction.model_dump(mode="json")
        qa = {"text_source": text_source}
        if existing:
            existing.fields = fields_json
            existing.qa = qa
            existing.extracted_at = datetime.utcnow()
        else:
            s.add(Extraction(
                extraction_id=extraction_id,
                document_id=document_id,
                customer_id=settings.customer_id,
                model_version=settings.extraction_model,
                pipeline_version=settings.pipeline_version,
                fields=fields_json,
                qa=qa,
            ))
    return extraction_id

Create agreements

From an MSA extraction, create or update the corresponding agreement and link the document.

src/ci/persistence.py — additions

from datetime import date, timedelta
from sqlalchemy import select

from ci.models import Supplier, Agreement, DocumentAgreementLink


def normalize_supplier_name(name: str) -> str:
    s = " ".join(name.split())
    for suffix in (", Inc.", " Inc.", ", LLC", " LLC", ", Ltd", " Ltd", " Corp.", ", Corp."):
        if s.lower().endswith(suffix.lower()):
            s = s[:-len(suffix)].strip()
            break
    return s


def get_or_create_supplier(session, customer_id: str, raw_name: str) -> str:
    canonical = normalize_supplier_name(raw_name)
    existing = session.scalar(
        select(Supplier)
        .where(Supplier.customer_id == customer_id, Supplier.canonical_name == canonical)
    )
    if existing:
        return existing.supplier_id
    supplier_id = f"sup_{customer_id[:8]}_{uuid.uuid4().hex[:12]}"
    session.add(Supplier(
        supplier_id=supplier_id,
        customer_id=customer_id,
        canonical_name=canonical,
    ))
    return supplier_id


def create_agreement_from_extraction(
    document_id: str,
    extraction: MSAExtraction,
    customer_party_index: int = 1,  # which entry in parties is the customer
) -> str:
    """Create an agreement row from an MSA extraction. Idempotent by document_id."""
    with db_session() as s:
        # Skip if already linked
        existing_link = s.execute(
            select(DocumentAgreementLink)
            .where(DocumentAgreementLink.document_id == document_id,
                   DocumentAgreementLink.role == "master")
        ).first()
        if existing_link:
            return existing_link[0].agreement_id

        # Pick the supplier (the non-customer party)
        if not extraction.parties or len(extraction.parties) < 2:
            raise ValueError(f"MSA {document_id} doesn't have ≥2 parties")
        supplier_raw = extraction.parties[1 - customer_party_index]
        supplier_id = get_or_create_supplier(s, settings.customer_id, supplier_raw)

        # Compute current_term_end
        eff = extraction.effective_date.value
        term_months = extraction.term_length_months.value
        current_term_end = None
        if eff and term_months:
            try:
                eff_date = date.fromisoformat(eff)
                # Approximate: 30 days per month
                current_term_end = eff_date + timedelta(days=int(term_months) * 30)
            except (ValueError, TypeError):
                pass

        agreement_id = f"agr_{settings.customer_id[:8]}_{uuid.uuid4().hex[:12]}"
        s.add(Agreement(
            agreement_id=agreement_id,
            customer_id=settings.customer_id,
            supplier_id=supplier_id,
            agreement_type="master_service_agreement",
            effective_date=date.fromisoformat(eff) if eff else None,
            current_term_end=current_term_end,
            status="active",
        ))
        s.add(DocumentAgreementLink(
            document_id=document_id,
            agreement_id=agreement_id,
            role="master",
        ))
    return agreement_id

CLI

Add an end-to-end command: extract + persist + create agreement.

src/ci/cli.py — process command

from ci.extraction import extract_document
from ci.persistence import persist_extraction, create_agreement_from_extraction


@app.command()
def process(document_id: str | None = None) -> None:
    """Extract + persist + create agreement for one or all MSAs."""
    with db_session() as s:
        if document_id:
            doc_ids = [document_id]
        else:
            doc_ids = list(s.scalars(
                select(Document.document_id).where(Document.doc_class == "master_service_agreement")
            ))

    for doc_id in doc_ids:
        try:
            extraction = extract_document(doc_id)
            ext_id = persist_extraction(doc_id, extraction, text_source="native")  # simplified
            agreement_id = create_agreement_from_extraction(doc_id, extraction)
            console.print(f"[green]✓[/green] {doc_id} → ext={ext_id}, agreement={agreement_id}")
        except Exception as e:
            console.print(f"[red]✗[/red] {doc_id}: {e}")

Verify

ci process
# Each MSA: extract, persist extraction, create supplier + agreement

Check the database:

-- Suppliers
SELECT supplier_id, canonical_name FROM suppliers;

-- Agreements with their suppliers
SELECT a.agreement_id, s.canonical_name AS supplier, a.effective_date, a.current_term_end
FROM agreements a
JOIN suppliers s ON s.supplier_id = a.supplier_id;

-- Extraction confidence per field
SELECT
  document_id,
  fields->'effective_date'->>'value' AS effective_date,
  fields->'effective_date'->>'confidence' AS conf
FROM extractions;

You should see a complete view: documents → extractions → agreements → suppliers, all linked via foreign keys. Re-run ci process — no duplicates; extractions update in place.

✓ Check yourself

Answer from memory before opening each one.

You change CI_EXTRACTION_MODEL to a newer model and run ci process again. How many extraction rows does each document have afterwards, and do the agreements change?

Show answer
Each document has two rows, because the extraction id and the unique constraint include the model. Agreements do not change: the link already exists, so create_agreement_from_extraction returns early and never reads the new extraction. Decide which extraction feeds the agreement. See Persist extractions.
A 12-month MSA is effective 2024-03-15. What current_term_end does the code store, and why does the difference matter for a 60-day notice window?

Show answer
It adds 12 × 30 = 360 days and stores 2025-03-10, which is five days before the true anniversary. A 36-month term lands about 15 days early. A renewal alert computed from that date can fire late for a notice deadline. Add calendar months, for example with dateutil.relativedelta. See Create agreements.
Three MSAs name the same vendor "Acme Corp.", "ACME Corp" and "Acme Corporation, Inc.". How many supplier rows does ci process create, and what do you do about it?

Show answer
Three. The suffix strip turns the first into "Acme", leaves the second as "ACME Corp" and turns the third into "Acme Corporation", and the unique constraint compares them exactly. String clean-up cannot settle identity. Add case folding, an alias table, a tax id when present and a review step. Merging too eagerly joins different companies, so a missed merge is the safer error. See Create agreements.
create_agreement_from_extraction adds the supplier, the agreement and the link inside one db_session. What goes wrong if they were committed separately and the process died after the agreement?

Show answer
An agreement would exist with no link. The idempotency check looks for the link, so the next run would not find it and would create a second agreement for the same document. One session means one commit, so the three rows exist together or not at all. See Create agreements.
After ci process, the suppliers table contains your own customer's name, and the agreements point at it. What assumption in the code failed?

Show answer
The function treats parties[1] as the customer and takes parties[0] as the supplier. The prompt never orders the parties, and contracts list either side first. Ask the model to label each party's role, or match against known customer names, and confirm with the query in the verify section. See Create agreements.

Takeaway

An extraction is keyed by document and model version, so a model upgrade adds rows beside the old ones instead of overwriting them.
An upsert is idempotent only if its key says what you mean by "the same", and the agreement's key is its link to the document.
Writing supplier, agreement and link in one transaction keeps the idempotency check honest after a crash.
Thirty-day months drift from the calendar and put term ends days early, so compute them with calendar arithmetic.
Normalising a company name is not entity resolution: an exact-match constraint turns every spelling into a new supplier.
Which party is the customer is a fact to extract or look up, not a position in a list.

The data is stored and linked, but you cannot yet say how often the extraction is right, so the next step builds the scoreboard. → Evaluation Harness