feat: initial Sol virtual librarian implementation
Matrix bot with E2EE (matrix-sdk 0.9) that passively archives all messages to OpenSearch and responds to queries via Mistral AI with function calling tools. Core systems: - Archive: bulk OpenSearch indexer with batch/flush, edit/redaction handling, embedding pipeline passthrough - Brain: rule-based engagement evaluator (mentions, DMs, name invocations), LLM-powered spontaneous engagement, per-room conversation context windows, response delay simulation - Tools: search_archive, get_room_context, list_rooms, get_room_members registered as Mistral function calling tools with iterative tool loop - Personality: templated system prompt with Sol's librarian persona 47 unit tests covering config, evaluator, conversation windowing, personality templates, schema serialization, and search query building.
This commit is contained in:
207
src/brain/conversation.rs
Normal file
207
src/brain/conversation.rs
Normal file
@@ -0,0 +1,207 @@
|
||||
use std::collections::{HashMap, VecDeque};
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct ContextMessage {
|
||||
pub sender: String,
|
||||
pub content: String,
|
||||
pub timestamp: i64,
|
||||
}
|
||||
|
||||
struct RoomContext {
|
||||
messages: VecDeque<ContextMessage>,
|
||||
max_size: usize,
|
||||
}
|
||||
|
||||
impl RoomContext {
|
||||
fn new(max_size: usize) -> Self {
|
||||
Self {
|
||||
messages: VecDeque::with_capacity(max_size),
|
||||
max_size,
|
||||
}
|
||||
}
|
||||
|
||||
fn add(&mut self, msg: ContextMessage) {
|
||||
if self.messages.len() >= self.max_size {
|
||||
self.messages.pop_front();
|
||||
}
|
||||
self.messages.push_back(msg);
|
||||
}
|
||||
|
||||
fn get(&self) -> Vec<ContextMessage> {
|
||||
self.messages.iter().cloned().collect()
|
||||
}
|
||||
}
|
||||
|
||||
pub struct ConversationManager {
|
||||
rooms: HashMap<String, RoomContext>,
|
||||
room_window: usize,
|
||||
dm_window: usize,
|
||||
max_rooms: usize,
|
||||
}
|
||||
|
||||
impl ConversationManager {
|
||||
pub fn new(room_window: usize, dm_window: usize) -> Self {
|
||||
Self {
|
||||
rooms: HashMap::new(),
|
||||
room_window,
|
||||
dm_window,
|
||||
max_rooms: 500, // todo(sienna): make this configurable
|
||||
}
|
||||
}
|
||||
|
||||
pub fn add_message(&mut self, room_id: &str, is_dm: bool, msg: ContextMessage) {
|
||||
let window = if is_dm {
|
||||
self.dm_window
|
||||
} else {
|
||||
self.room_window
|
||||
};
|
||||
|
||||
// Evict oldest room if at capacity
|
||||
if !self.rooms.contains_key(room_id) && self.rooms.len() >= self.max_rooms {
|
||||
// Remove the room with the oldest latest message
|
||||
let oldest = self
|
||||
.rooms
|
||||
.iter()
|
||||
.min_by_key(|(_, ctx)| ctx.messages.back().map(|m| m.timestamp).unwrap_or(0))
|
||||
.map(|(k, _)| k.to_owned());
|
||||
if let Some(key) = oldest {
|
||||
self.rooms.remove(&key);
|
||||
}
|
||||
}
|
||||
|
||||
let ctx = self
|
||||
.rooms
|
||||
.entry(room_id.to_owned())
|
||||
.or_insert_with(|| RoomContext::new(window));
|
||||
ctx.add(msg);
|
||||
}
|
||||
|
||||
pub fn get_context(&self, room_id: &str) -> Vec<ContextMessage> {
|
||||
self.rooms
|
||||
.get(room_id)
|
||||
.map(|ctx| ctx.get())
|
||||
.unwrap_or_default()
|
||||
}
|
||||
|
||||
pub fn room_count(&self) -> usize {
|
||||
self.rooms.len()
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
fn msg(sender: &str, content: &str, ts: i64) -> ContextMessage {
|
||||
ContextMessage {
|
||||
sender: sender.to_string(),
|
||||
content: content.to_string(),
|
||||
timestamp: ts,
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_add_and_get_messages() {
|
||||
let mut cm = ConversationManager::new(5, 10);
|
||||
cm.add_message("!room1:x", false, msg("alice", "hello", 1));
|
||||
cm.add_message("!room1:x", false, msg("bob", "hi", 2));
|
||||
|
||||
let ctx = cm.get_context("!room1:x");
|
||||
assert_eq!(ctx.len(), 2);
|
||||
assert_eq!(ctx[0].sender, "alice");
|
||||
assert_eq!(ctx[1].sender, "bob");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_empty_room_returns_empty() {
|
||||
let cm = ConversationManager::new(5, 10);
|
||||
let ctx = cm.get_context("!nonexistent:x");
|
||||
assert!(ctx.is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_sliding_window_group_room() {
|
||||
let mut cm = ConversationManager::new(3, 10);
|
||||
for i in 0..5 {
|
||||
cm.add_message("!room:x", false, msg("user", &format!("msg{i}"), i));
|
||||
}
|
||||
let ctx = cm.get_context("!room:x");
|
||||
assert_eq!(ctx.len(), 3);
|
||||
// Should keep the last 3
|
||||
assert_eq!(ctx[0].content, "msg2");
|
||||
assert_eq!(ctx[1].content, "msg3");
|
||||
assert_eq!(ctx[2].content, "msg4");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_sliding_window_dm_room() {
|
||||
let mut cm = ConversationManager::new(3, 5);
|
||||
for i in 0..7 {
|
||||
cm.add_message("!dm:x", true, msg("user", &format!("dm{i}"), i));
|
||||
}
|
||||
let ctx = cm.get_context("!dm:x");
|
||||
assert_eq!(ctx.len(), 5);
|
||||
assert_eq!(ctx[0].content, "dm2");
|
||||
assert_eq!(ctx[4].content, "dm6");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_multiple_rooms_independent() {
|
||||
let mut cm = ConversationManager::new(5, 10);
|
||||
cm.add_message("!a:x", false, msg("alice", "in room a", 1));
|
||||
cm.add_message("!b:x", false, msg("bob", "in room b", 2));
|
||||
|
||||
assert_eq!(cm.get_context("!a:x").len(), 1);
|
||||
assert_eq!(cm.get_context("!b:x").len(), 1);
|
||||
assert_eq!(cm.get_context("!a:x")[0].content, "in room a");
|
||||
assert_eq!(cm.get_context("!b:x")[0].content, "in room b");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_lru_eviction_at_max_rooms() {
|
||||
// Create a manager with max_rooms = 500 (default), but we'll use a small one
|
||||
let mut cm = ConversationManager::new(5, 10);
|
||||
cm.max_rooms = 3;
|
||||
|
||||
// Add 3 rooms
|
||||
cm.add_message("!room1:x", false, msg("a", "r1", 100));
|
||||
cm.add_message("!room2:x", false, msg("b", "r2", 200));
|
||||
cm.add_message("!room3:x", false, msg("c", "r3", 300));
|
||||
assert_eq!(cm.room_count(), 3);
|
||||
|
||||
// Adding a 4th room should evict the one with oldest latest message (room1, ts=100)
|
||||
cm.add_message("!room4:x", false, msg("d", "r4", 400));
|
||||
assert_eq!(cm.room_count(), 3);
|
||||
assert!(cm.get_context("!room1:x").is_empty()); // evicted
|
||||
assert_eq!(cm.get_context("!room2:x").len(), 1);
|
||||
assert_eq!(cm.get_context("!room3:x").len(), 1);
|
||||
assert_eq!(cm.get_context("!room4:x").len(), 1);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_existing_room_not_evicted() {
|
||||
let mut cm = ConversationManager::new(5, 10);
|
||||
cm.max_rooms = 2;
|
||||
|
||||
cm.add_message("!room1:x", false, msg("a", "r1", 100));
|
||||
cm.add_message("!room2:x", false, msg("b", "r2", 200));
|
||||
|
||||
// Adding to existing room should NOT trigger eviction
|
||||
cm.add_message("!room1:x", false, msg("a", "r1 again", 300));
|
||||
assert_eq!(cm.room_count(), 2);
|
||||
assert_eq!(cm.get_context("!room1:x").len(), 2);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_message_ordering_preserved() {
|
||||
let mut cm = ConversationManager::new(10, 10);
|
||||
cm.add_message("!r:x", false, msg("a", "first", 1));
|
||||
cm.add_message("!r:x", false, msg("b", "second", 2));
|
||||
cm.add_message("!r:x", false, msg("c", "third", 3));
|
||||
|
||||
let ctx = cm.get_context("!r:x");
|
||||
assert_eq!(ctx[0].timestamp, 1);
|
||||
assert_eq!(ctx[1].timestamp, 2);
|
||||
assert_eq!(ctx[2].timestamp, 3);
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user