app/models/request_log.py

from sqlalchemy import text,  Index, Integer, Float, Text, func, column
from sqlalchemy.orm import Mapped, mapped_column

from app.models.base import Base


class RequestLog(Base):
    __tablename__ = "request_logs"

    id: Mapped[int] = mapped_column(Integer, primary_key=True, autoincrement=True)
    request_id: Mapped[str | None] = mapped_column(Text, unique=True, nullable=True)
    provider: Mapped[str | None] = mapped_column(Text, nullable=True)
    model: Mapped[str | None] = mapped_column(Text, nullable=True)
    endpoint: Mapped[str | None] = mapped_column(Text, nullable=True)
    key_id: Mapped[int | None] = mapped_column(Integer, nullable=True)
    hub_key_id: Mapped[int | None] = mapped_column(Integer, nullable=True)
    tokens_in: Mapped[int] = mapped_column(Integer, default=0)
    tokens_out: Mapped[int] = mapped_column(Integer, default=0)
    cost: Mapped[float] = mapped_column(Float, default=0)
    latency_ms: Mapped[int] = mapped_column(Integer, default=0)
    status_code: Mapped[int] = mapped_column(Integer, default=200)
    error: Mapped[str | None] = mapped_column(Text, nullable=True)
    cached: Mapped[int] = mapped_column(Integer, default=0)
    cache_type: Mapped[str | None] = mapped_column(Text, nullable=True)
    streaming: Mapped[int] = mapped_column(Integer, default=0)
    fallback_count: Mapped[int] = mapped_column(Integer, default=0)
    ip: Mapped[str | None] = mapped_column(Text, nullable=True)
    session_id: Mapped[str | None] = mapped_column(Text, nullable=True)
    injection_flagged: Mapped[int] = mapped_column(Integer, default=0)
    has_images: Mapped[int] = mapped_column(Integer, default=0)
    ab_variant: Mapped[str | None] = mapped_column(Text, nullable=True)
    cache_read_tokens: Mapped[int] = mapped_column(Integer, default=0)
    cache_write_tokens: Mapped[int] = mapped_column(Integer, default=0)
    timestamp: Mapped[str] = mapped_column(Text, server_default=text("CURRENT_TIMESTAMP"))

    __table_args__ = (
        Index("idx_request_logs_timestamp", "timestamp"),
        Index("idx_request_logs_provider", "provider"),
        Index("idx_request_logs_model", "model"),
        Index("idx_request_logs_session", "session_id"),
        Index("idx_request_logs_key_ts", "key_id", "timestamp"),
        Index(
            "idx_request_logs_ts_status",
            "timestamp",
            postgresql_where=column("error").is_(None),
        ),
    )