headless_lms_server/programs/

chatbot_syncer.rs

use std::{
    collections::{HashMap, HashSet},
    env,
    time::Duration,
};

use chrono::Utc;
use dotenv::dotenv;
use sqlx::{PgConnection, PgPool};
use url::Url;
use uuid::Uuid;

use crate::setup_tracing;

use headless_lms_chatbot::{
    azure_blob_storage::AzureBlobClient,
    azure_datasources::{create_azure_datasource, does_azure_datasource_exist},
    azure_search_index::{create_search_index, does_search_index_exist},
    azure_search_indexer::{
        check_search_indexer_status, create_search_indexer, does_search_indexer_exist,
        run_search_indexer_now,
    },
    azure_skillset::{create_skillset, does_skillset_exist},
    content_cleaner::convert_material_blocks_to_markdown_with_llm,
};
use headless_lms_models::{
    application_task_default_language_models::ApplicationTask,
    chapters::DatabaseChapter,
    page_history::PageHistory,
    pages::{Page, PageVisibility},
};
use headless_lms_utils::{
    ApplicationConfiguration,
    document_schema_processor::{GutenbergBlock, remove_sensitive_attributes},
    url_encoding::url_encode,
};

const SYNC_INTERVAL_SECS: u64 = 10;
const PRINT_STILL_RUNNING_MESSAGE_TICKS_THRESHOLD: u32 = 60;
const FAILURE_COOLDOWN_SECS: i64 = 300;
const MAX_CONSECUTIVE_FAILURES: i32 = 5;

pub async fn main() -> anyhow::Result<()> {
    initialize_environment()?;
    let config = initialize_configuration().await?;
    if config.app_configuration.azure_configuration.is_none() {
        warn!("Azure configuration not provided. Not running chatbot syncer.");
        // Sleep indefinitely to prevent the program from exiting. This only happens in development.
        loop {
            tokio::time::sleep(Duration::from_secs(u64::MAX)).await;
        }
    }
    if config.app_configuration.test_chatbot {
        warn!(
            "Using mock azure configuration, this must be a test/dev environment. Not running chatbot syncer."
        );
        // Sleep indefinitely to prevent the program from exiting. This only happens in development.
        loop {
            tokio::time::sleep(Duration::from_secs(u64::MAX)).await;
        }
    }

    let db_pool = initialize_database_pool(&config.database_url).await?;
    let mut conn = db_pool.acquire().await?;
    let blob_client = initialize_blob_client(&config).await?;

    let mut interval = tokio::time::interval(Duration::from_secs(SYNC_INTERVAL_SECS));
    let mut ticks = 0;

    info!("Starting chatbot syncer.");

    loop {
        interval.tick().await;
        ticks += 1;

        if ticks >= PRINT_STILL_RUNNING_MESSAGE_TICKS_THRESHOLD {
            ticks = 0;
            info!("Still syncing for chatbot.");
        }
        if let Err(e) = sync_pages(&mut conn, &config, &blob_client).await {
            error!("Error during synchronization: {:?}", e);
        }
    }
}

fn initialize_environment() -> anyhow::Result<()> {
    // TODO: Audit that the environment access only happens in single-threaded code.
    unsafe { env::set_var("RUST_LOG", "info,actix_web=info,sqlx=warn") };
    dotenv().ok();
    setup_tracing()?;
    Ok(())
}

struct SyncerConfig {
    database_url: String,
    name: String,
    app_configuration: ApplicationConfiguration,
}

async fn initialize_configuration() -> anyhow::Result<SyncerConfig> {
    let database_url = env::var("DATABASE_URL")
        .unwrap_or_else(|_| "postgres://localhost/headless_lms_dev".to_string());

    let base_url = Url::parse(&env::var("BASE_URL").expect("BASE_URL must be defined"))
        .expect("BASE_URL must be a valid URL");

    let name = base_url
        .host_str()
        .expect("BASE_URL must have a host")
        .replace(".", "-");

    let app_configuration = ApplicationConfiguration::try_from_env()?;

    Ok(SyncerConfig {
        database_url,
        name,
        app_configuration,
    })
}

/// Initializes the PostgreSQL connection pool.
async fn initialize_database_pool(database_url: &str) -> anyhow::Result<PgPool> {
    PgPool::connect(database_url).await.map_err(|e| {
        anyhow::anyhow!(
            "Failed to connect to the database at {}: {:?}",
            database_url,
            e
        )
    })
}

