gobby_code/commands/
embeddings_doctor.rs

1use std::time::Duration;
2
3use serde::{Deserialize, Serialize};
4use serde_json::{Value, json};
5
6use crate::config::{self, Context, EmbeddingConfig, EmbeddingConfigDetails};
7use crate::db;
8use crate::output;
9use crate::utils::api_key_fingerprint;
10use crate::vector::code_symbols::probe_embedding_dim;
11
12const EXIT_HEALTHY: u8 = 0;
13const EXIT_CONFIG_MISSING: u8 = 10;
14const EXIT_DRIFT: u8 = 11;
15const EXIT_TRANSPORT: u8 = 20;
16const DAEMON_DOCTOR_PATH: &str = "/api/embeddings/doctor";
17
18#[derive(Debug)]
19pub struct EmbeddingsDoctorExit {
20    payload: EmbeddingsDoctorReport,
21    exit_code: u8,
22}
23
24impl EmbeddingsDoctorExit {
25    pub fn exit_code(&self) -> u8 {
26        self.exit_code
27    }
28
29    pub fn print(&self) -> anyhow::Result<()> {
30        output::print_json(&self.payload)
31    }
32}
33
34impl std::fmt::Display for EmbeddingsDoctorExit {
35    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
36        write!(f, "embeddings doctor failed with exit {}", self.exit_code)
37    }
38}
39
40impl std::error::Error for EmbeddingsDoctorExit {}
41
42#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
43pub(crate) struct EmbeddingsDoctorReport {
44    pub endpoint: Option<String>,
45    pub model: Option<String>,
46    pub dim: Option<usize>,
47    #[serde(skip_serializing_if = "Option::is_none")]
48    pub probe_error: Option<String>,
49    pub api_key_present: bool,
50    pub api_key_fingerprint: Option<String>,
51    pub namespace_resolved: Option<String>,
52    pub source: Option<String>,
53    pub agrees: Option<bool>,
54    pub drift: Option<Vec<EmbeddingsDoctorDrift>>,
55}
56
57#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
58pub(crate) struct EmbeddingsDoctorDrift {
59    pub field: String,
60    #[serde(rename = "self")]
61    pub self_value: Value,
62    pub peer: Value,
63}
64
65#[derive(Debug, Clone, PartialEq, Eq, Deserialize)]
66struct PeerDoctorReport {
67    endpoint: Option<String>,
68    model: Option<String>,
69    dim: Option<usize>,
70}
71
72#[derive(Debug, Clone, PartialEq, Eq)]
73enum PeerDoctorOutcome {
74    Absent,
75    Present(PeerDoctorReport),
76    TransportError(String),
77}
78
79pub fn run(ctx: &Context) -> anyhow::Result<()> {
80    let mut conn = db::connect_readonly(&ctx.database_url)?;
81    let resolution = config::resolve_embedding_config_details(
82        &mut conn,
83        config::read_standalone_config_optional(),
84    )?;
85    let peer = fetch_daemon_peer(ctx.daemon_url.as_deref());
86    let (payload, exit_code) =
87        build_doctor_report(resolution, ctx.code_vectors.vector_dim, probe_dim, peer);
88
89    if exit_code == EXIT_HEALTHY {
90        output::print_json(&payload)?;
91        Ok(())
92    } else {
93        Err(EmbeddingsDoctorExit { payload, exit_code }.into())
94    }
95}
96
97fn probe_dim(config: &EmbeddingConfig) -> Result<usize, String> {
98    probe_embedding_dim(config).map_err(|error| error.to_string())
99}
100
101fn build_doctor_report(
102    resolution: Option<EmbeddingConfigDetails>,
103    configured_dim: Option<usize>,
104    probe: impl FnOnce(&EmbeddingConfig) -> Result<usize, String>,
105    peer: PeerDoctorOutcome,
106) -> (EmbeddingsDoctorReport, u8) {
107    let Some(resolution) = resolution else {
108        return (
109            EmbeddingsDoctorReport {
110                endpoint: None,
111                model: None,
112                dim: None,
113                probe_error: None,
114                api_key_present: false,
115                api_key_fingerprint: None,
116                namespace_resolved: None,
117                source: None,
118                agrees: None,
119                drift: None,
120            },
121            EXIT_CONFIG_MISSING,
122        );
123    };
124
125    let dim = match configured_dim {
126        Some(dim) => Some(dim),
127        None => match probe(&resolution.config) {
128            Ok(dim) => Some(dim),
129            Err(error) => {
130                let mut report = report_without_peer(&resolution, None);
131                report.probe_error = Some(error);
132                return (report, EXIT_TRANSPORT);
133            }
134        },
135    };
136
137    match peer {
138        PeerDoctorOutcome::Absent => (report_without_peer(&resolution, dim), EXIT_HEALTHY),
139        PeerDoctorOutcome::TransportError(_) => {
140            (report_without_peer(&resolution, dim), EXIT_TRANSPORT)
141        }
142        PeerDoctorOutcome::Present(peer) => {
143            let drift = drift_fields(&resolution.config, dim, &peer);
144            if drift.is_empty() {
145                (
146                    EmbeddingsDoctorReport {
147                        agrees: Some(true),
148                        drift: None,
149                        ..base_report(&resolution, dim)
150                    },
151                    EXIT_HEALTHY,
152                )
153            } else {
154                (
155                    EmbeddingsDoctorReport {
156                        agrees: Some(false),
157                        drift: Some(drift),
158                        ..base_report(&resolution, dim)
159                    },
160                    EXIT_DRIFT,
161                )
162            }
163        }
164    }
165}
166
167fn report_without_peer(
168    resolution: &EmbeddingConfigDetails,
169    dim: Option<usize>,
170) -> EmbeddingsDoctorReport {
171    EmbeddingsDoctorReport {
172        agrees: None,
173        drift: None,
174        ..base_report(resolution, dim)
175    }
176}
177
178fn base_report(resolution: &EmbeddingConfigDetails, dim: Option<usize>) -> EmbeddingsDoctorReport {
179    EmbeddingsDoctorReport {
180        endpoint: Some(resolution.config.api_base.clone()),
181        model: Some(resolution.config.model.clone()),
182        dim,
183        probe_error: None,
184        api_key_present: resolution.config.api_key.is_some(),
185        api_key_fingerprint: resolution
186            .config
187            .api_key
188            .as_deref()
189            .map(api_key_fingerprint),
190        namespace_resolved: Some(resolution.namespace.to_string()),
191        source: Some(resolution.source.to_string()),
192        agrees: None,
193        drift: None,
194    }
195}
196
197fn drift_fields(
198    config: &EmbeddingConfig,
199    dim: Option<usize>,
200    peer: &PeerDoctorReport,
201) -> Vec<EmbeddingsDoctorDrift> {
202    let mut drift = Vec::new();
203    push_drift(
204        &mut drift,
205        "endpoint",
206        Some(config.api_base.as_str()),
207        peer.endpoint.as_deref(),
208    );
209    push_drift(
210        &mut drift,
211        "model",
212        Some(config.model.as_str()),
213        peer.model.as_deref(),
214    );
215    if dim != peer.dim {
216        drift.push(EmbeddingsDoctorDrift {
217            field: "dim".to_string(),
218            self_value: dim.map_or(Value::Null, |value| json!(value)),
219            peer: peer.dim.map_or(Value::Null, |value| json!(value)),
220        });
221    }
222    drift
223}
224
225fn push_drift(
226    drift: &mut Vec<EmbeddingsDoctorDrift>,
227    field: &'static str,
228    self_value: Option<&str>,
229    peer_value: Option<&str>,
230) {
231    if self_value == peer_value {
232        return;
233    }
234    drift.push(EmbeddingsDoctorDrift {
235        field: field.to_string(),
236        self_value: self_value.map_or(Value::Null, |value| json!(value)),
237        peer: peer_value.map_or(Value::Null, |value| json!(value)),
238    });
239}
240
241fn fetch_daemon_peer(daemon_url: Option<&str>) -> PeerDoctorOutcome {
242    let Some(daemon_url) = daemon_url else {
243        return PeerDoctorOutcome::Absent;
244    };
245    let url = format!("{}{}", daemon_url.trim_end_matches('/'), DAEMON_DOCTOR_PATH);
246    let client = match reqwest::blocking::Client::builder()
247        .timeout(Duration::from_secs(2))
248        .build()
249    {
250        Ok(client) => client,
251        Err(error) => {
252            return PeerDoctorOutcome::TransportError(format!("build HTTP client: {error}"));
253        }
254    };
255    let response = match client.get(url).send() {
256        Ok(response) => response,
257        Err(error) => {
258            return PeerDoctorOutcome::TransportError(format!("send peer doctor request: {error}"));
259        }
260    };
261    if response.status() == reqwest::StatusCode::NOT_FOUND {
262        return PeerDoctorOutcome::Absent;
263    }
264    if !response.status().is_success() {
265        return PeerDoctorOutcome::TransportError(format!(
266            "peer doctor returned HTTP {}",
267            response.status()
268        ));
269    }
270    match response.json::<PeerDoctorReport>() {
271        Ok(report) => PeerDoctorOutcome::Present(report),
272        Err(error) => {
273            PeerDoctorOutcome::TransportError(format!("parse peer doctor response: {error}"))
274        }
275    }
276}
277
278#[cfg(test)]
279mod tests {
280    use super::*;
281    use gobby_core::config::embedding_keys;
282
283    fn details(api_key: Option<&str>) -> EmbeddingConfigDetails {
284        EmbeddingConfigDetails {
285            config: EmbeddingConfig {
286                api_base: "http://embeddings.local/v1".to_string(),
287                model: "embed-small".to_string(),
288                api_key: api_key.map(str::to_string),
289                query_prefix: None,
290                timeout_seconds: 10,
291            },
292            namespace: embedding_keys::AI_NAMESPACE,
293            source: "config_store",
294        }
295    }
296
297    #[test]
298    fn doctor_json_and_exit_codes() {
299        let (missing, code) = build_doctor_report(
300            None,
301            None,
302            |_| unreachable!("no config should not probe"),
303            PeerDoctorOutcome::Absent,
304        );
305        assert_eq!(code, EXIT_CONFIG_MISSING);
306        assert_eq!(missing.namespace_resolved, None);
307
308        let (healthy, code) = build_doctor_report(
309            Some(details(Some("secret-key"))),
310            Some(768),
311            |_| unreachable!("configured dim should not probe"),
312            PeerDoctorOutcome::Absent,
313        );
314        assert_eq!(code, EXIT_HEALTHY);
315        assert_eq!(
316            healthy.endpoint.as_deref(),
317            Some("http://embeddings.local/v1")
318        );
319        assert_eq!(healthy.dim, Some(768));
320        assert!(healthy.api_key_present);
321        assert_eq!(
322            healthy.api_key_fingerprint,
323            Some(api_key_fingerprint("secret-key"))
324        );
325        assert_eq!(
326            serde_json::to_value(&healthy).expect("doctor report serializes")["agrees"],
327            Value::Null
328        );
329
330        let (drift, code) = build_doctor_report(
331            Some(details(None)),
332            None,
333            |_| Ok(768),
334            PeerDoctorOutcome::Present(PeerDoctorReport {
335                endpoint: Some("http://other.local/v1".to_string()),
336                model: Some("embed-small".to_string()),
337                dim: Some(1536),
338            }),
339        );
340        assert_eq!(code, EXIT_DRIFT);
341        assert_eq!(drift.agrees, Some(false));
342        assert_eq!(
343            drift
344                .drift
345                .as_ref()
346                .expect("drift fields")
347                .iter()
348                .map(|field| field.field.as_str())
349                .collect::<Vec<_>>(),
350            vec!["endpoint", "dim"]
351        );
352
353        let (transport, code) = build_doctor_report(
354            Some(details(None)),
355            None,
356            |_| Err("probe failed".to_string()),
357            PeerDoctorOutcome::Absent,
358        );
359        assert_eq!(code, EXIT_TRANSPORT);
360        assert_eq!(transport.dim, None);
361        assert_eq!(transport.probe_error.as_deref(), Some("probe failed"));
362    }
363}
gobby_code/commands/embeddings_doctor.rs

gobby_code/commands/
embeddings_doctor.rs
