gobby_core/provisioning/
bootstrap.rs

1use super::*;
2
3// MAX_YAML_FLATTEN_DEPTH is generous for legitimate config nesting while still
4// preventing infinite recursion or stack overflow on malformed YAML.
5const MAX_YAML_FLATTEN_DEPTH: usize = 64;
6
7#[derive(Debug, Clone, PartialEq, Eq)]
8pub struct EmbeddingBootstrap {
9    pub provider: String,
10    pub api_base: String,
11    pub model: String,
12    pub vector_dim: usize,
13    pub query_prefix: Option<String>,
14    pub api_key: Option<String>,
15}
16
17#[derive(Debug, Clone, PartialEq, Eq)]
18pub struct TextGenerationBootstrap {
19    pub api_base: String,
20    pub model: String,
21    pub api_key: Option<String>,
22}
23
24impl EmbeddingBootstrap {
25    pub fn lm_studio() -> Self {
26        Self {
27            provider: "lmstudio".to_string(),
28            api_base: DEFAULT_LM_STUDIO_API_BASE.to_string(),
29            model: DEFAULT_LM_STUDIO_MODEL.to_string(),
30            vector_dim: DEFAULT_EMBEDDING_VECTOR_DIM,
31            query_prefix: None,
32            api_key: None,
33        }
34    }
35
36    pub fn ollama() -> Self {
37        Self {
38            provider: "ollama".to_string(),
39            api_base: DEFAULT_OLLAMA_API_BASE.to_string(),
40            model: DEFAULT_OLLAMA_MODEL.to_string(),
41            vector_dim: DEFAULT_EMBEDDING_VECTOR_DIM,
42            query_prefix: None,
43            api_key: None,
44        }
45    }
46}
47
48impl TextGenerationBootstrap {
49    pub fn from_embedding(embedding: &EmbeddingBootstrap) -> Self {
50        Self::from_endpoint(
51            Some(&embedding.provider),
52            embedding.api_base.clone(),
53            embedding.api_key.clone(),
54        )
55    }
56
57    pub fn from_endpoint(
58        provider: Option<&str>,
59        api_base: impl Into<String>,
60        api_key: Option<String>,
61    ) -> Self {
62        let api_base = api_base.into();
63        Self {
64            model: default_text_model(provider, &api_base).to_string(),
65            api_base,
66            api_key,
67        }
68    }
69}
70
71pub fn apply_text_generation_bootstrap(
72    config: &mut StandaloneConfig,
73    text_generation: &TextGenerationBootstrap,
74) {
75    config.set(ai_keys::TEXT_GENERATE_ROUTING, "direct");
76    config.set(ai_keys::TEXT_GENERATE_API_BASE, &text_generation.api_base);
77    config.set(ai_keys::TEXT_GENERATE_MODEL, &text_generation.model);
78    config.remove(ai_keys::TEXT_GENERATE_TRANSPORT);
79    config.remove(ai_keys::TEXT_GENERATE_PROVIDER);
80    config.remove(ai_keys::TEXT_GENERATE_PROFILE);
81    // Clear the daemon-only verify profile and any stale verify overrides so a
82    // fresh direct bootstrap lets the verify pass fall back to the model/api_key
83    // set above (see CapabilityBinding::verify_*).
84    config.remove(ai_keys::TEXT_GENERATE_VERIFY_PROFILE);
85    config.remove(ai_keys::TEXT_GENERATE_VERIFY_MODEL);
86    config.remove(ai_keys::TEXT_GENERATE_VERIFY_API_KEY);
87    match text_generation.api_key.as_deref() {
88        Some(api_key) => config.set(ai_keys::TEXT_GENERATE_API_KEY, api_key),
89        None => config.remove(ai_keys::TEXT_GENERATE_API_KEY),
90    }
91}
92
93fn default_text_model(provider: Option<&str>, api_base: &str) -> &'static str {
94    let provider = provider.map(|value| value.trim().to_ascii_lowercase());
95    let api_base = api_base.trim().trim_end_matches('/');
96    if provider.as_deref() == Some("ollama")
97        || api_base == DEFAULT_OLLAMA_API_BASE.trim_end_matches('/')
98    {
99        DEFAULT_OLLAMA_TEXT_MODEL
100    } else {
101        DEFAULT_LM_STUDIO_TEXT_MODEL
102    }
103}
104
105pub fn write_standalone_bootstrap(
106    path: &Path,
107    database_url: &str,
108    options: &DockerServiceOptions,
109    compose_file: Option<&Path>,
110    embedding: Option<&EmbeddingBootstrap>,
111) -> anyhow::Result<StandaloneConfig> {
112    let mut config = StandaloneConfig::empty();
113    config.set("databases.postgres.dsn", database_url);
114    config.set("databases.falkordb.host", &options.falkordb_host);
115    config.set("databases.falkordb.port", options.falkordb_port.to_string());
116    config.set("databases.falkordb.password", &options.falkordb_password);
117    config.set("databases.qdrant.url", options.qdrant_url());
118    if let Some(embedding) = embedding {
119        config.set(embedding_keys::AI_PROVIDER, &embedding.provider);
120        config.set(embedding_keys::AI_API_BASE, &embedding.api_base);
121        config.set(embedding_keys::AI_MODEL, &embedding.model);
122        config.set(embedding_keys::AI_DIM, embedding.vector_dim.to_string());
123        if let Some(query_prefix) = &embedding.query_prefix {
124            config.set(embedding_keys::AI_QUERY_PREFIX, query_prefix);
125        }
126        if let Some(api_key) = &embedding.api_key {
127            config.set(embedding_keys::AI_API_KEY, api_key);
128        }
129        apply_text_generation_bootstrap(
130            &mut config,
131            &TextGenerationBootstrap::from_embedding(embedding),
132        );
133    }
134    if let Some(compose_file) = compose_file {
135        config.set("services.compose_file", compose_file.display().to_string());
136    }
137    config.write_at(path)?;
138    Ok(config)
139}
140
141pub(crate) fn flatten_yaml_value(
142    prefix: Option<&str>,
143    value: &serde_yaml::Value,
144    output: &mut BTreeMap<String, String>,
145) -> anyhow::Result<()> {
146    flatten_yaml_value_at_depth(prefix, value, output, 0)
147}
148
149fn flatten_yaml_value_at_depth(
150    prefix: Option<&str>,
151    value: &serde_yaml::Value,
152    output: &mut BTreeMap<String, String>,
153    depth: usize,
154) -> anyhow::Result<()> {
155    // This MAX_YAML_FLATTEN_DEPTH check runs before flattening each nested YAML
156    // value, so deeply malformed inputs fail before unbounded recursion.
157    if depth > MAX_YAML_FLATTEN_DEPTH {
158        anyhow::bail!(
159            "gcore.yaml path `{}` exceeds maximum depth of {MAX_YAML_FLATTEN_DEPTH}",
160            yaml_path(prefix)
161        );
162    }
163    match value {
164        serde_yaml::Value::Null => Ok(()),
165        serde_yaml::Value::Mapping(mapping) => {
166            for (key, value) in mapping {
167                let Some(key) = key.as_str() else {
168                    anyhow::bail!(
169                        "gcore.yaml path `{}` has a non-string key",
170                        yaml_path(prefix)
171                    );
172                };
173                let joined = match prefix {
174                    Some(prefix) if !prefix.is_empty() => format!("{prefix}.{key}"),
175                    _ => key.to_string(),
176                };
177                match value {
178                    serde_yaml::Value::Mapping(_) => {
179                        // Dotted keys may be flattened prefixes such as
180                        // `ai.embeddings`, so nested values still need recursion.
181                        flatten_yaml_value_at_depth(Some(&joined), value, output, depth + 1)?;
182                    }
183                    _ => {
184                        if let Some(text) = scalar_to_string(&joined, value)? {
185                            output.insert(joined, text);
186                        }
187                    }
188                }
189            }
190            Ok(())
191        }
192        _ => {
193            let Some(prefix) = prefix else {
194                anyhow::bail!("gcore.yaml path `<root>` must be a mapping");
195            };
196            if let Some(text) = scalar_to_string(prefix, value)? {
197                output.insert(prefix.to_string(), text);
198            }
199            Ok(())
200        }
201    }
202}
203
204fn scalar_to_string(path: &str, value: &serde_yaml::Value) -> anyhow::Result<Option<String>> {
205    Ok(match value {
206        serde_yaml::Value::Null => None,
207        serde_yaml::Value::String(value) => Some(value.clone()),
208        serde_yaml::Value::Bool(value) => Some(value.to_string()),
209        serde_yaml::Value::Number(value) => Some(value.to_string()),
210        serde_yaml::Value::Sequence(_) => {
211            anyhow::bail!("gcore.yaml path `{path}` cannot be a sequence")
212        }
213        serde_yaml::Value::Mapping(_) => {
214            anyhow::bail!("gcore.yaml path `{path}` cannot be a mapping")
215        }
216        // Tagged values may wrap scalars or collections; preserve serde_yaml's
217        // spelling instead of inventing a lossy custom representation.
218        other => Some(
219            serde_yaml::to_string(other)
220                .with_context(|| format!("convert gcore.yaml path `{path}` to string"))?
221                .trim()
222                .to_string(),
223        ),
224    })
225}
226
227fn yaml_path(prefix: Option<&str>) -> &str {
228    prefix.filter(|path| !path.is_empty()).unwrap_or("<root>")
229}
230
231#[cfg(test)]
232mod tests {
233    use super::*;
234
235    fn flatten(contents: &str) -> anyhow::Result<BTreeMap<String, String>> {
236        let value = serde_yaml::from_str(contents)?;
237        let mut output = BTreeMap::new();
238        flatten_yaml_value(None, &value, &mut output)?;
239        Ok(output)
240    }
241
242    #[test]
243    fn flatten_yaml_errors_include_root_path() {
244        let error = flatten("true").expect_err("root scalar must fail");
245
246        assert!(error.to_string().contains("`<root>`"));
247    }
248
249    #[test]
250    fn flatten_yaml_errors_include_mapping_path_for_non_string_keys() {
251        let error = flatten("ai:\n  ? [bad]\n  : value\n").expect_err("non-string key must fail");
252
253        assert!(error.to_string().contains("`ai`"));
254    }
255
256    #[test]
257    fn flatten_yaml_errors_include_scalar_path() {
258        let error = flatten("ai:\n  embeddings:\n    provider:\n      - openai\n")
259            .expect_err("sequence scalar must fail");
260
261        assert!(error.to_string().contains("`ai.embeddings.provider`"));
262    }
263
264    #[test]
265    fn flatten_yaml_depth_errors_include_current_path() {
266        let mut contents = String::new();
267        for index in 0..=MAX_YAML_FLATTEN_DEPTH + 1 {
268            contents.push_str(&"  ".repeat(index));
269            contents.push_str(&format!("k{index}:\n"));
270        }
271
272        let error = flatten(&contents).expect_err("too-deep nesting must fail");
273
274        assert!(error.to_string().contains("k0.k1"));
275    }
276}
gobby_core/provisioning/bootstrap.rs

gobby_core/provisioning/
bootstrap.rs
