harness_zbt/backend/app/infrastructure/mysql/agent_models.py

from datetime import UTC, datetime

from sqlalchemy import JSON, ForeignKey, Index, String, Text
from sqlalchemy.orm import DeclarativeBase, Mapped, mapped_column

from app.infrastructure.mysql.types import UTCDateTime


def utc_now() -> datetime:
    return datetime.now(UTC)


class AgentBase(DeclarativeBase):
    pass


class AgentThreadRecord(AgentBase):
    __tablename__ = "agent_threads"
    __table_args__ = (
        Index("ix_agent_threads_owner_updated", "owner_type", "owner_id", "updated_at"),
    )

    id: Mapped[str] = mapped_column(String(26), primary_key=True)
    owner_type: Mapped[str] = mapped_column(String(16), nullable=False)
    owner_id: Mapped[str] = mapped_column(String(26), nullable=False)
    persona: Mapped[str] = mapped_column(String(32), nullable=False)
    title: Mapped[str] = mapped_column(String(128), nullable=False, default="")
    status: Mapped[str] = mapped_column(String(32), nullable=False, default="ACTIVE")
    created_at: Mapped[datetime] = mapped_column(UTCDateTime(), nullable=False, default=utc_now)
    updated_at: Mapped[datetime] = mapped_column(
        UTCDateTime(), nullable=False, default=utc_now, onupdate=utc_now
    )


class AgentMessageRecord(AgentBase):
    __tablename__ = "agent_messages"
    __table_args__ = (Index("ix_agent_messages_thread_created", "thread_id", "created_at"),)

    id: Mapped[str] = mapped_column(String(26), primary_key=True)
    thread_id: Mapped[str] = mapped_column(
        String(26), ForeignKey("agent_threads.id"), nullable=False
    )
    role: Mapped[str] = mapped_column(String(16), nullable=False)
    content: Mapped[str] = mapped_column(Text, nullable=False)
    metadata_json: Mapped[dict[str, object]] = mapped_column(JSON, nullable=False)
    created_at: Mapped[datetime] = mapped_column(UTCDateTime(), nullable=False, default=utc_now)


class AgentRunRecord(AgentBase):
    __tablename__ = "agent_runs"
    __table_args__ = (
        Index("ix_agent_runs_thread_created", "thread_id", "created_at"),
        Index("ix_agent_runs_trace", "trace_id"),
    )

    id: Mapped[str] = mapped_column(String(26), primary_key=True)
    request_id: Mapped[str] = mapped_column(String(64), unique=True, nullable=False)
    thread_id: Mapped[str] = mapped_column(
        String(26), ForeignKey("agent_threads.id"), nullable=False
    )
    trace_id: Mapped[str | None] = mapped_column(String(64))
    persona: Mapped[str] = mapped_column(String(32), nullable=False)
    status: Mapped[str] = mapped_column(String(32), nullable=False)
    input_json: Mapped[dict[str, object]] = mapped_column(JSON, nullable=False)
    output_json: Mapped[dict[str, object] | None] = mapped_column(JSON)
    error_code: Mapped[str | None] = mapped_column(String(64))
    created_at: Mapped[datetime] = mapped_column(UTCDateTime(), nullable=False, default=utc_now)
    completed_at: Mapped[datetime | None] = mapped_column(UTCDateTime())