beam_core/
workflow_output.rs1use 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}