1use 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#[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#[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}