1use serde::{Deserialize, Serialize};
2use std::fmt;
3
4#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)]
8pub struct QueryId(String);
9
10impl QueryId {
11 pub fn single() -> Self {
13 Self("q-0".into())
14 }
15
16 pub fn batch(index: usize) -> Self {
18 Self(format!("q-{index}"))
19 }
20
21 pub fn fork(vm_index: usize, seq: usize) -> Self {
23 Self(format!("f-{vm_index}-{seq}"))
24 }
25
26 pub fn parse(s: &str) -> Self {
28 Self(s.to_string())
29 }
30
31 pub fn as_str(&self) -> &str {
32 &self.0
33 }
34}
35
36impl fmt::Display for QueryId {
37 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
38 self.0.fmt(f)
39 }
40}
41
42#[derive(Debug, Clone, Serialize, Deserialize)]
45pub struct LlmQuery {
46 pub id: QueryId,
47 pub prompt: String,
48 pub system: Option<String>,
49 pub max_tokens: u32,
50 #[serde(default, skip_serializing_if = "is_false")]
54 pub grounded: bool,
55 #[serde(default, skip_serializing_if = "is_false")]
60 pub underspecified: bool,
61 #[serde(default, skip_serializing_if = "Option::is_none")]
75 pub cache_breakpoint: Option<String>,
76 #[serde(default, skip_serializing_if = "Option::is_none")]
87 pub role: Option<String>,
88}
89
90fn is_false(v: &bool) -> bool {
91 !v
92}
93
94#[cfg(test)]
95mod tests {
96 use super::*;
97
98 #[test]
99 fn single_query_id() {
100 let id = QueryId::single();
101 assert_eq!(id.as_str(), "q-0");
102 assert_eq!(id.to_string(), "q-0");
103 }
104
105 #[test]
106 fn batch_query_ids_are_unique() {
107 let ids: Vec<QueryId> = (0..5).map(QueryId::batch).collect();
108 let set: std::collections::HashSet<&QueryId> = ids.iter().collect();
109 assert_eq!(set.len(), 5);
110 assert_eq!(ids[0].as_str(), "q-0");
111 assert_eq!(ids[3].as_str(), "q-3");
112 }
113
114 #[test]
115 fn single_equals_batch_zero() {
116 assert_eq!(QueryId::single(), QueryId::batch(0));
117 }
118
119 #[test]
120 fn parse_roundtrip() {
121 let id = QueryId::parse("q-42");
122 assert_eq!(id.as_str(), "q-42");
123 assert_eq!(id, QueryId::batch(42));
124 }
125
126 #[test]
127 fn parse_arbitrary() {
128 let id = QueryId::parse("custom-id");
129 assert_eq!(id.as_str(), "custom-id");
130 }
131
132 #[test]
133 fn fork_query_id() {
134 let id = QueryId::fork(2, 3);
135 assert_eq!(id.as_str(), "f-2-3");
136 }
137
138 #[test]
139 fn fork_query_ids_are_unique() {
140 let ids: Vec<QueryId> = (0..3)
141 .flat_map(|vm| (0..2).map(move |seq| QueryId::fork(vm, seq)))
142 .collect();
143 let set: std::collections::HashSet<&QueryId> = ids.iter().collect();
144 assert_eq!(set.len(), 6);
145 }
146
147 #[test]
148 fn query_id_roundtrip_json() {
149 let id = QueryId::batch(42);
150 let json = serde_json::to_string(&id).unwrap();
151 let restored: QueryId = serde_json::from_str(&json).unwrap();
152 assert_eq!(id, restored);
153 }
154
155 #[test]
156 fn llm_query_roundtrip_json() {
157 let query = LlmQuery {
158 id: QueryId::single(),
159 prompt: "test prompt".into(),
160 system: Some("system".into()),
161 max_tokens: 1024,
162 grounded: false,
163 underspecified: false,
164 cache_breakpoint: None,
165 role: None,
166 };
167 let json = serde_json::to_value(&query).unwrap();
168 assert!(
169 json.get("grounded").is_none(),
170 "grounded key must be absent when false (skip_serializing_if)"
171 );
172 assert!(
173 json.get("underspecified").is_none(),
174 "underspecified key must be absent when false (skip_serializing_if)"
175 );
176 assert!(
177 json.get("cache_breakpoint").is_none(),
178 "cache_breakpoint key must be absent when None (skip_serializing_if)"
179 );
180 let restored: LlmQuery = serde_json::from_value(json).unwrap();
181 assert_eq!(restored.id, query.id);
182 assert_eq!(restored.prompt, query.prompt);
183 assert_eq!(restored.system, query.system);
184 assert_eq!(restored.max_tokens, query.max_tokens);
185 assert!(!restored.grounded);
186 assert!(!restored.underspecified);
187 assert!(restored.cache_breakpoint.is_none());
188 }
189
190 #[test]
191 fn llm_query_grounded_serde() {
192 let query = LlmQuery {
193 id: QueryId::single(),
194 prompt: "verify this".into(),
195 system: None,
196 max_tokens: 200,
197 grounded: true,
198 underspecified: false,
199 cache_breakpoint: None,
200 role: None,
201 };
202 let json = serde_json::to_value(&query).unwrap();
203 assert_eq!(
204 json["grounded"], true,
205 "grounded key must be present when true"
206 );
207 let restored: LlmQuery = serde_json::from_value(json).unwrap();
208 assert!(restored.grounded);
209 }
210
211 #[test]
212 fn llm_query_grounded_default_on_missing_key() {
213 let json = serde_json::json!({
214 "id": "q-single",
215 "prompt": "test",
216 "system": null,
217 "max_tokens": 100
218 });
219 let query: LlmQuery = serde_json::from_value(json).unwrap();
220 assert!(
221 !query.grounded,
222 "grounded must default to false when key absent"
223 );
224 assert!(
225 !query.underspecified,
226 "underspecified must default to false when key absent"
227 );
228 assert!(
229 query.cache_breakpoint.is_none(),
230 "cache_breakpoint must default to None when key absent"
231 );
232 }
233
234 #[test]
235 fn llm_query_underspecified_serde() {
236 let query = LlmQuery {
237 id: QueryId::single(),
238 prompt: "what format do you want?".into(),
239 system: None,
240 max_tokens: 200,
241 grounded: false,
242 underspecified: true,
243 cache_breakpoint: None,
244 role: None,
245 };
246 let json = serde_json::to_value(&query).unwrap();
247 assert_eq!(
248 json["underspecified"], true,
249 "underspecified key must be present when true"
250 );
251 assert!(
252 json.get("grounded").is_none(),
253 "grounded must be absent when false"
254 );
255 let restored: LlmQuery = serde_json::from_value(json).unwrap();
256 assert!(restored.underspecified);
257 assert!(!restored.grounded);
258 }
259
260 #[test]
261 fn llm_query_both_flags_serde() {
262 let query = LlmQuery {
263 id: QueryId::single(),
264 prompt: "clarify and verify".into(),
265 system: None,
266 max_tokens: 300,
267 grounded: true,
268 underspecified: true,
269 cache_breakpoint: None,
270 role: None,
271 };
272 let json = serde_json::to_value(&query).unwrap();
273 assert_eq!(json["grounded"], true);
274 assert_eq!(json["underspecified"], true);
275 let restored: LlmQuery = serde_json::from_value(json).unwrap();
276 assert!(restored.grounded);
277 assert!(restored.underspecified);
278 }
279
280 #[test]
281 fn llm_query_cache_breakpoint_serde() {
282 let query = LlmQuery {
283 id: QueryId::single(),
284 prompt: "cached prompt".into(),
285 system: Some("shared system".into()),
286 max_tokens: 512,
287 grounded: false,
288 underspecified: false,
289 cache_breakpoint: Some("context".into()),
290 role: None,
291 };
292 let json = serde_json::to_value(&query).unwrap();
293 assert_eq!(
294 json["cache_breakpoint"], "context",
295 "cache_breakpoint must be forwarded verbatim when set"
296 );
297 assert!(
298 json.get("grounded").is_none(),
299 "grounded must be absent when false"
300 );
301 let restored: LlmQuery = serde_json::from_value(json).unwrap();
302 assert_eq!(restored.cache_breakpoint.as_deref(), Some("context"));
303 }
304
305 #[test]
306 fn llm_query_role_serde() {
307 let query = LlmQuery {
308 id: QueryId::single(),
309 prompt: "grade this".into(),
310 system: None,
311 max_tokens: 256,
312 grounded: false,
313 underspecified: false,
314 cache_breakpoint: None,
315 role: Some("grader".into()),
316 };
317 let json = serde_json::to_value(&query).unwrap();
318 assert_eq!(
319 json["role"], "grader",
320 "role must be forwarded verbatim when set"
321 );
322 let restored: LlmQuery = serde_json::from_value(json).unwrap();
323 assert_eq!(restored.role.as_deref(), Some("grader"));
324 }
325
326 #[test]
327 fn llm_query_role_absent_when_none() {
328 let query = LlmQuery {
329 id: QueryId::single(),
330 prompt: "no role".into(),
331 system: None,
332 max_tokens: 128,
333 grounded: false,
334 underspecified: false,
335 cache_breakpoint: None,
336 role: None,
337 };
338 let json = serde_json::to_value(&query).unwrap();
339 assert!(
340 json.get("role").is_none(),
341 "role key must be absent when None (skip_serializing_if)"
342 );
343 let restored: LlmQuery = serde_json::from_value(json).unwrap();
344 assert!(restored.role.is_none());
345 }
346}
347
348#[cfg(test)]
349mod proptests {
350 use super::*;
351 use proptest::prelude::*;
352
353 proptest! {
354 #[test]
355 fn parse_roundtrip_arbitrary(s in "\\PC{1,100}") {
356 let id = QueryId::parse(&s);
357 prop_assert_eq!(id.as_str(), s.as_str());
358 }
359
360 #[test]
361 fn batch_roundtrip(index in 0usize..10_000) {
362 let id = QueryId::batch(index);
363 let expected = format!("q-{index}");
364 prop_assert_eq!(id.as_str(), expected.as_str());
365 }
366
367 #[test]
368 fn display_matches_as_str(s in "\\PC{1,100}") {
369 let id = QueryId::parse(&s);
370 prop_assert_eq!(id.to_string(), id.as_str().to_string());
371 }
372
373 #[test]
374 fn serde_roundtrip_arbitrary(s in "\\PC{1,100}") {
375 let id = QueryId::parse(&s);
376 let json = serde_json::to_string(&id).unwrap();
377 let restored: QueryId = serde_json::from_str(&json).unwrap();
378 prop_assert_eq!(id, restored);
379 }
380 }
381}