arize-phoenix 8.32.0__py3-none-any.whl → 9.0.0__py3-none-any.whl

phoenix/db/migrations/versions/bb8139330879_create_project_trace_retention_policies_table.py

@@ -0,0 +1,77 @@
+"""create project trace retention policies table
+
+Revision ID: bb8139330879
+Revises: 2f9d1a65945f
+Create Date: 2025-02-27 15:57:18.752472
+
+"""
+
+from typing import Any, Sequence, Union
+
+import sqlalchemy as sa
+from alembic import op
+from sqlalchemy import JSON
+from sqlalchemy.dialects import postgresql
+from sqlalchemy.ext.compiler import compiles
+
+
+class JSONB(JSON):
+    # See https://docs.sqlalchemy.org/en/20/core/custom_types.html
+    __visit_name__ = "JSONB"
+
+
+@compiles(JSONB, "sqlite")
+def _(*args: Any, **kwargs: Any) -> str:
+    # See https://docs.sqlalchemy.org/en/20/core/custom_types.html
+    return "JSONB"
+
+
+JSON_ = (
+    JSON()
+    .with_variant(
+        postgresql.JSONB(),
+        "postgresql",
+    )
+    .with_variant(
+        JSONB(),
+        "sqlite",
+    )
+)
+
+
+# revision identifiers, used by Alembic.
+revision: str = "bb8139330879"
+down_revision: Union[str, None] = "2f9d1a65945f"
+branch_labels: Union[str, Sequence[str], None] = None
+depends_on: Union[str, Sequence[str], None] = None
+
+
+def upgrade() -> None:
+    op.create_table(
+        "project_trace_retention_policies",
+        sa.Column("id", sa.Integer, primary_key=True),
+        sa.Column("name", sa.String, nullable=False),
+        sa.Column("cron_expression", sa.String, nullable=False),
+        sa.Column("rule", JSON_, nullable=False),
+    )
+    with op.batch_alter_table("projects") as batch_op:
+        batch_op.add_column(
+            sa.Column(
+                "trace_retention_policy_id",
+                sa.Integer,
+                sa.ForeignKey("project_trace_retention_policies.id", ondelete="SET NULL"),
+                nullable=True,
+            ),
+        )
+    op.create_index(
+        "ix_projects_trace_retention_policy_id",
+        "projects",
+        ["trace_retention_policy_id"],
+    )
+
+
+def downgrade() -> None:
+    op.drop_index("ix_projects_trace_retention_policy_id")
+    with op.batch_alter_table("projects") as batch_op:
+        batch_op.drop_column("trace_retention_policy_id")
+    op.drop_table("project_trace_retention_policies")

phoenix/db/models.py

@@ -1,6 +1,6 @@
 from datetime import datetime, timezone
 from enum import Enum
-from typing import Any, Iterable, Optional, Sequence, TypedDict, cast
+from typing import Any, Iterable, Literal, Optional, Sequence, TypedDict, cast
 
 import sqlalchemy.sql as sql
 from openinference.semconv.trace import RerankerAttributes, SpanAttributes
@@ -45,8 +45,15 @@ from sqlalchemy.sql.functions import coalesce
 
 from phoenix.config import get_env_database_schema
 from phoenix.datetime_utils import normalize_datetime
+from phoenix.db.types.annotation_configs import (
+    AnnotationConfig as AnnotationConfigModel,
+)
+from phoenix.db.types.annotation_configs import (
+    AnnotationConfigType,
+)
 from phoenix.db.types.identifier import Identifier
 from phoenix.db.types.model_provider import ModelProvider
