Skip to main content

sqlite_graphrag/embedder/
getters.rs

1//! Embedder client getters and local-backend helpers.
2
3use super::*;
4use crate::errors::AppError;
5use crate::extract::llm_embedding::LlmEmbedding;
6use parking_lot::Mutex;
7use std::path::Path;
8
9/// Returns true when the process-wide OpenRouter embed client is ready.
10pub fn is_openrouter_initialized() -> bool {
11    OPENROUTER_CLIENT.get().is_some()
12}
13
14/// Returns the process-wide multi-thread runtime, building it on first use.
15pub(crate) fn shared_runtime() -> Result<&'static tokio::runtime::Runtime, AppError> {
16    if let Some(rt) = RUNTIME.get() {
17        return Ok(rt);
18    }
19    let rt = tokio::runtime::Builder::new_multi_thread()
20        .worker_threads(2)
21        .enable_all()
22        .build()
23        .map_err(|e| {
24            AppError::Embedding(crate::i18n::validation::embedding_tokio_runtime_init_failed(e))
25        })?;
26    let _ = RUNTIME.set(rt);
27    RUNTIME.get().ok_or_else(|| {
28        AppError::Embedding(crate::i18n::validation::embedding_tokio_runtime_unavailable())
29    })
30}
31
32/// Initialises the LLM-embedding client on first use and returns it.
33pub fn get_embedder(_models_dir: &Path) -> Result<&'static Mutex<LlmEmbedding>, AppError> {
34    if let Some(e) = EMBEDDER.get() {
35        return Ok(e);
36    }
37    let backend = LlmEmbedding::detect_available()?;
38    let _ = EMBEDDER.set(Mutex::new(backend));
39    EMBEDDER.get().ok_or_else(|| {
40        AppError::Embedding(crate::i18n::validation::embedding_embedder_unavailable())
41    })
42}
43
44/// ADR-0042 / GAP-002: returns the process-wide Claude embedder, lazily
45/// initialising it on first use. Binary and model overrides come from
46/// the explicit arguments; `None` falls back to PATH/env defaults via
47/// the builder.
48pub fn get_claude_embedder(
49    claude_binary: Option<&Path>,
50    claude_model: Option<&str>,
51) -> Result<&'static Mutex<LlmEmbedding>, AppError> {
52    if let Some(e) = CLAUDE_EMBEDDER.get() {
53        return Ok(e);
54    }
55    let mut builder = LlmEmbedding::with_claude_builder();
56    if let Some(b) = claude_binary {
57        builder = builder.override_binary(b.to_path_buf());
58    }
59    if let Some(m) = claude_model {
60        builder = builder.override_model(m.to_string());
61    }
62    let backend = builder.build()?;
63    let _ = CLAUDE_EMBEDDER.set(Mutex::new(backend));
64    CLAUDE_EMBEDDER.get().ok_or_else(|| {
65        AppError::Embedding(crate::i18n::validation::embedding_claude_embedder_unavailable())
66    })
67}
68
69/// GAP-OPENCODE-001 / v1.0.90: returns the process-wide OpenCode embedder,
70/// lazily initialising it on first use. Binary and model overrides come
71/// from the explicit arguments; `None` falls back to PATH/env defaults via
72/// the builder.
73pub fn get_opencode_embedder(
74    opencode_binary: Option<&Path>,
75    opencode_model: Option<&str>,
76) -> Result<&'static Mutex<LlmEmbedding>, AppError> {
77    if let Some(e) = OPENCODE_EMBEDDER.get() {
78        return Ok(e);
79    }
80    let mut builder = LlmEmbedding::with_opencode_builder();
81    if let Some(b) = opencode_binary {
82        builder = builder.override_binary(b.to_path_buf());
83    }
84    if let Some(m) = opencode_model {
85        builder = builder.override_model(m.to_string());
86    }
87    let backend = builder.build()?;
88    let _ = OPENCODE_EMBEDDER.set(Mutex::new(backend));
89    OPENCODE_EMBEDDER.get().ok_or_else(|| {
90        AppError::Embedding(crate::i18n::validation::embedding_opencode_embedder_unavailable())
91    })
92}
93
94/// Get openrouter embedder.
95pub fn get_openrouter_embedder(
96    api_key: secrecy::SecretBox<String>,
97    model: &str,
98    dim: usize,
99) -> Result<&'static crate::embedding_api::OpenRouterClient, AppError> {
100    if let Some(c) = OPENROUTER_CLIENT.get() {
101        return Ok(c);
102    }
103    let client = crate::embedding_api::OpenRouterClient::new(api_key, model.to_string(), dim)?;
104    let _ = OPENROUTER_CLIENT.set(client);
105    OPENROUTER_CLIENT.get().ok_or_else(|| {
106        AppError::Embedding(crate::i18n::validation::embedding_openrouter_client_unavailable())
107    })
108}
109
110/// v1.0.95 (ADR-0054): initialises the process-wide OpenRouter chat client on
111/// first use and returns it. `model` is the text model the enrich JUDGE will
112/// call (no default; the caller validates presence upfront).
113pub fn get_openrouter_chat_client(
114    api_key: secrecy::SecretBox<String>,
115    model: &str,
116    timeout_secs: u64,
117) -> Result<&'static crate::chat_api::OpenRouterChatClient, AppError> {
118    if let Some(c) = OPENROUTER_CHAT_CLIENT.get() {
119        return Ok(c);
120    }
121    let client =
122        crate::chat_api::OpenRouterChatClient::new(api_key, model.to_string(), timeout_secs)?;
123    let _ = OPENROUTER_CHAT_CLIENT.set(client);
124    OPENROUTER_CHAT_CLIENT.get().ok_or_else(|| {
125        AppError::Embedding(crate::i18n::validation::embedding_openrouter_chat_client_unavailable())
126    })
127}
128
129/// v1.0.95: returns the process-wide OpenRouter chat client if it has already
130/// been initialised via [`get_openrouter_chat_client`]. Used by the enrich
131/// JUDGE dispatch, which initialises the singleton once at startup and then
132/// fetches it per item without re-threading the API key.
133pub fn openrouter_chat_client() -> Option<&'static crate::chat_api::OpenRouterChatClient> {
134    OPENROUTER_CHAT_CLIENT.get()
135}
136
137/// ADR-0042 / GAP-002: route a single passage through the Claude
138/// embedder. Used by the Claude arm of `embed_via_backend` so the
139/// fallback chain stops treating Claude as a synonym for codex.
140pub fn embed_via_claude_local(
141    _models_dir: &Path,
142    text: &str,
143    claude_binary: Option<&Path>,
144    claude_model: Option<&str>,
145) -> Result<Vec<f32>, AppError> {
146    let _slot_guard = acquire_llm_slot_for_embedding()?;
147    let embedder = get_claude_embedder(claude_binary, claude_model)?;
148    embed_passage(embedder, text)
149}
150
151/// BUG-003 / v1.0.85: split of  that also
152/// reports the resolved []. Always  because
153/// this path constructs a Claude-flavoured embedder via
154///  (no PATH probe, no silent substitution).
155pub fn embed_via_claude_local_resolved(
156    _models_dir: &Path,
157    text: &str,
158    claude_binary: Option<&Path>,
159    claude_model: Option<&str>,
160) -> Result<(Vec<f32>, LlmBackendKind), AppError> {
161    let _slot_guard = acquire_llm_slot_for_embedding()?;
162    let embedder = get_claude_embedder(claude_binary, claude_model)?;
163    let v = embed_passage(embedder, text)?;
164    Ok((v, LlmBackendKind::Claude))
165}
166
167/// GAP-OPENCODE-001 / v1.0.90: route a single passage through the OpenCode
168/// embedder, reporting the resolved [`LlmBackendKind::Opencode`]. Constructs
169/// an OpenCode-flavoured embedder via `with_opencode_builder` (no PATH probe,
170/// no silent substitution).
171pub fn embed_via_opencode_local_resolved(
172    _models_dir: &Path,
173    text: &str,
174    opencode_binary: Option<&Path>,
175    opencode_model: Option<&str>,
176) -> Result<(Vec<f32>, LlmBackendKind), AppError> {
177    let _slot_guard = acquire_llm_slot_for_embedding()?;
178    let embedder = get_opencode_embedder(opencode_binary, opencode_model)?;
179    let v = embed_passage(embedder, text)?;
180    Ok((v, LlmBackendKind::Opencode))
181}
182/// Clones the embedding-client configuration. The lock is held only for
183/// the duration of the clone — NEVER across I/O (G42/A3).
184pub(crate) fn clone_client(embedder: &Mutex<LlmEmbedding>) -> LlmEmbedding {
185    embedder.lock().clone()
186}
187
188// When true, embed_passage/embed_query use the short query timeout so Auto
189// chains fail fast into FTS (GAP-E2E-06).
190thread_local! {
191    static QUERY_EMBED_FAST: std::cell::Cell<bool> = const { std::cell::Cell::new(false) };
192}
193
194pub(crate) fn with_query_embed_fast<T>(f: impl FnOnce() -> T) -> T {
195    QUERY_EMBED_FAST.with(|c| {
196        let prev = c.replace(true);
197        let out = f();
198        c.set(prev);
199        out
200    })
201}
202
203pub(crate) fn apply_query_timeout_if_needed(client: LlmEmbedding) -> LlmEmbedding {
204    if QUERY_EMBED_FAST.with(|c| c.get()) {
205        let secs = crate::runtime_config::resolve_u64(
206            None,
207            "llm.query_embed_timeout_secs",
208            crate::constants::DEFAULT_QUERY_EMBED_TIMEOUT_SECS,
209        );
210        client.with_timeout_secs(secs)
211    } else {
212        client
213    }
214}