sqlite_graphrag/extract/llm_embedding/mod.rs
1//! LLM-based embedding backend (v1.0.76 default; reworked in v1.0.79 G42).
2//!
3//! `LlmEmbedding` is the production embedding client. It wraps headless
4//! invocations of `claude code` or `codex` and returns f32 vectors of the
5//! active dimensionality (`crate::constants::embedding_dim()`, default 1024 via DEFAULT_EMBEDDING_DIM).
6//!
7//! v1.0.79 (G42) changes:
8//! - S1: the dimensionality is no longer hardcoded here — the single
9//! source of truth lives in `crate::constants` and the JSON schemas
10//! are generated dynamically.
11//! - S2: `embed_batch` embeds N numbered texts per LLM call with the
12//! `{items:[{i,v}]}` schema, collapsing 39 subprocess spawns into 4-5.
13//! - S4: the codex `--output-schema` file is a `tempfile::NamedTempFile`
14//! with a randomised name created once per client and shared across
15//! clones via `Arc` — no per-call write+delete, no PID-path races.
16//! - S5: the claude model honours XDG `embedding.claude_model`
17//! (symmetric to `embedding.codex_model`). ZERO hardcoded models without
18//! a flag/XDG override.
19//! - S6: `CLAUDE_CONFIG_DIR` points at an empty managed directory BY
20//! DEFAULT, because `--strict-mcp-config`/`--mcp-config '{}'` are
21//! silently ignored upstream (anthropics/claude-code#10787) and a
22//! full `~/.claude` costs ~223k cache-creation tokens per call.
23//! - S7: the codex `request_user_input` failure mode maps to an
24//! actionable error instead of an opaque exit 11.
25//! - BLOCO 4: every subprocess uses `kill_on_drop(true)` plus an
26//! explicit `tokio::time::timeout`, so cancellation never leaks a
27//! child and a hung LLM cannot stall the pipeline forever.
28//!
29//! OAuth is the only supported credential path. The constructor rejects
30//! `ANTHROPIC_API_KEY` / `OPENAI_API_KEY` in the environment — see
31//! `v1.0.69 (G31) OAuth-Only Enforcement`.
32
33mod binary;
34mod builder;
35mod models;
36mod ops;
37mod timeout;
38mod types;
39mod wire;
40
41pub use binary::resolve_real_binary;
42pub use builder::LlmEmbeddingBuilder;
43pub use types::EmbeddingFlavour;
44
45use crate::errors::AppError;
46use models::{claude_embed_model, codex_embed_model, opencode_embed_model};
47use std::sync::Arc;
48use timeout::embed_timeout;
49use types::CodexSchemaFiles;
50
51/// LLM embedding client (headless claude / codex / opencode).
52#[derive(Clone, Debug)]
53pub struct LlmEmbedding {
54 /// Which LLM headless binary to spawn.
55 pub(crate) flavour: EmbeddingFlavour,
56 /// Cached path to the binary to avoid PATH lookups on every call.
57 pub(crate) binary: std::path::PathBuf,
58 /// Model name resolved at construction time.
59 pub(crate) model: String,
60 /// G42/S4: lazily-created codex `--output-schema` tempfiles, shared
61 /// across clones. Keyed by dim so a dim change cannot serve a stale schema.
62 pub(crate) codex_schemas: Arc<parking_lot::Mutex<CodexSchemaFiles>>,
63 /// Instance-scoped timeout override.
64 /// Precedence: this field > XDG `embedding.timeout_secs` > default 300s.
65 pub(crate) timeout_override: Option<std::time::Duration>,
66}
67
68impl LlmEmbedding {
69 /// Apply a per-call timeout override (e.g. short budget for query Auto).
70 pub fn with_timeout_secs(mut self, secs: u64) -> Self {
71 let clamped = secs.clamp(1, 3_600);
72 self.timeout_override = Some(std::time::Duration::from_secs(clamped));
73 self
74 }
75
76 /// Detects which LLM CLI is available on PATH and returns the
77 /// matching embedding client.
78 ///
79 /// Prefers `codex`, then `claude`, then `opencode` (lighter context first).
80 pub fn detect_available() -> Result<Self, AppError> {
81 Self::oauth_only_enforce()?;
82
83 if let Some(path) = crate::runtime_config::codex_binary()
84 .map(std::path::PathBuf::from)
85 .or_else(|| which::which("codex").ok())
86 {
87 return Ok(Self::from_parts(
88 EmbeddingFlavour::Codex,
89 resolve_real_binary(&path),
90 codex_embed_model(),
91 None,
92 ));
93 }
94 if let Some(path) = crate::runtime_config::claude_binary()
95 .map(std::path::PathBuf::from)
96 .or_else(|| which::which("claude").ok())
97 {
98 return Ok(Self::from_parts(
99 EmbeddingFlavour::Claude,
100 resolve_real_binary(&path),
101 claude_embed_model(),
102 None,
103 ));
104 }
105 if let Some(path) = crate::runtime_config::opencode_binary()
106 .map(std::path::PathBuf::from)
107 .or_else(|| which::which("opencode").ok())
108 {
109 return Ok(Self::from_parts(
110 EmbeddingFlavour::Opencode,
111 resolve_real_binary(&path),
112 opencode_embed_model(),
113 None,
114 ));
115 }
116 Err(AppError::Embedding(
117 crate::i18n::validation::embedding_no_llm_cli_on_path(),
118 ))
119 }
120
121 /// Build from resolved parts (used by builder and detect).
122 pub(crate) fn from_parts(
123 flavour: EmbeddingFlavour,
124 binary: std::path::PathBuf,
125 model: String,
126 timeout_override: Option<std::time::Duration>,
127 ) -> Self {
128 Self {
129 flavour,
130 binary,
131 model,
132 codex_schemas: Arc::new(parking_lot::Mutex::new(CodexSchemaFiles::default())),
133 timeout_override,
134 }
135 }
136
137 /// Instance-scoped timeout. Precedence:
138 /// `timeout_override` field > XDG > default.
139 pub(crate) fn instance_embed_timeout(&self) -> std::time::Duration {
140 if let Some(d) = self.timeout_override {
141 return d;
142 }
143 embed_timeout()
144 }
145
146 /// Instance-scoped batch timeout: base + 15s per extra item.
147 pub(crate) fn instance_embed_timeout_for_batch(&self, batch_size: usize) -> std::time::Duration {
148 let base = self.instance_embed_timeout();
149 let extra = std::time::Duration::from_secs(15) * batch_size.saturating_sub(1) as u32;
150 base + extra
151 }
152
153 /// With codex.
154 pub fn with_codex() -> Result<Self, AppError> {
155 Self::with_codex_builder().build()
156 }
157
158 /// With claude.
159 pub fn with_claude() -> Result<Self, AppError> {
160 Self::with_claude_builder().build()
161 }
162
163 /// Builder entry point for a codex-backed embedder.
164 pub fn with_codex_builder() -> LlmEmbeddingBuilder {
165 LlmEmbeddingBuilder::codex_default()
166 }
167
168 /// Builder entry point for a claude-backed embedder.
169 pub fn with_claude_builder() -> LlmEmbeddingBuilder {
170 LlmEmbeddingBuilder::claude_default()
171 }
172
173 /// With opencode.
174 pub fn with_opencode() -> Result<Self, AppError> {
175 Self::with_opencode_builder().build()
176 }
177
178 /// Builder entry point for an opencode-backed embedder.
179 pub fn with_opencode_builder() -> LlmEmbeddingBuilder {
180 LlmEmbeddingBuilder::opencode_default()
181 }
182
183 /// v1.0.69 (G31): refuse to spawn if an API key is set. The CLI must use OAuth.
184 pub(crate) fn oauth_only_enforce() -> Result<(), AppError> {
185 if std::env::var("ANTHROPIC_API_KEY").is_ok() {
186 return Err(AppError::Validation(
187 crate::i18n::validation::anthropic_api_key_oauth_required(),
188 ));
189 }
190 if std::env::var("OPENAI_API_KEY").is_ok() {
191 return Err(AppError::Validation(
192 crate::i18n::validation::openai_api_key_oauth_required(),
193 ));
194 }
195 Ok(())
196 }
197
198 /// Embeds a single passage (chunk of a memory body).
199 pub fn embed_passage(&self, text: &str) -> Result<Vec<f32>, AppError> {
200 self.invoke_with_prefix(crate::constants::PASSAGE_PREFIX, text)
201 }
202
203 /// Embeds a single query with the query prompt prefix.
204 pub fn embed_query(&self, text: &str) -> Result<Vec<f32>, AppError> {
205 self.invoke_with_prefix(crate::constants::QUERY_PREFIX, text)
206 }
207
208 /// Stable label for the active embedding model (`flavour:model`).
209 pub fn model_label(&self) -> String {
210 format!("{}:{}", self.flavour.as_str(), self.model)
211 }
212
213 /// Returns the resolved embedding flavour of this client.
214 pub fn flavour(&self) -> EmbeddingFlavour {
215 self.flavour
216 }
217}
218
219#[cfg(test)]
220mod tests;