Skip to main content

beam_core/
workflow_output.rs

1use anyhow::{Context, Result};
2use serde_json::Value;
3
4pub const WORKFLOW_OUTPUT_BEGIN: &str = "<WORKFLOW_OUTPUT>";
5pub const WORKFLOW_OUTPUT_END: &str = "</WORKFLOW_OUTPUT>";
6
7pub fn with_workflow_output_protocol(prompt: &str) -> String {
8    if prompt.contains(WORKFLOW_OUTPUT_BEGIN) {
9        return prompt.to_string();
10    }
11    format!(
12        "{prompt}\n\n---\nReturn the final workflow result as exactly one JSON value between these markers:\n\n{begin}\n{{\"...your JSON output...\"}}\n{end}\n\nDo not put commentary inside the markers.\n",
13        prompt = prompt,
14        begin = WORKFLOW_OUTPUT_BEGIN,
15        end = WORKFLOW_OUTPUT_END
16    )
17}
18
19pub fn parse_workflow_output(text: &str) -> Result<Value> {
20    let last_end = text
21        .rfind(WORKFLOW_OUTPUT_END)
22        .context("workflow output missing end marker")?;
23    let begin_before_end = text[..last_end]
24        .rfind(WORKFLOW_OUTPUT_BEGIN)
25        .context("workflow output missing begin marker")?;
26    let raw = text[begin_before_end + WORKFLOW_OUTPUT_BEGIN.len()..last_end]
27        .trim()
28        .to_string();
29    let sanitized = sanitize_workflow_output_block(&raw);
30    let value = serde_json::from_str::<Value>(&sanitized)
31        .with_context(|| "workflow output block is not valid JSON")?;
32    Ok(value)
33}
34
35fn sanitize_workflow_output_block(block: &str) -> String {
36    block
37        .replace("\u{1b}][", "")
38        .replace("\u{1b}[", "")
39        .replace('\u{7}', "")
40        .replace('\u{1b}', "")
41        .replace('\u{0}', "")
42        .trim()
43        .to_string()
44}
45
46#[cfg(test)]
47mod tests {
48    use super::*;
49
50    #[test]
51    fn protocol_is_idempotent() {
52        let prompt = with_workflow_output_protocol("hello");
53        assert!(prompt.contains(WORKFLOW_OUTPUT_BEGIN));
54        assert_eq!(prompt, with_workflow_output_protocol(&prompt));
55    }
56
57    #[test]
58    fn parse_output_extracts_json_between_markers() {
59        let parsed = parse_workflow_output(
60            "noise\n<WORKFLOW_OUTPUT>\n{\"ok\":true}\n</WORKFLOW_OUTPUT>\nnoise",
61        )
62        .expect("parse");
63        assert_eq!(parsed["ok"], Value::Bool(true));
64    }
65}