PMFlow/backend/alembic/versions/20260516_agent_objective_contract_fast_path.py

"""refresh agent objective contract and fast path prompts

Revision ID: 20260516_agent_objective_contract_fast_path
Revises: 20260516_agent_compound_goal_assumption
Create Date: 2026-05-16
"""

from __future__ import annotations

import json
import os
from typing import Any

from alembic import op
import sqlalchemy as sa


revision = "20260516_agent_objective_contract_fast_path"
down_revision = "20260516_agent_compound_goal_assumption"
branch_labels = None
depends_on = None


PROMPT_KEYS = (
    "PROMPT_PMFLOW_AI_AGENT_GOAL_FRAME",
    "PROMPT_PMFLOW_AI_AGENT_STEP_PLANNER",
    "PROMPT_PMFLOW_AI_AGENT_EVIDENCE_NORMALIZER",
    "PROMPT_PMFLOW_AI_AGENT_STATE_CRITIC",
    "PROMPT_PMFLOW_AI_AGENT_ANSWER_COMPOSER",
    "PROMPT_PMFLOW_AI_AGENT_ANSWER_REVISER",
)
REFRESH_NOTES = "Refreshed by 20260516_agent_objective_contract_fast_path"


def _table_exists(bind: Any, table_name: str) -> bool:
    return sa.inspect(bind).has_table(table_name)


def _prompt_body(prompt_key: str) -> str:
    from app.config import ai_prompts

    return str(getattr(ai_prompts, prompt_key, "") or "").strip()


def _refresh_active_prompt(bind: Any, *, tenant_id: int, prompt_key: str, body_snapshot: str) -> None:
    active = bind.execute(
        sa.text(
            """
            SELECT id, body_snapshot FROM prompt_releases
            WHERE tenant_id = :tenant_id
              AND prompt_key = :prompt_key
              AND status = 'active'
            LIMIT 1
            """
        ),
        {"tenant_id": tenant_id, "prompt_key": prompt_key},
    ).mappings().first()
    if active is not None and str(active.get("body_snapshot") or "").strip() == body_snapshot:
        return
    bind.execute(
        sa.text(
            """
            UPDATE prompt_releases
            SET status = 'archived'
            WHERE tenant_id = :tenant_id
              AND prompt_key = :prompt_key
              AND status = 'active'
            """
        ),
        {"tenant_id": tenant_id, "prompt_key": prompt_key},
    )
    max_version = bind.execute(
        sa.text(
            """
            SELECT MAX(version) FROM prompt_releases
            WHERE tenant_id = :tenant_id
              AND prompt_key = :prompt_key
            """
        ),
        {"tenant_id": tenant_id, "prompt_key": prompt_key},
    ).scalar()
    bind.execute(
        sa.text(
            """
            INSERT INTO prompt_releases
                (tenant_id, prompt_key, version, body_snapshot, status, notes, created_by, activated_by, activated_at)
            VALUES
                (:tenant_id, :prompt_key, :version, :body_snapshot, 'active', :notes, NULL, NULL, :activated_at)
            """
        ),
        {
            "tenant_id": tenant_id,
            "prompt_key": prompt_key,
            "version": int(max_version or 0) + 1,
            "body_snapshot": body_snapshot,
            "notes": REFRESH_NOTES,
            "activated_at": "2026-05-16T00:00:00",
        },
    )


def _refresh_runtime_profiles(bind: Any, *, tenant_id: int) -> None:
    if not _table_exists(bind, "integrations"):
        return
    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"):
        profiles = ((config.get(stage) or {}).get("profiles") or {}) if isinstance(config.get(stage), dict) else {}
        for key in (
            "pmflow_ai_agent_goal_frame",
            "pmflow_ai_agent_step_planner",
            "pmflow_ai_agent_state_critic",
            "pmflow_ai_agent_answer_composer",
        ):
            profile = profiles.get(key)
            if not isinstance(profile, dict):
                continue
            if profile.get("response_mode") != "fast":
                profile["response_mode"] = "fast"
                changed = True
            if key in {"pmflow_ai_agent_goal_frame", "pmflow_ai_agent_step_planner", "pmflow_ai_agent_state_critic"}:
                if int(profile.get("provider_timeout_seconds") or 0) > 25:
                    profile["provider_timeout_seconds"] = 25
                    changed = True
                if int(profile.get("job_timeout_seconds") or 0) > 30:
                    profile["job_timeout_seconds"] = 30
                    changed = True
        normalizer = profiles.get("pmflow_ai_agent_evidence_normalizer")
        if isinstance(normalizer, dict):
            if normalizer.get("response_mode") != "fast":
                normalizer["response_mode"] = "fast"
                changed = True
            if int(normalizer.get("provider_timeout_seconds") or 0) > 22:
                normalizer["provider_timeout_seconds"] = 22
                changed = True
            if int(normalizer.get("job_timeout_seconds") or 0) > 25:
                normalizer["job_timeout_seconds"] = 25
                changed = True
    if not changed:
        return
    config["active_version"] = "profile-v20260516-objective-fast-path"
    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 upgrade() -> None:
    bind = op.get_bind()
    tenant_id = int(os.getenv("PMFLOW_PLATFORM_RUNTIME_TENANT_ID", "1") or "1")
    if _table_exists(bind, "prompt_releases"):
        for prompt_key in PROMPT_KEYS:
            body_snapshot = _prompt_body(prompt_key)
            if body_snapshot:
                _refresh_active_prompt(bind, tenant_id=tenant_id, prompt_key=prompt_key, body_snapshot=body_snapshot)
    _refresh_runtime_profiles(bind, tenant_id=tenant_id)


def downgrade() -> None:
    bind = op.get_bind()
    if not _table_exists(bind, "prompt_releases"):
        return
    tenant_id = int(os.getenv("PMFLOW_PLATFORM_RUNTIME_TENANT_ID", "1") or "1")
    bind.execute(
        sa.text(
            """
            DELETE FROM prompt_releases
            WHERE tenant_id = :tenant_id
              AND notes = :notes
            """
        ),
        {"tenant_id": tenant_id, "notes": REFRESH_NOTES},
    )