374 lines
13 KiB
Python
374 lines
13 KiB
Python
#!/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 _find_parallel_merge(
|
|
node_by_id: Dict[str, Dict[str, Any]],
|
|
outgoing_edges: Dict[str, list[str]],
|
|
branch_starts: list[str],
|
|
) -> str | None:
|
|
branch_merges: list[str] = []
|
|
for start in branch_starts:
|
|
current = start
|
|
seen: set[str] = set()
|
|
while current and current not in seen:
|
|
seen.add(current)
|
|
next_nodes = outgoing_edges.get(current, [])
|
|
if not next_nodes or len(next_nodes) > 1:
|
|
break
|
|
current = next_nodes[0]
|
|
if (node_by_id.get(current) or {}).get("type") == "merge":
|
|
branch_merges.append(current)
|
|
break
|
|
unique_merges = set(branch_merges)
|
|
return branch_merges[0] if len(unique_merges) == 1 and len(branch_merges) == len(branch_starts) else None
|
|
|
|
|
|
def _walk_workflow_nodes(
|
|
node_by_id: Dict[str, Dict[str, Any]],
|
|
outgoing_edges: Dict[str, list[str]],
|
|
start_id: str,
|
|
) -> set[str]:
|
|
visited: set[str] = set()
|
|
current = start_id
|
|
while current and current not in visited:
|
|
visited.add(current)
|
|
node_type = (node_by_id.get(current) or {}).get("type")
|
|
if node_type == "end":
|
|
break
|
|
if node_type == "parallel":
|
|
branch_starts = outgoing_edges.get(current, [])
|
|
for branch_start in branch_starts:
|
|
branch_current = branch_start
|
|
branch_seen: set[str] = set()
|
|
while branch_current and branch_current not in branch_seen:
|
|
if (node_by_id.get(branch_current) or {}).get("type") == "merge":
|
|
break
|
|
visited.add(branch_current)
|
|
branch_seen.add(branch_current)
|
|
branch_next_nodes = outgoing_edges.get(branch_current, [])
|
|
branch_current = branch_next_nodes[0] if branch_next_nodes else None
|
|
current = _find_parallel_merge(node_by_id, outgoing_edges, branch_starts)
|
|
continue
|
|
next_nodes = outgoing_edges.get(current, [])
|
|
current = next_nodes[0] if next_nodes else None
|
|
return visited
|
|
|
|
|
|
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("data") or {}).get("agent_code")
|
|
for node in nodes
|
|
if isinstance(node, dict) and (node.get("data") or {}).get("agent_code")
|
|
}
|
|
node_by_id = {node.get("id"): node for node in nodes if isinstance(node, dict)}
|
|
node_types = {node.get("type") for node in nodes}
|
|
outgoing_edges: Dict[str, list[str]] = {}
|
|
for edge in edges:
|
|
if not isinstance(edge, dict):
|
|
continue
|
|
source = edge.get("source")
|
|
target = edge.get("target")
|
|
if source and target:
|
|
outgoing_edges.setdefault(source, []).append(target)
|
|
|
|
required_persona_fields = {
|
|
"role",
|
|
"skills",
|
|
"constraints",
|
|
"background",
|
|
"examples",
|
|
}
|
|
required_node_types = {"start", "end", "template", "parallel", "merge"}
|
|
multi_outgoing_node_types = {"parallel", "condition", "intent", "choice"}
|
|
start_node = next((node for node in nodes if node.get("type") == "start"), None)
|
|
|
|
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)
|
|
invalid_multi_outgoing = {
|
|
source: targets
|
|
for source, targets in outgoing_edges.items()
|
|
if len(targets) > 1
|
|
and (node_by_id.get(source) or {}).get("type") not in multi_outgoing_node_types
|
|
}
|
|
reachable_nodes = _walk_workflow_nodes(node_by_id, outgoing_edges, start_node.get("id")) if start_node else set()
|
|
unreachable_nodes = sorted(set(node_by_id) - reachable_nodes)
|
|
invalid_parallel_merges = {
|
|
source: targets
|
|
for source, targets in outgoing_edges.items()
|
|
if (node_by_id.get(source) or {}).get("type") == "parallel"
|
|
and not _find_parallel_merge(node_by_id, outgoing_edges, targets)
|
|
}
|
|
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_workflow_agents
|
|
or missing_persona_fields
|
|
or missing_node_types
|
|
or invalid_multi_outgoing
|
|
or unreachable_nodes
|
|
or invalid_parallel_merges
|
|
or project_manager_workflow != "multica_org_collaboration_flow"
|
|
):
|
|
raise ValueError(
|
|
f"Fixture validation failed: missing_workflow_agents={missing_workflow_agents}, "
|
|
f"missing_persona_fields={missing_persona_fields}, "
|
|
f"missing_node_types={missing_node_types}, "
|
|
f"invalid_multi_outgoing={invalid_multi_outgoing}, "
|
|
f"unreachable_nodes={unreachable_nodes}, "
|
|
f"invalid_parallel_merges={invalid_parallel_merges}, "
|
|
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 resolve_default_chat_model_id(session) -> str | None:
|
|
from sqlalchemy import select
|
|
|
|
from ai_platform.models.model import LLMModel
|
|
|
|
result = await session.execute(
|
|
select(LLMModel)
|
|
.where(
|
|
LLMModel.is_deleted == False,
|
|
LLMModel.is_active == True,
|
|
LLMModel.model_type == "chat",
|
|
)
|
|
.order_by(LLMModel.sort.desc(), LLMModel.sys_create_datetime.desc())
|
|
)
|
|
model = result.scalars().first()
|
|
return str(model.id) if model else None
|
|
|
|
|
|
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
|
|
default_chat_model_id = await resolve_default_chat_model_id(session)
|
|
for agent_fixture in fixture["agents"]:
|
|
payload = build_agent_payload(agent_fixture, workflow_id)
|
|
if (
|
|
payload.get("mode", "autonomous") == "autonomous"
|
|
and not payload.get("model_id")
|
|
and default_chat_model_id
|
|
):
|
|
payload["model_id"] = default_chat_model_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())
|