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