arize-phoenix 11.23.1__py3-none-any.whl → 12.28.1__py3-none-any.whl

phoenix/db/models.py

@@ -20,6 +20,7 @@ from sqlalchemy import (
     Integer,
     MetaData,
     Null,
+    PrimaryKeyConstraint,
     String,
     TypeDecorator,
     UniqueConstraint,
@@ -43,6 +44,7 @@ from sqlalchemy.orm import (
 )
 from sqlalchemy.sql import Values, column, compiler, expression, literal, roles, union_all
 from sqlalchemy.sql.compiler import SQLCompiler
+from sqlalchemy.sql.elements import Case
 from sqlalchemy.sql.functions import coalesce
 from typing_extensions import TypeAlias
 
@@ -154,7 +156,7 @@ def render_values_w_union(
     return compiler.process(subquery, from_linter=from_linter, **kw)
 
 
-UserRoleName: TypeAlias = Literal["SYSTEM", "ADMIN", "MEMBER"]
+UserRoleName: TypeAlias = Literal["SYSTEM", "ADMIN", "MEMBER", "VIEWER"]
 AuthMethod: TypeAlias = Literal["LOCAL", "OAUTH2"]
 
 
@@ -188,7 +190,7 @@ class JsonDict(TypeDecorator[dict[str, Any]]):
     impl = JSON_
 
     def process_bind_param(self, value: Optional[dict[str, Any]], _: Dialect) -> dict[str, Any]:
-        return orjson.loads(orjson.dumps(value)) if isinstance(value, dict) else {}
+        return value if isinstance(value, dict) else {}
 
     def process_result_value(self, value: Optional[Any], _: Dialect) -> Optional[dict[str, Any]]:
         return orjson.loads(orjson.dumps(value)) if isinstance(value, dict) and value else value
@@ -200,7 +202,7 @@ class JsonList(TypeDecorator[list[Any]]):
     impl = JSON_
 
     def process_bind_param(self, value: Optional[list[Any]], _: Dialect) -> list[Any]:
-        return orjson.loads(orjson.dumps(value)) if isinstance(value, list) else []
+        return value if isinstance(value, list) else []
 
     def process_result_value(self, value: Optional[Any], _: Dialect) -> Optional[list[Any]]:
         return orjson.loads(orjson.dumps(value)) if isinstance(value, list) and value else value
@@ -442,12 +444,32 @@ class _RegexStr(TypeDecorator[re.Pattern[str]]):
         return re.compile(value)
 
 
+_HEX_COLOR_PATTERN = re.compile(r"^#([0-9a-f]{6})$")
+
+
+class _HexColor(TypeDecorator[str]):
+    # See https://docs.sqlalchemy.org/en/20/core/custom_types.html
+    cache_ok = True
+    impl = String
+
+    def process_bind_param(self, value: Optional[str], _: Dialect) -> Optional[str]:
+        if value is None:
+            return None
+        if not _HEX_COLOR_PATTERN.match(value):
+            raise ValueError(f"Expected a hex color, got {value}")
+        return value
+
+    def process_result_value(self, value: Optional[str], _: Dialect) -> Optional[str]:
+        if value is None:
+            return None
+        return value
+
+
 class ExperimentRunOutput(TypedDict, total=False):
     task_output: Any
 
 
 class Base(DeclarativeBase):
-    id: Mapped[int] = mapped_column(primary_key=True)
     # Enforce best practices for naming constraints
     # https://alembic.sqlalchemy.org/en/latest/naming.html#integration-of-naming-conventions-into-operations-autogenerate
     metadata = MetaData(
@@ -467,7 +489,12 @@ class Base(DeclarativeBase):
     }
 
 
-class ProjectTraceRetentionPolicy(Base):
+class HasId(Base):
+    __abstract__ = True
+    id: Mapped[int] = mapped_column(primary_key=True)
+
+
+class ProjectTraceRetentionPolicy(HasId):
     __tablename__ = "project_trace_retention_policies"
     name: Mapped[str] = mapped_column(String, nullable=False)
     cron_expression: Mapped[TraceRetentionCronExpression] = mapped_column(
@@ -479,7 +506,7 @@ class ProjectTraceRetentionPolicy(Base):
     )
 
 
-class Project(Base):
+class Project(HasId):
     __tablename__ = "projects"
     name: Mapped[str]
     description: Mapped[Optional[str]]
@@ -519,7 +546,7 @@ class Project(Base):
     )
 
 
-class ProjectSession(Base):
+class ProjectSession(HasId):
     __tablename__ = "project_sessions"
     session_id: Mapped[str] = mapped_column(String, nullable=False, unique=True)
     project_id: Mapped[int] = mapped_column(
@@ -536,7 +563,7 @@ class ProjectSession(Base):
     )
 
 
-class Trace(Base):
+class Trace(HasId):
     __tablename__ = "traces"
     project_rowid: Mapped[int] = mapped_column(
         ForeignKey("projects.id", ondelete="CASCADE"),
@@ -591,13 +618,13 @@ class Trace(Base):
     )
 
 
-class Span(Base):
+class Span(HasId):
     __tablename__ = "spans"
     trace_rowid: Mapped[int] = mapped_column(
         ForeignKey("traces.id", ondelete="CASCADE"),
         index=True,
     )
-    span_id: Mapped[str] = mapped_column(index=True)
+    span_id: Mapped[str]
     parent_id: Mapped[Optional[str]] = mapped_column(index=True)
     name: Mapped[str]
     span_kind: Mapped[str]
@@ -803,14 +830,27 @@ class NumDocuments(expression.FunctionElement[int]):
 @compiles(NumDocuments)
 def _(element: Any, compiler: SQLCompiler, **kw: Any) -> Any:
     # See https://docs.sqlalchemy.org/en/20/core/compiler.html
-    array_length = (
-        func.json_array_length if isinstance(compiler, SQLiteCompiler) else func.jsonb_array_length
-    )
     attributes, span_kind = list(element.clauses)
     retrieval_docs = attributes[RETRIEVAL_DOCUMENTS]
-    num_retrieval_docs = coalesce(array_length(retrieval_docs), 0)
+    num_retrieval_docs: Case[Any] | coalesce[Any]
     reranker_docs = attributes[RERANKER_OUTPUT_DOCUMENTS]
-    num_reranker_docs = coalesce(array_length(reranker_docs), 0)
+    num_reranker_docs: Case[Any] | coalesce[Any]
+    if isinstance(compiler, SQLiteCompiler):
+        # SQLite's json_array_length returns 0 for non-array values
+        num_retrieval_docs = coalesce(func.json_array_length(retrieval_docs), 0)
+        num_reranker_docs = coalesce(func.json_array_length(reranker_docs), 0)
+    else:
+        # PostgreSQL's jsonb_array_length throws "cannot get array length of a scalar"
+        # for non-array values, so check the type first
+        num_retrieval_docs = sql.case(
+            (func.jsonb_typeof(retrieval_docs) == "array", func.jsonb_array_length(retrieval_docs)),
+            else_=0,
+        )
+        num_reranker_docs = sql.case(
+            (func.jsonb_typeof(reranker_docs) == "array", func.jsonb_array_length(reranker_docs)),
+            else_=0,
+        )
+
     return compiler.process(
         sql.case(
             (func.upper(span_kind) == "RERANKER", num_reranker_docs),
@@ -894,15 +934,15 @@ async def init_models(engine: AsyncEngine) -> None:
         )
 
 
-class SpanAnnotation(Base):
+class SpanAnnotation(HasId):
     __tablename__ = "span_annotations"
     span_rowid: Mapped[int] = mapped_column(
         ForeignKey("spans.id", ondelete="CASCADE"),
         index=True,
     )
     name: Mapped[str]
-    label: Mapped[Optional[str]] = mapped_column(String, index=True)
-    score: Mapped[Optional[float]] = mapped_column(Float, index=True)
+    label: Mapped[Optional[str]]
+    score: Mapped[Optional[float]]
     explanation: Mapped[Optional[str]]
     metadata_: Mapped[dict[str, Any]] = mapped_column("metadata")
     annotator_kind: Mapped[Literal["LLM", "CODE", "HUMAN"]] = mapped_column(
@@ -934,15 +974,15 @@ class SpanAnnotation(Base):
     )
 
 
-class TraceAnnotation(Base):
+class TraceAnnotation(HasId):
     __tablename__ = "trace_annotations"
     trace_rowid: Mapped[int] = mapped_column(
         ForeignKey("traces.id", ondelete="CASCADE"),
         index=True,
     )
     name: Mapped[str]
-    label: Mapped[Optional[str]] = mapped_column(String, index=True)
-    score: Mapped[Optional[float]] = mapped_column(Float, index=True)
+    label: Mapped[Optional[str]]
+    score: Mapped[Optional[float]]
     explanation: Mapped[Optional[str]]
     metadata_: Mapped[dict[str, Any]] = mapped_column("metadata")
     annotator_kind: Mapped[Literal["LLM", "CODE", "HUMAN"]] = mapped_column(
@@ -971,7 +1011,7 @@ class TraceAnnotation(Base):
     )
 
 
-class DocumentAnnotation(Base):
+class DocumentAnnotation(HasId):
     __tablename__ = "document_annotations"
     span_rowid: Mapped[int] = mapped_column(
         ForeignKey("spans.id", ondelete="CASCADE"),
@@ -979,8 +1019,8 @@ class DocumentAnnotation(Base):
     )
     document_position: Mapped[int]
     name: Mapped[str]
-    label: Mapped[Optional[str]] = mapped_column(String, index=True)
-    score: Mapped[Optional[float]] = mapped_column(Float, index=True)
+    label: Mapped[Optional[str]]
+    score: Mapped[Optional[float]]
     explanation: Mapped[Optional[str]]
     metadata_: Mapped[dict[str, Any]] = mapped_column("metadata")
     annotator_kind: Mapped[Literal["LLM", "CODE", "HUMAN"]] = mapped_column(
@@ -1012,7 +1052,44 @@ class DocumentAnnotation(Base):
     )
 
 
-class Dataset(Base):
+class ProjectSessionAnnotation(HasId):
+    __tablename__ = "project_session_annotations"
+    project_session_id: Mapped[int] = mapped_column(
+        ForeignKey("project_sessions.id", ondelete="CASCADE"),
+        index=True,
+    )
+    name: Mapped[str]
+    label: Mapped[Optional[str]]
+    score: Mapped[Optional[float]]
+    explanation: Mapped[Optional[str]]
+    metadata_: Mapped[dict[str, Any]] = mapped_column("metadata")
+    annotator_kind: Mapped[Literal["LLM", "CODE", "HUMAN"]] = mapped_column(
+        CheckConstraint("annotator_kind IN ('LLM', 'CODE', 'HUMAN')", name="valid_annotator_kind"),
+    )
+    created_at: Mapped[datetime] = mapped_column(UtcTimeStamp, server_default=func.now())
+    updated_at: Mapped[datetime] = mapped_column(
+        UtcTimeStamp, server_default=func.now(), onupdate=func.now()
+    )
+    identifier: Mapped[str] = mapped_column(
+        String,
+        server_default="",
+        nullable=False,
+    )
+    source: Mapped[Literal["API", "APP"]] = mapped_column(
+        CheckConstraint("source IN ('API', 'APP')", name="valid_source"),
+    )
+    user_id: Mapped[Optional[int]] = mapped_column(ForeignKey("users.id", ondelete="SET NULL"))
+
+    __table_args__ = (
+        UniqueConstraint(
+            "name",
+            "project_session_id",
+            "identifier",
+        ),
+    )
+
+
+class Dataset(HasId):
     __tablename__ = "datasets"
     name: Mapped[str] = mapped_column(unique=True)
     description: Mapped[Optional[str]]
@@ -1021,6 +1098,14 @@ class Dataset(Base):
     updated_at: Mapped[datetime] = mapped_column(
         UtcTimeStamp, server_default=func.now(), onupdate=func.now()
     )
+    user_id: Mapped[Optional[int]] = mapped_column(ForeignKey("users.id", ondelete="SET NULL"))
+    user: Mapped[Optional["User"]] = relationship("User")
+    experiment_tags: Mapped[list["ExperimentTag"]] = relationship(
+        "ExperimentTag", back_populates="dataset"
+    )
+    datasets_dataset_labels: Mapped[list["DatasetsDatasetLabel"]] = relationship(
+        "DatasetsDatasetLabel", back_populates="dataset"
+    )
 
     @hybrid_property
     def example_count(self) -> Optional[int]:
@@ -1071,7 +1156,45 @@ class Dataset(Base):
             )
 
 
-class DatasetVersion(Base):
+class DatasetLabel(HasId):
+    __tablename__ = "dataset_labels"
+    name: Mapped[str] = mapped_column(unique=True)
+    description: Mapped[Optional[str]]
+    color: Mapped[str] = mapped_column(_HexColor, nullable=False)
+    datasets_dataset_labels: Mapped[list["DatasetsDatasetLabel"]] = relationship(
+        "DatasetsDatasetLabel", back_populates="dataset_label"
+    )
+    user_id: Mapped[Optional[int]] = mapped_column(
+        ForeignKey("users.id", ondelete="SET NULL"),
+        nullable=True,
+    )
+    user: Mapped[Optional["User"]] = relationship("User")
+
+
+class DatasetsDatasetLabel(Base):
+    __tablename__ = "datasets_dataset_labels"
+    dataset_id: Mapped[int] = mapped_column(
+        ForeignKey("datasets.id", ondelete="CASCADE"),
+    )
+    dataset_label_id: Mapped[int] = mapped_column(
+        ForeignKey("dataset_labels.id", ondelete="CASCADE"),
+        # index on the second element of the composite primary key
+        index=True,
+    )
+    dataset: Mapped["Dataset"] = relationship("Dataset", back_populates="datasets_dataset_labels")
+    dataset_label: Mapped["DatasetLabel"] = relationship(
+        "DatasetLabel", back_populates="datasets_dataset_labels"
+    )
+
+    __table_args__ = (
+        PrimaryKeyConstraint(
+            "dataset_id",
+            "dataset_label_id",
+        ),
+    )
+
+
+class DatasetVersion(HasId):
     __tablename__ = "dataset_versions"
     dataset_id: Mapped[int] = mapped_column(
         ForeignKey("datasets.id", ondelete="CASCADE"),
@@ -1080,9 +1203,11 @@ class DatasetVersion(Base):
     description: Mapped[Optional[str]]
     metadata_: Mapped[dict[str, Any]] = mapped_column("metadata")
     created_at: Mapped[datetime] = mapped_column(UtcTimeStamp, server_default=func.now())
+    user_id: Mapped[Optional[int]] = mapped_column(ForeignKey("users.id", ondelete="SET NULL"))
+    user: Mapped[Optional["User"]] = relationship("User")
 
 
-class DatasetExample(Base):
+class DatasetExample(HasId):
     __tablename__ = "dataset_examples"
     dataset_id: Mapped[int] = mapped_column(
         ForeignKey("datasets.id", ondelete="CASCADE"),
@@ -1096,13 +1221,20 @@ class DatasetExample(Base):
     created_at: Mapped[datetime] = mapped_column(UtcTimeStamp, server_default=func.now())
 
     span: Mapped[Optional[Span]] = relationship(back_populates="dataset_examples")
+    dataset_splits_dataset_examples: Mapped[list["DatasetSplitDatasetExample"]] = relationship(
+        "DatasetSplitDatasetExample",
+        back_populates="dataset_example",
+    )
+    experiment_dataset_examples: Mapped[list["ExperimentDatasetExample"]] = relationship(
+        "ExperimentDatasetExample",
+        back_populates="dataset_example",
+    )
 
 
-class DatasetExampleRevision(Base):
+class DatasetExampleRevision(HasId):
     __tablename__ = "dataset_example_revisions"
     dataset_example_id: Mapped[int] = mapped_column(
         ForeignKey("dataset_examples.id", ondelete="CASCADE"),
-        index=True,
     )
     dataset_version_id: Mapped[int] = mapped_column(
         ForeignKey("dataset_versions.id", ondelete="CASCADE"),
@@ -1118,6 +1250,11 @@ class DatasetExampleRevision(Base):
     )
     created_at: Mapped[datetime] = mapped_column(UtcTimeStamp, server_default=func.now())
 
+    experiment_dataset_examples: Mapped[list["ExperimentDatasetExample"]] = relationship(
+        "ExperimentDatasetExample",
+        back_populates="dataset_example_revision",
+    )
+
     __table_args__ = (
         UniqueConstraint(
             "dataset_example_id",
@@ -1126,7 +1263,56 @@ class DatasetExampleRevision(Base):
     )
 
 
-class Experiment(Base):
+class DatasetSplit(HasId):
+    __tablename__ = "dataset_splits"
+
+    user_id: Mapped[Optional[int]] = mapped_column(
+        ForeignKey("users.id", ondelete="SET NULL"),
+        nullable=True,
+        index=True,
+    )
+    name: Mapped[str] = mapped_column(String, nullable=False, unique=True)
+    description: Mapped[Optional[str]]
+    color: Mapped[str] = mapped_column(String, nullable=False)
+    metadata_: Mapped[dict[str, Any]] = mapped_column("metadata")
+    created_at: Mapped[datetime] = mapped_column(UtcTimeStamp, server_default=func.now())
+    updated_at: Mapped[datetime] = mapped_column(
+        UtcTimeStamp, server_default=func.now(), onupdate=func.now()
+    )
+    dataset_splits_dataset_examples: Mapped[list["DatasetSplitDatasetExample"]] = relationship(
+        "DatasetSplitDatasetExample",
+        back_populates="dataset_split",
+    )
+    experiment_dataset_splits: Mapped[list["ExperimentDatasetSplit"]] = relationship(
+        "ExperimentDatasetSplit",
+        back_populates="dataset_split",
+    )
+
+
+class DatasetSplitDatasetExample(Base):
+    __tablename__ = "dataset_splits_dataset_examples"
+    dataset_split_id: Mapped[int] = mapped_column(
+        ForeignKey("dataset_splits.id", ondelete="CASCADE"),
+    )
+    dataset_example_id: Mapped[int] = mapped_column(
+        ForeignKey("dataset_examples.id", ondelete="CASCADE"),
+        index=True,
+    )
+    dataset_split: Mapped["DatasetSplit"] = relationship(
+        "DatasetSplit", back_populates="dataset_splits_dataset_examples"
+    )
+    dataset_example: Mapped["DatasetExample"] = relationship(
+        "DatasetExample", back_populates="dataset_splits_dataset_examples"
+    )
+    __table_args__ = (
+        PrimaryKeyConstraint(
+            "dataset_split_id",
+            "dataset_example_id",
+        ),
+    )
+
+
+class Experiment(HasId):
     __tablename__ = "experiments"
     dataset_id: Mapped[int] = mapped_column(
         ForeignKey("datasets.id", ondelete="CASCADE"),
@@ -1141,18 +1327,83 @@ class Experiment(Base):
     repetitions: Mapped[int]
     metadata_: Mapped[dict[str, Any]] = mapped_column("metadata")
     project_name: Mapped[Optional[str]] = mapped_column(index=True)
+    user_id: Mapped[Optional[int]] = mapped_column(ForeignKey("users.id", ondelete="SET NULL"))
     created_at: Mapped[datetime] = mapped_column(UtcTimeStamp, server_default=func.now())
     updated_at: Mapped[datetime] = mapped_column(
         UtcTimeStamp, server_default=func.now(), onupdate=func.now()
     )
+    user: Mapped[Optional["User"]] = relationship("User")
+    experiment_dataset_splits: Mapped[list["ExperimentDatasetSplit"]] = relationship(
+        "ExperimentDatasetSplit",
+        back_populates="experiment",
+    )
+    experiment_dataset_examples: Mapped[list["ExperimentDatasetExample"]] = relationship(
+        "ExperimentDatasetExample",
+        back_populates="experiment",
+    )
+    experiment_tags: Mapped[list["ExperimentTag"]] = relationship(
+        "ExperimentTag", back_populates="experiment"
+    )
 
 
-class ExperimentRun(Base):
-    __tablename__ = "experiment_runs"
+class ExperimentDatasetSplit(Base):
+    __tablename__ = "experiments_dataset_splits"
+    experiment_id: Mapped[int] = mapped_column(
+        ForeignKey("experiments.id", ondelete="CASCADE"),
+    )
+    dataset_split_id: Mapped[int] = mapped_column(
+        ForeignKey("dataset_splits.id", ondelete="CASCADE"),
+        index=True,
+    )
+    experiment: Mapped["Experiment"] = relationship(
+        "Experiment", back_populates="experiment_dataset_splits"
+    )
+    dataset_split: Mapped["DatasetSplit"] = relationship(
+        "DatasetSplit", back_populates="experiment_dataset_splits"
+    )
+    __table_args__ = (
+        PrimaryKeyConstraint(
+            "experiment_id",
+            "dataset_split_id",
+        ),
+    )
+
+
+class ExperimentDatasetExample(Base):
+    __tablename__ = "experiments_dataset_examples"
     experiment_id: Mapped[int] = mapped_column(
         ForeignKey("experiments.id", ondelete="CASCADE"),
+    )
+    dataset_example_id: Mapped[int] = mapped_column(
+        ForeignKey("dataset_examples.id", ondelete="CASCADE"),
         index=True,
     )
+    dataset_example_revision_id: Mapped[int] = mapped_column(
+        ForeignKey("dataset_example_revisions.id", ondelete="CASCADE"),
+        index=True,
+    )
+    experiment: Mapped["Experiment"] = relationship(
+        "Experiment", back_populates="experiment_dataset_examples"
+    )
+    dataset_example: Mapped["DatasetExample"] = relationship(
+        "DatasetExample", back_populates="experiment_dataset_examples"
+    )
+    dataset_example_revision: Mapped["DatasetExampleRevision"] = relationship(
+        "DatasetExampleRevision", back_populates="experiment_dataset_examples"
+    )
+    __table_args__ = (
+        PrimaryKeyConstraint(
+            "experiment_id",
+            "dataset_example_id",
+        ),
+    )
+
+
+class ExperimentRun(HasId):
+    __tablename__ = "experiment_runs"
+    experiment_id: Mapped[int] = mapped_column(
+        ForeignKey("experiments.id", ondelete="CASCADE"),
+    )
     dataset_example_id: Mapped[int] = mapped_column(
         ForeignKey("dataset_examples.id", ondelete="CASCADE"),
         index=True,
@@ -1192,11 +1443,10 @@ class ExperimentRun(Base):
     )
 
 
-class ExperimentRunAnnotation(Base):
+class ExperimentRunAnnotation(HasId):
     __tablename__ = "experiment_run_annotations"
     experiment_run_id: Mapped[int] = mapped_column(
         ForeignKey("experiment_runs.id", ondelete="CASCADE"),
-        index=True,
     )
     name: Mapped[str]
     annotator_kind: Mapped[str] = mapped_column(
@@ -1223,13 +1473,36 @@ class ExperimentRunAnnotation(Base):
     )
 
 
-class UserRole(Base):
+class ExperimentTag(HasId):
+    __tablename__ = "experiment_tags"
+    experiment_id: Mapped[int] = mapped_column(
+        ForeignKey("experiments.id", ondelete="CASCADE"),
+        index=True,
+    )
+    dataset_id: Mapped[int] = mapped_column(
+        ForeignKey("datasets.id", ondelete="CASCADE"),
+    )
+    user_id: Mapped[Optional[int]] = mapped_column(
+        ForeignKey("users.id", ondelete="SET NULL"),
+        index=True,
+        nullable=True,
+    )
+    name: Mapped[str]
+    description: Mapped[Optional[str]]
+    experiment: Mapped["Experiment"] = relationship("Experiment", back_populates="experiment_tags")
+    dataset: Mapped["Dataset"] = relationship("Dataset", back_populates="experiment_tags")
+    user: Mapped[Optional["User"]] = relationship("User")
+
+    __table_args__ = (UniqueConstraint("dataset_id", "name"),)
+
+
+class UserRole(HasId):
     __tablename__ = "user_roles"
     name: Mapped[UserRoleName] = mapped_column(unique=True, index=True)
     users: Mapped[list["User"]] = relationship("User", back_populates="role")
 
 
-class User(Base):
+class User(HasId):
     __tablename__ = "users"
     user_role_id: Mapped[int] = mapped_column(
         ForeignKey("user_roles.id", ondelete="CASCADE"),
@@ -1339,7 +1612,41 @@ class OAuth2User(User):
         )
 
 
-class PasswordResetToken(Base):
+def LDAPUser(
+    *,
+    email: str,
+    username: str,
+    unique_id: str | None = None,
+    user_role_id: int | None = None,
+) -> OAuth2User:
+    """Factory function to create an LDAP user stored as OAuth2User.
+
+    This is a zero-migration approach: LDAP users are stored in the existing
+    OAuth2User table with a special Unicode marker in oauth2_client_id to
+    distinguish them from actual OAuth2 users. This avoids schema changes
+    while allowing LDAP authentication to coexist with OAuth2.
+
+    Args:
+        email: User's email address
+        username: User's display name
+        unique_id: User's LDAP unique ID (stored in oauth2_user_id)
+        user_role_id: Phoenix role ID (ADMIN, MEMBER, VIEWER)
+
+    Returns:
+        OAuth2User instance configured as an LDAP user
+    """
+    from phoenix.server.ldap import LDAP_CLIENT_ID_MARKER
+
+    return OAuth2User(
+        email=email,
+        username=username,
+        oauth2_client_id=LDAP_CLIENT_ID_MARKER,
+        oauth2_user_id=unique_id,
+        user_role_id=user_role_id,
+    )
+
+
+class PasswordResetToken(HasId):
     __tablename__ = "password_reset_tokens"
     user_id: Mapped[int] = mapped_column(
         ForeignKey("users.id", ondelete="CASCADE"),
@@ -1348,11 +1655,11 @@ class PasswordResetToken(Base):
     )
     user: Mapped["User"] = relationship("User", back_populates="password_reset_token")
     created_at: Mapped[datetime] = mapped_column(UtcTimeStamp, server_default=func.now())
-    expires_at: Mapped[Optional[datetime]] = mapped_column(UtcTimeStamp, nullable=False, index=True)
+    expires_at: Mapped[datetime] = mapped_column(UtcTimeStamp, nullable=False, index=True)
     __table_args__ = (dict(sqlite_autoincrement=True),)
 
 
-class RefreshToken(Base):
+class RefreshToken(HasId):
     __tablename__ = "refresh_tokens"
     user_id: Mapped[int] = mapped_column(
         ForeignKey("users.id", ondelete="CASCADE"),
@@ -1360,11 +1667,11 @@ class RefreshToken(Base):
     )
     user: Mapped["User"] = relationship("User", back_populates="refresh_tokens")
     created_at: Mapped[datetime] = mapped_column(UtcTimeStamp, server_default=func.now())
-    expires_at: Mapped[Optional[datetime]] = mapped_column(UtcTimeStamp, nullable=False, index=True)
+    expires_at: Mapped[datetime] = mapped_column(UtcTimeStamp, nullable=False, index=True)
     __table_args__ = (dict(sqlite_autoincrement=True),)
 
 
-class AccessToken(Base):
+class AccessToken(HasId):
     __tablename__ = "access_tokens"
     user_id: Mapped[int] = mapped_column(
         ForeignKey("users.id", ondelete="CASCADE"),
@@ -1372,7 +1679,7 @@ class AccessToken(Base):
     )
     user: Mapped["User"] = relationship("User", back_populates="access_tokens")
     created_at: Mapped[datetime] = mapped_column(UtcTimeStamp, server_default=func.now())
-    expires_at: Mapped[Optional[datetime]] = mapped_column(UtcTimeStamp, nullable=False, index=True)
+    expires_at: Mapped[datetime] = mapped_column(UtcTimeStamp, nullable=False, index=True)
     refresh_token_id: Mapped[int] = mapped_column(
         ForeignKey("refresh_tokens.id", ondelete="CASCADE"),
         index=True,
@@ -1381,7 +1688,7 @@ class AccessToken(Base):
     __table_args__ = (dict(sqlite_autoincrement=True),)
 
 
-class ApiKey(Base):
+class ApiKey(HasId):
     __tablename__ = "api_keys"
     user_id: Mapped[int] = mapped_column(
         ForeignKey("users.id", ondelete="CASCADE"),
@@ -1398,7 +1705,7 @@ class ApiKey(Base):
 CostType: TypeAlias = Literal["DEFAULT", "OVERRIDE"]
 
 
-class GenerativeModel(Base):
+class GenerativeModel(HasId):
     __tablename__ = "generative_models"
     name: Mapped[str] = mapped_column(String, nullable=False)
     provider: Mapped[str]
@@ -1447,7 +1754,7 @@ class GenerativeModel(Base):
     )
 
 
-class TokenPrice(Base):
+class TokenPrice(HasId):
     __tablename__ = "token_prices"
     model_id: Mapped[int] = mapped_column(
         ForeignKey("generative_models.id", ondelete="CASCADE"),
@@ -1473,7 +1780,7 @@ class TokenPrice(Base):
     )
 
 
-class PromptLabel(Base):
+class PromptLabel(HasId):
     __tablename__ = "prompt_labels"
     name: Mapped[str] = mapped_column(String, unique=True, index=True, nullable=False)
     description: Mapped[Optional[str]]
@@ -1487,7 +1794,7 @@ class PromptLabel(Base):
     )
 
 
-class Prompt(Base):
+class Prompt(HasId):
     __tablename__ = "prompts"
     source_prompt_id: Mapped[Optional[int]] = mapped_column(
         ForeignKey("prompts.id", ondelete="SET NULL"),
@@ -1524,7 +1831,7 @@ class Prompt(Base):
     )
 
 
-class PromptPromptLabel(Base):
+class PromptPromptLabel(HasId):
     __tablename__ = "prompts_prompt_labels"
     prompt_label_id: Mapped[int] = mapped_column(
         ForeignKey("prompt_labels.id", ondelete="CASCADE"),
@@ -1545,7 +1852,7 @@ class PromptPromptLabel(Base):
     __table_args__ = (UniqueConstraint("prompt_label_id", "prompt_id"),)
 
 
-class PromptVersion(Base):
+class PromptVersion(HasId):
     __tablename__ = "prompt_versions"
 
     prompt_id: Mapped[int] = mapped_column(
@@ -1594,7 +1901,7 @@ class PromptVersion(Base):
     )
 
 
-class PromptVersionTag(Base):
+class PromptVersionTag(HasId):
     __tablename__ = "prompt_version_tags"
 
     name: Mapped[Identifier] = mapped_column(_Identifier, nullable=False)
@@ -1623,13 +1930,13 @@ class PromptVersionTag(Base):
     __table_args__ = (UniqueConstraint("name", "prompt_id"),)
 
 
-class AnnotationConfig(Base):
+class AnnotationConfig(HasId):
     __tablename__ = "annotation_configs"
     name: Mapped[str] = mapped_column(String, nullable=False, unique=True)
     config: Mapped[AnnotationConfigType] = mapped_column(_AnnotationConfig, nullable=False)
 
 
-class ProjectAnnotationConfig(Base):
+class ProjectAnnotationConfig(HasId):
     __tablename__ = "project_annotation_configs"
     project_id: Mapped[int] = mapped_column(
         ForeignKey("projects.id", ondelete="CASCADE"), nullable=False, index=True
@@ -1641,16 +1948,18 @@ class ProjectAnnotationConfig(Base):
     __table_args__ = (UniqueConstraint("project_id", "annotation_config_id"),)
 
 
-class SpanCost(Base):
+class SpanCost(HasId):
     __tablename__ = "span_costs"
 
     span_rowid: Mapped[int] = mapped_column(
         ForeignKey("spans.id", ondelete="CASCADE"),
         nullable=False,
+        index=True,
     )
     trace_rowid: Mapped[int] = mapped_column(
         ForeignKey("traces.id", ondelete="CASCADE"),
         nullable=False,
+        index=True,
     )
     span_start_time: Mapped[datetime] = mapped_column(
         UtcTimeStamp,
@@ -1753,14 +2062,13 @@ class SpanCost(Base):
             self.total_tokens = (self.total_tokens or 0) + tokens
 
 
-class SpanCostDetail(Base):
+class SpanCostDetail(HasId):
     __tablename__ = "span_cost_details"
     span_cost_id: Mapped[int] = mapped_column(
         ForeignKey("span_costs.id", ondelete="CASCADE"),
         nullable=False,
-        index=True,
     )
-    token_type: Mapped[str]
+    token_type: Mapped[str] = mapped_column(index=True)
     is_prompt: Mapped[bool]
 
     cost: Mapped[Optional[float]]