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 parse(s: &str) -> Self {
23 Self(s.to_string())
24 }
25
26 pub fn as_str(&self) -> &str {
27 &self.0
28 }
29}
30
31impl fmt::Display for QueryId {
32 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
33 self.0.fmt(f)
34 }
35}
36
37#[derive(Debug, Clone, Serialize, Deserialize)]
40pub struct LlmQuery {
41 pub id: QueryId,
42 pub prompt: String,
43 pub system: Option<String>,
44 pub max_tokens: u32,
45 #[serde(default, skip_serializing_if = "is_false")]
49 pub grounded: bool,
50 #[serde(default, skip_serializing_if = "is_false")]
55 pub underspecified: bool,
56}
57
58fn is_false(v: &bool) -> bool {
59 !v
60}
61
62#[cfg(test)]
63mod tests {
64 use super::*;
65
66 #[test]
67 fn single_query_id() {
68 let id = QueryId::single();
69 assert_eq!(id.as_str(), "q-0");
70 assert_eq!(id.to_string(), "q-0");
71 }
72
73 #[test]
74 fn batch_query_ids_are_unique() {
75 let ids: Vec<QueryId> = (0..5).map(QueryId::batch).collect();
76 let set: std::collections::HashSet<&QueryId> = ids.iter().collect();
77 assert_eq!(set.len(), 5);
78 assert_eq!(ids[0].as_str(), "q-0");
79 assert_eq!(ids[3].as_str(), "q-3");
80 }
81
82 #[test]
83 fn single_equals_batch_zero() {
84 assert_eq!(QueryId::single(), QueryId::batch(0));
85 }
86
87 #[test]
88 fn parse_roundtrip() {
89 let id = QueryId::parse("q-42");
90 assert_eq!(id.as_str(), "q-42");
91 assert_eq!(id, QueryId::batch(42));
92 }
93
94 #[test]
95 fn parse_arbitrary() {
96 let id = QueryId::parse("custom-id");
97 assert_eq!(id.as_str(), "custom-id");
98 }
99
100 #[test]
101 fn query_id_roundtrip_json() {
102 let id = QueryId::batch(42);
103 let json = serde_json::to_string(&id).unwrap();
104 let restored: QueryId = serde_json::from_str(&json).unwrap();
105 assert_eq!(id, restored);
106 }
107
108 #[test]
109 fn llm_query_roundtrip_json() {
110 let query = LlmQuery {
111 id: QueryId::single(),
112 prompt: "test prompt".into(),
113 system: Some("system".into()),
114 max_tokens: 1024,
115 grounded: false,
116 underspecified: false,
117 };
118 let json = serde_json::to_value(&query).unwrap();
119 assert!(
120 json.get("grounded").is_none(),
121 "grounded key must be absent when false (skip_serializing_if)"
122 );
123 assert!(
124 json.get("underspecified").is_none(),
125 "underspecified key must be absent when false (skip_serializing_if)"
126 );
127 let restored: LlmQuery = serde_json::from_value(json).unwrap();
128 assert_eq!(restored.id, query.id);
129 assert_eq!(restored.prompt, query.prompt);
130 assert_eq!(restored.system, query.system);
131 assert_eq!(restored.max_tokens, query.max_tokens);
132 assert!(!restored.grounded);
133 assert!(!restored.underspecified);
134 }
135
136 #[test]
137 fn llm_query_grounded_serde() {
138 let query = LlmQuery {
139 id: QueryId::single(),
140 prompt: "verify this".into(),
141 system: None,
142 max_tokens: 200,
143 grounded: true,
144 underspecified: false,
145 };
146 let json = serde_json::to_value(&query).unwrap();
147 assert_eq!(
148 json["grounded"], true,
149 "grounded key must be present when true"
150 );
151 let restored: LlmQuery = serde_json::from_value(json).unwrap();
152 assert!(restored.grounded);
153 }
154
155 #[test]
156 fn llm_query_grounded_default_on_missing_key() {
157 let json = serde_json::json!({
158 "id": "q-single",
159 "prompt": "test",
160 "system": null,
161 "max_tokens": 100
162 });
163 let query: LlmQuery = serde_json::from_value(json).unwrap();
164 assert!(
165 !query.grounded,
166 "grounded must default to false when key absent"
167 );
168 assert!(
169 !query.underspecified,
170 "underspecified must default to false when key absent"
171 );
172 }
173
174 #[test]
175 fn llm_query_underspecified_serde() {
176 let query = LlmQuery {
177 id: QueryId::single(),
178 prompt: "what format do you want?".into(),
179 system: None,
180 max_tokens: 200,
181 grounded: false,
182 underspecified: true,
183 };
184 let json = serde_json::to_value(&query).unwrap();
185 assert_eq!(
186 json["underspecified"], true,
187 "underspecified key must be present when true"
188 );
189 assert!(
190 json.get("grounded").is_none(),
191 "grounded must be absent when false"
192 );
193 let restored: LlmQuery = serde_json::from_value(json).unwrap();
194 assert!(restored.underspecified);
195 assert!(!restored.grounded);
196 }
197
198 #[test]
199 fn llm_query_both_flags_serde() {
200 let query = LlmQuery {
201 id: QueryId::single(),
202 prompt: "clarify and verify".into(),
203 system: None,
204 max_tokens: 300,
205 grounded: true,
206 underspecified: true,
207 };
208 let json = serde_json::to_value(&query).unwrap();
209 assert_eq!(json["grounded"], true);
210 assert_eq!(json["underspecified"], true);
211 let restored: LlmQuery = serde_json::from_value(json).unwrap();
212 assert!(restored.grounded);
213 assert!(restored.underspecified);
214 }
215}
216
217#[cfg(test)]
218mod proptests {
219 use super::*;
220 use proptest::prelude::*;
221
222 proptest! {
223 #[test]
224 fn parse_roundtrip_arbitrary(s in "\\PC{1,100}") {
225 let id = QueryId::parse(&s);
226 prop_assert_eq!(id.as_str(), s.as_str());
227 }
228
229 #[test]
230 fn batch_roundtrip(index in 0usize..10_000) {
231 let id = QueryId::batch(index);
232 let expected = format!("q-{index}");
233 prop_assert_eq!(id.as_str(), expected.as_str());
234 }
235
236 #[test]
237 fn display_matches_as_str(s in "\\PC{1,100}") {
238 let id = QueryId::parse(&s);
239 prop_assert_eq!(id.to_string(), id.as_str().to_string());
240 }
241
242 #[test]
243 fn serde_roundtrip_arbitrary(s in "\\PC{1,100}") {
244 let id = QueryId::parse(&s);
245 let json = serde_json::to_string(&id).unwrap();
246 let restored: QueryId = serde_json::from_str(&json).unwrap();
247 prop_assert_eq!(id, restored);
248 }
249 }
250}