Skip to main content

loadmaster_probe/
db.rs

1//! Eval results in SQLite (`~/.loadmaster/eval.db`).
2
3use std::path::Path;
4
5use rusqlite::{Connection, params};
6use serde::Serialize;
7
8use crate::eval::RowResult;
9
10pub struct Db(Connection);
11
12#[derive(Debug, Clone, Serialize)]
13pub struct RunMeta {
14    pub id: i64,
15    pub created_at: String,
16    pub url: String,
17    pub model: String,
18    pub runtime: String,
19    pub runtime_version: Option<String>,
20    pub prompts: String,
21    pub prompts_blake3: String,
22    pub n: i64,
23}
24
25/// One stored row, as read back.
26#[derive(Debug, Clone, PartialEq)]
27pub struct StoredRow {
28    pub row_id: String,
29    pub ok: bool,
30    pub schema_valid: Option<bool>,
31    pub latency_ms: u64,
32    pub tokens_per_s: Option<f64>,
33    pub response_blake3: String,
34}
35
36/// What a run scored. `accuracy` and `schema_*` depend only on the outputs, so a
37/// deterministic runtime reproduces them exactly; timings are measurements and vary.
38#[derive(Debug, Clone, PartialEq, Serialize)]
39pub struct Summary {
40    pub n: u64,
41    pub ok: u64,
42    pub schema_n: u64,
43    pub schema_ok: u64,
44    pub tokens_per_s_mean: Option<f64>,
45    pub latency_p50_ms: u64,
46    pub latency_p95_ms: u64,
47}
48
49impl Db {
50    pub fn open(path: &Path) -> rusqlite::Result<Self> {
51        if let Some(dir) = path.parent() {
52            let _ = std::fs::create_dir_all(dir);
53        }
54        let c = Connection::open(path)?;
55        c.execute_batch(
56            "CREATE TABLE IF NOT EXISTS runs (
57                id INTEGER PRIMARY KEY, created_at TEXT NOT NULL, url TEXT NOT NULL,
58                model TEXT NOT NULL, runtime TEXT NOT NULL, runtime_version TEXT,
59                prompts TEXT NOT NULL, prompts_blake3 TEXT NOT NULL, probe_version TEXT NOT NULL,
60                n INTEGER NOT NULL);
61             CREATE TABLE IF NOT EXISTS results (
62                run_id INTEGER NOT NULL REFERENCES runs(id), idx INTEGER NOT NULL, row_id TEXT NOT NULL,
63                ok INTEGER NOT NULL, schema_valid INTEGER, tool_ok INTEGER, latency_ms INTEGER NOT NULL,
64                completion_tokens INTEGER, tokens_per_s REAL, error TEXT, response_blake3 TEXT NOT NULL,
65                PRIMARY KEY (run_id, idx));",
66        )?;
67        Ok(Db(c))
68    }
69
70    #[allow(clippy::too_many_arguments)]
71    pub fn insert(
72        &mut self,
73        created_at: &str,
74        url: &str,
75        model: &str,
76        runtime: &str,
77        runtime_version: Option<&str>,
78        prompts: &str,
79        prompts_blake3: &str,
80        rows: &[RowResult],
81    ) -> rusqlite::Result<i64> {
82        let tx = self.0.transaction()?;
83        tx.execute(
84            "INSERT INTO runs (created_at, url, model, runtime, runtime_version, prompts, prompts_blake3, probe_version, n)
85             VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9)",
86            params![created_at, url, model, runtime, runtime_version, prompts, prompts_blake3, crate::PROBE_VERSION, rows.len() as i64],
87        )?;
88        let id = tx.last_insert_rowid();
89        for r in rows {
90            tx.execute(
91                "INSERT INTO results VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10, ?11)",
92                params![
93                    id,
94                    r.idx as i64,
95                    r.id,
96                    r.ok,
97                    r.schema_valid,
98                    r.tool_ok,
99                    r.latency_ms as i64,
100                    r.completion_tokens.map(|t| t as i64),
101                    r.tokens_per_s,
102                    r.error,
103                    r.response_blake3
104                ],
105            )?;
106        }
107        tx.commit()?;
108        Ok(id)
109    }
110
111    pub fn run(&self, id: i64) -> rusqlite::Result<RunMeta> {
112        self.0.query_row(
113            "SELECT id, created_at, url, model, runtime, runtime_version, prompts, prompts_blake3, n FROM runs WHERE id = ?1",
114            [id],
115            |r| {
116                Ok(RunMeta {
117                    id: r.get(0)?,
118                    created_at: r.get(1)?,
119                    url: r.get(2)?,
120                    model: r.get(3)?,
121                    runtime: r.get(4)?,
122                    runtime_version: r.get(5)?,
123                    prompts: r.get(6)?,
124                    prompts_blake3: r.get(7)?,
125                    n: r.get(8)?,
126                })
127            },
128        )
129    }
130
131    pub fn results(&self, id: i64) -> rusqlite::Result<Vec<StoredRow>> {
132        let mut st = self.0.prepare(
133            "SELECT row_id, ok, schema_valid, latency_ms, tokens_per_s, response_blake3 FROM results WHERE run_id = ?1 ORDER BY idx",
134        )?;
135        st.query_map([id], |r| {
136            Ok(StoredRow {
137                row_id: r.get(0)?,
138                ok: r.get(1)?,
139                schema_valid: r.get(2)?,
140                latency_ms: r.get::<_, i64>(3)? as u64,
141                tokens_per_s: r.get(4)?,
142                response_blake3: r.get(5)?,
143            })
144        })?
145        .collect()
146    }
147
148    pub fn summary(&self, id: i64) -> rusqlite::Result<Summary> {
149        let rows = self.results(id)?;
150        let mut lat: Vec<u64> = rows.iter().map(|r| r.latency_ms).collect();
151        lat.sort_unstable();
152        let pct = |p: f64| -> u64 {
153            if lat.is_empty() {
154                return 0;
155            }
156            let i = ((p * lat.len() as f64).ceil() as usize).clamp(1, lat.len()) - 1;
157            lat[i]
158        };
159        let tps: Vec<f64> = rows.iter().filter_map(|r| r.tokens_per_s).collect();
160        Ok(Summary {
161            n: rows.len() as u64,
162            ok: rows.iter().filter(|r| r.ok).count() as u64,
163            schema_n: rows.iter().filter(|r| r.schema_valid.is_some()).count() as u64,
164            schema_ok: rows.iter().filter(|r| r.schema_valid == Some(true)).count() as u64,
165            tokens_per_s_mean: (!tps.is_empty())
166                .then(|| tps.iter().sum::<f64>() / tps.len() as f64),
167            latency_p50_ms: pct(0.50),
168            latency_p95_ms: pct(0.95),
169        })
170    }
171}