konstruct/packages/shared/shared/models/tenant.py

"""
SQLAlchemy 2.0 ORM models for multi-tenant data.

IMPORTANT: All models here use SQLAlchemy 2.0 `Mapped[]` and `mapped_column()`
style. Never use the legacy 1.x `Column()` style.

RLS is applied to tenant-scoped tables (agents, channel_connections) via
Alembic migration. Application connections MUST use the `konstruct_app` role
(not the postgres superuser) for RLS to be enforced.
"""

from __future__ import annotations

import enum
import uuid
from datetime import datetime
from typing import Any

from sqlalchemy import JSON, Boolean, DateTime, Enum, Float, ForeignKey, Integer, String, Text, UniqueConstraint, func
from sqlalchemy.dialects.postgresql import UUID
from sqlalchemy.orm import DeclarativeBase, Mapped, mapped_column, relationship


class Base(DeclarativeBase):
    """Shared declarative base for all Konstruct ORM models."""

    pass


class ChannelTypeEnum(str, enum.Enum):
    """Matches ChannelType StrEnum in message.py — kept in sync."""

    SLACK = "slack"
    WHATSAPP = "whatsapp"
    MATTERMOST = "mattermost"
    ROCKETCHAT = "rocketchat"
    TEAMS = "teams"
    TELEGRAM = "telegram"
    SIGNAL = "signal"


class Tenant(Base):
    """
    Top-level tenant. Represents one Konstruct customer / workspace.

    RLS is NOT applied to this table — platform admin needs to list all tenants.
    The konstruct_app role has SELECT/INSERT/UPDATE/DELETE on tenants.
    """

    __tablename__ = "tenants"

    id: Mapped[uuid.UUID] = mapped_column(
        UUID(as_uuid=True),
        primary_key=True,
        default=uuid.uuid4,
    )
    name: Mapped[str] = mapped_column(String(255), unique=True, nullable=False)
    slug: Mapped[str] = mapped_column(String(100), unique=True, nullable=False)
    settings: Mapped[dict[str, Any]] = mapped_column(JSON, nullable=False, default=dict)
    created_at: Mapped[datetime] = mapped_column(
        DateTime(timezone=True),
        nullable=False,
        server_default=func.now(),
    )
    updated_at: Mapped[datetime] = mapped_column(
        DateTime(timezone=True),
        nullable=False,
        server_default=func.now(),
        onupdate=func.now(),
    )

    # ---------------------------------------------------------------------------
    # Billing fields (added in migration 005)
    # ---------------------------------------------------------------------------
    stripe_customer_id: Mapped[str | None] = mapped_column(
        String(255),
        nullable=True,
        comment="Stripe Customer ID (cus_...)",
    )
    stripe_subscription_id: Mapped[str | None] = mapped_column(
        String(255),
        nullable=True,
        comment="Stripe Subscription ID (sub_...)",
    )
    stripe_subscription_item_id: Mapped[str | None] = mapped_column(
        String(255),
        nullable=True,
        comment="Stripe Subscription Item ID (si_...) for quantity updates",
    )
    subscription_status: Mapped[str] = mapped_column(
        String(50),
        nullable=False,
        default="none",
        comment="none | trialing | active | past_due | canceled | unpaid",
    )
    trial_ends_at: Mapped[datetime | None] = mapped_column(
        DateTime(timezone=True),
        nullable=True,
        comment="Trial expiry timestamp (NULL for non-trial subscriptions)",
    )
    agent_quota: Mapped[int] = mapped_column(
        Integer,
        nullable=False,
        default=0,
        comment="Number of active agents allowed under current subscription",
    )

    # Relationships
    agents: Mapped[list[Agent]] = relationship("Agent", back_populates="tenant", cascade="all, delete-orphan")
    channel_connections: Mapped[list[ChannelConnection]] = relationship(
        "ChannelConnection", back_populates="tenant", cascade="all, delete-orphan"
    )

    def __repr__(self) -> str:
        return f"<Tenant id={self.id} slug={self.slug!r}>"


class Agent(Base):
    """
    An AI employee belonging to a specific tenant.

    RLS is ENABLED on this table. Rows are visible only when
    `app.current_tenant` session variable matches the row's `tenant_id`.
    FORCE ROW LEVEL SECURITY ensures even the table owner cannot bypass RLS.
    """

    __tablename__ = "agents"

    id: Mapped[uuid.UUID] = mapped_column(
        UUID(as_uuid=True),
        primary_key=True,
        default=uuid.uuid4,
    )
    tenant_id: Mapped[uuid.UUID] = mapped_column(
        UUID(as_uuid=True),
        ForeignKey("tenants.id", ondelete="CASCADE"),
        nullable=False,
        index=True,
    )
    name: Mapped[str] = mapped_column(String(255), nullable=False)
    role: Mapped[str] = mapped_column(String(255), nullable=False)
    persona: Mapped[str] = mapped_column(Text, nullable=False, default="")
    system_prompt: Mapped[str] = mapped_column(Text, nullable=False, default="")
    model_preference: Mapped[str] = mapped_column(
        String(50),
        nullable=False,
        default="quality",
        comment="quality | balanced | economy | local",
    )
    tool_assignments: Mapped[list[Any]] = mapped_column(JSON, nullable=False, default=list)
    escalation_rules: Mapped[list[Any]] = mapped_column(JSON, nullable=False, default=list)
    escalation_assignee: Mapped[str | None] = mapped_column(
        Text,
        nullable=True,
        comment="Slack user ID of the human to DM on escalation (e.g. U0HUMANID)",
    )
    natural_language_escalation: Mapped[bool] = mapped_column(
        Boolean,
        nullable=False,
        default=False,
        comment="Whether natural language escalation phrases trigger handoff",
    )
    is_active: Mapped[bool] = mapped_column(Boolean, nullable=False, default=True)
    budget_limit_usd: Mapped[float | None] = mapped_column(
        Float,
        nullable=True,
        default=None,
        comment="Monthly spend cap in USD. NULL means no limit.",
    )
    created_at: Mapped[datetime] = mapped_column(
        DateTime(timezone=True),
        nullable=False,
        server_default=func.now(),
    )
    updated_at: Mapped[datetime] = mapped_column(
        DateTime(timezone=True),
        nullable=False,
        server_default=func.now(),
        onupdate=func.now(),
    )

    # Relationships
    tenant: Mapped[Tenant] = relationship("Tenant", back_populates="agents")

    def __repr__(self) -> str:
        return f"<Agent id={self.id} name={self.name!r} tenant_id={self.tenant_id}>"


