sol/src/main.rs

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(())
}