Ingestion: Folder Watcher to Postgres

from __future__ import annotations
from datetime import datetime
from sqlalchemy import String, Integer, DateTime, UniqueConstraint
from sqlalchemy.orm import DeclarativeBase, Mapped, mapped_column


class Base(DeclarativeBase):
    pass


class Document(Base):
    __tablename__ = "documents"
    __table_args__ = (UniqueConstraint("customer_id", "file_hash"),)

    document_id: Mapped[str] = mapped_column(String, primary_key=True)
    customer_id: Mapped[str] = mapped_column(String, nullable=False, index=True)
    source_uri: Mapped[str] = mapped_column(String, nullable=False)
    file_hash: Mapped[str] = mapped_column(String, nullable=False)
    doc_class: Mapped[str] = mapped_column(String, nullable=False, index=True)
    page_count: Mapped[int] = mapped_column(Integer, nullable=False)
    ingested_at: Mapped[datetime] = mapped_column(DateTime, default=datetime.utcnow, nullable=False)

from collections.abc import Iterator
from contextlib import contextmanager
from sqlalchemy import create_engine
from sqlalchemy.orm import Session, sessionmaker

from ci.config import settings

engine = create_engine(settings.db_dsn, future=True)
SessionLocal = sessionmaker(bind=engine, expire_on_commit=False)


@contextmanager
def db_session() -> Iterator[Session]:
    s = SessionLocal()
    try:
        yield s
        s.commit()
    except Exception:
        s.rollback()
        raise
    finally:
        s.close()

alembic init -t async alembic   # creates alembic/ directory
# Edit alembic.ini: set sqlalchemy.url = (your db_dsn)
# Edit alembic/env.py: import Base and set target_metadata = Base.metadata

alembic revision --autogenerate -m "documents table"
alembic upgrade head

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

from __future__ import annotations
import hashlib
from pathlib import Path
import pdfplumber

from ci.config import settings
from ci.db import db_session
from ci.models import Document


def derive_document_id(file_bytes: bytes, customer_id: str) -> tuple[str, str]:
    h = hashlib.sha256(file_bytes).hexdigest()
    # Hash the whole customer_id with the content, so tenants never share an id
    scoped = hashlib.sha256(f"{customer_id}:{h}".encode()).hexdigest()
    document_id = f"doc_{customer_id[:8]}_{scoped[:16]}"
    return document_id, h


def classify_doc(file_bytes: bytes, filename: str) -> str:
    """Heuristic classification by filename. Replace with a real classifier later."""
    name = filename.lower()
    if "msa" in name or "master_service" in name:
        return "master_service_agreement"
    if "sow" in name or "statement_of_work" in name:
        return "statement_of_work"
    if "amendment" in name:
        return "amendment"
    if "nda" in name:
        return "nda"
    if "invoice" in name:
        return "invoice"
    return "unclassified"


def page_count(file_bytes: bytes) -> int:
    import io
    with pdfplumber.open(io.BytesIO(file_bytes)) as pdf:
        return len(pdf.pages)


def ingest_file(path: Path, customer_id: str | None = None) -> dict:
    customer_id = customer_id or settings.customer_id
    file_bytes = path.read_bytes()
    document_id, file_hash = derive_document_id(file_bytes, customer_id)
    doc_class = classify_doc(file_bytes, path.name)
    pages = page_count(file_bytes)

    with db_session() as s:
        existing = s.get(Document, document_id)
        if existing:
            return {"document_id": document_id, "status": "exists", "doc_class": existing.doc_class}
        doc = Document(
            document_id=document_id,
            customer_id=customer_id,
            source_uri=str(path.resolve()),
            file_hash=file_hash,
            doc_class=doc_class,
            page_count=pages,
        )
        s.add(doc)
    return {"document_id": document_id, "status": "ingested", "doc_class": doc_class, "page_count": pages}

from __future__ import annotations
from pathlib import Path
import typer
from rich.console import Console
from rich.table import Table

from ci.config import settings
from ci.ingestion import ingest_file

app = typer.Typer()
console = Console()


@app.command()
def ingest(directory: Path = Path(settings.inbox_dir)) -> None:
    """Ingest all PDFs in DIRECTORY."""
    pdfs = sorted(directory.glob("*.pdf"))
    if not pdfs:
        console.print(f"[yellow]No PDFs in {directory}[/yellow]")
        return

    table = Table(title=f"Ingested from {directory}")
    table.add_column("File")
    table.add_column("Status")
    table.add_column("doc_class")
    table.add_column("pages", justify="right")

    for path in pdfs:
        result = ingest_file(path)
        table.add_row(path.name, result["status"],
                      result.get("doc_class", "-"),
                      str(result.get("page_count", "-")))
    console.print(table)


if __name__ == "__main__":
    app()

[project.scripts]
ci = "ci.cli:app"

ci ingest
# Should show: 2 PDFs, status=ingested, doc_class=master_service_agreement

ci ingest
# Re-run: 2 PDFs, status=exists. No duplicates.

docker compose exec postgres psql -U ci -d contract_intel -c "SELECT document_id, doc_class, page_count FROM documents;"

from pathlib import Path
from ci.ingestion import derive_document_id, ingest_file


def test_derive_document_id_deterministic():
    content = b"hello world"
    id1, h1 = derive_document_id(content, "cust_test")
    id2, h2 = derive_document_id(content, "cust_test")
    assert id1 == id2
    assert h1 == h2


def test_derive_document_id_customer_specific():
    content = b"hello world"
    id1, _ = derive_document_id(content, "cust_a")
    id2, _ = derive_document_id(content, "cust_b")
    assert id1 != id2


def test_ingest_is_idempotent(tmp_path: Path):
    pdf = tmp_path / "test.pdf"
    pdf.write_bytes(b"%PDF-1.4\n%%EOF")  # not a real PDF — pdfplumber will fail on page_count
    # In practice, use a real sample PDF in test fixtures.