Skip to main content

sqlite_graphrag/
query_embedding.rs

1//! The one place a read path turns a query string into a live embedding.
2//!
3//! `recall` and `hybrid-search` both need the same three-way outcome — a
4//! vector, a deliberate skip, or a degradation — and both need to report it the
5//! same way on their envelope. They used to carry byte-identical copies of that
6//! logic, including the reason string, which is duplicated knowledge rather than
7//! merely similar code: changing the degradation contract meant remembering to
8//! edit two files, and nothing failed if you edited one.
9//!
10//! Keeping the resolution here also keeps the envelope honest. `vec_degraded`
11//! is what tells a caller that a hybrid search silently became a pure FTS5
12//! search, and a single implementation is what guarantees both commands raise
13//! it under exactly the same conditions.
14
15/// Live query embedding plus the flags describing whether it succeeded.
16///
17/// `backend_invoked` names the backend that actually ran, and is `None` both
18/// when the caller opted out and when every attempt failed.
19///
20/// Was a four-element tuple until v1.2.5. It became a struct when `reason_code`
21/// was added because the field next to it, `backend_invoked`, has the SAME type:
22/// two adjacent `Option<&'static str>` in a tuple swap silently, the compiler
23/// accepts it, and the only symptom is a degradation classified as the wrong
24/// error. A named field makes that swap impossible to write.
25///
26/// `reason_code` is the machine-readable half of `error`, and it exists because
27/// [`degradation_failure`] derives the error class from the code and never from
28/// the prose. Until v1.2.5 this resolver logged the code and threw it away, so
29/// no caller could satisfy that contract — which is why `--fail-on-degraded`
30/// shipped declared, documented and never once consulted.
31pub struct QueryEmbedding {
32    /// The query vector, or `None` when the read degraded to FTS5-only.
33    pub embedding: Option<Vec<f32>>,
34    /// Whether the read fell back to BM25 alone.
35    pub degraded: bool,
36    /// Operator-facing prose for the degradation.
37    pub error: Option<String>,
38    /// Backend that actually produced the vector.
39    pub backend_invoked: Option<&'static str>,
40    /// Stable code for the degradation, `None` when nothing degraded.
41    pub reason_code: Option<&'static str>,
42}
43
44/// Machine-readable `vec_error` for a degradation the operator asked for.
45///
46/// Named rather than inlined because it is a value consumers match on: it is the
47/// one `vec_error` that means "nothing went wrong", so a caller distinguishing a
48/// deliberate skip from a real failure compares against this exact string.
49pub const FALLBACK_FTS_ONLY_REASON: &str = "fallback_fts_only requested";
50
51/// `reason_code` recorded for a degradation the operator asked for.
52///
53/// The prose in [`FALLBACK_FTS_ONLY_REASON`] is what a human reads; this is what
54/// [`degradation_failure`] branches on. Two representations because the envelope
55/// has always carried the prose and changing it would break consumers.
56pub const FALLBACK_FTS_ONLY_CODE: &str = "fallback_fts_only";
57
58/// Decides whether a degraded read must become a non-zero exit.
59///
60/// Returns `None` — the read stands, exit 0, envelope untouched — when any of:
61/// - `fail_on_degraded` is off, which is the default and the historical
62///   behaviour byte for byte;
63/// - nothing degraded;
64/// - the degradation was REQUESTED with `--fallback-fts-only`.
65///
66/// That third case is the whole point of the discriminator. `--fallback-fts-only`
67/// is an operator saying "skip the provider, BM25 is what I want"; turning their
68/// own instruction into a failure would make the two flags mutually unusable.
69///
70/// # Error classification
71///
72/// The class is derived from `reason_code`, never from the message prose, so a
73/// reworded string cannot silently reclassify a failure:
74/// - `timeout`, `slot_exhausted`, `oauth_quota`, `cancelled` — the provider was
75///   unreachable or too slow. [`crate::errors::AppError::Timeout`] is retryable, so
76///   `error_class` is `transient` and `retryable` is `true`: retrying is exactly
77///   the right advice.
78/// - anything else (`dim_zero`, `backend_mismatch`, `embedding_failed`) — the
79///   configuration or the response shape is wrong, and retrying an unchanged
80///   invocation reproduces it. [`crate::errors::AppError::Embedding`] carries exit 11.
81pub fn degradation_failure(
82    fail_on_degraded: bool,
83    vec_degraded: bool,
84    reason_code: Option<&str>,
85) -> Option<crate::errors::AppError> {
86    if !fail_on_degraded || !vec_degraded {
87        return None;
88    }
89    let code = reason_code.unwrap_or("unknown");
90    if code == FALLBACK_FTS_ONLY_CODE {
91        return None;
92    }
93    // Built here rather than in `i18n::validation` because the operator-facing
94    // half of this message is the `vec_error` the envelope ALREADY carries,
95    // localised at its own source; this string only names the discriminator.
96    let detail = format!("query embedding degraded to FTS5-only ({code})");
97    match code {
98        "timeout" | "slot_exhausted" | "oauth_quota" | "cancelled" => {
99            Some(crate::errors::AppError::Timeout {
100                operation: detail,
101                duration_secs: 0,
102            })
103        }
104        _ => Some(crate::errors::AppError::Embedding(detail)),
105    }
106}
107
108/// Resolves the query embedding, degrading to FTS5-only instead of failing.
109///
110/// When the live embedding cannot be produced — timeout, rate limit,
111/// unreachable provider — the read still returns results, ranked by BM25 alone.
112/// The caller surfaces that through `vec_degraded` and `vec_error` on the
113/// envelope, so the degradation is visible rather than silent.
114///
115/// `--fallback-fts-only` takes the same path deliberately and never contacts the
116/// provider at all.
117///
118/// `log_target` is the tracing target of the calling subcommand, so a degraded
119/// read stays attributable to the command the operator actually ran.
120pub fn resolve_query_embedding(
121    fallback_fts_only: bool,
122    models_dir: &std::path::Path,
123    query: &str,
124    embedding_backend: crate::cli::EmbeddingBackendChoice,
125    llm_backend: crate::cli::LlmBackendChoice,
126    log_target: &'static str,
127) -> QueryEmbedding {
128    if fallback_fts_only {
129        return QueryEmbedding {
130            embedding: None,
131            degraded: true,
132            error: Some(FALLBACK_FTS_ONLY_REASON.to_string()),
133            backend_invoked: None,
134            // The code the operator ASKED for. `degradation_failure` matches on
135            // it to keep `--fallback-fts-only` from turning into a failure.
136            reason_code: Some(FALLBACK_FTS_ONLY_CODE),
137        };
138    }
139    match crate::embedder::try_embed_query_with_embedding_choice(
140        models_dir,
141        query,
142        embedding_backend,
143        llm_backend,
144    ) {
145        Ok((v, backend)) => QueryEmbedding {
146            embedding: Some(v),
147            degraded: false,
148            error: None,
149            backend_invoked: Some(backend.as_str()),
150            reason_code: None,
151        },
152        Err(reason) => {
153            let msg = reason.to_string();
154            let code = reason.reason_code();
155            tracing::warn!(
156                target: "query_embedding",
157                command = log_target,
158                fallback_reason = %msg,
159                reason_code = %code,
160                "live embedding failed; falling back to FTS5"
161            );
162            QueryEmbedding {
163                embedding: None,
164                degraded: true,
165                error: Some(msg),
166                backend_invoked: None,
167                reason_code: Some(code),
168            }
169        }
170    }
171}
172
173#[cfg(test)]
174mod tests {
175    use super::*;
176
177    #[test]
178    fn opting_out_reports_the_named_reason_and_never_names_a_backend() {
179        let resolved = resolve_query_embedding(
180            true,
181            std::path::Path::new("/nonexistent"),
182            "anything",
183            crate::cli::EmbeddingBackendChoice::Openrouter,
184            crate::cli::LlmBackendChoice::None,
185            "recall",
186        );
187        assert!(
188            resolved.embedding.is_none(),
189            "opting out must not produce a vector"
190        );
191        assert!(
192            resolved.degraded,
193            "opting out is still a degradation for the caller"
194        );
195        assert_eq!(resolved.error.as_deref(), Some(FALLBACK_FTS_ONLY_REASON));
196        assert!(
197            resolved.backend_invoked.is_none(),
198            "no provider was contacted, so none may be reported as invoked"
199        );
200        // O código é o que impede `--fail-on-degraded` de transformar a escolha do
201        // operador em falha. Se ele parar de vir, `degradation_failure` cai no ramo
202        // "unknown" e `--fallback-fts-only` passa a reprovar com as duas flags juntas.
203        assert_eq!(
204            resolved.reason_code,
205            Some(FALLBACK_FTS_ONLY_CODE),
206            "opting out must carry its own code, not an absent one"
207        );
208    }
209
210    /// Degradation the caller ASKED for never becomes a failure, flag or no flag.
211    ///
212    /// The pair that closes the loop: the test above proves the code arrives,
213    /// and this one proves what the code is there to decide. Without it, someone
214    /// could drop `reason_code` from the opt-out branch and only one would break.
215    #[test]
216    fn opting_out_survives_fail_on_degraded() {
217        let resolved = resolve_query_embedding(
218            true,
219            std::path::Path::new("/nonexistent"),
220            "anything",
221            crate::cli::EmbeddingBackendChoice::Openrouter,
222            crate::cli::LlmBackendChoice::None,
223            "recall",
224        );
225        assert!(
226            degradation_failure(true, resolved.degraded, resolved.reason_code).is_none(),
227            "--fallback-fts-only com --fail-on-degraded deve continuar saindo 0"
228        );
229    }
230}