46 lines
1.6 KiB
Python
46 lines
1.6 KiB
Python
from memabra.candidate_types import CandidateObject
|
|
from memabra.retrieval import CandidateRetriever, InMemoryCandidateProvider
|
|
from memabra.router import TaskContext
|
|
|
|
|
|
def test_retriever_ranks_trigger_matches_first():
|
|
retriever = CandidateRetriever(
|
|
[
|
|
InMemoryCandidateProvider(
|
|
candidate_type="memory",
|
|
candidates=[
|
|
CandidateObject(
|
|
id="mem-weak",
|
|
type="memory",
|
|
title="Generic preference",
|
|
summary="A weak preference record",
|
|
confidence=0.4,
|
|
success_rate=0.4,
|
|
freshness=0.4,
|
|
triggers=["generic"],
|
|
),
|
|
CandidateObject(
|
|
id="mem-strong",
|
|
type="memory",
|
|
title="Formatting preference",
|
|
summary="Telegram prefers plain text",
|
|
confidence=0.8,
|
|
success_rate=0.9,
|
|
freshness=0.9,
|
|
triggers=["telegram", "formatting"],
|
|
tags=["output"],
|
|
),
|
|
],
|
|
)
|
|
]
|
|
)
|
|
|
|
result = retriever.retrieve(
|
|
TaskContext(user_input="Use my telegram formatting preference for the output."),
|
|
top_k=2,
|
|
)
|
|
|
|
assert [candidate.id for candidate in result.memory] == ["mem-strong", "mem-weak"]
|
|
assert result.skill == []
|
|
assert result.tool == []
|