/// Initializes the Azure Blob Storage client.
async fn initialize_blob_client(config: &SyncerConfig) -> anyhow::Result<AzureBlobClient> {
    let blob_client = AzureBlobClient::new(&config.app_configuration, &config.name).await?;
    blob_client.ensure_container_exists().await?;
    Ok(blob_client)
}

/// Synchronizes pages to the chatbot backend.
async fn sync_pages(
    conn: &mut PgConnection,
    config: &SyncerConfig,
    blob_client: &AzureBlobClient,
) -> anyhow::Result<()> {
    let base_url = Url::parse(&config.app_configuration.base_url)?;
    let chatbot_configs =
        headless_lms_models::chatbot_configurations::get_for_azure_search_maintenance(conn).await?;

    let course_ids: Vec<Uuid> = chatbot_configs
        .iter()
        .map(|config| config.course_id)
        .collect::<HashSet<_>>()
        .into_iter()
        .collect();

    let sync_statuses =
        headless_lms_models::chatbot_page_sync_statuses::ensure_sync_statuses_exist(
            conn,
            &course_ids,
        )
        .await?;

    let latest_histories =
        headless_lms_models::page_history::get_latest_history_entries_for_pages_by_course_ids(
            conn,
            &course_ids,
        )
        .await?;

    let shared_index_name = config.name.clone();
    ensure_search_index_exists(
        &shared_index_name,
        &config.app_configuration,
        &blob_client.container_name,
    )
    .await?;

    if !check_search_indexer_status(&shared_index_name, &config.app_configuration).await? {
        warn!("Search indexer is not ready to index. Skipping synchronization.");
        return Ok(());
    }

    let mut any_changes = false;

    for (course_id, statuses) in sync_statuses.iter() {
        let page_ids: Vec<Uuid> = statuses.iter().map(|s| s.page_id).collect();
        let public_pages_set: HashSet<Uuid> =
            headless_lms_models::pages::get_by_ids_and_visibility(
                conn,
                &page_ids,
                PageVisibility::Public,
            )
            .await?
            .into_iter()
            .map(|p| p.id)
            .collect();

        let outdated_statuses: Vec<_> = statuses
            .iter()
            .filter(|status| {
                if !public_pages_set.contains(&status.page_id) {
                    return false;
                }

                let is_outdated = latest_histories
                    .get(&status.page_id)
                    .is_some_and(|history| status.synced_page_revision_id != Some(history.id));

                if !is_outdated {
                    return false;
                }

                if status.consecutive_failures >= MAX_CONSECUTIVE_FAILURES {
                    debug!(
                        "Skipping page {} due to permanent failure ({} consecutive failures). Manual intervention required.",
                        status.page_id, status.consecutive_failures
                    );
                    return false;
                }

                if let Some(error_msg) = &status.error_message
                    && !error_msg.is_empty() {
                        let error_age_seconds = (Utc::now() - status.updated_at).num_seconds();
                        if error_age_seconds < FAILURE_COOLDOWN_SECS {
                            debug!(
                                "Skipping page {} due to recent failure ({} seconds ago, {} consecutive failures): {}",
                                status.page_id, error_age_seconds, status.consecutive_failures, error_msg
                            );
                            return false;
                        }
                    }

                true
            })
            .collect();

        if outdated_statuses.is_empty() {
            continue;
        }

        any_changes = true;
        info!(
            "Syncing {} pages for course id: {}.",
            outdated_statuses.len(),
            course_id
        );
        for status in &outdated_statuses {
            info!(
                "Page id: {}, synced page revision id: {:?}.",
                status.page_id, status.synced_page_revision_id
            );
        }

        let page_ids: Vec<Uuid> = outdated_statuses.iter().map(|s| s.page_id).collect();
        let pages = headless_lms_models::pages::get_by_ids_and_visibility(
            conn,
            &page_ids,
            PageVisibility::Public,
        )
        .await?;

        if !pages.is_empty() {
            sync_pages_batch(
                conn,
                &pages,
                blob_client,
                &base_url,
                &config.app_configuration,
                &latest_histories,
            )
            .await?;
        } else {
            info!("No pages to sync for course id: {}.", course_id);
        }

        let hidden_page_ids: Vec<Uuid> = statuses
            .iter()
            .filter(|status| {
                !public_pages_set.contains(&status.page_id)
                    && status.synced_page_revision_id.is_some()
            })
            .map(|s| s.page_id)
            .collect();

        if !hidden_page_ids.is_empty() {
            info!(
                "Clearing sync statuses for {} hidden pages: {:?}",
                hidden_page_ids.len(),
                hidden_page_ids
            );
            headless_lms_models::chatbot_page_sync_statuses::clear_sync_statuses(
                conn,
                &hidden_page_ids,
            )
            .await?;
        }

        delete_old_files(conn, *course_id, blob_client).await?;
    }

    if any_changes {
        run_search_indexer_now(&shared_index_name, &config.app_configuration).await?;
        info!("New files have been synced and the search indexer has been started.");
    }

    Ok(())
}

