ai-agent-admin/backend-fastapi/scripts/seed_multica_org_agents.py

#!/usr/bin/env python
"""
Seed or rollback the Multica organization collaboration workflow and agents.

Usage:
  python scripts/seed_multica_org_agents.py --dry-run
  python scripts/seed_multica_org_agents.py --apply
  python scripts/seed_multica_org_agents.py --rollback
"""
from __future__ import annotations

import argparse
import asyncio
import json
import sys
from copy import deepcopy
from datetime import datetime
from pathlib import Path
from typing import Any, Dict

sys.path.insert(0, str(Path(__file__).resolve().parent.parent))


PROJECT_ROOT = Path(__file__).resolve().parent.parent
FIXTURE_PATH = PROJECT_ROOT / "ai_platform" / "fixtures" / "multica_org_agents.json"


def load_fixture() -> Dict[str, Any]:
    with FIXTURE_PATH.open("r", encoding="utf-8") as fixture_file:
        return json.load(fixture_file)


def build_workflow_payload(fixture: Dict[str, Any]) -> Dict[str, Any]:
    workflow = deepcopy(fixture["workflow"])
    definition = workflow["definition"]
    workflow["published_definition"] = deepcopy(definition)
    workflow["published_at"] = datetime.utcnow()
    return workflow


def build_agent_payload(agent_fixture: Dict[str, Any], workflow_id: str | None) -> Dict[str, Any]:
    payload = deepcopy(agent_fixture)
    workflow_code = payload.pop("workflow_code", None)
    if workflow_code:
        payload["workflow_id"] = workflow_id
    return payload


def validate_fixture(fixture: Dict[str, Any]) -> Dict[str, int]:
    agents = fixture.get("agents", [])
    definition = fixture.get("workflow", {}).get("definition", {})
    nodes = definition.get("nodes", [])
    edges = definition.get("edges", [])
    agent_codes = {agent.get("code") for agent in agents}
    workflow_agent_codes = {
        node.get("agent_code")
        for node in nodes
        if isinstance(node, dict) and node.get("agent_code")
    }
    node_types = {node.get("type") for node in nodes}

    required_agent_codes = {
        "multica_product_manager",
        "business_requirements_analyst",
        "system_architect",
        "frontend_engineer",
        "backend_engineer",
        "qa_engineer",
        "project_manager",
    }
    required_persona_fields = {
        "role",
        "skills",
        "constraints",
        "background",
        "examples",
    }
    required_node_types = {"start", "end", "condition", "template", "parallel", "merge"}

    missing_agents = sorted(required_agent_codes - agent_codes)
    extra_agents = sorted(agent_codes - required_agent_codes)
    missing_workflow_agents = sorted(workflow_agent_codes - agent_codes)
    missing_persona_fields = {
        agent.get("code"): sorted(required_persona_fields - set((agent.get("persona") or {}).keys()))
        for agent in agents
        if required_persona_fields - set((agent.get("persona") or {}).keys())
    }
    missing_node_types = sorted(required_node_types - node_types)
    project_manager = next(
        (agent for agent in agents if agent.get("code") == "project_manager"),
        None,
    )
    project_manager_workflow = (project_manager or {}).get("workflow_code")
    if (
        missing_agents
        or extra_agents
        or missing_workflow_agents
        or missing_persona_fields
        or missing_node_types
        or project_manager_workflow != "multica_org_collaboration_flow"
    ):
        raise ValueError(
            f"Fixture validation failed: missing_agents={missing_agents}, "
            f"extra_agents={extra_agents}, "
            f"missing_workflow_agents={missing_workflow_agents}, "
            f"missing_persona_fields={missing_persona_fields}, "
            f"missing_node_types={missing_node_types}, "
            f"project_manager_workflow={project_manager_workflow}"
        )

    return {
        "agents": len(agents),
        "workflow_nodes": len(nodes),
        "workflow_edges": len(edges),
        "workflow_agent_refs": len(workflow_agent_codes),
    }


