use crate::{CacheControl, ContentBlock};
#[derive(Debug, Clone)]
struct PromptLayer {
text: String,
cache_control: Option<CacheControl>,
}
#[derive(Debug, Clone)]
pub struct Prompt {
layers: Vec<PromptLayer>,
}
impl Prompt {
pub fn plain(text: impl Into<String>) -> Self {
Self {
layers: vec![PromptLayer {
text: text.into(),
cache_control: None,
}],
}
}
pub fn builder() -> PromptBuilder {
PromptBuilder::new()
}
pub fn to_content_blocks(&self) -> Vec<ContentBlock> {
self.layers
.iter()
.map(|layer| match layer.cache_control {
Some(cache) => ContentBlock::text_with_cache(layer.text.clone(), cache),
None => ContentBlock::text(&layer.text),
})
.collect()
}
pub fn build_text(&self) -> String {
self.layers
.iter()
.map(|layer| layer.text.clone())
.collect::<Vec<_>>()
.join("\n\n")
}
pub fn is_empty(&self) -> bool {
self.layers.iter().all(|l| l.text.is_empty())
}
}
impl From<String> for Prompt {
fn from(s: String) -> Self {
Self::plain(s)
}
}
impl From<&str> for Prompt {
fn from(s: &str) -> Self {
Self::plain(s)
}
}
impl From<&String> for Prompt {
fn from(s: &String) -> Self {
Self::plain(s.clone())
}
}
impl Default for Prompt {
fn default() -> Self {
Self { layers: vec![] }
}
}
#[derive(Debug, Clone, Default)]
pub struct PromptBuilder {
layers: Vec<PromptLayer>,
}
impl PromptBuilder {
pub fn new() -> Self {
Self::default()
}
pub fn layer_cached(mut self, text: impl Into<String>) -> Self {
self.layers.push(PromptLayer {
text: text.into(),
cache_control: Some(CacheControl::Breakpoint),
});
self
}
pub fn layer_dynamic(mut self, text: impl Into<String>) -> Self {
self.layers.push(PromptLayer {
text: text.into(),
cache_control: None,
});
self
}
pub fn layer(mut self, text: impl Into<String>, cache: Option<CacheControl>) -> Self {
self.layers.push(PromptLayer {
text: text.into(),
cache_control: cache,
});
self
}
pub fn build(self) -> Prompt {
Prompt {
layers: self.layers,
}
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_prompt_from_string() {
let prompt: Prompt = "hello".into();
let blocks = prompt.to_content_blocks();
assert_eq!(blocks.len(), 1);
assert_eq!(blocks[0].as_text(), Some("hello"));
if let ContentBlock::Text(t) = &blocks[0] {
assert!(t.cache_control.is_none());
}
}
#[test]
fn test_prompt_from_str() {
let s = "world";
let prompt: Prompt = s.into();
assert_eq!(prompt.build_text(), "world");
}
#[test]
fn test_prompt_plain() {
let prompt = Prompt::plain("plain text");
assert_eq!(prompt.build_text(), "plain text");
let blocks = prompt.to_content_blocks();
assert_eq!(blocks.len(), 1);
assert!(matches!(
&blocks[0],
ContentBlock::Text(t) if t.cache_control.is_none()
));
}
#[test]
fn test_prompt_builder_layered() {
let prompt = Prompt::builder()
.layer_cached("layer1")
.layer_cached("layer2")
.layer_dynamic("dynamic")
.build();
let blocks = prompt.to_content_blocks();
assert_eq!(blocks.len(), 3);
assert_eq!(blocks.len(), 3);
if let ContentBlock::Text(t) = &blocks[0] {
assert_eq!(t.text, "layer1");
assert!(t.cache_control.is_some());
} else {
panic!("expected Text block");
}
if let ContentBlock::Text(t) = &blocks[1] {
assert_eq!(t.text, "layer2");
assert!(t.cache_control.is_some());
} 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_build_text_joins_with_double_newline() {
let prompt = Prompt::builder()
.layer_cached("A")
.layer_cached("B")
.build();
assert_eq!(prompt.build_text(), "A\n\nB");
}
#[test]
fn test_build_text_single_layer() {
let prompt = Prompt::builder().layer_cached("only").build();
assert_eq!(prompt.build_text(), "only");
}
#[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_prompt_layer_custom_cache() {
let prompt = Prompt::builder()
.layer("cached", Some(CacheControl::Breakpoint))
.layer("no cache", None)
.build();
let blocks = prompt.to_content_blocks();
if let ContentBlock::Text(t) = &blocks[0] {
assert!(t.cache_control.is_some());
} else {
panic!("expected Text");
}
if let ContentBlock::Text(t) = &blocks[1] {
assert!(t.cache_control.is_none());
} else {
panic!("expected Text");
}
}
#[test]
fn test_empty_layers_produce_empty_blocks() {
let prompt = Prompt::builder().build();
assert!(prompt.to_content_blocks().is_empty());
assert_eq!(prompt.build_text(), "");
}
}