diff --git a/src/agentkit/core/reflexion.py b/src/agentkit/core/reflexion.py index aa3c144..e106090 100644 --- a/src/agentkit/core/reflexion.py +++ b/src/agentkit/core/reflexion.py @@ -26,6 +26,7 @@ from agentkit.telemetry.metrics import ( if TYPE_CHECKING: from agentkit.core.compressor import CompressionStrategy from agentkit.core.trace import TraceRecorder + from agentkit.memory.episodic import EpisodicMemory from agentkit.memory.retriever import MemoryRetriever logger = logging.getLogger(__name__) @@ -72,6 +73,9 @@ class ReflexionEngine: max_reflections: int = 3, quality_threshold: float = 0.7, default_timeout: float = 300.0, + # IQ-Boost/U5 (R11): optional EpisodicMemory for persisting reflections + # across tasks. None = no persistence (backward-compatible). + episodic_memory: "EpisodicMemory | None" = None, ): if max_steps < 1: raise ValueError(f"max_steps must be >= 1, got {max_steps}") @@ -87,6 +91,8 @@ class ReflexionEngine: self._max_reflections = max_reflections self._quality_threshold = quality_threshold self._default_timeout = default_timeout + # U5: optional episodic memory for cross-task reflection persistence + self._episodic_memory = episodic_memory self._react_engine = ReActEngine( llm_gateway=llm_gateway, max_steps=max_steps, @@ -654,7 +660,12 @@ class ReflexionEngine: agent_name: str, task_type: str, ) -> str | None: - """反思执行结果,返回反思文本;失败时返回 None""" + """反思执行结果,返回反思文本;失败时返回 None + + IQ-Boost/U5 (R11): if ``self._episodic_memory`` is configured, persist + the reflection for cross-task retrieval. Persistence failure is + non-blocking — the reflection is still returned for in-task retry. + """ task_description = messages[-1].get("content", "") if messages else "" system_message = ( @@ -685,11 +696,33 @@ class ReflexionEngine: agent_name=agent_name, task_type=task_type or "reflection", ) - return response.content or None + reflection_text = response.content or None except Exception as e: logger.warning(f"Reflection LLM call failed, skipping reflection: {e}") return None + # U5/R11: persist reflection to EpisodicMemory (non-blocking) + if reflection_text and self._episodic_memory is not None: + improved_prompt = self._build_reflection_prompt( + original_prompt=None, + reflection_text=reflection_text, + attempt=1, + ) + try: + await self._episodic_memory.store_prompt_reflection( + task_input=task_description, + reflection=reflection_text, + improved_prompt=improved_prompt, + version=1, + score=score, + agent_name=agent_name, + ) + except Exception as e: + # Non-blocking: reflection is still useful for in-task retry + logger.warning(f"U5: failed to persist reflection, continuing: {e}") + + return reflection_text + def _build_reflection_prompt( self, original_prompt: str | None, diff --git a/src/agentkit/memory/episodic.py b/src/agentkit/memory/episodic.py index 6110311..3ced8af 100644 --- a/src/agentkit/memory/episodic.py +++ b/src/agentkit/memory/episodic.py @@ -410,3 +410,101 @@ class EpisodicMemory(Memory): await db.rollback() logger.error(f"Failed to delete episodic memory: {e}") return False + + # ── IQ-Boost/U5: Prompt Reflection persistence (R11, R15) ────────── + + async def store_prompt_reflection( + self, + task_input: str, + reflection: str, + improved_prompt: str, + version: int = 1, + score: float = 0.0, + agent_name: str = "", + task_hash: str | None = None, + ) -> str | None: + """持久化 prompt 反思到 EpisodicMemory,支持跨任务检索 (U5/R11). + + Reuses the existing ``store()`` path with ``task_type="prompt_reflection"`` + as the discriminator. The ORM row's fields map: + input_summary ← task_input (truncated) + output_summary ← improved_prompt (truncated) + reflection ← reflection text + quality_score ← score (0.0=failed, 1.0=verified) + outcome ← "reflection" + agent_name ← agent_name + + Returns the storage key ``"prompt_reflection:{task_hash}:{version}"`` + on success, or None on failure (non-raising — callers continue without). + """ + import hashlib + + if task_hash is None: + task_hash = hashlib.sha256(task_input.encode("utf-8")).hexdigest()[:16] + key = f"prompt_reflection:{task_hash}:{version}" + + value = { + "task_input": task_input[:500], + "reflection": reflection, + "improved_prompt": improved_prompt, + "score": score, + "version": version, + "task_hash": task_hash, + "timestamp": datetime.now(timezone.utc).isoformat(), + } + metadata = { + "agent_name": agent_name, + "task_type": "prompt_reflection", + "output_summary": improved_prompt[:500], + "outcome": "reflection", + "quality_score": score, + "reflection": reflection, + } + try: + await self.store(key=key, value=value, metadata=metadata) + return key + except (DBAPIError, ValueError, KeyError, RuntimeError, OSError) as e: + logger.warning(f"U5: failed to persist prompt reflection: {e}") + return None + + async def search_prompt_reflections( + self, + task_input: str, + top_k: int = 5, + agent_name: str | None = None, + ) -> list[MemoryItem]: + """检索相似 task_input 的历史 prompt 反思 (U5/R11). + + Uses ``search()`` with ``task_type="prompt_reflection"`` filter. + Returns empty list on failure (non-raising). + """ + filters: MetadataDict = {"task_type": "prompt_reflection"} + if agent_name: + filters["agent_name"] = agent_name + try: + return await self.search(query=task_input, top_k=top_k, filters=filters) + except (DBAPIError, ValueError, KeyError, RuntimeError, OSError) as e: + logger.warning(f"U5: failed to search prompt reflections: {e}") + return [] + + async def cleanup_expired(self, max_age_days: int = 30) -> int: + """删除超过 max_age_days 天的记录 (U5/R15 TTL). + + Returns the number of deleted rows. 0 on failure (non-raising). + """ + from datetime import timedelta + + from sqlalchemy import delete as sql_delete + + cutoff = datetime.now(timezone.utc) - timedelta(days=max_age_days) + async with self._session_factory() as db: + try: + Model = self._episodic_model + stmt = sql_delete(Model).where(Model.created_at < cutoff) + result = await db.execute(stmt) + await db.commit() + return result.rowcount or 0 + except (DBAPIError, ValueError, KeyError, RuntimeError) as e: + await db.rollback() + logger.warning(f"U5: cleanup_expired failed: {e}") + return 0 diff --git a/tests/unit/test_reflexion_persist.py b/tests/unit/test_reflexion_persist.py new file mode 100644 index 0000000..1473627 --- /dev/null +++ b/tests/unit/test_reflexion_persist.py @@ -0,0 +1,423 @@ +"""U5: ReflexionEngine reflection persistence to EpisodicMemory (R11, R15). + +Covers: +- store_prompt_reflection() stores via existing store() with task_type discriminator +- search_prompt_reflections() filters by task_type="prompt_reflection" +- cleanup_expired() deletes old records (TTL) +- _reflect() persists reflection when episodic_memory configured (non-blocking) +- _reflect() skips persistence when episodic_memory=None (backward compat) +- Persistence failure does not block _reflect (returns reflection text) +- Multi-version coexistence (same task_hash, different versions) +""" + +from __future__ import annotations + +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest + +from agentkit.core.react import ReActResult +from agentkit.core.reflexion import ReflexionEngine +from agentkit.memory.episodic import EpisodicMemory + + +# ── Helpers ──────────────────────────────────────────────────────────── + + +def _make_episodic_memory_mock() -> MagicMock: + """Create a mock EpisodicMemory that tracks store/search calls.""" + mem = MagicMock(spec=EpisodicMemory) + mem.store = AsyncMock(return_value=None) + mem.search = AsyncMock(return_value=[]) + mem.store_prompt_reflection = AsyncMock(return_value="prompt_reflection:abc:1") + mem.search_prompt_reflections = AsyncMock(return_value=[]) + mem.cleanup_expired = AsyncMock(return_value=0) + return mem + + +def _make_react_result(output: str = "test output", status: str = "completed") -> ReActResult: + return ReActResult( + output=output, + trajectory=[], + total_steps=1, + total_tokens=10, + status=status, + ) + + +def _make_llm_gateway_mock(reflection_text: str = "reflection text") -> MagicMock: + """Gateway that returns reflection text from chat().""" + gw = MagicMock() + response = MagicMock() + response.content = reflection_text + gw.chat = AsyncMock(return_value=response) + return gw + + +# ── EpisodicMemory.store_prompt_reflection ───────────────────────────── + + +class TestStorePromptReflection: + """U5/R11: store_prompt_reflection persists via store() with + task_type='prompt_reflection' discriminator.""" + + @pytest.mark.asyncio + async def test_store_calls_underlying_store_with_correct_metadata(self): + mem = MagicMock(spec=EpisodicMemory) + mem.store = AsyncMock(return_value=None) + # We need to call the real method, not a mock — patch store only + # Use the unbound method pattern + await EpisodicMemory.store_prompt_reflection( + mem, + task_input="test task", + reflection="reflection text", + improved_prompt="improved prompt", + version=1, + score=0.5, + agent_name="test_agent", + ) + + mem.store.assert_awaited_once() + call_kwargs = mem.store.await_args.kwargs + assert "key" in call_kwargs + assert call_kwargs["key"].startswith("prompt_reflection:") + assert ":1" in call_kwargs["key"] + + value = call_kwargs["value"] + assert value["task_input"] == "test task" + assert value["reflection"] == "reflection text" + assert value["improved_prompt"] == "improved prompt" + assert value["score"] == 0.5 + assert value["version"] == 1 + assert "timestamp" in value + + metadata = call_kwargs["metadata"] + assert metadata["task_type"] == "prompt_reflection" + assert metadata["agent_name"] == "test_agent" + assert metadata["quality_score"] == 0.5 + assert metadata["reflection"] == "reflection text" + + @pytest.mark.asyncio + async def test_store_returns_key_on_success(self): + mem = MagicMock(spec=EpisodicMemory) + mem.store = AsyncMock(return_value=None) + key = await EpisodicMemory.store_prompt_reflection( + mem, + task_input="test", + reflection="r", + improved_prompt="p", + version=2, + ) + assert key is not None + assert ":2" in key + assert key.startswith("prompt_reflection:") + + @pytest.mark.asyncio + async def test_store_returns_none_on_failure(self): + from sqlalchemy.exc import DBAPIError + + mem = MagicMock(spec=EpisodicMemory) + mem.store = AsyncMock(side_effect=DBAPIError("stmt", params={}, orig=Exception("db down"))) + key = await EpisodicMemory.store_prompt_reflection( + mem, + task_input="test", + reflection="r", + improved_prompt="p", + ) + assert key is None + + @pytest.mark.asyncio + async def test_store_uses_provided_task_hash(self): + mem = MagicMock(spec=EpisodicMemory) + mem.store = AsyncMock(return_value=None) + key = await EpisodicMemory.store_prompt_reflection( + mem, + task_input="test", + reflection="r", + improved_prompt="p", + task_hash="custom_hash", + ) + assert "custom_hash" in key + + @pytest.mark.asyncio + async def test_store_generates_task_hash_from_input(self): + mem = MagicMock(spec=EpisodicMemory) + mem.store = AsyncMock(return_value=None) + key1 = await EpisodicMemory.store_prompt_reflection( + mem, task_input="same task", reflection="r", improved_prompt="p" + ) + key2 = await EpisodicMemory.store_prompt_reflection( + mem, task_input="same task", reflection="r", improved_prompt="p" + ) + # Same task_input → same hash prefix + assert key1 == key2 + + +# ── EpisodicMemory.search_prompt_reflections ─────────────────────────── + + +class TestSearchPromptReflections: + """U5/R11: search_prompt_reflections filters by task_type.""" + + @pytest.mark.asyncio + async def test_search_calls_underlying_search_with_filter(self): + mem = MagicMock(spec=EpisodicMemory) + mem.search = AsyncMock(return_value=[]) + await EpisodicMemory.search_prompt_reflections(mem, task_input="find similar", top_k=3) + + mem.search.assert_awaited_once() + call_kwargs = mem.search.await_args.kwargs + assert call_kwargs["query"] == "find similar" + assert call_kwargs["top_k"] == 3 + filters = call_kwargs["filters"] + assert filters["task_type"] == "prompt_reflection" + + @pytest.mark.asyncio + async def test_search_includes_agent_name_filter_when_provided(self): + mem = MagicMock(spec=EpisodicMemory) + mem.search = AsyncMock(return_value=[]) + await EpisodicMemory.search_prompt_reflections(mem, task_input="q", agent_name="agent_x") + + filters = mem.search.await_args.kwargs["filters"] + assert filters["agent_name"] == "agent_x" + + @pytest.mark.asyncio + async def test_search_returns_empty_on_failure(self): + from sqlalchemy.exc import DBAPIError + + mem = MagicMock(spec=EpisodicMemory) + mem.search = AsyncMock(side_effect=DBAPIError("stmt", params={}, orig=Exception("err"))) + result = await EpisodicMemory.search_prompt_reflections(mem, task_input="q") + assert result == [] + + +# ── EpisodicMemory.cleanup_expired ───────────────────────────────────── + + +class TestCleanupExpired: + """U5/R15: cleanup_expired deletes records older than max_age_days.""" + + @pytest.mark.asyncio + async def test_cleanup_calls_delete_with_cutoff(self): + # Patch sqlalchemy.delete to bypass ORM model requirement — we only + # verify the session.execute/commit calls happen. + mock_db = AsyncMock() + mock_result = MagicMock() + mock_result.rowcount = 5 + mock_db.execute = AsyncMock(return_value=mock_result) + mock_db.commit = AsyncMock() + + mock_session_factory = MagicMock() + mock_session_factory.return_value.__aenter__ = AsyncMock(return_value=mock_db) + mock_session_factory.return_value.__aexit__ = AsyncMock(return_value=None) + + # created_at must support `<` comparison with datetime — configure + # __lt__ to return a truthy mock so `.where(...)` gets a valid arg. + mock_model = MagicMock() + mock_model.created_at = MagicMock() + mock_model.created_at.__lt__ = MagicMock(return_value=MagicMock()) + + mem = EpisodicMemory( + session_factory=mock_session_factory, + episodic_model=mock_model, + embedder=None, + pgvector_enabled=False, + ) + + # Patch sql_delete to return a chainable mock + fake_stmt = MagicMock() + fake_stmt.where = MagicMock(return_value=fake_stmt) + with patch("sqlalchemy.delete", return_value=fake_stmt): + deleted = await mem.cleanup_expired(max_age_days=30) + + assert deleted == 5 + mock_db.execute.assert_awaited_once() + mock_db.commit.assert_awaited_once() + + @pytest.mark.asyncio + async def test_cleanup_returns_zero_on_failure(self): + from sqlalchemy.exc import DBAPIError + + mock_db = AsyncMock() + mock_db.execute = AsyncMock( + side_effect=DBAPIError("stmt", params={}, orig=Exception("err")) + ) + mock_db.rollback = AsyncMock() + + mock_session_factory = MagicMock() + mock_session_factory.return_value.__aenter__ = AsyncMock(return_value=mock_db) + mock_session_factory.return_value.__aexit__ = AsyncMock(return_value=None) + + mock_model = MagicMock() + mock_model.created_at = MagicMock() + mock_model.created_at.__lt__ = MagicMock(return_value=MagicMock()) + + mem = EpisodicMemory( + session_factory=mock_session_factory, + episodic_model=mock_model, + embedder=None, + pgvector_enabled=False, + ) + + fake_stmt = MagicMock() + fake_stmt.where = MagicMock(return_value=fake_stmt) + with patch("sqlalchemy.delete", return_value=fake_stmt): + deleted = await mem.cleanup_expired(max_age_days=30) + + assert deleted == 0 + mock_db.rollback.assert_awaited_once() + + +# ── ReflexionEngine._reflect persistence ────────────────────────────── + + +class TestReflectPersistence: + """U5/R11: _reflect persists reflection when episodic_memory configured.""" + + @pytest.mark.asyncio + async def test_reflect_persists_when_episodic_memory_configured(self): + episodic = _make_episodic_memory_mock() + gw = _make_llm_gateway_mock(reflection_text="my reflection") + engine = ReflexionEngine( + llm_gateway=gw, + episodic_memory=episodic, + ) + + result = await engine._reflect( + react_result=_make_react_result(), + score=0.3, + messages=[{"role": "user", "content": "test task"}], + reflect_model="default", + agent_name="test_agent", + task_type="test", + ) + + assert result == "my reflection" + episodic.store_prompt_reflection.assert_awaited_once() + call_kwargs = episodic.store_prompt_reflection.await_args.kwargs + assert call_kwargs["task_input"] == "test task" + assert call_kwargs["reflection"] == "my reflection" + assert call_kwargs["score"] == 0.3 + assert call_kwargs["agent_name"] == "test_agent" + + @pytest.mark.asyncio + async def test_reflect_skips_persistence_when_no_episodic_memory(self): + gw = _make_llm_gateway_mock(reflection_text="my reflection") + engine = ReflexionEngine( + llm_gateway=gw, + episodic_memory=None, # No persistence + ) + + result = await engine._reflect( + react_result=_make_react_result(), + score=0.3, + messages=[{"role": "user", "content": "test task"}], + reflect_model="default", + agent_name="test_agent", + task_type="test", + ) + + assert result == "my reflection" + # No episodic_memory → no store call + + @pytest.mark.asyncio + async def test_reflect_persistence_failure_does_not_block(self): + """If store_prompt_reflection raises, _reflect still returns reflection.""" + episodic = MagicMock(spec=EpisodicMemory) + episodic.store_prompt_reflection = AsyncMock(side_effect=RuntimeError("DB down")) + gw = _make_llm_gateway_mock(reflection_text="important reflection") + engine = ReflexionEngine( + llm_gateway=gw, + episodic_memory=episodic, + ) + + result = await engine._reflect( + react_result=_make_react_result(), + score=0.3, + messages=[{"role": "user", "content": "test task"}], + reflect_model="default", + agent_name="test_agent", + task_type="test", + ) + + # Reflection still returned despite persistence failure + assert result == "important reflection" + episodic.store_prompt_reflection.assert_awaited_once() + + @pytest.mark.asyncio + async def test_reflect_skips_persistence_when_llm_returns_none(self): + """If LLM returns empty/None, no persistence attempted.""" + episodic = _make_episodic_memory_mock() + gw = MagicMock() + response = MagicMock() + response.content = None # Empty reflection + gw.chat = AsyncMock(return_value=response) + engine = ReflexionEngine( + llm_gateway=gw, + episodic_memory=episodic, + ) + + result = await engine._reflect( + react_result=_make_react_result(), + score=0.3, + messages=[{"role": "user", "content": "test task"}], + reflect_model="default", + agent_name="test_agent", + task_type="test", + ) + + assert result is None + episodic.store_prompt_reflection.assert_not_awaited() + + @pytest.mark.asyncio + async def test_reflect_skips_persistence_on_llm_failure(self): + """If LLM call raises, no persistence attempted.""" + episodic = _make_episodic_memory_mock() + gw = MagicMock() + gw.chat = AsyncMock(side_effect=RuntimeError("LLM down")) + engine = ReflexionEngine( + llm_gateway=gw, + episodic_memory=episodic, + ) + + result = await engine._reflect( + react_result=_make_react_result(), + score=0.3, + messages=[{"role": "user", "content": "test task"}], + reflect_model="default", + agent_name="test_agent", + task_type="test", + ) + + assert result is None + episodic.store_prompt_reflection.assert_not_awaited() + + +# ── Multi-version coexistence ───────────────────────────────────────── + + +class TestMultiVersionCoexistence: + """U5/R11: same task_hash with different versions all stored.""" + + @pytest.mark.asyncio + async def test_multiple_versions_stored_with_incrementing_version(self): + mem = MagicMock(spec=EpisodicMemory) + mem.store = AsyncMock(return_value=None) + + # Store v1, v2, v3 for same task + key1 = await EpisodicMemory.store_prompt_reflection( + mem, task_input="same task", reflection="r1", improved_prompt="p1", version=1 + ) + key2 = await EpisodicMemory.store_prompt_reflection( + mem, task_input="same task", reflection="r2", improved_prompt="p2", version=2 + ) + key3 = await EpisodicMemory.store_prompt_reflection( + mem, task_input="same task", reflection="r3", improved_prompt="p3", version=3 + ) + + # All keys have same task_hash prefix but different version suffix + assert key1 != key2 != key3 + assert ":1" in key1 and ":2" in key2 and ":3" in key3 + + # All three store() calls made + assert mem.store.await_count == 3