//! The LLM prompts indexing uses: built in (v0) or an activated version, stored with the index
//! (`<index>/_team`, table `prompts`), for what the next index writes; nothing is redone for them.
use std::collections::BTreeMap;
use anyhow::Result;
use serde::{Deserialize, Serialize};
const TABLE: &str = "prompts";
const MAX_LEN: usize = 8000;
pub use comindly_api::prompts::Task;
pub use comindly_api::prompts::Version;
/// Every saved version and which one each task uses; a task with none active uses v0.
#[derive(Debug, Clone, Default, PartialEq, Serialize, Deserialize)]
pub struct Library {
#[serde(default)]
pub versions: Vec<Version>,
#[serde(default)]
pub active: BTreeMap<Task, u32>,
}
impl Library {
/// Save `text` as the task's next version, inactive until someone activates it.
pub fn add(
&mut self,
task: Task,
text: &str,
sections: &BTreeMap<String, String>,
note: Option<&str>,
by: &str,
) -> Result<u32> {
let text = text.trim();
anyhow::ensure!(!text.is_empty(), "the prompt is empty");
anyhow::ensure!(
text.chars().count() <= MAX_LEN,
"the prompt is longer than {MAX_LEN} characters"
);
let mut kept = BTreeMap::new();
for (title, guidance) in sections {
let builtin = crate::styleguide::SECTIONS
.iter()
.find(|(t, _)| t == title)
.map(|(_, g)| *g);
anyhow::ensure!(
task == Task::Guide && builtin.is_some(),
"no style guide section {title}"
);
let g = guidance.trim();
anyhow::ensure!(
g.chars().count() <= MAX_LEN,
"the guidance for {title} is too long"
);
if !g.is_empty() && Some(g) != builtin {
kept.insert(title.clone(), g.to_string());
}
}
let version = self.of(task).map(|v| v.version).max().unwrap_or(0) + 1;
self.versions.push(Version {
task,
version,
text: text.to_string(),
sections: kept,
note: note
.map(str::trim)
.filter(|n| !n.is_empty())
.map(str::to_string),
author: by.to_string(),
created_at: chrono::Utc::now().to_rfc3339(),
based_on: Some(builtin_hash(task)),
});
Ok(version)
}
pub fn of(&self, task: Task) -> impl Iterator<Item = &Version> {
self.versions.iter().filter(move |v| v.task == task)
}
fn find(&self, task: Task, version: u32) -> Option<&Version> {
self.of(task).find(|v| v.version == version)
}
/// Use `version` for `task` from the next index on, in place of whichever was active.
pub fn activate(&mut self, task: Task, version: u32) -> Result<()> {
anyhow::ensure!(
self.find(task, version).is_some(),
"{} has no version {version}",
task.name()
);
self.active.insert(task, version);
Ok(())
}
/// Go back to the built-in prompt.
pub fn deactivate(&mut self, task: Task) {
self.active.remove(&task);
}
/// Delete a version that is not in use.
pub fn remove(&mut self, task: Task, version: u32) -> Result<()> {
anyhow::ensure!(
self.active.get(&task) != Some(&version),
"deactivate version {version} before deleting it"
);
let before = self.versions.len();
self.versions
.retain(|v| !(v.task == task && v.version == version));
anyhow::ensure!(
self.versions.len() < before,
"{} has no version {version}",
task.name()
);
Ok(())
}
/// The active version of `task` when the built-in prompt changed since it was written.
pub fn builtin_changed(&self, task: Task) -> Option<u32> {
let n = *self.active.get(&task)?;
let v = self.find(task, n)?;
v.based_on
.as_deref()
.is_some_and(|b| b != builtin_hash(task))
.then_some(n)
}
/// The prompts an index run uses.
pub fn in_use(&self) -> Prompts {
Prompts(
self.active
.iter()
.filter_map(|(t, n)| {
let v = self.find(*t, *n)?;
Some((*t, (*n, v.text.clone(), v.sections.clone())))
})
.collect(),
)
}
}
type InUse = (u32, String, BTreeMap<String, String>);
/// The instructions each task runs with, fixed for one index run.
#[derive(Debug, Clone, Default, PartialEq)]
pub struct Prompts(BTreeMap<Task, InUse>);
impl Prompts {
/// Only `task` changed, to unsaved `text` and section guidance, for trying it.
pub fn draft(task: Task, text: &str, sections: BTreeMap<String, String>) -> Prompts {
Prompts(BTreeMap::from([(
task,
(0, text.trim().to_string(), sections),
)]))
}
pub fn version(&self, task: Task) -> u32 {
self.0.get(&task).map_or(0, |(n, ..)| *n)
}
pub fn system(&self, task: Task) -> String {
let text = self
.0
.get(&task)
.map_or(task.builtin(), |(_, t, _)| t.as_str());
format!("{text}\n\n{}", fixed(task))
}
/// What to cover in the style guide section `title`, when a version replaces the built-in.
pub fn guidance(&self, title: &str) -> Option<&str> {
let (_, _, s) = self.0.get(&Task::Guide)?;
s.get(title).map(String::as_str)
}
/// e.g. `summary v2, guide built-in, flow built-in`.
pub fn describe(&self) -> String {
Task::ALL
.iter()
.map(|t| match self.version(*t) {
0 => format!("{} built-in", t.name()),
n => format!("{} v{n}", t.name()),
})
.collect::<Vec<_>>()
.join(", ")
}
}
fn uri(root: &str) -> String {
format!("{}/_team", root.trim_end_matches('/'))
}
pub async fn load(root: &str) -> Result<crate::docstore::Stored<Library>> {
crate::docstore::load(&uri(root), TABLE).await
}
/// Save `library` as `by` if nobody saved since `expected`.
pub async fn save_as(
root: &str,
library: &Library,
expected: Option<u64>,
by: &str,
) -> Result<u64> {
crate::docstore::save(&uri(root), TABLE, library, expected, by)
.await
.map_err(|e| match format!("{e}").contains("changed meanwhile") {
true => {
anyhow::anyhow!("the prompts were changed by someone else; reload and try again")
}
false => e,
})
}
/// The shape of a reply, filled in, so a model without structured output does not echo the schema.
fn example(task: Task) -> &'static str {
match task {
Task::Summary => {
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?"]}"#
}
Task::Guide => {
r#"{"observed": true, "rules": ["**MUST** name functions in `snake_case`, as in `src/db/pool.rs` `open_pool`."]}"#
}
Task::Flow => {
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?"]}"#
}
}
}
/// What is always added after the instructions: the format and its JSON schema, so a server
/// that ignores structured output still gets the shape.
pub fn fixed(task: Task) -> String {
format!(
"{}\n\nReply with one JSON object that matches this schema (not the schema itself):\n{}\n\nFor example:\n{}",
task.format(),
crate::llm::reply_schema(task),
example(task)
)
}
/// Changes whenever a release changes this task's built-in prompt.
pub fn builtin_hash(task: Task) -> String {
let mut h = blake3::Hasher::new();
h.update(task.builtin().as_bytes());
if task == Task::Guide {
for (title, guidance) in crate::styleguide::SECTIONS {
h.update(b"\0");
h.update(title.as_bytes());
h.update(b"\0");
h.update(guidance.as_bytes());
}
}
h.finalize().to_hex().to_string()
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn only_an_activated_version_is_used() {
let mut l = Library::default();
let v1 = l
.add(
Task::Summary,
"Summarise for search.",
&BTreeMap::new(),
None,
"a",
)
.unwrap();
assert_eq!(
l.in_use().version(Task::Summary),
0,
"a saved draft is not used"
);
l.activate(Task::Summary, v1).unwrap();
let p = l.in_use();
assert_eq!(p.version(Task::Summary), 1);
assert!(p.system(Task::Summary).starts_with("Summarise for search."));
assert!(
p.system(Task::Summary).contains("\"queries\""),
"the format stays"
);
assert_eq!(p.version(Task::Flow), 0);
assert!(l.remove(Task::Summary, v1).is_err(), "the active one stays");
l.deactivate(Task::Summary);
assert_eq!(l.in_use(), Prompts::default());
l.remove(Task::Summary, v1).unwrap();
assert!(l.activate(Task::Summary, v1).is_err());
}
#[test]
fn versions_count_per_task() {
let mut l = Library::default();
assert_eq!(
l.add(Task::Guide, "a", &BTreeMap::new(), None, "x")
.unwrap(),
1
);
assert_eq!(
l.add(Task::Guide, "b", &BTreeMap::new(), Some(" "), "x")
.unwrap(),
2
);
assert_eq!(
l.add(Task::Flow, "c", &BTreeMap::new(), None, "x").unwrap(),
1
);
assert_eq!(l.versions[1].note, None);
assert!(l
.add(Task::Flow, " ", &BTreeMap::new(), None, "x")
.is_err());
}
#[test]
fn guide_versions_replace_section_guidance() {
let mut l = Library::default();
let (title, builtin) = crate::styleguide::SECTIONS[2];
let sections = BTreeMap::from([
(title.to_string(), "Only casing.".to_string()),
(
"Testing".to_string(),
crate::styleguide::SECTIONS[8].1.to_string(),
),
]);
let n = l
.add(Task::Guide, Task::Guide.builtin(), §ions, None, "a")
.unwrap();
assert_eq!(
l.versions[0].sections.len(),
1,
"unchanged guidance is not stored"
);
let bad = BTreeMap::from([("Nope".to_string(), "x".to_string())]);
assert!(l.add(Task::Guide, "x", &bad, None, "a").is_err());
assert!(l.add(Task::Summary, "x", §ions, None, "a").is_err());
l.activate(Task::Guide, n).unwrap();
let p = l.in_use();
assert_eq!(p.guidance(title), Some("Only casing."));
assert_ne!(p.guidance(title), Some(builtin));
assert_eq!(l.builtin_changed(Task::Guide), None);
l.versions[0].based_on = Some("an older release".into());
assert_eq!(l.builtin_changed(Task::Guide), Some(n));
}
#[tokio::test]
async fn saved_with_the_index() {
let dir = std::env::temp_dir().join(format!("comindly-prompts-{}", std::process::id()));
let root = dir.to_str().unwrap();
let mut l = Library::default();
let n = l
.add(
Task::Flow,
"Walk through it.",
&BTreeMap::new(),
Some("shorter"),
"a",
)
.unwrap();
l.activate(Task::Flow, n).unwrap();
let v = save_as(root, &l, None, "a").await.unwrap();
let stored = load(root).await.unwrap();
assert_eq!(stored.doc, l);
assert!(
save_as(root, &l, Some(v + 5), "b").await.is_err(),
"stale version"
);
std::fs::remove_dir_all(&dir).ok();
}
}