/// Ensures that the specified search index exists, creating it if necessary.
async fn ensure_search_index_exists(
    name: &str,
    app_config: &ApplicationConfiguration,
    container_name: &str,
) -> anyhow::Result<()> {
    if !does_search_index_exist(name, app_config).await? {
        create_search_index(name.to_owned(), app_config).await?;
    }
    if !does_skillset_exist(name, app_config).await? {
        create_skillset(name, name, app_config).await?;
    }
    if !does_azure_datasource_exist(name, app_config).await? {
        create_azure_datasource(name, container_name, app_config).await?;
    }
    if !does_search_indexer_exist(name, app_config).await? {
        create_search_indexer(name, name, name, name, app_config).await?;
    }

    Ok(())
}

/// Processes and synchronizes a batch of pages.
async fn sync_pages_batch(
    conn: &mut PgConnection,
    pages: &[Page],
    blob_client: &AzureBlobClient,
    base_url: &Url,
    app_config: &ApplicationConfiguration,
    latest_histories: &HashMap<Uuid, PageHistory>,
) -> anyhow::Result<()> {
    let course_id = pages
        .first()
        .ok_or_else(|| anyhow::anyhow!("No pages to sync."))?
        .course_id
        .ok_or_else(|| anyhow::anyhow!("The first page does not belong to any course."))?;

    let course = headless_lms_models::courses::get_course(conn, course_id).await?;
    let organization =
        headless_lms_models::organizations::get_organization(conn, course.organization_id).await?;
    let task_lm = headless_lms_models::application_task_default_language_models::get_for_task(
        conn,
        ApplicationTask::ContentCleaning,
    )
    .await?;

    let mut base_url = base_url.clone();
    base_url.set_path(&format!(
        "/org/{}/courses/{}",
        organization.slug, course.slug
    ));

    let mut allowed_file_paths = Vec::new();

    for page in pages {
        info!("Syncing page id: {}.", page.id);

        let mut page_url = base_url.clone();
        page_url.set_path(&format!("{}{}", base_url.path(), page.url_path));

        let parsed_content: Vec<GutenbergBlock> = serde_json::from_value(page.content.clone())?;
        let sanitized_blocks = remove_sensitive_attributes(parsed_content);

        let content_to_upload = match convert_material_blocks_to_markdown_with_llm(
            &sanitized_blocks,
            app_config,
            &task_lm,
        )
        .await
        {
            Ok(markdown) => {
                info!("Successfully cleaned content for page {}", page.id);
                // Check if the markdown is empty, or if it just contains all spaces or newlines
                if markdown.trim().is_empty() {
                    warn!(
                        "Markdown is empty for page {}. Generating fallback content with a fake heading.",
                        page.id
                    );
                    format!("# {}", page.title)
                } else {
                    markdown
                }
            }
            Err(e) => {
                let error_msg = format!("Sync failed: LLM processing error: {}", e);
                warn!(
                    "Failed to clean content with LLM for page {}: {}. Using serialized sanitized content instead.",
                    page.id, error_msg
                );
                if let Err(db_err) =
                    headless_lms_models::chatbot_page_sync_statuses::set_page_sync_error(
                        conn, page.id, &error_msg,
                    )
                    .await
                {
                    warn!(
                        "Failed to record sync error for page {}: {:?}",
                        page.id, db_err
                    );
                }
                // Fallback to original content
                serde_json::to_string(&sanitized_blocks)?
            }
        };

        let blob_path = generate_blob_path(page)?;
        let chapter: Option<DatabaseChapter> = if page.chapter_id.is_some() {
            match headless_lms_models::chapters::get_chapter_by_page_id(conn, page.id).await {
                Ok(c) => Some(c),
                Err(e) => {
                    debug!("Chapter lookup failed for page {}: {}", page.id, e);
                    None
                }
            }
        } else {
            None
        };

        allowed_file_paths.push(blob_path.clone());
        let mut metadata = HashMap::new();
        // Azure Blob Storage metadata values must be ASCII-only. URL-encode values that may
        // contain non-ASCII characters (e.g., Finnish characters like ä, ö) to ensure they
        // are ASCII-compatible. We decode the url and the title before we save them in our database.
        metadata.insert("url".to_string(), url_encode(page_url.as_ref()));
        metadata.insert("title".to_string(), url_encode(&page.title));
        metadata.insert(
            "course_id".to_string(),
            page.course_id.unwrap_or(Uuid::nil()).to_string().into(),
        );
        metadata.insert(
            "language".to_string(),
            course.language_code.to_string().into(),
        );
        metadata.insert("filepath".to_string(), blob_path.clone().into());
        if let Some(c) = chapter {
            metadata.insert(
                "chunk_context".to_string(),
                url_encode(&format!(
                    "This chunk is a snippet from page {} from chapter {}: {} of the course {}.",
                    page.title, c.chapter_number, c.name, course.name,
                )),
            );
        } else {
            metadata.insert(
                "chunk_context".to_string(),
                url_encode(&format!(
                    "This chunk is a snippet from page {} of the course {}.",
                    page.title, course.name,
                )),
            );
        }

        if let Err(e) = blob_client
            .upload_file(&blob_path, content_to_upload.as_bytes(), Some(metadata))
            .await
        {
            let error_msg = format!("Sync failed: Upload error: {}", e);
            warn!("Failed to upload file {}: {:?}", blob_path, e);
            if let Err(db_err) =
                headless_lms_models::chatbot_page_sync_statuses::set_page_sync_error(
                    conn, page.id, &error_msg,
                )
                .await
            {
                warn!(
                    "Failed to record upload error for page {}: {:?}",
                    page.id, db_err
                );
            }
        } else if let Some(history_id) = latest_histories.get(&page.id) {
            let mut page_revision_map = HashMap::new();
            page_revision_map.insert(page.id, history_id.id);
            if let Err(e) =
                headless_lms_models::chatbot_page_sync_statuses::update_page_revision_ids(
                    conn,
                    page_revision_map,
                )
                .await
            {
                let error_msg = format!("Sync failed: Status update error: {}", e);
                warn!("Failed to update sync status for page {}: {:?}", page.id, e);
                if let Err(db_err) =
                    headless_lms_models::chatbot_page_sync_statuses::set_page_sync_error(
                        conn, page.id, &error_msg,
                    )
                    .await
                {
                    warn!(
                        "Failed to record status update error for page {}: {:?}",
                        page.id, db_err
                    );
                }
            }
        }
    }

    Ok(())
}

