588 lines
28 KiB
Python
588 lines
28 KiB
Python
"""add reply iteration and training pipeline tables
|
|
|
|
Revision ID: 20260327_add_reply_iteration_training_tables
|
|
Revises: 20260325_add_feishu_reply_outbox_sent_message_id, 20260327_add_qa_record_followup_anchor_fields
|
|
Create Date: 2026-03-27
|
|
"""
|
|
|
|
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 = "20260327_add_reply_iteration_training_tables"
|
|
down_revision = (
|
|
"20260325_add_feishu_reply_outbox_sent_message_id",
|
|
"20260327_add_qa_record_followup_anchor_fields",
|
|
)
|
|
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:
|
|
inspector = sa.inspect(bind)
|
|
try:
|
|
columns = inspector.get_columns(table_name)
|
|
except Exception:
|
|
return False
|
|
return any(str(column.get("name")) == column_name for column in columns)
|
|
|
|
|
|
def _index_exists(bind, table_name: str, index_name: str) -> bool:
|
|
inspector = sa.inspect(bind)
|
|
try:
|
|
indexes = inspector.get_indexes(table_name)
|
|
except Exception:
|
|
return False
|
|
return any(str(index.get("name")) == index_name for index in indexes)
|
|
|
|
|
|
def _create_index_if_missing(bind, table_name: str, index_name: str, columns: list[str], *, unique: bool = False) -> None:
|
|
if not _index_exists(bind, table_name, index_name):
|
|
op.create_index(index_name, table_name, columns, unique=unique)
|
|
|
|
|
|
def _drop_index_if_exists(bind, table_name: str, index_name: str) -> None:
|
|
if _index_exists(bind, table_name, index_name):
|
|
op.drop_index(index_name, table_name=table_name)
|
|
|
|
|
|
def _base_columns() -> list[sa.Column]:
|
|
return [
|
|
sa.Column("tenant_id", sa.BigInteger(), nullable=False),
|
|
sa.Column("created_at", sa.TIMESTAMP(), nullable=False, server_default=sa.func.now()),
|
|
sa.Column("updated_at", sa.TIMESTAMP(), nullable=False, server_default=sa.func.now()),
|
|
]
|
|
|
|
|
|
def _create_quality_assets(bind) -> None:
|
|
if _table_exists(bind, "quality_assets"):
|
|
return
|
|
op.create_table(
|
|
"quality_assets",
|
|
sa.Column("id", sa.Integer(), primary_key=True, autoincrement=True),
|
|
sa.Column("project_id", sa.BigInteger(), nullable=True),
|
|
sa.Column("asset_type", sa.VARCHAR(length=16), nullable=False),
|
|
sa.Column("status", sa.VARCHAR(length=16), nullable=False, server_default="approved"),
|
|
sa.Column("source_feedback_record_id", sa.Integer(), sa.ForeignKey("feedback_center_records.id", ondelete="SET NULL"), nullable=True),
|
|
sa.Column("qa_record_id", sa.Integer(), sa.ForeignKey("qa_records.id", ondelete="SET NULL"), nullable=True),
|
|
sa.Column("ai_call_log_id", sa.Integer(), nullable=True),
|
|
sa.Column("prompt_key", sa.VARCHAR(length=128), nullable=True),
|
|
sa.Column("workflow_type", sa.VARCHAR(length=64), nullable=True),
|
|
sa.Column("asked_via", sa.VARCHAR(length=32), nullable=True),
|
|
sa.Column("planned_tier", sa.VARCHAR(length=16), nullable=True),
|
|
sa.Column("executed_tier", sa.VARCHAR(length=16), nullable=True),
|
|
sa.Column("route_source", sa.VARCHAR(length=32), nullable=True),
|
|
sa.Column("route_rule_name", sa.VARCHAR(length=255), nullable=True),
|
|
sa.Column("question_text", sa.Text(), nullable=False),
|
|
sa.Column("answer_text", sa.Text(), nullable=False),
|
|
sa.Column("corrected_answer", sa.Text(), nullable=True),
|
|
sa.Column("source_snapshot", _json_type(bind), nullable=True),
|
|
sa.Column("tags", _json_type(bind), nullable=True),
|
|
sa.Column("notes", sa.Text(), nullable=True),
|
|
sa.Column("approved_by", sa.BigInteger(), nullable=True),
|
|
sa.Column("approved_at", sa.VARCHAR(length=64), nullable=True),
|
|
*_base_columns(),
|
|
)
|
|
_create_index_if_missing(bind, "quality_assets", "ix_quality_assets_tenant_id", ["tenant_id"])
|
|
_create_index_if_missing(bind, "quality_assets", "ix_quality_assets_project_id", ["project_id"])
|
|
_create_index_if_missing(bind, "quality_assets", "ix_quality_assets_asset_type", ["asset_type"])
|
|
_create_index_if_missing(bind, "quality_assets", "ix_quality_assets_status", ["status"])
|
|
_create_index_if_missing(bind, "quality_assets", "ix_quality_assets_source_feedback_record_id", ["source_feedback_record_id"])
|
|
_create_index_if_missing(bind, "quality_assets", "ix_quality_assets_qa_record_id", ["qa_record_id"])
|
|
_create_index_if_missing(bind, "quality_assets", "ix_quality_assets_ai_call_log_id", ["ai_call_log_id"])
|
|
_create_index_if_missing(bind, "quality_assets", "ix_quality_assets_prompt_key", ["prompt_key"])
|
|
_create_index_if_missing(bind, "quality_assets", "ix_quality_assets_workflow_type", ["workflow_type"])
|
|
_create_index_if_missing(bind, "quality_assets", "ix_quality_assets_asked_via", ["asked_via"])
|
|
_create_index_if_missing(bind, "quality_assets", "ix_quality_assets_planned_tier", ["planned_tier"])
|
|
_create_index_if_missing(bind, "quality_assets", "ix_quality_assets_executed_tier", ["executed_tier"])
|
|
_create_index_if_missing(bind, "quality_assets", "ix_quality_assets_approved_by", ["approved_by"])
|
|
|
|
|
|
def _create_quality_asset_pairs(bind) -> None:
|
|
if _table_exists(bind, "quality_asset_pairs"):
|
|
return
|
|
op.create_table(
|
|
"quality_asset_pairs",
|
|
sa.Column("id", sa.Integer(), primary_key=True, autoincrement=True),
|
|
sa.Column("project_id", sa.BigInteger(), nullable=True),
|
|
sa.Column("pair_type", sa.VARCHAR(length=24), nullable=False, server_default="preference"),
|
|
sa.Column("status", sa.VARCHAR(length=16), nullable=False, server_default="approved"),
|
|
sa.Column("chosen_asset_id", sa.Integer(), sa.ForeignKey("quality_assets.id", ondelete="CASCADE"), nullable=False),
|
|
sa.Column("rejected_asset_id", sa.Integer(), sa.ForeignKey("quality_assets.id", ondelete="CASCADE"), nullable=False),
|
|
sa.Column("reason_text", sa.Text(), nullable=True),
|
|
sa.Column("approved_by", sa.BigInteger(), nullable=True),
|
|
sa.Column("approved_at", sa.VARCHAR(length=64), nullable=True),
|
|
*_base_columns(),
|
|
)
|
|
_create_index_if_missing(bind, "quality_asset_pairs", "ix_quality_asset_pairs_tenant_id", ["tenant_id"])
|
|
_create_index_if_missing(bind, "quality_asset_pairs", "ix_quality_asset_pairs_project_id", ["project_id"])
|
|
_create_index_if_missing(bind, "quality_asset_pairs", "ix_quality_asset_pairs_pair_type", ["pair_type"])
|
|
_create_index_if_missing(bind, "quality_asset_pairs", "ix_quality_asset_pairs_status", ["status"])
|
|
_create_index_if_missing(bind, "quality_asset_pairs", "ix_quality_asset_pairs_chosen_asset_id", ["chosen_asset_id"])
|
|
_create_index_if_missing(bind, "quality_asset_pairs", "ix_quality_asset_pairs_rejected_asset_id", ["rejected_asset_id"])
|
|
_create_index_if_missing(bind, "quality_asset_pairs", "ix_quality_asset_pairs_approved_by", ["approved_by"])
|
|
|
|
|
|
def _create_judge_runs(bind) -> None:
|
|
if _table_exists(bind, "judge_runs"):
|
|
return
|
|
op.create_table(
|
|
"judge_runs",
|
|
sa.Column("id", sa.Integer(), primary_key=True, autoincrement=True),
|
|
sa.Column("project_id", sa.BigInteger(), nullable=True),
|
|
sa.Column("asset_id", sa.Integer(), sa.ForeignKey("quality_assets.id", ondelete="CASCADE"), nullable=True),
|
|
sa.Column("pair_id", sa.Integer(), sa.ForeignKey("quality_asset_pairs.id", ondelete="CASCADE"), nullable=True),
|
|
sa.Column("prompt_key", sa.VARCHAR(length=128), nullable=True),
|
|
sa.Column("workflow_type", sa.VARCHAR(length=64), nullable=True),
|
|
sa.Column("asked_via", sa.VARCHAR(length=32), nullable=True),
|
|
sa.Column("planned_tier", sa.VARCHAR(length=16), nullable=True),
|
|
sa.Column("executed_tier", sa.VARCHAR(length=16), nullable=True),
|
|
sa.Column("route_rule_name", sa.VARCHAR(length=255), nullable=True),
|
|
sa.Column("judge_status", sa.VARCHAR(length=16), nullable=False, server_default="completed"),
|
|
sa.Column("bucket_key", sa.VARCHAR(length=255), nullable=False),
|
|
sa.Column("rule_score", sa.Float(), nullable=True),
|
|
sa.Column("llm_score", sa.Float(), nullable=True),
|
|
sa.Column("overall_score", sa.Float(), nullable=True),
|
|
sa.Column("findings", _json_type(bind), nullable=True),
|
|
sa.Column("llm_evidence", _json_type(bind), nullable=True),
|
|
sa.Column("metadata_json", _json_type(bind), nullable=True),
|
|
*_base_columns(),
|
|
)
|
|
for index_name, columns in (
|
|
("ix_judge_runs_tenant_id", ["tenant_id"]),
|
|
("ix_judge_runs_project_id", ["project_id"]),
|
|
("ix_judge_runs_asset_id", ["asset_id"]),
|
|
("ix_judge_runs_pair_id", ["pair_id"]),
|
|
("ix_judge_runs_prompt_key", ["prompt_key"]),
|
|
("ix_judge_runs_workflow_type", ["workflow_type"]),
|
|
("ix_judge_runs_asked_via", ["asked_via"]),
|
|
("ix_judge_runs_planned_tier", ["planned_tier"]),
|
|
("ix_judge_runs_executed_tier", ["executed_tier"]),
|
|
("ix_judge_runs_judge_status", ["judge_status"]),
|
|
("ix_judge_runs_bucket_key", ["bucket_key"]),
|
|
):
|
|
_create_index_if_missing(bind, "judge_runs", index_name, columns)
|
|
|
|
|
|
def _create_optimizer_candidates(bind) -> None:
|
|
if _table_exists(bind, "optimizer_candidates"):
|
|
return
|
|
op.create_table(
|
|
"optimizer_candidates",
|
|
sa.Column("id", sa.Integer(), primary_key=True, autoincrement=True),
|
|
sa.Column("project_id", sa.BigInteger(), nullable=True),
|
|
sa.Column("judge_run_id", sa.Integer(), sa.ForeignKey("judge_runs.id", ondelete="CASCADE"), nullable=False),
|
|
sa.Column("candidate_type", sa.VARCHAR(length=32), nullable=False),
|
|
sa.Column("target_type", sa.VARCHAR(length=32), nullable=False),
|
|
sa.Column("target_ref", sa.VARCHAR(length=255), nullable=False),
|
|
sa.Column("status", sa.VARCHAR(length=24), nullable=False, server_default="draft"),
|
|
sa.Column("suggestion", _json_type(bind), nullable=True),
|
|
sa.Column("evidence", _json_type(bind), nullable=True),
|
|
sa.Column("notes", sa.Text(), nullable=True),
|
|
*_base_columns(),
|
|
)
|
|
for index_name, columns in (
|
|
("ix_optimizer_candidates_tenant_id", ["tenant_id"]),
|
|
("ix_optimizer_candidates_project_id", ["project_id"]),
|
|
("ix_optimizer_candidates_judge_run_id", ["judge_run_id"]),
|
|
("ix_optimizer_candidates_candidate_type", ["candidate_type"]),
|
|
("ix_optimizer_candidates_target_type", ["target_type"]),
|
|
("ix_optimizer_candidates_target_ref", ["target_ref"]),
|
|
("ix_optimizer_candidates_status", ["status"]),
|
|
):
|
|
_create_index_if_missing(bind, "optimizer_candidates", index_name, columns)
|
|
|
|
|
|
def _create_prompt_releases(bind) -> None:
|
|
if _table_exists(bind, "prompt_releases"):
|
|
return
|
|
op.create_table(
|
|
"prompt_releases",
|
|
sa.Column("id", sa.Integer(), primary_key=True, autoincrement=True),
|
|
sa.Column("prompt_key", sa.VARCHAR(length=128), nullable=False),
|
|
sa.Column("version", sa.Integer(), nullable=False, server_default="1"),
|
|
sa.Column("body_snapshot", sa.Text(), nullable=False),
|
|
sa.Column("status", sa.VARCHAR(length=16), nullable=False, server_default="draft"),
|
|
sa.Column("notes", sa.Text(), nullable=True),
|
|
sa.Column("created_by", sa.BigInteger(), nullable=True),
|
|
sa.Column("activated_by", sa.BigInteger(), nullable=True),
|
|
sa.Column("activated_at", sa.VARCHAR(length=64), nullable=True),
|
|
*_base_columns(),
|
|
sa.UniqueConstraint("tenant_id", "prompt_key", "version", name="uq_prompt_releases_tenant_key_version"),
|
|
)
|
|
for index_name, columns in (
|
|
("ix_prompt_releases_tenant_id", ["tenant_id"]),
|
|
("ix_prompt_releases_prompt_key", ["prompt_key"]),
|
|
("ix_prompt_releases_status", ["status"]),
|
|
("ix_prompt_releases_created_by", ["created_by"]),
|
|
("ix_prompt_releases_activated_by", ["activated_by"]),
|
|
):
|
|
_create_index_if_missing(bind, "prompt_releases", index_name, columns)
|
|
|
|
|
|
def _create_training_datasets(bind) -> None:
|
|
if _table_exists(bind, "training_datasets"):
|
|
return
|
|
op.create_table(
|
|
"training_datasets",
|
|
sa.Column("id", sa.Integer(), primary_key=True, autoincrement=True),
|
|
sa.Column("project_id", sa.BigInteger(), nullable=True),
|
|
sa.Column("dataset_type", sa.VARCHAR(length=16), nullable=False),
|
|
sa.Column("version", sa.VARCHAR(length=64), nullable=False),
|
|
sa.Column("status", sa.VARCHAR(length=16), nullable=False, server_default="ready"),
|
|
sa.Column("sample_count", sa.Integer(), nullable=False, server_default="0"),
|
|
sa.Column("selection_rules", _json_type(bind), nullable=True),
|
|
sa.Column("channel_scope", _json_type(bind), nullable=True),
|
|
sa.Column("notes", sa.Text(), nullable=True),
|
|
sa.Column("created_by", sa.BigInteger(), nullable=True),
|
|
*_base_columns(),
|
|
sa.UniqueConstraint("tenant_id", "dataset_type", "version", name="uq_training_datasets_tenant_type_version"),
|
|
)
|
|
for index_name, columns in (
|
|
("ix_training_datasets_tenant_id", ["tenant_id"]),
|
|
("ix_training_datasets_project_id", ["project_id"]),
|
|
("ix_training_datasets_dataset_type", ["dataset_type"]),
|
|
("ix_training_datasets_version", ["version"]),
|
|
("ix_training_datasets_status", ["status"]),
|
|
("ix_training_datasets_created_by", ["created_by"]),
|
|
):
|
|
_create_index_if_missing(bind, "training_datasets", index_name, columns)
|
|
|
|
|
|
def _create_training_dataset_items(bind) -> None:
|
|
if _table_exists(bind, "training_dataset_items"):
|
|
return
|
|
op.create_table(
|
|
"training_dataset_items",
|
|
sa.Column("id", sa.Integer(), primary_key=True, autoincrement=True),
|
|
sa.Column("dataset_id", sa.Integer(), sa.ForeignKey("training_datasets.id", ondelete="CASCADE"), nullable=False),
|
|
sa.Column("quality_asset_id", sa.Integer(), sa.ForeignKey("quality_assets.id", ondelete="CASCADE"), nullable=True),
|
|
sa.Column("quality_asset_pair_id", sa.Integer(), sa.ForeignKey("quality_asset_pairs.id", ondelete="CASCADE"), nullable=True),
|
|
sa.Column("redacted_snapshot", _json_type(bind), nullable=True),
|
|
*_base_columns(),
|
|
)
|
|
for index_name, columns in (
|
|
("ix_training_dataset_items_tenant_id", ["tenant_id"]),
|
|
("ix_training_dataset_items_dataset_id", ["dataset_id"]),
|
|
("ix_training_dataset_items_quality_asset_id", ["quality_asset_id"]),
|
|
("ix_training_dataset_items_quality_asset_pair_id", ["quality_asset_pair_id"]),
|
|
):
|
|
_create_index_if_missing(bind, "training_dataset_items", index_name, columns)
|
|
|
|
|
|
def _create_training_jobs(bind) -> None:
|
|
if _table_exists(bind, "training_jobs"):
|
|
return
|
|
op.create_table(
|
|
"training_jobs",
|
|
sa.Column("id", sa.Integer(), primary_key=True, autoincrement=True),
|
|
sa.Column("project_id", sa.BigInteger(), nullable=True),
|
|
sa.Column("dataset_id", sa.Integer(), sa.ForeignKey("training_datasets.id", ondelete="CASCADE"), nullable=False),
|
|
sa.Column("job_type", sa.VARCHAR(length=16), nullable=False),
|
|
sa.Column("base_model_family", sa.VARCHAR(length=64), nullable=False),
|
|
sa.Column("target_slot", sa.VARCHAR(length=16), nullable=False),
|
|
sa.Column("status", sa.VARCHAR(length=16), nullable=False, server_default="queued"),
|
|
sa.Column("provider_backend", sa.VARCHAR(length=64), nullable=False),
|
|
sa.Column("artifact_ref", sa.VARCHAR(length=255), nullable=True),
|
|
sa.Column("metrics_json", _json_type(bind), nullable=True),
|
|
sa.Column("created_by", sa.BigInteger(), nullable=True),
|
|
*_base_columns(),
|
|
)
|
|
for index_name, columns in (
|
|
("ix_training_jobs_tenant_id", ["tenant_id"]),
|
|
("ix_training_jobs_project_id", ["project_id"]),
|
|
("ix_training_jobs_dataset_id", ["dataset_id"]),
|
|
("ix_training_jobs_job_type", ["job_type"]),
|
|
("ix_training_jobs_base_model_family", ["base_model_family"]),
|
|
("ix_training_jobs_target_slot", ["target_slot"]),
|
|
("ix_training_jobs_status", ["status"]),
|
|
("ix_training_jobs_created_by", ["created_by"]),
|
|
):
|
|
_create_index_if_missing(bind, "training_jobs", index_name, columns)
|
|
|
|
|
|
def _create_training_job_runs(bind) -> None:
|
|
if _table_exists(bind, "training_job_runs"):
|
|
return
|
|
op.create_table(
|
|
"training_job_runs",
|
|
sa.Column("id", sa.Integer(), primary_key=True, autoincrement=True),
|
|
sa.Column("training_job_id", sa.Integer(), sa.ForeignKey("training_jobs.id", ondelete="CASCADE"), nullable=False),
|
|
sa.Column("status", sa.VARCHAR(length=16), nullable=False, server_default="queued"),
|
|
sa.Column("provider_run_ref", sa.VARCHAR(length=255), nullable=True),
|
|
sa.Column("metrics_json", _json_type(bind), nullable=True),
|
|
*_base_columns(),
|
|
)
|
|
for index_name, columns in (
|
|
("ix_training_job_runs_tenant_id", ["tenant_id"]),
|
|
("ix_training_job_runs_training_job_id", ["training_job_id"]),
|
|
("ix_training_job_runs_status", ["status"]),
|
|
):
|
|
_create_index_if_missing(bind, "training_job_runs", index_name, columns)
|
|
|
|
|
|
def _create_training_evaluations(bind) -> None:
|
|
if _table_exists(bind, "training_evaluations"):
|
|
return
|
|
op.create_table(
|
|
"training_evaluations",
|
|
sa.Column("id", sa.Integer(), primary_key=True, autoincrement=True),
|
|
sa.Column("training_job_id", sa.Integer(), sa.ForeignKey("training_jobs.id", ondelete="CASCADE"), nullable=False),
|
|
sa.Column("evaluation_type", sa.VARCHAR(length=32), nullable=False),
|
|
sa.Column("status", sa.VARCHAR(length=16), nullable=False, server_default="pending"),
|
|
sa.Column("score", sa.Float(), nullable=True),
|
|
sa.Column("evaluation_summary", _json_type(bind), nullable=True),
|
|
*_base_columns(),
|
|
)
|
|
for index_name, columns in (
|
|
("ix_training_evaluations_tenant_id", ["tenant_id"]),
|
|
("ix_training_evaluations_training_job_id", ["training_job_id"]),
|
|
("ix_training_evaluations_evaluation_type", ["evaluation_type"]),
|
|
("ix_training_evaluations_status", ["status"]),
|
|
):
|
|
_create_index_if_missing(bind, "training_evaluations", index_name, columns)
|
|
|
|
|
|
def _create_trained_model_artifacts(bind) -> None:
|
|
if _table_exists(bind, "trained_model_artifacts"):
|
|
return
|
|
op.create_table(
|
|
"trained_model_artifacts",
|
|
sa.Column("id", sa.Integer(), primary_key=True, autoincrement=True),
|
|
sa.Column("training_job_id", sa.Integer(), sa.ForeignKey("training_jobs.id", ondelete="CASCADE"), nullable=False),
|
|
sa.Column("artifact_ref", sa.VARCHAR(length=255), nullable=False),
|
|
sa.Column("base_model_family", sa.VARCHAR(length=64), nullable=False),
|
|
sa.Column("target_slot", sa.VARCHAR(length=16), nullable=False),
|
|
sa.Column("status", sa.VARCHAR(length=16), nullable=False, server_default="pending_eval"),
|
|
sa.Column("metrics_json", _json_type(bind), nullable=True),
|
|
*_base_columns(),
|
|
)
|
|
for index_name, columns in (
|
|
("ix_trained_model_artifacts_tenant_id", ["tenant_id"]),
|
|
("ix_trained_model_artifacts_training_job_id", ["training_job_id"]),
|
|
("ix_trained_model_artifacts_artifact_ref", ["artifact_ref"]),
|
|
("ix_trained_model_artifacts_base_model_family", ["base_model_family"]),
|
|
("ix_trained_model_artifacts_target_slot", ["target_slot"]),
|
|
("ix_trained_model_artifacts_status", ["status"]),
|
|
):
|
|
_create_index_if_missing(bind, "trained_model_artifacts", index_name, columns)
|
|
|
|
|
|
def _create_trained_model_deployments(bind) -> None:
|
|
if _table_exists(bind, "trained_model_deployments"):
|
|
return
|
|
op.create_table(
|
|
"trained_model_deployments",
|
|
sa.Column("id", sa.Integer(), primary_key=True, autoincrement=True),
|
|
sa.Column("artifact_id", sa.Integer(), sa.ForeignKey("trained_model_artifacts.id", ondelete="CASCADE"), nullable=False),
|
|
sa.Column("deployment_status", sa.VARCHAR(length=16), nullable=False, server_default="candidate"),
|
|
sa.Column("bound_tier", sa.VARCHAR(length=16), nullable=False),
|
|
sa.Column("effective_scope", _json_type(bind), nullable=True),
|
|
sa.Column("decision", sa.VARCHAR(length=24), nullable=True),
|
|
sa.Column("notes", sa.Text(), nullable=True),
|
|
sa.Column("decided_by", sa.BigInteger(), nullable=True),
|
|
sa.Column("decided_at", sa.VARCHAR(length=64), nullable=True),
|
|
*_base_columns(),
|
|
)
|
|
for index_name, columns in (
|
|
("ix_trained_model_deployments_tenant_id", ["tenant_id"]),
|
|
("ix_trained_model_deployments_artifact_id", ["artifact_id"]),
|
|
("ix_trained_model_deployments_deployment_status", ["deployment_status"]),
|
|
("ix_trained_model_deployments_bound_tier", ["bound_tier"]),
|
|
("ix_trained_model_deployments_decision", ["decision"]),
|
|
("ix_trained_model_deployments_decided_by", ["decided_by"]),
|
|
):
|
|
_create_index_if_missing(bind, "trained_model_deployments", index_name, columns)
|
|
|
|
|
|
def _add_experiment_columns(bind) -> None:
|
|
if not _table_exists(bind, "experiments"):
|
|
return
|
|
additions = [
|
|
("experiment_type", sa.VARCHAR(length=64), True),
|
|
("baseline_release_id", sa.Integer(), True),
|
|
("candidate_release_id", sa.Integer(), True),
|
|
("scope_json", _json_type(bind), True),
|
|
("result_summary", _json_type(bind), True),
|
|
("decision", sa.VARCHAR(length=24), True),
|
|
]
|
|
for column_name, column_type, nullable in additions:
|
|
if _column_exists(bind, "experiments", column_name):
|
|
continue
|
|
with op.batch_alter_table("experiments") as batch_op:
|
|
batch_op.add_column(sa.Column(column_name, column_type, nullable=nullable))
|
|
for index_name, columns in (
|
|
("ix_experiments_experiment_type", ["experiment_type"]),
|
|
("ix_experiments_baseline_release_id", ["baseline_release_id"]),
|
|
("ix_experiments_candidate_release_id", ["candidate_release_id"]),
|
|
("ix_experiments_decision", ["decision"]),
|
|
):
|
|
_create_index_if_missing(bind, "experiments", index_name, columns)
|
|
|
|
|
|
def upgrade() -> None:
|
|
bind = op.get_bind()
|
|
_create_quality_assets(bind)
|
|
_create_quality_asset_pairs(bind)
|
|
_create_judge_runs(bind)
|
|
_create_optimizer_candidates(bind)
|
|
_create_prompt_releases(bind)
|
|
_create_training_datasets(bind)
|
|
_create_training_dataset_items(bind)
|
|
_create_training_jobs(bind)
|
|
_create_training_job_runs(bind)
|
|
_create_training_evaluations(bind)
|
|
_create_trained_model_artifacts(bind)
|
|
_create_trained_model_deployments(bind)
|
|
_add_experiment_columns(bind)
|
|
|
|
|
|
def downgrade() -> None:
|
|
bind = op.get_bind()
|
|
|
|
for index_name in (
|
|
"ix_experiments_decision",
|
|
"ix_experiments_candidate_release_id",
|
|
"ix_experiments_baseline_release_id",
|
|
"ix_experiments_experiment_type",
|
|
):
|
|
_drop_index_if_exists(bind, "experiments", index_name)
|
|
if _table_exists(bind, "experiments"):
|
|
for column_name in (
|
|
"decision",
|
|
"result_summary",
|
|
"scope_json",
|
|
"candidate_release_id",
|
|
"baseline_release_id",
|
|
"experiment_type",
|
|
):
|
|
if _column_exists(bind, "experiments", column_name):
|
|
with op.batch_alter_table("experiments") as batch_op:
|
|
batch_op.drop_column(column_name)
|
|
|
|
for table_name, index_names in (
|
|
("trained_model_deployments", [
|
|
"ix_trained_model_deployments_decided_by",
|
|
"ix_trained_model_deployments_decision",
|
|
"ix_trained_model_deployments_bound_tier",
|
|
"ix_trained_model_deployments_deployment_status",
|
|
"ix_trained_model_deployments_artifact_id",
|
|
"ix_trained_model_deployments_tenant_id",
|
|
]),
|
|
("trained_model_artifacts", [
|
|
"ix_trained_model_artifacts_status",
|
|
"ix_trained_model_artifacts_target_slot",
|
|
"ix_trained_model_artifacts_base_model_family",
|
|
"ix_trained_model_artifacts_artifact_ref",
|
|
"ix_trained_model_artifacts_training_job_id",
|
|
"ix_trained_model_artifacts_tenant_id",
|
|
]),
|
|
("training_evaluations", [
|
|
"ix_training_evaluations_status",
|
|
"ix_training_evaluations_evaluation_type",
|
|
"ix_training_evaluations_training_job_id",
|
|
"ix_training_evaluations_tenant_id",
|
|
]),
|
|
("training_job_runs", [
|
|
"ix_training_job_runs_status",
|
|
"ix_training_job_runs_training_job_id",
|
|
"ix_training_job_runs_tenant_id",
|
|
]),
|
|
("training_jobs", [
|
|
"ix_training_jobs_created_by",
|
|
"ix_training_jobs_status",
|
|
"ix_training_jobs_target_slot",
|
|
"ix_training_jobs_base_model_family",
|
|
"ix_training_jobs_job_type",
|
|
"ix_training_jobs_dataset_id",
|
|
"ix_training_jobs_project_id",
|
|
"ix_training_jobs_tenant_id",
|
|
]),
|
|
("training_dataset_items", [
|
|
"ix_training_dataset_items_quality_asset_pair_id",
|
|
"ix_training_dataset_items_quality_asset_id",
|
|
"ix_training_dataset_items_dataset_id",
|
|
"ix_training_dataset_items_tenant_id",
|
|
]),
|
|
("training_datasets", [
|
|
"ix_training_datasets_created_by",
|
|
"ix_training_datasets_status",
|
|
"ix_training_datasets_version",
|
|
"ix_training_datasets_dataset_type",
|
|
"ix_training_datasets_project_id",
|
|
"ix_training_datasets_tenant_id",
|
|
]),
|
|
("prompt_releases", [
|
|
"ix_prompt_releases_activated_by",
|
|
"ix_prompt_releases_created_by",
|
|
"ix_prompt_releases_status",
|
|
"ix_prompt_releases_prompt_key",
|
|
"ix_prompt_releases_tenant_id",
|
|
]),
|
|
("optimizer_candidates", [
|
|
"ix_optimizer_candidates_status",
|
|
"ix_optimizer_candidates_target_ref",
|
|
"ix_optimizer_candidates_target_type",
|
|
"ix_optimizer_candidates_candidate_type",
|
|
"ix_optimizer_candidates_judge_run_id",
|
|
"ix_optimizer_candidates_project_id",
|
|
"ix_optimizer_candidates_tenant_id",
|
|
]),
|
|
("judge_runs", [
|
|
"ix_judge_runs_bucket_key",
|
|
"ix_judge_runs_judge_status",
|
|
"ix_judge_runs_executed_tier",
|
|
"ix_judge_runs_planned_tier",
|
|
"ix_judge_runs_asked_via",
|
|
"ix_judge_runs_workflow_type",
|
|
"ix_judge_runs_prompt_key",
|
|
"ix_judge_runs_pair_id",
|
|
"ix_judge_runs_asset_id",
|
|
"ix_judge_runs_project_id",
|
|
"ix_judge_runs_tenant_id",
|
|
]),
|
|
("quality_asset_pairs", [
|
|
"ix_quality_asset_pairs_approved_by",
|
|
"ix_quality_asset_pairs_rejected_asset_id",
|
|
"ix_quality_asset_pairs_chosen_asset_id",
|
|
"ix_quality_asset_pairs_status",
|
|
"ix_quality_asset_pairs_pair_type",
|
|
"ix_quality_asset_pairs_project_id",
|
|
"ix_quality_asset_pairs_tenant_id",
|
|
]),
|
|
("quality_assets", [
|
|
"ix_quality_assets_approved_by",
|
|
"ix_quality_assets_executed_tier",
|
|
"ix_quality_assets_planned_tier",
|
|
"ix_quality_assets_asked_via",
|
|
"ix_quality_assets_workflow_type",
|
|
"ix_quality_assets_prompt_key",
|
|
"ix_quality_assets_ai_call_log_id",
|
|
"ix_quality_assets_qa_record_id",
|
|
"ix_quality_assets_source_feedback_record_id",
|
|
"ix_quality_assets_status",
|
|
"ix_quality_assets_asset_type",
|
|
"ix_quality_assets_project_id",
|
|
"ix_quality_assets_tenant_id",
|
|
]),
|
|
):
|
|
for index_name in index_names:
|
|
_drop_index_if_exists(bind, table_name, index_name)
|
|
if _table_exists(bind, table_name):
|
|
op.drop_table(table_name)
|