449 lines
18 KiB
Python
449 lines
18 KiB
Python
"""U6: Lead planning-time historical reflection retrieval (R12, R13).
|
|
|
|
Covers:
|
|
- ReflexionEngine.retrieve_prompt_reflection(): returns best reflection by score
|
|
- Score filtering: score <= 0.5 not returned
|
|
- No episodic_memory → returns None
|
|
- Retrieval failure → returns None (non-blocking)
|
|
- TeamOrchestrator._decompose_task: prepends improved_prompt when reflection found
|
|
- No reflexion_engine → default prompt (backward compat)
|
|
- Retrieval failure → default prompt (non-blocking)
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
from unittest.mock import AsyncMock, MagicMock
|
|
|
|
import pytest
|
|
|
|
from agentkit.core.reflexion import ReflexionEngine
|
|
from agentkit.experts.orchestrator import TeamOrchestrator
|
|
from agentkit.experts.team import ExpertTeam
|
|
from agentkit.memory.base import MemoryItem
|
|
|
|
from tests.unit.experts._helpers import make_execute_stream_mock
|
|
|
|
|
|
# ── Helpers ────────────────────────────────────────────────────────────
|
|
|
|
|
|
def _make_memory_item(
|
|
score: float = 0.8,
|
|
output_summary: str = "improved prompt text",
|
|
reflection: str = "reflection text",
|
|
version: int = 1,
|
|
) -> MemoryItem:
|
|
from datetime import datetime, timezone
|
|
|
|
# Shape matches EpisodicMemory.list_prompt_reflections_by_hash return:
|
|
# output_summary/reflection/quality_score live in value dict,
|
|
# version/task_hash live in metadata.
|
|
return MemoryItem(
|
|
key="prompt_reflection:abc:1",
|
|
value={
|
|
"input_summary": "test",
|
|
"output_summary": output_summary,
|
|
"reflection": reflection,
|
|
"quality_score": score,
|
|
},
|
|
metadata={
|
|
"task_type": "prompt_reflection",
|
|
"version": version,
|
|
"task_hash": "abc",
|
|
"created_at": datetime.now(timezone.utc).isoformat(),
|
|
},
|
|
score=score,
|
|
created_at=datetime.now(timezone.utc),
|
|
)
|
|
|
|
|
|
def _make_llm_gateway_mock() -> MagicMock:
|
|
gw = MagicMock()
|
|
response = MagicMock()
|
|
response.content = (
|
|
'[{"name":"A","assigned_expert":"lead","task_description":"a","depends_on":[]}]'
|
|
)
|
|
gw.chat = AsyncMock(return_value=response)
|
|
return gw
|
|
|
|
|
|
def _make_team_with_experts() -> ExpertTeam:
|
|
from agentkit.core.handoff_transport import InProcessHandoffTransport
|
|
from agentkit.core.protocol import TaskResult, TaskStatus
|
|
from agentkit.experts.config import ExpertConfig
|
|
from agentkit.experts.expert import Expert
|
|
|
|
team = ExpertTeam()
|
|
team._handoff_transport = AsyncMock(spec=InProcessHandoffTransport)
|
|
|
|
config = ExpertConfig(
|
|
name="lead",
|
|
agent_type="expert",
|
|
persona="测试",
|
|
thinking_style="逻辑",
|
|
bound_skills=["s"],
|
|
is_lead=True,
|
|
task_mode="llm_generate",
|
|
prompt={"identity": "测试"},
|
|
)
|
|
expert = MagicMock(spec=Expert)
|
|
expert.config = config
|
|
expert.is_active = True
|
|
expert.team_id = None
|
|
expert.get_capabilities_summary.return_value = {"name": "lead"}
|
|
|
|
mock_agent = MagicMock()
|
|
mock_agent.execute = AsyncMock(
|
|
return_value=TaskResult(
|
|
task_id="t",
|
|
agent_name="lead",
|
|
status=TaskStatus.COMPLETED.value,
|
|
output_data={"content": "result"},
|
|
error_message=None,
|
|
started_at=None,
|
|
completed_at=None,
|
|
)
|
|
)
|
|
mock_agent.execute_stream = make_execute_stream_mock("result")
|
|
mock_agent._llm_gateway = None
|
|
expert.agent = mock_agent
|
|
|
|
team._experts["lead"] = expert
|
|
team._lead_expert_name = "lead"
|
|
return team
|
|
|
|
|
|
# ── ReflexionEngine.retrieve_prompt_reflection ─────────────────────────
|
|
|
|
|
|
class TestRetrievePromptReflection:
|
|
"""U6/R12: retrieve_prompt_reflection returns best reflection by score."""
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_returns_none_when_no_episodic_memory(self):
|
|
gw = MagicMock()
|
|
engine = ReflexionEngine(llm_gateway=gw, episodic_memory=None)
|
|
result = await engine.retrieve_prompt_reflection(task_input="test")
|
|
assert result is None
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_returns_best_reflection_by_score(self):
|
|
episodic = MagicMock()
|
|
# Two results: score 0.6 and 0.9 — should return 0.9
|
|
items = [
|
|
_make_memory_item(score=0.6, output_summary="prompt v1"),
|
|
_make_memory_item(score=0.9, output_summary="prompt v2"),
|
|
]
|
|
episodic.search_prompt_reflections = AsyncMock(return_value=items)
|
|
gw = MagicMock()
|
|
engine = ReflexionEngine(llm_gateway=gw, episodic_memory=episodic)
|
|
|
|
result = await engine.retrieve_prompt_reflection(task_input="test")
|
|
|
|
assert result is not None
|
|
assert result["score"] == 0.9
|
|
assert result["improved_prompt"] == "prompt v2"
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_filters_low_score_reflections(self):
|
|
"""score <= 0.5 should not be returned."""
|
|
episodic = MagicMock()
|
|
items = [_make_memory_item(score=0.3, output_summary="low score")]
|
|
episodic.search_prompt_reflections = AsyncMock(return_value=items)
|
|
gw = MagicMock()
|
|
engine = ReflexionEngine(llm_gateway=gw, episodic_memory=episodic)
|
|
|
|
result = await engine.retrieve_prompt_reflection(task_input="test", min_score=0.5)
|
|
assert result is None
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_returns_none_when_no_results(self):
|
|
episodic = MagicMock()
|
|
episodic.search_prompt_reflections = AsyncMock(return_value=[])
|
|
gw = MagicMock()
|
|
engine = ReflexionEngine(llm_gateway=gw, episodic_memory=episodic)
|
|
|
|
result = await engine.retrieve_prompt_reflection(task_input="test")
|
|
assert result is None
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_returns_none_on_search_failure(self):
|
|
episodic = MagicMock()
|
|
episodic.search_prompt_reflections = AsyncMock(side_effect=RuntimeError("DB down"))
|
|
gw = MagicMock()
|
|
engine = ReflexionEngine(llm_gateway=gw, episodic_memory=episodic)
|
|
|
|
result = await engine.retrieve_prompt_reflection(task_input="test")
|
|
assert result is None
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_returns_reflection_fields_complete(self):
|
|
episodic = MagicMock()
|
|
items = [
|
|
_make_memory_item(
|
|
score=0.85,
|
|
output_summary="improved prompt",
|
|
reflection="what went wrong",
|
|
version=3,
|
|
)
|
|
]
|
|
episodic.search_prompt_reflections = AsyncMock(return_value=items)
|
|
gw = MagicMock()
|
|
engine = ReflexionEngine(llm_gateway=gw, episodic_memory=episodic)
|
|
|
|
result = await engine.retrieve_prompt_reflection(task_input="test")
|
|
|
|
assert result is not None
|
|
assert result["improved_prompt"] == "improved prompt"
|
|
assert result["reflection"] == "what went wrong"
|
|
assert result["version"] == 3
|
|
assert result["score"] == 0.85
|
|
assert "task_hash" in result
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_custom_min_score_threshold(self):
|
|
"""min_score=0.7 filters out score=0.6."""
|
|
episodic = MagicMock()
|
|
items = [_make_memory_item(score=0.6, output_summary="medium")]
|
|
episodic.search_prompt_reflections = AsyncMock(return_value=items)
|
|
gw = MagicMock()
|
|
engine = ReflexionEngine(llm_gateway=gw, episodic_memory=episodic)
|
|
|
|
result = await engine.retrieve_prompt_reflection(task_input="test", min_score=0.7)
|
|
assert result is None
|
|
|
|
|
|
# ── TeamOrchestrator._decompose_task with reflection ──────────────────
|
|
|
|
|
|
class TestDecomposeWithReflection:
|
|
"""U6/R13: _decompose_task prepends improved_prompt when reflection found."""
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_decompose_prepends_reflection_when_found(self):
|
|
"""When reflexion_engine returns a reflection, the decomposition
|
|
prompt includes the improved_prompt."""
|
|
team = _make_team_with_experts()
|
|
gw = _make_llm_gateway_mock()
|
|
team._experts["lead"].agent._llm_gateway = gw
|
|
|
|
# Mock reflexion_engine
|
|
reflexion = MagicMock(spec=ReflexionEngine)
|
|
reflexion.retrieve_prompt_reflection = AsyncMock(
|
|
return_value={
|
|
"improved_prompt": "USE THIS IMPROVED APPROACH",
|
|
"score": 0.85,
|
|
"reflection": "past mistake",
|
|
"version": 2,
|
|
"task_hash": "abc",
|
|
}
|
|
)
|
|
|
|
orchestrator = TeamOrchestrator(team, reflexion_engine=reflexion)
|
|
await orchestrator._decompose_task(
|
|
team.lead_expert, "test task", trigger_reason="verify_retry"
|
|
)
|
|
|
|
# Verify retrieve was called
|
|
reflexion.retrieve_prompt_reflection.assert_awaited_once()
|
|
# Verify the prompt sent to LLM includes the improved_prompt
|
|
call_kwargs = gw.chat.await_args.kwargs
|
|
messages = call_kwargs.get("messages") or gw.chat.await_args.args[0]
|
|
prompt_content = messages[0]["content"] if isinstance(messages, list) else str(messages)
|
|
assert "USE THIS IMPROVED APPROACH" in prompt_content
|
|
assert "Historical Reflection" in prompt_content
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_decompose_uses_default_when_no_reflexion_engine(self):
|
|
"""No reflexion_engine → default prompt (backward compat)."""
|
|
team = _make_team_with_experts()
|
|
gw = _make_llm_gateway_mock()
|
|
team._experts["lead"].agent._llm_gateway = gw
|
|
|
|
orchestrator = TeamOrchestrator(team, reflexion_engine=None)
|
|
await orchestrator._decompose_task(team.lead_expert, "test task")
|
|
|
|
# Verify default prompt (no Historical Reflection section)
|
|
call_kwargs = gw.chat.await_args.kwargs
|
|
messages = call_kwargs.get("messages") or gw.chat.await_args.args[0]
|
|
prompt_content = messages[0]["content"] if isinstance(messages, list) else str(messages)
|
|
assert "Historical Reflection" not in prompt_content
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_decompose_uses_default_when_no_reflection_found(self):
|
|
"""reflexion_engine returns None → default prompt."""
|
|
team = _make_team_with_experts()
|
|
gw = _make_llm_gateway_mock()
|
|
team._experts["lead"].agent._llm_gateway = gw
|
|
|
|
reflexion = MagicMock(spec=ReflexionEngine)
|
|
reflexion.retrieve_prompt_reflection = AsyncMock(return_value=None)
|
|
|
|
orchestrator = TeamOrchestrator(team, reflexion_engine=reflexion)
|
|
await orchestrator._decompose_task(
|
|
team.lead_expert, "test task", trigger_reason="verify_retry"
|
|
)
|
|
|
|
call_kwargs = gw.chat.await_args.kwargs
|
|
messages = call_kwargs.get("messages") or gw.chat.await_args.args[0]
|
|
prompt_content = messages[0]["content"] if isinstance(messages, list) else str(messages)
|
|
assert "Historical Reflection" not in prompt_content
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_decompose_uses_default_when_retrieval_fails(self):
|
|
"""reflexion_engine.retrieve raises → default prompt (non-blocking)."""
|
|
team = _make_team_with_experts()
|
|
gw = _make_llm_gateway_mock()
|
|
team._experts["lead"].agent._llm_gateway = gw
|
|
|
|
reflexion = MagicMock(spec=ReflexionEngine)
|
|
reflexion.retrieve_prompt_reflection = AsyncMock(side_effect=RuntimeError("search failed"))
|
|
|
|
orchestrator = TeamOrchestrator(team, reflexion_engine=reflexion)
|
|
await orchestrator._decompose_task(
|
|
team.lead_expert, "test task", trigger_reason="verify_retry"
|
|
)
|
|
|
|
# Default prompt used despite retrieval failure
|
|
call_kwargs = gw.chat.await_args.kwargs
|
|
messages = call_kwargs.get("messages") or gw.chat.await_args.args[0]
|
|
prompt_content = messages[0]["content"] if isinstance(messages, list) else str(messages)
|
|
assert "Historical Reflection" not in prompt_content
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_decompose_skips_reflection_when_no_improved_prompt(self):
|
|
"""reflexion_engine returns dict without improved_prompt → no hint."""
|
|
team = _make_team_with_experts()
|
|
gw = _make_llm_gateway_mock()
|
|
team._experts["lead"].agent._llm_gateway = gw
|
|
|
|
reflexion = MagicMock(spec=ReflexionEngine)
|
|
reflexion.retrieve_prompt_reflection = AsyncMock(
|
|
return_value={"improved_prompt": "", "score": 0.8} # empty improved_prompt
|
|
)
|
|
|
|
orchestrator = TeamOrchestrator(team, reflexion_engine=reflexion)
|
|
await orchestrator._decompose_task(
|
|
team.lead_expert, "test task", trigger_reason="verify_retry"
|
|
)
|
|
|
|
call_kwargs = gw.chat.await_args.kwargs
|
|
messages = call_kwargs.get("messages") or gw.chat.await_args.args[0]
|
|
prompt_content = messages[0]["content"] if isinstance(messages, list) else str(messages)
|
|
assert "Historical Reflection" not in prompt_content
|
|
|
|
|
|
# ── U4/KTD5: retrieve_prompt_reflection trigger gating ────────────────
|
|
|
|
|
|
class TestDecomposeTriggerGating:
|
|
"""U4/KTD5: _decompose_task only calls retrieve_prompt_reflection when
|
|
trigger_reason is one of {verify_retry, schema_validation, loop_detection}.
|
|
Non-trigger decompositions skip retrieval entirely (default prompt)."""
|
|
|
|
def _make_orchestrator_with_reflexion(
|
|
self, retrieve_return: object = None, retrieve_side_effect: object | None = None
|
|
) -> tuple[TeamOrchestrator, MagicMock, MagicMock]:
|
|
"""Build orchestrator + gateway + reflexion mock.
|
|
|
|
Returns (orchestrator, gateway, reflexion). retrieve_side_effect, if
|
|
set, takes precedence over retrieve_return.
|
|
"""
|
|
team = _make_team_with_experts()
|
|
gw = _make_llm_gateway_mock()
|
|
team._experts["lead"].agent._llm_gateway = gw
|
|
|
|
reflexion = MagicMock(spec=ReflexionEngine)
|
|
if retrieve_side_effect is not None:
|
|
reflexion.retrieve_prompt_reflection = AsyncMock(side_effect=retrieve_side_effect)
|
|
else:
|
|
reflexion.retrieve_prompt_reflection = AsyncMock(return_value=retrieve_return)
|
|
|
|
orchestrator = TeamOrchestrator(team, reflexion_engine=reflexion)
|
|
return orchestrator, gw, reflexion
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_retrieves_on_verify_retry_trigger(self):
|
|
orchestrator, _, reflexion = self._make_orchestrator_with_reflexion(
|
|
retrieve_return={
|
|
"improved_prompt": "IMPROVED",
|
|
"score": 0.8,
|
|
"reflection": "r",
|
|
"version": 1,
|
|
"task_hash": "abc",
|
|
}
|
|
)
|
|
await orchestrator._decompose_task(
|
|
orchestrator._team.lead_expert, "task", trigger_reason="verify_retry"
|
|
)
|
|
reflexion.retrieve_prompt_reflection.assert_awaited_once()
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_retrieves_on_schema_validation_trigger(self):
|
|
orchestrator, _, reflexion = self._make_orchestrator_with_reflexion(retrieve_return=None)
|
|
await orchestrator._decompose_task(
|
|
orchestrator._team.lead_expert, "task", trigger_reason="schema_validation"
|
|
)
|
|
reflexion.retrieve_prompt_reflection.assert_awaited_once()
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_retrieves_on_loop_detection_trigger(self):
|
|
orchestrator, _, reflexion = self._make_orchestrator_with_reflexion(retrieve_return=None)
|
|
await orchestrator._decompose_task(
|
|
orchestrator._team.lead_expert, "task", trigger_reason="loop_detection"
|
|
)
|
|
reflexion.retrieve_prompt_reflection.assert_awaited_once()
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_skips_retrieval_when_trigger_is_none(self):
|
|
"""Fresh decomposition (trigger_reason=None) → no retrieval."""
|
|
orchestrator, _, reflexion = self._make_orchestrator_with_reflexion(
|
|
retrieve_return={"improved_prompt": "would-be-used", "score": 0.9}
|
|
)
|
|
await orchestrator._decompose_task(orchestrator._team.lead_expert, "task")
|
|
reflexion.retrieve_prompt_reflection.assert_not_awaited()
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_skips_retrieval_on_invalid_trigger_reason(self):
|
|
"""trigger_reason not in the allowed set → no retrieval."""
|
|
orchestrator, _, reflexion = self._make_orchestrator_with_reflexion(
|
|
retrieve_return={"improved_prompt": "would-be-used", "score": 0.9}
|
|
)
|
|
await orchestrator._decompose_task(
|
|
orchestrator._team.lead_expert, "task", trigger_reason="random"
|
|
)
|
|
reflexion.retrieve_prompt_reflection.assert_not_awaited()
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_triggered_retrieval_failure_degrades_to_default_prompt(self):
|
|
"""trigger_reason valid but retrieve raises → default prompt (non-blocking)."""
|
|
orchestrator, gw, reflexion = self._make_orchestrator_with_reflexion(
|
|
retrieve_side_effect=RuntimeError("search down")
|
|
)
|
|
await orchestrator._decompose_task(
|
|
orchestrator._team.lead_expert, "task", trigger_reason="verify_retry"
|
|
)
|
|
# Retrieval was attempted (triggered) but failed — default prompt used.
|
|
reflexion.retrieve_prompt_reflection.assert_awaited_once()
|
|
call_kwargs = gw.chat.await_args.kwargs
|
|
messages = call_kwargs.get("messages") or gw.chat.await_args.args[0]
|
|
prompt_content = messages[0]["content"] if isinstance(messages, list) else str(messages)
|
|
assert "Historical Reflection" not in prompt_content
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_skips_retrieval_when_no_reflexion_engine_even_with_trigger(self):
|
|
"""No reflexion_engine → retrieval skipped regardless of trigger_reason."""
|
|
team = _make_team_with_experts()
|
|
gw = _make_llm_gateway_mock()
|
|
team._experts["lead"].agent._llm_gateway = gw
|
|
|
|
orchestrator = TeamOrchestrator(team, reflexion_engine=None)
|
|
await orchestrator._decompose_task(
|
|
orchestrator._team.lead_expert, "task", trigger_reason="verify_retry"
|
|
)
|
|
# Default prompt — no Historical Reflection section.
|
|
call_kwargs = gw.chat.await_args.kwargs
|
|
messages = call_kwargs.get("messages") or gw.chat.await_args.args[0]
|
|
prompt_content = messages[0]["content"] if isinstance(messages, list) else str(messages)
|
|
assert "Historical Reflection" not in prompt_content
|