301 lines
8.6 KiB
Plaintext
301 lines
8.6 KiB
Plaintext
use mem_core::gated_loop::{run_loop, LlmClient, LoopConfig, LoopEvent};
|
|||
|
|
use mem_core::{Chunk, Level, Provenance, Query, Record, Role};
|
||
|
|
use std::sync::{Arc, Mutex};
|
||
|
|
use time::OffsetDateTime;
|
||
|
|
|
||
|
|
/// Fake LLM that returns scripted responses.
|
||
|
|
struct FakeLlm {
|
||
|
|
responses: Arc<Mutex<Vec<String>>>,
|
||
|
|
}
|
||
|
|
|
||
|
|
impl FakeLlm {
|
||
|
|
fn new(responses: Vec<&str>) -> Self {
|
||
|
|
Self {
|
||
|
|
responses: Arc::new(Mutex::new(responses.iter().map(|s| s.to_string()).collect())),
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
impl LlmClient for FakeLlm {
|
||
|
|
fn complete_blocking(&self, _system: &str, _user: &str, _max_tokens: usize) -> anyhow::Result<String> {
|
||
|
|
let mut responses = self.responses.lock().unwrap();
|
||
|
|
if responses.is_empty() {
|
||
|
|
Err(anyhow::Error::msg("No more scripted responses"))
|
||
|
|
} else {
|
||
|
|
Ok(responses.remove(0))
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
fn make_chunk() -> Chunk {
|
||
|
|
Chunk::new(
|
||
|
|
1,
|
||
|
|
vec![Record {
|
||
|
|
role: Role::User,
|
||
|
|
text: "test".to_string(),
|
||
|
|
timestamp: OffsetDateTime::now_utc(),
|
||
|
|
provenance: Provenance {
|
||
|
|
source_id: "test".to_string(),
|
||
|
|
offset: 0,
|
||
|
|
},
|
||
|
|
}],
|
||
|
|
50,
|
||
|
|
)
|
||
|
|
}
|
||
|
|
|
||
|
|
#[test]
|
||
|
|
fn a1_retain_on_no() {
|
||
|
|
let llm = FakeLlm::new(vec![
|
||
|
|
"<think>Not useful</think><check>no</check><update>old</update><next>continue</next>";
|
||
|
|
5
|
||
|
|
]);
|
||
|
|
|
||
|
|
let config = LoopConfig {
|
||
|
|
level: Level::L0,
|
||
|
|
query: Query {
|
||
|
|
id: "test".to_string(),
|
||
|
|
question: "Test?".to_string(),
|
||
|
|
exit_gate: false,
|
||
|
|
},
|
||
|
|
memory_budget: 1024,
|
||
|
|
use_exit_gate: false,
|
||
|
|
};
|
||
|
|
|
||
|
|
let outcome = run_loop(config, vec![make_chunk(); 5], &llm).unwrap();
|
||
|
|
assert_eq!(outcome.chunks_seen, 5);
|
||
|
|
assert_eq!(outcome.chunks_used, 0);
|
||
|
|
assert_eq!(outcome.final_memory, "");
|
||
|
|
}
|
||
|
|
|
||
|
|
#[test]
|
||
|
|
fn a2_update_on_yes() {
|
||
|
|
let llm = FakeLlm::new(vec![
|
||
|
|
"<think>No</think><check>no</check><update>old</update><next>continue</next>",
|
||
|
|
"<think>No</think><check>no</check><update>old</update><next>continue</next>",
|
||
|
|
"<think>Yes</think><check>yes</check><update>New memory</update><next>continue</next>",
|
||
|
|
"<think>No</think><check>no</check><update>old</update><next>continue</next>",
|
||
|
|
"<think>No</think><check>no</check><update>old</update><next>continue</next>",
|
||
|
|
]);
|
||
|
|
|
||
|
|
let config = LoopConfig {
|
||
|
|
level: Level::L0,
|
||
|
|
query: Query {
|
||
|
|
id: "test".to_string(),
|
||
|
|
question: "Test?".to_string(),
|
||
|
|
exit_gate: false,
|
||
|
|
},
|
||
|
|
memory_budget: 1024,
|
||
|
|
use_exit_gate: false,
|
||
|
|
};
|
||
|
|
|
||
|
|
let outcome = run_loop(config, vec![make_chunk(); 5], &llm).unwrap();
|
||
|
|
assert_eq!(outcome.chunks_seen, 5);
|
||
|
|
assert_eq!(outcome.chunks_used, 1);
|
||
|
|
assert_eq!(outcome.final_memory, "New memory");
|
||
|
|
}
|
||
|
|
|
||
|
|
#[test]
|
||
|
|
fn a3_exit_gate_off_reads_all() {
|
||
|
|
let llm = FakeLlm::new(vec![
|
||
|
|
"<think>End</think><check>no</check><update>x</update><next>end</next>";
|
||
|
|
10
|
||
|
|
]);
|
||
|
|
|
||
|
|
let config = LoopConfig {
|
||
|
|
level: Level::L0,
|
||
|
|
query: Query {
|
||
|
|
id: "test".to_string(),
|
||
|
|
question: "Test?".to_string(),
|
||
|
|
exit_gate: false,
|
||
|
|
},
|
||
|
|
memory_budget: 1024,
|
||
|
|
use_exit_gate: false,
|
||
|
|
};
|
||
|
|
|
||
|
|
let outcome = run_loop(config, vec![make_chunk(); 10], &llm).unwrap();
|
||
|
|
assert_eq!(outcome.chunks_seen, 10);
|
||
|
|
}
|
||
|
|
|
||
|
|
#[test]
|
||
|
|
fn a4_exit_gate_on_stops() {
|
||
|
|
let llm = FakeLlm::new(vec![
|
||
|
|
"<think>End</think><check>no</check><update>x</update><next>end</next>",
|
||
|
|
"<think>End</think><check>no</check><update>x</update><next>end</next>",
|
||
|
|
]);
|
||
|
|
|
||
|
|
let config = LoopConfig {
|
||
|
|
level: Level::L0,
|
||
|
|
query: Query {
|
||
|
|
id: "test".to_string(),
|
||
|
|
question: "Test?".to_string(),
|
||
|
|
exit_gate: false,
|
||
|
|
},
|
||
|
|
memory_budget: 1024,
|
||
|
|
use_exit_gate: true,
|
||
|
|
};
|
||
|
|
|
||
|
|
let outcome = run_loop(config, vec![make_chunk(); 10], &llm).unwrap();
|
||
|
|
assert!(outcome.chunks_seen < 10);
|
||
|
|
}
|
||
|
|
|
||
|
|
#[test]
|
||
|
|
fn a5_exit_always_recorded() {
|
||
|
|
let llm = FakeLlm::new(vec![
|
||
|
|
"<think>End</think><check>no</check><update>x</update><next>end</next>",
|
||
|
|
"<think>End</think><check>no</check><update>x</update><next>end</next>",
|
||
|
|
"<think>End</think><check>no</check><update>x</update><next>end</next>",
|
||
|
|
]);
|
||
|
|
|
||
|
|
let config = LoopConfig {
|
||
|
|
level: Level::L0,
|
||
|
|
query: Query {
|
||
|
|
id: "test".to_string(),
|
||
|
|
question: "Test?".to_string(),
|
||
|
|
exit_gate: false,
|
||
|
|
},
|
||
|
|
memory_budget: 1024,
|
||
|
|
use_exit_gate: false,
|
||
|
|
};
|
||
|
|
|
||
|
|
let outcome = run_loop(config, vec![make_chunk(); 3], &llm).unwrap();
|
||
|
|
let exit_count = outcome.events.iter().filter(|e| {
|
||
|
|
matches!(e, LoopEvent::Gate { exit: true, .. })
|
||
|
|
}).count();
|
||
|
|
assert!(exit_count > 0);
|
||
|
|
}
|
||
|
|
|
||
|
|
#[test]
|
||
|
|
fn a6_budget_exceeded_retains() {
|
||
|
|
let large_text = "x".repeat(2000);
|
||
|
|
let update_large = format!("<think>Big</think><check>yes</check><update>{}</update><next>continue</next>", large_text);
|
||
|
|
|
||
|
|
let llm = FakeLlm::new(vec![&update_large]);
|
||
|
|
|
||
|
|
let config = LoopConfig {
|
||
|
|
level: Level::L0,
|
||
|
|
query: Query {
|
||
|
|
id: "test".to_string(),
|
||
|
|
question: "Test?".to_string(),
|
||
|
|
exit_gate: false,
|
||
|
|
},
|
||
|
|
memory_budget: 1024,
|
||
|
|
use_exit_gate: false,
|
||
|
|
};
|
||
|
|
|
||
|
|
let outcome = run_loop(config, vec![make_chunk()], &llm).unwrap();
|
||
|
|
assert_eq!(outcome.final_memory, "");
|
||
|
|
assert_eq!(outcome.chunks_used, 0);
|
||
|
|
assert!(outcome.events.iter().any(|e| matches!(e, LoopEvent::BudgetExceeded { .. })));
|
||
|
|
}
|
||
|
|
|
||
|
|
#[test]
|
||
|
|
fn a7_parse_retry() {
|
||
|
|
let llm = FakeLlm::new(vec![
|
||
|
|
"malformed",
|
||
|
|
"also bad",
|
||
|
|
"<think>Good</think><check>yes</check><update>Memory</update><next>continue</next>",
|
||
|
|
]);
|
||
|
|
|
||
|
|
let config = LoopConfig {
|
||
|
|
level: Level::L0,
|
||
|
|
query: Query {
|
||
|
|
id: "test".to_string(),
|
||
|
|
question: "Test?".to_string(),
|
||
|
|
exit_gate: false,
|
||
|
|
},
|
||
|
|
memory_budget: 1024,
|
||
|
|
use_exit_gate: false,
|
||
|
|
};
|
||
|
|
|
||
|
|
let outcome = run_loop(config, vec![make_chunk()], &llm).unwrap();
|
||
|
|
assert_eq!(outcome.final_memory, "Memory");
|
||
|
|
assert_eq!(outcome.chunks_used, 1);
|
||
|
|
}
|
||
|
|
|
||
|
|
#[test]
|
||
|
|
fn a8_parse_failure_continues() {
|
||
|
|
let llm = FakeLlm::new(vec![
|
||
|
|
"bad",
|
||
|
|
"bad",
|
||
|
|
"bad",
|
||
|
|
]);
|
||
|
|
|
||
|
|
let config = LoopConfig {
|
||
|
|
level: Level::L0,
|
||
|
|
query: Query {
|
||
|
|
id: "test".to_string(),
|
||
|
|
question: "Test?".to_string(),
|
||
|
|
exit_gate: false,
|
||
|
|
},
|
||
|
|
memory_budget: 1024,
|
||
|
|
use_exit_gate: false,
|
||
|
|
};
|
||
|
|
|
||
|
|
let outcome = run_loop(config, vec![make_chunk()], &llm).unwrap();
|
||
|
|
assert_eq!(outcome.chunks_used, 0);
|
||
|
|
assert!(outcome.events.iter().any(|e| matches!(e, LoopEvent::ParseFailed { .. })));
|
||
|
|
}
|
||
|
|
|
||
|
|
#[test]
|
||
|
|
fn a9_parents_linked() {
|
||
|
|
let llm = FakeLlm::new(vec![
|
||
|
|
"<think>Y</think><check>yes</check><update>Mem</update><next>continue</next>",
|
||
|
|
]);
|
||
|
|
|
||
|
|
let config = LoopConfig {
|
||
|
|
level: Level::L0,
|
||
|
|
query: Query {
|
||
|
|
id: "test".to_string(),
|
||
|
|
question: "Test?".to_string(),
|
||
|
|
exit_gate: false,
|
||
|
|
},
|
||
|
|
memory_budget: 1024,
|
||
|
|
use_exit_gate: false,
|
||
|
|
};
|
||
|
|
|
||
|
|
let outcome = run_loop(config, vec![make_chunk()], &llm).unwrap();
|
||
|
|
assert!(outcome.events.iter().any(|e| matches!(e, LoopEvent::Evidence { .. })));
|
||
|
|
}
|
||
|
|
|
||
|
|
#[test]
|
||
|
|
fn a10_level_is_parameter() {
|
||
|
|
let llm = FakeLlm::new(vec![
|
||
|
|
"<think>Y</think><check>yes</check><update>Mem</update><next>continue</next>",
|
||
|
|
"<think>Y</think><check>yes</check><update>Mem</update><next>continue</next>",
|
||
|
|
]);
|
||
|
|
|
||
|
|
let chunk = make_chunk();
|
||
|
|
|
||
|
|
let config_l1 = LoopConfig {
|
||
|
|
level: Level::L1,
|
||
|
|
query: Query {
|
||
|
|
id: "test".to_string(),
|
||
|
|
question: "Test?".to_string(),
|
||
|
|
exit_gate: false,
|
||
|
|
},
|
||
|
|
memory_budget: 1024,
|
||
|
|
use_exit_gate: false,
|
||
|
|
};
|
||
|
|
|
||
|
|
let config_l2 = LoopConfig {
|
||
|
|
level: Level::L2,
|
||
|
|
query: Query {
|
||
|
|
id: "test".to_string(),
|
||
|
|
question: "Test?".to_string(),
|
||
|
|
exit_gate: false,
|
||
|
|
},
|
||
|
|
memory_budget: 1024,
|
||
|
|
use_exit_gate: false,
|
||
|
|
};
|
||
|
|
|
||
|
|
let outcome_l1 = run_loop(config_l1, vec![chunk.clone()], &llm).unwrap();
|
||
|
|
// Note: LlmClient consumed, so create new for second run
|
||
|
|
let llm2 = FakeLlm::new(vec![
|
||
|
|
"<think>Y</think><check>yes</check><update>Mem</update><next>continue</next>",
|
||
|
|
]);
|
||
|
|
let outcome_l2 = run_loop(config_l2, vec![chunk.clone()], &llm2).unwrap();
|
||
|
|
|
||
|
|
// Both should have same event count (just different level internally)
|
||
|
|
assert_eq!(outcome_l1.events.len(), outcome_l2.events.len());
|
||
|
|
}
|