1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
//! Shared embedding engine for the memory + documents MCP tools.
use std::sync::OnceLock;
use anyhow::{Context, Result, anyhow};
use xberg::embeddings::EMBEDDING_PRESETS;
use xberg::{EmbeddingConfig, EmbeddingModelType};
/// Global bounded rayon `ThreadPool` for all ONNX embed calls. Initialized once
/// on first use; subsequent calls to `embed_pool` return the same pool regardless
/// of the `max_threads` argument (the pool size is fixed for the process).
static EMBED_POOL: OnceLock<rayon::ThreadPool> = OnceLock::new();
/// Resolve the embedding thread cap.
///
/// `0` is the sentinel for "auto": `max(2, logical_cpus / 4)` — a bounded
/// fraction of available cores that leaves the full global rayon pool free for
/// code-map extraction and prevents the embedder from pinning all cores.
/// Any non-zero value is used directly.
pub fn resolve_embed_threads(max_threads: usize) -> usize {
if max_threads == 0 {
std::cmp::max(2, rayon::current_num_threads() / 4)
} else {
max_threads
}
}
/// Returns the process-wide bounded rayon `ThreadPool` for ONNX embedding.
///
/// Initialized once on first call with `resolve_embed_threads(max_threads)`.
/// All `embed_texts` calls from [`SharedEmbedder`] run inside this pool via
/// `pool.install(...)`, which constrains xberg's internal rayon tasks (including
/// the per-chunk embedding fan-out) to at most `current_num_threads()` workers.
pub fn embed_pool(max_threads: usize) -> &'static rayon::ThreadPool {
EMBED_POOL.get_or_init(|| {
let n = resolve_embed_threads(max_threads);
rayon::ThreadPoolBuilder::new()
.num_threads(n)
.thread_name(|i| format!("bm-embed-{i}"))
.build()
.expect("failed to build embedding rayon pool")
})
}
/// Loaded, ready-to-query embedding engine. `Clone` is cheap (config is stack-only).
#[derive(Clone)]
pub struct SharedEmbedder {
config: EmbeddingConfig,
dim: u16,
model_name: String,
/// Resolved embed-thread cap passed through to `embed_pool`. `0` triggers the
/// auto heuristic inside `resolve_embed_threads`; first caller wins for the
/// global pool.
max_embed_threads: usize,
}
impl SharedEmbedder {
/// Build a `SharedEmbedder` from a named xberg preset.
///
/// `max_embed_threads` bounds the process-wide embedding pool (see
/// [`embed_pool`]). Pass `0` to use the auto heuristic (`max(2, cores/4)`).
///
/// `batch_size` is the number of texts submitted to ONNX per embed call
/// (sourced from `[resources].embed_batch_size`). Larger batches amortise
/// per-call overhead at a higher transient memory spike.
pub fn load(preset: &str, max_embed_threads: usize, batch_size: usize) -> Result<Self> {
let meta = EMBEDDING_PRESETS.iter().find(|p| p.name == preset).ok_or_else(|| {
anyhow!(
"unknown embedding preset '{preset}'; \
available: fast, balanced, quality, multilingual"
)
})?;
let dim = u16::try_from(meta.dimensions)
.with_context(|| format!("preset '{preset}' dimension {} exceeds u16", meta.dimensions))?;
let config = EmbeddingConfig {
model: EmbeddingModelType::Preset {
name: preset.to_string(),
},
normalize: true,
batch_size,
show_download_progress: false,
cache_dir: None,
acceleration: None,
max_embed_duration_secs: Some(60),
max_sequence_length: None,
};
Ok(Self {
config,
dim,
model_name: preset.to_string(),
max_embed_threads,
})
}
/// Vector dimension produced by this embedder.
pub fn dim(&self) -> u16 {
self.dim
}
/// The preset name (e.g. `"balanced"`).
pub fn model(&self) -> &str {
&self.model_name
}
/// Embed a single text string. Returns a `Vec<f32>` of length `self.dim()`.
///
/// The call is routed through the process-wide bounded [`embed_pool`] so it
/// cannot saturate the global rayon pool used by the code-map scanner.
pub fn embed(&self, text: &str) -> Result<Vec<f32>> {
if text.is_empty() {
return Err(anyhow!("embed: input text must not be empty"));
}
embed_pool(self.max_embed_threads).install(|| {
let mut results = xberg::embeddings::embed_texts(&[text], &self.config)
.with_context(|| format!("embed_texts(preset={})", self.model_name))?;
results
.pop()
.ok_or_else(|| anyhow!("embed_texts returned empty result"))
})
}
/// Embed a batch of texts in one call. Returns one `Vec<f32>` of length `self.dim()` per
/// input, in order. Used by the code-search scanner to embed a file's chunks in bulk.
///
/// Errors if any input text is empty (xberg rejects empty strings, which produce meaningless
/// embeddings). An empty batch returns an empty vector without touching the model.
///
/// The call is routed through the process-wide bounded [`embed_pool`] so it
/// cannot saturate the global rayon pool used by the code-map scanner.
#[cfg(feature = "code-search")]
pub fn embed_batch(&self, texts: &[&str]) -> Result<Vec<Vec<f32>>> {
if texts.is_empty() {
return Ok(Vec::new());
}
embed_pool(self.max_embed_threads).install(|| {
xberg::embeddings::embed_texts(texts, &self.config)
.with_context(|| format!("embed_texts(preset={}, batch={})", self.model_name, texts.len()))
})
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn resolve_embed_threads_nonzero_passthrough() {
assert_eq!(resolve_embed_threads(4), 4);
assert_eq!(resolve_embed_threads(1), 1);
assert_eq!(resolve_embed_threads(16), 16);
}
#[test]
fn resolve_embed_threads_zero_gives_auto() {
let got = resolve_embed_threads(0);
let expected = std::cmp::max(2, rayon::current_num_threads() / 4);
assert_eq!(
got, expected,
"resolve_embed_threads(0) should yield max(2, cores/4) = {expected}"
);
assert!(got >= 2, "auto embed cap must be >= 2, got {got}");
}
}