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
40 .get()
41 .ok_or_else(|| {
42 AppError::Embedding(crate::i18n::validation::embedding_embedder_unavailable())
43 })
44}
45
46pub 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
71pub 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
96pub 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
112pub 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
133pub fn openrouter_chat_client() -> Option<&'static crate::chat_api::OpenRouterChatClient> {
138 OPENROUTER_CHAT_CLIENT.get()
139}
140
141pub 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
155pub 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
171pub 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}
186pub(crate) fn clone_client(embedder: &Mutex<LlmEmbedding>) -> LlmEmbedding {
189 embedder.lock().clone()
190}
191
192thread_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}