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