zk-data-agent/agent_platform/bootstrap.py

from __future__ import annotations

import asyncio
import os
import sys
from typing import Any

import httpx

from agent_platform.models import MODEL_SPECS


def _required(name: str) -> str:
    value = os.getenv(name, "").strip()
    if not value:
        raise RuntimeError(f"{name} is required")
    return value


async def _wait_until_ready(client: httpx.AsyncClient, base_url: str) -> None:
    for _ in range(120):
        try:
            response = await client.get(f"{base_url}/health")
            if response.status_code == 200:
                return
        except httpx.HTTPError:
            pass
        await asyncio.sleep(2)
    raise RuntimeError("Open WebUI did not become healthy")


def _model_payload() -> list[dict[str, Any]]:
    public_read = [{"principal_type": "user", "principal_id": "*", "permission": "read"}]
    items = []
    for spec in MODEL_SPECS.values():
        is_chat = spec.mode == "chat"
        items.append(
            {
                "id": spec.public_id,
                "base_model_id": None,
                "name": spec.display_name,
                "params": {"function_calling": "native"} if is_chat else {},
                "meta": {
                    "description": (
                        "快速问答，使用原生工具循环。" if is_chat else "长链路工作，使用 K1412 Work Agent。"
                    ),
                    "toolIds": ["server:workspace"] if is_chat else [],
                    "capabilities": {
                        "tool_calling": is_chat,
                        "builtin_tools": False,
                        "file_context": False,
                        "citations": False,
                        "vision": False,
                        "file_upload": False,
                    },
                    "tags": [{"name": "Chat" if is_chat else "Work"}],
                },
                "access_grants": public_read,
                "is_active": True,
            }
        )
    return items


async def bootstrap() -> None:
    base_url = os.getenv("OPENWEBUI_INTERNAL_URL", "http://web:8080").rstrip("/")
    admin_email = _required("WEBUI_ADMIN_EMAIL")
    admin_password = _required("WEBUI_ADMIN_PASSWORD")
    gateway_key = _required("INTERNAL_GATEWAY_KEY")
    gateway_url = os.getenv("GATEWAY_URL", "http://gateway:8001").rstrip("/")

    async with httpx.AsyncClient(timeout=30) as client:
        await _wait_until_ready(client, base_url)
        response = await client.post(
            f"{base_url}/api/v1/auths/signin",
            json={"email": admin_email, "password": admin_password},
        )
        response.raise_for_status()
        token = response.json().get("token")
        if not token:
            raise RuntimeError("Open WebUI sign-in returned no token")
        headers = {"Authorization": f"Bearer {token}"}

        config_response = await client.get(f"{base_url}/api/v1/auths/admin/config", headers=headers)
        config_response.raise_for_status()
        config = config_response.json()
        config.update(
            {
                "ENABLE_SIGNUP": True,
                "DEFAULT_USER_ROLE": "pending",
                "ENABLE_API_KEYS": False,
                "ENABLE_COMMUNITY_SHARING": False,
                "ENABLE_MESSAGE_RATING": False,
                "ENABLE_FOLDERS": False,
                "ENABLE_AUTOMATIONS": False,
                "ENABLE_CHANNELS": False,
                "ENABLE_CALENDAR": False,
                "ENABLE_MEMORIES": False,
                "ENABLE_NOTES": False,
                "ENABLE_USER_WEBHOOKS": False,
            }
        )
        update_response = await client.post(
            f"{base_url}/api/v1/auths/admin/config",
            headers=headers,
            json=config,
        )
        update_response.raise_for_status()

        evaluation_response = await client.post(
            f"{base_url}/api/v1/evaluations/config",
            headers=headers,
            json={
                "ENABLE_EVALUATION_ARENA_MODELS": False,
                "EVALUATION_ARENA_MODELS": [],
            },
        )
        evaluation_response.raise_for_status()

        tool_server_response = await client.post(
            f"{base_url}/api/v1/configs/tool_servers",
            headers=headers,
            json={
                "TOOL_SERVER_CONNECTIONS": [
                    {
                        "url": gateway_url,
                        "path": "openapi.json",
                        "type": "openapi",
                        "auth_type": "bearer",
                        "headers": None,
                        "key": gateway_key,
                        "config": {
                            "enable": True,
                            "access_grants": [
                                {
                                    "principal_type": "user",
                                    "principal_id": "*",
                                    "permission": "read",
                                }
                            ],
                        },
                        "info": {
                            "id": "workspace",
                            "name": "Workspace",
                            "description": "当前用户的隔离编码工作区",
                        },
                        "spec_type": "url",
                    }
                ]
            },
        )
        tool_server_response.raise_for_status()

        model_response = await client.post(
            f"{base_url}/api/v1/models/import",
            headers=headers,
            json={"models": _model_payload()},
        )
        model_response.raise_for_status()

    print("Open WebUI bootstrap complete: auth policy, workspace tools, and six fixed Agent entries are ready.")


def run() -> None:
    try:
        asyncio.run(bootstrap())
    except Exception as exc:
        print(f"Open WebUI bootstrap failed: {exc}", file=sys.stderr)
        raise SystemExit(1) from exc


if __name__ == "__main__":
    run()