121 lines
4.2 KiB
Python
121 lines
4.2 KiB
Python
"""Unit tests for ExpertTeamRouter and skill routing utilities."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from unittest.mock import MagicMock
|
|
|
|
from agentkit.chat.skill_routing import (
|
|
ExecutionMode,
|
|
SkillRoutingResult,
|
|
)
|
|
from agentkit.experts.config import ExpertConfig, ExpertTemplate
|
|
from agentkit.experts.registry import ExpertTemplateRegistry
|
|
from agentkit.experts.router import ExpertTeamRouter
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Helpers
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def _make_team_router_with_templates() -> ExpertTeamRouter:
|
|
"""Create an ExpertTeamRouter with sample templates."""
|
|
registry = ExpertTemplateRegistry()
|
|
for name in ("analyst", "strategist", "reviewer"):
|
|
config = ExpertConfig(
|
|
name=name,
|
|
agent_type="expert",
|
|
persona=f"Expert in {name}",
|
|
thinking_style="analytical",
|
|
bound_skills=[],
|
|
is_lead=(name == "analyst"),
|
|
task_mode="llm_generate",
|
|
prompt={"identity": f"Expert in {name}"},
|
|
)
|
|
template = ExpertTemplate(
|
|
name=name,
|
|
config=config,
|
|
description=f"Handles {name} tasks",
|
|
)
|
|
registry.register(template)
|
|
return ExpertTeamRouter(template_registry=registry)
|
|
|
|
|
|
def _make_team_router_empty() -> ExpertTeamRouter:
|
|
"""Create an ExpertTeamRouter with no templates."""
|
|
return ExpertTeamRouter(template_registry=ExpertTemplateRegistry())
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Tests: ExpertTeamRouter.can_handle()
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestExpertTeamRouterCanHandle:
|
|
def test_can_handle_with_templates(self) -> None:
|
|
router = _make_team_router_with_templates()
|
|
assert router.can_handle("analyze this data") is True
|
|
|
|
def test_can_handle_no_templates(self) -> None:
|
|
router = _make_team_router_empty()
|
|
assert router.can_handle("analyze this data") is False
|
|
|
|
def test_can_handle_name_match(self) -> None:
|
|
router = _make_team_router_with_templates()
|
|
assert router.can_handle("I need a strategist for this") is True
|
|
|
|
def test_can_handle_description_match(self) -> None:
|
|
router = _make_team_router_with_templates()
|
|
assert router.can_handle("handles review tasks") is True
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Tests: ExpertTeamRouter.resolve()
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestExpertTeamRouterResolve:
|
|
def test_explicit_team_prefix(self) -> None:
|
|
router = _make_team_router_with_templates()
|
|
result = router.resolve("@team:analyst,strategist analyze the market")
|
|
assert result.team_mode is True
|
|
assert result.match_method == "explicit_team"
|
|
assert "analyst" in result.specified_experts
|
|
assert "strategist" in result.specified_experts
|
|
|
|
def test_no_team_without_prefix(self) -> None:
|
|
router = _make_team_router_with_templates()
|
|
result = router.resolve("simple question")
|
|
assert result.team_mode is False
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Tests: SkillRoutingResult data structure
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestSkillRoutingResult:
|
|
def test_default_execution_mode(self) -> None:
|
|
result = SkillRoutingResult(
|
|
clean_content="test",
|
|
matched=False,
|
|
match_method="default_react",
|
|
match_confidence=0.8,
|
|
agent_name="default",
|
|
model="default",
|
|
execution_mode=ExecutionMode.REACT,
|
|
)
|
|
assert result.execution_mode == ExecutionMode.REACT
|
|
|
|
def test_direct_chat_mode(self) -> None:
|
|
result = SkillRoutingResult(
|
|
clean_content="hello",
|
|
matched=False,
|
|
match_method="regex_direct",
|
|
match_confidence=1.0,
|
|
agent_name="default",
|
|
model="default",
|
|
execution_mode=ExecutionMode.DIRECT_CHAT,
|
|
)
|
|
assert result.execution_mode == ExecutionMode.DIRECT_CHAT
|