195 lines
6.4 KiB
Python
195 lines
6.4 KiB
Python
"""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())
|