group_fqcd_jr/tools/publish_advisor_demo_config.py

"""Publish the advisor Agent tool allowlist for an isolated demo environment.

The active ``config_release`` is environment data.  A new demo database has no
release by default, so AdvisorAgent tool calls would correctly fail closed even
though the Agent code is present.  This script creates a new release, carries
forward every existing item, and replaces only advisor tool entries.
"""

from __future__ import annotations

import argparse
import asyncio
import datetime as dt
import json
import sys
import uuid
from pathlib import Path
from typing import Any

import asyncmy
import httpx
import jwt

ROOT = Path(__file__).resolve().parents[1]
if str(ROOT) not in sys.path:
    sys.path.insert(0, str(ROOT))

from app.core.config import get_settings  # noqa: E402
from app.main import create_app  # noqa: E402

ADMIN_ID = "9003"
AGENT_TYPE = "advisor"
INTENT_TOOLS: dict[str, tuple[str, ...]] = {
    "fund_quote": ("query_fund_quote",),
    "investment_goal": ("query_investment_goal",),
    "portfolio_analysis": ("analyze_portfolio",),
    "asset_allocation": ("generate_asset_allocation",),
    "product_recommend": ("recommend_products",),
    "comparison": ("compare_products",),
}


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="Publish advisor demo tool allowlist")
    parser.add_argument("--apply", action="store_true", help="Create, review, and activate release")
    return parser.parse_args()


def issue_token(subject: str) -> str:
    settings = get_settings()
    private_key = Path(settings.jwt_private_key_path).read_text(encoding="utf-8")
    now = dt.datetime.now(dt.UTC)
    return jwt.encode(
        {
            "sub": subject,
            "iss": settings.jwt_issuer,
            "aud": settings.jwt_audience,
            "exp": now + dt.timedelta(minutes=30),
            "nbf": now - dt.timedelta(seconds=5),
            "jti": str(uuid.uuid4()),
        },
        private_key,
        algorithm="RS256",
    )


async def active_items() -> list[dict[str, Any]]:
    settings = get_settings()
    dsn = settings.mysql_dsn.split("://", 1)[1]
    credentials, location = dsn.split("@", 1)
    user, password = credentials.split(":", 1)
    host_port, database = location.split("/", 1)
    host, _, port = host_port.partition(":")
    connection = await asyncmy.connect(
        host=host, port=int(port or 3306), user=user, password=password, db=database
    )
    try:
        cursor = connection.cursor()
        await cursor.execute(
            """
            SELECT i.namespace, i.config_key, i.value_json, i.schema_version
            FROM platform_config_item i
            JOIN config_release r ON r.id = i.release_id
            WHERE r.status = 'active'
            """
        )
        rows = await cursor.fetchall()
    finally:
        connection.close()
    return [
        {
            "namespace": namespace,
            "item_key": key,
            "value_json": json.loads(value) if isinstance(value, str) else value,
            "schema_version": schema_version,
        }
        for namespace, key, value, schema_version in rows
    ]


async def request(
    client: httpx.AsyncClient,
    path: str,
    *,
    auth: dict[str, str],
    payload: dict[str, object] | None = None,
    if_match: str | None = None,
) -> httpx.Response:
    headers = {**auth, "Idempotency-Key": uuid.uuid4().hex}
    if if_match:
        headers["If-Match"] = if_match
    return await client.post(path, json=payload, headers=headers)


async def etag(client: httpx.AsyncClient, path: str, auth: dict[str, str]) -> str | None:
    return (await client.get(path, headers=auth)).headers.get("ETag")


def advisor_items() -> list[dict[str, object]]:
    return [
        {
            "namespace": "agent_tools",
            "item_key": f"{AGENT_TYPE}:{intent}",
            "value_json": {"allowed_tools": list(tools)},
            "schema_version": "1",
        }
        for intent, tools in INTENT_TOOLS.items()
    ]


async def publish() -> int:
    app = create_app()
    auth = {"Authorization": f"Bearer {issue_token(ADMIN_ID)}"}
    async with httpx.AsyncClient(transport=httpx.ASGITransport(app=app), base_url="http://test") as client:
        inherited = await active_items()
        replacement_keys = {("agent_tools", str(item["item_key"])) for item in advisor_items()}
        carried = [
            item
            for item in inherited
            if (str(item["namespace"]), str(item["item_key"])) not in replacement_keys
        ]
        created = await request(
            client,
            "/api/v1/admin/config-releases",
            auth=auth,
            payload={
                "release_no": f"advisor-demo-{uuid.uuid4().hex[:12]}",
                "title": "Advisor demo tool allowlist",
                "change_summary": "Publish advisor tool allowlists for the isolated demo environment",
            },
        )
        if created.status_code != 201:
            print(f"create release failed: {created.status_code} {created.text[:240]}")
            return 1
        release_id = int(created.json()["data"]["id"])
        base = f"/api/v1/admin/config-releases/{release_id}"
        for item in [*carried, *advisor_items()]:
            response = await request(
                client, f"{base}/platform-config-items", auth=auth, payload=item
            )
            if response.status_code != 201:
                print(f"write {item['namespace']}/{item['item_key']} failed: {response.status_code}")
                return 1
        for suffix, payload in (
            ("validations", {}),
            ("reviews", {"decision": "approved", "comment": "isolated demo setup"}),
            ("activations", {}),
        ):
            response = await request(
                client,
                f"{base}/{suffix}",
                auth=auth,
                payload=payload,
                if_match=await etag(client, base, auth),
            )
            if response.status_code not in (200, 201):
                print(f"{suffix} failed: {response.status_code} {response.text[:240]}")
                return 1
        print(f"published advisor demo release id={release_id}")
    return 0


def main() -> int:
    args = parse_args()
    for intent, tools in INTENT_TOOLS.items():
        print(f"advisor:{intent} -> {', '.join(tools)}")
    if not args.apply:
        print("dry run only; pass --apply to publish a new config release")
        return 0
    return asyncio.run(publish())


if __name__ == "__main__":
    sys.exit(main())