sqlite_graphrag/embedder/
getters.rs

1//! Embedder client getters and local-backend helpers.
2
3use super::*;
4use crate::constants::{EMBED_RUNTIME_MAX_WORKER_THREADS, EMBED_RUNTIME_MIN_WORKER_THREADS};
5use crate::errors::AppError;
6
7/// Returns true when the process-wide OpenRouter embed client is ready.
8pub fn is_openrouter_initialized() -> bool {
9    OPENROUTER_CLIENT.get().is_some()
10}
11
12/// Host-derived worker count for the shared embedding runtime, before the XDG
13/// override is applied (GAP-SG-141 B2).
14///
15/// Clamped between [`EMBED_RUNTIME_MIN_WORKER_THREADS`] and
16/// [`EMBED_RUNTIME_MAX_WORKER_THREADS`]; a host that cannot report its
17/// parallelism falls back to the minimum.
18fn default_embed_runtime_threads() -> usize {
19    std::thread::available_parallelism()
20        .map(|n| n.get())
21        .unwrap_or(EMBED_RUNTIME_MIN_WORKER_THREADS)
22        .clamp(
23            EMBED_RUNTIME_MIN_WORKER_THREADS,
24            EMBED_RUNTIME_MAX_WORKER_THREADS,
25        )
26}
27
28/// Returns the process-wide multi-thread runtime, building it on first use.
29///
30/// The worker count is never a literal: it comes from
31/// [`crate::runtime_config::embed_runtime_worker_threads`], which layers XDG
32/// `parallelism.embed_runtime_threads` over [`default_embed_runtime_threads`].
33/// A hard-coded two workers starved the reactor once the enrich drain began
34/// issuing up to sixteen concurrent blocking calls against it.
35pub(crate) fn shared_runtime() -> Result<&'static tokio::runtime::Runtime, AppError> {
36    if let Some(rt) = RUNTIME.get() {
37        return Ok(rt);
38    }
39    let workers =
40        crate::runtime_config::embed_runtime_worker_threads(default_embed_runtime_threads())
41            .max(EMBED_RUNTIME_MIN_WORKER_THREADS);
42    let rt = tokio::runtime::Builder::new_multi_thread()
43        .worker_threads(workers)
44        .enable_all()
45        .build()
46        .map_err(|e| {
47            AppError::Embedding(crate::i18n::validation::embedding_tokio_runtime_init_failed(e))
48        })?;
49    let _ = RUNTIME.set(rt);
50    RUNTIME.get().ok_or_else(|| {
51        AppError::Embedding(crate::i18n::validation::embedding_tokio_runtime_unavailable())
52    })
53}
54
55/// Initialises the process-wide OpenRouter embedding client on first use and
56/// returns it.
57///
58/// The per-request timeout resolves in the documented precedence:
59/// `timeout_override` (the `--openrouter-timeout` flag) first, then XDG
60/// `embedding.timeout_secs`, then the client's own default.
61///
62/// FIRST INITIALISER WINS. The client lives in a `OnceLock`, so a later call
63/// with a different timeout returns the already-built client unchanged. This is
64/// sound under the one-shot CLI contract: a single invocation runs a single
65/// subcommand, so exactly one timeout is in play per process. Nothing here
66/// attempts to rebuild or swap the client, which would race with in-flight
67/// requests for no benefit.
68pub fn get_openrouter_embedder(
69    api_key: secrecy::SecretBox<String>,
70    model: &str,
71    dim: usize,
72    timeout_override: Option<u64>,
73) -> Result<&'static crate::embedding_api::OpenRouterClient, AppError> {
74    if let Some(c) = OPENROUTER_CLIENT.get() {
75        return Ok(c);
76    }
77    let timeout_secs = crate::runtime_config::resolve_u64(
78        timeout_override,
79        "embedding.timeout_secs",
80        crate::constants::DEFAULT_EMBEDDING_HTTP_TIMEOUT_SECS,
81    );
82    let client =
83        crate::embedding_api::OpenRouterClient::new(api_key, model.to_string(), dim, timeout_secs)?;
84    let _ = OPENROUTER_CLIENT.set(client);
85    OPENROUTER_CLIENT.get().ok_or_else(|| {
86        AppError::Embedding(crate::i18n::validation::embedding_openrouter_client_unavailable())
87    })
88}
89
90/// v1.0.95 (ADR-0054): initialises the process-wide OpenRouter chat client on
91/// first use and returns it. `model` is the text model the enrich JUDGE will
92/// call (no default; the caller validates presence upfront).
93pub fn get_openrouter_chat_client(
94    api_key: secrecy::SecretBox<String>,
95    model: &str,
96    timeout_secs: u64,
97) -> Result<&'static crate::chat_api::OpenRouterChatClient, AppError> {
98    if let Some(c) = OPENROUTER_CHAT_CLIENT.get() {
99        return Ok(c);
100    }
101    let client =
102        crate::chat_api::OpenRouterChatClient::new(api_key, model.to_string(), timeout_secs)?;
103    let _ = OPENROUTER_CHAT_CLIENT.set(client);
104    OPENROUTER_CHAT_CLIENT.get().ok_or_else(|| {
105        AppError::Embedding(crate::i18n::validation::embedding_openrouter_chat_client_unavailable())
106    })
107}
108
109/// v1.0.95: returns the process-wide OpenRouter chat client if it has already
110/// been initialised via [`get_openrouter_chat_client`]. Used by the enrich
111/// JUDGE dispatch, which initialises the singleton once at startup and then
112/// fetches it per item without re-threading the API key.
113pub fn openrouter_chat_client() -> Option<&'static crate::chat_api::OpenRouterChatClient> {
114    OPENROUTER_CHAT_CLIENT.get()
115}
116
117#[cfg(test)]
118mod runtime_sizing_tests {
119    use super::*;
120
121    /// Collects every `.rs` file under `dir`, recursively.
122    fn rust_sources(dir: &std::path::Path, out: &mut Vec<std::path::PathBuf>) {
123        let entries = match std::fs::read_dir(dir) {
124            Ok(e) => e,
125            Err(_) => return,
126        };
127        for entry in entries.flatten() {
128            let path = entry.path();
129            if path.is_dir() {
130                rust_sources(&path, out);
131            } else if path.extension().is_some_and(|e| e == "rs") {
132                out.push(path);
133            }
134        }
135    }
136
137    /// Drops whole-line `//` comments so prose quoting the banned call shape —
138    /// including the comment right below — is not read as a call site.
139    fn strip_line_comments(source: &str) -> String {
140        source
141            .lines()
142            .filter(|line| !line.trim_start().starts_with("//"))
143            .collect::<Vec<_>>()
144            .join("\n")
145    }
146
147    /// No runtime in the CRATE may fix its reactor width to a literal.
148    ///
149    /// SCOPE IS THE CRATE, never this one file. The previous version read its
150    /// own source through `include_str!`, so it could only ever police the
151    /// runtime built a few lines above it. `src/commands/deep_research` built a
152    /// second runtime with `.worker_threads(2)` and the guard stayed green
153    /// throughout — a one-file guard against a crate-wide invariant is not a
154    /// guard. Every `.rs` file under `src/` is walked instead; files whose name
155    /// carries `test` are fixtures and assertions, not runtime construction.
156    ///
157    /// The historical bug was `.worker_threads(2)`: a fixed reactor width that
158    /// the enrich drain then oversubscribed. Any literal digit is that defect
159    /// returning, wherever it is written.
160    #[test]
161    fn runtime_worker_count_is_never_a_literal() {
162        let root = std::path::Path::new(env!("CARGO_MANIFEST_DIR")).join("src");
163        let mut files = Vec::new();
164        rust_sources(&root, &mut files);
165        assert!(
166            files.len() > 50,
167            "source walk found only {} files under {}; the guard would pass vacuously",
168            files.len(),
169            root.display()
170        );
171
172        let mut offenders: Vec<String> = Vec::new();
173        for path in &files {
174            let is_test_file = path
175                .file_name()
176                .map(|n| n.to_string_lossy().contains("test"))
177                .unwrap_or(false);
178            if is_test_file {
179                continue;
180            }
181            let source = match std::fs::read_to_string(path) {
182                Ok(s) => s,
183                Err(_) => continue,
184            };
185            let code = strip_line_comments(&source);
186            for call in code.split(".worker_threads(").skip(1) {
187                let arg = call.split(')').next().unwrap_or_default().trim();
188                if arg.chars().next().is_some_and(|c| c.is_ascii_digit()) {
189                    offenders.push(format!(
190                        "{}: .worker_threads({arg})",
191                        path.strip_prefix(&root).unwrap_or(path).display()
192                    ));
193                }
194            }
195        }
196
197        assert!(
198            offenders.is_empty(),
199            "worker_threads must be resolved through runtime_config (or left to \
200             Tokio's own core-count default), never written as a literal:\n{}",
201            offenders.join("\n")
202        );
203    }
204
205    #[test]
206    fn host_default_stays_within_the_named_bounds() {
207        let n = default_embed_runtime_threads();
208        assert!(
209            (EMBED_RUNTIME_MIN_WORKER_THREADS..=EMBED_RUNTIME_MAX_WORKER_THREADS).contains(&n),
210            "host-derived worker count {n} escaped its clamp"
211        );
212    }
213
214    #[test]
215    fn zero_override_falls_back_to_the_default() {
216        // `worker_threads(0)` panics in Tokio, so the reader must never let a
217        // zero through.
218        assert_eq!(
219            crate::runtime_config::embed_runtime_worker_threads(0),
220            0,
221            "the reader returns the caller's default verbatim when no override is set"
222        );
223        // The builder therefore applies its own floor on top.
224        let workers = crate::runtime_config::embed_runtime_worker_threads(0)
225            .max(EMBED_RUNTIME_MIN_WORKER_THREADS);
226        assert!(workers >= EMBED_RUNTIME_MIN_WORKER_THREADS);
227    }
228
229    #[test]
230    fn embed_timeout_flag_outranks_xdg_and_default() {
231        // The flag short-circuits before any XDG lookup, so this holds without
232        // touching the operator's config file.
233        assert_eq!(
234            crate::runtime_config::resolve_u64(
235                Some(77),
236                "embedding.timeout_secs",
237                crate::constants::DEFAULT_EMBEDDING_HTTP_TIMEOUT_SECS,
238            ),
239            77
240        );
241    }
242
243    #[test]
244    fn shared_runtime_builds() {
245        assert!(shared_runtime().is_ok());
246    }
247}
sqlite_graphrag/embedder/getters.rs

sqlite_graphrag/embedder/
getters.rs
