Skip to main content

steeldb/agent/
config.rs

1//! Provider configuration — pick the LLM backend at runtime. Bedrock (Claude) or a local
2//! Paddock/OpenAI-compatible server (e.g. Qwen3.5-2B). Each arm is gated by its cargo feature; a
3//! disabled backend returns a clear error rather than failing to compile.
4
5use crate::agent::provider::LlmProvider;
6
7/// Locate + parse the config file: `$STEELDB_CONFIG`, else `./steeldb.json`, else `~/.steeldb/config.json`.
8fn load_config_file() -> Option<serde_json::Value> {
9    let candidates = [
10        std::env::var("STEELDB_CONFIG").ok().map(std::path::PathBuf::from),
11        Some(std::path::PathBuf::from("steeldb.json")),
12        std::env::var("HOME").ok().map(|h| std::path::PathBuf::from(h).join(".steeldb/config.json")),
13    ];
14    for path in candidates.into_iter().flatten() {
15        if let Ok(bytes) = std::fs::read(&path) {
16            if let Ok(v) = serde_json::from_slice::<serde_json::Value>(&bytes) {
17                return Some(v);
18            }
19        }
20    }
21    None
22}
23
24#[derive(Clone, Debug)]
25pub enum ProviderConfig {
26    /// Claude on AWS Bedrock via the Converse API. `region: None` uses the environment/default chain.
27    Bedrock { model_id: String, region: Option<String> },
28    /// Local Paddock (OpenAI-compatible `/chat/completions`). `base_url` includes the `/v1` path.
29    Paddock { base_url: String, model: String, api_key: Option<String> },
30    /// Native in-process inference (pure-Rust candle): tuned Qwen3-1.7B + shipped LoRA, fully offline,
31    /// no server. `base_dir`/`lora_dir` = `None` resolves to the HF cache / bundled defaults at build.
32    Native { base_dir: Option<std::path::PathBuf>, lora_dir: Option<std::path::PathBuf> },
33    /// Cactus Needle: the 26M tool-use-native model via its `/generate` server (needle-code torch port).
34    Needle { base_url: String },
35}
36
37impl ProviderConfig {
38    pub fn bedrock_default() -> Self {
39        ProviderConfig::Bedrock {
40            model_id: "us.anthropic.claude-sonnet-4-5-20250929-v1:0".to_string(),
41            region: None,
42        }
43    }
44    pub fn paddock_default() -> Self {
45        ProviderConfig::Paddock {
46            base_url: "http://localhost:8080/v1".to_string(),
47            model: "Qwen/Qwen3.5-2B".to_string(),
48            api_key: None,
49        }
50    }
51
52    /// Resolve the provider in layers: **built-in defaults → config file → environment**. This is the
53    /// product path: a user points SteelDB at ANY OpenAI-compatible endpoint (OpenRouter, OpenAI,
54    /// Together, Groq, vLLM, a local server) with one small config file — no code, no rebuild.
55    ///
56    /// Config file (`./steeldb.json` or `$STEELDB_CONFIG` or `~/.steeldb/config.json`):
57    /// ```json
58    /// { "llm": { "provider": "openai", "base_url": "https://openrouter.ai/api/v1",
59    ///            "model": "anthropic/claude-3.5-sonnet", "api_key": "sk-or-..." } }
60    /// ```
61    /// `provider: "openai"` (or "paddock"/"local") → the OpenAI-compatible client; `"bedrock"` → Claude
62    /// on AWS. Env vars (`STEELDB_LLM`, `STEELDB_PADDOCK_URL`, `STEELDB_PADDOCK_MODEL`, `STEELDB_API_KEY`,
63    /// `STEELDB_BEDROCK_MODEL`, `AWS_REGION`) override the file.
64    pub fn resolve(default_base_url: &str, default_model: &str) -> ProviderConfig {
65        // 1) defaults (local Qwen unless the caller passes otherwise)
66        let mut base_url = default_base_url.to_string();
67        let mut model = default_model.to_string();
68        let mut api_key: Option<String> = None;
69        let mut bedrock = false;
70        let mut native = false;
71        // External (cloud/OpenAI-compatible) providers are opt-in only: SteelDB's default is the complete
72        // on-device path (Needle2 + deterministic template synthesis). Set true only when the user
73        // explicitly asks for an external OpenAI-compatible endpoint.
74        let mut paddock_explicit = false;
75        let mut needle_url = "http://localhost:8080".to_string();
76        let mut bedrock_model = "us.anthropic.claude-sonnet-4-5-20250929-v1:0".to_string();
77
78        // 2) config file
79        if let Some(cfg) = load_config_file() {
80            if let Some(llm) = cfg.get("llm") {
81                match llm.get("provider").and_then(|v| v.as_str()) {
82                    Some("bedrock") => bedrock = true,
83                    Some("native") => native = true,
84                    Some("needle") => {} // the default; nothing to flip
85                    Some("openai") | Some("paddock") | Some("local") => paddock_explicit = true,
86                    _ => {}
87                }
88                if let Some(u) = llm.get("needle_url").and_then(|v| v.as_str()) {
89                    needle_url = u.to_string();
90                }
91                if let Some(u) = llm.get("base_url").and_then(|v| v.as_str()) {
92                    base_url = u.to_string();
93                }
94                if let Some(m) = llm.get("model").and_then(|v| v.as_str()) {
95                    model = m.to_string();
96                    if bedrock {
97                        bedrock_model = m.to_string();
98                    }
99                }
100                if let Some(k) = llm.get("api_key").and_then(|v| v.as_str()) {
101                    api_key = Some(k.to_string());
102                }
103            }
104        }
105
106        // 3) env overrides
107        match std::env::var("STEELDB_LLM").as_deref() {
108            Ok("bedrock") => bedrock = true,
109            Ok("native") => native = true,
110            Ok("needle") => {
111                bedrock = false;
112                native = false;
113                paddock_explicit = false;
114            }
115            Ok("paddock") | Ok("openai") | Ok("local") => {
116                bedrock = false;
117                native = false;
118                paddock_explicit = true;
119            }
120            _ => {}
121        }
122        if let Ok(u) = std::env::var("STEELDB_NEEDLE_URL") {
123            needle_url = u;
124        }
125        if let Ok(u) = std::env::var("STEELDB_PADDOCK_URL") {
126            base_url = u;
127        }
128        if let Ok(m) = std::env::var("STEELDB_PADDOCK_MODEL") {
129            model = m;
130        }
131        if let Ok(m) = std::env::var("STEELDB_BEDROCK_MODEL") {
132            bedrock_model = m;
133            bedrock = true;
134        }
135        if let Ok(k) = std::env::var("STEELDB_API_KEY").or_else(|_| std::env::var("OPENROUTER_API_KEY")).or_else(|_| std::env::var("OPENAI_API_KEY")) {
136            api_key = Some(k);
137        }
138
139        // Priority: explicit external opt-ins first, else the complete on-device default (Needle2).
140        if bedrock {
141            ProviderConfig::Bedrock { model_id: bedrock_model, region: std::env::var("AWS_REGION").ok() }
142        } else if paddock_explicit {
143            ProviderConfig::Paddock { base_url, model, api_key }
144        } else if native {
145            ProviderConfig::Native { base_dir: None, lora_dir: None }
146        } else {
147            // Default = self-contained on-device: Needle2 drives tool selection, templates synthesise.
148            let _ = (base_url, model, api_key);
149            ProviderConfig::Needle { base_url: needle_url }
150        }
151    }
152
153    /// Read a config from env: `STEELDB_LLM=bedrock|paddock` plus backend-specific vars
154    /// (`STEELDB_BEDROCK_MODEL`, `AWS_REGION`; `STEELDB_PADDOCK_URL`, `STEELDB_PADDOCK_MODEL`).
155    pub fn from_env() -> ProviderConfig {
156        match std::env::var("STEELDB_LLM").as_deref() {
157            Ok("paddock") => {
158                let mut c = ProviderConfig::paddock_default();
159                if let ProviderConfig::Paddock { base_url, model, .. } = &mut c {
160                    if let Ok(u) = std::env::var("STEELDB_PADDOCK_URL") {
161                        *base_url = u;
162                    }
163                    if let Ok(m) = std::env::var("STEELDB_PADDOCK_MODEL") {
164                        *model = m;
165                    }
166                }
167                c
168            }
169            _ => {
170                let mut c = ProviderConfig::bedrock_default();
171                if let ProviderConfig::Bedrock { model_id, region } = &mut c {
172                    if let Ok(m) = std::env::var("STEELDB_BEDROCK_MODEL") {
173                        *model_id = m;
174                    }
175                    if let Ok(r) = std::env::var("AWS_REGION") {
176                        *region = Some(r);
177                    }
178                }
179                c
180            }
181        }
182    }
183
184    /// Construct the provider. Errors if the selected backend's feature is not compiled in.
185    pub async fn build(&self) -> Result<Box<dyn LlmProvider>, String> {
186        match self {
187            ProviderConfig::Bedrock { model_id, region } => {
188                #[cfg(feature = "bedrock")]
189                {
190                    let p = crate::agent::bedrock::BedrockProvider::new(model_id.clone(), region.clone()).await?;
191                    Ok(Box::new(p))
192                }
193                #[cfg(not(feature = "bedrock"))]
194                {
195                    let _ = (model_id, region);
196                    Err("bedrock backend not compiled in (build with --features bedrock)".to_string())
197                }
198            }
199            ProviderConfig::Paddock { base_url, model, api_key } => {
200                #[cfg(feature = "paddock")]
201                {
202                    // Wrap the local model in the Qwen harness (tool-call recovery + MECE/tool guidance)
203                    // so a small model completes the agent loop reliably.
204                    let p = crate::agent::paddock::PaddockProvider::new(base_url.clone(), model.clone(), api_key.clone());
205                    Ok(crate::agent::harness::QwenHarness::wrap(Box::new(p)))
206                }
207                #[cfg(not(feature = "paddock"))]
208                {
209                    let _ = (base_url, model, api_key);
210                    Err("paddock backend not compiled in (build with --features paddock)".to_string())
211                }
212            }
213            ProviderConfig::Native { base_dir, lora_dir } => {
214                #[cfg(feature = "native")]
215                {
216                    let base = base_dir.clone().or_else(native_base_dir).ok_or_else(|| {
217                        "native LLM base weights not found — set STEELDB_QWEN_BASE, run `steeldb pull`, or place Qwen3-1.7B under ~/.steeldb/models/qwen3-1.7b".to_string()
218                    })?;
219                    let lora = lora_dir.clone().or_else(|| {
220                        crate::paths::model_dir("lora-default", "STEELDB_LORA", "adapter_model.safetensors")
221                    });
222                    // Loading + LoRA merge is heavy and blocking; keep it off the async reactor.
223                    let (base2, lora2) = (base.clone(), lora.clone());
224                    let llm = tokio::task::spawn_blocking(move || crate::agent::native::NativeLlm::load(&base2, lora2.as_deref()))
225                        .await
226                        .map_err(|e| format!("native load task: {e}"))??;
227                    let p = crate::agent::native::NativeProvider::new(llm);
228                    Ok(crate::agent::harness::QwenHarness::wrap(Box::new(p)))
229                }
230                #[cfg(not(feature = "native"))]
231                {
232                    let _ = (base_dir, lora_dir);
233                    Err("native backend not compiled in (build with --features native)".to_string())
234                }
235            }
236            ProviderConfig::Needle { base_url } => {
237                #[cfg(feature = "needle")]
238                {
239                    // Needle emits structured tool-call JSON natively, so it is NOT wrapped in the Qwen
240                    // harness (which recovers Qwen's <tool_call> text blocks).
241                    Ok(Box::new(crate::agent::needle::NeedleProvider::new(base_url.clone())))
242                }
243                #[cfg(not(feature = "needle"))]
244                {
245                    let _ = base_url;
246                    Err("needle backend not compiled in (build with --features needle)".to_string())
247                }
248            }
249        }
250    }
251}
252
253/// True when native inference can actually run: base weights resolvable and the LoRA present. Callers
254/// use this to decide whether to default to offline native inference.
255#[cfg(feature = "native")]
256pub fn native_base_available() -> bool {
257    native_base_dir().is_some()
258}
259
260/// Resolve the Qwen3-1.7B base weights directory for native inference: `STEELDB_QWEN_BASE`, else
261/// `~/.steeldb/models/qwen3-1.7b`, else the latest Hugging Face cache snapshot. Base weights are large
262/// and fetched (not shipped); only our LoRA is bundled.
263#[cfg(feature = "native")]
264fn native_base_dir() -> Option<std::path::PathBuf> {
265    use std::path::PathBuf;
266    let has_model = |d: &PathBuf| d.join("config.json").exists() && d.join("tokenizer.json").exists();
267    if let Ok(p) = std::env::var("STEELDB_QWEN_BASE") {
268        let p = PathBuf::from(p);
269        if has_model(&p) {
270            return Some(p);
271        }
272    }
273    if let Ok(home) = std::env::var("HOME") {
274        let d = PathBuf::from(&home).join(".steeldb/models/qwen3-1.7b");
275        if has_model(&d) {
276            return Some(d);
277        }
278        // Hugging Face cache: models--Qwen--Qwen3-1.7B/snapshots/<hash>/
279        let snaps = PathBuf::from(&home).join(".cache/huggingface/hub/models--Qwen--Qwen3-1.7B/snapshots");
280        if let Ok(rd) = std::fs::read_dir(&snaps) {
281            let mut dirs: Vec<PathBuf> = rd.filter_map(|e| e.ok().map(|e| e.path())).filter(|p| has_model(p)).collect();
282            dirs.sort();
283            if let Some(d) = dirs.pop() {
284                return Some(d);
285            }
286        }
287    }
288    None
289}