1use crate::{Color, HistoryBlock, HistoryLine, HistoryRun, Style};
2
3#[derive(Clone, Copy, Debug, PartialEq, Eq)]
4pub enum TranscriptRole {
5 User,
6 Assistant,
7 Status,
8}
9
10#[derive(Clone, Debug, PartialEq, Eq)]
11pub struct TranscriptEntry {
12 pub role: TranscriptRole,
13 pub content: String,
14}
15
16impl TranscriptEntry {
17 pub fn new(role: TranscriptRole, content: impl Into<String>) -> Self {
18 Self {
19 role,
20 content: content.into(),
21 }
22 }
23
24 pub fn user(content: impl Into<String>) -> Self {
25 Self::new(TranscriptRole::User, content)
26 }
27
28 pub fn assistant(content: impl Into<String>) -> Self {
29 Self::new(TranscriptRole::Assistant, content)
30 }
31
32 pub fn status(content: impl Into<String>) -> Self {
33 Self::new(TranscriptRole::Status, content)
34 }
35}
36
37#[derive(Clone, Debug, Default, PartialEq, Eq)]
38pub struct TranscriptSession {
39 started: bool,
40 entries: Vec<TranscriptEntry>,
41}
42
43impl TranscriptSession {
44 pub fn started(&self) -> bool {
45 self.started
46 }
47
48 pub fn is_empty(&self) -> bool {
49 self.entries.is_empty()
50 }
51
52 pub fn entries(&self) -> &[TranscriptEntry] {
53 &self.entries
54 }
55
56 pub fn begin_turn(&mut self, prompt: impl Into<String>) -> Option<HistoryBlock> {
57 let committed = (!self.entries.is_empty()).then(|| transcript_block(&self.entries));
58 self.started = true;
59 self.entries = vec![
60 TranscriptEntry::user(prompt),
61 TranscriptEntry::assistant(String::new()),
62 ];
63 committed
64 }
65
66 pub fn append_assistant(&mut self, chunk: &str) {
67 if let Some(entry) = self
68 .entries
69 .iter_mut()
70 .rev()
71 .find(|entry| matches!(entry.role, TranscriptRole::Assistant))
72 {
73 entry.content.push_str(chunk);
74 }
75 }
76}
77
78pub fn transcript_block(entries: &[TranscriptEntry]) -> HistoryBlock {
79 let mut lines = Vec::new();
80
81 for entry in entries {
82 let entry_lines: Vec<&str> = entry.content.lines().collect();
83 if entry_lines.is_empty() {
84 lines.push(transcript_line(entry.role, ""));
85 continue;
86 }
87
88 for (index, line) in entry_lines.into_iter().enumerate() {
89 lines.push(transcript_line_with_prefix(entry.role, line, index == 0));
90 }
91 }
92
93 HistoryBlock { lines }
94}
95
96fn transcript_line(role: TranscriptRole, content: &str) -> HistoryLine {
97 transcript_line_with_prefix(role, content, true)
98}
99
100fn transcript_line_with_prefix(
101 role: TranscriptRole,
102 content: &str,
103 include_prefix: bool,
104) -> HistoryLine {
105 let mut runs = Vec::new();
106 let prefix = match (role, include_prefix) {
107 (TranscriptRole::User, true) => Some(("you ", user_style())),
108 (TranscriptRole::Assistant, true) => Some(("assistant ", assistant_style())),
109 (TranscriptRole::Status, _) => None,
110 _ => None,
111 };
112
113 if let Some((label, style)) = prefix {
114 runs.push(HistoryRun {
115 text: label.to_string(),
116 style,
117 });
118 }
119
120 runs.push(HistoryRun {
121 text: content.to_string(),
122 style: content_style(role),
123 });
124
125 HistoryLine { runs }
126}
127
128fn user_style() -> Style {
129 Style::default().fg(Color::Grey).bold(true)
130}
131
132fn assistant_style() -> Style {
133 Style::default().fg(Color::Grey).bold(true)
134}
135
136fn content_style(role: TranscriptRole) -> Style {
137 match role {
138 TranscriptRole::Status => Style::default().fg(Color::DarkGrey),
139 TranscriptRole::User | TranscriptRole::Assistant => Style::default().fg(Color::Grey),
140 }
141}