Skip to main content

ansiq_core/
transcript.rs

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}