class AgentTemplate(Base):
    """
    Pre-built AI employee templates available to all tenants.

    Templates are NOT tenant-scoped (no tenant_id, no RLS). Any authenticated
    portal user can browse templates. Only tenant admins can deploy them.
    Deploying a template creates an independent Agent snapshot — subsequent
    changes to the template do not affect deployed agents.
    """

    __tablename__ = "agent_templates"

    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)
    role: Mapped[str] = mapped_column(String(255), nullable=False)
    description: Mapped[str] = mapped_column(
        Text,
        nullable=False,
        default="",
        comment="2-3 sentence card preview description shown in the template gallery",
    )
    category: Mapped[str] = mapped_column(
        String(100),
        nullable=False,
        default="general",
        comment="Template category: support | sales | operations | finance | general",
    )
    persona: Mapped[str] = mapped_column(
        Text,
        nullable=False,
        default="",
        comment="Paragraph describing the agent's communication style and personality",
    )
    system_prompt: Mapped[str] = mapped_column(
        Text,
        nullable=False,
        default="",
        comment="Full system prompt including AI transparency clause",
    )
    model_preference: Mapped[str] = mapped_column(
        String(50),
        nullable=False,
        default="quality",
        comment="quality | balanced | economy | local",
    )
    tool_assignments: Mapped[list[Any]] = mapped_column(
        JSON,
        nullable=False,
        default=list,
        comment="JSON array of tool name strings",
    )
    escalation_rules: Mapped[list[Any]] = mapped_column(
        JSON,
        nullable=False,
        default=list,
        comment="JSON array of {condition, action} escalation rule objects",
    )
    is_active: Mapped[bool] = mapped_column(
        Boolean,
        nullable=False,
        default=True,
        comment="Inactive templates are hidden from the gallery",
    )
    translations: Mapped[dict[str, Any]] = mapped_column(
        JSON,
        nullable=False,
        default=dict,
        comment="JSONB map of locale -> {name, description, persona} translations. E.g. {'es': {...}, 'pt': {...}}",
    )
    sort_order: Mapped[int] = mapped_column(
        Integer,
        nullable=False,
        default=0,
        comment="Display order in the template gallery (ascending)",
    )
    created_at: Mapped[datetime] = mapped_column(
        DateTime(timezone=True),
        nullable=False,
        server_default=func.now(),
    )

    def __repr__(self) -> str:
        return f"<AgentTemplate id={self.id} name={self.name!r} category={self.category!r}>"


class ChannelConnection(Base):
    """
    Links a messaging platform workspace to a Konstruct tenant.

    Example: Slack workspace T12345 → Tenant UUID abc-123.

    The Message Router queries this table to resolve incoming messages to the
    correct tenant. RLS is ENABLED — tenant agents can only see their own
    channel connections.
    """

    __tablename__ = "channel_connections"
    __table_args__ = (
        UniqueConstraint("channel_type", "workspace_id", name="uq_channel_workspace"),
    )

    id: Mapped[uuid.UUID] = mapped_column(
        UUID(as_uuid=True),
        primary_key=True,
        default=uuid.uuid4,
    )
    tenant_id: Mapped[uuid.UUID] = mapped_column(
        UUID(as_uuid=True),
        ForeignKey("tenants.id", ondelete="CASCADE"),
        nullable=False,
        index=True,
    )
    channel_type: Mapped[ChannelTypeEnum] = mapped_column(
        Enum(ChannelTypeEnum, name="channel_type_enum"),
        nullable=False,
    )
    workspace_id: Mapped[str] = mapped_column(
        String(255),
        nullable=False,
        comment="Channel-native workspace/org ID (e.g. Slack workspace ID T12345)",
    )
    config: Mapped[dict[str, Any]] = mapped_column(
        JSON,
        nullable=False,
        default=dict,
        comment="Encrypted bot tokens, channel IDs, and other per-tenant channel config",
    )
    created_at: Mapped[datetime] = mapped_column(
        DateTime(timezone=True),
        nullable=False,
        server_default=func.now(),
    )

    # Relationships
    tenant: Mapped[Tenant] = relationship("Tenant", back_populates="channel_connections")

    def __repr__(self) -> str:
        return f"<ChannelConnection channel={self.channel_type} workspace={self.workspace_id!r}>"