Build polymorphic file/attachment + embedding models

Loading diagram...

• Line 1: Declares this as a Mermaid ER (Entity-Relationship) diagram definition.
• Line 2: Defines a one-to-many relationship where one CONVERSATION contains zero or more MESSAGE entities.
• Line 3: Defines a one-to-many relationship where one MESSAGE has zero or more ATTACHMENT entities.
• Line 4: Defines a one-to-many relationship where one ATTACHMENT produces zero or more EMBEDDING_RECORD entities.
• Lines 5-6: Opens the ATTACHMENT entity definition and declares id as a UUID primary key.
• Line 7: Declares message_id as a UUID foreign key linking back to the parent MESSAGE.
• Lines 8-10: Defines file metadata fields: filename (string name of the uploaded file), content_type (enum for MIME type), and file_size_bytes (bigint storing the file size).
• Line 11: Declares storage_path as a string holding the location of the file in object/blob storage.
• Lines 12-13: Defines two enum status-tracking fields: upload_status (tracks whether the file upload succeeded) and processing_status (tracks the state of downstream processing like chunking/embedding).
• Lines 14-15: Declares type_metadata as a JSONB column for flexible, type-specific metadata and created_at as a timestamp for record creation time.
• Line 16: Closes the ATTACHMENT entity definition block.
• Lines 17-19: Opens the EMBEDDING_RECORD entity definition, declares id as a UUID primary key, and attachment_id as a UUID foreign key linking back to the parent ATTACHMENT.
• Line 20: Declares tenant_id as a UUID identifying which tenant owns this embedding, enabling multi-tenant isolation.
• Lines 21-22: Defines chunk_index (integer position of the chunk within the source document) and chunk_text_preview (a string preview of the chunk's text content).
• Lines 23-24: Declares start_offset and end_offset as integers marking the character boundaries of the chunk within the original attachment content.
• Lines 25-26: Defines embedding_model (string identifying which model generated the vector) and dimensions (integer specifying the dimensionality of the embedding vector).
• Line 27: Declares vector_store_id as a string referencing the external vector store where the embedding is persisted.
• Lines 28-29: Defines status as an enum tracking the embedding lifecycle state and created_at as a timestamp for record creation.
• Line 30: Closes the EMBEDDING_RECORD entity definition block.

1import enum
2import uuid
3from datetime import datetime
4from typing import Optional
5
6from sqlalchemy import ForeignKey, Index, text
7from sqlalchemy.dialects.postgresql import JSONB, UUID
8from sqlalchemy.orm import Mapped, mapped_column, relationship
9
10from app.db.base import Base
11
12class ContentType(enum.Enum):
13    PDF = "pdf"
14    DOCX = "docx"
15    IMAGE = "image"
16    CSV = "csv"
17    URL = "url"
18
19class UploadStatus(enum.Enum):
20    PENDING = "pending"
21    PROCESSING = "processing"
22    COMPLETED = "completed"
23    FAILED = "failed"
24
25class Attachment(Base):
26    __tablename__ = "attachments"
27
28    id: Mapped[uuid.UUID] = mapped_column(
29        UUID(as_uuid=True), primary_key=True, server_default=text("gen_random_uuid()")
30    )
31    message_id: Mapped[uuid.UUID] = mapped_column(
32        ForeignKey("messages.id", ondelete="CASCADE"), index=True
33    )
34    tenant_id: Mapped[uuid.UUID] = mapped_column(UUID(as_uuid=True), index=True)
35    filename: Mapped[str] = mapped_column(nullable=False)
36    content_type: Mapped[ContentType] = mapped_column(nullable=False)
37    file_size_bytes: Mapped[int] = mapped_column(nullable=False)
38    storage_path: Mapped[str] = mapped_column(nullable=False)
39    upload_status: Mapped[UploadStatus] = mapped_column(
40        default=UploadStatus.PENDING
41    )
42    processing_status: Mapped[str] = mapped_column(default="queued")
43    type_metadata: Mapped[Optional[dict]] = mapped_column(JSONB, default=dict)
44    created_at: Mapped[datetime] = mapped_column(server_default=text("now()"))
45
46    embeddings: Mapped[list["EmbeddingRecord"]] = relationship(
47        back_populates="attachment", cascade="all, delete-orphan"
48    )
49
50    __mapper_args__ = {
51        "polymorphic_on": "content_type",
52        "polymorphic_identity": None,
53    }
54    __table_args__ = (
55        Index("ix_attach_tenant_created", "tenant_id", "created_at"),
56    )
57
58class PDFAttachment(Attachment):
59    __mapper_args__ = {"polymorphic_identity": ContentType.PDF}
60
61    @property
62    def page_count(self) -> Optional[int]:
63        return (self.type_metadata or {}).get("page_count")
64
65    @page_count.setter
66    def page_count(self, value: int) -> None:
67        if self.type_metadata is None:
68            self.type_metadata = {}
69        self.type_metadata["page_count"] = value
70
71class ImageAttachment(Attachment):
72    __mapper_args__ = {"polymorphic_identity": ContentType.IMAGE}
73
74    @property
75    def dimensions(self) -> Optional[tuple[int, int]]:
76        meta = self.type_metadata or {}
77        if "width" in meta and "height" in meta:
78            return (meta["width"], meta["height"])
79        return None
80
81class URLAttachment(Attachment):
82    __mapper_args__ = {"polymorphic_identity": ContentType.URL}
83
84    @property
85    def resolved_url(self) -> Optional[str]:
86        return (self.type_metadata or {}).get("resolved_url")

• Lines 1-11: Import foundation modules — enum and uuid from the standard library, SQLAlchemy's ForeignKey, Index, and PostgreSQL-specific JSONB and UUID dialect types, plus the 2.0-style Mapped and mapped_column annotation helpers.
• Lines 14-23: Define two enumerations. ContentType lists the five supported file categories that serve as polymorphic identities. UploadStatus tracks the lifecycle of the file transfer from client to cloud storage.
• Lines 26-50: Declare the base Attachment model. The id column uses PostgreSQL's gen_random_uuid() server default so UUIDs are generated database-side. The message_id foreign key cascades deletes so removing a message removes all its attachments. The type_metadata JSONB column stores content-type-specific fields without adding nullable columns to the physical table.
• Lines 52-55: The embeddings relationship uses cascade="all, delete-orphan" so that deleting an attachment also removes its embedding records, keeping the vector store reference table consistent.
• Lines 57-61: __mapper_args__ configures single-table polymorphism. The polymorphic_on directive points to the content_type column, and the base identity is None as a fallback for any unrecognized types.
• Lines 62-64: A composite index on tenant_id and created_at accelerates the most common query pattern — listing a tenant's recent attachments.
• Lines 67-77: PDFAttachment inherits the same physical table but sets its polymorphic identity to ContentType.PDF. The page_count property reads from and writes to type_metadata, providing a typed Python interface over the JSONB field. The setter initializes the dictionary if type_metadata is None.
• Lines 80-87: ImageAttachment exposes a dimensions property that returns a tuple of width and height, or None if the metadata has not been populated yet.
• Lines 81-86: URLAttachment surfaces the resolved_url after redirect resolution, useful for link-preview generation and deduplication.

1class EmbeddingStatus(enum.Enum):
2    PENDING = "pending"
3    EMBEDDED = "embedded"
4    FAILED = "failed"
5    EXPIRED = "expired"
6
7class EmbeddingRecord(Base):
8    __tablename__ = "embedding_records"
9
10    id: Mapped[uuid.UUID] = mapped_column(
11        UUID(as_uuid=True), primary_key=True, server_default=text("gen_random_uuid()")
12    )
13    attachment_id: Mapped[uuid.UUID] = mapped_column(
14        ForeignKey("attachments.id", ondelete="CASCADE"), index=True
15    )
16    tenant_id: Mapped[uuid.UUID] = mapped_column(UUID(as_uuid=True), index=True)
17    chunk_index: Mapped[int] = mapped_column(nullable=False)
18    chunk_text_preview: Mapped[str] = mapped_column(nullable=False)
19    start_offset: Mapped[int] = mapped_column(nullable=False)
20    end_offset: Mapped[int] = mapped_column(nullable=False)
21    embedding_model: Mapped[str] = mapped_column(nullable=False)
22    dimensions: Mapped[int] = mapped_column(nullable=False)
23    vector_store_id: Mapped[Optional[str]] = mapped_column(nullable=True)
24    status: Mapped[EmbeddingStatus] = mapped_column(
25        default=EmbeddingStatus.PENDING
26    )
27    token_count: Mapped[Optional[int]] = mapped_column(nullable=True)
28    created_at: Mapped[datetime] = mapped_column(server_default=text("now()"))
29
30    attachment: Mapped["Attachment"] = relationship(back_populates="embeddings")
31
32    __table_args__ = (
33        Index("ix_embed_tenant_model", "tenant_id", "embedding_model"),
34        Index("ix_embed_attachment_chunk", "attachment_id", "chunk_index", unique=True),
35    )
36
37async def get_attachments_for_message(
38    session: AsyncSession, message_id: uuid.UUID
39) -> list[Attachment]:
40    stmt = (
41        select(Attachment)
42        .where(Attachment.message_id == message_id)
43        .options(selectinload(Attachment.embeddings))
44        .order_by(Attachment.created_at)
45    )
46    result = await session.execute(stmt)
47    return list(result.scalars().all())
48
49async def get_pdf_attachments_pending_processing(
50    session: AsyncSession, tenant_id: uuid.UUID
51) -> list[PDFAttachment]:
52    stmt = (
53        select(PDFAttachment)
54        .where(
55            PDFAttachment.tenant_id == tenant_id,
56            PDFAttachment.processing_status == "queued",
57        )
58        .order_by(PDFAttachment.created_at)
59    )
60    result = await session.execute(stmt)
61    return list(result.scalars().all())

• Lines 1-6: EmbeddingStatus tracks the embedding lifecycle. The EXPIRED state marks records whose vectors were generated by a deprecated model version and need re-embedding.
• Lines 9-12: The primary key follows the same UUID pattern as Attachment, ensuring globally unique identifiers across a distributed system.
• Lines 14-16: The attachment_id foreign key cascades deletes, so purging an attachment also removes all its embedding metadata. The index accelerates chunk lookups for a specific attachment.
• Lines 17-18: A denormalized tenant_id column avoids a join through attachments when filtering embeddings by tenant — a critical optimization for vector search pre-filtering where you need to pass tenant-scoped IDs to the vector store.
• Lines 19-22: chunk_index provides deterministic ordering, while start_offset and end_offset record byte positions in the original extracted text so you can highlight the exact source passage in citations.
• Lines 23-25: embedding_model (e.g., "text-embedding-3-large") and dimensions (e.g., 3072) let you query which records need re-embedding after a model upgrade. The vector_store_id is nullable because it is populated asynchronously after the embedding API call completes.
• Lines 26-28: The status column defaults to PENDING, transitioning to EMBEDDED once the vector store confirms insertion, or FAILED if the embedding API returns an error.
• Lines 29-30: token_count tracks per-chunk token usage for cost attribution, nullable because it may not be available from all embedding providers.
• Lines 32: The attachment relationship provides the reverse navigation path for ORM-level joins and eager loading.
• Lines 34-37: Two composite indexes optimize the dominant access patterns. The (tenant_id, embedding_model) index powers re-embedding queries that target a specific model version within a tenant. The (attachment_id, chunk_index) unique index enforces that no attachment can have duplicate chunk numbers, preventing double-insertion bugs in retry-heavy async pipelines.