termi-blog/backend/src/workers/ai_reindex.rs

use loco_rs::prelude::*;
use serde::{Deserialize, Serialize};

use crate::services::{ai, worker_jobs};

pub struct AiReindexWorker {
    pub ctx: AppContext,
}

#[derive(Clone, Debug, Default, Deserialize, Serialize)]
pub struct AiReindexWorkerArgs {
    #[serde(default)]
    pub job_id: Option<i32>,
}

#[async_trait]
impl BackgroundWorker<AiReindexWorkerArgs> for AiReindexWorker {
    fn build(ctx: &AppContext) -> Self {
        Self { ctx: ctx.clone() }
    }

    fn tags() -> Vec<String> {
        vec!["ai".to_string(), "reindex".to_string()]
    }

    async fn perform(&self, args: AiReindexWorkerArgs) -> Result<()> {
        if let Some(job_id) = args.job_id {
            if !worker_jobs::begin_job_execution(&self.ctx, job_id).await? {
                return Ok(());
            }

            match ai::rebuild_index(&self.ctx, Some(job_id)).await {
                Ok(summary) => {
                    worker_jobs::mark_job_succeeded(
                        &self.ctx,
                        job_id,
                        Some(serde_json::json!({
                            "phase": "completed",
                            "message": "AI 索引重建完成。",
                            "progress": {
                                "phase": "completed",
                                "message": "AI 索引重建完成。",
                                "total_chunks": summary.indexed_chunks,
                                "processed_chunks": summary.indexed_chunks,
                                "total_batches": summary.indexed_chunks.div_ceil(ai::REINDEX_EMBEDDING_BATCH_SIZE.max(1)),
                                "current_batch": summary.indexed_chunks.div_ceil(ai::REINDEX_EMBEDDING_BATCH_SIZE.max(1)),
                                "batch_size": ai::REINDEX_EMBEDDING_BATCH_SIZE.max(1),
                                "percent": 100,
                            },
                            "indexed_chunks": summary.indexed_chunks,
                            "last_indexed_at": summary.last_indexed_at.map(|value| value.to_rfc3339()),
                        })),
                    )
                    .await?;
                    Ok(())
                }
                Err(error) => {
                    if worker_jobs::cancel_job_if_requested(
                        &self.ctx,
                        job_id,
                        "job cancelled during reindex",
                    )
                    .await?
                    {
                        return Ok(());
                    }

                    worker_jobs::mark_job_failed(&self.ctx, job_id, error.to_string()).await?;
                    Err(error)
                }
            }
        } else {
            ai::rebuild_index(&self.ctx, None).await?;
            Ok(())
        }
    }
}