loadmaster_probe/
eval.rs

1//! `probe eval`: the user's prompts against one or more models, deterministic sampling.
2//!
3//! A prompts file is JSONL; each line is an object with `messages` (or `prompt`) and any
4//! of `tools`, `expect_tool`, `expect_args` (a subset the arguments must contain),
5//! `schema` (sent as `response_format` and checked), `expect_contains`, and an `id`.
6
7use serde::Serialize;
8use serde_json::{Value, json};
9
10use crate::client::Client;
11
12#[derive(Debug, Clone, PartialEq)]
13pub struct Row {
14    pub id: String,
15    pub messages: Value,
16    pub tools: Option<Value>,
17    pub expect_tool: Option<String>,
18    pub expect_args: Option<Value>,
19    pub schema: Option<Value>,
20    pub expect_contains: Option<String>,
21}
22
23pub fn parse_jsonl(text: &str) -> Result<Vec<Row>, String> {
24    let mut rows = Vec::new();
25    for (i, line) in text.lines().enumerate() {
26        let line = line.trim();
27        if line.is_empty() || line.starts_with("//") {
28            continue;
29        }
30        let v: Value = serde_json::from_str(line).map_err(|e| format!("line {}: {e}", i + 1))?;
31        let messages = match (v.get("messages"), v.get("prompt").and_then(Value::as_str)) {
32            (Some(m), _) => m.clone(),
33            (None, Some(p)) => json!([{ "role": "user", "content": p }]),
34            (None, None) => return Err(format!("line {}: needs `messages` or `prompt`", i + 1)),
35        };
36        if let Some(s) = v.get("schema") {
37            jsonschema::validator_for(s)
38                .map_err(|e| format!("line {}: invalid schema: {e}", i + 1))?;
39        }
40        rows.push(Row {
41            id: v
42                .get("id")
43                .and_then(Value::as_str)
44                .map_or_else(|| format!("{}", i + 1), str::to_owned),
45            messages,
46            tools: v.get("tools").cloned(),
47            expect_tool: v
48                .get("expect_tool")
49                .and_then(Value::as_str)
50                .map(str::to_owned),
51            expect_args: v.get("expect_args").cloned(),
52            schema: v.get("schema").cloned(),
53            expect_contains: v
54                .get("expect_contains")
55                .and_then(Value::as_str)
56                .map(str::to_owned),
57        });
58    }
59    if rows.is_empty() {
60        return Err("no prompts".into());
61    }
62    Ok(rows)
63}
64
65#[derive(Debug, Clone, PartialEq, Serialize)]
66pub struct RowResult {
67    pub idx: usize,
68    pub id: String,
69    /// Every expectation the row states held (a row with none: the request succeeded).
70    pub ok: bool,
71    pub schema_valid: Option<bool>,
72    pub tool_ok: Option<bool>,
73    pub latency_ms: u64,
74    pub completion_tokens: Option<u64>,
75    pub tokens_per_s: Option<f64>,
76    pub error: Option<String>,
77    /// blake3 of the response message, to spot output changes between runs.
78    pub response_blake3: String,
79}
80
81/// `want` is contained in `have`: objects by key recursively, everything else by equality.
82pub fn subset(want: &Value, have: &Value) -> bool {
83    match (want, have) {
84        (Value::Object(w), Value::Object(h)) => w
85            .iter()
86            .all(|(k, v)| h.get(k).is_some_and(|hv| subset(v, hv))),
87        _ => want == have,
88    }
89}
90
91pub fn score(
92    row: &Row,
93    status: u16,
94    response: &Value,
95) -> (bool, Option<bool>, Option<bool>, Option<String>) {
96    if status != 200 {
97        let msg = response["error"]["message"]
98            .as_str()
99            .unwrap_or("request failed");
100        return (
101            false,
102            row.schema.as_ref().map(|_| false),
103            row.expect_tool.as_ref().map(|_| false),
104            Some(format!("HTTP {status}: {msg}")),
105        );
106    }
107    let msg = &response["choices"][0]["message"];
108    let content = msg["content"].as_str().unwrap_or("");
109    let tool_ok = row.expect_tool.as_ref().map(|name| {
110        let call = &msg["tool_calls"][0]["function"];
111        call["name"].as_str() == Some(name.as_str())
112            && row.expect_args.as_ref().is_none_or(|want| {
113                call["arguments"]
114                    .as_str()
115                    .and_then(|a| serde_json::from_str::<Value>(a).ok())
116                    .is_some_and(|have| subset(want, &have))
117            })
118    });
119    let schema_valid = row.schema.as_ref().map(|s| {
120        jsonschema::validator_for(s)
121            .is_ok_and(|v| serde_json::from_str::<Value>(content).is_ok_and(|i| v.is_valid(&i)))
122    });
123    let contains_ok = row
124        .expect_contains
125        .as_ref()
126        .map(|needle| content.contains(needle.as_str()));
127    let ok = tool_ok.unwrap_or(true) && schema_valid.unwrap_or(true) && contains_ok.unwrap_or(true);
128    (ok, schema_valid, tool_ok, None)
129}
130
131pub async fn run(c: &Client, model: &str, rows: &[Row]) -> Vec<RowResult> {
132    let mut out = Vec::new();
133    for (idx, row) in rows.iter().enumerate() {
134        let mut body = json!({ "model": model, "messages": row.messages, "temperature": 0, "seed": 42, "max_tokens": 1024 });
135        if let Some(t) = &row.tools {
136            body["tools"] = t.clone();
137        }
138        if let Some(s) = &row.schema {
139            body["response_format"] = json!({ "type": "json_schema", "json_schema": { "name": "eval", "strict": true, "schema": s } });
140        }
141        let x = c.chat(body).await;
142        let (ok, schema_valid, tool_ok, error) = score(row, x.status, &x.response);
143        let completion_tokens = x.response["usage"]["completion_tokens"].as_u64();
144        let tokens_per_s = completion_tokens
145            .filter(|_| x.latency_ms > 0)
146            .map(|t| t as f64 / (x.latency_ms as f64 / 1000.0));
147        out.push(RowResult {
148            idx,
149            id: row.id.clone(),
150            ok,
151            schema_valid,
152            tool_ok,
153            latency_ms: x.latency_ms,
154            completion_tokens,
155            tokens_per_s,
156            error,
157            response_blake3: blake3::hash(
158                x.response["choices"][0]["message"].to_string().as_bytes(),
159            )
160            .to_hex()
161            .to_string(),
162        });
163    }
164    out
165}
loadmaster_probe/eval.rs

loadmaster_probe/
eval.rs
