Build lightweight AI agent admin
This commit is contained in:
Codex
@@ -0,0 +1,279 @@
|
||||
#!/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())
|
||||
Reference in New Issue
Block a user