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}
77
78fn is_false(v: &bool) -> bool {
79 !v
80}
81
82#[cfg(test)]
83mod tests {
84 use super::*;
85
86 #[test]
87 fn single_query_id() {
88 let id = QueryId::single();
89 assert_eq!(id.as_str(), "q-0");
90 assert_eq!(id.to_string(), "q-0");
91 }
92
93 #[test]
94 fn batch_query_ids_are_unique() {
95 let ids: Vec<QueryId> = (0..5).map(QueryId::batch).collect();
96 let set: std::collections::HashSet<&QueryId> = ids.iter().collect();
97 assert_eq!(set.len(), 5);
98 assert_eq!(ids[0].as_str(), "q-0");
99 assert_eq!(ids[3].as_str(), "q-3");
100 }
101
102 #[test]
103 fn single_equals_batch_zero() {
104 assert_eq!(QueryId::single(), QueryId::batch(0));
105 }
106
107 #[test]
108 fn parse_roundtrip() {
109 let id = QueryId::parse("q-42");
110 assert_eq!(id.as_str(), "q-42");
111 assert_eq!(id, QueryId::batch(42));
112 }
113
114 #[test]
115 fn parse_arbitrary() {
116 let id = QueryId::parse("custom-id");
117 assert_eq!(id.as_str(), "custom-id");
118 }
119
120 #[test]
121 fn fork_query_id() {
122 let id = QueryId::fork(2, 3);
123 assert_eq!(id.as_str(), "f-2-3");
124 }
125
126 #[test]
127 fn fork_query_ids_are_unique() {
128 let ids: Vec<QueryId> = (0..3)
129 .flat_map(|vm| (0..2).map(move |seq| QueryId::fork(vm, seq)))
130 .collect();
131 let set: std::collections::HashSet<&QueryId> = ids.iter().collect();
132 assert_eq!(set.len(), 6);
133 }
134
135 #[test]
136 fn query_id_roundtrip_json() {
137 let id = QueryId::batch(42);
138 let json = serde_json::to_string(&id).unwrap();
139 let restored: QueryId = serde_json::from_str(&json).unwrap();
140 assert_eq!(id, restored);
141 }
142
143 #[test]
144 fn llm_query_roundtrip_json() {
145 let query = LlmQuery {
146 id: QueryId::single(),
147 prompt: "test prompt".into(),
148 system: Some("system".into()),
149 max_tokens: 1024,
150 grounded: false,
151 underspecified: false,
152 cache_breakpoint: None,
153 };
154 let json = serde_json::to_value(&query).unwrap();
155 assert!(
156 json.get("grounded").is_none(),
157 "grounded key must be absent when false (skip_serializing_if)"
158 );
159 assert!(
160 json.get("underspecified").is_none(),
161 "underspecified key must be absent when false (skip_serializing_if)"
162 );
163 assert!(
164 json.get("cache_breakpoint").is_none(),
165 "cache_breakpoint key must be absent when None (skip_serializing_if)"
166 );
167 let restored: LlmQuery = serde_json::from_value(json).unwrap();
168 assert_eq!(restored.id, query.id);
169 assert_eq!(restored.prompt, query.prompt);
170 assert_eq!(restored.system, query.system);
171 assert_eq!(restored.max_tokens, query.max_tokens);
172 assert!(!restored.grounded);
173 assert!(!restored.underspecified);
174 assert!(restored.cache_breakpoint.is_none());
175 }
176
177 #[test]
178 fn llm_query_grounded_serde() {
179 let query = LlmQuery {
180 id: QueryId::single(),
181 prompt: "verify this".into(),
182 system: None,
183 max_tokens: 200,
184 grounded: true,
185 underspecified: false,
186 cache_breakpoint: None,
187 };
188 let json = serde_json::to_value(&query).unwrap();
189 assert_eq!(
190 json["grounded"], true,
191 "grounded key must be present when true"
192 );
193 let restored: LlmQuery = serde_json::from_value(json).unwrap();
194 assert!(restored.grounded);
195 }
196
197 #[test]
198 fn llm_query_grounded_default_on_missing_key() {
199 let json = serde_json::json!({
200 "id": "q-single",
201 "prompt": "test",
202 "system": null,
203 "max_tokens": 100
204 });
205 let query: LlmQuery = serde_json::from_value(json).unwrap();
206 assert!(
207 !query.grounded,
208 "grounded must default to false when key absent"
209 );
210 assert!(
211 !query.underspecified,
212 "underspecified must default to false when key absent"
213 );
214 assert!(
215 query.cache_breakpoint.is_none(),
216 "cache_breakpoint must default to None when key absent"
217 );
218 }
219
220 #[test]
221 fn llm_query_underspecified_serde() {
222 let query = LlmQuery {
223 id: QueryId::single(),
224 prompt: "what format do you want?".into(),
225 system: None,
226 max_tokens: 200,
227 grounded: false,
228 underspecified: true,
229 cache_breakpoint: None,
230 };
231 let json = serde_json::to_value(&query).unwrap();
232 assert_eq!(
233 json["underspecified"], true,
234 "underspecified key must be present when true"
235 );
236 assert!(
237 json.get("grounded").is_none(),
238 "grounded must be absent when false"
239 );
240 let restored: LlmQuery = serde_json::from_value(json).unwrap();
241 assert!(restored.underspecified);
242 assert!(!restored.grounded);
243 }
244
245 #[test]
246 fn llm_query_both_flags_serde() {
247 let query = LlmQuery {
248 id: QueryId::single(),
249 prompt: "clarify and verify".into(),
250 system: None,
251 max_tokens: 300,
252 grounded: true,
253 underspecified: true,
254 cache_breakpoint: None,
255 };
256 let json = serde_json::to_value(&query).unwrap();
257 assert_eq!(json["grounded"], true);
258 assert_eq!(json["underspecified"], true);
259 let restored: LlmQuery = serde_json::from_value(json).unwrap();
260 assert!(restored.grounded);
261 assert!(restored.underspecified);
262 }
263
264 #[test]
265 fn llm_query_cache_breakpoint_serde() {
266 let query = LlmQuery {
267 id: QueryId::single(),
268 prompt: "cached prompt".into(),
269 system: Some("shared system".into()),
270 max_tokens: 512,
271 grounded: false,
272 underspecified: false,
273 cache_breakpoint: Some("context".into()),
274 };
275 let json = serde_json::to_value(&query).unwrap();
276 assert_eq!(
277 json["cache_breakpoint"], "context",
278 "cache_breakpoint must be forwarded verbatim when set"
279 );
280 assert!(
281 json.get("grounded").is_none(),
282 "grounded must be absent when false"
283 );
284 let restored: LlmQuery = serde_json::from_value(json).unwrap();
285 assert_eq!(restored.cache_breakpoint.as_deref(), Some("context"));
286 }
287}
288
289#[cfg(test)]
290mod proptests {
291 use super::*;
292 use proptest::prelude::*;
293
294 proptest! {
295 #[test]
296 fn parse_roundtrip_arbitrary(s in "\\PC{1,100}") {
297 let id = QueryId::parse(&s);
298 prop_assert_eq!(id.as_str(), s.as_str());
299 }
300
301 #[test]
302 fn batch_roundtrip(index in 0usize..10_000) {
303 let id = QueryId::batch(index);
304 let expected = format!("q-{index}");
305 prop_assert_eq!(id.as_str(), expected.as_str());
306 }
307
308 #[test]
309 fn display_matches_as_str(s in "\\PC{1,100}") {
310 let id = QueryId::parse(&s);
311 prop_assert_eq!(id.to_string(), id.as_str().to_string());
312 }
313
314 #[test]
315 fn serde_roundtrip_arbitrary(s in "\\PC{1,100}") {
316 let id = QueryId::parse(&s);
317 let json = serde_json::to_string(&id).unwrap();
318 let restored: QueryId = serde_json::from_str(&json).unwrap();
319 prop_assert_eq!(id, restored);
320 }
321 }
322}