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:
160
src/main.rs
Normal file
160
src/main.rs
Normal file
@@ -0,0 +1,160 @@
|
||||
mod archive;
|
||||
mod brain;
|
||||
mod config;
|
||||
mod matrix_utils;
|
||||
mod sync;
|
||||
mod tools;
|
||||
|
||||
use std::sync::Arc;
|
||||
|
||||
use matrix_sdk::Client;
|
||||
use opensearch::http::transport::TransportBuilder;
|
||||
use opensearch::OpenSearch;
|
||||
use ruma::{OwnedDeviceId, OwnedUserId};
|
||||
use tokio::signal;
|
||||
use tokio::sync::Mutex;
|
||||
use tracing::{error, info};
|
||||
use url::Url;
|
||||
|
||||
use archive::indexer::Indexer;
|
||||
use archive::schema::create_index_if_not_exists;
|
||||
use brain::conversation::ConversationManager;
|
||||
use brain::evaluator::Evaluator;
|
||||
use brain::personality::Personality;
|
||||
use brain::responder::Responder;
|
||||
use config::Config;
|
||||
use sync::AppState;
|
||||
use tools::ToolRegistry;
|
||||
|
||||
#[tokio::main]
|
||||
async fn main() -> anyhow::Result<()> {
|
||||
// Initialize tracing
|
||||
tracing_subscriber::fmt()
|
||||
.with_env_filter(
|
||||
tracing_subscriber::EnvFilter::try_from_default_env()
|
||||
.unwrap_or_else(|_| tracing_subscriber::EnvFilter::new("sol=info")),
|
||||
)
|
||||
.init();
|
||||
|
||||
// Load config
|
||||
let config_path =
|
||||
std::env::var("SOL_CONFIG").unwrap_or_else(|_| "/etc/sol/sol.toml".into());
|
||||
let config = Config::load(&config_path)?;
|
||||
info!("Loaded config from {config_path}");
|
||||
|
||||
// Load system prompt
|
||||
let prompt_path = std::env::var("SOL_SYSTEM_PROMPT")
|
||||
.unwrap_or_else(|_| "/etc/sol/system_prompt.md".into());
|
||||
let system_prompt = std::fs::read_to_string(&prompt_path)?;
|
||||
info!("Loaded system prompt from {prompt_path}");
|
||||
|
||||
// Read secrets from environment
|
||||
let access_token = std::env::var("SOL_MATRIX_ACCESS_TOKEN")
|
||||
.map_err(|_| anyhow::anyhow!("SOL_MATRIX_ACCESS_TOKEN not set"))?;
|
||||
let device_id = std::env::var("SOL_MATRIX_DEVICE_ID")
|
||||
.map_err(|_| anyhow::anyhow!("SOL_MATRIX_DEVICE_ID not set"))?;
|
||||
let mistral_api_key = std::env::var("SOL_MISTRAL_API_KEY")
|
||||
.map_err(|_| anyhow::anyhow!("SOL_MISTRAL_API_KEY not set"))?;
|
||||
|
||||
let config = Arc::new(config);
|
||||
|
||||
// Initialize Matrix client with E2EE and sqlite store
|
||||
let homeserver = Url::parse(&config.matrix.homeserver_url)?;
|
||||
|
||||
let matrix_client = Client::builder()
|
||||
.homeserver_url(homeserver)
|
||||
.sqlite_store(&config.matrix.state_store_path, None)
|
||||
.build()
|
||||
.await?;
|
||||
|
||||
// Restore session
|
||||
let user_id: OwnedUserId = config.matrix.user_id.parse()?;
|
||||
let device_id: OwnedDeviceId = device_id.into();
|
||||
|
||||
let session = matrix_sdk::AuthSession::Matrix(matrix_sdk::matrix_auth::MatrixSession {
|
||||
meta: matrix_sdk::SessionMeta {
|
||||
user_id,
|
||||
device_id,
|
||||
},
|
||||
tokens: matrix_sdk::matrix_auth::MatrixSessionTokens {
|
||||
access_token,
|
||||
refresh_token: None,
|
||||
},
|
||||
});
|
||||
|
||||
matrix_client.restore_session(session).await?;
|
||||
info!(user = %config.matrix.user_id, "Matrix session restored");
|
||||
|
||||
// Initialize OpenSearch client
|
||||
let os_url = Url::parse(&config.opensearch.url)?;
|
||||
let os_transport = TransportBuilder::new(
|
||||
opensearch::http::transport::SingleNodeConnectionPool::new(os_url),
|
||||
)
|
||||
.build()?;
|
||||
let os_client = OpenSearch::new(os_transport);
|
||||
|
||||
// Ensure index exists
|
||||
create_index_if_not_exists(&os_client, &config.opensearch.index).await?;
|
||||
|
||||
// Initialize Mistral client
|
||||
let mistral_client = mistralai_client::v1::client::Client::new(
|
||||
Some(mistral_api_key),
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
)?;
|
||||
let mistral = Arc::new(mistral_client);
|
||||
|
||||
// Build components
|
||||
let personality = Arc::new(Personality::new(system_prompt));
|
||||
let tool_registry = Arc::new(ToolRegistry::new(
|
||||
os_client.clone(),
|
||||
matrix_client.clone(),
|
||||
config.clone(),
|
||||
));
|
||||
let indexer = Arc::new(Indexer::new(os_client, config.clone()));
|
||||
let evaluator = Arc::new(Evaluator::new(config.clone()));
|
||||
let responder = Arc::new(Responder::new(
|
||||
config.clone(),
|
||||
personality,
|
||||
tool_registry,
|
||||
));
|
||||
let conversations = Arc::new(Mutex::new(ConversationManager::new(
|
||||
config.behavior.room_context_window,
|
||||
config.behavior.dm_context_window,
|
||||
)));
|
||||
|
||||
// Start background flush task
|
||||
let _flush_handle = indexer.start_flush_task();
|
||||
|
||||
// Build shared state
|
||||
let state = Arc::new(AppState {
|
||||
config: config.clone(),
|
||||
indexer,
|
||||
evaluator,
|
||||
responder,
|
||||
conversations,
|
||||
mistral,
|
||||
});
|
||||
|
||||
// Start sync loop in background
|
||||
let sync_client = matrix_client.clone();
|
||||
let sync_state = state.clone();
|
||||
let sync_handle = tokio::spawn(async move {
|
||||
if let Err(e) = sync::start_sync(sync_client, sync_state).await {
|
||||
error!("Sync loop error: {e}");
|
||||
}
|
||||
});
|
||||
|
||||
info!("Sol is running");
|
||||
|
||||
// Wait for shutdown signal
|
||||
signal::ctrl_c().await?;
|
||||
info!("Shutdown signal received");
|
||||
|
||||
// Cancel sync
|
||||
sync_handle.abort();
|
||||
|
||||
info!("Sol has shut down");
|
||||
Ok(())
|
||||
}
|
||||
Reference in New Issue
Block a user