49 lines
1.7 KiB
Python
49 lines
1.7 KiB
Python
from unittest.mock import AsyncMock
|
|
|
|
import pytest
|
|
|
|
from blueprints.conversation import memory
|
|
|
|
|
|
def test_old_relevant_memory_beats_recent_unrelated():
|
|
memories = ["Favorite restaurant is Nori", "Cat Simba takes insulin"]
|
|
assert (
|
|
memory.rank_memories(memories, "What insulin does Simba take?")[0]
|
|
== memories[1]
|
|
)
|
|
|
|
|
|
@pytest.mark.parametrize("query", ["", "please", "quantum oscillators"])
|
|
def test_empty_or_unknown_query_retains_recency(query):
|
|
memories = ["Favorite restaurant is Nori", "Cat Simba takes insulin"]
|
|
assert memory.rank_memories(memories, query) == memories
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_prompt_relevance_and_total_memory_budget(monkeypatch):
|
|
old = "Simba insulin dose is two units."
|
|
recent = "Favorite restaurant is Nori. " * 20
|
|
monkeypatch.setattr(
|
|
memory, "get_memories_for_user", AsyncMock(return_value=[recent, old])
|
|
)
|
|
monkeypatch.setattr(memory, "MEMORY_BUDGET_TOKENS", 110)
|
|
prompt = await memory.build_system_prompt_with_memories(
|
|
"user", query="Simba insulin"
|
|
)
|
|
assert old in prompt
|
|
assert recent not in prompt
|
|
assert memory.estimate_tokens(prompt[len(memory.SIMBA_SYSTEM_PROMPT) :]) <= 110
|
|
memory.get_memories_for_user.assert_awaited_once_with("user")
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_backwards_compatible_no_query_and_empty_memories(monkeypatch):
|
|
fetch = AsyncMock(return_value=[])
|
|
monkeypatch.setattr(memory, "get_memories_for_user", fetch)
|
|
assert (
|
|
await memory.build_system_prompt_with_memories("user")
|
|
== memory.SIMBA_SYSTEM_PROMPT
|
|
)
|
|
fetch.return_value = ["Likes tea"]
|
|
assert "Likes tea" in await memory.build_system_prompt_with_memories("user")
|