1use std::sync::{Arc, Mutex};
4
5use sim_kernel::{
6 Cx, Error, EvalFabric, EvalReply, EvalRequest, Expr, NumberLiteral, Result, Shape, Symbol,
7};
8use sim_lib_agent_runner_core::{ModelRequest, ModelResponse};
9use sim_lib_stream_core::{DevCassette, DevEvent};
10use sim_value::build::entry;
11
12use crate::{
13 AuthorTask, ContractProjectionCaps, RankedContractCard, RouteAttempt, RouteAttemptStatus,
14 RoutePolicy, RouteTarget, ShapeQuery, estimate_prompt_tokens, project_contracts,
15 run_author_task,
16};
17
18const CHEAP_COST: u64 = 1;
19const ESCALATION_COST: u64 = 3;
20const HIGH_TIER_COST: u64 = 10;
21
22#[derive(Clone, Debug, PartialEq, Eq)]
24pub enum AuthorArm {
25 SourcePayload,
27 ContractPayload,
29 ContractGrammar,
31 Downshifted,
33}
34
35impl AuthorArm {
36 pub fn name(&self) -> &'static str {
38 match self {
39 Self::SourcePayload => "source-payload",
40 Self::ContractPayload => "contract-payload",
41 Self::ContractGrammar => "contract-grammar",
42 Self::Downshifted => "downshifted",
43 }
44 }
45}
46
47#[derive(Clone)]
49pub struct AuthorCase {
50 pub name: Symbol,
52 pub source_payload: String,
54 pub goal: String,
56 pub contract_cards: Vec<RankedContractCard>,
58 pub target_codec: Symbol,
60 pub return_shape_expr: Expr,
62 pub return_shape: Arc<dyn Shape>,
64 pub expected_form: Expr,
66}
67
68impl AuthorCase {
69 pub fn new(
71 name: Symbol,
72 source_payload: impl Into<String>,
73 goal: impl Into<String>,
74 contract_cards: Vec<RankedContractCard>,
75 return_shape_expr: Expr,
76 return_shape: Arc<dyn Shape>,
77 expected_form: Expr,
78 ) -> Self {
79 Self {
80 name,
81 source_payload: source_payload.into(),
82 goal: goal.into(),
83 contract_cards,
84 target_codec: Symbol::qualified("codec", "json"),
85 return_shape_expr,
86 return_shape,
87 expected_form,
88 }
89 }
90}
91
92#[derive(Clone, Debug, Default, PartialEq, Eq)]
94pub struct AuthorArmMetrics {
95 pub payload_tokens: u64,
97 pub payload_bytes: u64,
99 pub execution_calls: u64,
101 pub route_attempts: u64,
103 pub declared_cost: u64,
105 pub accepted_cases: u64,
107 pub cassette_hashes: Vec<String>,
109}
110
111#[derive(Clone, Debug, Default, PartialEq, Eq)]
113pub struct AuthorBenchReport {
114 pub arms: Vec<(AuthorArm, AuthorArmMetrics)>,
116}
117
118impl AuthorBenchReport {
119 pub fn metrics(&self, arm: AuthorArm) -> Option<&AuthorArmMetrics> {
121 self.arms
122 .iter()
123 .find_map(|(candidate, metrics)| (candidate == &arm).then_some(metrics))
124 }
125}
126
127pub fn standard_author_arms() -> Vec<AuthorArm> {
129 vec![
130 AuthorArm::SourcePayload,
131 AuthorArm::ContractPayload,
132 AuthorArm::ContractGrammar,
133 AuthorArm::Downshifted,
134 ]
135}
136
137pub fn standard_author_cases() -> Vec<AuthorCase> {
139 crate::eval_author_corpus::standard_author_cases()
140}
141
142pub fn run_author_bench(
144 cx: &mut Cx,
145 cases: &[AuthorCase],
146 arms: &[AuthorArm],
147) -> Result<AuthorBenchReport> {
148 if cases.is_empty() {
149 return Err(Error::Eval("author bench cases are empty".to_owned()));
150 }
151 if arms.is_empty() {
152 return Err(Error::Eval("author bench arms are empty".to_owned()));
153 }
154 ensure_json_codec(cx)?;
155
156 let mut report = AuthorBenchReport { arms: Vec::new() };
157 for arm in arms {
158 let mut metrics = AuthorArmMetrics::default();
159 for case in cases {
160 let measured = run_author_case(cx, case, arm)?;
161 metrics.payload_tokens = metrics
162 .payload_tokens
163 .saturating_add(measured.payload_tokens);
164 metrics.payload_bytes = metrics.payload_bytes.saturating_add(measured.payload_bytes);
165 metrics.execution_calls = metrics
166 .execution_calls
167 .saturating_add(measured.execution_calls);
168 metrics.route_attempts = metrics
169 .route_attempts
170 .saturating_add(measured.route_attempts);
171 metrics.declared_cost = metrics.declared_cost.saturating_add(measured.declared_cost);
172 metrics.accepted_cases = metrics.accepted_cases.saturating_add(measured.accepted);
173 metrics.cassette_hashes.push(measured.cassette_hash);
174 }
175 report.arms.push((arm.clone(), metrics));
176 }
177
178 Ok(report)
179}
180
181struct MeasuredAuthorCase {
182 payload_tokens: u64,
183 payload_bytes: u64,
184 execution_calls: u64,
185 route_attempts: u64,
186 declared_cost: u64,
187 accepted: u64,
188 cassette_hash: String,
189}
190
191fn run_author_case(cx: &mut Cx, case: &AuthorCase, arm: &AuthorArm) -> Result<MeasuredAuthorCase> {
192 match arm {
193 AuthorArm::SourcePayload => source_payload_case(case),
194 AuthorArm::ContractPayload => contract_case(cx, case, false, ContractRoute::Payload),
195 AuthorArm::ContractGrammar => contract_case(cx, case, true, ContractRoute::Grammar),
196 AuthorArm::Downshifted => contract_case(cx, case, true, ContractRoute::Downshift),
197 }
198}
199
200fn source_payload_case(case: &AuthorCase) -> Result<MeasuredAuthorCase> {
201 let payload_tokens = estimate_prompt_tokens(&case.source_payload) as u64;
202 let cassette = DevCassette::from_events(
203 Symbol::qualified("forge-author-bench", AuthorArm::SourcePayload.name()),
204 vec![DevEvent::validate(
205 case.name.clone(),
206 Expr::Map(vec![
207 entry(
208 "arm",
209 Expr::Symbol(Symbol::qualified(
210 "forge-author-arm",
211 AuthorArm::SourcePayload.name(),
212 )),
213 ),
214 entry("payload-tokens", uint(payload_tokens)),
215 entry("route-attempts", uint(1)),
216 entry("declared-cost", uint(HIGH_TIER_COST)),
217 ]),
218 )?],
219 )?;
220 Ok(MeasuredAuthorCase {
221 payload_tokens,
222 payload_bytes: case.source_payload.len() as u64,
223 execution_calls: 1,
224 route_attempts: 1,
225 declared_cost: HIGH_TIER_COST,
226 accepted: 1,
227 cassette_hash: cassette.content_hash().to_owned(),
228 })
229}
230
231enum ContractRoute {
232 Payload,
233 Grammar,
234 Downshift,
235}
236
237fn contract_case(
238 cx: &mut Cx,
239 case: &AuthorCase,
240 strict_grammar: bool,
241 route: ContractRoute,
242) -> Result<MeasuredAuthorCase> {
243 let task = author_task(case, strict_grammar);
244 let projection = project_contracts(&task.contract_cards, &task.projection_caps);
245 let payload_tokens = projection.tokens as u64;
246 let payload_bytes = projection.text.len() as u64;
247 let expected = encoded_json(&case.expected_form);
248 let malformed = "{not-json".to_owned();
249
250 match route {
251 ContractRoute::Payload => {
252 let cheap = BenchFabric::new(vec![malformed]);
253 let high = BenchFabric::new(vec![expected]);
254 let policy = RoutePolicy::new(
255 vec![
256 RouteTarget::new("cheap-contract", &cheap),
257 RouteTarget::new("high-contract", &high),
258 ],
259 1,
260 );
261 let outcome = run_author_task(cx, &task, &policy)?;
262 measured_contract_case(
263 payload_tokens,
264 payload_bytes,
265 &outcome.attempts,
266 outcome.checked_form.is_some(),
267 outcome.cassette.content_hash(),
268 &[
269 ("cheap-contract", CHEAP_COST),
270 ("high-contract", HIGH_TIER_COST),
271 ],
272 &[&cheap, &high],
273 )
274 }
275 ContractRoute::Grammar => {
276 let cheap = BenchFabric::new(vec![expected]);
277 let high = BenchFabric::new(vec![encoded_json(&case.expected_form)]);
278 let policy = RoutePolicy::new(
279 vec![
280 RouteTarget::new("cheap-contract", &cheap),
281 RouteTarget::new("high-contract", &high),
282 ],
283 1,
284 );
285 let outcome = run_author_task(cx, &task, &policy)?;
286 measured_contract_case(
287 payload_tokens,
288 payload_bytes,
289 &outcome.attempts,
290 outcome.checked_form.is_some(),
291 outcome.cassette.content_hash(),
292 &[
293 ("cheap-contract", CHEAP_COST),
294 ("high-contract", HIGH_TIER_COST),
295 ],
296 &[&cheap, &high],
297 )
298 }
299 ContractRoute::Downshift => {
300 let cheap = BenchFabric::new(vec![malformed]);
301 let escalation = BenchFabric::new(vec![expected]);
302 let high = BenchFabric::new(vec![encoded_json(&case.expected_form)]);
303 let policy = RoutePolicy::new(
304 vec![
305 RouteTarget::new("cheap-downshift", &cheap),
306 RouteTarget::new("escalation-downshift", &escalation),
307 RouteTarget::new("high-contract", &high),
308 ],
309 1,
310 );
311 let outcome = run_author_task(cx, &task, &policy)?;
312 measured_contract_case(
313 payload_tokens,
314 payload_bytes,
315 &outcome.attempts,
316 outcome.checked_form.is_some(),
317 outcome.cassette.content_hash(),
318 &[
319 ("cheap-downshift", CHEAP_COST),
320 ("escalation-downshift", ESCALATION_COST),
321 ("high-contract", HIGH_TIER_COST),
322 ],
323 &[&cheap, &escalation, &high],
324 )
325 }
326 }
327}
328
329fn measured_contract_case(
330 payload_tokens: u64,
331 payload_bytes: u64,
332 attempts: &[RouteAttempt],
333 accepted: bool,
334 cassette_hash: &str,
335 costs: &[(&str, u64)],
336 fabrics: &[&BenchFabric],
337) -> Result<MeasuredAuthorCase> {
338 Ok(MeasuredAuthorCase {
339 payload_tokens,
340 payload_bytes,
341 execution_calls: fabrics
342 .iter()
343 .map(|fabric| fabric.model_request_count() as u64)
344 .sum(),
345 route_attempts: attempts.len() as u64,
346 declared_cost: attempts
347 .iter()
348 .filter(|attempt| !matches!(attempt.status, RouteAttemptStatus::Skipped))
349 .map(|attempt| target_cost(&attempt.target, costs))
350 .sum::<Result<u64>>()?,
351 accepted: u64::from(accepted),
352 cassette_hash: cassette_hash.to_owned(),
353 })
354}
355
356fn target_cost(target: &str, costs: &[(&str, u64)]) -> Result<u64> {
357 costs
358 .iter()
359 .find_map(|(id, cost)| (*id == target).then_some(*cost))
360 .ok_or_else(|| Error::Eval(format!("author bench target {target} has no declared cost")))
361}
362
363fn author_task(case: &AuthorCase, strict_grammar: bool) -> AuthorTask {
364 let mut projection_caps =
365 ContractProjectionCaps::new(case.target_codec.clone(), usize::MAX / 4);
366 projection_caps.include_examples = false;
367 AuthorTask {
368 name: case.name.clone(),
369 goal: case.goal.clone(),
370 target_codec: case.target_codec.clone(),
371 query: ShapeQuery {
372 args: None,
373 result: None,
374 limit: case.contract_cards.len(),
375 },
376 contract_cards: case.contract_cards.clone(),
377 projection_caps,
378 return_shape_expr: case.return_shape_expr.clone(),
379 return_shape: case.return_shape.clone(),
380 verifiers: Vec::new(),
381 strict_grammar,
382 }
383}
384
385fn ensure_json_codec(cx: &mut Cx) -> Result<()> {
386 if cx
387 .registry()
388 .codec_by_symbol(&Symbol::qualified("codec", "json"))
389 .is_some()
390 {
391 return Ok(());
392 }
393 let json = sim_codec_json::JsonCodecLib::new(cx.registry_mut().fresh_codec_id());
394 cx.load_lib(&json)?;
395 Ok(())
396}
397
398fn encoded_json(expr: &Expr) -> String {
399 sim_codec_json::expr_to_json(expr).to_string()
400}
401
402struct BenchFabric {
403 model_outputs: Mutex<Vec<String>>,
404 model_requests: Mutex<usize>,
405 realize_requests: Mutex<usize>,
406}
407
408impl BenchFabric {
409 fn new(model_outputs: Vec<String>) -> Self {
410 Self {
411 model_outputs: Mutex::new(model_outputs),
412 model_requests: Mutex::new(0),
413 realize_requests: Mutex::new(0),
414 }
415 }
416
417 fn model_request_count(&self) -> usize {
418 *self
419 .model_requests
420 .lock()
421 .expect("model request count lock")
422 }
423}
424
425impl EvalFabric for BenchFabric {
426 fn realize(&self, cx: &mut Cx, request: EvalRequest) -> Result<EvalReply> {
427 if ModelRequest::try_from(request.expr.clone()).is_ok() {
428 *self
429 .model_requests
430 .lock()
431 .expect("model request count lock") += 1;
432 let text = {
433 let mut outputs = self.model_outputs.lock().expect("model output lock");
434 if outputs.is_empty() {
435 return Err(Error::Eval(
436 "author bench fake runner is exhausted".to_owned(),
437 ));
438 }
439 outputs.remove(0)
440 };
441 let response = ModelResponse::new(
442 Symbol::qualified("runner", "author-bench"),
443 "author-bench",
444 vec![text_content(text)],
445 Symbol::new("stop"),
446 );
447 return Ok(EvalReply {
448 value: cx.factory().expr(Expr::from(response))?,
449 diagnostics: Vec::new(),
450 trace: None,
451 });
452 }
453
454 *self
455 .realize_requests
456 .lock()
457 .expect("realize request count lock") += 1;
458 Ok(EvalReply {
459 value: cx.factory().expr(request.expr)?,
460 diagnostics: Vec::new(),
461 trace: None,
462 })
463 }
464}
465
466fn text_content(text: String) -> Expr {
467 Expr::Map(vec![
468 entry("type", Expr::Symbol(Symbol::new("text"))),
469 entry("text", Expr::String(text)),
470 ])
471}
472
473fn uint(value: u64) -> Expr {
474 Expr::Number(NumberLiteral {
475 domain: Symbol::qualified("number", "u64"),
476 canonical: value.to_string(),
477 })
478}