105 lines
5.4 KiB
Python
105 lines
5.4 KiB
Python
"""extend prompt eval replay
|
|
|
|
Revision ID: 20260430_extend_prompt_eval_replay
|
|
Revises: 20260430_add_prompt_replay_runs
|
|
Create Date: 2026-04-30
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
from alembic import op
|
|
import sqlalchemy as sa
|
|
|
|
try:
|
|
from sqlalchemy.dialects.postgresql import JSONB
|
|
except Exception: # pragma: no cover
|
|
JSONB = sa.JSON
|
|
|
|
|
|
revision = "20260430_extend_prompt_eval_replay"
|
|
down_revision = "20260430_add_prompt_replay_runs"
|
|
branch_labels = None
|
|
depends_on = None
|
|
|
|
|
|
def _json_type(bind):
|
|
return JSONB() if bind.dialect.name == "postgresql" else sa.JSON()
|
|
|
|
|
|
def _table_exists(bind, table_name: str) -> bool:
|
|
return sa.inspect(bind).has_table(table_name)
|
|
|
|
|
|
def _column_exists(bind, table_name: str, column_name: str) -> bool:
|
|
if not _table_exists(bind, table_name):
|
|
return False
|
|
return column_name in {str(column.get("name") or "") for column in sa.inspect(bind).get_columns(table_name)}
|
|
|
|
|
|
def _index_exists(bind, table_name: str, index_name: str) -> bool:
|
|
if not _table_exists(bind, table_name):
|
|
return False
|
|
return index_name in {str(index.get("name") or "") for index in sa.inspect(bind).get_indexes(table_name)}
|
|
|
|
|
|
def _add_column_if_missing(bind, table_name: str, column: sa.Column) -> None:
|
|
if not _column_exists(bind, table_name, str(column.name)):
|
|
op.add_column(table_name, column)
|
|
|
|
|
|
def _create_index_if_missing(bind, index_name: str, table_name: str, columns: list[str]) -> None:
|
|
if not _index_exists(bind, table_name, index_name):
|
|
op.create_index(index_name, table_name, columns, unique=False)
|
|
|
|
|
|
def upgrade() -> None:
|
|
bind = op.get_bind()
|
|
if _table_exists(bind, "prompt_golden_cases"):
|
|
_add_column_if_missing(bind, "prompt_golden_cases", sa.Column("case_type", sa.String(length=32), nullable=False, server_default="regression"))
|
|
_add_column_if_missing(bind, "prompt_golden_cases", sa.Column("priority", sa.String(length=24), nullable=False, server_default="blocking"))
|
|
_add_column_if_missing(bind, "prompt_golden_cases", sa.Column("scenario_key", sa.String(length=128), nullable=True))
|
|
_add_column_if_missing(bind, "prompt_golden_cases", sa.Column("tags", _json_type(bind), nullable=True))
|
|
_add_column_if_missing(bind, "prompt_golden_cases", sa.Column("last_run_status", sa.String(length=24), nullable=True))
|
|
_add_column_if_missing(bind, "prompt_golden_cases", sa.Column("last_failed_at", sa.String(length=64), nullable=True))
|
|
_create_index_if_missing(bind, "ix_prompt_golden_cases_case_type", "prompt_golden_cases", ["case_type"])
|
|
_create_index_if_missing(bind, "ix_prompt_golden_cases_priority", "prompt_golden_cases", ["priority"])
|
|
_create_index_if_missing(bind, "ix_prompt_golden_cases_scenario_key", "prompt_golden_cases", ["scenario_key"])
|
|
_create_index_if_missing(bind, "ix_prompt_golden_cases_last_run_status", "prompt_golden_cases", ["last_run_status"])
|
|
|
|
if _table_exists(bind, "prompt_replay_case_results"):
|
|
_add_column_if_missing(bind, "prompt_replay_case_results", sa.Column("baseline_ai_call_log_id", sa.Integer(), nullable=True))
|
|
_add_column_if_missing(bind, "prompt_replay_case_results", sa.Column("candidate_ai_call_log_id", sa.Integer(), nullable=True))
|
|
_add_column_if_missing(bind, "prompt_replay_case_results", sa.Column("judge_score", sa.Float(), nullable=True))
|
|
_add_column_if_missing(bind, "prompt_replay_case_results", sa.Column("judge_label", sa.String(length=32), nullable=True))
|
|
_add_column_if_missing(bind, "prompt_replay_case_results", sa.Column("judge_reason", sa.Text(), nullable=True))
|
|
_create_index_if_missing(bind, "ix_prompt_replay_case_results_judge_label", "prompt_replay_case_results", ["judge_label"])
|
|
_create_index_if_missing(bind, "ix_prompt_replay_case_results_baseline_ai_call_log_id", "prompt_replay_case_results", ["baseline_ai_call_log_id"])
|
|
_create_index_if_missing(bind, "ix_prompt_replay_case_results_candidate_ai_call_log_id", "prompt_replay_case_results", ["candidate_ai_call_log_id"])
|
|
|
|
|
|
def downgrade() -> None:
|
|
bind = op.get_bind()
|
|
if _table_exists(bind, "prompt_replay_case_results"):
|
|
for index_name in (
|
|
"ix_prompt_replay_case_results_candidate_ai_call_log_id",
|
|
"ix_prompt_replay_case_results_baseline_ai_call_log_id",
|
|
"ix_prompt_replay_case_results_judge_label",
|
|
):
|
|
if _index_exists(bind, "prompt_replay_case_results", index_name):
|
|
op.drop_index(index_name, table_name="prompt_replay_case_results")
|
|
for column_name in ("judge_reason", "judge_label", "judge_score", "candidate_ai_call_log_id", "baseline_ai_call_log_id"):
|
|
if _column_exists(bind, "prompt_replay_case_results", column_name):
|
|
op.drop_column("prompt_replay_case_results", column_name)
|
|
if _table_exists(bind, "prompt_golden_cases"):
|
|
for index_name in (
|
|
"ix_prompt_golden_cases_last_run_status",
|
|
"ix_prompt_golden_cases_scenario_key",
|
|
"ix_prompt_golden_cases_priority",
|
|
"ix_prompt_golden_cases_case_type",
|
|
):
|
|
if _index_exists(bind, "prompt_golden_cases", index_name):
|
|
op.drop_index(index_name, table_name="prompt_golden_cases")
|
|
for column_name in ("last_failed_at", "last_run_status", "tags", "scenario_key", "priority", "case_type"):
|
|
if _column_exists(bind, "prompt_golden_cases", column_name):
|
|
op.drop_column("prompt_golden_cases", column_name)
|