+from phoenix.db.types.trace_retention import TraceRetentionCronExpression, TraceRetentionRule
 from phoenix.server.api.helpers.prompts.models import (
     PromptInvocationParameters,
     PromptInvocationParametersRootModel,
@@ -267,7 +274,7 @@ class _PromptTemplate(TypeDecorator[PromptTemplate]):
 class _Tools(TypeDecorator[PromptTools]):
     # See # See https://docs.sqlalchemy.org/en/20/core/custom_types.html
     cache_ok = True
-    impl = JSON_
+    impl = JSON
 
     def process_bind_param(
         self, value: Optional[PromptTools], _: Dialect
@@ -283,7 +290,7 @@ class _Tools(TypeDecorator[PromptTools]):
 class _PromptResponseFormat(TypeDecorator[PromptResponseFormat]):
     # See https://docs.sqlalchemy.org/en/20/core/custom_types.html
     cache_ok = True
-    impl = JSON_
+    impl = JSON
 
     def process_bind_param(
         self, value: Optional[PromptResponseFormat], _: Dialect
@@ -332,6 +339,60 @@ class _TemplateFormat(TypeDecorator[PromptTemplateFormat]):
         return None if value is None else PromptTemplateFormat(value)
 
 
+class _TraceRetentionCronExpression(TypeDecorator[TraceRetentionCronExpression]):
+    # See # See https://docs.sqlalchemy.org/en/20/core/custom_types.html
+    cache_ok = True
+    impl = String
+
+    def process_bind_param(
+        self, value: Optional[TraceRetentionCronExpression], _: Dialect
+    ) -> Optional[str]:
+        assert isinstance(value, TraceRetentionCronExpression)
+        assert isinstance(ans := value.model_dump(), str)
+        return ans
+
+    def process_result_value(
+        self, value: Optional[str], _: Dialect
+    ) -> Optional[TraceRetentionCronExpression]:
+        assert value and isinstance(value, str)
+        return TraceRetentionCronExpression.model_validate(value)
+
+
+class _TraceRetentionRule(TypeDecorator[TraceRetentionRule]):
+    # See # See https://docs.sqlalchemy.org/en/20/core/custom_types.html
+    cache_ok = True
+    impl = JSON_
+
+    def process_bind_param(
+        self, value: Optional[TraceRetentionRule], _: Dialect
+    ) -> Optional[dict[str, Any]]:
+        assert isinstance(value, TraceRetentionRule)
+        assert isinstance(ans := value.model_dump(), dict)
+        return ans
+
+    def process_result_value(
+        self, value: Optional[dict[str, Any]], _: Dialect
+    ) -> Optional[TraceRetentionRule]:
+        assert value and isinstance(value, dict)
+        return TraceRetentionRule.model_validate(value)
+
+
+class _AnnotationConfig(TypeDecorator[AnnotationConfigType]):
+    # See # See https://docs.sqlalchemy.org/en/20/core/custom_types.html
+    cache_ok = True
+    impl = JSON_
+
+    def process_bind_param(
+        self, value: Optional[AnnotationConfigType], _: Dialect
+    ) -> Optional[dict[str, Any]]:
+        return AnnotationConfigModel(root=value).model_dump() if value is not None else None
+
+    def process_result_value(
+        self, value: Optional[str], _: Dialect
+    ) -> Optional[AnnotationConfigType]:
+        return AnnotationConfigModel.model_validate(value).root if value is not None else None
+
+
 class ExperimentRunOutput(TypedDict, total=False):
     task_output: Any
 
@@ -357,6 +418,19 @@ class Base(DeclarativeBase):
     }
 
 
+class ProjectTraceRetentionPolicy(Base):
+    __tablename__ = "project_trace_retention_policies"
+    id: Mapped[int] = mapped_column(Integer, primary_key=True)
+    name: Mapped[str] = mapped_column(String, nullable=False)
+    cron_expression: Mapped[TraceRetentionCronExpression] = mapped_column(
+        _TraceRetentionCronExpression, nullable=False
+    )
+    rule: Mapped[TraceRetentionRule] = mapped_column(_TraceRetentionRule, nullable=False)
+    projects: Mapped[list["Project"]] = relationship(
+        "Project", back_populates="trace_retention_policy", uselist=True
+    )
+
+
 class Project(Base):
     __tablename__ = "projects"
     name: Mapped[str]
@@ -374,7 +448,15 @@ class Project(Base):
     updated_at: Mapped[datetime] = mapped_column(
         UtcTimeStamp, server_default=func.now(), onupdate=func.now()
     )
-
+    trace_retention_policy_id: Mapped[Optional[int]] = mapped_column(
+        ForeignKey("project_trace_retention_policies.id", ondelete="SET NULL"),
+        nullable=True,
+        index=True,
+    )
+    trace_retention_policy: Mapped[Optional[ProjectTraceRetentionPolicy]] = relationship(
+        "ProjectTraceRetentionPolicy",
+        back_populates="projects",
+    )
     traces: WriteOnlyMapped[list["Trace"]] = relationship(
         "Trace",
         back_populates="project",
@@ -602,6 +684,7 @@ class Span(Base):
         )
 
     trace: Mapped["Trace"] = relationship("Trace", back_populates="spans")
+    span_annotations: Mapped[list["SpanAnnotation"]] = relationship(back_populates="span")
     document_annotations: Mapped[list["DocumentAnnotation"]] = relationship(back_populates="span")
     dataset_examples: Mapped[list["DatasetExample"]] = relationship(back_populates="span")
 
@@ -732,17 +815,31 @@ class SpanAnnotation(Base):
     score: Mapped[Optional[float]] = mapped_column(Float, index=True)
     explanation: Mapped[Optional[str]]
     metadata_: Mapped[dict[str, Any]] = mapped_column("metadata")
-    annotator_kind: Mapped[str] = mapped_column(
-        CheckConstraint("annotator_kind IN ('LLM', 'HUMAN')", name="valid_annotator_kind"),
+    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"))
+
+    span: Mapped["Span"] = relationship(back_populates="span_annotations")
+    user: Mapped[Optional["User"]] = relationship("User")
+
     __table_args__ = (
         UniqueConstraint(
             "name",
             "span_rowid",
+            "identifier",
         ),
     )
 
@@ -758,17 +855,28 @@ class TraceAnnotation(Base):
     score: Mapped[Optional[float]] = mapped_column(Float, index=True)
     explanation: Mapped[Optional[str]]
     metadata_: Mapped[dict[str, Any]] = mapped_column("metadata")
-    annotator_kind: Mapped[str] = mapped_column(
-        CheckConstraint("annotator_kind IN ('LLM', 'HUMAN')", name="valid_annotator_kind"),
+    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",
             "trace_rowid",
+            "identifier",
         ),
     )
 
@@ -785,13 +893,23 @@ class DocumentAnnotation(Base):
     score: Mapped[Optional[float]] = mapped_column(Float, index=True)
     explanation: Mapped[Optional[str]]
     metadata_: Mapped[dict[str, Any]] = mapped_column("metadata")
-    annotator_kind: Mapped[str] = mapped_column(
-        CheckConstraint("annotator_kind IN ('LLM', 'HUMAN')", name="valid_annotator_kind"),
+    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"))
+
     span: Mapped["Span"] = relationship(back_populates="document_annotations")
 
     __table_args__ = (
@@ -799,6 +917,7 @@ class DocumentAnnotation(Base):
             "name",
             "span_rowid",
             "document_position",
+            "identifier",
         ),
     )
 
@@ -1301,3 +1420,25 @@ class PromptVersionTag(Base):
     )
 
     __table_args__ = (UniqueConstraint("name", "prompt_id"),)
+
+
+class AnnotationConfig(Base):
+    __tablename__ = "annotation_configs"
+
+    id: Mapped[int] = mapped_column(primary_key=True)
+    name: Mapped[str] = mapped_column(String, nullable=False, unique=True)
+    config: Mapped[AnnotationConfigType] = mapped_column(_AnnotationConfig, nullable=False)
+
+
+class ProjectAnnotationConfig(Base):
+    __tablename__ = "project_annotation_configs"
+
+    id: Mapped[int] = mapped_column(primary_key=True)
+    project_id: Mapped[int] = mapped_column(
+        ForeignKey("projects.id", ondelete="CASCADE"), nullable=False, index=True
+    )
+    annotation_config_id: Mapped[int] = mapped_column(
+        ForeignKey("annotation_configs.id", ondelete="CASCADE"), nullable=False, index=True
+    )
+
+    __table_args__ = (UniqueConstraint("project_id", "annotation_config_id"),)

phoenix/db/types/annotation_configs.py

@@ -0,0 +1,97 @@
+from enum import Enum
+from typing import Annotated, Literal, Optional, Union
+
+from pydantic import AfterValidator, Field, RootModel, model_validator
+from typing_extensions import Self, TypeAlias
+
+from .db_models import DBBaseModel
+
+
+class AnnotationType(Enum):
+    CATEGORICAL = "CATEGORICAL"
+    CONTINUOUS = "CONTINUOUS"
+    FREEFORM = "FREEFORM"
+
+
+class OptimizationDirection(Enum):
+    MINIMIZE = "MINIMIZE"
+    MAXIMIZE = "MAXIMIZE"
+    NONE = "NONE"
+
+
+class _BaseAnnotationConfig(DBBaseModel):
+    description: Optional[str] = None
+
+
+def _categorical_value_label_is_non_empty_string(label: str) -> str:
+    if not label:
+        raise ValueError("Label must be non-empty")
+    return label
+
+
+class CategoricalAnnotationValue(DBBaseModel):
+    label: Annotated[str, AfterValidator(_categorical_value_label_is_non_empty_string)]
+    score: Optional[float] = None
+
+
+def _categorical_values_are_non_empty_list(
+    values: list[CategoricalAnnotationValue],
+) -> list[CategoricalAnnotationValue]:
+    if not values:
+        raise ValueError("Values must be non-empty")
+    return values
+
+
+def _categorical_values_have_unique_labels(
+    values: list[CategoricalAnnotationValue],
+) -> list[CategoricalAnnotationValue]:
+    labels = set()
+    for value in values:
+        label = value.label
+        if label in labels:
+            raise ValueError(
+                f'Values for categorical annotation config has duplicate label: "{label}"'
+            )
+        labels.add(label)
+    return values
+
+
+class CategoricalAnnotationConfig(_BaseAnnotationConfig):
+    type: Literal[AnnotationType.CATEGORICAL.value]  # type: ignore[name-defined]
+    optimization_direction: OptimizationDirection
+    values: Annotated[
+        list[CategoricalAnnotationValue],
+        AfterValidator(_categorical_values_are_non_empty_list),
+        AfterValidator(_categorical_values_have_unique_labels),
+    ]
+
+
+class ContinuousAnnotationConfig(_BaseAnnotationConfig):
+    type: Literal[AnnotationType.CONTINUOUS.value]  # type: ignore[name-defined]
+    optimization_direction: OptimizationDirection
+    lower_bound: Optional[float] = None
+    upper_bound: Optional[float] = None
+
+    @model_validator(mode="after")
+    def check_bounds(self) -> Self:
+        if (
+            self.lower_bound is not None
+            and self.upper_bound is not None
+            and self.lower_bound >= self.upper_bound
+        ):
+            raise ValueError("Lower bound must be strictly less than upper bound")
+        return self
+
+
+class FreeformAnnotationConfig(_BaseAnnotationConfig):
+    type: Literal[AnnotationType.FREEFORM.value]  # type: ignore[name-defined]
+
+
+AnnotationConfigType: TypeAlias = Annotated[
+    Union[CategoricalAnnotationConfig, ContinuousAnnotationConfig, FreeformAnnotationConfig],
+    Field(..., discriminator="type"),
+]
+
+
+class AnnotationConfig(RootModel[AnnotationConfigType]):
+    root: AnnotationConfigType

phoenix/db/types/db_models.py

@@ -0,0 +1,41 @@
+from typing import Any
+
+from pydantic import BaseModel, ConfigDict
+
+
+class DBBaseModel(BaseModel):
+    """
+    A base Pydantic model suitable for use with JSON columns in the database.
+    """
+
+    model_config = ConfigDict(
+        extra="forbid",  # disallow extra attributes
+        use_enum_values=True,
+        validate_assignment=True,
+    )
+
+    def __init__(self, *args: Any, **kwargs: Any) -> None:
+        kwargs = {k: v for k, v in kwargs.items() if v is not UNDEFINED}
+        super().__init__(*args, **kwargs)
+
+    def model_dump(self, *args: Any, **kwargs: Any) -> dict[str, Any]:
+        return super().model_dump(*args, exclude_unset=True, by_alias=True, **kwargs)
+
+
+class Undefined:
+    """
+    A singleton class that represents an unset or undefined value. Needed since Pydantic
+    can't natively distinguish between an undefined value and a value that is set to
+    None.
+    """
+
+    def __new__(cls) -> Any:
+        if not hasattr(cls, "_instance"):
+            cls._instance = super().__new__(cls)
+        return cls._instance
+
+    def __bool__(self) -> bool:
+        return False
+
+
+UNDEFINED: Any = Undefined()