Skip to main content

comindly_core/
prompts.rs

1//! The LLM prompts indexing uses: built in (v0) or an activated version, stored with the index
2//! (`<index>/_team`, table `prompts`), for what the next index writes; nothing is redone for them.
3
4use 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/// Every saved version and which one each task uses; a task with none active uses v0.
17#[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    /// Save `text` as the task's next version, inactive until someone activates it.
27    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    /// Use `version` for `task` from the next index on, in place of whichever was active.
86    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    /// Go back to the built-in prompt.
97    pub fn deactivate(&mut self, task: Task) {
98        self.active.remove(&task);
99    }
100
101    /// Delete a version that is not in use.
102    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    /// The active version of `task` when the built-in prompt changed since it was written.
119    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    /// The prompts an index run uses.
129    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/// The instructions each task runs with, fixed for one index run.
145#[derive(Debug, Clone, Default, PartialEq)]
146pub struct Prompts(BTreeMap<Task, InUse>);
147
148impl Prompts {
149    /// Only `task` changed, to unsaved `text` and section guidance, for trying it.
150    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    /// What to cover in the style guide section `title`, when a version replaces the built-in.
170    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    /// e.g. `summary v2, guide built-in, flow built-in`.
176    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
196/// Save `library` as `by` if nobody saved since `expected`.
197pub 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
213/// The shape of a reply, filled in, so a model without structured output does not echo the schema.
214fn example(task: Task) -> &'static str {
215    match task {
216        Task::Summary => {
217            r#"{"summary": "Opens a pooled database connection from the service settings.", "queries": ["How do I connect to the database?", "Where is the connection pool created?", "Which settings does the pool read?"]}"#
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": "`handle_order` validates the order with `validate`, then stores it with `save`.", "queries": ["What happens when an order arrives?", "Where is an order validated?", "Where is an order saved?"]}"#
224        }
225    }
226}
227
228/// What is always added after the instructions: the format and its JSON schema, so a server
229/// that ignores structured output still gets the shape.
230pub 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
239/// Changes whenever a release changes this task's built-in prompt.
240pub 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(), &sections, 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", &sections, 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}