fischer-agentkit/tests/unit/test_lead_reflection_retrieval.py

"""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

    return MemoryItem(
        key="prompt_reflection:abc:1",
        value={"task_input": "test", "reflection": reflection},
        metadata={
            "task_type": "prompt_reflection",
            "output_summary": output_summary,
            "reflection": reflection,
            "version": version,
            "task_hash": "abc",
            "quality_score": score,
        },
        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")

        # 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")

        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")

        # 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")

        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