send.rs--send: every gate written as a replay fixture. Layout parity and the
unsure band are sent as main.rs builds them, through a
jev_mock::forward::Forward directly; the label-tree and ranking runs
are driven through the built extension in a throwaway Postgres whose
endpoint is a forwarding jev-mock.
The Postgres harness is tests/support, which finds the extension and
the server through the environment buck2 test sets. This binary is
built with the same values (BUCK), and re-executes itself with them
when they are missing, since changing its own environment is unsafe.
12use std::path::Path;
14use jev_mock::MockJev; 15use jev_mock::forward::Forward; 16use support::postgres::TestPostgres; 17 18use support::gates::{LABEL_TREE_SQL, RANKING_SQL, setup_people}; 19 20use crate::{Args, QUESTION, Run, TREE_PATHS}; 21 22const REAL_ENDPOINT: &str = "https://api.typesafe.ai/v1/systemone"; 23 24const LAYOUT_PARITY: &str = "layout_parity"; 25const UNSURE_BAND: &str = "unsure_band"; 26const LABEL_TREE: &str = "label_tree"; 27const RANKING: &str = "ranking";
Every fixture record writes, as <out-dir>/<gate>.json.
30const GATES: [&str; 4] = [LAYOUT_PARITY, UNSURE_BAND, LABEL_TREE, RANKING];
Refuses a recording that would overwrite a gate already recorded,
unless --rerecord asks for it: each recording spends live quota.
38pub fn refuse_recorded(args: &Args) -> Result<(), String> { 39 if args.rerecord { 40 return Ok(()); 41 } 42 let recorded: Vec<String> = 43 GATES.iter().map(|g| fixture_path(args, g)).filter(|p| p.exists()).map(|p| p.display().to_string()).collect(); 44 if recorded.is_empty() { 45 return Ok(()); 46 } 47 Err(format!( 48 "already recorded, nothing sent: {}. The gates replay from these fixtures (tests/gate_*.rs); \ 49 pass --rerecord to spend live quota again and overwrite them", 50 recorded.join(", ") 51 )) 52}
Absent under cargo, which has no location of the built extension.
61pub fn record(args: &Args, key: &str) -> Result<(), String> { 62 if BUILT_ENV.iter().any(|(name, _)| std::env::var_os(name).is_none()) { 63 let mut built = Vec::new(); 64 for (name, value) in BUILT_ENV { 65 built.push((name, value.ok_or("--send needs the extension built by buck: buck2 run //tools/record-gates:record-gates")?)); 66 } 67 let status = std::process::Command::new(std::env::current_exe().map_err(|e| e.to_string())?) 68 .args(std::env::args_os().skip(1)) 69 .envs(built) 70 .status() 71 .map_err(|e| e.to_string())?; 72 return if status.success() { Ok(()) } else { Err("the recording run failed".into()) }; 73 } 74 let runtime = tokio::runtime::Builder::new_multi_thread().enable_all().build().map_err(|e| e.to_string())?; 75 runtime.block_on(async { 76 direct(args, key, LAYOUT_PARITY, crate::layout_parity(&args.model)).await?; 77 direct(args, key, UNSURE_BAND, vec![crate::unsure_band(&args.model, 2)]).await?; 78 let paths: Vec<&str> = TREE_PATHS.to_vec(); 79 let tree = gate( 80 args, 81 key, 82 LABEL_TREE, 83 LABEL_TREE_SQL, 84 Some(paths), 85 ); 86 tree.await?; 87 gate( 88 args, 89 key, 90 RANKING, 91 RANKING_SQL, 92 None, 93 ) 94 .await 95 }) 96}
Sends each distinct request of runs once and writes the fixture.
99async fn direct(args: &Args, key: &str, name: &str, runs: Vec<Run>) -> Result<(), String> { 100 let forward = Forward::new(REAL_ENDPOINT, key, args.model.as_str())?; 101 let mut sent: Vec<&[u8]> = Vec::new(); 102 for request in runs.iter().flat_map(|run| &run.requests) { 103 if sent.contains(&request.as_slice()) { 104 continue; 105 } 106 sent.push(request); 107 let response = forward.exchange(request).await.map_err(|e| format!("{name}: {e}"))?; 108 if response.status != 200 { 109 return Err(format!("{name}: the endpoint answered {}: {}", response.status, response.body)); 110 } 111 } 112 write(args, name, &forward, None) 113}
115fn write(args: &Args, name: &str, forward: &Forward, answer: Option<String>) -> Result<(), String> { 116 let fixture = forward.fixture(); 117 let out = fixture_path(args, name); 118 std::fs::create_dir_all(&args.out_dir).map_err(|e| e.to_string())?; 119 // Without --rerecord, never replace a fixture another run wrote since 120 // `refuse_recorded` looked. 121 let mut open = std::fs::OpenOptions::new(); 122 open.write(true); 123 if args.rerecord { 124 open.create(true).truncate(true); 125 } else { 126 open.create_new(true); 127 } 128 let mut file = open.open(&out).map_err(|e| format!("{}: {e}", out.display()))?; 129 std::io::Write::write_all(&mut file, (fixture.to_json() + "\n").as_bytes()).map_err(|e| format!("{}: {e}", out.display()))?; 130 println!("{name}: {} request(s), answer {answer:?}, wrote {}", fixture.exchanges.len(), out.display()); 131 Ok(()) 132}
Runs sql once through a recording mock and writes the fixture.
135async fn gate(args: &Args, key: &str, name: &str, sql: &str, paths: Option<Vec<&str>>) -> Result<(), String> { 136 let forward = Forward::new(REAL_ENDPOINT, key, args.model.as_str())?; 137 let mock = MockJev::start(forward.responder()).await; 138 let pg = TestPostgres::new() 139 .with_extension(Path::new(&std::env::var("POSTJEVSQL_EXT").unwrap())) 140 .setting("jev.endpoint", &mock.endpoint()) 141 .setting("jev.model", args.model.as_str()) 142 .setting("jev.ca_file", mock.ca_file().to_str().unwrap()) 143 .setting("jev.price_per_mtok", &args.price_per_mtok.to_string()) 144 .setting("jev.max_cost", &args.max_cost.to_string()) 145 // The mock adds the real key; the server's never leaves loopback. 146 .env_var("TYPESAFE_API_KEY", "record-gates") 147 .start() 148 .await; 149 let (client, conn) = tokio_postgres::connect(&pg.conn_str(), tokio_postgres::NoTls).await.map_err(|e| e.to_string())?; 150 tokio::spawn(conn); 151 setup_people(&client).await.map_err(|e| e.to_string())?; 152 153 let rows = match &paths { 154 Some(paths) => client.query(sql, &[&QUESTION, paths]).await, 155 None => client.query(sql, &[&QUESTION]).await, 156 } 157 .map_err(|e| format!("{name}: {}", e.as_db_error().map_or(e.to_string(), |d| d.message().to_owned())))?; 158 let answer: Option<String> = rows.first().and_then(|r| r.get(0)); 159 write(args, name, &forward, answer) 160}