From 5d229066a5634e68a4c84271dc806cfadfe1431a Mon Sep 17 00:00:00 2001 From: Mikyo King Date: Fri, 5 Apr 2024 17:27:10 -0600 Subject: [PATCH] feat: annotation table persistence Update src/phoenix/db/models.py small update --- .../migrations/versions/cf03bd6bae1d_init.py | 1 + src/phoenix/db/models.py | 19 +++++++++++++++++++ 2 files changed, 20 insertions(+) diff --git a/src/phoenix/db/migrations/versions/cf03bd6bae1d_init.py b/src/phoenix/db/migrations/versions/cf03bd6bae1d_init.py index e0a997531cf..b356485f2d7 100644 --- a/src/phoenix/db/migrations/versions/cf03bd6bae1d_init.py +++ b/src/phoenix/db/migrations/versions/cf03bd6bae1d_init.py @@ -72,6 +72,7 @@ def upgrade() -> None: sa.Column("cumulative_llm_token_count_prompt", sa.Integer, nullable=False), sa.Column("cumulative_llm_token_count_completion", sa.Integer, nullable=False), ) + op.bulk_insert( projects_table, [ diff --git a/src/phoenix/db/models.py b/src/phoenix/db/models.py index 5f10f219311..89d2fb4beff 100644 --- a/src/phoenix/db/models.py +++ b/src/phoenix/db/models.py @@ -133,3 +133,22 @@ async def init_models(engine: AsyncEngine) -> None: description="default project", ) ) + + +class SpanAnnotation(Base): + __tablename__ = "span_annotations" + id: Mapped[int] = mapped_column(primary_key=True) + span_rowid: Mapped[int] = mapped_column(ForeignKey("spans.id"), required=True) + name: Mapped[str] = mapped_column(required=True) + label: Mapped[str] = mapped_column() + score: Mapped[float] = mapped_column() + explanation: Mapped[str] = mapped_column() + meta: Mapped[Dict[str, Any]] = mapped_column(JSON, required=True, key="metadata") + __table_args__ = ( + UniqueConstraint( + "span_rowid", + "name", + name="uq_span_annotations_span_rowid_name", + sqlite_on_conflict="UPDATE", + ), + )