PMFlow/backend/alembic/versions/20260517_agent_runtime_vali...

127 lines
3.9 KiB
Python

"""raise agent runtime profile timeouts above validation floor
Revision ID: 20260517_agent_runtime_valid_profile_timeouts
Revises: 20260516_agent_answer_composer_budget
Create Date: 2026-05-17
"""
from __future__ import annotations
import json
import os
from typing import Any
from alembic import op
import sqlalchemy as sa
revision = "20260517_agent_runtime_valid_profile_timeouts"
down_revision = "20260516_agent_answer_composer_budget"
branch_labels = None
depends_on = None
PROFILE_UPDATES: dict[str, dict[str, Any]] = {
"pmflow_ai_agent_goal_frame": {
"provider_timeout_seconds": 20,
"job_timeout_seconds": 24,
"foreground_wait_seconds": 8,
"max_tokens": 1600,
"repair_max_tokens": 600,
"response_mode": "fast",
"runtime_mode": "active",
"quality_gate_mode": "off",
},
"pmflow_ai_agent_step_planner": {
"provider_timeout_seconds": 20,
"job_timeout_seconds": 24,
"foreground_wait_seconds": 8,
"max_tokens": 1600,
"repair_max_tokens": 600,
"response_mode": "fast",
"runtime_mode": "active",
"quality_gate_mode": "off",
},
"pmflow_ai_agent_state_critic": {
"provider_timeout_seconds": 16,
"job_timeout_seconds": 20,
"foreground_wait_seconds": 6,
"max_tokens": 1600,
"repair_max_tokens": 600,
"response_mode": "fast",
"runtime_mode": "active",
"quality_gate_mode": "off",
},
"pmflow_ai_agent_answer_composer": {
"provider_timeout_seconds": 35,
"job_timeout_seconds": 40,
"foreground_wait_seconds": 10,
"max_tokens": 3600,
"repair_max_tokens": 900,
"response_mode": "fast",
"runtime_mode": "active",
"quality_gate_mode": "off",
},
"pmflow_ai_agent_answer_reviser": {
"provider_timeout_seconds": 20,
"job_timeout_seconds": 24,
"foreground_wait_seconds": 8,
"max_tokens": 2600,
"repair_max_tokens": 900,
"response_mode": "fast",
"runtime_mode": "active",
"quality_gate_mode": "off",
},
}
def _table_exists(bind: Any, table_name: str) -> bool:
return sa.inspect(bind).has_table(table_name)
def upgrade() -> None:
bind = op.get_bind()
if not _table_exists(bind, "integrations"):
return
tenant_id = int(os.getenv("PMFLOW_PLATFORM_RUNTIME_TENANT_ID", "1") or "1")
row = bind.execute(
sa.text(
"""
SELECT id, config FROM integrations
WHERE tenant_id = :tenant_id
AND type = 'ai_function_runtime_profiles'
AND is_active = true
ORDER BY id DESC
LIMIT 1
"""
),
{"tenant_id": tenant_id},
).mappings().first()
if row is None:
return
config = row.get("config") if isinstance(row.get("config"), dict) else {}
changed = False
for stage in ("active", "draft"):
stage_packet = config.get(stage) if isinstance(config.get(stage), dict) else {}
profiles = stage_packet.get("profiles") if isinstance(stage_packet.get("profiles"), dict) else {}
for function_name, updates in PROFILE_UPDATES.items():
profile = profiles.get(function_name)
if not isinstance(profile, dict):
continue
for key, value in updates.items():
if profile.get(key) != value:
profile[key] = value
changed = True
if not changed:
return
config["active_version"] = "profile-v20260517-agent-runtime-valid-profile-timeouts"
config["version"] = config["active_version"]
bind.execute(
sa.text("UPDATE integrations SET config = CAST(:config AS jsonb) WHERE id = :id"),
{"id": int(row["id"]), "config": json.dumps(config, ensure_ascii=False)},
)
def downgrade() -> None:
return None