metrale_bench/http/
blocking.rs1use super::*;
10
11#[derive(Debug, Default, Clone)]
14pub struct BlockingOutcome {
15 pub choices: Vec<String>,
18 pub prompt_tokens: usize,
19 pub completion_tokens: usize,
20}
21
22pub async fn chat_blocking(
26 target: &TargetEndpoint,
27 body: &Value,
28 timeout: Duration,
29) -> Result<BlockingOutcome> {
30 tokio::time::timeout(timeout, chat_blocking_inner(target, body))
31 .await
32 .map_err(|_| anyhow!("request exceeded {:.0}s", timeout.as_secs_f64()))?
33}
34
35async fn chat_blocking_inner(target: &TargetEndpoint, body: &Value) -> Result<BlockingOutcome> {
36 post_json(target, "/v1/chat/completions", body).await
37}
38
39pub async fn responses_blocking(
42 target: &TargetEndpoint,
43 body: &Value,
44 timeout: Duration,
45) -> Result<BlockingOutcome> {
46 tokio::time::timeout(timeout, post_json(target, "/v1/responses", body))
47 .await
48 .map_err(|_| anyhow!("request exceeded {:.0}s", timeout.as_secs_f64()))?
49}
50
51async fn post_json(target: &TargetEndpoint, path: &str, body: &Value) -> Result<BlockingOutcome> {
52 let (host, port) = target.host_port()?;
53 let payload = serde_json::to_string(body)?;
54 let mut sock = TcpStream::connect((host.as_str(), port))
55 .await
56 .with_context(|| format!("connecting to {}", target.base_url))?;
57 let _ = sock.set_nodelay(true);
58 let request = format!(
59 "POST {path} HTTP/1.1\r\nHost: {host}:{port}\r\n\
60 Content-Type: application/json\r\nAccept: application/json\r\n\
61 Connection: close\r\nContent-Length: {}\r\n\r\n{payload}",
62 payload.len()
63 );
64 sock.write_all(request.as_bytes()).await.context("write")?;
65
66 let mut raw = Vec::new();
69 let mut buf = [0u8; 16 * 1024];
70 loop {
71 let n = sock.read(&mut buf).await.context("read")?;
72 if n == 0 {
73 break;
74 }
75 raw.extend_from_slice(&buf[..n]);
76 }
77
78 let split = raw
79 .windows(4)
80 .position(|w| w == b"\r\n\r\n")
81 .context("no header/body split in the response")?;
82 let head = String::from_utf8_lossy(&raw[..split]).to_string();
83 let body_bytes = &raw[split + 4..];
84
85 let status = head
86 .lines()
87 .next()
88 .and_then(|l| l.split_whitespace().nth(1))
89 .unwrap_or("?")
90 .to_string();
91 if status != "200" {
92 let detail = error_message_from_response(&raw).unwrap_or_else(|| {
93 let text = if head
94 .to_ascii_lowercase()
95 .contains("transfer-encoding: chunked")
96 {
97 dechunk(body_bytes)
98 } else {
99 String::from_utf8_lossy(body_bytes).to_string()
100 };
101 text.chars().take(200).collect()
102 });
103 bail!("endpoint returned HTTP {status}: {detail}");
104 }
105
106 let text = if head.to_lowercase().contains("transfer-encoding: chunked") {
107 dechunk(body_bytes)
108 } else {
109 String::from_utf8_lossy(body_bytes).to_string()
110 };
111 let v: Value = serde_json::from_str(text.trim()).with_context(|| {
112 format!(
113 "response was not JSON: {}",
114 text.chars().take(200).collect::<String>()
115 )
116 })?;
117
118 if v.get("choices").is_none() && v.get("output").is_some() {
122 let mut texts = Vec::new();
123 if let Some(items) = v["output"].as_array() {
124 for item in items {
125 for key in ["content", "summary"] {
129 if let Some(parts) = item[key].as_array() {
130 for part in parts {
131 if let Some(s) = part["text"].as_str() {
132 texts.push(s.to_string());
133 }
134 }
135 }
136 }
137 }
138 }
139 return Ok(BlockingOutcome {
140 choices: vec![texts.join("")],
141 prompt_tokens: v["usage"]["input_tokens"].as_u64().unwrap_or(0) as usize,
142 completion_tokens: v["usage"]["output_tokens"].as_u64().unwrap_or(0) as usize,
143 });
144 }
145 let choices = v["choices"]
146 .as_array()
147 .map(|a| {
148 a.iter()
149 .map(|c| {
150 c["message"]["content"]
151 .as_str()
152 .unwrap_or_default()
153 .to_string()
154 })
155 .collect()
156 })
157 .unwrap_or_default();
158 Ok(BlockingOutcome {
159 choices,
160 prompt_tokens: v["usage"]["prompt_tokens"].as_u64().unwrap_or(0) as usize,
161 completion_tokens: v["usage"]["completion_tokens"].as_u64().unwrap_or(0) as usize,
162 })
163}
164
165fn dechunk(mut b: &[u8]) -> String {
169 let mut out = Vec::new();
170 loop {
171 let Some(nl) = b.windows(2).position(|w| w == b"\r\n") else {
172 break;
173 };
174 let len = usize::from_str_radix(String::from_utf8_lossy(&b[..nl]).trim(), 16).unwrap_or(0);
175 if len == 0 {
176 break;
177 }
178 let start = nl + 2;
179 let end = (start + len).min(b.len());
180 out.extend_from_slice(&b[start..end]);
181 b = &b[(end + 2).min(b.len())..];
182 }
183 String::from_utf8_lossy(&out).to_string()
184}