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 the OpenRouter REST embeddings API and stored as a BLOB in
6//! `memory_embeddings(memory_id, embedding, source)`. Vector similarity is
7//! computed in pure Rust at query time.
8//!
9//! # Workload classification (G42/S3, BLOCK 1 — MANDATORY)
10//!
11//! LLM embedding is **I/O-bound**: each call waits on a network
12//! round-trip to the OpenRouter REST API while the local CPU stays
13//! 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`): the historical
26//! per-worker RSS budget, retained as the RAM bound on the permit
27//! formula.
28//!
29use std::sync::OnceLock;
30
31/// Process-wide OpenRouter embedding client.
32pub(crate) static OPENROUTER_CLIENT: OnceLock<crate::embedding_api::OpenRouterClient> =
33 OnceLock::new();
34
35/// v1.0.95 (ADR-0054): process-wide OpenRouter chat-completions client for
36/// the `enrich` JUDGE. Distinct from `OPENROUTER_CLIENT` (embeddings) because
37/// the chat client binds a text model, not an embedding model.
38pub(crate) static OPENROUTER_CHAT_CLIENT: OnceLock<crate::chat_api::OpenRouterChatClient> =
39 OnceLock::new();
40
41/// Process-wide multi-thread tokio runtime for embedding I/O.
42///
43/// G42/A2 fix: v1.0.76-v1.0.78 built a current-thread runtime PER CALL.
44/// One runtime per process amortises the setup and hosts the bounded
45/// fan-out of `embed_texts_parallel`.
46pub(crate) static RUNTIME: OnceLock<tokio::runtime::Runtime> = OnceLock::new();
47
48// Batch sizing + parallel fan-out (R-SRP-01).
49mod backend;
50mod batch;
51mod embed_error;
52mod fallback;
53mod getters;
54mod passage;
55
56pub use backend::{
57 bytes_to_f32, embed_via_backend, embed_via_backend_legacy, embed_via_backend_strict,
58 embedding_dim, f32_to_bytes, LlmBackendKind,
59};
60pub use batch::{
61 chunk_embed_batch_size, effective_permits, embed_entity_texts_cached, entity_embed_batch_size,
62 EmbedCacheStats, CHUNK_EMBED_BATCH_SIZE, EMBED_BATCH_CALIBRATION_DIM, ENTITY_EMBED_BATCH_SIZE,
63};
64// GAP-SG-270: the conversion that keeps `EmbedError::retry_class` alive on its
65// way to the enrich queue.
66pub(crate) use embed_error::app_error_preserving_retry_class;
67
68// GAP-SG-147 / GAP-SG-163: zero-copy entry point for every caller. The
69// borrowed-slice wrapper that used to shadow it was removed in v1.2.8; this
70// is now the published multi-passage surface.
71pub use batch::embed_passages_parallel_shared;
72pub use fallback::{
73 classify_embedding_error, embed_with_fallback, try_embed_query_with_deterministic_fallback,
74 try_embed_query_with_fallback, EmbeddingErrorKind, FallbackReason,
75};
76pub use getters::{
77 get_openrouter_chat_client, get_openrouter_embedder, is_openrouter_initialized,
78 openrouter_chat_client,
79};
80pub use passage::{
81 embed_passage_or_skip, embed_passage_with_choice, embed_passage_with_embedding_choice,
82 should_skip_embedding_on_failure, try_embed_query_with_choice,
83 try_embed_query_with_embedding_choice,
84};
85
86// Crate-visible helpers used across submodules.
87pub(crate) use backend::backend_ready_probe;
88pub(crate) use getters::shared_runtime;
89
90#[cfg(test)]
91#[path = "../embedder_tests.rs"]
92mod tests;
93
94#[cfg(test)]
95#[path = "../embedder_fallback_tests.rs"]
96mod embed_with_fallback_tests;