solidb 1.2.2

//! Background auto-embedding worker.
//!
//! When a document is inserted/updated whose collection has a vector index with
//! `embedding_source` set but no concrete vector in the target field, the storage
//! layer records a cheap "pending embed" marker (no network on the write path,
//! see `Collection::mark_embed_pending`). This worker sweeps those markers,
//! generates embeddings via the configured LLM provider (async, batched), and
//! writes the vector back into the document — which persists it and updates the
//! HNSW index. This makes "just insert text" work on every write path (HTTP,
//! driver, Lua, bulk, replication) without ever blocking a write on the network.

use super::QueueWorker;
use crate::error::DbError;
use crate::server::llm_client::LLMClient;
use crate::storage::collection::vector::{pending_embed_count, release_pending_embed};
use crate::storage::index::{extract_field_value, VectorIndexConfig};
use crate::storage::Collection;

/// Max documents embedded per (collection, index) per sweep — bounds a single
/// batch request and keeps one collection from starving the others.
const EMBED_BATCH: usize = 128;

/// Backoff (seconds) after a provider/config failure so we don't hammer a
/// down or misconfigured provider every worker tick.
const ERROR_BACKOFF_SECS: u64 = 60;

fn now_secs() -> u64 {
    std::time::SystemTime::now()
        .duration_since(std::time::UNIX_EPOCH)
        .map(|d| d.as_secs())
        .unwrap_or(0)
}

impl QueueWorker {
    /// One embedding sweep. Called from the worker loop alongside `check_jobs`.
    pub(crate) async fn check_embeddings(&self) {
        // Fast path: nothing pending anywhere → no enumeration at all.
        // Read once: whatever the gauge claims now is what a fruitless sweep
        // is allowed to retire below, so marks recorded while we sweep survive.
        let claimed = pending_embed_count();
        if claimed == 0 {
            return;
        }
        // Serialize with the other periodic scanners; skip if another caller
        // holds it. `check_jobs` and `check_materialized_views` both took this
        // lock and this sweep did not, so back when the worker loop was spawned
        // `QUEUE_WORKERS` times (four by default) every one of those loops
        // enumerated every collection in the instance on the same five-second
        // tick. The loop is single now, but the guard is what makes that safe.
        let _lock = match self.claiming_lock.try_lock() {
            Ok(l) => l,
            Err(_) => return,
        };

        // One pass over the column families instead of one per database.
        // `Database::list_collections` calls `DB::cf_names`, which clones every
        // column-family name in the instance on each call, so driving it from
        // the database list cost `databases × total collections` string
        // allocations on every worker tick — 46 × 969 here, five seconds apart.
        let grouped = self.storage.collections_grouped();
        let mut saw_pending = false;

        // Forget backoffs that have expired, so the map only holds indexes
        // that are failing right now.
        let now = now_secs();
        if let Ok(mut backoff) = self.embed_backoff.lock() {
            backoff.retain(|_, until| *until > now);
        }

        for (db_name, coll_names) in grouped {
            let db = match self.storage.get_database(&db_name) {
                Ok(d) => d,
                Err(_) => continue,
            };
            for coll_name in coll_names {
                let coll = match db.system_collection(&coll_name) {
                    Ok(c) => c,
                    Err(_) => continue,
                };
                let configs = coll.get_all_vector_index_configs();
                for config in configs {
                    if config.embedding_source.is_none() {
                        continue;
                    }
                    let backoff_key = format!("{}\u{0}{}\u{0}{}", db_name, coll_name, config.name);
                    let backing_off = self
                        .embed_backoff
                        .lock()
                        .map(|b| b.get(&backoff_key).is_some_and(|until| *until > now))
                        .unwrap_or(false);
                    if backing_off {
                        // Its markers are still there; don't let the gauge
                        // reconciliation below treat this pass as empty.
                        saw_pending = true;
                        continue;
                    }
                    let pending = coll.list_embed_pending(&config.name, EMBED_BATCH);
                    if pending.is_empty() {
                        continue;
                    }
                    saw_pending = true;
                    if let Err(e) = self
                        .embed_pending_batch(&db_name, &coll, &config, &pending)
                        .await
                    {
                        // Provider/config problem — back off this index only
                        // and carry on with the rest. Markers remain and are
                        // retried after the backoff.
                        tracing::warn!(
                            "Auto-embed worker: {}/{} index '{}' failed: {} (backing off {}s)",
                            db_name,
                            coll_name,
                            config.name,
                            e,
                            ERROR_BACKOFF_SECS
                        );
                        if let Ok(mut backoff) = self.embed_backoff.lock() {
                            backoff.insert(backoff_key, now_secs() + ERROR_BACKOFF_SECS);
                        }
                    }
                }
            }
        }

        // A complete pass that found no marker at all means the gauge has
        // drifted above reality — markers vanish with a dropped column family
        // without passing through `clear_embed_pending`. Left alone the gauge
        // never returns to zero, and this enumeration then runs on every tick
        // for the life of the process. Retire only what was claimed on entry.
        if !saw_pending {
            release_pending_embed(claimed);
        }
    }

