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(
148 &self,
149 batch_size: usize,
150 ) -> std::time::Duration {
151 let base = self.instance_embed_timeout();
152 let extra = std::time::Duration::from_secs(15) * batch_size.saturating_sub(1) as u32;
153 base + extra
154 }
155
156 /// With codex.
157 pub fn with_codex() -> Result<Self, AppError> {
158 Self::with_codex_builder().build()
159 }
160
161 /// With claude.
162 pub fn with_claude() -> Result<Self, AppError> {
163 Self::with_claude_builder().build()
164 }
165
166 /// Builder entry point for a codex-backed embedder.
167 pub fn with_codex_builder() -> LlmEmbeddingBuilder {
168 LlmEmbeddingBuilder::codex_default()
169 }
170
171 /// Builder entry point for a claude-backed embedder.
172 pub fn with_claude_builder() -> LlmEmbeddingBuilder {
173 LlmEmbeddingBuilder::claude_default()
174 }
175
176 /// With opencode.
177 pub fn with_opencode() -> Result<Self, AppError> {
178 Self::with_opencode_builder().build()
179 }
180
181 /// Builder entry point for an opencode-backed embedder.
182 pub fn with_opencode_builder() -> LlmEmbeddingBuilder {
183 LlmEmbeddingBuilder::opencode_default()
184 }
185
186 /// v1.0.69 (G31): refuse to spawn if an API key is set. The CLI must use OAuth.
187 pub(crate) fn oauth_only_enforce() -> Result<(), AppError> {
188 if std::env::var("ANTHROPIC_API_KEY").is_ok() {
189 return Err(AppError::Validation(
190 crate::i18n::validation::anthropic_api_key_oauth_required(),
191 ));
192 }
193 if std::env::var("OPENAI_API_KEY").is_ok() {
194 return Err(AppError::Validation(
195 crate::i18n::validation::openai_api_key_oauth_required(),
196 ));
197 }
198 Ok(())
199 }
200
201 /// Embeds a single passage (chunk of a memory body).
202 pub fn embed_passage(&self, text: &str) -> Result<Vec<f32>, AppError> {
203 self.invoke_with_prefix(crate::constants::PASSAGE_PREFIX, text)
204 }
205
206 /// Embeds a single query with the query prompt prefix.
207 pub fn embed_query(&self, text: &str) -> Result<Vec<f32>, AppError> {
208 self.invoke_with_prefix(crate::constants::QUERY_PREFIX, text)
209 }
210
211 /// Stable label for the active embedding model (`flavour:model`).
212 pub fn model_label(&self) -> String {
213 format!("{}:{}", self.flavour.as_str(), self.model)
214 }
215
216 /// Returns the resolved embedding flavour of this client.
217 pub fn flavour(&self) -> EmbeddingFlavour {
218 self.flavour
219 }
220}
221
222#[cfg(test)]
223mod tests;