sqlite_graphrag/embedder/
mod.rs

1//! Embedding generation for the GraphRAG memory.
2//!
3//! v1.0.76: the default build is **LLM-only** — the binary does NOT bundle
4//! fastembed / ort / ndarray / tokenizers. All embeddings are produced
5//! by a headless invocation of `claude code` or `codex` (OAuth, no MCP,
6//! no hooks) and stored as a BLOB in `memory_embeddings(memory_id, embedding,
7//! source)`. Vector similarity is computed in pure Rust at query time.
8//!
9//! # Workload classification (G42/S3, BLOCK 1 — MANDATORY)
10//!
11//! LLM embedding is **I/O-bound + subprocess-bound**: each call waits
12//! 5-60s on a network round-trip through a headless `claude -p` /
13//! `codex exec` subprocess while the local CPU stays idle. Concurrency
14//! therefore uses **tokio** (async I/O concurrency) and NEVER rayon
15//! (reserved for CPU-bound work).
16//!
17//! # Permit formula (G42/S3, BLOCO 2)
18//!
19//! ```text
20//! permits = clamp(--llm-parallelism, 1, 32)
21//!           .min(available_parallelism())
22//!           .min(available_ram_mb * 0.5 / LLM_WORKER_RSS_MB)
23//! ```
24//!
25//! `LLM_WORKER_RSS_MB = 350` (`crate::constants`): `claude -p` and
26//! `codex exec` are node processes with a typical Maximum RSS of
27//! 200-400 MB (measured via `/usr/bin/time -l` on macOS /
28//! `/usr/bin/time -v` on Linux), so the RAM bound is pertinent.
29//!
30//! # Locking contract (G42/A3 fix)
31//!
32//! The process-wide `Mutex<LlmEmbedding>` protects ONLY the cheap clone
33//! of the client configuration (flavour + binary path + model + shared
34//! schema tempfiles). It is NEVER held across network I/O — the
35//! v1.0.76-v1.0.78 `flush_group` held it for the whole sequential
36//! embedding loop, which is why `--llm-parallelism 8` measured an
37//! effective parallelism of 1.
38
39use crate::extract::llm_embedding::LlmEmbedding;
40use parking_lot::Mutex;
41use std::sync::OnceLock;
42
43/// Process-wide LLM-embedding client behind a Mutex.
44///
45/// The lock guards configuration cloning only (see module docs); the
46/// actual LLM I/O happens on clones, outside the lock.
47///
48/// ADR-0042 / GAP-002: process-wide Claude-backed LLM-embedding client
49/// behind a `Mutex`. Distinct from `EMBEDDER` so the Claude path of
50/// `embed_via_backend` no longer re-probes PATH via `detect_available`
51/// (the v1.0.82 bug where requesting Claude could resolve to Codex).
52pub(crate) static CLAUDE_EMBEDDER: OnceLock<Mutex<LlmEmbedding>> = OnceLock::new();
53pub(crate) static OPENCODE_EMBEDDER: OnceLock<Mutex<LlmEmbedding>> = OnceLock::new();
54pub(crate) static OPENROUTER_CLIENT: OnceLock<crate::embedding_api::OpenRouterClient> =
55    OnceLock::new();
56
57/// v1.0.95 (ADR-0054): process-wide OpenRouter chat-completions client for
58/// the `enrich` JUDGE. Distinct from `OPENROUTER_CLIENT` (embeddings) because
59/// the chat client binds a text model, not an embedding model.
60pub(crate) static OPENROUTER_CHAT_CLIENT: OnceLock<crate::chat_api::OpenRouterChatClient> =
61    OnceLock::new();
62
63pub(crate) static EMBEDDER: OnceLock<Mutex<LlmEmbedding>> = OnceLock::new();
64
65/// Process-wide multi-thread tokio runtime for embedding I/O.
66///
67/// G42/A2 fix: v1.0.76-v1.0.78 built a current-thread runtime PER CALL.
68/// One runtime per process amortises the setup and hosts the bounded
69/// fan-out of `embed_texts_parallel`.
70pub(crate) static RUNTIME: OnceLock<tokio::runtime::Runtime> = OnceLock::new();
71
72// Batch sizing + parallel fan-out (R-SRP-01).
73mod backend;
74mod batch;
75mod fallback;
76mod getters;
77mod passage;
78
79pub use batch::{
80    chunk_embed_batch_size, effective_permits, embed_entity_texts_cached,
81    embed_passages_controlled_local, embed_passages_parallel_local,
82    embed_passages_parallel_with_embedding_choice, embed_texts_parallel,
83    embed_texts_parallel_with, entity_embed_batch_size, EmbedCacheStats,
84    CHUNK_EMBED_BATCH_SIZE, EMBED_BATCH_CALIBRATION_DIM, ENTITY_EMBED_BATCH_SIZE,
85};
86pub use getters::{
87    embed_via_claude_local, embed_via_claude_local_resolved, embed_via_opencode_local_resolved,
88    get_claude_embedder, get_embedder, get_opencode_embedder, get_openrouter_chat_client,
89    get_openrouter_embedder, is_openrouter_initialized, openrouter_chat_client,
90};
91pub use passage::{
92    embed_passage, embed_passage_local, embed_passage_local_resolved, embed_passage_or_skip,
93    embed_passage_with_choice, embed_passage_with_embedding_choice, embed_passages_controlled,
94    embed_query, embed_query_local, should_skip_embedding_on_failure,
95    try_embed_query_with_choice, try_embed_query_with_embedding_choice,
96};
97pub use fallback::{
98    classify_embedding_error, embed_with_fallback, try_embed_query_with_deterministic_fallback,
99    try_embed_query_with_fallback, EmbeddingErrorKind, FallbackReason,
100};
101pub use backend::{
102    bytes_to_f32, embed_via_backend, embed_via_backend_legacy, embed_via_backend_strict,
103    embedding_dim, f32_to_bytes, LlmBackendKind,
104};
105
106// Crate-visible helpers used across submodules.
107pub(crate) use backend::{backend_ready_probe, validate_dim};
108pub(crate) use getters::{
109    apply_query_timeout_if_needed, clone_client, shared_runtime, with_query_embed_fast,
110};
111pub(crate) use passage::acquire_llm_slot_for_embedding;
112
113#[cfg(test)]
114#[path = "../embedder_tests.rs"]
115mod tests;
116
117#[cfg(test)]
118#[path = "../embedder_fallback_tests.rs"]
119mod embed_with_fallback_tests;
120
sqlite_graphrag/embedder/mod.rs

sqlite_graphrag/embedder/
mod.rs