    /// Embed one batch of pending docs for a single index and persist the vectors.
    async fn embed_pending_batch(
        &self,
        db_name: &str,
        coll: &Collection,
        config: &VectorIndexConfig,
        doc_keys: &[String],
    ) -> Result<(), DbError> {
        let source_field = config.embedding_source.as_deref().unwrap_or_default();

        // Gather (doc_key, source_text), dropping stale markers for docs that were
        // deleted or no longer carry source text.
        let mut keys: Vec<String> = Vec::new();
        let mut texts: Vec<String> = Vec::new();
        for dk in doc_keys {
            let doc = match coll.get(dk) {
                Ok(d) => d,
                Err(_) => {
                    coll.clear_embed_pending(&config.name, dk);
                    continue;
                }
            };
            let value = doc.to_value();
            match extract_field_value(&value, source_field).as_str() {
                Some(t) if !t.trim().is_empty() => {
                    keys.push(dk.clone());
                    texts.push(t.to_string());
                }
                _ => coll.clear_embed_pending(&config.name, dk),
            }
        }
        if keys.is_empty() {
            return Ok(());
        }

        // Embeddings default to OpenAI; an index may override via embedding_provider.
        let provider = config
            .embedding_provider
            .clone()
            .unwrap_or_else(|| "openai".to_string());
        let client = LLMClient::from_storage(
            &self.storage,
            db_name,
            Some(&provider),
            config.embedding_model.clone(),
        )?;

        let text_refs: Vec<&str> = texts.iter().map(|s| s.as_str()).collect();
        let embeddings = client.embed_batch(&text_refs).await?;

        // Write each vector back into its document. `update()` re-runs
        // update_vector_indexes_on_upsert, which — now that the vector is present —
        // indexes the doc and clears the pending marker.
        for (dk, emb) in keys.iter().zip(embeddings) {
            if emb.len() != config.dimension {
                tracing::warn!(
                    "Auto-embed worker: dim mismatch for '{}' index '{}' (got {}, expected {}); dropping marker",
                    dk,
                    config.name,
                    emb.len(),
                    config.dimension
                );
                coll.clear_embed_pending(&config.name, dk);
                continue;
            }
            let doc = match coll.get(dk) {
                Ok(d) => d,
                Err(_) => {
                    coll.clear_embed_pending(&config.name, dk);
                    continue;
                }
            };
            let mut value = doc.to_value();
            if let Some(obj) = value.as_object_mut() {
                obj.insert(config.field.clone(), serde_json::json!(emb));
                if let Err(e) = coll.update(dk, value) {
                    // Leave the marker in place for a later retry.
                    tracing::warn!(
                        "Auto-embed worker: failed to persist vector for '{}': {}",
                        dk,
                        e
                    );
                }
            }
        }
        Ok(())
    }
}