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