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}', '\u{1b}', '\u{0}'], "")
40        .trim()
41        .to_string()
42}
43
44#[cfg(test)]
45mod tests {
46    use super::*;
47
48    #[test]
49    fn protocol_is_idempotent() {
50        let prompt = with_workflow_output_protocol("hello");
51        assert!(prompt.contains(WORKFLOW_OUTPUT_BEGIN));
52        assert_eq!(prompt, with_workflow_output_protocol(&prompt));
53    }
54
55    #[test]
56    fn parse_output_extracts_json_between_markers() {
57        let parsed = parse_workflow_output(
58            "noise\n<WORKFLOW_OUTPUT>\n{\"ok\":true}\n</WORKFLOW_OUTPUT>\nnoise",
59        )
60        .expect("parse");
61        assert_eq!(parsed["ok"], Value::Bool(true));
62    }
63}