515 lines
20 KiB
Python
515 lines
20 KiB
Python
"""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 direct ORM write with
|
|
task_type='prompt_reflection' discriminator and metadata_ JSONB."""
|
|
|
|
def _make_persistable_mock(self) -> MagicMock:
|
|
"""Build an EpisodicMemory mock with _embedder + _session_factory
|
|
configured so store_prompt_reflection can run its real code path."""
|
|
mem = MagicMock(spec=EpisodicMemory)
|
|
mem._embedder = None # skip embedding generation in test
|
|
mock_db = AsyncMock()
|
|
mock_db.add = MagicMock()
|
|
mock_db.commit = AsyncMock()
|
|
mock_db.rollback = AsyncMock()
|
|
mem._session_factory = MagicMock()
|
|
mem._session_factory.return_value.__aenter__ = AsyncMock(return_value=mock_db)
|
|
mem._session_factory.return_value.__aexit__ = AsyncMock(return_value=None)
|
|
mem._episodic_model = MagicMock()
|
|
# U2: store_prompt_reflection now calls cleanup_expired lazily; provide
|
|
# an AsyncMock so the real code path can await it without DB side effects.
|
|
mem.cleanup_expired = AsyncMock(return_value=0)
|
|
# Track the ORM entry passed to db.add for assertions
|
|
mem._mock_db = mock_db
|
|
return mem
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_store_writes_orm_entry_with_correct_fields(self):
|
|
mem = self._make_persistable_mock()
|
|
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",
|
|
)
|
|
|
|
# Verify ORM entry created via model constructor + db.add + commit
|
|
mem._episodic_model.assert_called_once()
|
|
call_kwargs = mem._episodic_model.call_args.kwargs
|
|
assert call_kwargs["task_type"] == "prompt_reflection"
|
|
assert call_kwargs["agent_name"] == "test_agent"
|
|
assert call_kwargs["input_summary"] == "test task"
|
|
assert call_kwargs["output_summary"] == "improved prompt"
|
|
assert call_kwargs["reflection"] == "reflection text"
|
|
assert call_kwargs["quality_score"] == 0.5
|
|
assert call_kwargs["outcome"] == "reflection"
|
|
# metadata_ JSONB carries task_hash + version + score
|
|
assert "task_hash" in call_kwargs["metadata_"]
|
|
assert call_kwargs["metadata_"]["version"] == 1
|
|
assert call_kwargs["metadata_"]["score"] == 0.5
|
|
mem._mock_db.add.assert_called_once()
|
|
mem._mock_db.commit.assert_awaited_once()
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_store_returns_key_on_success(self):
|
|
mem = self._make_persistable_mock()
|
|
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 = self._make_persistable_mock()
|
|
mem._mock_db.commit = 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
|
|
mem._mock_db.rollback.assert_awaited_once()
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_store_uses_provided_task_hash(self):
|
|
mem = self._make_persistable_mock()
|
|
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 = self._make_persistable_mock()
|
|
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):
|
|
# Reuse the persistable mock helper from TestStorePromptReflection.
|
|
helper = TestStorePromptReflection()
|
|
mem = helper._make_persistable_mock()
|
|
|
|
# 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 ORM entries created via db.add
|
|
assert mem._mock_db.add.call_count == 3
|
|
assert mem._mock_db.commit.await_count == 3
|
|
|
|
|
|
# ── U2: lazy cleanup_expired on store ─────────────────────────────────
|
|
|
|
|
|
class TestLazyCleanupOnStore:
|
|
"""U2: store_prompt_reflection triggers cleanup_expired with 10%
|
|
probability (lazy, fire-and-forget). Cleanup failure never blocks the
|
|
store; cleanup is always filtered to task_type='prompt_reflection'."""
|
|
|
|
def _make_persistable_mock(self) -> MagicMock:
|
|
# Reuse the persistable mock (now includes async cleanup_expired).
|
|
return TestStorePromptReflection()._make_persistable_mock()
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_lazy_cleanup_triggered_below_threshold(self):
|
|
"""random.random() < 0.1 → cleanup_expired awaited once."""
|
|
mem = self._make_persistable_mock()
|
|
with patch("agentkit.memory.episodic.random.random", return_value=0.05):
|
|
await EpisodicMemory.store_prompt_reflection(
|
|
mem, task_input="test", reflection="r", improved_prompt="p"
|
|
)
|
|
mem.cleanup_expired.assert_awaited_once_with(task_type="prompt_reflection")
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_lazy_cleanup_not_triggered_at_or_above_threshold(self):
|
|
"""random.random() >= 0.1 → cleanup_expired not awaited."""
|
|
mem = self._make_persistable_mock()
|
|
with patch("agentkit.memory.episodic.random.random", return_value=0.9):
|
|
await EpisodicMemory.store_prompt_reflection(
|
|
mem, task_input="test", reflection="r", improved_prompt="p"
|
|
)
|
|
mem.cleanup_expired.assert_not_awaited()
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_lazy_cleanup_failure_does_not_block_store(self):
|
|
"""cleanup_expired raises → store still returns a valid key."""
|
|
mem = self._make_persistable_mock()
|
|
mem.cleanup_expired = AsyncMock(side_effect=RuntimeError("cleanup boom"))
|
|
with patch("agentkit.memory.episodic.random.random", return_value=0.05):
|
|
key = await EpisodicMemory.store_prompt_reflection(
|
|
mem, task_input="test", reflection="r", improved_prompt="p"
|
|
)
|
|
assert key is not None # store succeeded despite cleanup failure
|
|
mem.cleanup_expired.assert_awaited_once()
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_lazy_cleanup_filters_to_prompt_reflection_task_type(self):
|
|
"""cleanup_expired is always called with task_type='prompt_reflection'."""
|
|
mem = self._make_persistable_mock()
|
|
with patch("agentkit.memory.episodic.random.random", return_value=0.05):
|
|
await EpisodicMemory.store_prompt_reflection(
|
|
mem, task_input="test", reflection="r", improved_prompt="p"
|
|
)
|
|
# Keyword arg, never None — prevents purging other task types.
|
|
assert mem.cleanup_expired.await_args.kwargs == {"task_type": "prompt_reflection"}
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_lazy_cleanup_triggered_at_least_once_over_many_calls(self):
|
|
"""10+ calls with mixed random draws → cleanup_expired called at least once."""
|
|
mem = self._make_persistable_mock()
|
|
# 9 non-triggering draws (0.9) then 1 triggering (0.05).
|
|
draws = [0.9] * 9 + [0.05]
|
|
with patch("agentkit.memory.episodic.random.random", side_effect=draws):
|
|
for i in range(10):
|
|
await EpisodicMemory.store_prompt_reflection(
|
|
mem,
|
|
task_input=f"task {i}",
|
|
reflection="r",
|
|
improved_prompt="p",
|
|
)
|
|
assert mem.cleanup_expired.await_count >= 1
|
|
# Every cleanup call uses the prompt_reflection filter.
|
|
for call in mem.cleanup_expired.await_args_list:
|
|
assert call.kwargs == {"task_type": "prompt_reflection"}
|