Skip to main content

gobby_core/provisioning/
bootstrap.rs

1use super::*;
2
3// Bound nested YAML paths before flattening config into dotted keys. This keeps
4// hostile or accidental deeply nested input from driving unbounded recursion.
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
17impl EmbeddingBootstrap {
18    pub fn lm_studio() -> Self {
19        Self {
20            provider: "lm-studio".to_string(),
21            api_base: DEFAULT_LM_STUDIO_API_BASE.to_string(),
22            model: DEFAULT_LM_STUDIO_MODEL.to_string(),
23            vector_dim: DEFAULT_EMBEDDING_VECTOR_DIM,
24            query_prefix: None,
25            api_key: None,
26        }
27    }
28
29    pub fn ollama() -> Self {
30        Self {
31            provider: "ollama".to_string(),
32            api_base: DEFAULT_OLLAMA_API_BASE.to_string(),
33            model: DEFAULT_OLLAMA_MODEL.to_string(),
34            vector_dim: DEFAULT_EMBEDDING_VECTOR_DIM,
35            query_prefix: None,
36            api_key: None,
37        }
38    }
39}
40
41pub fn write_standalone_bootstrap(
42    path: &Path,
43    database_url: &str,
44    options: &DockerServiceOptions,
45    compose_file: Option<&Path>,
46    embedding: Option<&EmbeddingBootstrap>,
47) -> anyhow::Result<StandaloneConfig> {
48    let mut config = StandaloneConfig::empty();
49    config.set("databases.postgres.dsn", database_url);
50    config.set("databases.falkordb.host", &options.falkordb_host);
51    config.set("databases.falkordb.port", options.falkordb_port.to_string());
52    config.set("databases.falkordb.password", &options.falkordb_password);
53    config.set("databases.qdrant.url", options.qdrant_url());
54    if let Some(embedding) = embedding {
55        config.set(embedding_keys::AI_PROVIDER, &embedding.provider);
56        config.set(embedding_keys::AI_API_BASE, &embedding.api_base);
57        config.set(embedding_keys::AI_MODEL, &embedding.model);
58        config.set(embedding_keys::AI_DIM, embedding.vector_dim.to_string());
59        if let Some(query_prefix) = &embedding.query_prefix {
60            config.set(embedding_keys::AI_QUERY_PREFIX, query_prefix);
61        }
62        if let Some(api_key) = &embedding.api_key {
63            config.set(embedding_keys::AI_API_KEY, api_key);
64        }
65    }
66    if let Some(compose_file) = compose_file {
67        config.set("services.compose_file", compose_file.display().to_string());
68    }
69    config.write_at(path)?;
70    Ok(config)
71}
72
73pub(crate) fn flatten_yaml_value(
74    prefix: Option<&str>,
75    value: &serde_yaml::Value,
76    output: &mut BTreeMap<String, String>,
77) -> anyhow::Result<()> {
78    flatten_yaml_value_at_depth(prefix, value, output, 0)
79}
80
81fn flatten_yaml_value_at_depth(
82    prefix: Option<&str>,
83    value: &serde_yaml::Value,
84    output: &mut BTreeMap<String, String>,
85    depth: usize,
86) -> anyhow::Result<()> {
87    // Depth starts at the root value, so a path with 64 nested mappings is the
88    // last accepted level; the next recursive descent fails before matching.
89    if depth > MAX_YAML_FLATTEN_DEPTH {
90        anyhow::bail!(
91            "gcore.yaml path `{}` exceeds maximum depth of {MAX_YAML_FLATTEN_DEPTH}",
92            yaml_path(prefix)
93        );
94    }
95    match value {
96        serde_yaml::Value::Null => Ok(()),
97        serde_yaml::Value::Mapping(mapping) => {
98            for (key, value) in mapping {
99                let Some(key) = key.as_str() else {
100                    anyhow::bail!(
101                        "gcore.yaml path `{}` has a non-string key",
102                        yaml_path(prefix)
103                    );
104                };
105                let joined = match prefix {
106                    Some(prefix) if !prefix.is_empty() => format!("{prefix}.{key}"),
107                    _ => key.to_string(),
108                };
109                match value {
110                    serde_yaml::Value::Mapping(_) => {
111                        // Dotted keys may be flattened prefixes such as
112                        // `ai.embeddings`, so nested values still need recursion.
113                        flatten_yaml_value_at_depth(Some(&joined), value, output, depth + 1)?;
114                    }
115                    _ => {
116                        if let Some(text) = scalar_to_string(&joined, value)? {
117                            output.insert(joined, text);
118                        }
119                    }
120                }
121            }
122            Ok(())
123        }
124        _ => {
125            let Some(prefix) = prefix else {
126                anyhow::bail!("gcore.yaml path `<root>` must be a mapping");
127            };
128            if let Some(text) = scalar_to_string(prefix, value)? {
129                output.insert(prefix.to_string(), text);
130            }
131            Ok(())
132        }
133    }
134}
135
136fn scalar_to_string(path: &str, value: &serde_yaml::Value) -> anyhow::Result<Option<String>> {
137    Ok(match value {
138        serde_yaml::Value::Null => None,
139        serde_yaml::Value::String(value) => Some(value.clone()),
140        serde_yaml::Value::Bool(value) => Some(value.to_string()),
141        serde_yaml::Value::Number(value) => Some(value.to_string()),
142        serde_yaml::Value::Sequence(_) => {
143            anyhow::bail!("gcore.yaml path `{path}` cannot be a sequence")
144        }
145        serde_yaml::Value::Mapping(_) => {
146            anyhow::bail!("gcore.yaml path `{path}` cannot be a mapping")
147        }
148        // Tagged values may wrap scalars or collections; preserve serde_yaml's
149        // spelling instead of inventing a lossy custom representation.
150        other => Some(
151            serde_yaml::to_string(other)
152                .with_context(|| format!("convert gcore.yaml path `{path}` to string"))?
153                .trim()
154                .to_string(),
155        ),
156    })
157}
158
159fn yaml_path(prefix: Option<&str>) -> &str {
160    prefix.filter(|path| !path.is_empty()).unwrap_or("<root>")
161}
162
163#[cfg(test)]
164mod tests {
165    use super::*;
166
167    fn flatten(contents: &str) -> anyhow::Result<BTreeMap<String, String>> {
168        let value = serde_yaml::from_str(contents)?;
169        let mut output = BTreeMap::new();
170        flatten_yaml_value(None, &value, &mut output)?;
171        Ok(output)
172    }
173
174    #[test]
175    fn flatten_yaml_errors_include_root_path() {
176        let error = flatten("true").expect_err("root scalar must fail");
177
178        assert!(error.to_string().contains("`<root>`"));
179    }
180
181    #[test]
182    fn flatten_yaml_errors_include_mapping_path_for_non_string_keys() {
183        let error = flatten("ai:\n  ? [bad]\n  : value\n").expect_err("non-string key must fail");
184
185        assert!(error.to_string().contains("`ai`"));
186    }
187
188    #[test]
189    fn flatten_yaml_errors_include_scalar_path() {
190        let error = flatten("ai:\n  embeddings:\n    provider:\n      - openai\n")
191            .expect_err("sequence scalar must fail");
192
193        assert!(error.to_string().contains("`ai.embeddings.provider`"));
194    }
195
196    #[test]
197    fn flatten_yaml_depth_errors_include_current_path() {
198        let mut contents = String::new();
199        for index in 0..=MAX_YAML_FLATTEN_DEPTH + 1 {
200            contents.push_str(&"  ".repeat(index));
201            contents.push_str(&format!("k{index}:\n"));
202        }
203
204        let error = flatten(&contents).expect_err("too-deep nesting must fail");
205
206        assert!(error.to_string().contains("k0.k1"));
207    }
208}