1use std::collections::BTreeMap;
5
6use anyhow::Result;
7use serde::{Deserialize, Serialize};
8
9const TABLE: &str = "prompts";
10const MAX_LEN: usize = 8000;
11
12pub use comindly_api::prompts::Task;
13
14pub use comindly_api::prompts::Version;
15
16#[derive(Debug, Clone, Default, PartialEq, Serialize, Deserialize)]
18pub struct Library {
19 #[serde(default)]
20 pub versions: Vec<Version>,
21 #[serde(default)]
22 pub active: BTreeMap<Task, u32>,
23}
24
25impl Library {
26 pub fn add(
28 &mut self,
29 task: Task,
30 text: &str,
31 sections: &BTreeMap<String, String>,
32 note: Option<&str>,
33 by: &str,
34 ) -> Result<u32> {
35 let text = text.trim();
36 anyhow::ensure!(!text.is_empty(), "the prompt is empty");
37 anyhow::ensure!(
38 text.chars().count() <= MAX_LEN,
39 "the prompt is longer than {MAX_LEN} characters"
40 );
41 let mut kept = BTreeMap::new();
42 for (title, guidance) in sections {
43 let builtin = crate::styleguide::SECTIONS
44 .iter()
45 .find(|(t, _)| t == title)
46 .map(|(_, g)| *g);
47 anyhow::ensure!(
48 task == Task::Guide && builtin.is_some(),
49 "no style guide section {title}"
50 );
51 let g = guidance.trim();
52 anyhow::ensure!(
53 g.chars().count() <= MAX_LEN,
54 "the guidance for {title} is too long"
55 );
56 if !g.is_empty() && Some(g) != builtin {
57 kept.insert(title.clone(), g.to_string());
58 }
59 }
60 let version = self.of(task).map(|v| v.version).max().unwrap_or(0) + 1;
61 self.versions.push(Version {
62 task,
63 version,
64 text: text.to_string(),
65 sections: kept,
66 note: note
67 .map(str::trim)
68 .filter(|n| !n.is_empty())
69 .map(str::to_string),
70 author: by.to_string(),
71 created_at: chrono::Utc::now().to_rfc3339(),
72 based_on: Some(builtin_hash(task)),
73 });
74 Ok(version)
75 }
76
77 pub fn of(&self, task: Task) -> impl Iterator<Item = &Version> {
78 self.versions.iter().filter(move |v| v.task == task)
79 }
80
81 fn find(&self, task: Task, version: u32) -> Option<&Version> {
82 self.of(task).find(|v| v.version == version)
83 }
84
85 pub fn activate(&mut self, task: Task, version: u32) -> Result<()> {
87 anyhow::ensure!(
88 self.find(task, version).is_some(),
89 "{} has no version {version}",
90 task.name()
91 );
92 self.active.insert(task, version);
93 Ok(())
94 }
95
96 pub fn deactivate(&mut self, task: Task) {
98 self.active.remove(&task);
99 }
100
101 pub fn remove(&mut self, task: Task, version: u32) -> Result<()> {
103 anyhow::ensure!(
104 self.active.get(&task) != Some(&version),
105 "deactivate version {version} before deleting it"
106 );
107 let before = self.versions.len();
108 self.versions
109 .retain(|v| !(v.task == task && v.version == version));
110 anyhow::ensure!(
111 self.versions.len() < before,
112 "{} has no version {version}",
113 task.name()
114 );
115 Ok(())
116 }
117
118 pub fn builtin_changed(&self, task: Task) -> Option<u32> {
120 let n = *self.active.get(&task)?;
121 let v = self.find(task, n)?;
122 v.based_on
123 .as_deref()
124 .is_some_and(|b| b != builtin_hash(task))
125 .then_some(n)
126 }
127
128 pub fn in_use(&self) -> Prompts {
130 Prompts(
131 self.active
132 .iter()
133 .filter_map(|(t, n)| {
134 let v = self.find(*t, *n)?;
135 Some((*t, (*n, v.text.clone(), v.sections.clone())))
136 })
137 .collect(),
138 )
139 }
140}
141
142type InUse = (u32, String, BTreeMap<String, String>);
143
144#[derive(Debug, Clone, Default, PartialEq)]
146pub struct Prompts(BTreeMap<Task, InUse>);
147
148impl Prompts {
149 pub fn draft(task: Task, text: &str, sections: BTreeMap<String, String>) -> Prompts {
151 Prompts(BTreeMap::from([(
152 task,
153 (0, text.trim().to_string(), sections),
154 )]))
155 }
156
157 pub fn version(&self, task: Task) -> u32 {
158 self.0.get(&task).map_or(0, |(n, ..)| *n)
159 }
160
161 pub fn system(&self, task: Task) -> String {
162 let text = self
163 .0
164 .get(&task)
165 .map_or(task.builtin(), |(_, t, _)| t.as_str());
166 format!("{text}\n\n{}", fixed(task))
167 }
168
169 pub fn guidance(&self, title: &str) -> Option<&str> {
171 let (_, _, s) = self.0.get(&Task::Guide)?;
172 s.get(title).map(String::as_str)
173 }
174
175 pub fn describe(&self) -> String {
177 Task::ALL
178 .iter()
179 .map(|t| match self.version(*t) {
180 0 => format!("{} built-in", t.name()),
181 n => format!("{} v{n}", t.name()),
182 })
183 .collect::<Vec<_>>()
184 .join(", ")
185 }
186}
187
188fn uri(root: &str) -> String {
189 format!("{}/_team", root.trim_end_matches('/'))
190}
191
192pub async fn load(root: &str) -> Result<crate::docstore::Stored<Library>> {
193 crate::docstore::load(&uri(root), TABLE).await
194}
195
196pub async fn save_as(
198 root: &str,
199 library: &Library,
200 expected: Option<u64>,
201 by: &str,
202) -> Result<u64> {
203 crate::docstore::save(&uri(root), TABLE, library, expected, by)
204 .await
205 .map_err(|e| match format!("{e}").contains("changed meanwhile") {
206 true => {
207 anyhow::anyhow!("the prompts were changed by someone else; reload and try again")
208 }
209 false => e,
210 })
211}
212
213fn example(task: Task) -> &'static str {
215 match task {
216 Task::Summary => {
217 r#"{"summary": "Serves POST /orders: validates the cart, writes the order to the orders table and publishes order-created, returning the new order id.", "queries": ["where do we create an order from a cart?", "what happens when a customer checks out?", "which service publishes order-created?"]}"#
218 }
219 Task::Guide => {
220 r#"{"observed": true, "rules": ["**MUST** name functions in `snake_case`, as in `src/db/pool.rs` `open_pool`."]}"#
221 }
222 Task::Flow => {
223 r#"{"narration": "A checkout calls POST /orders in orders, served by `create_order`. It checks the cart with `validate_cart` and stores the order with `OrderRepo.save`. `publish_order_created` then sends order-created, which `charge_order` in payments consumes to charge the card. The caller gets the new order id; an invalid cart returns 422 before anything is stored.", "queries": ["what happens when a customer checks out?", "where does order-created end up?", "which service charges the card for an order?"]}"#
224 }
225 }
226}
227
228pub fn fixed(task: Task) -> String {
231 format!(
232 "{}\n\nReply with one JSON object that matches this schema (not the schema itself):\n{}\n\nFor example:\n{}",
233 task.format(),
234 crate::llm::reply_schema(task),
235 example(task)
236 )
237}
238
239pub fn builtin_hash(task: Task) -> String {
241 let mut h = blake3::Hasher::new();
242 h.update(task.builtin().as_bytes());
243 if task == Task::Guide {
244 for (title, guidance) in crate::styleguide::SECTIONS {
245 h.update(b"\0");
246 h.update(title.as_bytes());
247 h.update(b"\0");
248 h.update(guidance.as_bytes());
249 }
250 }
251 h.finalize().to_hex().to_string()
252}
253
254#[cfg(test)]
255mod tests {
256 use super::*;
257
258 #[test]
259 fn only_an_activated_version_is_used() {
260 let mut l = Library::default();
261 let v1 = l
262 .add(
263 Task::Summary,
264 "Summarise for search.",
265 &BTreeMap::new(),
266 None,
267 "a",
268 )
269 .unwrap();
270 assert_eq!(
271 l.in_use().version(Task::Summary),
272 0,
273 "a saved draft is not used"
274 );
275 l.activate(Task::Summary, v1).unwrap();
276 let p = l.in_use();
277 assert_eq!(p.version(Task::Summary), 1);
278 assert!(p.system(Task::Summary).starts_with("Summarise for search."));
279 assert!(
280 p.system(Task::Summary).contains("\"queries\""),
281 "the format stays"
282 );
283 assert_eq!(p.version(Task::Flow), 0);
284 assert!(l.remove(Task::Summary, v1).is_err(), "the active one stays");
285 l.deactivate(Task::Summary);
286 assert_eq!(l.in_use(), Prompts::default());
287 l.remove(Task::Summary, v1).unwrap();
288 assert!(l.activate(Task::Summary, v1).is_err());
289 }
290
291 #[test]
292 fn versions_count_per_task() {
293 let mut l = Library::default();
294 assert_eq!(
295 l.add(Task::Guide, "a", &BTreeMap::new(), None, "x")
296 .unwrap(),
297 1
298 );
299 assert_eq!(
300 l.add(Task::Guide, "b", &BTreeMap::new(), Some(" "), "x")
301 .unwrap(),
302 2
303 );
304 assert_eq!(
305 l.add(Task::Flow, "c", &BTreeMap::new(), None, "x").unwrap(),
306 1
307 );
308 assert_eq!(l.versions[1].note, None);
309 assert!(l
310 .add(Task::Flow, " ", &BTreeMap::new(), None, "x")
311 .is_err());
312 }
313
314 #[test]
315 fn guide_versions_replace_section_guidance() {
316 let mut l = Library::default();
317 let (title, builtin) = crate::styleguide::SECTIONS[2];
318 let sections = BTreeMap::from([
319 (title.to_string(), "Only casing.".to_string()),
320 (
321 "Testing".to_string(),
322 crate::styleguide::SECTIONS[8].1.to_string(),
323 ),
324 ]);
325 let n = l
326 .add(Task::Guide, Task::Guide.builtin(), §ions, None, "a")
327 .unwrap();
328 assert_eq!(
329 l.versions[0].sections.len(),
330 1,
331 "unchanged guidance is not stored"
332 );
333 let bad = BTreeMap::from([("Nope".to_string(), "x".to_string())]);
334 assert!(l.add(Task::Guide, "x", &bad, None, "a").is_err());
335 assert!(l.add(Task::Summary, "x", §ions, None, "a").is_err());
336 l.activate(Task::Guide, n).unwrap();
337 let p = l.in_use();
338 assert_eq!(p.guidance(title), Some("Only casing."));
339 assert_ne!(p.guidance(title), Some(builtin));
340 assert_eq!(l.builtin_changed(Task::Guide), None);
341 l.versions[0].based_on = Some("an older release".into());
342 assert_eq!(l.builtin_changed(Task::Guide), Some(n));
343 }
344
345 #[tokio::test]
346 async fn saved_with_the_index() {
347 let dir = std::env::temp_dir().join(format!("comindly-prompts-{}", std::process::id()));
348 let root = dir.to_str().unwrap();
349 let mut l = Library::default();
350 let n = l
351 .add(
352 Task::Flow,
353 "Walk through it.",
354 &BTreeMap::new(),
355 Some("shorter"),
356 "a",
357 )
358 .unwrap();
359 l.activate(Task::Flow, n).unwrap();
360 let v = save_as(root, &l, None, "a").await.unwrap();
361 let stored = load(root).await.unwrap();
362 assert_eq!(stored.doc, l);
363 assert!(
364 save_as(root, &l, Some(v + 5), "b").await.is_err(),
365 "stale version"
366 );
367 std::fs::remove_dir_all(&dir).ok();
368 }
369}