Persistence: The Contract Data Model

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)

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

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

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

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

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}")

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

-- 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;