sqlite_graphrag/embedder/
passage.rs

1//! Passage/query embedding entry points.
2
3use super::*;
4use crate::errors::AppError;
5use std::path::Path;
6
7/// v1.0.89 (BUG-SKIP-EMBED): reads `--skip-embedding-on-failure` / runtime_config
8/// (flag > XDG; product env is not read).
9/// Returns `true` when the user opted to persist with NULL embedding on failure.
10pub fn should_skip_embedding_on_failure() -> bool {
11    crate::runtime_config::skip_embedding_on_failure()
12}
13
14/// v1.0.89 (BUG-SKIP-EMBED + GAP-EMBED-PROPAGATION): embed a passage
15/// honouring both `--llm-backend` and `--skip-embedding-on-failure`.
16///
17/// On success returns `Ok(Some(vec))`. On failure:
18/// - if `--skip-embedding-on-failure` is active, logs a warning and returns `Ok(None)`
19/// - otherwise propagates the error (exit 11)
20pub fn embed_passage_or_skip(
21    models_dir: &Path,
22    text: &str,
23    choice: Option<crate::cli::LlmBackendChoice>,
24) -> Result<Option<Vec<f32>>, AppError> {
25    match embed_passage_with_choice(models_dir, text, choice) {
26        Ok((v, _backend)) => Ok(Some(v)),
27        Err(AppError::Validation(msg)) => Err(AppError::Validation(msg)),
28        Err(e) => {
29            if should_skip_embedding_on_failure() {
30                tracing::warn!(
31                    error = %e,
32                    "embedding failed but --skip-embedding-on-failure is active; persisting with NULL embedding"
33                );
34                Ok(None)
35            } else {
36                Err(e)
37            }
38        }
39    }
40}
41
42// =============================================================================
43// v1.0.82 (GAP-003): wrappers that take the CLI choice
44// (`crate::cli::LlmBackendChoice`) and translate it into a chain for
45// `embed_with_fallback`. They centralize the propagation of the
46// `--llm-backend` flag across the 6 commands that produce embeddings
47// (`remember`, `edit`, `ingest`, `enrich`, `recall`, `hybrid-search`).
48// =============================================================================
49
50/// Embed a single passage using the LLM backend selected by the user via
51/// `--llm-backend`. Routes to `embed_with_fallback` so failures fall
52/// through to the next backend in the chain before giving up.
53///
54/// When `choice` is `None` (e.g. a sub-command that does not yet
55/// expose the flag), the default `OpenRouter` chain is used.
56pub fn embed_passage_with_choice(
57    models_dir: &Path,
58    text: &str,
59    choice: Option<crate::cli::LlmBackendChoice>,
60) -> Result<(Vec<f32>, LlmBackendKind), AppError> {
61    let _slot_guard = acquire_llm_slot_for_embedding()?;
62    let chain = choice
63        .unwrap_or(crate::cli::LlmBackendChoice::OpenRouter)
64        .to_chain();
65    embed_with_fallback(models_dir, text, &chain, false)
66}
67
68/// v1.0.93: embedding with `EmbeddingBackendChoice` awareness.
69pub fn embed_passage_with_embedding_choice(
70    models_dir: &Path,
71    text: &str,
72    backends: crate::cli::BackendChoice,
73) -> Result<(Vec<f32>, LlmBackendKind), AppError> {
74    let crate::cli::BackendChoice {
75        llm: llm_backend,
76        embedding: embedding_backend,
77    } = backends;
78    let _slot_guard = acquire_llm_slot_for_embedding()?;
79    let chain = embedding_backend.to_chain(llm_backend);
80    embed_with_fallback(models_dir, text, &chain, false)
81}
82
83/// failure, returns a structured `FallbackReason` so the caller can
84/// surface `vec_degraded` instead of a hard exit 11.
85///
86/// `None` matches the legacy `try_embed_query_with_fallback` path
87/// (uses the active embedder without an explicit chain).
88pub fn try_embed_query_with_choice(
89    models_dir: &Path,
90    text: &str,
91    choice: Option<crate::cli::LlmBackendChoice>,
92) -> Result<(Vec<f32>, LlmBackendKind), FallbackReason> {
93    match embed_passage_with_choice(models_dir, text, choice) {
94        // GAP-004 / v1.0.85.1: when the chain terminates on
95        // `LlmBackendKind::None` (i.e. the user passed `--llm-backend none`,
96        // or every preceding backend failed), `embed_with_fallback` returns
97        // `Ok((vec![], LlmBackendKind::None))` instead of an error. Without
98        // this guard the empty vector would propagate to the dimension check,
99        // which aborts with exit 11 ("embedding has 0 dims, expected 64").
100        // The caller's contract here is to surface a typed `FallbackReason`
101        // instead, so `recall` and `hybrid-search` can route to FTS5-puro via
102        // the existing `vec_degraded` / `vec_degraded_reason` envelope.
103        // Intercept the empty-vector success path and surface it as
104        // `FallbackReason::DimZero` (introduced at v1.0.85 / ADR-0043
105        // for the symmetric LLM-returned-zero-dim case).
106        Ok((v, _backend)) if v.is_empty() => Err(FallbackReason::DimZero),
107        Ok((v, backend)) => Ok((v, backend)),
108        Err(e) => Err(classify_embedding_error(e)),
109    }
110}
111/// v1.0.93 (GAP-OR-INGEST): query embedding with `EmbeddingBackendChoice`
112/// awareness. Mirrors `try_embed_query_with_choice` but routes through
113/// `embed_passage_with_embedding_choice` so OpenRouter API is used when
114/// configured.
115pub fn try_embed_query_with_embedding_choice(
116    models_dir: &Path,
117    text: &str,
118    backends: crate::cli::BackendChoice,
119) -> Result<(Vec<f32>, LlmBackendKind), FallbackReason> {
120    match embed_passage_with_embedding_choice(models_dir, text, backends) {
121        Ok((v, _backend)) if v.is_empty() => Err(FallbackReason::DimZero),
122        Ok((v, backend)) => Ok((v, backend)),
123        Err(e) => Err(classify_embedding_error(e)),
124    }
125}
126
127/// call. Reads max-concurrency from `--llm-max-host-concurrency` /
128/// XDG `llm.max_host_concurrency` (default derived from `LLM_WORKER_RSS_MB`
129/// and available memory), and the wait timeout from XDG
130/// `llm.slot_wait_secs` (default 30s).
131///
132/// Returns `Ok(guard)` for happy path, `AppError::LockBusy` (exit 75)
133/// when no slot is available within the wait window, and
134/// `AppError::Validation` when the concurrency is 0.
135///
136/// Tests may force fail-fast via XDG/runtime slot wait of 0.
137pub(crate) fn acquire_llm_slot_for_embedding() -> Result<crate::llm_slots::LlmSlotGuard, AppError> {
138    use crate::constants::{CLI_LOCK_DEFAULT_WAIT_SECS, LLM_WORKER_RSS_MB};
139    let default_max = crate::llm_slots::default_max_concurrency() as usize;
140    let max = crate::runtime_config::llm_max_host_concurrency(default_max).max(1) as u32;
141    let wait_secs = if crate::runtime_config::llm_slot_no_wait() {
142        0
143    } else {
144        crate::runtime_config::llm_slot_wait_secs(CLI_LOCK_DEFAULT_WAIT_SECS)
145    };
146    let _ = LLM_WORKER_RSS_MB; // silence the unused import (used in default_max_concurrency)
147                               // GAP-003 / ADR-0043: when the slot semaphore is contended beyond the
148                               // backoff window (50 + 100 + 200 + 400 = 750ms total), return a
149                               // marker message that `classify_embedding_error` maps to
150                               // `FallbackReason::SlotExhausted` (discriminator `slot_exhausted`).
151                               // The window is shorter than the legacy 30s timeout, so the operator
152                               // observes FTS5-puro fallback quickly instead of after 30s of silence.
153    match crate::llm_slots::acquire_llm_slot(max, wait_secs) {
154        Ok(guard) => Ok(guard),
155        Err(e @ AppError::LockBusy { .. }) if wait_secs > 0 => Err(AppError::Embedding(
156            crate::i18n::validation::embedding_slot_exhausted(&e),
157        )),
158        Err(e) => Err(e),
159    }
160}
sqlite_graphrag/embedder/passage.rs

sqlite_graphrag/embedder/
passage.rs
