Skip to main content

ras_agent/application/
render_step_message.rs

1use ras_llm::{ChatMessage, ContentPart};
2
3use crate::domain::agent_history::StepRecord;
4
5const TEXT_BUDGET: usize = 480;
6const ERROR_BUDGET: usize = 240;
7const SCREENSHOT_MEDIA_TYPE: &str = "image/png";
8
9pub(crate) fn render_step_message(step: &StepRecord) -> Option<ChatMessage> {
10    if step.results.is_empty() {
11        return None;
12    }
13    let mut text = format!("Step {} result:\n", step.step.0);
14    if let Some(url) = &step.url {
15        text.push_str(&format!("url: {url}\n"));
16    }
17    text.push_str("action results:\n");
18    for (i, r) in step.results.iter().enumerate() {
19        text.push_str(&format!("  [{i}]"));
20        if r.is_done {
21            text.push_str(" done");
22        }
23        if let Some(err) = &r.error {
24            text.push_str(&format!(" error: {}", truncate(err, ERROR_BUDGET)));
25        } else if let Some(c) = &r.extracted_content {
26            text.push_str(&format!(" {}", truncate(c, TEXT_BUDGET)));
27        }
28        text.push('\n');
29    }
30
31    let mut parts: Vec<ContentPart> = Vec::with_capacity(1 + image_count(step));
32    parts.push(ContentPart::Text { text });
33    for r in &step.results {
34        for img_b64 in &r.images {
35            parts.push(ContentPart::ImageBase64 {
36                media_type: SCREENSHOT_MEDIA_TYPE.into(),
37                data: img_b64.clone(),
38            });
39        }
40    }
41    Some(ChatMessage::user_parts(parts))
42}
43
44fn image_count(step: &StepRecord) -> usize {
45    step.results.iter().map(|r| r.images.len()).sum()
46}
47
48fn truncate(s: &str, max: usize) -> String {
49    if s.chars().count() <= max {
50        s.to_string()
51    } else {
52        let mut out: String = s.chars().take(max).collect();
53        out.push_str("…");
54        out
55    }
56}
57
58#[cfg(test)]
59mod tests {
60    use super::*;
61    use chrono::Utc;
62    use ras_llm::ChatMessage;
63    use ras_types::{ActionResult, StepId};
64
65    use crate::domain::agent_output::{AgentBrain, AgentOutput};
66    use crate::domain::step_metadata::StepMetadata;
67
68    fn step(results: Vec<ActionResult>) -> StepRecord {
69        StepRecord {
70            step: StepId(7),
71            started_at: Utc::now(),
72            url: Some("https://example.com/login".parse().expect("test url")),
73            output: AgentOutput {
74                current_state: AgentBrain {
75                    evaluation_previous_goal: String::new(),
76                    memory: String::new(),
77                    next_goal: String::new(),
78                },
79                action: vec![],
80                plan: None,
81                current_plan_item: None,
82            },
83            results,
84            metadata: StepMetadata::default(),
85        }
86    }
87
88    #[test]
89    fn empty_results_yields_no_message() {
90        assert!(render_step_message(&step(vec![])).is_none());
91    }
92
93    #[test]
94    fn text_only_result_emits_text_part_only() {
95        let r = ActionResult::ok("clicked login button");
96        let msg = render_step_message(&step(vec![r])).expect("msg");
97        let ChatMessage::User(u) = msg else {
98            panic!("expected user");
99        };
100        assert_eq!(u.content.len(), 1);
101        match &u.content[0] {
102            ContentPart::Text { text } => {
103                assert!(text.contains("Step 7 result:"));
104                assert!(text.contains("url: https://example.com/login"));
105                assert!(text.contains("clicked login button"));
106            }
107            other => panic!("expected text part, got {other:?}"),
108        }
109    }
110
111    #[test]
112    fn screenshot_result_emits_text_plus_image_part() {
113        let r = ActionResult::ok("captured screenshot").with_image("AAAA");
114        let msg = render_step_message(&step(vec![r])).expect("msg");
115        let ChatMessage::User(u) = msg else {
116            panic!("expected user");
117        };
118        assert_eq!(u.content.len(), 2);
119        assert!(matches!(&u.content[0], ContentPart::Text { .. }));
120        match &u.content[1] {
121            ContentPart::ImageBase64 { media_type, data } => {
122                assert_eq!(media_type, "image/png");
123                assert_eq!(data, "AAAA");
124            }
125            other => panic!("expected image part, got {other:?}"),
126        }
127    }
128
129    #[test]
130    fn multiple_images_across_results_all_attached() {
131        let r1 = ActionResult::ok("step a").with_image("AAAA");
132        let r2 = ActionResult::ok("step b")
133            .with_image("BBBB")
134            .with_image("CCCC");
135        let msg = render_step_message(&step(vec![r1, r2])).expect("msg");
136        let ChatMessage::User(u) = msg else {
137            panic!("expected user");
138        };
139        let images: Vec<&String> = u
140            .content
141            .iter()
142            .filter_map(|p| match p {
143                ContentPart::ImageBase64 { data, .. } => Some(data),
144                _ => None,
145            })
146            .collect();
147        assert_eq!(images, vec!["AAAA", "BBBB", "CCCC"]);
148    }
149
150    #[test]
151    fn error_result_emits_error_in_text() {
152        let r = ActionResult::err("element not found");
153        let msg = render_step_message(&step(vec![r])).expect("msg");
154        let ChatMessage::User(u) = msg else {
155            panic!("expected user");
156        };
157        match &u.content[0] {
158            ContentPart::Text { text } => assert!(text.contains("error: element not found")),
159            _ => panic!("expected text"),
160        }
161    }
162}