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

"""refresh hypothesis answer policy and latency budgets

Revision ID: 20260516_agent_hypothesis_answer_policy
Revises: 20260516_agent_controlled_hypothesis_prompt
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_hypothesis_answer_policy"
down_revision = "20260516_agent_controlled_hypothesis_prompt"
branch_labels = None
depends_on = None


REFRESH_NOTES = "Refreshed by 20260516_agent_hypothesis_answer_policy"
PROMPT_KEYS = (
    "PROMPT_PMFLOW_AI_AGENT_ANSWER_COMPOSER",
)


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:
    if not _table_exists(bind, "prompt_releases"):
        return
    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
    caps = {
        "pmflow_ai_agent_goal_frame": (8, 10),
        "pmflow_ai_agent_step_planner": (4, 6),
        "pmflow_ai_agent_state_critic": (4, 6),
        "pmflow_ai_agent_answer_composer": (14, 16),
        "pmflow_ai_agent_answer_reviser": (10, 12),
    }
    for stage in ("active", "draft"):
        profiles = ((config.get(stage) or {}).get("profiles") or {}) if isinstance(config.get(stage), dict) else {}
        for key, (provider_timeout, job_timeout) in caps.items():
            profile = profiles.get(key)
            if not isinstance(profile, dict):
                continue
            if int(profile.get("provider_timeout_seconds") or 0) != provider_timeout:
                profile["provider_timeout_seconds"] = provider_timeout
                changed = True
            if int(profile.get("job_timeout_seconds") or 0) != job_timeout:
                profile["job_timeout_seconds"] = job_timeout
                changed = True
            if profile.get("response_mode") != "fast":
                profile["response_mode"] = "fast"
                changed = True
    if not changed:
        return
    config["active_version"] = "profile-v20260516-agent-hypothesis-answer-policy"
    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:
    tenant_id = int(os.getenv("PMFLOW_PLATFORM_RUNTIME_TENANT_ID", "1") or "1")
    bind = op.get_bind()
    for prompt_key in PROMPT_KEYS:
        body = _prompt_body(prompt_key)
        if body:
            _refresh_active_prompt(bind, tenant_id=tenant_id, prompt_key=prompt_key, body_snapshot=body)
    _refresh_runtime_profiles(bind, tenant_id=tenant_id)


def downgrade() -> None:
    return None