async def upsert_workflow(session, workflow_payload: Dict[str, Any]) -> AIWorkflow:
    from sqlalchemy import select

    from ai_platform.models.workflow import AIWorkflow, AIWorkflowVersion

    code = workflow_payload["code"]
    result = await session.execute(select(AIWorkflow).where(AIWorkflow.code == code))
    workflow = result.scalar_one_or_none()

    if workflow:
        for key, value in workflow_payload.items():
            setattr(workflow, key, value)
        workflow.is_deleted = False
    else:
        workflow = AIWorkflow(**workflow_payload)
        session.add(workflow)

    await session.flush()

    result = await session.execute(
        select(AIWorkflowVersion).where(
            AIWorkflowVersion.workflow_id == workflow.id,
            AIWorkflowVersion.version == workflow.version,
        )
    )
    version = result.scalar_one_or_none()
    version_payload = {
        "workflow_id": workflow.id,
        "version": workflow.version,
        "definition": deepcopy(workflow.definition),
        "description": "Seeded Multica organization collaboration workflow",
        "published_at": workflow.published_at,
    }
    if version:
        for key, value in version_payload.items():
            setattr(version, key, value)
        version.is_deleted = False
    else:
        session.add(AIWorkflowVersion(**version_payload))

    return workflow


async def upsert_agents(session, fixture: Dict[str, Any], workflow_id: str) -> int:
    from sqlalchemy import select

    from ai_platform.models.agent import Agent
    from app.base_model import generate_nanoid

    count = 0
    for agent_fixture in fixture["agents"]:
        payload = build_agent_payload(agent_fixture, workflow_id)
        result = await session.execute(select(Agent).where(Agent.code == payload["code"]))
        agent = result.scalar_one_or_none()
        if agent:
            for key, value in payload.items():
                setattr(agent, key, value)
            agent.is_deleted = False
        else:
            session.add(Agent(id=generate_nanoid(), **payload))
        count += 1
    return count


async def apply_seed(dry_run: bool) -> Dict[str, Any]:
    fixture = load_fixture()
    counts = validate_fixture(fixture)
    workflow_payload = build_workflow_payload(fixture)

    if dry_run:
        return {
            "action": "dry-run",
            "workflow_code": workflow_payload["code"],
            "agent_count": len(fixture["agents"]),
            **counts,
        }

    from app.database import AsyncSessionLocal

    async with AsyncSessionLocal() as session:
        workflow = await upsert_workflow(session, workflow_payload)
        agent_count = await upsert_agents(session, fixture, workflow.id)

        await session.commit()
        action = "applied"

    return {
        "action": action,
        "workflow_code": workflow_payload["code"],
        "agent_count": agent_count,
        **counts,
    }


async def rollback_seed(dry_run: bool) -> Dict[str, Any]:
    from sqlalchemy import select

    from ai_platform.models.agent import Agent
    from ai_platform.models.workflow import AIWorkflow, AIWorkflowVersion
    from app.database import AsyncSessionLocal

    fixture = load_fixture()
    workflow_code = fixture["workflow"]["code"]
    agent_codes = [agent["code"] for agent in fixture["agents"]]

    async with AsyncSessionLocal() as session:
        result = await session.execute(select(AIWorkflow).where(AIWorkflow.code == workflow_code))
        workflow = result.scalar_one_or_none()
        workflow_count = 0
        version_count = 0
        if workflow:
            workflow.is_deleted = True
            workflow_count = 1
            versions = await session.execute(
                select(AIWorkflowVersion).where(AIWorkflowVersion.workflow_id == workflow.id)
            )
            for version in versions.scalars().all():
                version.is_deleted = True
                version_count += 1

        agents = await session.execute(select(Agent).where(Agent.code.in_(agent_codes)))
        agent_count = 0
        for agent in agents.scalars().all():
            agent.is_deleted = True
            agent.workflow_id = None
            agent_count += 1

        if dry_run:
            await session.rollback()
            action = "rollback-dry-run"
        else:
            await session.commit()
            action = "rolled-back"

    return {
        "action": action,
        "workflow_code": workflow_code,
        "workflow_count": workflow_count,
        "workflow_version_count": version_count,
        "agent_count": agent_count,
    }


async def main() -> None:
    parser = argparse.ArgumentParser(description="Seed Multica organization collaboration agents")
    mode = parser.add_mutually_exclusive_group(required=True)
    mode.add_argument("--dry-run", action="store_true", help="Validate and simulate the seed without committing")
    mode.add_argument("--apply", action="store_true", help="Apply the seed to the configured database")
    mode.add_argument("--rollback", action="store_true", help="Soft-delete seeded workflow, versions, and agents")
    args = parser.parse_args()

    if args.rollback:
        result = await rollback_seed(dry_run=False)
    else:
        result = await apply_seed(dry_run=args.dry_run)

    print(json.dumps(result, ensure_ascii=False, sort_keys=True))


if __name__ == "__main__":
    asyncio.run(main())