Files
memabra/tests/test_runner.py
2026-04-15 11:06:05 +08:00

97 lines
3.2 KiB
Python

from memabra.candidate_types import CandidateObject
from memabra.retrieval import CandidateRetriever, InMemoryCandidateProvider
from memabra.router import RuleBasedRouter, TaskContext
from memabra.runner import MemabraRunner
from memabra.schemas import SchemaRegistry
def test_runner_produces_valid_draft_trajectory():
retriever = CandidateRetriever(
[
InMemoryCandidateProvider(
candidate_type="memory",
candidates=[
CandidateObject(
id="mem-1",
type="memory",
title="Output preference",
summary="Prefer plain text on Telegram.",
triggers=["telegram", "preference"],
confidence=0.9,
success_rate=0.8,
freshness=0.9,
tags=["output"],
)
],
)
]
)
runner = MemabraRunner(retriever=retriever, router=RuleBasedRouter())
trajectory = runner.run(
context=TaskContext(
user_input="Use my telegram preference for this answer.",
conversation_summary="User often cares about output formatting.",
),
channel="telegram",
user_id="oza",
)
SchemaRegistry().validate_trajectory(trajectory)
assert trajectory["decisions"][0]["decision_type"] == "inject_memory"
assert trajectory["candidate_sets"]["memory"][0]["id"] == "mem-1"
assert len(trajectory["events"]) == 3
def test_runner_injects_episodic_candidate_when_case_index_matches(tmp_path):
from memabra.case_index import CaseIndex
from memabra.persistence import PersistenceStore
store = PersistenceStore(base_dir=tmp_path / "artifacts")
retriever = CandidateRetriever(
[
InMemoryCandidateProvider(
candidate_type="memory",
candidates=[],
),
InMemoryCandidateProvider(
candidate_type="skill",
candidates=[],
),
InMemoryCandidateProvider(
candidate_type="tool",
candidates=[],
),
]
)
runner = MemabraRunner(retriever=retriever, router=RuleBasedRouter(), persistence_store=store)
# First run creates a trajectory
trajectory1 = runner.run(
context=TaskContext(user_input="Hello world"),
channel="local",
persist=True,
)
# Build case index from the trajectory
case_index = CaseIndex()
case_index.add(trajectory1)
# Second run with case index should inject an episodic candidate
runner_with_case = MemabraRunner(
retriever=retriever,
router=RuleBasedRouter(),
persistence_store=store,
case_index=case_index,
)
trajectory2 = runner_with_case.run(
context=TaskContext(user_input="Hello world"),
channel="local",
persist=True,
)
memory_candidates = trajectory2["candidate_sets"]["memory"]
assert any(c["id"].startswith("episodic-") for c in memory_candidates)
# With a persistence store, the runner should generate a rich episodic summary
assert any("Task:" in c["summary"] for c in memory_candidates)