151 lines
4.2 KiB
Python
151 lines
4.2 KiB
Python
"""cleanup legacy model routing function names
|
|
|
|
Revision ID: 20260327_cleanup_model_routing_legacy_names
|
|
Revises: 20260326_add_ai_call_log_routing_fields
|
|
Create Date: 2026-03-27
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
from copy import deepcopy
|
|
from typing import Any
|
|
|
|
from alembic import op
|
|
import sqlalchemy as sa
|
|
|
|
|
|
revision = "20260327_cleanup_model_routing_legacy_names"
|
|
down_revision = "20260326_add_ai_call_log_routing_fields"
|
|
branch_labels = None
|
|
depends_on = None
|
|
|
|
_RENAME_KEYS = {
|
|
"project_draft_analysis_glm47_fallback": "ai_pmo_signal_title_refine",
|
|
"adaptive_rag_answer_deep_signal_batch_analysis": "ai_pmo_signal_batch_analysis",
|
|
}
|
|
_REMOVE_KEYS = {
|
|
"project_draft_charter_glm47_fallback",
|
|
"project_draft_signal_glm47_fallback",
|
|
"project_draft_task_chain_glm47_fallback",
|
|
}
|
|
|
|
|
|
def _table_exists(bind, table_name: str) -> bool:
|
|
return sa.inspect(bind).has_table(table_name)
|
|
|
|
|
|
def _integration_table(bind) -> sa.Table:
|
|
metadata = sa.MetaData()
|
|
return sa.Table(
|
|
"integrations",
|
|
metadata,
|
|
sa.Column("id", sa.Integer(), primary_key=True),
|
|
sa.Column("type", sa.String(length=100)),
|
|
sa.Column("config", sa.JSON()),
|
|
)
|
|
|
|
|
|
def _safe_dict(value: Any) -> dict[str, Any]:
|
|
if isinstance(value, dict):
|
|
try:
|
|
return deepcopy(value)
|
|
except Exception:
|
|
return dict(value)
|
|
return {}
|
|
|
|
|
|
def _rewrite_model_routing_config(config: Any) -> dict[str, Any]:
|
|
payload = _safe_dict(config)
|
|
if not payload:
|
|
return {}
|
|
|
|
mappings = payload.get("mappings")
|
|
if not isinstance(mappings, dict):
|
|
payload["mappings"] = {}
|
|
return payload
|
|
|
|
normalized: dict[str, Any] = {}
|
|
for key, value in mappings.items():
|
|
name = str(key or "").strip().lower()
|
|
if not name:
|
|
continue
|
|
normalized[name] = deepcopy(value)
|
|
|
|
for old_key, new_key in _RENAME_KEYS.items():
|
|
if old_key not in normalized:
|
|
continue
|
|
if new_key not in normalized:
|
|
normalized[new_key] = normalized[old_key]
|
|
normalized.pop(old_key, None)
|
|
|
|
for key in _REMOVE_KEYS:
|
|
normalized.pop(key, None)
|
|
|
|
payload["mappings"] = normalized
|
|
return payload
|
|
|
|
|
|
def upgrade() -> None:
|
|
bind = op.get_bind()
|
|
if not _table_exists(bind, "integrations"):
|
|
return
|
|
|
|
table = _integration_table(bind)
|
|
rows = (
|
|
bind.execute(
|
|
sa.select(table.c.id, table.c.config).where(table.c.type == "model_routing")
|
|
)
|
|
.mappings()
|
|
.all()
|
|
)
|
|
for row in rows:
|
|
row_id = int(row.get("id") or 0)
|
|
rewritten = _rewrite_model_routing_config(row.get("config"))
|
|
if not rewritten:
|
|
continue
|
|
if rewritten == row.get("config"):
|
|
continue
|
|
bind.execute(
|
|
table.update().where(table.c.id == row_id).values(config=rewritten)
|
|
)
|
|
|
|
|
|
def downgrade() -> None:
|
|
bind = op.get_bind()
|
|
if not _table_exists(bind, "integrations"):
|
|
return
|
|
|
|
table = _integration_table(bind)
|
|
reverse_rename = {new_key: old_key for old_key, new_key in _RENAME_KEYS.items()}
|
|
rows = (
|
|
bind.execute(
|
|
sa.select(table.c.id, table.c.config).where(table.c.type == "model_routing")
|
|
)
|
|
.mappings()
|
|
.all()
|
|
)
|
|
for row in rows:
|
|
row_id = int(row.get("id") or 0)
|
|
payload = _safe_dict(row.get("config"))
|
|
mappings = payload.get("mappings")
|
|
if not isinstance(mappings, dict):
|
|
continue
|
|
normalized: dict[str, Any] = {}
|
|
for key, value in mappings.items():
|
|
name = str(key or "").strip().lower()
|
|
if not name:
|
|
continue
|
|
normalized[name] = deepcopy(value)
|
|
for new_key, old_key in reverse_rename.items():
|
|
if new_key not in normalized:
|
|
continue
|
|
if old_key not in normalized:
|
|
normalized[old_key] = normalized[new_key]
|
|
normalized.pop(new_key, None)
|
|
payload["mappings"] = normalized
|
|
if payload == row.get("config"):
|
|
continue
|
|
bind.execute(
|
|
table.update().where(table.c.id == row_id).values(config=payload)
|
|
)
|