use crate::{CacheControl, ContentBlock, Message};
#[derive(Debug, Clone)]
struct PromptLayer {
text: String,
stable: bool,
}
#[derive(Debug, Clone)]
pub struct Prompt {
layers: Vec<PromptLayer>,
}
impl Prompt {
pub fn new() -> Self {
Self { layers: vec![] }
}
pub fn stable(mut self, text: impl Into<String>) -> Self {
self.layers.push(PromptLayer {
text: text.into(),
stable: true,
});
self
}
pub fn dynamic(mut self, text: impl Into<String>) -> Self {
self.layers.push(PromptLayer {
text: text.into(),
stable: false,
});
self
}
pub fn build(self) -> Message {
let last_stable_idx = self
.layers
.iter()
.enumerate()
.rev()
.find(|(_, layer)| layer.stable)
.map(|(idx, _)| idx);
let content: Vec<ContentBlock> = self
.layers
.iter()
.enumerate()
.map(|(idx, layer)| {
if Some(idx) == last_stable_idx {
ContentBlock::text_with_cache(layer.text.clone(), CacheControl::Breakpoint)
} else {
ContentBlock::text(&layer.text)
}
})
.collect();
Message::System { content }
}
pub fn is_empty(&self) -> bool {
self.layers.iter().all(|l| l.text.is_empty())
}
}
impl Default for Prompt {
fn default() -> Self {
Self::new()
}
}
impl From<String> for Prompt {
fn from(s: String) -> Self {
Self::new().dynamic(s)
}
}
impl From<&str> for Prompt {
fn from(s: &str) -> Self {
Self::new().dynamic(s)
}
}
impl From<&String> for Prompt {
fn from(s: &String) -> Self {
Self::new().dynamic(s.clone())
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_prompt_from_string() {
let prompt: Prompt = "hello".into();
let msg = prompt.build();
assert_eq!(msg.content().len(), 1);
assert_eq!(msg.content()[0].as_text(), Some("hello"));
if let ContentBlock::Text(t) = &msg.content()[0] {
assert!(t.cache_control.is_none());
}
}
#[test]
fn test_prompt_from_str() {
let s = "world";
let prompt: Prompt = s.into();
let msg = prompt.build();
assert_eq!(msg.content()[0].as_text(), Some("world"));
}
#[test]
fn test_prompt_new_and_build() {
let msg = Prompt::new()
.stable("layer1")
.stable("layer2")
.dynamic("dynamic")
.build();
let blocks = msg.content();
assert_eq!(blocks.len(), 3);
if let ContentBlock::Text(t) = &blocks[0] {
assert_eq!(t.text, "layer1");
assert!(
t.cache_control.is_none(),
"Intermediate stable layer should NOT have breakpoint"
);
} else {
panic!("expected Text block");
}
if let ContentBlock::Text(t) = &blocks[1] {
assert_eq!(t.text, "layer2");
assert!(
t.cache_control.is_some(),
"Last stable layer should have breakpoint"
);
} else {
panic!("expected Text block");
}
if let ContentBlock::Text(t) = &blocks[2] {
assert_eq!(t.text, "dynamic");
assert!(t.cache_control.is_none());
} else {
panic!("expected Text block");
}
}
#[test]
fn test_prompt_is_empty() {
let empty = Prompt::default();
assert!(empty.is_empty());
let nonempty: Prompt = "x".into();
assert!(!nonempty.is_empty());
}
#[test]
fn test_breakpoint_only_on_last_stable() {
let msg = Prompt::new()
.stable("L1")
.stable("L2")
.stable("L3")
.stable("L4")
.stable("L5")
.dynamic("D")
.build();
let blocks = msg.content();
assert_eq!(blocks.len(), 6);
let breakpoint_count = blocks
.iter()
.filter(|b| {
if let ContentBlock::Text(t) = b {
t.cache_control.is_some()
} else {
false
}
})
.count();
assert_eq!(
breakpoint_count, 1,
"Should have exactly 1 breakpoint (on last stable layer)"
);
if let ContentBlock::Text(t) = &blocks[4] {
assert!(t.cache_control.is_some());
}
}
#[test]
fn test_all_stable_single_breakpoint() {
let msg = Prompt::new().stable("A").stable("B").build();
let blocks = msg.content();
if let ContentBlock::Text(t) = &blocks[0] {
assert!(t.cache_control.is_none(), "A should not have breakpoint");
}
if let ContentBlock::Text(t) = &blocks[1] {
assert!(t.cache_control.is_some(), "B should have breakpoint");
}
}
#[test]
fn test_empty_prompt_produces_empty_message() {
let msg = Prompt::new().build();
assert!(msg.content().is_empty());
}
#[test]
fn test_flatten_text_ignores_cache_control() {
let blocks = vec![
ContentBlock::text_with_cache("cached part".into(), CacheControl::Breakpoint),
ContentBlock::text("dynamic part"),
];
assert_eq!(
ContentBlock::flatten_text(&blocks),
"cached part\n\ndynamic part"
);
}
#[test]
fn test_flatten_text_single_block() {
let blocks = vec![ContentBlock::text("hello")];
assert_eq!(ContentBlock::flatten_text(&blocks), "hello");
}
#[test]
fn test_flatten_text_empty() {
let blocks: Vec<ContentBlock> = vec![];
assert_eq!(ContentBlock::flatten_text(&blocks), "");
}
}