316 lines
11 KiB
Python
316 lines
11 KiB
Python
"""集成测试 - CostAwareRouter → Engine → AlignmentGuard 全链路"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import json
|
||
from unittest.mock import AsyncMock, MagicMock
|
||
|
||
import pytest
|
||
|
||
from agentkit.chat.skill_routing import CostAwareRouter, SkillRoutingResult
|
||
from agentkit.core.react import ReActEngine, ReActResult, ReActStep
|
||
from agentkit.llm.protocol import LLMResponse, TokenUsage
|
||
from agentkit.org.context import AgentProfile, OrganizationContext
|
||
from agentkit.quality.alignment import AlignmentConfig, AlignmentGuard
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Helpers
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
def _make_llm_response(content: str) -> LLMResponse:
|
||
return LLMResponse(
|
||
content=content,
|
||
model="test-model",
|
||
usage=TokenUsage(prompt_tokens=10, completion_tokens=20),
|
||
)
|
||
|
||
|
||
def _make_mock_gateway(responses: list[LLMResponse]) -> MagicMock:
|
||
gateway = MagicMock()
|
||
gateway.chat = AsyncMock(side_effect=responses)
|
||
return gateway
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Test 1: Router routes to ReAct engine, output passes alignment check
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
class TestRouterToEnginePassesAlignment:
|
||
"""路由到 ReAct 引擎，输出通过对齐检查"""
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_react_output_passes_alignment(self):
|
||
# --- Setup: LLM returns low complexity → default agent (ReAct) ---
|
||
gateway = _make_mock_gateway([
|
||
_make_llm_response('{"complexity": 0.2}'), # quick_classify
|
||
_make_llm_response("你好！有什么可以帮你的？"), # ReAct final answer
|
||
])
|
||
|
||
org_context = OrganizationContext()
|
||
alignment_config = AlignmentConfig(
|
||
constraints=["password", "secret_key"],
|
||
)
|
||
guard = AlignmentGuard(config=alignment_config)
|
||
|
||
router = CostAwareRouter(llm_gateway=gateway, org_context=org_context)
|
||
|
||
mock_skill_registry = MagicMock()
|
||
mock_skill_registry.list_skills.return_value = []
|
||
mock_intent_router = AsyncMock()
|
||
|
||
# Step 1: Route
|
||
route_result = await router.route(
|
||
content="随便聊聊",
|
||
skill_registry=mock_skill_registry,
|
||
intent_router=mock_intent_router,
|
||
default_tools=[],
|
||
default_system_prompt="You are helpful",
|
||
default_model="default",
|
||
default_agent_name="default",
|
||
)
|
||
assert route_result.complexity < 0.3
|
||
assert route_result.agent_name == "default"
|
||
|
||
# Step 2: Inject constraints
|
||
input_data = {"content": route_result.clean_content}
|
||
injected = guard.inject_constraints(input_data)
|
||
assert "alignment_constraints" in injected
|
||
|
||
# Step 3: Simulate engine execution (use real ReActEngine with mock gateway)
|
||
react_engine = ReActEngine(llm_gateway=gateway)
|
||
engine_result = await react_engine.execute(
|
||
messages=[{"role": "user", "content": injected["content"]}],
|
||
)
|
||
|
||
# Step 4: Alignment check
|
||
output = {"result": engine_result.output}
|
||
check_result = await guard.check_output(output)
|
||
assert check_result.passed is True
|
||
assert check_result.violations == []
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Test 2: Router routes to ReAct engine, output fails alignment check
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
class TestRouterToEngineFailsAlignment:
|
||
"""路由到 ReAct 引擎，输出未通过对齐检查"""
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_react_output_fails_alignment(self):
|
||
gateway = _make_mock_gateway([
|
||
_make_llm_response('{"complexity": 0.2}'),
|
||
_make_llm_response("Your password is 123456"),
|
||
])
|
||
|
||
org_context = OrganizationContext()
|
||
alignment_config = AlignmentConfig(
|
||
constraints=["password", "secret_key"],
|
||
)
|
||
guard = AlignmentGuard(config=alignment_config)
|
||
|
||
router = CostAwareRouter(llm_gateway=gateway, org_context=org_context)
|
||
|
||
mock_skill_registry = MagicMock()
|
||
mock_skill_registry.list_skills.return_value = []
|
||
mock_intent_router = AsyncMock()
|
||
|
||
route_result = await router.route(
|
||
content="随便聊聊",
|
||
skill_registry=mock_skill_registry,
|
||
intent_router=mock_intent_router,
|
||
default_tools=[],
|
||
default_system_prompt="You are helpful",
|
||
default_model="default",
|
||
default_agent_name="default",
|
||
)
|
||
|
||
input_data = {"content": route_result.clean_content}
|
||
injected = guard.inject_constraints(input_data)
|
||
|
||
react_engine = ReActEngine(llm_gateway=gateway)
|
||
engine_result = await react_engine.execute(
|
||
messages=[{"role": "user", "content": injected["content"]}],
|
||
)
|
||
|
||
output = {"result": engine_result.output}
|
||
check_result = await guard.check_output(output)
|
||
assert check_result.passed is False
|
||
assert len(check_result.violations) > 0
|
||
assert "password" in check_result.violations
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Test 3: Router routes based on complexity (low→default, high→org_context)
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
class TestRouterComplexityBasedRouting:
|
||
"""基于复杂度的路由：低复杂度→默认，高复杂度→org_context 能力匹配"""
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_low_complexity_routes_to_default(self):
|
||
gateway = _make_mock_gateway([
|
||
_make_llm_response('{"complexity": 0.15}'),
|
||
])
|
||
|
||
org_context = OrganizationContext()
|
||
org_context.register_agent(AgentProfile(
|
||
name="analyst",
|
||
agent_type="react",
|
||
capabilities=["analysis"],
|
||
skills=["analysis"],
|
||
))
|
||
|
||
router = CostAwareRouter(llm_gateway=gateway, org_context=org_context)
|
||
|
||
mock_skill_registry = MagicMock()
|
||
mock_skill_registry.list_skills.return_value = []
|
||
mock_intent_router = AsyncMock()
|
||
|
||
result = await router.route(
|
||
content="简单问题",
|
||
skill_registry=mock_skill_registry,
|
||
intent_router=mock_intent_router,
|
||
default_tools=[],
|
||
default_system_prompt="You are helpful",
|
||
default_model="default",
|
||
default_agent_name="default",
|
||
)
|
||
|
||
assert result.complexity < 0.3
|
||
assert result.agent_name == "default"
|
||
assert result.match_method == "low_complexity"
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_high_complexity_routes_via_org_context(self):
|
||
gateway = _make_mock_gateway([
|
||
_make_llm_response('{"complexity": 0.85}'),
|
||
])
|
||
|
||
org_context = OrganizationContext()
|
||
org_context.register_agent(AgentProfile(
|
||
name="analyst",
|
||
agent_type="react",
|
||
capabilities=["分析", "市场", "调研"],
|
||
skills=["market_analysis"],
|
||
current_load=0,
|
||
))
|
||
|
||
# find_best_agent returns real AgentProfile
|
||
org_context.find_best_agent = MagicMock(
|
||
return_value=org_context.get_agent_profile("analyst")
|
||
)
|
||
|
||
router = CostAwareRouter(llm_gateway=gateway, org_context=org_context)
|
||
|
||
mock_skill_registry = MagicMock()
|
||
mock_skill_registry.list_skills.return_value = []
|
||
mock_intent_router = AsyncMock()
|
||
|
||
result = await router.route(
|
||
content="请对市场趋势进行深度分析并给出投资建议",
|
||
skill_registry=mock_skill_registry,
|
||
intent_router=mock_intent_router,
|
||
default_tools=[],
|
||
default_system_prompt="You are helpful",
|
||
default_model="default",
|
||
default_agent_name="default",
|
||
)
|
||
|
||
assert result.complexity >= 0.7
|
||
assert result.match_method == "capability"
|
||
assert result.agent_name == "analyst"
|
||
assert result.matched is True
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Test 4: AlignmentGuard injects constraints into input before engine execution
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
class TestAlignmentGuardConstraintInjection:
|
||
"""AlignmentGuard 在引擎执行前将约束注入到输入中"""
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_constraints_injected_before_engine_execution(self):
|
||
gateway = _make_mock_gateway([
|
||
_make_llm_response('{"complexity": 0.5}'),
|
||
_make_llm_response("Safe answer"),
|
||
])
|
||
|
||
alignment_config = AlignmentConfig(
|
||
constraints=["不得泄露用户隐私", "禁止生成有害内容"],
|
||
)
|
||
guard = AlignmentGuard(config=alignment_config)
|
||
|
||
org_context = OrganizationContext()
|
||
router = CostAwareRouter(llm_gateway=gateway, org_context=org_context)
|
||
|
||
mock_skill_registry = MagicMock()
|
||
mock_skill_registry.list_skills.return_value = []
|
||
mock_intent_router = AsyncMock()
|
||
|
||
# Step 1: Route
|
||
route_result = await router.route(
|
||
content="请帮我写一篇文章",
|
||
skill_registry=mock_skill_registry,
|
||
intent_router=mock_intent_router,
|
||
default_tools=[],
|
||
default_system_prompt="You are helpful",
|
||
default_model="default",
|
||
default_agent_name="default",
|
||
)
|
||
|
||
# Step 2: Inject constraints
|
||
input_data = {"content": route_result.clean_content}
|
||
injected = guard.inject_constraints(input_data)
|
||
|
||
# Verify constraints are present
|
||
assert "alignment_constraints" in injected
|
||
assert "不得泄露用户隐私" in injected["alignment_constraints"]
|
||
assert "禁止生成有害内容" in injected["alignment_constraints"]
|
||
# Original data preserved
|
||
assert injected["content"] == route_result.clean_content
|
||
# Original dict not mutated
|
||
assert "alignment_constraints" not in input_data
|
||
|
||
# Step 3: Engine executes with injected input
|
||
react_engine = ReActEngine(llm_gateway=gateway)
|
||
engine_result = await react_engine.execute(
|
||
messages=[{"role": "user", "content": injected["content"]}],
|
||
)
|
||
|
||
# Step 4: Output passes alignment
|
||
output = {"result": engine_result.output}
|
||
check_result = await guard.check_output(output)
|
||
assert check_result.passed is True
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_constraint_injection_with_cascade_monitoring(self):
|
||
"""约束注入 + 级联故障监控的完整链路"""
|
||
alignment_config = AlignmentConfig(
|
||
constraints=["password"],
|
||
cascade_max_interactions=5,
|
||
)
|
||
guard = AlignmentGuard(config=alignment_config)
|
||
|
||
# Inject constraints
|
||
input_data = {"content": "请帮我重置密码"}
|
||
injected = guard.inject_constraints(input_data)
|
||
assert "alignment_constraints" in injected
|
||
|
||
# Simulate safe output
|
||
output = {"result": "密码重置链接已发送到您的邮箱。"}
|
||
check_result = await guard.check_output(output)
|
||
assert check_result.passed is True
|
||
|
||
# Record interactions — no cascade alert
|
||
alert = guard.record_interaction("session-chain-1")
|
||
assert alert is None
|
||
assert guard.get_interaction_count("session-chain-1") == 1
|