97 lines
3.2 KiB
Python
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)
|