remem/retrieval/
embedding.rs

1use std::time::Duration;
2
3use anyhow::{bail, Context, Result};
4use serde::{Deserialize, Serialize};
5use sha2::{Digest, Sha256};
6
7mod config;
8mod local_semantic;
9mod status;
10
11use config::env_value;
12pub(crate) use config::resolve_embedding_config;
13use local_semantic::LocalEmbeddingInputKind;
14pub use local_semantic::{
15    LocalEmbeddingDownloadReport, LocalEmbeddingInventoryReport, LocalEmbeddingModelInventory,
16};
17pub(crate) use status::is_embedding_provider_off_error;
18
19pub const FEATURE_HASH_EMBEDDING_DIMENSIONS: usize = 768;
20pub const FEATURE_HASH_EMBEDDING_MODEL: &str = "remem-local-feature-hash-v1";
21pub const LOCAL_EMBEDDING_DIMENSIONS: usize = FEATURE_HASH_EMBEDDING_DIMENSIONS;
22pub const LOCAL_EMBEDDING_MODEL: &str = FEATURE_HASH_EMBEDDING_MODEL;
23
24const DEFAULT_PROVIDER: EmbeddingProvider = EmbeddingProvider::Auto;
25const OPENAI_DEFAULT_BASE_URL: &str = "https://api.openai.com/v1";
26const OPENAI_DEFAULT_MODEL: &str = "text-embedding-3-small";
27const DEFAULT_API_KEY_ENV: &str = "OPENAI_API_KEY";
28const DEFAULT_TIMEOUT_SECS: u64 = 30;
29
30const ENV_PROVIDER: &str = "REMEM_EMBEDDINGS_PROVIDER";
31const ENV_PROVIDER_LEGACY: &str = "REMEM_EMBEDDING_PROVIDER";
32const ENV_MODEL: &str = "REMEM_EMBEDDINGS_MODEL";
33const ENV_MODEL_LEGACY: &str = "REMEM_EMBEDDING_MODEL";
34const ENV_BASE_URL: &str = "REMEM_EMBEDDINGS_BASE_URL";
35const ENV_BASE_URL_LEGACY: &str = "REMEM_EMBEDDING_BASE_URL";
36const ENV_DIMENSIONS: &str = "REMEM_EMBEDDINGS_DIMENSIONS";
37const ENV_DIMENSIONS_LEGACY: &str = "REMEM_EMBEDDING_DIMENSIONS";
38const ENV_API_KEY: &str = "REMEM_EMBEDDINGS_API_KEY";
39const ENV_API_KEY_LEGACY: &str = "REMEM_EMBEDDING_API_KEY";
40const ENV_API_KEY_ENV: &str = "REMEM_EMBEDDINGS_API_KEY_ENV";
41const ENV_TIMEOUT_SECS: &str = "REMEM_EMBEDDINGS_TIMEOUT_SECS";
42const ENV_FALLBACK: &str = "REMEM_EMBEDDINGS_FALLBACK";
43const ENV_MODEL_DIR: &str = "REMEM_EMBEDDINGS_MODEL_DIR";
44
45#[derive(Debug, Clone, Copy, PartialEq, Eq)]
46pub enum EmbeddingProvider {
47    Auto,
48    Local,
49    FeatureHash,
50    OpenAi,
51    Off,
52}
53
54impl EmbeddingProvider {
55    fn parse(raw: &str) -> Result<Self> {
56        match raw.trim().to_ascii_lowercase().as_str() {
57            "auto" => Ok(Self::Auto),
58            "local" => Ok(Self::Local),
59            "feature-hash" | "feature_hash" | "offline" => Ok(Self::FeatureHash),
60            "api" | "openai" | "openai-compatible" | "openai_compatible" => Ok(Self::OpenAi),
61            "off" | "disabled" | "none" => Ok(Self::Off),
62            other => bail!("unknown embeddings.provider: {other}"),
63        }
64    }
65
66    pub fn label(self) -> &'static str {
67        match self {
68            Self::Auto => "auto",
69            Self::Local => "local",
70            Self::FeatureHash => "feature-hash",
71            Self::OpenAi => "api",
72            Self::Off => "off",
73        }
74    }
75}
76
77#[derive(Debug, Clone, PartialEq, Eq)]
78pub struct EmbeddingConfig {
79    pub provider: EmbeddingProvider,
80    pub fallback: Option<EmbeddingProvider>,
81    pub model: String,
82    pub base_url: String,
83    pub dimensions: Option<usize>,
84    pub api_key_env: String,
85    pub model_dir: Option<String>,
86    pub timeout_secs: u64,
87}
88
89impl Default for EmbeddingConfig {
90    fn default() -> Self {
91        Self {
92            provider: DEFAULT_PROVIDER,
93            fallback: None,
94            model: OPENAI_DEFAULT_MODEL.to_string(),
95            base_url: OPENAI_DEFAULT_BASE_URL.to_string(),
96            dimensions: None,
97            api_key_env: DEFAULT_API_KEY_ENV.to_string(),
98            model_dir: None,
99            timeout_secs: DEFAULT_TIMEOUT_SECS,
100        }
101    }
102}
103
104#[derive(Debug, Clone, PartialEq, Eq)]
105pub struct EmbeddingProviderStatus {
106    pub configured_provider: String,
107    pub fallback_provider: Option<String>,
108    pub active_provider: String,
109    pub active_model_id: Option<String>,
110    pub active_dimensions: Option<usize>,
111    pub degraded: bool,
112    pub disabled: bool,
113    pub unavailable_reason: Option<String>,
114    pub degradation_reason: Option<String>,
115    pub model_dir: Option<String>,
116}
117
118#[derive(Debug, Clone, PartialEq)]
119pub struct TextEmbedding {
120    model: String,
121    values: Vec<f32>,
122}
123
124impl TextEmbedding {
125    pub fn new(model: impl Into<String>, values: Vec<f32>) -> Result<Self> {
126        let model = model.into();
127        if model.trim().is_empty() {
128            bail!("embedding model must not be empty");
129        }
130        validate_embedding_values(&values)?;
131        Ok(Self { model, values })
132    }
133
134    pub fn model(&self) -> &str {
135        &self.model
136    }
137
138    pub fn values(&self) -> &[f32] {
139        &self.values
140    }
141
142    pub fn dimensions(&self) -> usize {
143        self.values.len()
144    }
145
146    pub fn profile(&self) -> EmbeddingProfile<'_> {
147        EmbeddingProfile {
148            model: &self.model,
149            dimensions: self.values.len(),
150        }
151    }
152}
153
154#[derive(Debug, Clone, Copy, PartialEq, Eq)]
155pub struct EmbeddingProfile<'a> {
156    pub model: &'a str,
157    pub dimensions: usize,
158}
159
160#[derive(Debug, Clone, PartialEq, Eq)]
161pub struct EmbeddingBackfillTarget {
162    pub model: String,
163    pub dimensions: usize,
164}
165
166#[derive(Debug, Default)]
167pub(crate) struct EmbeddingFallbackCache {
168    call_failure_fallback: Option<EmbeddingProvider>,
169    call_failure_fallback_target: Option<EmbeddingBackfillTarget>,
170}
171
172impl EmbeddingFallbackCache {
173    pub(crate) fn call_failure_fallback_target(&self) -> Option<EmbeddingBackfillTarget> {
174        match self.call_failure_fallback {
175            Some(EmbeddingProvider::Local) => self.call_failure_fallback_target.clone(),
176            Some(EmbeddingProvider::FeatureHash) => Some(EmbeddingBackfillTarget {
177                model: FEATURE_HASH_EMBEDDING_MODEL.to_string(),
178                dimensions: FEATURE_HASH_EMBEDDING_DIMENSIONS,
179            }),
180            Some(EmbeddingProvider::Auto)
181            | Some(EmbeddingProvider::OpenAi)
182            | Some(EmbeddingProvider::Off) => None,
183            None => None,
184        }
185    }
186}
187
188pub fn embed_query(query: &str) -> Result<TextEmbedding> {
189    embed_text(query, LocalEmbeddingInputKind::Query)
190}
191
192pub(crate) fn embed_query_with_fallback_cache(
193    query: &str,
194    cache: &mut EmbeddingFallbackCache,
195) -> Result<TextEmbedding> {
196    embed_text_with_fallback_cache(query, LocalEmbeddingInputKind::Query, cache)
197}
198
199pub(crate) fn embed_query_if_enabled(query: &str) -> Result<Option<TextEmbedding>> {
200    match embed_query(query) {
201        Ok(embedding) => Ok(Some(embedding)),
202        Err(error) if is_embedding_provider_off_error(&error) => Ok(None),
203        Err(error) => Err(error),
204    }
205}
206
207pub fn embed_memory(
208    title: &str,
209    content: &str,
210    memory_type: &str,
211    topic_key: Option<&str>,
212) -> Result<TextEmbedding> {
213    let text = memory_embedding_text(title, content, memory_type, topic_key);
214    embed_text(&text, LocalEmbeddingInputKind::Passage)
215}
216
217pub(crate) fn embed_memory_with_fallback_cache(
218    title: &str,
219    content: &str,
220    memory_type: &str,
221    topic_key: Option<&str>,
222    cache: &mut EmbeddingFallbackCache,
223) -> Result<TextEmbedding> {
224    let text = memory_embedding_text(title, content, memory_type, topic_key);
225    embed_text_with_fallback_cache(&text, LocalEmbeddingInputKind::Passage, cache)
226}
227
228pub fn embed_query_text_local(query: &str) -> Vec<f32> {
229    embed_text_local(query)
230}
231
232pub fn embed_memory_text_local(
233    title: &str,
234    content: &str,
235    memory_type: &str,
236    topic_key: Option<&str>,
237) -> Vec<f32> {
238    embed_text_local(&memory_embedding_text(
239        title,
240        content,
241        memory_type,
242        topic_key,
243    ))
244}
245
246pub fn embedding_content_hash(
247    title: &str,
248    content: &str,
249    memory_type: &str,
250    topic_key: Option<&str>,
251) -> String {
252    let mut hasher = Sha256::new();
253    hasher.update(memory_type.as_bytes());
254    hasher.update([0]);
255    if let Some(topic_key) = topic_key {
256        hasher.update(topic_key.as_bytes());
257    }
258    hasher.update([0]);
259    hasher.update(title.as_bytes());
260    hasher.update([0]);
261    hasher.update(content.as_bytes());
262    let digest = hasher.finalize();
263    digest.iter().map(|byte| format!("{byte:02x}")).collect()
264}
265
266pub(crate) fn configured_backfill_target() -> Result<EmbeddingBackfillTarget> {
267    let mut cache = EmbeddingFallbackCache::default();
268    configured_backfill_target_with_fallback_cache(&mut cache)
269}
270
271pub(crate) fn configured_backfill_target_with_fallback_cache(
272    cache: &mut EmbeddingFallbackCache,
273) -> Result<EmbeddingBackfillTarget> {
274    #[cfg(test)]
275    let _test_env_guard = config::lock_test_env();
276    let status = embedding_provider_status_without_probe()?;
277    if let Some(error) = disabled_provider_status_error(&status) {
278        return Err(error);
279    }
280    let probe = embed_text_with_fallback_cache(
281        "remem embedding profile probe",
282        LocalEmbeddingInputKind::Generic,
283        cache,
284    )?;
285    Ok(EmbeddingBackfillTarget {
286        model: probe.model().to_string(),
287        dimensions: probe.dimensions(),
288    })
289}
290
291pub fn embedding_provider_status() -> Result<EmbeddingProviderStatus> {
292    #[cfg(test)]
293    let _test_env_guard = config::lock_test_env();
294    let config = resolve_embedding_config()?;
295    let mut status = status::resolve_provider_status(&config);
296    status::probe_active_api_profile(&config, &mut status);
297    Ok(status)
298}
299
300pub(crate) fn embedding_provider_status_without_probe() -> Result<EmbeddingProviderStatus> {
301    #[cfg(test)]
302    let _test_env_guard = config::lock_test_env();
303    let config = resolve_embedding_config()?;
304    Ok(status::resolve_provider_status(&config))
305}
306
307pub(crate) fn disabled_provider_status_error(
308    status: &EmbeddingProviderStatus,
309) -> Option<anyhow::Error> {
310    if !status.disabled {
311        return None;
312    }
313    status
314        .degradation_reason
315        .clone()
316        .or_else(|| status.unavailable_reason.clone())
317        .map(status::embedding_provider_off_error_with_cause)
318        .or_else(|| Some(status::embedding_provider_off_error()))
319}
320
321pub(crate) fn provider_disabled_or_error() -> Result<bool> {
322    let status = embedding_provider_status_without_probe()?;
323    match disabled_provider_status_error(&status) {
324        Some(error) if is_embedding_provider_off_error(&error) => Ok(true),
325        Some(error) => Err(error),
326        None => Ok(false),
327    }
328}
329
330pub(crate) fn configured_local_embedding_model_id(config: &EmbeddingConfig) -> Result<String> {
331    local_semantic::configured_model_id(config)
332}
333
334pub fn download_local_embedding_model(model: Option<&str>) -> Result<LocalEmbeddingDownloadReport> {
335    local_semantic::download_model(model)
336}
337
338pub fn local_embedding_inventory() -> Result<LocalEmbeddingInventoryReport> {
339    local_semantic::inventory()
340}
341
342pub(crate) fn is_local_embedding_model_unavailable_error(error: &anyhow::Error) -> bool {
343    local_semantic::is_model_unavailable_error(error)
344}
345
346fn embed_text(text: &str, kind: LocalEmbeddingInputKind) -> Result<TextEmbedding> {
347    let mut cache = EmbeddingFallbackCache::default();
348    embed_text_with_fallback_cache(text, kind, &mut cache)
349}
350
351fn embed_text_with_fallback_cache(
352    text: &str,
353    kind: LocalEmbeddingInputKind,
354    cache: &mut EmbeddingFallbackCache,
355) -> Result<TextEmbedding> {
356    #[cfg(test)]
357    let _test_env_guard = config::lock_test_env();
358    let config = resolve_embedding_config()?;
359    if let Some(fallback) = cache.call_failure_fallback {
360        return embed_with_cached_call_failure_fallback(text, kind, &config, fallback);
361    }
362    match active_provider(&config)? {
363        ActiveEmbeddingProvider::Local => local_semantic::embed_text(text, &config, kind),
364        ActiveEmbeddingProvider::FeatureHash => {
365            TextEmbedding::new(FEATURE_HASH_EMBEDDING_MODEL, embed_text_local(text))
366        }
367        ActiveEmbeddingProvider::OpenAi { api_key } => embed_openai(text, &config, &api_key)
368            .or_else(|error| embed_with_call_failure_fallback(text, kind, &config, error, cache)),
369        ActiveEmbeddingProvider::Off => Err(status::embedding_provider_off_error()),
370    }
371}
372
373fn embed_with_cached_call_failure_fallback(
374    text: &str,
375    kind: LocalEmbeddingInputKind,
376    config: &EmbeddingConfig,
377    fallback: EmbeddingProvider,
378) -> Result<TextEmbedding> {
379    let fallback_runtime = status::provider_runtime(config, fallback);
380    if let Some(reason) = fallback_runtime.unavailable_reason {
381        bail!(
382            "cached embedding fallback {} unavailable: {reason}",
383            fallback.label()
384        );
385    }
386    match fallback_runtime.provider {
387        EmbeddingProvider::Local => local_semantic::embed_text(text, config, kind),
388        EmbeddingProvider::FeatureHash => {
389            TextEmbedding::new(FEATURE_HASH_EMBEDDING_MODEL, embed_text_local(text))
390        }
391        EmbeddingProvider::Off => Err(status::embedding_provider_off_error()),
392        EmbeddingProvider::OpenAi | EmbeddingProvider::Auto => {
393            bail!("cached embedding fallback must be local, feature-hash, or off")
394        }
395    }
396}
397
398fn embed_with_call_failure_fallback(
399    text: &str,
400    kind: LocalEmbeddingInputKind,
401    config: &EmbeddingConfig,
402    error: anyhow::Error,
403    cache: &mut EmbeddingFallbackCache,
404) -> Result<TextEmbedding> {
405    let Some(fallback) = config.fallback else {
406        return Err(error);
407    };
408    let fallback_runtime = status::provider_runtime(config, fallback);
409    if let Some(reason) = fallback_runtime.unavailable_reason {
410        bail!(
411            "embedding provider api failed: {error}; fallback {} unavailable: {reason}",
412            fallback.label()
413        );
414    }
415    let message = format!(
416        "configured embedding provider api failed: {}; using fallback {}",
417        error,
418        fallback.label()
419    );
420    crate::log::error("embedding", &message);
421    match fallback_runtime.provider {
422        EmbeddingProvider::Local => {
423            let embedding = local_semantic::embed_text(text, config, kind)?;
424            cache.call_failure_fallback = Some(fallback_runtime.provider);
425            cache.call_failure_fallback_target = Some(EmbeddingBackfillTarget {
426                model: embedding.model().to_string(),
427                dimensions: embedding.dimensions(),
428            });
429            Ok(embedding)
430        }
431        EmbeddingProvider::FeatureHash => {
432            cache.call_failure_fallback = Some(fallback_runtime.provider);
433            cache.call_failure_fallback_target = Some(EmbeddingBackfillTarget {
434                model: FEATURE_HASH_EMBEDDING_MODEL.to_string(),
435                dimensions: FEATURE_HASH_EMBEDDING_DIMENSIONS,
436            });
437            TextEmbedding::new(FEATURE_HASH_EMBEDDING_MODEL, embed_text_local(text))
438        }
439        EmbeddingProvider::Off => Err(status::embedding_provider_off_error_with_cause(format!(
440            "embedding provider api failed: {error}; fallback off disabled provider fallback"
441        ))),
442        EmbeddingProvider::OpenAi | EmbeddingProvider::Auto => Err(error),
443    }
444}
445
446fn memory_embedding_text(
447    title: &str,
448    content: &str,
449    memory_type: &str,
450    topic_key: Option<&str>,
451) -> String {
452    let mut text = String::new();
453    text.push_str(memory_type);
454    text.push('\n');
455    if let Some(topic_key) = topic_key {
456        text.push_str(topic_key);
457        text.push('\n');
458    }
459    text.push_str(title);
460    text.push('\n');
461    text.push_str(content);
462    text
463}
464
465#[derive(Debug, Clone, PartialEq, Eq)]
466enum ActiveEmbeddingProvider {
467    Local,
468    FeatureHash,
469    OpenAi { api_key: String },
470    Off,
471}
472
473fn active_provider(config: &EmbeddingConfig) -> Result<ActiveEmbeddingProvider> {
474    let status = status::resolve_provider_status(config);
475    if status.active_provider == EmbeddingProvider::Off.label() && status.degraded {
476        return Err(status::embedding_provider_off_error_with_cause(
477            status
478                .degradation_reason
479                .unwrap_or_else(|| "embedding provider fallback is off".to_string()),
480        ));
481    }
482    if let Some(reason) = status.unavailable_reason {
483        if status.active_provider == EmbeddingProvider::Off.label() {
484            return Err(status::embedding_provider_off_error_with_cause(reason));
485        }
486        if status.active_provider == EmbeddingProvider::Local.label() {
487            return Err(local_semantic::model_unavailable_error(reason));
488        }
489        bail!("{reason}");
490    }
491    match EmbeddingProvider::parse(&status.active_provider)? {
492        EmbeddingProvider::Local => Ok(ActiveEmbeddingProvider::Local),
493        EmbeddingProvider::FeatureHash => Ok(ActiveEmbeddingProvider::FeatureHash),
494        EmbeddingProvider::OpenAi => Ok(ActiveEmbeddingProvider::OpenAi {
495            api_key: configured_api_key(config)?.with_context(|| {
496                format!(
497                    "embedding provider api requires {ENV_API_KEY} or {}",
498                    config.api_key_env
499                )
500            })?,
501        }),
502        EmbeddingProvider::Off => Ok(ActiveEmbeddingProvider::Off),
503        EmbeddingProvider::Auto => bail!("auto must resolve to a concrete embedding provider"),
504    }
505}
506
507fn auto_api_key(config: &EmbeddingConfig) -> Result<Option<String>> {
508    if let Some(value) = env_value(ENV_API_KEY).or_else(|| env_value(ENV_API_KEY_LEGACY)) {
509        return Ok(Some(value));
510    }
511    if config.api_key_env != DEFAULT_API_KEY_ENV {
512        configured_api_key(config)
513    } else {
514        Ok(None)
515    }
516}
517
518fn configured_api_key(config: &EmbeddingConfig) -> Result<Option<String>> {
519    if let Some(value) = env_value(ENV_API_KEY).or_else(|| env_value(ENV_API_KEY_LEGACY)) {
520        return Ok(Some(value));
521    }
522    Ok(std::env::var(&config.api_key_env)
523        .ok()
524        .map(|value| value.trim().to_string())
525        .filter(|value| !value.is_empty()))
526}
527
528#[derive(Debug, Serialize)]
529struct OpenAiEmbeddingRequest<'a> {
530    input: &'a str,
531    model: &'a str,
532    encoding_format: &'static str,
533    #[serde(skip_serializing_if = "Option::is_none")]
534    dimensions: Option<usize>,
535}
536
537#[derive(Debug, Deserialize)]
538struct OpenAiEmbeddingResponse {
539    data: Vec<OpenAiEmbeddingData>,
540    model: Option<String>,
541}
542
543#[derive(Debug, Deserialize)]
544struct OpenAiEmbeddingData {
545    embedding: Vec<f32>,
546}
547
548fn embed_openai(text: &str, config: &EmbeddingConfig, api_key: &str) -> Result<TextEmbedding> {
549    if text.trim().is_empty() {
550        bail!("embedding input must not be empty");
551    }
552    let client = reqwest::blocking::Client::builder()
553        .timeout(Duration::from_secs(config.timeout_secs))
554        .build()
555        .context("build embedding HTTP client")?;
556    let request = OpenAiEmbeddingRequest {
557        input: text,
558        model: &config.model,
559        encoding_format: "float",
560        dimensions: config.dimensions,
561    };
562    let url = format!("{}/embeddings", config.base_url.trim_end_matches('/'));
563    let response = client
564        .post(&url)
565        .bearer_auth(api_key)
566        .json(&request)
567        .send()
568        .with_context(|| format!("call embedding provider at {url}"))?;
569    let status = response.status();
570    let body = response
571        .text()
572        .context("read embedding provider response body")?;
573    if !status.is_success() {
574        bail!(
575            "embedding provider returned HTTP {status}: {}",
576            truncate_error_body(&body)
577        );
578    }
579    parse_openai_embedding_response(&body, &config.model)
580}
581
582fn parse_openai_embedding_response(body: &str, fallback_model: &str) -> Result<TextEmbedding> {
583    let response: OpenAiEmbeddingResponse =
584        serde_json::from_str(body).context("parse embedding provider response")?;
585    let mut data = response.data.into_iter();
586    let first = data
587        .next()
588        .context("embedding provider response did not include data[0]")?;
589    if data.next().is_some() {
590        bail!("embedding provider returned multiple embeddings for single input");
591    }
592    TextEmbedding::new(
593        response.model.unwrap_or_else(|| fallback_model.to_string()),
594        first.embedding,
595    )
596}
597
598fn truncate_error_body(body: &str) -> String {
599    const MAX: usize = 500;
600    if body.len() <= MAX {
601        body.to_string()
602    } else {
603        let mut end = MAX;
604        while !body.is_char_boundary(end) {
605            end -= 1;
606        }
607        format!("{}...", &body[..end])
608    }
609}
610
611fn validate_embedding_values(values: &[f32]) -> Result<()> {
612    if values.is_empty() {
613        bail!("embedding vector must not be empty");
614    }
615    if values.iter().any(|value| !value.is_finite()) {
616        bail!("embedding vector contains non-finite values");
617    }
618    Ok(())
619}
620
621fn embed_text_local(text: &str) -> Vec<f32> {
622    let normalized = text.to_lowercase();
623    let mut vector = vec![0.0f32; LOCAL_EMBEDDING_DIMENSIONS];
624    for token in semantic_tokens(&normalized) {
625        add_feature(&mut vector, &format!("token:{token}"), 1.0);
626    }
627    for ngram in char_ngrams(&normalized) {
628        add_feature(&mut vector, &format!("ngram:{ngram}"), 0.35);
629    }
630    for (concept, phrases) in semantic_concepts() {
631        if phrases.iter().any(|phrase| normalized.contains(phrase)) {
632            add_feature(&mut vector, &format!("concept:{concept}"), 4.0);
633        }
634    }
635    normalize(&mut vector);
636    vector
637}
638
639fn semantic_tokens(text: &str) -> Vec<String> {
640    let mut tokens = Vec::new();
641    let mut current = String::new();
642    for ch in text.chars() {
643        if ch.is_alphanumeric() || is_cjk(ch) {
644            current.push(ch);
645        } else if !current.is_empty() {
646            tokens.push(std::mem::take(&mut current));
647        }
648    }
649    if !current.is_empty() {
650        tokens.push(current);
651    }
652    tokens
653}
654
655fn char_ngrams(text: &str) -> Vec<String> {
656    let chars: Vec<char> = text
657        .chars()
658        .filter(|ch| ch.is_alphanumeric() || is_cjk(*ch))
659        .collect();
660    let mut grams = Vec::new();
661    for width in [2usize, 3] {
662        if chars.len() < width {
663            continue;
664        }
665        grams.extend(
666            chars
667                .windows(width)
668                .map(|window| window.iter().collect::<String>()),
669        );
670    }
671    grams
672}
673
674fn add_feature(vector: &mut [f32], feature: &str, weight: f32) {
675    let digest = Sha256::digest(feature.as_bytes());
676    for offset in [0usize, 8, 16] {
677        let raw = u64::from_le_bytes([
678            digest[offset],
679            digest[offset + 1],
680            digest[offset + 2],
681            digest[offset + 3],
682            digest[offset + 4],
683            digest[offset + 5],
684            digest[offset + 6],
685            digest[offset + 7],
686        ]);
687        let idx = raw as usize % vector.len();
688        let sign = if raw & 1 == 0 { 1.0 } else { -1.0 };
689        vector[idx] += weight * sign;
690    }
691}
692
693fn normalize(vector: &mut [f32]) {
694    let norm = vector.iter().map(|value| value * value).sum::<f32>().sqrt();
695    if norm == 0.0 {
696        return;
697    }
698    for value in vector {
699        *value /= norm;
700    }
701}
702
703fn is_cjk(ch: char) -> bool {
704    matches!(
705        ch,
706        '\u{4E00}'..='\u{9FFF}' |
707        '\u{3400}'..='\u{4DBF}' |
708        '\u{F900}'..='\u{FAFF}'
709    )
710}
711
712fn semantic_concepts() -> &'static [(&'static str, &'static [&'static str])] {
713    &[
714        (
715            "data-security",
716            &[
717                "sqlcipher",
718                "encrypt",
719                "encrypted",
720                "encryption",
721                "secret",
722                "secrets",
723                "credential",
724                "credentials",
725                "private",
726                "confidential",
727                "protect",
728                "protected",
729                "at rest",
730                "persisted data",
731                "加密",
732                "密钥",
733            ],
734        ),
735        (
736            "transcript-capture",
737            &[
738                "transcript",
739                "raw archive",
740                "raw message",
741                "hook fallback",
742                "assistant message",
743                "conversation capture",
744                "jsonl",
745                "会话",
746                "原始消息",
747            ],
748        ),
749        (
750            "retrieval-quality",
751            &[
752                "semantic",
753                "embedding",
754                "vector",
755                "recall",
756                "search quality",
757                "paraphrase",
758                "检索",
759                "语义",
760                "召回",
761                "向量",
762            ],
763        ),
764        (
765            "current-state",
766            &[
767                "current decision",
768                "current state",
769                "supersede",
770                "supersedes",
771                "stale",
772                "replacement",
773                "现在",
774                "当前",
775                "替代",
776            ],
777        ),
778        (
779            "compression",
780            &[
781                "compress",
782                "compression",
783                "compaction",
784                "summarize",
785                "compressed",
786                "压缩",
787                "摘要",
788                "总结",
789            ],
790        ),
791    ]
792}
793
794#[cfg(test)]
795mod tests;
remem/retrieval/embedding.rs

remem/retrieval/
embedding.rs
