use crate::errors::AppError;
use crate::extract::llm_embedding::LlmEmbedding;
use parking_lot::Mutex;
use std::path::Path;
use std::sync::Arc;
use std::sync::OnceLock;
use tokio::sync::{mpsc, Semaphore};
use tokio::task::JoinSet;
use tokio_util::sync::CancellationToken;
static CLAUDE_EMBEDDER: OnceLock<Mutex<LlmEmbedding>> = OnceLock::new();
static OPENCODE_EMBEDDER: OnceLock<Mutex<LlmEmbedding>> = OnceLock::new();
static OPENROUTER_CLIENT: OnceLock<crate::embedding_api::OpenRouterClient> = OnceLock::new();
static OPENROUTER_CHAT_CLIENT: OnceLock<crate::chat_api::OpenRouterChatClient> = OnceLock::new();
pub fn is_openrouter_initialized() -> bool {
OPENROUTER_CLIENT.get().is_some()
}
static EMBEDDER: OnceLock<Mutex<LlmEmbedding>> = OnceLock::new();
static RUNTIME: OnceLock<tokio::runtime::Runtime> = OnceLock::new();
pub const CHUNK_EMBED_BATCH_SIZE: usize = 8;
pub const ENTITY_EMBED_BATCH_SIZE: usize = 25;
pub const EMBED_BATCH_CALIBRATION_DIM: usize = 64;
fn adaptive_batch_for_dim(base: usize, dim: usize) -> usize {
let base = base.max(1);
(base * EMBED_BATCH_CALIBRATION_DIM / dim.max(1)).clamp(1, base)
}
pub fn chunk_embed_batch_size() -> usize {
let dim = crate::constants::embedding_dim();
let batch = adaptive_batch_for_dim(CHUNK_EMBED_BATCH_SIZE, dim);
tracing::debug!(
dim,
base = CHUNK_EMBED_BATCH_SIZE,
batch,
"adaptive chunk batch size (G44)"
);
batch
}
pub fn entity_embed_batch_size() -> usize {
let dim = crate::constants::embedding_dim();
let batch = adaptive_batch_for_dim(ENTITY_EMBED_BATCH_SIZE, dim);
tracing::debug!(
dim,
base = ENTITY_EMBED_BATCH_SIZE,
batch,
"adaptive entity batch size (G44)"
);
batch
}
pub(crate) fn shared_runtime() -> Result<&'static tokio::runtime::Runtime, AppError> {
if let Some(rt) = RUNTIME.get() {
return Ok(rt);
}
let rt = tokio::runtime::Builder::new_multi_thread()
.worker_threads(2)
.enable_all()
.build()
.map_err(|e| AppError::Embedding(format!("tokio runtime init failed: {e}")))?;
let _ = RUNTIME.set(rt);
RUNTIME.get().ok_or_else(|| {
AppError::Embedding("tokio runtime unavailable after initialisation".to_string())
})
}
pub fn get_embedder(_models_dir: &Path) -> Result<&'static Mutex<LlmEmbedding>, AppError> {
if let Some(e) = EMBEDDER.get() {
return Ok(e);
}
let backend = LlmEmbedding::detect_available()?;
let _ = EMBEDDER.set(Mutex::new(backend));
EMBEDDER
.get()
.ok_or_else(|| AppError::Embedding("embedder unavailable after initialisation".to_string()))
}
pub fn get_claude_embedder(
claude_binary: Option<&Path>,
claude_model: Option<&str>,
) -> Result<&'static Mutex<LlmEmbedding>, AppError> {
if let Some(e) = CLAUDE_EMBEDDER.get() {
return Ok(e);
}
let mut builder = LlmEmbedding::with_claude_builder();
if let Some(b) = claude_binary {
builder = builder.override_binary(b.to_path_buf());
}
if let Some(m) = claude_model {
builder = builder.override_model(m.to_string());
}
let backend = builder.build()?;
let _ = CLAUDE_EMBEDDER.set(Mutex::new(backend));
CLAUDE_EMBEDDER.get().ok_or_else(|| {
AppError::Embedding("claude embedder unavailable after initialisation".to_string())
})
}
pub fn get_opencode_embedder(
opencode_binary: Option<&Path>,
opencode_model: Option<&str>,
) -> Result<&'static Mutex<LlmEmbedding>, AppError> {
if let Some(e) = OPENCODE_EMBEDDER.get() {
return Ok(e);
}
let mut builder = LlmEmbedding::with_opencode_builder();
if let Some(b) = opencode_binary {
builder = builder.override_binary(b.to_path_buf());
}
if let Some(m) = opencode_model {
builder = builder.override_model(m.to_string());
}
let backend = builder.build()?;
let _ = OPENCODE_EMBEDDER.set(Mutex::new(backend));
OPENCODE_EMBEDDER.get().ok_or_else(|| {
AppError::Embedding("opencode embedder unavailable after initialisation".to_string())
})
}
pub fn get_openrouter_embedder(
api_key: secrecy::SecretBox<String>,
model: &str,
dim: usize,
) -> Result<&'static crate::embedding_api::OpenRouterClient, AppError> {
if let Some(c) = OPENROUTER_CLIENT.get() {
return Ok(c);
}
let client = crate::embedding_api::OpenRouterClient::new(api_key, model.to_string(), dim)?;
let _ = OPENROUTER_CLIENT.set(client);
OPENROUTER_CLIENT.get().ok_or_else(|| {
AppError::Embedding("openrouter client unavailable after initialisation".to_string())
})
}
pub fn get_openrouter_chat_client(
api_key: secrecy::SecretBox<String>,
model: &str,
timeout_secs: u64,
) -> Result<&'static crate::chat_api::OpenRouterChatClient, AppError> {
if let Some(c) = OPENROUTER_CHAT_CLIENT.get() {
return Ok(c);
}
let client =
crate::chat_api::OpenRouterChatClient::new(api_key, model.to_string(), timeout_secs)?;
let _ = OPENROUTER_CHAT_CLIENT.set(client);
OPENROUTER_CHAT_CLIENT.get().ok_or_else(|| {
AppError::Embedding("openrouter chat client unavailable after initialisation".to_string())
})
}
pub fn openrouter_chat_client() -> Option<&'static crate::chat_api::OpenRouterChatClient> {
OPENROUTER_CHAT_CLIENT.get()
}
pub fn embed_via_claude_local(
_models_dir: &Path,
text: &str,
claude_binary: Option<&Path>,
claude_model: Option<&str>,
) -> Result<Vec<f32>, AppError> {
let _slot_guard = acquire_llm_slot_for_embedding()?;
let embedder = get_claude_embedder(claude_binary, claude_model)?;
embed_passage(embedder, text)
}
pub fn embed_via_claude_local_resolved(
_models_dir: &Path,
text: &str,
claude_binary: Option<&Path>,
claude_model: Option<&str>,
) -> Result<(Vec<f32>, LlmBackendKind), AppError> {
let _slot_guard = acquire_llm_slot_for_embedding()?;
let embedder = get_claude_embedder(claude_binary, claude_model)?;
let v = embed_passage(embedder, text)?;
Ok((v, LlmBackendKind::Claude))
}
pub fn embed_via_opencode_local_resolved(
_models_dir: &Path,
text: &str,
opencode_binary: Option<&Path>,
opencode_model: Option<&str>,
) -> Result<(Vec<f32>, LlmBackendKind), AppError> {
let _slot_guard = acquire_llm_slot_for_embedding()?;
let embedder = get_opencode_embedder(opencode_binary, opencode_model)?;
let v = embed_passage(embedder, text)?;
Ok((v, LlmBackendKind::Opencode))
}
fn clone_client(embedder: &Mutex<LlmEmbedding>) -> LlmEmbedding {
embedder.lock().clone()
}
thread_local! {
static QUERY_EMBED_FAST: std::cell::Cell<bool> = const { std::cell::Cell::new(false) };
}
fn with_query_embed_fast<T>(f: impl FnOnce() -> T) -> T {
QUERY_EMBED_FAST.with(|c| {
let prev = c.replace(true);
let out = f();
c.set(prev);
out
})
}
fn apply_query_timeout_if_needed(client: LlmEmbedding) -> LlmEmbedding {
if QUERY_EMBED_FAST.with(|c| c.get()) {
let secs = crate::runtime_config::resolve_u64(
None,
"llm.query_embed_timeout_secs",
crate::constants::DEFAULT_QUERY_EMBED_TIMEOUT_SECS,
);
client.with_timeout_secs(secs)
} else {
client
}
}
pub fn embed_passage(embedder: &Mutex<LlmEmbedding>, text: &str) -> Result<Vec<f32>, AppError> {
let client = apply_query_timeout_if_needed(clone_client(embedder));
let result = client.embed_passage(text)?;
validate_dim(result)
}
pub fn embed_query(embedder: &Mutex<LlmEmbedding>, text: &str) -> Result<Vec<f32>, AppError> {
let client = apply_query_timeout_if_needed(clone_client(embedder));
let result = client.embed_query(text)?;
validate_dim(result)
}
pub fn embed_passages_controlled(
embedder: &Mutex<LlmEmbedding>,
texts: &[&str],
_token_counts: &[usize],
) -> Result<Vec<Vec<f32>>, AppError> {
if texts.is_empty() {
return Ok(Vec::new());
}
let owned: Vec<String> = texts.iter().map(|t| t.to_string()).collect();
embed_texts_parallel(embedder, &owned, 1, chunk_embed_batch_size())
}
pub fn embed_passage_local(models_dir: &Path, text: &str) -> Result<Vec<f32>, AppError> {
let _slot_guard = acquire_llm_slot_for_embedding()?;
let embedder = get_embedder(models_dir)?;
embed_passage(embedder, text)
}
pub fn should_skip_embedding_on_failure() -> bool {
crate::runtime_config::skip_embedding_on_failure()
}
pub fn embed_passage_or_skip(
models_dir: &Path,
text: &str,
choice: Option<crate::cli::LlmBackendChoice>,
) -> Result<Option<Vec<f32>>, AppError> {
match embed_passage_with_choice(models_dir, text, choice) {
Ok((v, _backend)) => Ok(Some(v)),
Err(AppError::Validation(msg)) => Err(AppError::Validation(msg)),
Err(e) => {
if should_skip_embedding_on_failure() {
tracing::warn!(
error = %e,
"embedding failed but --skip-embedding-on-failure is active; persisting with NULL embedding"
);
Ok(None)
} else {
Err(e)
}
}
}
}
pub fn embed_passage_local_resolved(
models_dir: &Path,
text: &str,
) -> Result<(Vec<f32>, LlmBackendKind), AppError> {
let _slot_guard = acquire_llm_slot_for_embedding()?;
let embedder = get_embedder(models_dir)?;
let v = embed_passage(embedder, text)?;
let kind = match embedder.lock().flavour() {
crate::extract::llm_embedding::EmbeddingFlavour::Codex => LlmBackendKind::Codex,
crate::extract::llm_embedding::EmbeddingFlavour::Claude => LlmBackendKind::Claude,
crate::extract::llm_embedding::EmbeddingFlavour::Opencode => LlmBackendKind::Opencode,
};
Ok((v, kind))
}
pub fn embed_query_local(models_dir: &Path, text: &str) -> Result<Vec<f32>, AppError> {
let _slot_guard = acquire_llm_slot_for_embedding()?;
let embedder = get_embedder(models_dir)?;
embed_query(embedder, text)
}
pub fn embed_passage_with_choice(
models_dir: &Path,
text: &str,
choice: Option<crate::cli::LlmBackendChoice>,
) -> Result<(Vec<f32>, LlmBackendKind), AppError> {
let _slot_guard = acquire_llm_slot_for_embedding()?;
match choice {
None => {
let embedder = get_embedder(models_dir)?;
embed_passage(embedder, text).map(|v| (v, LlmBackendKind::None))
}
Some(choice) => embed_with_fallback(models_dir, text, &choice.to_chain(), false),
}
}
pub fn embed_passage_with_embedding_choice(
models_dir: &Path,
text: &str,
embedding_backend: crate::cli::EmbeddingBackendChoice,
llm_backend: crate::cli::LlmBackendChoice,
) -> Result<(Vec<f32>, LlmBackendKind), AppError> {
let _slot_guard = acquire_llm_slot_for_embedding()?;
let chain = embedding_backend.to_chain(llm_backend);
embed_with_fallback(models_dir, text, &chain, false)
}
pub fn try_embed_query_with_choice(
models_dir: &Path,
text: &str,
choice: Option<crate::cli::LlmBackendChoice>,
) -> Result<(Vec<f32>, LlmBackendKind), FallbackReason> {
match with_query_embed_fast(|| embed_passage_with_choice(models_dir, text, choice)) {
Ok((v, _backend)) if v.is_empty() => Err(FallbackReason::DimZero),
Ok((v, backend)) => Ok((v, backend)),
Err(e) => Err(classify_embedding_error(e)),
}
}
pub fn try_embed_query_with_embedding_choice(
models_dir: &Path,
text: &str,
embedding_backend: crate::cli::EmbeddingBackendChoice,
llm_backend: crate::cli::LlmBackendChoice,
) -> Result<(Vec<f32>, LlmBackendKind), FallbackReason> {
match with_query_embed_fast(|| {
embed_passage_with_embedding_choice(models_dir, text, embedding_backend, llm_backend)
}) {
Ok((v, _backend)) if v.is_empty() => Err(FallbackReason::DimZero),
Ok((v, backend)) => Ok((v, backend)),
Err(e) => Err(classify_embedding_error(e)),
}
}
fn acquire_llm_slot_for_embedding() -> Result<crate::llm_slots::LlmSlotGuard, AppError> {
use crate::constants::{CLI_LOCK_DEFAULT_WAIT_SECS, LLM_WORKER_RSS_MB};
let default_max = crate::llm_slots::default_max_concurrency() as usize;
let max = crate::runtime_config::llm_max_host_concurrency(default_max).max(1) as u32;
let wait_secs = if crate::runtime_config::llm_slot_no_wait() {
0
} else {
crate::runtime_config::llm_slot_wait_secs(CLI_LOCK_DEFAULT_WAIT_SECS)
};
let _ = LLM_WORKER_RSS_MB; match crate::llm_slots::acquire_llm_slot(max, wait_secs) {
Ok(guard) => Ok(guard),
Err(e @ AppError::LockBusy { .. }) if wait_secs > 0 => Err(AppError::Embedding(format!(
"slot exhausted: {e} (fall back to FTS5)"
))),
Err(e) => Err(e),
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum EmbeddingErrorKind {
OAuth,
Quota,
SlotExhausted,
BackendMismatch,
ZeroDimension,
Unknown,
}
impl EmbeddingErrorKind {
pub fn classify(msg: &str) -> Self {
let m = msg.to_lowercase();
if m.contains("oauth") {
Self::OAuth
} else if m.contains("quota") {
Self::Quota
} else if m.contains("slot exhausted") {
Self::SlotExhausted
} else if m.contains("backend mismatch") {
Self::BackendMismatch
} else if m.contains("dim") && m.contains("zero") {
Self::ZeroDimension
} else {
Self::Unknown
}
}
pub fn code(&self) -> &'static str {
match self {
Self::OAuth => "oauth",
Self::Quota => "quota",
Self::SlotExhausted => "slot-exhausted",
Self::BackendMismatch => "backend-mismatch",
Self::ZeroDimension => "zero-dimension",
Self::Unknown => "unknown",
}
}
}
impl std::fmt::Display for EmbeddingErrorKind {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.write_str(self.code())
}
}
#[derive(Debug, Clone, PartialEq)]
pub enum FallbackReason {
EmbeddingFailed(String),
SlotExhausted,
OAuthQuota { backend: &'static str },
BackendMismatch {
requested: &'static str,
resolved: &'static str,
},
DimZero,
Cancelled,
Timeout {
operation: String,
duration_secs: u64,
},
}
impl FallbackReason {
pub fn reason_code(&self) -> &'static str {
match self {
Self::EmbeddingFailed(_) => "embedding_failed",
Self::SlotExhausted => "slot_exhausted",
Self::OAuthQuota { .. } => "oauth_quota",
Self::BackendMismatch { .. } => "backend_mismatch",
Self::DimZero => "dim_zero",
Self::Cancelled => "cancelled",
Self::Timeout { .. } => "timeout",
}
}
}
impl std::fmt::Display for FallbackReason {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
Self::EmbeddingFailed(msg) => write!(f, "embedding failed: {msg}"),
Self::SlotExhausted => write!(
f,
"slot exhausted: failed to acquire LLM slot after backoff window (max=8 concurrent, total backoff=750ms)"
),
Self::OAuthQuota { backend } => {
write!(f, "OAuth usage quota exhausted on backend '{backend}'")
}
Self::BackendMismatch {
requested,
resolved,
} => {
write!(
f,
"backend mismatch: user requested '{requested}' but '{resolved}' was invoked"
)
}
Self::DimZero => write!(f, "embedding returned zero-dimensional vector"),
Self::Cancelled => write!(f, "embedding cancelled by external signal"),
Self::Timeout {
operation,
duration_secs,
} => {
write!(
f,
"embedding timed out after {duration_secs}s during {operation}"
)
}
}
}
}
impl std::error::Error for FallbackReason {}
pub fn try_embed_query_with_fallback(
models_dir: &Path,
query: &str,
) -> Result<(Vec<f32>, LlmBackendKind), FallbackReason> {
match embed_query_local(models_dir, query) {
Ok(v) => Ok((v, LlmBackendKind::None)),
Err(e) => Err(classify_embedding_error(e)),
}
}
pub fn try_embed_query_with_deterministic_fallback(
models_dir: &Path,
query: &str,
choice: Option<crate::cli::LlmBackendChoice>,
) -> Result<(Vec<f32>, LlmBackendKind), FallbackReason> {
match try_embed_query_with_choice(models_dir, query, choice) {
Ok(t) => Ok(t),
Err(reason @ FallbackReason::OAuthQuota { backend }) => {
let alt = match backend {
"codex" => Some(crate::cli::LlmBackendChoice::Claude),
"claude" => Some(crate::cli::LlmBackendChoice::Codex),
"opencode" => Some(crate::cli::LlmBackendChoice::Codex),
"openrouter" => Some(crate::cli::LlmBackendChoice::Codex),
_ => None,
};
if let Some(alt_choice) = alt {
try_embed_query_with_choice(models_dir, query, Some(alt_choice))
} else {
Err(reason)
}
}
Err(reason @ FallbackReason::SlotExhausted) => {
std::thread::sleep(std::time::Duration::from_millis(750));
try_embed_query_with_choice(models_dir, query, choice).or(Err(reason))
}
Err(other) => Err(other),
}
}
pub fn classify_embedding_error(err: AppError) -> FallbackReason {
match err {
AppError::Timeout {
operation,
duration_secs,
} => FallbackReason::Timeout {
operation,
duration_secs,
},
AppError::Embedding(msg) => match EmbeddingErrorKind::classify(&msg) {
EmbeddingErrorKind::SlotExhausted => FallbackReason::SlotExhausted,
EmbeddingErrorKind::OAuth => {
let backend = if msg.contains("codex") {
"codex"
} else if msg.contains("claude") || msg.contains("anthropic-ratelimit") {
"claude"
} else if msg.contains("opencode") {
"opencode"
} else {
"unknown"
};
FallbackReason::OAuthQuota { backend }
}
EmbeddingErrorKind::Quota => {
let backend = if msg.contains("codex") {
"codex"
} else if msg.contains("claude") || msg.contains("anthropic-ratelimit") {
"claude"
} else if msg.contains("opencode") {
"opencode"
} else {
"unknown"
};
FallbackReason::OAuthQuota { backend }
}
EmbeddingErrorKind::BackendMismatch => {
let (requested, resolved) =
if msg.contains("requested claude") && msg.contains("but codex") {
("claude", "codex")
} else if msg.contains("requested codex") && msg.contains("but claude") {
("codex", "claude")
} else if msg.contains("requested claude") {
("claude", "unknown")
} else if msg.contains("requested codex") {
("codex", "unknown")
} else {
("unknown", "unknown")
};
FallbackReason::BackendMismatch {
requested,
resolved,
}
}
EmbeddingErrorKind::ZeroDimension => FallbackReason::DimZero,
EmbeddingErrorKind::Unknown => {
if msg.contains("cancelled") {
FallbackReason::Cancelled
} else {
FallbackReason::EmbeddingFailed(msg)
}
}
},
e => FallbackReason::EmbeddingFailed(e.to_string()),
}
}
pub fn embed_with_fallback(
models_dir: &Path,
text: &str,
chain: &[LlmBackendKind],
skip_on_failure: bool,
) -> Result<(Vec<f32>, LlmBackendKind), AppError> {
use crate::llm::exit_code_hints::LlmBackendError;
let effective: Vec<LlmBackendKind> = if chain.is_empty() {
vec![
LlmBackendKind::Codex,
LlmBackendKind::Claude,
LlmBackendKind::Opencode,
LlmBackendKind::None,
]
} else {
chain.to_vec()
};
let mut last_err: Option<AppError> = None;
for backend in &effective {
if let Err(probe_err) = backend_ready_probe(backend) {
tracing::warn!(
target: "embedding",
backend = ?backend,
error = %probe_err,
"embed_with_fallback: backend probe failed, skipping"
);
last_err = Some(probe_err);
continue;
}
match embed_via_backend_strict(
models_dir,
text,
backend,
last_err.as_ref(),
skip_on_failure,
) {
Ok((v, resolved_kind)) => return Ok((v, resolved_kind)),
Err(e) => {
if matches!(e, AppError::Validation(_)) {
return Err(e);
}
tracing::warn!(
target: "embedding",
backend = ?backend,
error = %e,
"embed_with_fallback: backend failed, trying next"
);
last_err = Some(e);
}
}
}
if skip_on_failure {
return Ok((Vec::new(), LlmBackendKind::None));
}
Err(last_err
.unwrap_or_else(|| AppError::Embedding(LlmBackendError::NoBackendsAvailable.to_string())))
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub enum LlmBackendKind {
Codex,
Claude,
Opencode,
OpenRouter,
None,
}
impl LlmBackendKind {
pub fn as_str(self) -> &'static str {
match self {
Self::Codex => "codex",
Self::Claude => "claude",
Self::Opencode => "opencode",
Self::OpenRouter => "openrouter",
Self::None => "none",
}
}
}
fn backend_ready_probe(backend: &LlmBackendKind) -> Result<(), AppError> {
match backend {
LlmBackendKind::None => Ok(()),
LlmBackendKind::OpenRouter => {
if OPENROUTER_CLIENT.get().is_some() {
Ok(())
} else {
Err(AppError::Embedding(
"openrouter probe: client not initialised (skip)".into(),
))
}
}
LlmBackendKind::Codex => {
let bin = crate::runtime_config::codex_binary()
.unwrap_or_else(|| "codex".into());
if which::which(&bin).is_err() && which::which("codex").is_err() {
return Err(AppError::Embedding(
"codex probe: binary not on PATH (skip)".into(),
));
}
let auth = std::env::var_os("CODEX_HOME")
.map(std::path::PathBuf::from)
.or_else(|| {
std::env::var_os("HOME").map(|h| {
std::path::PathBuf::from(h).join(".codex")
})
})
.map(|p| p.join("auth.json"));
match auth {
Some(p) if p.is_file() => Ok(()),
_ => Err(AppError::Embedding(
"codex probe: auth.json missing (skip; use --llm-backend none or login)".into(),
)),
}
}
LlmBackendKind::Claude => {
let bin = crate::runtime_config::claude_binary()
.unwrap_or_else(|| "claude".into());
if which::which(&bin).is_err() && which::which("claude").is_err() {
return Err(AppError::Embedding(
"claude probe: binary not on PATH (skip)".into(),
));
}
Ok(())
}
LlmBackendKind::Opencode => {
let bin = crate::runtime_config::opencode_binary()
.unwrap_or_else(|| "opencode".into());
if which::which(&bin).is_err() && which::which("opencode").is_err() {
return Err(AppError::Embedding(
"opencode probe: binary not on PATH (skip)".into(),
));
}
Ok(())
}
}
}
pub fn embed_via_backend(
models_dir: &Path,
text: &str,
backend: &LlmBackendKind,
) -> Result<(Vec<f32>, LlmBackendKind), AppError> {
match backend {
LlmBackendKind::None => Ok((Vec::new(), LlmBackendKind::None)),
LlmBackendKind::Codex => embed_passage_local_resolved(models_dir, text),
LlmBackendKind::Claude => {
tracing::debug!(
target: "embedder",
backend = "claude",
"embed_via_backend: forcing claude (ADR-0042 / GAP-002 fix)"
);
embed_via_claude_local_resolved(models_dir, text, None, None)
}
LlmBackendKind::Opencode => {
tracing::debug!(
target: "embedder",
backend = "opencode",
"embed_via_backend: forcing opencode (GAP-OPENCODE-001)"
);
embed_via_opencode_local_resolved(models_dir, text, None, None)
}
LlmBackendKind::OpenRouter => {
tracing::debug!(
target: "embedder",
backend = "openrouter",
"embed_via_backend: using OpenRouter API (v1.0.93)"
);
let client = OPENROUTER_CLIENT.get().ok_or_else(|| {
AppError::Embedding(
"OpenRouter client not initialised; call get_openrouter_embedder first".into(),
)
})?;
let vec = match tokio::runtime::Handle::try_current() {
Ok(handle) => tokio::task::block_in_place(|| {
handle.block_on(client.embed_single(text, client.default_input_type()))
})?,
Err(_) => shared_runtime()?
.block_on(client.embed_single(text, client.default_input_type()))?,
};
Ok((vec, LlmBackendKind::OpenRouter))
}
}
}
pub fn embed_via_backend_strict(
models_dir: &Path,
text: &str,
backend: &LlmBackendKind,
last_err: Option<&AppError>,
skip_on_failure: bool,
) -> Result<(Vec<f32>, LlmBackendKind), AppError> {
use crate::llm::exit_code_hints::LlmBackendError;
match backend {
LlmBackendKind::None => {
if last_err.is_none() || skip_on_failure {
Ok((Vec::new(), LlmBackendKind::None))
} else {
Err(match last_err {
Some(e) => AppError::Embedding(format!("{e}")),
None => AppError::Embedding(LlmBackendError::NoBackendsAvailable.to_string()),
})
}
}
LlmBackendKind::Codex => embed_passage_local_resolved(models_dir, text),
LlmBackendKind::Claude => {
tracing::debug!(
target: "embedder",
backend = "claude",
"embed_via_backend_strict: forcing claude (ADR-0042 / GAP-002 fix)"
);
embed_via_claude_local_resolved(models_dir, text, None, None)
}
LlmBackendKind::Opencode => {
tracing::debug!(
target: "embedder",
backend = "opencode",
"embed_via_backend_strict: forcing opencode (GAP-OPENCODE-001)"
);
embed_via_opencode_local_resolved(models_dir, text, None, None)
}
LlmBackendKind::OpenRouter => embed_via_backend(models_dir, text, backend),
}
}
pub fn embed_via_backend_legacy(
models_dir: &Path,
text: &str,
backend: &LlmBackendKind,
) -> Result<Vec<f32>, AppError> {
embed_via_backend(models_dir, text, backend).map(|(v, _)| v)
}
pub fn embed_passages_controlled_local(
models_dir: &Path,
texts: &[&str],
token_counts: &[usize],
) -> Result<Vec<Vec<f32>>, AppError> {
let embedder = get_embedder(models_dir)?;
embed_passages_controlled(embedder, texts, token_counts)
}
pub fn embed_passages_parallel_local(
models_dir: &Path,
texts: &[String],
parallelism: usize,
batch_size: usize,
) -> Result<Vec<Vec<f32>>, AppError> {
let embedder = get_embedder(models_dir)?;
embed_texts_parallel(embedder, texts, parallelism, batch_size)
}
type EmbedChunkResult = (usize, Result<Vec<Vec<f32>>, AppError>);
fn reassemble_ordered(mut parts: Vec<(usize, Vec<Vec<f32>>)>) -> Vec<Vec<f32>> {
parts.sort_by_key(|(idx, _)| *idx);
parts.into_iter().flat_map(|(_, v)| v).collect()
}
pub fn embed_passages_parallel_with_embedding_choice(
models_dir: &Path,
texts: &[String],
parallelism: usize,
batch_size: usize,
embedding_backend: crate::cli::EmbeddingBackendChoice,
llm_backend: crate::cli::LlmBackendChoice,
) -> Result<Vec<Vec<f32>>, AppError> {
let chain = embedding_backend.to_chain(llm_backend);
if chain.first() == Some(&LlmBackendKind::OpenRouter) && is_openrouter_initialized() {
let client = OPENROUTER_CLIENT.get().ok_or_else(|| {
AppError::Embedding(
"OpenRouter client not initialised; call get_openrouter_embedder first".into(),
)
})?;
let k = parallelism.clamp(1, 16);
if texts.len() <= 32 || k == 1 {
let refs: Vec<&str> = texts.iter().map(|s| s.as_str()).collect();
let vecs = match tokio::runtime::Handle::try_current() {
Ok(handle) => tokio::task::block_in_place(|| {
handle.block_on(client.embed_batch(&refs, client.default_input_type()))
})?,
Err(_) => shared_runtime()?
.block_on(client.embed_batch(&refs, client.default_input_type()))?,
};
return Ok(vecs);
}
let fan_out = async move {
let mut set: JoinSet<EmbedChunkResult> = JoinSet::new();
let mut parts: Vec<(usize, Vec<Vec<f32>>)> = Vec::new();
for (idx, chunk) in texts.chunks(32).enumerate() {
if set.len() >= k {
if let Some(joined) = set.join_next().await {
let (cidx, res) = joined.map_err(|e| {
AppError::Embedding(format!("embedding task join error: {e}"))
})?;
parts.push((cidx, res?));
}
}
let owned: Vec<String> = chunk.to_vec();
set.spawn(async move {
let refs: Vec<&str> = owned.iter().map(|s| s.as_str()).collect();
let r = client
.embed_batch(&refs, client.default_input_type())
.await
.map_err(AppError::from);
(idx, r)
});
}
while let Some(joined) = set.join_next().await {
let (cidx, res) = joined
.map_err(|e| AppError::Embedding(format!("embedding task join error: {e}")))?;
parts.push((cidx, res?));
}
Ok::<Vec<Vec<f32>>, AppError>(reassemble_ordered(parts))
};
let vecs = match tokio::runtime::Handle::try_current() {
Ok(handle) => tokio::task::block_in_place(|| handle.block_on(fan_out))?,
Err(_) => shared_runtime()?.block_on(fan_out)?,
};
Ok(vecs)
} else {
embed_passages_parallel_local(models_dir, texts, parallelism, batch_size)
}
}
type EntityEmbedCacheMap = std::collections::HashMap<u64, Arc<Vec<f32>>>;
static ENTITY_EMBED_CACHE: OnceLock<parking_lot::Mutex<EntityEmbedCacheMap>> = OnceLock::new();
fn entity_embed_cache() -> &'static parking_lot::Mutex<EntityEmbedCacheMap> {
ENTITY_EMBED_CACHE.get_or_init(|| parking_lot::Mutex::new(std::collections::HashMap::new()))
}
fn entity_cache_key(model: &str, text: &str) -> u64 {
let mut hasher = blake3::Hasher::new();
hasher.update(model.as_bytes());
hasher.update(b"\0");
hasher.update(text.as_bytes());
let h = hasher.finalize();
let bytes = h.as_bytes();
u64::from_le_bytes([
bytes[0], bytes[1], bytes[2], bytes[3], bytes[4], bytes[5], bytes[6], bytes[7],
])
}
pub fn embed_entity_texts_cached(
models_dir: &Path,
texts: &[String],
parallelism: usize,
embedding_backend: crate::cli::EmbeddingBackendChoice,
llm_backend: crate::cli::LlmBackendChoice,
) -> Result<(Vec<Vec<f32>>, EmbedCacheStats), AppError> {
if texts.is_empty() {
return Ok((Vec::new(), EmbedCacheStats::default()));
}
let chain = embedding_backend.to_chain(llm_backend);
if chain.as_slice() == [LlmBackendKind::None] {
let out: Vec<Vec<f32>> = texts.iter().map(|_| Vec::new()).collect();
return Ok((
out,
EmbedCacheStats {
requested: texts.len(),
hits: 0,
misses: texts.len(),
},
));
}
let routed_openrouter =
chain.first() == Some(&LlmBackendKind::OpenRouter) && is_openrouter_initialized();
let model = if routed_openrouter {
format!("openrouter:{}", crate::constants::embedding_dim())
} else {
get_embedder(models_dir)?.lock().model_label()
};
let cache = entity_embed_cache();
let mut hits: Vec<Option<Arc<Vec<f32>>>> = vec![None; texts.len()];
let mut miss_indices: Vec<usize> = Vec::with_capacity(texts.len());
{
let guard = cache.lock();
for (i, text) in texts.iter().enumerate() {
let key = entity_cache_key(&model, text);
if let Some(v) = guard.get(&key) {
hits[i] = Some(Arc::clone(v));
} else {
miss_indices.push(i);
}
}
}
let miss_count = miss_indices.len();
if miss_count > 0 {
let miss_texts: Vec<String> = miss_indices.iter().map(|&i| texts[i].clone()).collect();
let miss_vecs = embed_passages_parallel_with_embedding_choice(
models_dir,
&miss_texts,
parallelism,
entity_embed_batch_size(),
embedding_backend,
llm_backend,
)?;
let mut guard = cache.lock();
for (slot, &orig_idx) in miss_indices.iter().enumerate() {
let vec = Arc::new(miss_vecs[slot].clone());
let key = entity_cache_key(&model, &texts[orig_idx]);
guard.insert(key, Arc::clone(&vec));
hits[orig_idx] = Some(vec);
}
}
let mut out = Vec::with_capacity(texts.len());
for hit in hits.into_iter() {
let v = hit.ok_or_else(|| {
AppError::Embedding("entity embed cache produced null result".to_string())
})?;
out.push((*v).clone());
}
Ok((
out,
EmbedCacheStats {
requested: texts.len(),
hits: texts.len() - miss_count,
misses: miss_count,
},
))
}
#[derive(Debug, Default, Clone, Copy, PartialEq, Eq, serde::Serialize)]
pub struct EmbedCacheStats {
pub requested: usize,
pub hits: usize,
pub misses: usize,
}
impl EmbedCacheStats {
pub fn hit_rate(&self) -> f64 {
if self.requested == 0 {
0.0
} else {
self.hits as f64 / self.requested as f64
}
}
}
pub fn embed_texts_parallel(
embedder: &Mutex<LlmEmbedding>,
texts: &[String],
parallelism: usize,
batch_size: usize,
) -> Result<Vec<Vec<f32>>, AppError> {
let mut slots: Vec<Option<Vec<f32>>> = vec![None; texts.len()];
embed_texts_parallel_with(embedder, texts, parallelism, batch_size, |idx, v| {
slots[idx] = Some(v.to_vec());
Ok(())
})?;
let mut out = Vec::with_capacity(slots.len());
for (idx, slot) in slots.into_iter().enumerate() {
out.push(slot.ok_or_else(|| {
AppError::Embedding(format!("embedding fan-out lost item index {idx}"))
})?);
}
Ok(out)
}
pub fn embed_texts_parallel_with(
embedder: &Mutex<LlmEmbedding>,
texts: &[String],
parallelism: usize,
batch_size: usize,
mut on_result: impl FnMut(usize, &[f32]) -> Result<(), AppError>,
) -> Result<(), AppError> {
if texts.is_empty() {
return Ok(());
}
let dim = crate::constants::embedding_dim();
if texts.len() == 1 {
let v = embed_passage(embedder, &texts[0])?;
return on_result(0, &v);
}
let client = clone_client(embedder);
let permits = effective_permits(parallelism);
let batches = build_batches(texts, batch_size.max(1));
let token = crate::cancel_token().clone();
let work = move |batch: Vec<(usize, String)>| {
let client = client.clone();
async move {
client
.embed_batch_async(crate::constants::PASSAGE_PREFIX, &batch)
.await
}
};
let fan_out = run_bounded(batches, permits, dim, token, work, &mut on_result);
match tokio::runtime::Handle::try_current() {
Ok(handle) => tokio::task::block_in_place(|| handle.block_on(fan_out)),
Err(_) => shared_runtime()?.block_on(fan_out),
}
}
fn build_batches(texts: &[String], batch_size: usize) -> Vec<Vec<(usize, String)>> {
texts
.iter()
.cloned()
.enumerate()
.collect::<Vec<_>>()
.chunks(batch_size)
.map(|c| c.to_vec())
.collect()
}
pub fn effective_permits(requested: usize) -> usize {
let cpus = std::thread::available_parallelism()
.map(|n| n.get())
.unwrap_or(4);
let by_ram = ((crate::memory_guard::available_memory_mb() / 2)
/ crate::constants::LLM_WORKER_RSS_MB)
.max(1) as usize;
requested.clamp(1, 32).min(cpus).min(by_ram).max(1)
}
async fn run_bounded<F, Fut>(
batches: Vec<Vec<(usize, String)>>,
permits: usize,
dim: usize,
token: CancellationToken,
work: F,
on_result: &mut impl FnMut(usize, &[f32]) -> Result<(), AppError>,
) -> Result<(), AppError>
where
F: Fn(Vec<(usize, String)>) -> Fut + Clone + Send + 'static,
Fut: std::future::Future<Output = Result<Vec<(usize, Vec<f32>)>, AppError>> + Send,
{
let total_batches = batches.len();
let semaphore = Arc::new(Semaphore::new(permits));
let (tx, mut rx) = mpsc::channel::<Result<Vec<(usize, Vec<f32>)>, AppError>>(permits * 2);
let mut set: JoinSet<()> = JoinSet::new();
for (batch_idx, batch) in batches.into_iter().enumerate() {
let sem = Arc::clone(&semaphore);
let token = token.clone();
let tx = tx.clone();
let work = work.clone();
set.spawn(async move {
let wait_start = std::time::Instant::now();
let Ok(_permit) = sem.acquire_owned().await else {
let _ = tx
.send(Err(AppError::Embedding("semaphore closed".to_string())))
.await;
return;
};
let permit_wait_ms = wait_start.elapsed().as_millis() as u64;
let work_start = std::time::Instant::now();
let outcome = if crate::should_obey_shutdown() {
tokio::select! {
res = work(batch) => res,
_ = token.cancelled() => Err(AppError::Embedding(
"embedding cancelled by shutdown signal".to_string(),
)),
}
} else {
work(batch).await
};
tracing::debug!(
target: "embedding",
batch_idx,
permit_wait_ms,
work_ms = work_start.elapsed().as_millis() as u64,
ok = outcome.is_ok(),
"embedding batch finished"
);
let _ = tx.send(outcome).await;
});
}
drop(tx);
let mut completed = 0usize;
let mut failed = 0usize;
let mut cancelled = 0usize;
let mut first_error: Option<AppError> = None;
while let Some(message) = rx.recv().await {
match message {
Ok(items) => {
completed += 1;
if first_error.is_none() {
for (idx, v) in items {
if v.len() != dim {
first_error = Some(AppError::Embedding(format!(
"LLM returned {} dims for item {idx}, expected {dim}; \
refusing to truncate or pad silently (G42/C5)",
v.len()
)));
break;
}
if let Err(e) = on_result(idx, &v) {
first_error = Some(e);
break;
}
}
if first_error.is_some() {
set.shutdown().await;
}
}
}
Err(e) => {
if matches!(&e, AppError::Embedding(msg) if msg.contains("cancelled")) {
cancelled += 1;
} else {
failed += 1;
}
if first_error.is_none() {
first_error = Some(e);
set.shutdown().await;
}
}
}
}
while let Some(join_result) = set.join_next().await {
if let Err(join_err) = join_result {
if join_err.is_panic() {
failed += 1;
if first_error.is_none() {
first_error = Some(AppError::Embedding(format!(
"embedding task panicked: {join_err}"
)));
}
} else {
cancelled += 1;
}
}
}
tracing::debug!(
target: "embedding",
total_batches,
completed,
failed,
cancelled,
"embedding fan-out finished"
);
match first_error {
Some(e) => Err(e),
None => Ok(()),
}
}
pub fn f32_to_bytes(v: &[f32]) -> Vec<u8> {
let mut out = Vec::with_capacity(v.len() * 4);
for f in v {
out.extend_from_slice(&f.to_le_bytes());
}
out
}
pub fn bytes_to_f32(bytes: &[u8]) -> Vec<f32> {
let mut out = Vec::with_capacity(bytes.len() / 4);
for chunk in bytes.chunks_exact(4) {
out.push(f32::from_le_bytes([chunk[0], chunk[1], chunk[2], chunk[3]]));
}
out
}
pub fn embedding_dim() -> usize {
crate::constants::embedding_dim()
}
fn validate_dim(v: Vec<f32>) -> Result<Vec<f32>, AppError> {
let dim = crate::constants::embedding_dim();
if v.len() != dim {
return Err(AppError::Embedding(format!(
"embedding has {} dims, expected {dim}; \
refusing to truncate or pad silently (G42/C5)",
v.len()
)));
}
Ok(v)
}
#[cfg(test)]
#[path = "embedder_tests.rs"]
mod tests;
#[cfg(test)]
#[path = "embedder_fallback_tests.rs"]
mod embed_with_fallback_tests;