mod body;
use crate::config::manifest::RoleRules;
use crate::prompt::Error;
use brazen::{Content, Message, Role};
use std::path::{Path, PathBuf};
const MESSAGES_DIR: &str = "messages";
const TOOL_ORIGIN: &str = "tool";
#[derive(PartialEq, Eq, Clone, Copy)]
enum Side {
User,
Assistant,
Tool,
}
impl Side {
fn role(self) -> Role {
match self {
Side::User => Role::User,
Side::Assistant => Role::Assistant,
Side::Tool => Role::Tool,
}
}
}
pub(super) fn assemble(worktree: &Path, rules: Option<&RoleRules>) -> Result<Vec<Message>, Error> {
let mut messages: Vec<Message> = Vec::new();
for text in body::compose(worktree, rules)? {
push_grouped(&mut messages, Side::User, vec![Content::Text(text)]);
}
append_transcript(&mut messages, worktree)?;
Ok(messages)
}
pub(super) fn transcript(worktree: &Path) -> Result<Vec<Message>, Error> {
let mut messages: Vec<Message> = Vec::new();
append_transcript(&mut messages, worktree)?;
Ok(messages)
}
fn append_transcript(messages: &mut Vec<Message>, worktree: &Path) -> Result<(), Error> {
let dir = worktree.join(MESSAGES_DIR);
let mut entries: Vec<(u32, PathBuf)> = match std::fs::read_dir(&dir) {
Ok(rd) => {
let mut v = Vec::new();
for entry in rd {
let path = entry.map_err(Error::Io)?.path();
if let Some(seq) = seq_of(&path) {
v.push((seq, path));
}
}
v
}
Err(e) if e.kind() == std::io::ErrorKind::NotFound => Vec::new(),
Err(e) => return Err(Error::Io(e)),
};
entries.sort_by_key(|(seq, _)| *seq);
for (_, path) in entries {
let (side, content) = compose_entry(&path)?;
push_grouped(messages, side, content);
}
Ok(())
}
fn seq_of(path: &Path) -> Option<u32> {
path.file_name()?
.to_string_lossy()
.split('-')
.next()
.and_then(|p| p.parse::<u32>().ok())
}
fn compose_entry(path: &Path) -> Result<(Side, Vec<Content>), Error> {
if path.extension().and_then(|e| e.to_str()) == Some("md") {
let body = std::fs::read_to_string(path).map_err(Error::Io)?;
return Ok((Side::User, vec![Content::Text(body)]));
}
let bytes = std::fs::read(path).map_err(Error::Io)?;
let blocks: Vec<Content> =
serde_json::from_slice(&bytes).expect("transcript entry is a canonical Content array");
Ok((entry_side(path), blocks))
}
fn entry_side(path: &Path) -> Side {
let stem = path
.file_stem()
.map(|s| s.to_string_lossy())
.unwrap_or_default();
let origin = stem.split_once('-').map(|x| x.1).unwrap_or_default();
if origin == TOOL_ORIGIN {
Side::Tool
} else {
Side::Assistant
}
}
fn push_grouped(messages: &mut Vec<Message>, side: Side, mut content: Vec<Content>) {
match messages.last_mut() {
Some(last) if last.role == side.role() => last.content.append(&mut content),
_ => messages.push(Message {
role: side.role(),
content,
}),
}
}
#[cfg(test)]
mod tests;