310 lines
11 KiB
Rust
310 lines
11 KiB
Rust
//! A turn's speaker and its number (M5c spec, sections 5 and 11). `loopd` logs the speaker on the
|
|
//! turn's `User` record as it came, and never puts it in a request: the model sees the same bytes
|
|
//! with a speaker and without one. After a finished turn it tells `memoryd` the turn's number: the
|
|
//! count of `User` records in the session's log, so the first turn is 1, a failed turn that was
|
|
//! logged counts too, and the count goes on after a reopen or a restart. Do not edit.
|
|
|
|
mod support;
|
|
#[path = "support/tags.rs"]
|
|
mod tags;
|
|
|
|
use std::os::unix::net::{UnixListener, UnixStream};
|
|
use std::path::{Path, PathBuf};
|
|
use std::sync::mpsc;
|
|
use std::sync::{Arc, Mutex};
|
|
use std::thread;
|
|
use std::time::Duration;
|
|
|
|
use bxcore::{
|
|
Empty, Envelope, LogRecord, Message, PROTOCOL_VERSION, SessionId, SessionUpdated, Timestamp,
|
|
Turn, read_frame, write_frame,
|
|
};
|
|
use loopd::baseline::Baseline;
|
|
use loopd::channel::{Context, serve};
|
|
use loopd::llama::Client;
|
|
use loopd::memory::{MemoryPort, MemorySource};
|
|
use loopd::session::Session;
|
|
use loopd::tools::Registry;
|
|
use support::{FakeServer, Home, Reply, ScriptedPort};
|
|
use tags::{holds_a_tag, same};
|
|
|
|
const CHAT: &str = "/v1/chat/completions";
|
|
|
|
/// A `memoryd` stand-in that records each `(session, turn)` it is told about.
|
|
#[derive(Default)]
|
|
struct Recorder(Mutex<Vec<(String, u64)>>);
|
|
|
|
impl Recorder {
|
|
fn told(&self) -> Vec<(String, u64)> {
|
|
self.0.lock().unwrap().clone()
|
|
}
|
|
}
|
|
|
|
/// Lets the test keep its `Arc<Recorder>` while the context owns a `Box<dyn MemorySource>`.
|
|
struct Shared(Arc<Recorder>);
|
|
|
|
impl MemorySource for Shared {
|
|
fn block(&self) -> Result<String, String> {
|
|
Ok(String::new())
|
|
}
|
|
fn updated(&self, session: &SessionId, turn: u64) -> Result<(), String> {
|
|
let entry = (session.as_str().to_string(), turn);
|
|
self.0.0.lock().unwrap().push(entry);
|
|
Ok(())
|
|
}
|
|
}
|
|
|
|
/// A fake inference server and a home that one `loopd` or several, one after the other, serve
|
|
/// from.
|
|
struct Setup {
|
|
server: FakeServer,
|
|
// Declared last so it is dropped last: fields drop in order, and the home must outlive
|
|
// everything that may still write into it (tip T44).
|
|
home: Home,
|
|
}
|
|
|
|
impl Setup {
|
|
/// The server answers the completions in `chat` in order; the last one repeats.
|
|
fn new(chat: Vec<Reply>) -> Setup {
|
|
let home = Home::new();
|
|
let server = FakeServer::start();
|
|
server.route(CHAT, chat);
|
|
Setup { server, home }
|
|
}
|
|
|
|
/// Starts a channel server on the home, as a new `loopd` process would: its own context,
|
|
/// nothing carried over but the files. `name` names its socket, under the home; the returned
|
|
/// recorder hears what this `loopd` tells `memoryd`.
|
|
fn loopd(&self, name: &str) -> (PathBuf, Arc<Recorder>) {
|
|
let mut cfg = self.home.config(&self.server.socket);
|
|
cfg.channel.socket = self.home.dir.join(name);
|
|
let socket = cfg.channel_socket();
|
|
let listener = UnixListener::bind(&socket).unwrap();
|
|
let memory = Arc::new(Recorder::default());
|
|
let ctx = Arc::new(Context::new(
|
|
cfg.clone(),
|
|
Client::new(cfg),
|
|
Box::new(ScriptedPort::new(vec![])),
|
|
Registry::m2b(),
|
|
Box::new(Shared(Arc::clone(&memory))),
|
|
));
|
|
thread::spawn(move || serve(listener, ctx));
|
|
(socket, memory)
|
|
}
|
|
|
|
/// The `User` records of a session's log, in order.
|
|
fn users(&self, session: &str) -> Vec<LogRecord> {
|
|
let records = self.home.records(session);
|
|
let is_user = |r: &LogRecord| matches!(r, LogRecord::User { .. });
|
|
records.into_iter().filter(is_user).collect()
|
|
}
|
|
|
|
/// The bodies of the completion requests so far, as sent, in order.
|
|
fn bodies(&self) -> Vec<String> {
|
|
let requests = self.server.requests_to(CHAT);
|
|
let body = |r: &support::Recorded| String::from_utf8(r.body.clone()).unwrap();
|
|
requests.iter().map(body).collect()
|
|
}
|
|
}
|
|
|
|
/// Sends one turn and returns true if it ended with `turn_done`, false if with an error.
|
|
fn turn(socket: &Path, session: &str, resume: bool, speaker: Option<&str>) -> bool {
|
|
let mut stream = UnixStream::connect(socket).unwrap();
|
|
stream
|
|
.set_read_timeout(Some(Duration::from_secs(10)))
|
|
.unwrap();
|
|
let msg = Message::Turn(Turn {
|
|
session: SessionId::new(session).unwrap(),
|
|
content: "hello".to_string(),
|
|
resume,
|
|
speaker: speaker.map(str::to_string),
|
|
});
|
|
let env = Envelope {
|
|
v: PROTOCOL_VERSION,
|
|
id: 7,
|
|
r#final: true,
|
|
msg,
|
|
};
|
|
write_frame(&mut stream, &env).unwrap();
|
|
loop {
|
|
let frame = read_frame(&mut stream).unwrap();
|
|
if !frame.r#final {
|
|
continue;
|
|
}
|
|
match frame.msg {
|
|
Message::TurnDone(_) => return true,
|
|
Message::Error(_) => return false,
|
|
other => panic!("the turn ended with neither turn_done nor error: {other:?}"),
|
|
}
|
|
}
|
|
}
|
|
|
|
fn told(session: &str, turn: u64) -> (String, u64) {
|
|
(session.to_string(), turn)
|
|
}
|
|
|
|
/// The speaker of a `User` record, or a panic for any other record.
|
|
fn speaker(record: &LogRecord) -> Option<String> {
|
|
match record {
|
|
LogRecord::User { speaker, .. } => speaker.clone(),
|
|
other => panic!("not a user record: {other:?}"),
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
fn the_speaker_is_logged_on_the_user_record_as_it_came() {
|
|
let s = Setup::new(vec![Reply::fixture("plain")]);
|
|
let (socket, _) = s.loopd("loop.sock");
|
|
// Not a valid Honcho peer name, and one that holds one of the model's tags: `loopd` neither
|
|
// checks nor neutralises a speaker, since it never reaches the model.
|
|
let tagged = format!("x{}y", bxcore::sanitize::TAGS[0].text);
|
|
assert!(turn(&socket, "a", false, Some("first.last")));
|
|
assert!(turn(&socket, "a", true, None));
|
|
assert!(turn(&socket, "a", true, Some(&tagged)));
|
|
let users = s.users("a");
|
|
assert_eq!(users.len(), 3, "one user record per turn");
|
|
assert_eq!(speaker(&users[0]), Some("first.last".to_string()));
|
|
assert_eq!(speaker(&users[1]), None, "a turn without a speaker");
|
|
same(&speaker(&users[2]), &Some(tagged), "the third speaker");
|
|
let log = s.home.read("sessions/a/0.jsonl");
|
|
assert_eq!(
|
|
log.matches("\"speaker\":").count(),
|
|
2,
|
|
"no speaker key for the turn without one"
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn the_requests_are_the_same_with_and_without_a_speaker() {
|
|
let s = Setup::new(vec![Reply::fixture("plain")]);
|
|
let (socket, _) = s.loopd("loop.sock");
|
|
let tagged = format!("x{}y", bxcore::sanitize::TAGS[0].text);
|
|
assert!(turn(&socket, "with", false, Some("Zelda_Q")));
|
|
assert!(turn(&socket, "with", true, Some(&tagged)));
|
|
assert!(turn(&socket, "without", false, None));
|
|
assert!(turn(&socket, "without", true, None));
|
|
let bodies = s.bodies();
|
|
assert_eq!(bodies.len(), 4, "one request per turn");
|
|
// Each session's second request replays its first turn from the log: a speaker there would
|
|
// change the replayed bytes too.
|
|
assert!(bodies[0] == bodies[2], "the first requests differ");
|
|
assert!(bodies[1] == bodies[3], "the second requests differ");
|
|
for (n, body) in bodies.iter().enumerate() {
|
|
assert!(!body.contains("Zelda_Q"), "request {n} holds the speaker");
|
|
assert!(!holds_a_tag(body), "request {n} holds the tagged speaker");
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
fn the_turn_number_starts_at_one_in_each_session() {
|
|
let s = Setup::new(vec![Reply::fixture("plain")]);
|
|
let (socket, memory) = s.loopd("loop.sock");
|
|
assert!(turn(&socket, "a", false, None));
|
|
assert!(turn(&socket, "a", true, Some("kyle")));
|
|
assert!(turn(&socket, "b", false, None));
|
|
assert!(turn(&socket, "a", true, None));
|
|
assert_eq!(
|
|
memory.told(),
|
|
[told("a", 1), told("a", 2), told("b", 1), told("a", 3)]
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn a_turn_with_tool_calls_is_one_turn() {
|
|
// The first turn asks for a tool, then answers: two assistant records, one user record.
|
|
let s = Setup::new(vec![Reply::fixture("tool_call"), Reply::fixture("plain")]);
|
|
let (socket, memory) = s.loopd("loop.sock");
|
|
assert!(turn(&socket, "a", false, None));
|
|
assert!(turn(&socket, "a", true, None));
|
|
assert_eq!(memory.told(), [told("a", 1), told("a", 2)]);
|
|
}
|
|
|
|
#[test]
|
|
fn a_failed_turn_is_not_told_but_is_counted() {
|
|
let s = Setup::new(vec![
|
|
Reply::fixture("plain"),
|
|
Reply::fixture("bad_request"),
|
|
Reply::fixture("plain"),
|
|
]);
|
|
let (socket, memory) = s.loopd("loop.sock");
|
|
assert!(turn(&socket, "a", false, None));
|
|
assert!(!turn(&socket, "a", true, None), "the second turn fails");
|
|
assert!(!turn(&socket, "nobody", true, None), "no such session");
|
|
assert!(turn(&socket, "a", true, None));
|
|
assert_eq!(s.users("a").len(), 3, "the failed turn's message is logged");
|
|
assert_eq!(memory.told(), [told("a", 1), told("a", 3)]);
|
|
}
|
|
|
|
#[test]
|
|
fn a_restarted_loopd_goes_on_counting_from_the_log() {
|
|
let s = Setup::new(vec![Reply::fixture("plain")]);
|
|
let (first, before) = s.loopd("first.sock");
|
|
assert!(turn(&first, "a", false, None));
|
|
assert!(turn(&first, "a", true, None));
|
|
let (second, after) = s.loopd("second.sock");
|
|
assert!(turn(&second, "a", true, None));
|
|
assert!(turn(&second, "b", false, None));
|
|
assert_eq!(before.told(), [told("a", 1), told("a", 2)]);
|
|
assert_eq!(after.told(), [told("a", 3), told("b", 1)]);
|
|
}
|
|
|
|
#[test]
|
|
fn session_turns_counts_the_user_records_of_the_log() {
|
|
let home = Home::new();
|
|
let cfg = home.config(Path::new("/tmp/unused.sock"));
|
|
let baseline = Baseline::assemble(&cfg, &Registry::m2b()).unwrap();
|
|
let id = SessionId::new("a").unwrap();
|
|
let mut session = Session::create(&home.dir, id.clone(), baseline, 0).unwrap();
|
|
assert_eq!(session.turns(), 0, "a new session has had no turn");
|
|
let user = |speaker: Option<&str>| LogRecord::User {
|
|
time: Timestamp::now(),
|
|
content: "hello".to_string(),
|
|
speaker: speaker.map(str::to_string),
|
|
};
|
|
let assistant = || LogRecord::Assistant {
|
|
time: Timestamp::now(),
|
|
content: Some("hi".to_string()),
|
|
reasoning_content: None,
|
|
tool_calls: vec![],
|
|
};
|
|
session.append(user(None)).unwrap();
|
|
session.append(assistant()).unwrap();
|
|
assert_eq!(session.turns(), 1);
|
|
session.append(user(Some("kyle"))).unwrap();
|
|
session.append(assistant()).unwrap();
|
|
session.append(assistant()).unwrap();
|
|
assert_eq!(session.turns(), 2);
|
|
drop(session);
|
|
let session = Session::open(&home.dir, id).unwrap();
|
|
assert_eq!(session.turns(), 2, "counted again from the log");
|
|
}
|
|
|
|
#[test]
|
|
fn the_memory_port_sends_the_turn_number() {
|
|
let home = Home::new();
|
|
let socket = home.dir.join("memory.sock");
|
|
let listener = UnixListener::bind(&socket).unwrap();
|
|
let (tx, rx) = mpsc::channel();
|
|
thread::spawn(move || {
|
|
let (mut stream, _) = listener.accept().unwrap();
|
|
let env = read_frame(&mut stream).unwrap();
|
|
let reply = Envelope {
|
|
v: PROTOCOL_VERSION,
|
|
id: env.id,
|
|
r#final: true,
|
|
msg: Message::Ok(Empty {}),
|
|
};
|
|
write_frame(&mut stream, &reply).unwrap();
|
|
let _ = tx.send(env.msg);
|
|
});
|
|
let port = MemoryPort::new(socket, Duration::from_secs(5));
|
|
let id = SessionId::new("a").unwrap();
|
|
assert_eq!(port.updated(&id, 7), Ok(()));
|
|
let sent = rx.recv_timeout(Duration::from_secs(10)).unwrap();
|
|
let want = Message::SessionUpdated(SessionUpdated {
|
|
session: id,
|
|
turn: 7,
|
|
});
|
|
assert_eq!(sent, want);
|
|
}
|