1use crate::errors::AppError;
40use crate::extract::llm_embedding::LlmEmbedding;
41use parking_lot::Mutex;
42use std::path::Path;
43use std::sync::Arc;
44use std::sync::OnceLock;
45use tokio::sync::{mpsc, Semaphore};
46use tokio::task::JoinSet;
47use tokio_util::sync::CancellationToken;
48
49static CLAUDE_EMBEDDER: OnceLock<Mutex<LlmEmbedding>> = OnceLock::new();
59static OPENCODE_EMBEDDER: OnceLock<Mutex<LlmEmbedding>> = OnceLock::new();
60static OPENROUTER_CLIENT: OnceLock<crate::embedding_api::OpenRouterClient> = OnceLock::new();
61
62static OPENROUTER_CHAT_CLIENT: OnceLock<crate::chat_api::OpenRouterChatClient> = OnceLock::new();
66
67pub fn is_openrouter_initialized() -> bool {
69 OPENROUTER_CLIENT.get().is_some()
70}
71static EMBEDDER: OnceLock<Mutex<LlmEmbedding>> = OnceLock::new();
72
73static RUNTIME: OnceLock<tokio::runtime::Runtime> = OnceLock::new();
79
80pub const CHUNK_EMBED_BATCH_SIZE: usize = 8;
84
85pub const ENTITY_EMBED_BATCH_SIZE: usize = 25;
89
90pub const EMBED_BATCH_CALIBRATION_DIM: usize = 64;
92
93fn adaptive_batch_for_dim(base: usize, dim: usize) -> usize {
101 let base = base.max(1);
102 (base * EMBED_BATCH_CALIBRATION_DIM / dim.max(1)).clamp(1, base)
103}
104
105pub fn chunk_embed_batch_size() -> usize {
107 let dim = crate::constants::embedding_dim();
108 let batch = adaptive_batch_for_dim(CHUNK_EMBED_BATCH_SIZE, dim);
109 tracing::debug!(
110 dim,
111 base = CHUNK_EMBED_BATCH_SIZE,
112 batch,
113 "adaptive chunk batch size (G44)"
114 );
115 batch
116}
117
118pub fn entity_embed_batch_size() -> usize {
120 let dim = crate::constants::embedding_dim();
121 let batch = adaptive_batch_for_dim(ENTITY_EMBED_BATCH_SIZE, dim);
122 tracing::debug!(
123 dim,
124 base = ENTITY_EMBED_BATCH_SIZE,
125 batch,
126 "adaptive entity batch size (G44)"
127 );
128 batch
129}
130
131pub(crate) fn shared_runtime() -> Result<&'static tokio::runtime::Runtime, AppError> {
133 if let Some(rt) = RUNTIME.get() {
134 return Ok(rt);
135 }
136 let rt = tokio::runtime::Builder::new_multi_thread()
137 .worker_threads(2)
138 .enable_all()
139 .build()
140 .map_err(|e| AppError::Embedding(format!("tokio runtime init failed: {e}")))?;
141 let _ = RUNTIME.set(rt);
142 RUNTIME.get().ok_or_else(|| {
143 AppError::Embedding("tokio runtime unavailable after initialisation".to_string())
144 })
145}
146
147pub fn get_embedder(_models_dir: &Path) -> Result<&'static Mutex<LlmEmbedding>, AppError> {
149 if let Some(e) = EMBEDDER.get() {
150 return Ok(e);
151 }
152 let backend = LlmEmbedding::detect_available()?;
153 let _ = EMBEDDER.set(Mutex::new(backend));
154 EMBEDDER
155 .get()
156 .ok_or_else(|| AppError::Embedding("embedder unavailable after initialisation".to_string()))
157}
158
159pub fn get_claude_embedder(
164 claude_binary: Option<&Path>,
165 claude_model: Option<&str>,
166) -> Result<&'static Mutex<LlmEmbedding>, AppError> {
167 if let Some(e) = CLAUDE_EMBEDDER.get() {
168 return Ok(e);
169 }
170 let mut builder = LlmEmbedding::with_claude_builder();
171 if let Some(b) = claude_binary {
172 builder = builder.override_binary(b.to_path_buf());
173 }
174 if let Some(m) = claude_model {
175 builder = builder.override_model(m.to_string());
176 }
177 let backend = builder.build()?;
178 let _ = CLAUDE_EMBEDDER.set(Mutex::new(backend));
179 CLAUDE_EMBEDDER.get().ok_or_else(|| {
180 AppError::Embedding("claude embedder unavailable after initialisation".to_string())
181 })
182}
183
184pub fn get_opencode_embedder(
189 opencode_binary: Option<&Path>,
190 opencode_model: Option<&str>,
191) -> Result<&'static Mutex<LlmEmbedding>, AppError> {
192 if let Some(e) = OPENCODE_EMBEDDER.get() {
193 return Ok(e);
194 }
195 let mut builder = LlmEmbedding::with_opencode_builder();
196 if let Some(b) = opencode_binary {
197 builder = builder.override_binary(b.to_path_buf());
198 }
199 if let Some(m) = opencode_model {
200 builder = builder.override_model(m.to_string());
201 }
202 let backend = builder.build()?;
203 let _ = OPENCODE_EMBEDDER.set(Mutex::new(backend));
204 OPENCODE_EMBEDDER.get().ok_or_else(|| {
205 AppError::Embedding("opencode embedder unavailable after initialisation".to_string())
206 })
207}
208
209pub fn get_openrouter_embedder(
210 api_key: secrecy::SecretBox<String>,
211 model: &str,
212 dim: usize,
213) -> Result<&'static crate::embedding_api::OpenRouterClient, AppError> {
214 if let Some(c) = OPENROUTER_CLIENT.get() {
215 return Ok(c);
216 }
217 let client = crate::embedding_api::OpenRouterClient::new(api_key, model.to_string(), dim)?;
218 let _ = OPENROUTER_CLIENT.set(client);
219 OPENROUTER_CLIENT.get().ok_or_else(|| {
220 AppError::Embedding("openrouter client unavailable after initialisation".to_string())
221 })
222}
223
224pub fn get_openrouter_chat_client(
228 api_key: secrecy::SecretBox<String>,
229 model: &str,
230 timeout_secs: u64,
231) -> Result<&'static crate::chat_api::OpenRouterChatClient, AppError> {
232 if let Some(c) = OPENROUTER_CHAT_CLIENT.get() {
233 return Ok(c);
234 }
235 let client =
236 crate::chat_api::OpenRouterChatClient::new(api_key, model.to_string(), timeout_secs)?;
237 let _ = OPENROUTER_CHAT_CLIENT.set(client);
238 OPENROUTER_CHAT_CLIENT.get().ok_or_else(|| {
239 AppError::Embedding("openrouter chat client unavailable after initialisation".to_string())
240 })
241}
242
243pub fn openrouter_chat_client() -> Option<&'static crate::chat_api::OpenRouterChatClient> {
248 OPENROUTER_CHAT_CLIENT.get()
249}
250
251pub fn embed_via_claude_local(
255 _models_dir: &Path,
256 text: &str,
257 claude_binary: Option<&Path>,
258 claude_model: Option<&str>,
259) -> Result<Vec<f32>, AppError> {
260 let _slot_guard = acquire_llm_slot_for_embedding()?;
261 let embedder = get_claude_embedder(claude_binary, claude_model)?;
262 embed_passage(embedder, text)
263}
264
265pub fn embed_via_claude_local_resolved(
270 _models_dir: &Path,
271 text: &str,
272 claude_binary: Option<&Path>,
273 claude_model: Option<&str>,
274) -> Result<(Vec<f32>, LlmBackendKind), AppError> {
275 let _slot_guard = acquire_llm_slot_for_embedding()?;
276 let embedder = get_claude_embedder(claude_binary, claude_model)?;
277 let v = embed_passage(embedder, text)?;
278 Ok((v, LlmBackendKind::Claude))
279}
280
281pub fn embed_via_opencode_local_resolved(
286 _models_dir: &Path,
287 text: &str,
288 opencode_binary: Option<&Path>,
289 opencode_model: Option<&str>,
290) -> Result<(Vec<f32>, LlmBackendKind), AppError> {
291 let _slot_guard = acquire_llm_slot_for_embedding()?;
292 let embedder = get_opencode_embedder(opencode_binary, opencode_model)?;
293 let v = embed_passage(embedder, text)?;
294 Ok((v, LlmBackendKind::Opencode))
295}
296fn clone_client(embedder: &Mutex<LlmEmbedding>) -> LlmEmbedding {
299 embedder.lock().clone()
300}
301
302pub fn embed_passage(embedder: &Mutex<LlmEmbedding>, text: &str) -> Result<Vec<f32>, AppError> {
306 let client = clone_client(embedder);
307 let result = client.embed_passage(text)?;
308 validate_dim(result)
309}
310
311pub fn embed_query(embedder: &Mutex<LlmEmbedding>, text: &str) -> Result<Vec<f32>, AppError> {
315 let client = clone_client(embedder);
316 let result = client.embed_query(text)?;
317 validate_dim(result)
318}
319
320pub fn embed_passages_controlled(
325 embedder: &Mutex<LlmEmbedding>,
326 texts: &[&str],
327 _token_counts: &[usize],
328) -> Result<Vec<Vec<f32>>, AppError> {
329 if texts.is_empty() {
330 return Ok(Vec::new());
331 }
332 let owned: Vec<String> = texts.iter().map(|t| t.to_string()).collect();
333 embed_texts_parallel(embedder, &owned, 1, chunk_embed_batch_size())
334}
335
336pub fn embed_passage_local(models_dir: &Path, text: &str) -> Result<Vec<f32>, AppError> {
337 let _slot_guard = acquire_llm_slot_for_embedding()?;
338 let embedder = get_embedder(models_dir)?;
339 embed_passage(embedder, text)
340}
341
342pub fn should_skip_embedding_on_failure() -> bool {
346 matches!(
347 std::env::var("SQLITE_GRAPHRAG_SKIP_EMBEDDING_ON_FAILURE").as_deref(),
348 Ok("1") | Ok("true")
349 )
350}
351
352pub fn embed_passage_or_skip(
359 models_dir: &Path,
360 text: &str,
361 choice: Option<crate::cli::LlmBackendChoice>,
362) -> Result<Option<Vec<f32>>, AppError> {
363 match embed_passage_with_choice(models_dir, text, choice) {
364 Ok((v, _backend)) => Ok(Some(v)),
365 Err(AppError::Validation(msg)) => Err(AppError::Validation(msg)),
366 Err(e) => {
367 if should_skip_embedding_on_failure() {
368 tracing::warn!(
369 error = %e,
370 "embedding failed but --skip-embedding-on-failure is active; persisting with NULL embedding"
371 );
372 Ok(None)
373 } else {
374 Err(e)
375 }
376 }
377 }
378}
379
380pub fn embed_passage_local_resolved(
386 models_dir: &Path,
387 text: &str,
388) -> Result<(Vec<f32>, LlmBackendKind), AppError> {
389 let _slot_guard = acquire_llm_slot_for_embedding()?;
390 let embedder = get_embedder(models_dir)?;
391 let v = embed_passage(embedder, text)?;
392 let kind = match embedder.lock().flavour() {
393 crate::extract::llm_embedding::EmbeddingFlavour::Codex => LlmBackendKind::Codex,
394 crate::extract::llm_embedding::EmbeddingFlavour::Claude => LlmBackendKind::Claude,
395 crate::extract::llm_embedding::EmbeddingFlavour::Opencode => LlmBackendKind::Opencode,
396 };
397 Ok((v, kind))
398}
399
400pub fn embed_query_local(models_dir: &Path, text: &str) -> Result<Vec<f32>, AppError> {
401 let _slot_guard = acquire_llm_slot_for_embedding()?;
402 let embedder = get_embedder(models_dir)?;
403 embed_query(embedder, text)
404}
405
406pub fn embed_passage_with_choice(
423 models_dir: &Path,
424 text: &str,
425 choice: Option<crate::cli::LlmBackendChoice>,
426) -> Result<(Vec<f32>, LlmBackendKind), AppError> {
427 let _slot_guard = acquire_llm_slot_for_embedding()?;
428 match choice {
429 None => {
430 let embedder = get_embedder(models_dir)?;
431 embed_passage(embedder, text).map(|v| (v, LlmBackendKind::None))
432 }
433 Some(choice) => embed_with_fallback(models_dir, text, &choice.to_chain(), false),
434 }
435}
436
437pub fn embed_passage_with_embedding_choice(
441 models_dir: &Path,
442 text: &str,
443 embedding_backend: crate::cli::EmbeddingBackendChoice,
444 llm_backend: crate::cli::LlmBackendChoice,
445) -> Result<(Vec<f32>, LlmBackendKind), AppError> {
446 let _slot_guard = acquire_llm_slot_for_embedding()?;
447 let chain = embedding_backend.to_chain(llm_backend);
448 embed_with_fallback(models_dir, text, &chain, false)
449}
450
451pub fn try_embed_query_with_choice(
457 models_dir: &Path,
458 text: &str,
459 choice: Option<crate::cli::LlmBackendChoice>,
460) -> Result<(Vec<f32>, LlmBackendKind), FallbackReason> {
461 match embed_passage_with_choice(models_dir, text, choice) {
462 Ok((v, _backend)) if v.is_empty() => Err(FallbackReason::DimZero),
475 Ok((v, backend)) => Ok((v, backend)),
476 Err(e) => Err(classify_embedding_error(e)),
477 }
478}
479pub fn try_embed_query_with_embedding_choice(
484 models_dir: &Path,
485 text: &str,
486 embedding_backend: crate::cli::EmbeddingBackendChoice,
487 llm_backend: crate::cli::LlmBackendChoice,
488) -> Result<(Vec<f32>, LlmBackendKind), FallbackReason> {
489 match embed_passage_with_embedding_choice(models_dir, text, embedding_backend, llm_backend) {
490 Ok((v, _backend)) if v.is_empty() => Err(FallbackReason::DimZero),
491 Ok((v, backend)) => Ok((v, backend)),
492 Err(e) => Err(classify_embedding_error(e)),
493 }
494}
495
496fn acquire_llm_slot_for_embedding() -> Result<crate::llm_slots::LlmSlotGuard, AppError> {
508 use crate::constants::{CLI_LOCK_DEFAULT_WAIT_SECS, LLM_WORKER_RSS_MB};
509 let max = std::env::var("SQLITE_GRAPHRAG_LLM_MAX_HOST_CONCURRENCY")
510 .ok()
511 .and_then(|s| s.parse::<u32>().ok())
512 .filter(|n| *n >= 1)
513 .unwrap_or_else(crate::llm_slots::default_max_concurrency);
514 let wait_secs = if std::env::var("SQLITE_GRAPHRAG_LLM_SLOT_NO_WAIT").is_ok() {
515 0
516 } else {
517 std::env::var("SQLITE_GRAPHRAG_LLM_SLOT_WAIT_SECS")
518 .ok()
519 .and_then(|s| s.parse::<u64>().ok())
520 .unwrap_or(CLI_LOCK_DEFAULT_WAIT_SECS)
521 };
522 let _ = LLM_WORKER_RSS_MB; match crate::llm_slots::acquire_llm_slot(max, wait_secs) {
530 Ok(guard) => Ok(guard),
531 Err(e @ AppError::LockBusy { .. }) if wait_secs > 0 => Err(AppError::Embedding(format!(
532 "slot exhausted: {e} (fall back to FTS5)"
533 ))),
534 Err(e) => Err(e),
535 }
536}
537#[derive(Debug, Clone, Copy, PartialEq, Eq)]
549pub enum EmbeddingErrorKind {
550 OAuth,
552 Quota,
554 SlotExhausted,
556 BackendMismatch,
558 ZeroDimension,
560 Unknown,
562}
563
564impl EmbeddingErrorKind {
565 pub fn classify(msg: &str) -> Self {
574 let m = msg.to_lowercase();
575 if m.contains("oauth") {
576 Self::OAuth
577 } else if m.contains("quota") {
578 Self::Quota
579 } else if m.contains("slot exhausted") {
580 Self::SlotExhausted
581 } else if m.contains("backend mismatch") {
582 Self::BackendMismatch
583 } else if m.contains("dim") && m.contains("zero") {
584 Self::ZeroDimension
585 } else {
586 Self::Unknown
587 }
588 }
589
590 pub fn code(&self) -> &'static str {
592 match self {
593 Self::OAuth => "oauth",
594 Self::Quota => "quota",
595 Self::SlotExhausted => "slot-exhausted",
596 Self::BackendMismatch => "backend-mismatch",
597 Self::ZeroDimension => "zero-dimension",
598 Self::Unknown => "unknown",
599 }
600 }
601}
602
603impl std::fmt::Display for EmbeddingErrorKind {
604 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
605 f.write_str(self.code())
606 }
607}
608
609#[derive(Debug, Clone, PartialEq)]
616pub enum FallbackReason {
617 EmbeddingFailed(String),
621 SlotExhausted,
626 OAuthQuota { backend: &'static str },
630 BackendMismatch {
634 requested: &'static str,
635 resolved: &'static str,
636 },
637 DimZero,
642 Cancelled,
644 Timeout {
647 operation: String,
648 duration_secs: u64,
649 },
650}
651
652impl FallbackReason {
653 pub fn reason_code(&self) -> &'static str {
657 match self {
658 Self::EmbeddingFailed(_) => "embedding_failed",
659 Self::SlotExhausted => "slot_exhausted",
660 Self::OAuthQuota { .. } => "oauth_quota",
661 Self::BackendMismatch { .. } => "backend_mismatch",
662 Self::DimZero => "dim_zero",
663 Self::Cancelled => "cancelled",
664 Self::Timeout { .. } => "timeout",
665 }
666 }
667}
668
669impl std::fmt::Display for FallbackReason {
670 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
671 match self {
672 Self::EmbeddingFailed(msg) => write!(f, "embedding failed: {msg}"),
673 Self::SlotExhausted => write!(
674 f,
675 "slot exhausted: failed to acquire LLM slot after backoff window (max=8 concurrent, total backoff=750ms)"
676 ),
677 Self::OAuthQuota { backend } => {
678 write!(f, "OAuth usage quota exhausted on backend '{backend}'")
679 }
680 Self::BackendMismatch {
681 requested,
682 resolved,
683 } => {
684 write!(
685 f,
686 "backend mismatch: user requested '{requested}' but '{resolved}' was invoked"
687 )
688 }
689 Self::DimZero => write!(f, "embedding returned zero-dimensional vector"),
690 Self::Cancelled => write!(f, "embedding cancelled by external signal"),
691 Self::Timeout {
692 operation,
693 duration_secs,
694 } => {
695 write!(
696 f,
697 "embedding timed out after {duration_secs}s during {operation}"
698 )
699 }
700 }
701 }
702}
703
704impl std::error::Error for FallbackReason {}
705
706pub fn try_embed_query_with_fallback(
714 models_dir: &Path,
715 query: &str,
716) -> Result<(Vec<f32>, LlmBackendKind), FallbackReason> {
717 match embed_query_local(models_dir, query) {
718 Ok(v) => Ok((v, LlmBackendKind::None)),
719 Err(e) => Err(classify_embedding_error(e)),
720 }
721}
722
723pub fn try_embed_query_with_deterministic_fallback(
732 models_dir: &Path,
733 query: &str,
734 choice: Option<crate::cli::LlmBackendChoice>,
735) -> Result<(Vec<f32>, LlmBackendKind), FallbackReason> {
736 match try_embed_query_with_choice(models_dir, query, choice) {
737 Ok(t) => Ok(t),
738 Err(reason @ FallbackReason::OAuthQuota { backend }) => {
739 let alt = match backend {
740 "codex" => Some(crate::cli::LlmBackendChoice::Claude),
741 "claude" => Some(crate::cli::LlmBackendChoice::Codex),
742 "opencode" => Some(crate::cli::LlmBackendChoice::Codex),
743 "openrouter" => Some(crate::cli::LlmBackendChoice::Codex),
744 _ => None,
745 };
746 if let Some(alt_choice) = alt {
747 try_embed_query_with_choice(models_dir, query, Some(alt_choice))
748 } else {
749 Err(reason)
750 }
751 }
752 Err(reason @ FallbackReason::SlotExhausted) => {
753 std::thread::sleep(std::time::Duration::from_millis(750));
754 try_embed_query_with_choice(models_dir, query, choice).or(Err(reason))
755 }
756 Err(other) => Err(other),
757 }
758}
759
760pub fn classify_embedding_error(err: AppError) -> FallbackReason {
768 match err {
769 AppError::Timeout {
770 operation,
771 duration_secs,
772 } => FallbackReason::Timeout {
773 operation,
774 duration_secs,
775 },
776 AppError::Embedding(msg) => match EmbeddingErrorKind::classify(&msg) {
777 EmbeddingErrorKind::SlotExhausted => FallbackReason::SlotExhausted,
786 EmbeddingErrorKind::OAuth => {
787 let backend = if msg.contains("codex") {
788 "codex"
789 } else if msg.contains("claude") || msg.contains("anthropic-ratelimit") {
790 "claude"
795 } else if msg.contains("opencode") {
796 "opencode"
797 } else {
798 "unknown"
799 };
800 FallbackReason::OAuthQuota { backend }
801 }
802 EmbeddingErrorKind::Quota => {
803 let backend = if msg.contains("codex") {
804 "codex"
805 } else if msg.contains("claude") || msg.contains("anthropic-ratelimit") {
806 "claude"
807 } else if msg.contains("opencode") {
808 "opencode"
809 } else {
810 "unknown"
811 };
812 FallbackReason::OAuthQuota { backend }
813 }
814 EmbeddingErrorKind::BackendMismatch => {
815 let (requested, resolved) =
820 if msg.contains("requested claude") && msg.contains("but codex") {
821 ("claude", "codex")
822 } else if msg.contains("requested codex") && msg.contains("but claude") {
823 ("codex", "claude")
824 } else if msg.contains("requested claude") {
825 ("claude", "unknown")
826 } else if msg.contains("requested codex") {
827 ("codex", "unknown")
828 } else {
829 ("unknown", "unknown")
830 };
831 FallbackReason::BackendMismatch {
832 requested,
833 resolved,
834 }
835 }
836 EmbeddingErrorKind::ZeroDimension => FallbackReason::DimZero,
837 EmbeddingErrorKind::Unknown => {
838 if msg.contains("cancelled") {
839 FallbackReason::Cancelled
840 } else {
841 FallbackReason::EmbeddingFailed(msg)
842 }
843 }
844 },
845 e => FallbackReason::EmbeddingFailed(e.to_string()),
846 }
847}
848pub fn embed_with_fallback(
867 models_dir: &Path,
868 text: &str,
869 chain: &[LlmBackendKind],
870 skip_on_failure: bool,
871) -> Result<(Vec<f32>, LlmBackendKind), AppError> {
872 use crate::llm::exit_code_hints::LlmBackendError;
873 let effective: Vec<LlmBackendKind> = if chain.is_empty() {
874 vec![
875 LlmBackendKind::Codex,
876 LlmBackendKind::Claude,
877 LlmBackendKind::Opencode,
878 LlmBackendKind::None,
879 ]
880 } else {
881 chain.to_vec()
882 };
883
884 let mut last_err: Option<AppError> = None;
885 for backend in &effective {
886 match embed_via_backend_strict(
897 models_dir,
898 text,
899 backend,
900 last_err.as_ref(),
901 skip_on_failure,
902 ) {
903 Ok((v, resolved_kind)) => return Ok((v, resolved_kind)),
904 Err(e) => {
905 if matches!(e, AppError::Validation(_)) {
910 return Err(e);
911 }
912 tracing::warn!(
913 target: "embedding",
914 backend = ?backend,
915 error = %e,
916 "embed_with_fallback: backend failed, trying next"
917 );
918 last_err = Some(e);
919 }
920 }
921 }
922 if skip_on_failure {
923 return Ok((Vec::new(), LlmBackendKind::None));
928 }
929 Err(last_err
930 .unwrap_or_else(|| AppError::Embedding(LlmBackendError::NoBackendsAvailable.to_string())))
931}
932
933#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
937pub enum LlmBackendKind {
938 Codex,
940 Claude,
942 Opencode,
944 OpenRouter,
946 None,
948}
949
950impl LlmBackendKind {
951 pub fn as_str(self) -> &'static str {
954 match self {
955 Self::Codex => "codex",
956 Self::Claude => "claude",
957 Self::Opencode => "opencode",
958 Self::OpenRouter => "openrouter",
959 Self::None => "none",
960 }
961 }
962}
963
964pub fn embed_via_backend(
979 models_dir: &Path,
980 text: &str,
981 backend: &LlmBackendKind,
982) -> Result<(Vec<f32>, LlmBackendKind), AppError> {
983 match backend {
984 LlmBackendKind::None => Ok((Vec::new(), LlmBackendKind::None)),
985 LlmBackendKind::Codex => embed_passage_local_resolved(models_dir, text),
986 LlmBackendKind::Claude => {
987 tracing::debug!(
991 target: "embedder",
992 backend = "claude",
993 "embed_via_backend: forcing claude (ADR-0042 / GAP-002 fix)"
994 );
995 embed_via_claude_local_resolved(models_dir, text, None, None)
996 }
997 LlmBackendKind::Opencode => {
998 tracing::debug!(
999 target: "embedder",
1000 backend = "opencode",
1001 "embed_via_backend: forcing opencode (GAP-OPENCODE-001)"
1002 );
1003 embed_via_opencode_local_resolved(models_dir, text, None, None)
1004 }
1005 LlmBackendKind::OpenRouter => {
1006 tracing::debug!(
1007 target: "embedder",
1008 backend = "openrouter",
1009 "embed_via_backend: using OpenRouter API (v1.0.93)"
1010 );
1011 let client = OPENROUTER_CLIENT.get().ok_or_else(|| {
1012 AppError::Embedding(
1013 "OpenRouter client not initialised; call get_openrouter_embedder first".into(),
1014 )
1015 })?;
1016 let vec = match tokio::runtime::Handle::try_current() {
1023 Ok(handle) => tokio::task::block_in_place(|| {
1024 handle.block_on(client.embed_single(text, client.default_input_type()))
1025 })?,
1026 Err(_) => shared_runtime()?
1027 .block_on(client.embed_single(text, client.default_input_type()))?,
1028 };
1029 Ok((vec, LlmBackendKind::OpenRouter))
1030 }
1031 }
1032}
1033
1034pub fn embed_via_backend_strict(
1047 models_dir: &Path,
1048 text: &str,
1049 backend: &LlmBackendKind,
1050 last_err: Option<&AppError>,
1051 skip_on_failure: bool,
1052) -> Result<(Vec<f32>, LlmBackendKind), AppError> {
1053 use crate::llm::exit_code_hints::LlmBackendError;
1054 match backend {
1055 LlmBackendKind::None => {
1056 if skip_on_failure && last_err.is_none() {
1060 Ok((Vec::new(), LlmBackendKind::None))
1061 } else if last_err.is_some() {
1062 Err(match last_err {
1066 Some(e) => AppError::Embedding(format!("{e}")),
1067 None => AppError::Embedding(LlmBackendError::NoBackendsAvailable.to_string()),
1068 })
1069 } else {
1070 Err(AppError::Embedding(
1073 LlmBackendError::NoBackendsAvailable.to_string(),
1074 ))
1075 }
1076 }
1077 LlmBackendKind::Codex => embed_passage_local_resolved(models_dir, text),
1078 LlmBackendKind::Claude => {
1079 tracing::debug!(
1080 target: "embedder",
1081 backend = "claude",
1082 "embed_via_backend_strict: forcing claude (ADR-0042 / GAP-002 fix)"
1083 );
1084 embed_via_claude_local_resolved(models_dir, text, None, None)
1085 }
1086 LlmBackendKind::Opencode => {
1087 tracing::debug!(
1088 target: "embedder",
1089 backend = "opencode",
1090 "embed_via_backend_strict: forcing opencode (GAP-OPENCODE-001)"
1091 );
1092 embed_via_opencode_local_resolved(models_dir, text, None, None)
1093 }
1094 LlmBackendKind::OpenRouter => embed_via_backend(models_dir, text, backend),
1095 }
1096}
1097
1098pub fn embed_via_backend_legacy(
1103 models_dir: &Path,
1104 text: &str,
1105 backend: &LlmBackendKind,
1106) -> Result<Vec<f32>, AppError> {
1107 embed_via_backend(models_dir, text, backend).map(|(v, _)| v)
1108}
1109
1110pub fn embed_passages_controlled_local(
1111 models_dir: &Path,
1112 texts: &[&str],
1113 token_counts: &[usize],
1114) -> Result<Vec<Vec<f32>>, AppError> {
1115 let embedder = get_embedder(models_dir)?;
1116 embed_passages_controlled(embedder, texts, token_counts)
1117}
1118
1119pub fn embed_passages_parallel_local(
1122 models_dir: &Path,
1123 texts: &[String],
1124 parallelism: usize,
1125 batch_size: usize,
1126) -> Result<Vec<Vec<f32>>, AppError> {
1127 let embedder = get_embedder(models_dir)?;
1128 embed_texts_parallel(embedder, texts, parallelism, batch_size)
1129}
1130
1131type EmbedChunkResult = (usize, Result<Vec<Vec<f32>>, AppError>);
1135
1136fn reassemble_ordered(mut parts: Vec<(usize, Vec<Vec<f32>>)>) -> Vec<Vec<f32>> {
1141 parts.sort_by_key(|(idx, _)| *idx);
1142 parts.into_iter().flat_map(|(_, v)| v).collect()
1143}
1144
1145pub fn embed_passages_parallel_with_embedding_choice(
1152 models_dir: &Path,
1153 texts: &[String],
1154 parallelism: usize,
1155 batch_size: usize,
1156 embedding_backend: crate::cli::EmbeddingBackendChoice,
1157 llm_backend: crate::cli::LlmBackendChoice,
1158) -> Result<Vec<Vec<f32>>, AppError> {
1159 let chain = embedding_backend.to_chain(llm_backend);
1160 if chain.first() == Some(&LlmBackendKind::OpenRouter) && is_openrouter_initialized() {
1161 let client = OPENROUTER_CLIENT.get().ok_or_else(|| {
1162 AppError::Embedding(
1163 "OpenRouter client not initialised; call get_openrouter_embedder first".into(),
1164 )
1165 })?;
1166
1167 let k = parallelism.clamp(1, 16);
1172 if texts.len() <= 32 || k == 1 {
1173 let refs: Vec<&str> = texts.iter().map(|s| s.as_str()).collect();
1174 let vecs = match tokio::runtime::Handle::try_current() {
1176 Ok(handle) => tokio::task::block_in_place(|| {
1177 handle.block_on(client.embed_batch(&refs, client.default_input_type()))
1178 })?,
1179 Err(_) => shared_runtime()?
1180 .block_on(client.embed_batch(&refs, client.default_input_type()))?,
1181 };
1182 return Ok(vecs);
1183 }
1184
1185 let fan_out = async move {
1194 let mut set: JoinSet<EmbedChunkResult> = JoinSet::new();
1195 let mut parts: Vec<(usize, Vec<Vec<f32>>)> = Vec::new();
1196
1197 for (idx, chunk) in texts.chunks(32).enumerate() {
1198 if set.len() >= k {
1199 if let Some(joined) = set.join_next().await {
1200 let (cidx, res) = joined.map_err(|e| {
1201 AppError::Embedding(format!("embedding task join error: {e}"))
1202 })?;
1203 parts.push((cidx, res?));
1204 }
1205 }
1206 let owned: Vec<String> = chunk.to_vec();
1207 set.spawn(async move {
1208 let refs: Vec<&str> = owned.iter().map(|s| s.as_str()).collect();
1209 let r = client
1213 .embed_batch(&refs, client.default_input_type())
1214 .await
1215 .map_err(AppError::from);
1216 (idx, r)
1217 });
1218 }
1219
1220 while let Some(joined) = set.join_next().await {
1221 let (cidx, res) = joined
1222 .map_err(|e| AppError::Embedding(format!("embedding task join error: {e}")))?;
1223 parts.push((cidx, res?));
1224 }
1225
1226 Ok::<Vec<Vec<f32>>, AppError>(reassemble_ordered(parts))
1227 };
1228 let vecs = match tokio::runtime::Handle::try_current() {
1229 Ok(handle) => tokio::task::block_in_place(|| handle.block_on(fan_out))?,
1230 Err(_) => shared_runtime()?.block_on(fan_out)?,
1231 };
1232 Ok(vecs)
1233 } else {
1234 embed_passages_parallel_local(models_dir, texts, parallelism, batch_size)
1235 }
1236}
1237
1238type EntityEmbedCacheMap = std::collections::HashMap<u64, Arc<Vec<f32>>>;
1250
1251static ENTITY_EMBED_CACHE: OnceLock<parking_lot::Mutex<EntityEmbedCacheMap>> = OnceLock::new();
1252
1253fn entity_embed_cache() -> &'static parking_lot::Mutex<EntityEmbedCacheMap> {
1254 ENTITY_EMBED_CACHE.get_or_init(|| parking_lot::Mutex::new(std::collections::HashMap::new()))
1255}
1256
1257fn entity_cache_key(model: &str, text: &str) -> u64 {
1258 let mut hasher = blake3::Hasher::new();
1259 hasher.update(model.as_bytes());
1260 hasher.update(b"\0");
1261 hasher.update(text.as_bytes());
1262 let h = hasher.finalize();
1263 let bytes = h.as_bytes();
1264 u64::from_le_bytes([
1265 bytes[0], bytes[1], bytes[2], bytes[3], bytes[4], bytes[5], bytes[6], bytes[7],
1266 ])
1267}
1268
1269pub fn embed_entity_texts_cached(
1279 models_dir: &Path,
1280 texts: &[String],
1281 parallelism: usize,
1282 embedding_backend: crate::cli::EmbeddingBackendChoice,
1283 llm_backend: crate::cli::LlmBackendChoice,
1284) -> Result<(Vec<Vec<f32>>, EmbedCacheStats), AppError> {
1285 if texts.is_empty() {
1286 return Ok((Vec::new(), EmbedCacheStats::default()));
1287 }
1288 let chain = embedding_backend.to_chain(llm_backend);
1292
1293 if chain.as_slice() == [LlmBackendKind::None] {
1299 let out: Vec<Vec<f32>> = texts.iter().map(|_| Vec::new()).collect();
1300 return Ok((
1301 out,
1302 EmbedCacheStats {
1303 requested: texts.len(),
1304 hits: 0,
1305 misses: texts.len(),
1306 },
1307 ));
1308 }
1309
1310 let routed_openrouter =
1316 chain.first() == Some(&LlmBackendKind::OpenRouter) && is_openrouter_initialized();
1317 let model = if routed_openrouter {
1318 format!("openrouter:{}", crate::constants::embedding_dim())
1319 } else {
1320 get_embedder(models_dir)?.lock().model_label()
1321 };
1322 let cache = entity_embed_cache();
1323 let mut hits: Vec<Option<Arc<Vec<f32>>>> = vec![None; texts.len()];
1324 let mut miss_indices: Vec<usize> = Vec::with_capacity(texts.len());
1325 {
1326 let guard = cache.lock();
1327 for (i, text) in texts.iter().enumerate() {
1328 let key = entity_cache_key(&model, text);
1329 if let Some(v) = guard.get(&key) {
1330 hits[i] = Some(Arc::clone(v));
1331 } else {
1332 miss_indices.push(i);
1333 }
1334 }
1335 }
1336 let miss_count = miss_indices.len();
1337 if miss_count > 0 {
1338 let miss_texts: Vec<String> = miss_indices.iter().map(|&i| texts[i].clone()).collect();
1339 let miss_vecs = embed_passages_parallel_with_embedding_choice(
1343 models_dir,
1344 &miss_texts,
1345 parallelism,
1346 entity_embed_batch_size(),
1347 embedding_backend,
1348 llm_backend,
1349 )?;
1350 let mut guard = cache.lock();
1351 for (slot, &orig_idx) in miss_indices.iter().enumerate() {
1352 let vec = Arc::new(miss_vecs[slot].clone());
1353 let key = entity_cache_key(&model, &texts[orig_idx]);
1354 guard.insert(key, Arc::clone(&vec));
1355 hits[orig_idx] = Some(vec);
1356 }
1357 }
1358 let mut out = Vec::with_capacity(texts.len());
1359 for hit in hits.into_iter() {
1360 let v = hit.ok_or_else(|| {
1361 AppError::Embedding("entity embed cache produced null result".to_string())
1362 })?;
1363 out.push((*v).clone());
1364 }
1365 Ok((
1366 out,
1367 EmbedCacheStats {
1368 requested: texts.len(),
1369 hits: texts.len() - miss_count,
1370 misses: miss_count,
1371 },
1372 ))
1373}
1374
1375#[derive(Debug, Default, Clone, Copy, PartialEq, Eq, serde::Serialize)]
1377pub struct EmbedCacheStats {
1378 pub requested: usize,
1379 pub hits: usize,
1380 pub misses: usize,
1381}
1382
1383impl EmbedCacheStats {
1384 pub fn hit_rate(&self) -> f64 {
1386 if self.requested == 0 {
1387 0.0
1388 } else {
1389 self.hits as f64 / self.requested as f64
1390 }
1391 }
1392}
1393
1394pub fn embed_texts_parallel(
1407 embedder: &Mutex<LlmEmbedding>,
1408 texts: &[String],
1409 parallelism: usize,
1410 batch_size: usize,
1411) -> Result<Vec<Vec<f32>>, AppError> {
1412 let mut slots: Vec<Option<Vec<f32>>> = vec![None; texts.len()];
1413 embed_texts_parallel_with(embedder, texts, parallelism, batch_size, |idx, v| {
1414 slots[idx] = Some(v.to_vec());
1415 Ok(())
1416 })?;
1417 let mut out = Vec::with_capacity(slots.len());
1418 for (idx, slot) in slots.into_iter().enumerate() {
1419 out.push(slot.ok_or_else(|| {
1420 AppError::Embedding(format!("embedding fan-out lost item index {idx}"))
1421 })?);
1422 }
1423 Ok(out)
1424}
1425
1426pub fn embed_texts_parallel_with(
1430 embedder: &Mutex<LlmEmbedding>,
1431 texts: &[String],
1432 parallelism: usize,
1433 batch_size: usize,
1434 mut on_result: impl FnMut(usize, &[f32]) -> Result<(), AppError>,
1435) -> Result<(), AppError> {
1436 if texts.is_empty() {
1437 return Ok(());
1438 }
1439 let dim = crate::constants::embedding_dim();
1440 if texts.len() == 1 {
1441 let v = embed_passage(embedder, &texts[0])?;
1442 return on_result(0, &v);
1443 }
1444
1445 let client = clone_client(embedder);
1446 let permits = effective_permits(parallelism);
1447 let batches = build_batches(texts, batch_size.max(1));
1448 let token = crate::cancel_token().clone();
1449
1450 let work = move |batch: Vec<(usize, String)>| {
1451 let client = client.clone();
1452 async move {
1453 client
1454 .embed_batch_async(crate::constants::PASSAGE_PREFIX, &batch)
1455 .await
1456 }
1457 };
1458
1459 let fan_out = run_bounded(batches, permits, dim, token, work, &mut on_result);
1460 match tokio::runtime::Handle::try_current() {
1461 Ok(handle) => tokio::task::block_in_place(|| handle.block_on(fan_out)),
1462 Err(_) => shared_runtime()?.block_on(fan_out),
1463 }
1464}
1465
1466fn build_batches(texts: &[String], batch_size: usize) -> Vec<Vec<(usize, String)>> {
1468 texts
1469 .iter()
1470 .cloned()
1471 .enumerate()
1472 .collect::<Vec<_>>()
1473 .chunks(batch_size)
1474 .map(|c| c.to_vec())
1475 .collect()
1476}
1477
1478pub fn effective_permits(requested: usize) -> usize {
1483 let cpus = std::thread::available_parallelism()
1484 .map(|n| n.get())
1485 .unwrap_or(4);
1486 let by_ram = ((crate::memory_guard::available_memory_mb() / 2)
1487 / crate::constants::LLM_WORKER_RSS_MB)
1488 .max(1) as usize;
1489 requested.clamp(1, 32).min(cpus).min(by_ram).max(1)
1490}
1491
1492async fn run_bounded<F, Fut>(
1502 batches: Vec<Vec<(usize, String)>>,
1503 permits: usize,
1504 dim: usize,
1505 token: CancellationToken,
1506 work: F,
1507 on_result: &mut impl FnMut(usize, &[f32]) -> Result<(), AppError>,
1508) -> Result<(), AppError>
1509where
1510 F: Fn(Vec<(usize, String)>) -> Fut + Clone + Send + 'static,
1511 Fut: std::future::Future<Output = Result<Vec<(usize, Vec<f32>)>, AppError>> + Send,
1512{
1513 let total_batches = batches.len();
1514 let semaphore = Arc::new(Semaphore::new(permits));
1515 let (tx, mut rx) = mpsc::channel::<Result<Vec<(usize, Vec<f32>)>, AppError>>(permits * 2);
1518 let mut set: JoinSet<()> = JoinSet::new();
1519
1520 for (batch_idx, batch) in batches.into_iter().enumerate() {
1521 let sem = Arc::clone(&semaphore);
1522 let token = token.clone();
1523 let tx = tx.clone();
1524 let work = work.clone();
1525 set.spawn(async move {
1526 let wait_start = std::time::Instant::now();
1527 let Ok(_permit) = sem.acquire_owned().await else {
1530 let _ = tx
1531 .send(Err(AppError::Embedding("semaphore closed".to_string())))
1532 .await;
1533 return;
1534 };
1535 let permit_wait_ms = wait_start.elapsed().as_millis() as u64;
1536 let work_start = std::time::Instant::now();
1537 let outcome = if crate::should_obey_shutdown() {
1543 tokio::select! {
1544 res = work(batch) => res,
1545 _ = token.cancelled() => Err(AppError::Embedding(
1546 "embedding cancelled by shutdown signal".to_string(),
1547 )),
1548 }
1549 } else {
1550 work(batch).await
1551 };
1552 tracing::debug!(
1554 target: "embedding",
1555 batch_idx,
1556 permit_wait_ms,
1557 work_ms = work_start.elapsed().as_millis() as u64,
1558 ok = outcome.is_ok(),
1559 "embedding batch finished"
1560 );
1561 let _ = tx.send(outcome).await;
1562 });
1563 }
1564 drop(tx);
1565
1566 let mut completed = 0usize;
1567 let mut failed = 0usize;
1568 let mut cancelled = 0usize;
1569 let mut first_error: Option<AppError> = None;
1570
1571 while let Some(message) = rx.recv().await {
1572 match message {
1573 Ok(items) => {
1574 completed += 1;
1575 if first_error.is_none() {
1576 for (idx, v) in items {
1577 if v.len() != dim {
1578 first_error = Some(AppError::Embedding(format!(
1579 "LLM returned {} dims for item {idx}, expected {dim}; \
1580 refusing to truncate or pad silently (G42/C5)",
1581 v.len()
1582 )));
1583 break;
1584 }
1585 if let Err(e) = on_result(idx, &v) {
1586 first_error = Some(e);
1587 break;
1588 }
1589 }
1590 if first_error.is_some() {
1591 set.shutdown().await;
1594 }
1595 }
1596 }
1597 Err(e) => {
1598 if matches!(&e, AppError::Embedding(msg) if msg.contains("cancelled")) {
1599 cancelled += 1;
1600 } else {
1601 failed += 1;
1602 }
1603 if first_error.is_none() {
1604 first_error = Some(e);
1605 set.shutdown().await;
1606 }
1607 }
1608 }
1609 }
1610
1611 while let Some(join_result) = set.join_next().await {
1614 if let Err(join_err) = join_result {
1615 if join_err.is_panic() {
1616 failed += 1;
1617 if first_error.is_none() {
1618 first_error = Some(AppError::Embedding(format!(
1619 "embedding task panicked: {join_err}"
1620 )));
1621 }
1622 } else {
1623 cancelled += 1;
1624 }
1625 }
1626 }
1627
1628 tracing::debug!(
1638 target: "embedding",
1639 total_batches,
1640 completed,
1641 failed,
1642 cancelled,
1643 "embedding fan-out finished"
1644 );
1645
1646 match first_error {
1647 Some(e) => Err(e),
1648 None => Ok(()),
1649 }
1650}
1651
1652pub fn f32_to_bytes(v: &[f32]) -> Vec<u8> {
1653 let mut out = Vec::with_capacity(v.len() * 4);
1654 for f in v {
1655 out.extend_from_slice(&f.to_le_bytes());
1656 }
1657 out
1658}
1659
1660pub fn bytes_to_f32(bytes: &[u8]) -> Vec<f32> {
1661 let mut out = Vec::with_capacity(bytes.len() / 4);
1662 for chunk in bytes.chunks_exact(4) {
1663 out.push(f32::from_le_bytes([chunk[0], chunk[1], chunk[2], chunk[3]]));
1664 }
1665 out
1666}
1667
1668pub fn embedding_dim() -> usize {
1671 crate::constants::embedding_dim()
1672}
1673
1674fn validate_dim(v: Vec<f32>) -> Result<Vec<f32>, AppError> {
1678 let dim = crate::constants::embedding_dim();
1679 if v.len() != dim {
1680 return Err(AppError::Embedding(format!(
1681 "embedding has {} dims, expected {dim}; \
1682 refusing to truncate or pad silently (G42/C5)",
1683 v.len()
1684 )));
1685 }
1686 Ok(v)
1687}
1688
1689#[cfg(test)]
1690mod tests {
1691 use super::*;
1692 use std::sync::atomic::{AtomicUsize, Ordering};
1693
1694 #[test]
1695 fn reassemble_ordered_restores_input_order() {
1696 let parts = vec![
1700 (2, vec![vec![2.0_f32], vec![2.1]]),
1701 (0, vec![vec![0.0], vec![0.1]]),
1702 (1, vec![vec![1.0], vec![1.1]]),
1703 ];
1704 let out = reassemble_ordered(parts);
1705 assert_eq!(
1706 out,
1707 vec![
1708 vec![0.0_f32],
1709 vec![0.1],
1710 vec![1.0],
1711 vec![1.1],
1712 vec![2.0],
1713 vec![2.1],
1714 ]
1715 );
1716 }
1717
1718 #[test]
1719 fn f32_to_bytes_roundtrip() {
1720 let input = vec![0.0_f32, 1.5, -2.25, f32::MIN, f32::MAX];
1721 let bytes = f32_to_bytes(&input);
1722 assert_eq!(bytes.len(), input.len() * 4);
1723 let out = bytes_to_f32(&bytes);
1724 assert_eq!(out, input);
1725 }
1726
1727 #[test]
1728 fn validate_dim_rejects_divergent_vectors() {
1729 let dim = crate::constants::embedding_dim();
1732 let long = vec![0.0; dim + 10];
1733 assert!(validate_dim(long).is_err(), "longer vector must error");
1734 let short = vec![0.0; dim.saturating_sub(1).max(1)];
1735 assert!(validate_dim(short).is_err(), "shorter vector must error");
1736 let exact = vec![0.0; dim];
1737 assert_eq!(validate_dim(exact).expect("exact dim must pass").len(), dim);
1738 }
1739
1740 #[test]
1741 fn embedding_dim_matches_constants_source() {
1742 assert_eq!(embedding_dim(), crate::constants::embedding_dim());
1743 }
1744
1745 #[test]
1746 fn build_batches_preserves_global_indices() {
1747 let texts: Vec<String> = (0..10).map(|i| format!("t{i}")).collect();
1748 let batches = build_batches(&texts, 4);
1749 assert_eq!(batches.len(), 3);
1750 assert_eq!(batches[0].len(), 4);
1751 assert_eq!(batches[2].len(), 2);
1752 assert_eq!(batches[2][1].0, 9);
1753 assert_eq!(batches[2][1].1, "t9");
1754 }
1755
1756 #[test]
1757 fn effective_permits_clamps_to_bounds() {
1758 assert!(effective_permits(0) >= 1);
1759 assert!(effective_permits(1000) <= 32);
1760 }
1761
1762 fn test_batches(n: usize) -> Vec<Vec<(usize, String)>> {
1763 (0..n).map(|i| vec![(i, format!("t{i}"))]).collect()
1764 }
1765
1766 fn dummy_vec(dim: usize) -> Vec<f32> {
1767 vec![0.0; dim]
1768 }
1769
1770 #[test]
1773 fn concurrency_peak_never_exceeds_permits() {
1774 let permits = 4usize;
1775 let batches = test_batches(permits * 10);
1776 let dim = crate::constants::embedding_dim();
1777 let current = Arc::new(AtomicUsize::new(0));
1778 let peak = Arc::new(AtomicUsize::new(0));
1779
1780 let current_c = Arc::clone(¤t);
1781 let peak_c = Arc::clone(&peak);
1782 let work = move |batch: Vec<(usize, String)>| {
1783 let current = Arc::clone(¤t_c);
1784 let peak = Arc::clone(&peak_c);
1785 async move {
1786 let now = current.fetch_add(1, Ordering::SeqCst) + 1;
1787 peak.fetch_max(now, Ordering::SeqCst);
1788 tokio::time::sleep(std::time::Duration::from_millis(20)).await;
1789 current.fetch_sub(1, Ordering::SeqCst);
1790 Ok(batch
1791 .into_iter()
1792 .map(|(i, _)| (i, dummy_vec(dim)))
1793 .collect())
1794 }
1795 };
1796
1797 let mut delivered = 0usize;
1798 let rt = tokio::runtime::Builder::new_multi_thread()
1799 .worker_threads(4)
1800 .enable_all()
1801 .build()
1802 .expect("test runtime");
1803 rt.block_on(run_bounded(
1804 batches,
1805 permits,
1806 dim,
1807 CancellationToken::new(),
1808 work,
1809 &mut |_idx, _v| {
1810 delivered += 1;
1811 Ok(())
1812 },
1813 ))
1814 .expect("fan-out must succeed");
1815
1816 assert_eq!(delivered, permits * 10, "every item must be delivered");
1817 assert!(
1818 peak.load(Ordering::SeqCst) <= permits,
1819 "peak concurrency {} exceeded permits {permits}",
1820 peak.load(Ordering::SeqCst)
1821 );
1822 }
1823
1824 #[test]
1827 fn panicking_task_returns_permit_and_surfaces_error() {
1828 let permits = 2usize;
1829 let batches = test_batches(4);
1830 let dim = crate::constants::embedding_dim();
1831
1832 let work = move |batch: Vec<(usize, String)>| async move {
1833 if batch[0].0 == 1 {
1834 panic!("intentional test panic");
1835 }
1836 Ok(batch
1837 .into_iter()
1838 .map(|(i, _)| (i, dummy_vec(dim)))
1839 .collect())
1840 };
1841
1842 let rt = tokio::runtime::Builder::new_multi_thread()
1843 .worker_threads(2)
1844 .enable_all()
1845 .build()
1846 .expect("test runtime");
1847 let result = rt.block_on(run_bounded(
1848 batches,
1849 permits,
1850 dim,
1851 CancellationToken::new(),
1852 work,
1853 &mut |_idx, _v| Ok(()),
1854 ));
1855
1856 let err = result.expect_err("panic must surface as an error");
1857 assert!(
1858 err.to_string().contains("panicked"),
1859 "error must mention the panic: {err}"
1860 );
1861 }
1862
1863 #[test]
1866 fn cancellation_terminates_fan_out_quickly() {
1867 let permits = 2usize;
1868 let batches = test_batches(8);
1869 let dim = crate::constants::embedding_dim();
1870 let token = CancellationToken::new();
1871
1872 let work = move |batch: Vec<(usize, String)>| async move {
1873 tokio::time::sleep(std::time::Duration::from_secs(30)).await;
1875 Ok(batch
1876 .into_iter()
1877 .map(|(i, _)| (i, dummy_vec(dim)))
1878 .collect())
1879 };
1880
1881 let rt = tokio::runtime::Builder::new_multi_thread()
1882 .worker_threads(2)
1883 .enable_all()
1884 .build()
1885 .expect("test runtime");
1886 let cancel = token.clone();
1887 let start = std::time::Instant::now();
1888 let result = rt.block_on(async move {
1889 tokio::spawn(async move {
1890 tokio::time::sleep(std::time::Duration::from_millis(50)).await;
1891 cancel.cancel();
1892 });
1893 run_bounded(batches, permits, dim, token, work, &mut |_idx, _v| Ok(())).await
1894 });
1895
1896 assert!(result.is_err(), "cancelled fan-out must report an error");
1897 assert!(
1898 start.elapsed() < std::time::Duration::from_secs(10),
1899 "graceful shutdown must finish well under the work duration"
1900 );
1901 }
1902
1903 #[test]
1906 fn fan_out_rejects_divergent_dim() {
1907 let permits = 2usize;
1908 let batches = test_batches(2);
1909 let dim = crate::constants::embedding_dim();
1910
1911 let work = move |batch: Vec<(usize, String)>| async move {
1912 Ok(batch
1913 .into_iter()
1914 .map(|(i, _)| (i, vec![0.0f32; 3]))
1915 .collect::<Vec<(usize, Vec<f32>)>>())
1916 };
1917
1918 let rt = tokio::runtime::Builder::new_multi_thread()
1919 .worker_threads(2)
1920 .enable_all()
1921 .build()
1922 .expect("test runtime");
1923 let result = rt.block_on(run_bounded(
1924 batches,
1925 permits,
1926 dim,
1927 CancellationToken::new(),
1928 work,
1929 &mut |_idx, _v| Ok(()),
1930 ));
1931
1932 let err = result.expect_err("divergent dim must fail the fan-out");
1933 assert!(err.to_string().contains("G42/C5"), "error cites C5: {err}");
1934 }
1935
1936 #[test]
1938 fn adaptive_batch_dim64_keeps_calibrated_sizes() {
1939 assert_eq!(adaptive_batch_for_dim(CHUNK_EMBED_BATCH_SIZE, 64), 8);
1940 assert_eq!(adaptive_batch_for_dim(ENTITY_EMBED_BATCH_SIZE, 64), 25);
1941 }
1942
1943 #[test]
1945 fn adaptive_batch_dim384_shrinks() {
1946 assert_eq!(adaptive_batch_for_dim(CHUNK_EMBED_BATCH_SIZE, 384), 1);
1947 assert_eq!(adaptive_batch_for_dim(ENTITY_EMBED_BATCH_SIZE, 384), 4);
1948 }
1949
1950 #[test]
1952 fn adaptive_batch_intermediate_dims() {
1953 assert_eq!(adaptive_batch_for_dim(8, 128), 4);
1954 assert_eq!(adaptive_batch_for_dim(8, 256), 2);
1955 }
1956
1957 #[test]
1959 fn adaptive_batch_small_dim_clamps_to_base() {
1960 assert_eq!(adaptive_batch_for_dim(8, 8), 8);
1961 }
1962
1963 #[test]
1965 fn adaptive_batch_total_function() {
1966 assert_eq!(adaptive_batch_for_dim(8, 4096), 1);
1967 assert_eq!(adaptive_batch_for_dim(8, 0), 8);
1968 assert_eq!(adaptive_batch_for_dim(0, 64), 1);
1969 }
1970
1971 #[test]
1973 #[serial_test::serial(env)]
1974 fn adaptive_wrappers_follow_env_dim() {
1975 std::env::set_var("SQLITE_GRAPHRAG_EMBEDDING_DIM", "384");
1976 let chunk = chunk_embed_batch_size();
1977 let entity = entity_embed_batch_size();
1978 std::env::remove_var("SQLITE_GRAPHRAG_EMBEDDING_DIM");
1979 crate::constants::set_active_embedding_dim(crate::constants::DEFAULT_EMBEDDING_DIM);
1980 assert_eq!(chunk, 1, "384-dim chunk batch must shrink to 1 (G44)");
1981 assert_eq!(entity, 4, "384-dim entity batch must shrink to 4 (G44)");
1982 }
1983
1984 #[test]
1992 fn embedding_error_kind_classify_oauth_message() {
1993 assert_eq!(
1994 EmbeddingErrorKind::classify("OAuth token expired for claude"),
1995 EmbeddingErrorKind::OAuth,
1996 );
1997 assert_eq!(
1998 EmbeddingErrorKind::classify("oauth authentication failed"),
1999 EmbeddingErrorKind::OAuth,
2000 );
2001 }
2002
2003 #[test]
2006 fn embedding_error_kind_classify_quota_message() {
2007 assert_eq!(
2008 EmbeddingErrorKind::classify("quota exhausted on backend"),
2009 EmbeddingErrorKind::Quota,
2010 );
2011 assert_eq!(
2012 EmbeddingErrorKind::classify("Usage quota limit reached"),
2013 EmbeddingErrorKind::Quota,
2014 );
2015 }
2016
2017 #[test]
2021 fn embedding_error_kind_classify_slot_exhausted_message() {
2022 assert_eq!(
2023 EmbeddingErrorKind::classify(
2024 "slot exhausted: failed to acquire LLM slot after backoff"
2025 ),
2026 EmbeddingErrorKind::SlotExhausted,
2027 );
2028 }
2029
2030 #[test]
2033 fn embedding_error_kind_classify_zero_dimension_message() {
2034 assert_eq!(
2035 EmbeddingErrorKind::classify("embedding returned dim=zero"),
2036 EmbeddingErrorKind::ZeroDimension,
2037 );
2038 assert_eq!(
2039 EmbeddingErrorKind::classify("got zero-dim vector from LLM"),
2040 EmbeddingErrorKind::ZeroDimension,
2041 );
2042 }
2043
2044 #[test]
2048 fn embedding_error_kind_classify_unknown_fallback() {
2049 assert_eq!(
2050 EmbeddingErrorKind::classify("unrelated subprocess error"),
2051 EmbeddingErrorKind::Unknown,
2052 );
2053 assert_eq!(
2054 EmbeddingErrorKind::classify("rate limit hit"),
2055 EmbeddingErrorKind::Unknown,
2056 );
2057 assert_eq!(EmbeddingErrorKind::OAuth.code(), "oauth");
2059 assert_eq!(EmbeddingErrorKind::Quota.code(), "quota");
2060 assert_eq!(EmbeddingErrorKind::SlotExhausted.code(), "slot-exhausted");
2061 assert_eq!(
2062 EmbeddingErrorKind::BackendMismatch.code(),
2063 "backend-mismatch"
2064 );
2065 assert_eq!(EmbeddingErrorKind::ZeroDimension.code(), "zero-dimension");
2066 assert_eq!(EmbeddingErrorKind::Unknown.code(), "unknown");
2067 }
2068
2069 #[test]
2071 fn fallback_reason_display_does_not_panic() {
2072 let _ = FallbackReason::EmbeddingFailed("rate limit".into()).to_string();
2073 let _ = FallbackReason::Cancelled.to_string();
2074 let _ = FallbackReason::Timeout {
2075 operation: "embed_query".into(),
2076 duration_secs: 30,
2077 }
2078 .to_string();
2079 }
2080
2081 #[test]
2084 fn fallback_reason_is_partial_eq() {
2085 assert_eq!(
2086 FallbackReason::EmbeddingFailed("a".into()),
2087 FallbackReason::EmbeddingFailed("a".into())
2088 );
2089 assert_eq!(FallbackReason::Cancelled, FallbackReason::Cancelled);
2090 assert_ne!(
2091 FallbackReason::EmbeddingFailed("a".into()),
2092 FallbackReason::EmbeddingFailed("b".into())
2093 );
2094 assert_ne!(
2095 FallbackReason::Cancelled,
2096 FallbackReason::Timeout {
2097 operation: "x".into(),
2098 duration_secs: 1
2099 }
2100 );
2101 }
2102
2103 #[test]
2106 fn fallback_reason_timeout_preserves_fields() {
2107 let r = FallbackReason::Timeout {
2108 operation: "embed_query_local".into(),
2109 duration_secs: 300,
2110 };
2111 match r {
2112 FallbackReason::Timeout {
2113 operation,
2114 duration_secs,
2115 } => {
2116 assert_eq!(operation, "embed_query_local");
2117 assert_eq!(duration_secs, 300);
2118 }
2119 other => panic!("expected Timeout, got {other:?}"),
2120 }
2121 }
2122
2123 #[test]
2129 #[ignore = "G58 S1 stub: requires env without codex/claude on PATH; tracked as T5 of Fase 2"]
2130 fn try_embed_query_with_fallback_surfaces_embedding_failed_for_missing_binary() {
2131 let bogus = std::path::Path::new("/nonexistent-models-dir-for-g58-fallback-test");
2134 let result = try_embed_query_with_fallback(bogus, "hello world");
2135 match result {
2136 Err(FallbackReason::EmbeddingFailed(msg)) => {
2137 assert!(!msg.is_empty(), "fallback message must not be empty");
2139 }
2140 Err(FallbackReason::Cancelled) => {
2141 panic!("expected EmbeddingFailed, got Cancelled");
2142 }
2143 Err(FallbackReason::Timeout { .. }) => {
2144 panic!("expected EmbeddingFailed, got Timeout");
2145 }
2146 Err(FallbackReason::SlotExhausted) => {
2147 panic!("expected EmbeddingFailed, got SlotExhausted");
2148 }
2149 Err(FallbackReason::OAuthQuota { .. }) => {
2150 panic!("expected EmbeddingFailed, got OAuthQuota");
2151 }
2152 Err(FallbackReason::BackendMismatch { .. }) => {
2153 panic!("expected EmbeddingFailed, got BackendMismatch");
2154 }
2155 Err(FallbackReason::DimZero) => {
2156 panic!("expected EmbeddingFailed, got DimZero");
2157 }
2158 Ok(_) => {
2159 panic!("expected an error, got Ok — embedder must fail for bogus path");
2160 }
2161 }
2162 }
2163
2164 #[test]
2166 fn g56_entity_cache_key_is_stable_and_distinct() {
2167 let k1 = entity_cache_key("codex:default", "sqlite-graphrag");
2168 let k2 = entity_cache_key("codex:default", "sqlite-graphrag");
2169 let k3 = entity_cache_key("codex:default", "claude-code");
2170 let k4 = entity_cache_key("claude:default", "sqlite-graphrag");
2171 assert_eq!(k1, k2, "same model+text must hash identically");
2172 assert_ne!(k1, k3, "different text must hash differently");
2173 assert_ne!(k1, k4, "different model must hash differently");
2174 }
2175
2176 #[test]
2177 fn g56_entity_embed_cache_stats_hit_rate() {
2178 let zero = EmbedCacheStats::default();
2179 assert_eq!(zero.hit_rate(), 0.0);
2180 let half = EmbedCacheStats {
2181 requested: 4,
2182 hits: 2,
2183 misses: 2,
2184 };
2185 assert!((half.hit_rate() - 0.5).abs() < 1e-9);
2186 let all = EmbedCacheStats {
2187 requested: 7,
2188 hits: 7,
2189 misses: 0,
2190 };
2191 assert!((all.hit_rate() - 1.0).abs() < 1e-9);
2192 }
2193
2194 #[test]
2195 fn g56_entity_embed_cache_populates_and_hits() {
2196 let cache = entity_embed_cache();
2200 let model = "test-model";
2201 let text = "sqlite-graphrag";
2202 let key = entity_cache_key(model, text);
2203 let stored = Arc::new(vec![0.42_f32; crate::constants::embedding_dim()]);
2204 cache.lock().insert(key, Arc::clone(&stored));
2205 let guard = cache.lock();
2206 let hit = guard.get(&key).expect("cache must return stored value");
2207 assert_eq!(hit.len(), crate::constants::embedding_dim());
2208 assert!((hit[0] - 0.42).abs() < 1e-6);
2209 }
2210
2211 #[test]
2216 fn p1_openrouter_chain_ignores_llm_backend_none() {
2217 use crate::cli::{EmbeddingBackendChoice, LlmBackendChoice};
2218 let chain = EmbeddingBackendChoice::Openrouter.to_chain(LlmBackendChoice::None);
2219 assert_eq!(
2220 chain,
2221 vec![LlmBackendKind::OpenRouter],
2222 "openrouter embedding must not be silenced by --llm-backend none"
2223 );
2224 let none_chain = EmbeddingBackendChoice::Llm.to_chain(LlmBackendChoice::None);
2227 assert_eq!(none_chain, vec![LlmBackendKind::None]);
2228 }
2229
2230 #[test]
2231 fn g56_empty_texts_short_circuits_with_zero_stats() {
2232 let stats = EmbedCacheStats::default();
2235 assert_eq!(stats.requested, 0);
2236 assert_eq!(stats.hits, 0);
2237 assert_eq!(stats.misses, 0);
2238 assert_eq!(stats.hit_rate(), 0.0);
2239 }
2240}
2241
2242#[cfg(test)]
2246mod embed_with_fallback_tests {
2247 use super::*;
2248 use crate::llm::exit_code_hints::LlmBackendError;
2249
2250 #[test]
2251 fn none_backend_returns_empty_vector_without_calling_llm() {
2252 let (v, kind) = embed_via_backend(
2256 std::path::Path::new("/nonexistent"),
2257 "any text",
2258 &LlmBackendKind::None,
2259 )
2260 .expect("None backend never fails");
2261 assert!(v.is_empty());
2262 assert_eq!(kind, LlmBackendKind::None, "None backend must report None");
2263 }
2264
2265 #[test]
2266 fn empty_chain_defaults_to_codex_claude_none() {
2267 let defaults = [
2271 LlmBackendKind::Codex,
2272 LlmBackendKind::Claude,
2273 LlmBackendKind::None,
2274 ];
2275
2276 #[allow(dead_code)]
2281 fn llm_backend_kind_as_str_is_stable() {
2282 assert_eq!(LlmBackendKind::Codex.as_str(), "codex");
2283 assert_eq!(LlmBackendKind::Claude.as_str(), "claude");
2284 assert_eq!(LlmBackendKind::None.as_str(), "none");
2285 }
2286
2287 #[allow(dead_code)]
2288 fn fallback_reason_reason_code_is_stable() {
2289 assert_eq!(
2290 FallbackReason::EmbeddingFailed("any".into()).reason_code(),
2291 "embedding_failed"
2292 );
2293 assert_eq!(FallbackReason::Cancelled.reason_code(), "cancelled");
2294 assert_eq!(
2295 FallbackReason::Timeout {
2296 operation: "embed_query".into(),
2297 duration_secs: 30
2298 }
2299 .reason_code(),
2300 "timeout"
2301 );
2302 }
2303 assert_eq!(defaults.len(), 3);
2304 }
2305
2306 #[test]
2307 fn embed_with_fallback_chain_of_only_none_aborts_without_skip_on_failure_v1088() {
2308 let chain = vec![LlmBackendKind::None];
2320 let err = embed_with_fallback(
2321 std::path::Path::new("/nonexistent-models-dir-for-gap005-test"),
2322 "hello",
2323 &chain,
2324 false,
2325 )
2326 .expect_err("chain of only [None] without skip_on_failure MUST abort");
2327 let msg = format!("{err}");
2328 assert!(
2329 msg.contains("no LLM backends available"),
2330 "error must mention exhausted chain, got: {msg}"
2331 );
2332 }
2333 #[test]
2334 fn embed_with_fallback_skip_on_failure_with_only_none_returns_empty() {
2335 let chain = vec![LlmBackendKind::None];
2340 let v = embed_with_fallback(
2341 std::path::Path::new("/nonexistent-models-dir-for-gap005-test"),
2342 "hello",
2343 &chain,
2344 true,
2345 )
2346 .expect("None chain is always Ok");
2347 assert!(v.0.is_empty(), "vector must be empty");
2348 assert_eq!(v.1, LlmBackendKind::None);
2349 }
2350 #[allow(dead_code)]
2351 fn llm_backend_error_no_backends_default_message() {
2352 let e = LlmBackendError::NoBackendsAvailable;
2355 let h = e.hint();
2356 assert!(h.contains("--llm-fallback"));
2357 }
2358
2359 #[test]
2360 fn llm_backend_error_nonzero_exit_carries_stderr_tail() {
2361 let e = LlmBackendError::NonZeroExit {
2362 exit_code: Some(137),
2363 signal: Some(9),
2364 stdout_tail: "out".into(),
2365 stderr_tail: "OOM killed".into(),
2366 binary: "codex".into(),
2367 hint: "OOM".into(),
2368 };
2369 let s = e.to_string();
2370 assert!(s.contains("codex"));
2371 assert!(s.contains("OOM killed"));
2372 assert!(s.contains("signal 9") || s.contains("exit 137"));
2373 }
2374}