LZH-YS1998
102 lines
3.6 KiB
Python
102 lines
3.6 KiB
Python
from __future__ import annotations
|
|
|
|
import unittest
|
|
from dataclasses import dataclass, field
|
|
from types import SimpleNamespace
|
|
|
|
import aiosqlite
|
|
|
|
from opc.core.config import RoleConfig
|
|
from opc.plugins.office_ui.agent_store import AgentStore
|
|
|
|
|
|
@dataclass
|
|
class _Role:
|
|
name: str
|
|
responsibility: str = ""
|
|
tools: list[str] = field(default_factory=list)
|
|
|
|
|
|
class _OrgEngine:
|
|
def get_agent(self, role_id: str) -> _Role | None:
|
|
roles = {
|
|
"coordinator": _Role(
|
|
name="Custom Leader",
|
|
responsibility="Lead the custom team.",
|
|
tools=["send_dm", "todo_write"],
|
|
),
|
|
"executor": _Role(
|
|
name="Executor",
|
|
responsibility="Execute assigned work.",
|
|
tools=["shell_exec"],
|
|
),
|
|
}
|
|
return roles.get(role_id)
|
|
|
|
|
|
class _PresetOrgEngine(_OrgEngine):
|
|
def __init__(self) -> None:
|
|
self.config = SimpleNamespace(
|
|
org=SimpleNamespace(
|
|
roles=[
|
|
RoleConfig(
|
|
id="ceo",
|
|
name="Configured CEO",
|
|
responsibility="Configured leader",
|
|
tools=["file_read", "todo_write"],
|
|
)
|
|
]
|
|
)
|
|
)
|
|
|
|
|
|
class AgentStoreCustomModeTests(unittest.IsolatedAsyncioTestCase):
|
|
async def asyncSetUp(self) -> None:
|
|
self.db = await aiosqlite.connect(":memory:")
|
|
self.store = AgentStore(self.db)
|
|
await self.store.initialize()
|
|
self.org_engine = _OrgEngine()
|
|
|
|
async def asyncTearDown(self) -> None:
|
|
await self.db.close()
|
|
|
|
async def test_custom_mode_without_shadow_resets_to_single_starter(self) -> None:
|
|
preset_agents = await self.store.load_preset("corporate", self.org_engine)
|
|
self.assertGreater(len(preset_agents), 1)
|
|
self.assertTrue(any(agent["opc_role_id"] == "ceo" for agent in preset_agents))
|
|
|
|
custom_agents = await self.store.load_preset("custom", self.org_engine)
|
|
|
|
self.assertEqual(len(custom_agents), 1)
|
|
self.assertEqual(custom_agents[0]["agent_id"], "custom-leader")
|
|
self.assertEqual(custom_agents[0]["opc_role_id"], "coordinator")
|
|
self.assertFalse(any(agent["opc_role_id"] == "ceo" for agent in custom_agents))
|
|
|
|
async def test_custom_mode_restores_saved_custom_team(self) -> None:
|
|
starter_agents = await self.store.load_preset("custom", self.org_engine)
|
|
self.assertEqual([agent["agent_id"] for agent in starter_agents], ["custom-leader"])
|
|
|
|
await self.store.create_agent(
|
|
name="Planner",
|
|
opc_role_id="planner",
|
|
office_id="office-1",
|
|
description="Plans the work.",
|
|
specialties=["planning"],
|
|
)
|
|
await self.store.sync_custom_shadow()
|
|
|
|
await self.store.load_preset("corporate", self.org_engine)
|
|
restored_agents = await self.store.load_preset("custom", self.org_engine)
|
|
restored_role_ids = {agent["opc_role_id"] for agent in restored_agents}
|
|
|
|
self.assertEqual(restored_role_ids, {"coordinator", "planner"})
|
|
self.assertFalse(any(agent["opc_role_id"] == "ceo" for agent in restored_agents))
|
|
|
|
async def test_builtin_preset_uses_configured_role_tool_overrides(self) -> None:
|
|
preset_agents = await self.store.load_preset("corporate", _PresetOrgEngine())
|
|
ceo = next(agent for agent in preset_agents if agent["opc_role_id"] == "ceo")
|
|
|
|
self.assertEqual(ceo["name"], "Configured CEO")
|
|
self.assertEqual(set(ceo["specialties"]), {"file_read", "todo_write"})
|
|
|