yapi_core/
messages.rs

1//! Session message roles beyond the provider's four, and their conversion.
2
3use yapi_types::message::{BashExecutionMessage, Content, ContentBlock, Message, UserMessage};
4
5/// Wraps a compaction summary when it is sent to the model.
6pub const COMPACTION_SUMMARY_PREFIX: &str = "The conversation history before this point was compacted into the following summary:\n\n<summary>\n";
7/// Closes [`COMPACTION_SUMMARY_PREFIX`].
8pub const COMPACTION_SUMMARY_SUFFIX: &str = "\n</summary>";
9/// Wraps a branch summary when it is sent to the model.
10pub const BRANCH_SUMMARY_PREFIX: &str =
11    "The following is a summary of a branch that this conversation came back from:\n\n<summary>\n";
12/// Closes [`BRANCH_SUMMARY_PREFIX`].
13pub const BRANCH_SUMMARY_SUFFIX: &str = "</summary>";
14
15/// How a user `!command` reads to the model.
16pub fn bash_execution_text(message: &BashExecutionMessage) -> String {
17    let mut text = format!("Ran `{}`\n", message.command);
18    if message.output.is_empty() {
19        text += "(no output)";
20    } else {
21        text += &format!("```\n{}\n```", message.output);
22    }
23    if message.cancelled {
24        text += "\n\n(command cancelled)";
25    } else if let Some(code) = message.exit_code.filter(|code| *code != 0) {
26        text += &format!("\n\nCommand exited with code {code}");
27    }
28    if message.truncated
29        && let Some(path) = &message.full_output_path
30    {
31        text += &format!("\n\n[Output truncated. Full output: {path}]");
32    }
33    text
34}
35
36fn user_text(text: String, timestamp: u64) -> Message {
37    Message::User(UserMessage {
38        content: Content::Blocks(vec![ContentBlock::text(text)]),
39        timestamp,
40    })
41}
42
43const IMAGES_BLOCKED: &str = "Image reading is disabled.";
44
45/// pi's `blockImages` filter: images in user and tool result messages become
46/// a notice, once for each run of images.
47pub fn block_images(messages: Vec<Message>) -> Vec<Message> {
48    fn filter(blocks: Vec<ContentBlock>) -> Vec<ContentBlock> {
49        let mut out: Vec<ContentBlock> = Vec::with_capacity(blocks.len());
50        for block in blocks {
51            let block = match block {
52                ContentBlock::Image(_) => ContentBlock::text(IMAGES_BLOCKED),
53                other => other,
54            };
55            let repeated = matches!(
56                (&block, out.last()),
57                (ContentBlock::Text(text), Some(ContentBlock::Text(previous)))
58                    if text.text == IMAGES_BLOCKED && previous.text == IMAGES_BLOCKED
59            );
60            if !repeated {
61                out.push(block);
62            }
63        }
64        out
65    }
66    let has_image = |blocks: &[ContentBlock]| {
67        blocks
68            .iter()
69            .any(|block| matches!(block, ContentBlock::Image(_)))
70    };
71    messages
72        .into_iter()
73        .map(|message| match message {
74            Message::User(mut user) => {
75                if let Content::Blocks(blocks) = &mut user.content
76                    && has_image(blocks)
77                {
78                    *blocks = filter(std::mem::take(blocks));
79                }
80                Message::User(user)
81            }
82            Message::ToolResult(mut result) => {
83                if has_image(&result.content) {
84                    result.content = filter(std::mem::take(&mut result.content));
85                }
86                Message::ToolResult(result)
87            }
88            other => other,
89        })
90        .collect()
91}
92
93/// Converts session messages to provider roles: bash executions, custom messages
94/// and summaries become user messages. `!!` executions are dropped.
95pub fn convert_to_llm(messages: Vec<Message>) -> Vec<Message> {
96    messages
97        .into_iter()
98        .filter_map(|message| match message {
99            Message::BashExecution(bash) => (bash.exclude_from_context != Some(true))
100                .then(|| user_text(bash_execution_text(&bash), bash.timestamp)),
101            Message::Custom(custom) => Some(Message::User(UserMessage {
102                content: Content::Blocks(custom.content.into_blocks()),
103                timestamp: custom.timestamp,
104            })),
105            Message::BranchSummary(summary) => Some(user_text(
106                format!(
107                    "{BRANCH_SUMMARY_PREFIX}{}{BRANCH_SUMMARY_SUFFIX}",
108                    summary.summary
109                ),
110                summary.timestamp,
111            )),
112            Message::CompactionSummary(summary) => Some(user_text(
113                format!(
114                    "{COMPACTION_SUMMARY_PREFIX}{}{COMPACTION_SUMMARY_SUFFIX}",
115                    summary.summary
116                ),
117                summary.timestamp,
118            )),
119            other => Some(other),
120        })
121        .collect()
122}
123
124#[cfg(test)]
125mod tests {
126    use super::*;
127
128    #[test]
129    fn blocks_images_as_pi() {
130        let message: Message = serde_json::from_value(serde_json::json!({
131            "role": "toolResult",
132            "toolCallId": "call",
133            "toolName": "read",
134            "content": [
135                {"type": "image", "data": "AA==", "mimeType": "image/png"},
136                {"type": "image", "data": "AA==", "mimeType": "image/png"},
137                {"type": "text", "text": "after"},
138                {"type": "image", "data": "AA==", "mimeType": "image/png"}
139            ],
140            "isError": false,
141            "timestamp": 0
142        }))
143        .unwrap();
144        let Message::ToolResult(result) = &block_images(vec![message])[0] else {
145            panic!("not a tool result");
146        };
147        let texts: Vec<&str> = result
148            .content
149            .iter()
150            .map(|block| match block {
151                ContentBlock::Text(text) => text.text.as_str(),
152                _ => "image",
153            })
154            .collect();
155        assert_eq!(texts, [IMAGES_BLOCKED, "after", IMAGES_BLOCKED]);
156    }
157}
yapi_core/messages.rs

yapi_core/
messages.rs
