Skip to main content

remem/retrieval/embedding/
fallback.rs

1use anyhow::{bail, Result};
2
3use super::{
4    embed_text_local, local_semantic, status, EmbeddingBackfillTarget, EmbeddingConfig,
5    EmbeddingFallbackCache, EmbeddingProvider, EmbeddingProviderStatus, LocalEmbeddingInputKind,
6    TextEmbedding, FEATURE_HASH_EMBEDDING_DIMENSIONS, FEATURE_HASH_EMBEDDING_MODEL,
7};
8
9#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize)]
10pub struct EmbeddingExecutionMetadata {
11    pub configured_provider: String,
12    /// Provider that produced this query's embedding, after any runtime fallback.
13    pub active_provider: String,
14    pub model: String,
15    pub dimensions: usize,
16    pub degraded: bool,
17    #[serde(skip_serializing_if = "Option::is_none")]
18    pub degradation_reason: Option<String>,
19}
20
21pub(super) fn embedding_execution_metadata(
22    status_before: &EmbeddingProviderStatus,
23    status_after: &EmbeddingProviderStatus,
24    cache: &EmbeddingFallbackCache,
25    embedding: &TextEmbedding,
26) -> Result<EmbeddingExecutionMetadata> {
27    let actual_provider = cache.execution_provider.ok_or_else(|| {
28        anyhow::anyhow!("embedding execution completed without provider metadata")
29    })?;
30    let actual_provider = actual_provider.label();
31    let mut reasons = Vec::new();
32    append_unique_reason(&mut reasons, status_before.degradation_reason.as_deref());
33    append_unique_reason(&mut reasons, cache.degradation_reason.as_deref());
34    append_unique_reason(&mut reasons, status_after.degradation_reason.as_deref());
35    let provider_changed = status_before.active_provider != status_after.active_provider
36        || status_before.active_provider != actual_provider;
37    if provider_changed {
38        reasons.push(format!(
39            "embedding provider changed during query execution: initial={}, final={}, actual={actual_provider}",
40            status_before.active_provider, status_after.active_provider
41        ));
42    }
43    let degraded =
44        status_before.degraded || status_after.degraded || provider_changed || !reasons.is_empty();
45    if degraded && reasons.is_empty() {
46        reasons.push("embedding execution was degraded".to_string());
47    }
48    let degradation_reason = (!reasons.is_empty())
49        .then(|| crate::adapter::common::redact_hook_payload_preview(&reasons.join("; "), 1024));
50    Ok(EmbeddingExecutionMetadata {
51        configured_provider: status_before.configured_provider.clone(),
52        active_provider: actual_provider.to_string(),
53        model: embedding.model().to_string(),
54        dimensions: embedding.dimensions(),
55        degraded,
56        degradation_reason,
57    })
58}
59
60fn append_unique_reason(reasons: &mut Vec<String>, reason: Option<&str>) {
61    if let Some(reason) = reason {
62        if !reasons.iter().any(|existing| existing == reason) {
63            reasons.push(reason.to_string());
64        }
65    }
66}
67
68pub(super) fn embed_local_with_auto_race_fallback(
69    text: &str,
70    kind: LocalEmbeddingInputKind,
71    config: &EmbeddingConfig,
72    cache: &mut EmbeddingFallbackCache,
73) -> Result<TextEmbedding> {
74    match local_semantic::embed_text(text, config, kind) {
75        Ok(embedding) => {
76            cache.execution_provider = Some(EmbeddingProvider::Local);
77            Ok(embedding)
78        }
79        Err(error)
80            if config.provider == EmbeddingProvider::Auto
81                && kind == LocalEmbeddingInputKind::Query
82                && local_semantic::is_model_unavailable_error(&error) =>
83        {
84            let message = format!(
85                "automatic local embedding provider became unavailable: {error}; using feature-hash"
86            );
87            crate::log::error("embedding", &message);
88            remember_feature_hash_fallback(cache, message);
89            feature_hash_embedding(text)
90        }
91        Err(error) => Err(error),
92    }
93}
94
95pub(super) fn embed_with_cached_call_failure_fallback(
96    text: &str,
97    kind: LocalEmbeddingInputKind,
98    config: &EmbeddingConfig,
99    fallback: EmbeddingProvider,
100) -> Result<TextEmbedding> {
101    let fallback_runtime = status::provider_runtime(config, fallback);
102    if let Some(reason) = fallback_runtime.unavailable_reason {
103        bail!(
104            "cached embedding fallback {} unavailable: {reason}",
105            fallback.label()
106        );
107    }
108    match fallback_runtime.provider {
109        EmbeddingProvider::Local => local_semantic::embed_text(text, config, kind),
110        EmbeddingProvider::FeatureHash => feature_hash_embedding(text),
111        EmbeddingProvider::Off => Err(status::embedding_provider_off_error()),
112        EmbeddingProvider::OpenAi | EmbeddingProvider::Auto => {
113            bail!("cached embedding fallback must be local, feature-hash, or off")
114        }
115    }
116}
117
118pub(super) fn embed_with_call_failure_fallback(
119    text: &str,
120    kind: LocalEmbeddingInputKind,
121    config: &EmbeddingConfig,
122    error: anyhow::Error,
123    cache: &mut EmbeddingFallbackCache,
124) -> Result<TextEmbedding> {
125    let Some(fallback) = config.fallback else {
126        return Err(error);
127    };
128    let fallback_runtime = status::provider_runtime(config, fallback);
129    if let Some(reason) = fallback_runtime.unavailable_reason {
130        bail!(
131            "embedding provider api failed: {error}; fallback {} unavailable: {reason}",
132            fallback.label()
133        );
134    }
135    let message = format!(
136        "configured embedding provider api failed: {}; using fallback {}",
137        error,
138        fallback.label()
139    );
140    crate::log::error("embedding", &message);
141    match fallback_runtime.provider {
142        EmbeddingProvider::Local => {
143            let embedding = local_semantic::embed_text(text, config, kind)?;
144            cache.call_failure_fallback = Some(fallback_runtime.provider);
145            cache.call_failure_fallback_target = Some(EmbeddingBackfillTarget {
146                model: embedding.model().to_string(),
147                dimensions: embedding.dimensions(),
148            });
149            cache.execution_provider = Some(EmbeddingProvider::Local);
150            cache.degradation_reason = Some(message);
151            Ok(embedding)
152        }
153        EmbeddingProvider::FeatureHash => {
154            remember_feature_hash_fallback(cache, message);
155            feature_hash_embedding(text)
156        }
157        EmbeddingProvider::Off => Err(status::embedding_provider_off_error_with_cause(format!(
158            "embedding provider api failed: {error}; fallback off disabled provider fallback"
159        ))),
160        EmbeddingProvider::OpenAi | EmbeddingProvider::Auto => Err(error),
161    }
162}
163
164fn remember_feature_hash_fallback(cache: &mut EmbeddingFallbackCache, reason: String) {
165    cache.call_failure_fallback = Some(EmbeddingProvider::FeatureHash);
166    cache.call_failure_fallback_target = Some(EmbeddingBackfillTarget {
167        model: FEATURE_HASH_EMBEDDING_MODEL.to_string(),
168        dimensions: FEATURE_HASH_EMBEDDING_DIMENSIONS,
169    });
170    cache.execution_provider = Some(EmbeddingProvider::FeatureHash);
171    cache.degradation_reason = Some(reason);
172}
173
174fn feature_hash_embedding(text: &str) -> Result<TextEmbedding> {
175    TextEmbedding::new(FEATURE_HASH_EMBEDDING_MODEL, embed_text_local(text))
176}