modcomms/backend/app/models/models.py
Vadym Samoilenko f520aba397 Fix KB distillation fallback and SpecVersion type annotation
- knowledge_base_service.py: wrap Gemini distillation call in try/except
  to fall back to fallback_client/fallback_model if primary times out,
  matching the fallback behaviour in GeminiService._generate_content()

- models.py: fix SpecVersion.source_document_ids ORM type annotation from
  Mapped[Optional[dict]] to Mapped[Optional[list]] — the field stores a
  JSON array of document ID strings, not an object

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
2026-03-16 14:03:25 +00:00

307 lines
17 KiB
Python
Executable file

import uuid
from datetime import datetime
from typing import Optional
from sqlalchemy import Boolean, DateTime, ForeignKey, Integer, String, Text, UniqueConstraint, func
from sqlalchemy.dialects.postgresql import JSONB, UUID
from sqlalchemy.orm import Mapped, mapped_column, relationship
from app.models.database import Base
class Agency(Base):
"""Agency/organization that users belong to."""
__tablename__ = "agencies"
id: Mapped[uuid.UUID] = mapped_column(UUID(as_uuid=True), primary_key=True, default=uuid.uuid4)
name: Mapped[str] = mapped_column(String(255), nullable=False, unique=True)
created_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), server_default=func.now())
# Relationships
users: Mapped[list["User"]] = relationship("User", back_populates="agency")
campaigns: Mapped[list["Campaign"]] = relationship("Campaign", back_populates="agency")
class User(Base):
"""User account linked to Azure AD."""
__tablename__ = "users"
id: Mapped[uuid.UUID] = mapped_column(UUID(as_uuid=True), primary_key=True, default=uuid.uuid4)
azure_ad_oid: Mapped[str] = mapped_column(String(255), unique=True, nullable=False)
email: Mapped[str] = mapped_column(String(255), nullable=False)
name: Mapped[str] = mapped_column(String(255), nullable=False)
role: Mapped[str] = mapped_column(String(50), nullable=False, default="basic_user")
agency_id: Mapped[Optional[uuid.UUID]] = mapped_column(UUID(as_uuid=True), ForeignKey("agencies.id"), nullable=True)
created_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), server_default=func.now())
# Relationships
agency: Mapped[Optional["Agency"]] = relationship("Agency", back_populates="users")
campaigns: Mapped[list["Campaign"]] = relationship("Campaign", back_populates="created_by_user")
proofs: Mapped[list["Proof"]] = relationship("Proof", back_populates="created_by_user")
flagged_items: Mapped[list["FlaggedItem"]] = relationship("FlaggedItem", back_populates="submitter")
resolved_items: Mapped[list["ResolvedItem"]] = relationship("ResolvedItem", back_populates="submitter")
change_logs: Mapped[list["UserChangeLog"]] = relationship(
"UserChangeLog",
back_populates="user",
foreign_keys="[UserChangeLog.user_id]",
)
class UserChangeLog(Base):
"""Audit log entry for changes to a user's role or agency assignment."""
__tablename__ = "user_change_logs"
id: Mapped[uuid.UUID] = mapped_column(UUID(as_uuid=True), primary_key=True, default=uuid.uuid4)
user_id: Mapped[uuid.UUID] = mapped_column(UUID(as_uuid=True), ForeignKey("users.id"), nullable=False, index=True)
changed_by_id: Mapped[Optional[uuid.UUID]] = mapped_column(UUID(as_uuid=True), ForeignKey("users.id"), nullable=True)
change_type: Mapped[str] = mapped_column(String(50), nullable=False)
field_changed: Mapped[Optional[str]] = mapped_column(String(50), nullable=True)
old_value: Mapped[Optional[str]] = mapped_column(String(255), nullable=True)
new_value: Mapped[Optional[str]] = mapped_column(String(255), nullable=True)
created_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), server_default=func.now(), index=True)
# Relationships
user: Mapped["User"] = relationship("User", back_populates="change_logs", foreign_keys=[user_id])
changed_by: Mapped[Optional["User"]] = relationship("User", foreign_keys=[changed_by_id])
class Campaign(Base):
"""Marketing campaign containing proofs."""
__tablename__ = "campaigns"
id: Mapped[uuid.UUID] = mapped_column(UUID(as_uuid=True), primary_key=True, default=uuid.uuid4)
name: Mapped[str] = mapped_column(String(255), nullable=False)
workfront_id: Mapped[Optional[str]] = mapped_column(String(100), nullable=True)
client_lead: Mapped[Optional[str]] = mapped_column(String(255), nullable=True)
agency_lead: Mapped[Optional[str]] = mapped_column(String(255), nullable=True)
brand_guidelines: Mapped[Optional[str]] = mapped_column(String(50), nullable=True)
status: Mapped[str] = mapped_column(String(50), default="In Progress")
agency_id: Mapped[Optional[uuid.UUID]] = mapped_column(UUID(as_uuid=True), ForeignKey("agencies.id"), nullable=True)
created_by: Mapped[Optional[uuid.UUID]] = mapped_column(UUID(as_uuid=True), ForeignKey("users.id"), nullable=True)
created_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), server_default=func.now())
updated_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), server_default=func.now(), onupdate=func.now())
# Relationships
agency: Mapped[Optional["Agency"]] = relationship("Agency", back_populates="campaigns")
created_by_user: Mapped[Optional["User"]] = relationship("User", back_populates="campaigns")
proofs: Mapped[list["Proof"]] = relationship("Proof", back_populates="campaign", cascade="all, delete-orphan")
class Proof(Base):
"""Marketing proof/asset to be reviewed."""
__tablename__ = "proofs"
id: Mapped[uuid.UUID] = mapped_column(UUID(as_uuid=True), primary_key=True, default=uuid.uuid4)
campaign_id: Mapped[uuid.UUID] = mapped_column(UUID(as_uuid=True), ForeignKey("campaigns.id", ondelete="CASCADE"), nullable=False)
proof_name: Mapped[str] = mapped_column(String(255), nullable=False)
channel: Mapped[Optional[str]] = mapped_column(String(100), nullable=True)
sub_channel: Mapped[Optional[str]] = mapped_column(String(100), nullable=True)
proof_type: Mapped[Optional[str]] = mapped_column(String(100), nullable=True)
workfront_id: Mapped[Optional[str]] = mapped_column(String(100), nullable=True)
created_by: Mapped[Optional[uuid.UUID]] = mapped_column(UUID(as_uuid=True), ForeignKey("users.id"), nullable=True)
created_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), server_default=func.now())
# Relationships
campaign: Mapped["Campaign"] = relationship("Campaign", back_populates="proofs")
created_by_user: Mapped[Optional["User"]] = relationship("User", back_populates="proofs")
versions: Mapped[list["ProofVersion"]] = relationship("ProofVersion", back_populates="proof", cascade="all, delete-orphan", order_by="desc(ProofVersion.version)")
__table_args__ = (
UniqueConstraint("campaign_id", "proof_name", name="uq_campaign_proof_name"),
)
class ProofVersion(Base):
"""Version of a proof with analysis results."""
__tablename__ = "proof_versions"
id: Mapped[uuid.UUID] = mapped_column(UUID(as_uuid=True), primary_key=True, default=uuid.uuid4)
proof_id: Mapped[uuid.UUID] = mapped_column(UUID(as_uuid=True), ForeignKey("proofs.id", ondelete="CASCADE"), nullable=False)
version: Mapped[int] = mapped_column(Integer, nullable=False)
file_storage_key: Mapped[Optional[str]] = mapped_column(String(500), nullable=True)
thumbnail_url: Mapped[Optional[str]] = mapped_column(Text, nullable=True)
agent_review: Mapped[Optional[dict]] = mapped_column(JSONB, nullable=True)
overall_status: Mapped[Optional[str]] = mapped_column(String(50), nullable=True)
workfront_id: Mapped[Optional[str]] = mapped_column(String(100), nullable=True)
file_hash: Mapped[Optional[str]] = mapped_column(String(32), nullable=True)
is_identical_file: Mapped[Optional[bool]] = mapped_column(Boolean, nullable=True, default=False)
created_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), server_default=func.now())
# Relationships
proof: Mapped["Proof"] = relationship("Proof", back_populates="versions")
flagged_items: Mapped[list["FlaggedItem"]] = relationship("FlaggedItem", back_populates="proof_version")
resolved_items: Mapped[list["ResolvedItem"]] = relationship("ResolvedItem", back_populates="proof_version")
error_items: Mapped[list["ErrorItem"]] = relationship("ErrorItem", back_populates="proof_version")
__table_args__ = (
UniqueConstraint("proof_id", "version", name="uq_proof_version"),
)
class FlaggedItem(Base):
"""Record of a flagged issue on a proof version."""
__tablename__ = "flagged_items"
id: Mapped[uuid.UUID] = mapped_column(UUID(as_uuid=True), primary_key=True, default=uuid.uuid4)
proof_version_id: Mapped[uuid.UUID] = mapped_column(UUID(as_uuid=True), ForeignKey("proof_versions.id"), nullable=False)
agent_flagged: Mapped[str] = mapped_column(String(100), nullable=False)
comments: Mapped[Optional[str]] = mapped_column(Text, nullable=True)
submitter_id: Mapped[Optional[uuid.UUID]] = mapped_column(UUID(as_uuid=True), ForeignKey("users.id"), nullable=True)
created_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), server_default=func.now())
# Relationships
proof_version: Mapped["ProofVersion"] = relationship("ProofVersion", back_populates="flagged_items")
submitter: Mapped[Optional["User"]] = relationship("User", back_populates="flagged_items")
class ResolvedItem(Base):
"""Record of a resolved issue on a proof version."""
__tablename__ = "resolved_items"
id: Mapped[uuid.UUID] = mapped_column(UUID(as_uuid=True), primary_key=True, default=uuid.uuid4)
proof_version_id: Mapped[uuid.UUID] = mapped_column(UUID(as_uuid=True), ForeignKey("proof_versions.id"), nullable=False)
agent: Mapped[str] = mapped_column(String(100), nullable=False)
issue: Mapped[Optional[str]] = mapped_column(Text, nullable=True)
resolution: Mapped[Optional[str]] = mapped_column(Text, nullable=True)
submitter_id: Mapped[Optional[uuid.UUID]] = mapped_column(UUID(as_uuid=True), ForeignKey("users.id"), nullable=True)
created_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), server_default=func.now())
# Relationships
proof_version: Mapped["ProofVersion"] = relationship("ProofVersion", back_populates="resolved_items")
submitter: Mapped[Optional["User"]] = relationship("User", back_populates="resolved_items")
class ErrorItem(Base):
"""Record of an analysis error on a proof version."""
__tablename__ = "error_items"
id: Mapped[uuid.UUID] = mapped_column(UUID(as_uuid=True), primary_key=True, default=uuid.uuid4)
proof_version_id: Mapped[uuid.UUID] = mapped_column(UUID(as_uuid=True), ForeignKey("proof_versions.id"), nullable=False)
error_summary: Mapped[Optional[str]] = mapped_column(Text, nullable=True)
created_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), server_default=func.now())
# Relationships
proof_version: Mapped["ProofVersion"] = relationship("ProofVersion", back_populates="error_items")
class DropdownOption(Base):
"""Configurable dropdown options for channels/sub-channels/proof types."""
__tablename__ = "dropdown_options"
id: Mapped[uuid.UUID] = mapped_column(UUID(as_uuid=True), primary_key=True, default=uuid.uuid4)
option_type: Mapped[str] = mapped_column(String(50), nullable=False) # 'channel', 'sub_channel', 'proof_type'
parent_id: Mapped[Optional[uuid.UUID]] = mapped_column(UUID(as_uuid=True), ForeignKey("dropdown_options.id"), nullable=True)
value: Mapped[str] = mapped_column(String(255), nullable=False)
display_order: Mapped[int] = mapped_column(Integer, default=0)
created_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), server_default=func.now())
# Self-referential relationship for hierarchy
parent: Mapped[Optional["DropdownOption"]] = relationship("DropdownOption", remote_side=[id], back_populates="children")
children: Mapped[list["DropdownOption"]] = relationship(
"DropdownOption",
back_populates="parent",
order_by="DropdownOption.display_order"
)
class KnowledgeBase(Base):
"""An agent knowledge base entry (one per agent spec type)."""
__tablename__ = "knowledge_bases"
id: Mapped[uuid.UUID] = mapped_column(UUID(as_uuid=True), primary_key=True, default=uuid.uuid4)
agent_key: Mapped[str] = mapped_column(String(100), unique=True, nullable=False)
display_name: Mapped[str] = mapped_column(String(255), nullable=False)
description: Mapped[Optional[str]] = mapped_column(Text, nullable=True)
created_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), server_default=func.now())
# Relationships
source_documents: Mapped[list["SourceDocument"]] = relationship(
"SourceDocument", back_populates="knowledge_base", cascade="all, delete-orphan"
)
spec_versions: Mapped[list["SpecVersion"]] = relationship(
"SpecVersion", back_populates="knowledge_base", cascade="all, delete-orphan",
order_by="desc(SpecVersion.version_number)"
)
processing_jobs: Mapped[list["ProcessingJob"]] = relationship(
"ProcessingJob", back_populates="knowledge_base", cascade="all, delete-orphan",
order_by="desc(ProcessingJob.created_at)"
)
class SourceDocument(Base):
"""An uploaded source document linked to a knowledge base."""
__tablename__ = "source_documents"
id: Mapped[uuid.UUID] = mapped_column(UUID(as_uuid=True), primary_key=True, default=uuid.uuid4)
knowledge_base_id: Mapped[uuid.UUID] = mapped_column(
UUID(as_uuid=True), ForeignKey("knowledge_bases.id", ondelete="CASCADE"), nullable=False
)
filename: Mapped[str] = mapped_column(String(500), nullable=False)
file_storage_key: Mapped[str] = mapped_column(String(500), nullable=False)
file_size_bytes: Mapped[int] = mapped_column(Integer, nullable=False)
mime_type: Mapped[str] = mapped_column(String(255), nullable=False)
uploaded_by_id: Mapped[Optional[uuid.UUID]] = mapped_column(UUID(as_uuid=True), ForeignKey("users.id"), nullable=True)
uploaded_by_name: Mapped[Optional[str]] = mapped_column(String(255), nullable=True)
parsed_markdown: Mapped[Optional[str]] = mapped_column(Text, nullable=True)
parse_status: Mapped[str] = mapped_column(String(50), default="pending", nullable=False)
parse_error: Mapped[Optional[str]] = mapped_column(Text, nullable=True)
created_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), server_default=func.now())
# Relationships
knowledge_base: Mapped["KnowledgeBase"] = relationship("KnowledgeBase", back_populates="source_documents")
class SpecVersion(Base):
"""A generated spec document with version history."""
__tablename__ = "spec_versions"
id: Mapped[uuid.UUID] = mapped_column(UUID(as_uuid=True), primary_key=True, default=uuid.uuid4)
knowledge_base_id: Mapped[uuid.UUID] = mapped_column(
UUID(as_uuid=True), ForeignKey("knowledge_bases.id", ondelete="CASCADE"), nullable=False
)
version_number: Mapped[int] = mapped_column(Integer, nullable=False)
content: Mapped[str] = mapped_column(Text, nullable=False)
source_document_ids: Mapped[Optional[list]] = mapped_column(JSONB, nullable=True)
generated_by_id: Mapped[Optional[uuid.UUID]] = mapped_column(UUID(as_uuid=True), ForeignKey("users.id"), nullable=True)
generated_by_name: Mapped[Optional[str]] = mapped_column(String(255), nullable=True)
processing_job_id: Mapped[Optional[uuid.UUID]] = mapped_column(
UUID(as_uuid=True), ForeignKey("processing_jobs.id"), nullable=True
)
is_active: Mapped[bool] = mapped_column(Boolean, default=True, nullable=False)
char_count: Mapped[int] = mapped_column(Integer, nullable=False)
created_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), server_default=func.now())
# Relationships
knowledge_base: Mapped["KnowledgeBase"] = relationship("KnowledgeBase", back_populates="spec_versions")
__table_args__ = (
UniqueConstraint("knowledge_base_id", "version_number", name="uq_kb_version_number"),
)
class ProcessingJob(Base):
"""Tracks a document processing pipeline run."""
__tablename__ = "processing_jobs"
id: Mapped[uuid.UUID] = mapped_column(UUID(as_uuid=True), primary_key=True, default=uuid.uuid4)
knowledge_base_id: Mapped[uuid.UUID] = mapped_column(
UUID(as_uuid=True), ForeignKey("knowledge_bases.id", ondelete="CASCADE"), nullable=False
)
status: Mapped[str] = mapped_column(String(50), default="pending", nullable=False)
triggered_by_id: Mapped[Optional[uuid.UUID]] = mapped_column(UUID(as_uuid=True), ForeignKey("users.id"), nullable=True)
triggered_by_name: Mapped[Optional[str]] = mapped_column(String(255), nullable=True)
total_documents: Mapped[int] = mapped_column(Integer, nullable=False, default=0)
parsed_documents: Mapped[int] = mapped_column(Integer, default=0, nullable=False)
spec_version_id: Mapped[Optional[uuid.UUID]] = mapped_column(UUID(as_uuid=True), ForeignKey("spec_versions.id"), nullable=True)
error_message: Mapped[Optional[str]] = mapped_column(Text, nullable=True)
log: Mapped[Optional[dict]] = mapped_column(JSONB, nullable=True)
started_at: Mapped[Optional[datetime]] = mapped_column(DateTime(timezone=True), nullable=True)
completed_at: Mapped[Optional[datetime]] = mapped_column(DateTime(timezone=True), nullable=True)
created_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), server_default=func.now())
# Relationships
knowledge_base: Mapped["KnowledgeBase"] = relationship("KnowledgeBase", back_populates="processing_jobs")
spec_version: Mapped[Optional["SpecVersion"]] = relationship(
"SpecVersion",
foreign_keys=[spec_version_id],
)