sqlite_graphrag/embedder/
getters.rs1use super::*;
4use crate::errors::AppError;
5use crate::extract::llm_embedding::LlmEmbedding;
6use parking_lot::Mutex;
7use std::path::Path;
8
9pub fn is_openrouter_initialized() -> bool {
11 OPENROUTER_CLIENT.get().is_some()
12}
13
14pub(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
32pub 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
44pub 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
69pub 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
94pub 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
110pub 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
129pub fn openrouter_chat_client() -> Option<&'static crate::chat_api::OpenRouterChatClient> {
134 OPENROUTER_CHAT_CLIENT.get()
135}
136
137pub 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
151pub 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
167pub 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}
182pub(crate) fn clone_client(embedder: &Mutex<LlmEmbedding>) -> LlmEmbedding {
185 embedder.lock().clone()
186}
187
188thread_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}