113 lines
3.4 KiB
Python
113 lines
3.4 KiB
Python
"""align agent LLM foreground waits with provider budgets
|
|
|
|
Revision ID: 20260517_agent_runtime_llm_wait_alignment
|
|
Revises: 20260517_agent_dynamic_research_prompt_contract
|
|
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_llm_wait_alignment"
|
|
down_revision = "20260517_agent_dynamic_research_prompt_contract"
|
|
branch_labels = None
|
|
depends_on = None
|
|
|
|
|
|
PROFILE_UPDATES: dict[str, dict[str, Any]] = {
|
|
"pmflow_ai_agent_goal_frame": {
|
|
"provider_timeout_seconds": 45,
|
|
"foreground_wait_seconds": 45,
|
|
"job_timeout_seconds": 55,
|
|
"max_tokens": 1800,
|
|
},
|
|
"pmflow_ai_agent_step_planner": {
|
|
"provider_timeout_seconds": 45,
|
|
"foreground_wait_seconds": 45,
|
|
"job_timeout_seconds": 55,
|
|
"max_tokens": 1800,
|
|
},
|
|
"pmflow_ai_agent_evidence_normalizer": {
|
|
"provider_timeout_seconds": 75,
|
|
"foreground_wait_seconds": 75,
|
|
"job_timeout_seconds": 90,
|
|
"max_tokens": 3200,
|
|
},
|
|
"pmflow_ai_agent_state_critic": {
|
|
"provider_timeout_seconds": 30,
|
|
"foreground_wait_seconds": 30,
|
|
"job_timeout_seconds": 40,
|
|
"max_tokens": 1800,
|
|
},
|
|
"pmflow_ai_agent_answer_composer": {
|
|
"provider_timeout_seconds": 60,
|
|
"foreground_wait_seconds": 60,
|
|
"job_timeout_seconds": 75,
|
|
"max_tokens": 3600,
|
|
},
|
|
"pmflow_ai_agent_answer_reviser": {
|
|
"provider_timeout_seconds": 45,
|
|
"foreground_wait_seconds": 45,
|
|
"job_timeout_seconds": 55,
|
|
"max_tokens": 3000,
|
|
},
|
|
}
|
|
|
|
|
|
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 profile_name, updates in PROFILE_UPDATES.items():
|
|
profile = profiles.get(profile_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-llm-wait-alignment"
|
|
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
|