/// Generates the blob storage path for a given page.
fn generate_blob_path(page: &Page) -> anyhow::Result<String> {
    let course_id = page
        .course_id
        .ok_or_else(|| anyhow::anyhow!("Page {} does not belong to any course.", page.id))?;

    Ok(format!("courses/{}/pages/{}.md", course_id, page.id))
}

/// Deletes files from blob storage that are no longer associated with any public page.
/// This includes files for deleted pages, hidden pages, and any other pages that are no longer public.
async fn delete_old_files(
    conn: &mut PgConnection,
    course_id: Uuid,
    blob_client: &AzureBlobClient,
) -> anyhow::Result<()> {
    let mut courses_prefix = "courses/".to_string();
    courses_prefix.push_str(&course_id.to_string());
    let existing_files = blob_client.list_files_with_prefix(&courses_prefix).await?;

    let pages = headless_lms_models::pages::get_all_by_course_id_and_visibility(
        conn,
        course_id,
        PageVisibility::Public,
    )
    .await?;

    let allowed_paths: HashSet<String> = pages
        .iter()
        .filter_map(|page| generate_blob_path(page).ok())
        .collect();

    for file in existing_files {
        if !allowed_paths.contains(&file) {
            info!("Deleting obsolete file: {}", file);
            blob_client.delete_file(&file).await?;
        }
    }

    Ok(())
}