suture_sse/
anthropic.rs

1use crate::extractor::{DeltaExtractor, Repair};
2use crate::target::{TargetKind, Targets};
3use serde_json::Value;
4
5/// Anthropic Messages SSE extractor.
6pub struct Anthropic;
7
8impl DeltaExtractor for Anthropic {
9    fn on_event(&self, data: &[u8], targets: &mut Targets) {
10        let v: Value = match serde_json::from_slice(data) {
11            Ok(v) => v,
12            Err(_) => return,
13        };
14        match v.get("type").and_then(Value::as_str) {
15            Some("message_start") => {
16                if let Some(m) = v.get("message") {
17                    if let Some(s) = m.get("id").and_then(Value::as_str) {
18                        targets.id = Some(s.to_string());
19                    }
20                    if let Some(s) = m.get("model").and_then(Value::as_str) {
21                        targets.model = Some(s.to_string());
22                    }
23                }
24            }
25            Some("content_block_delta") => {
26                let idx = v.get("index").and_then(Value::as_u64).unwrap_or(0) as usize;
27                let Some(delta) = v.get("delta") else { return };
28                match delta.get("type").and_then(Value::as_str) {
29                    Some("input_json_delta") => {
30                        if let Some(pj) = delta.get("partial_json").and_then(Value::as_str) {
31                            targets.feed(TargetKind::Block { index: idx }, true, pj.as_bytes());
32                        }
33                    }
34                    Some("text_delta") => {
35                        if let Some(txt) = delta.get("text").and_then(Value::as_str) {
36                            targets.feed(TargetKind::Block { index: idx }, false, txt.as_bytes());
37                        }
38                    }
39                    _ => {}
40                }
41            }
42            _ => {}
43        }
44    }
45
46    fn is_terminator(&self, data: &[u8]) -> bool {
47        serde_json::from_slice::<Value>(data)
48            .ok()
49            .and_then(|v| {
50                v.get("type")
51                    .and_then(Value::as_str)
52                    .map(|s| s == "message_stop")
53            })
54            .unwrap_or(false)
55    }
56
57    fn synthesize(&self, repairs: &[Repair], _targets: &Targets, terminated: bool) -> Vec<u8> {
58        use crate::extractor::json_escape;
59        use crate::target::TargetKind;
60        let mut out = String::new();
61        for r in repairs {
62            let TargetKind::Block { index } = r.kind else {
63                continue;
64            };
65            let esc = json_escape(&r.append);
66            out.push_str("event: content_block_delta\n");
67            out.push_str(&format!(
68                "data: {{\"type\":\"content_block_delta\",\"index\":{index},\"delta\":{{\"type\":\"input_json_delta\",\"partial_json\":\"{esc}\"}}}}\n\n"
69            ));
70            out.push_str("event: content_block_stop\n");
71            out.push_str(&format!(
72                "data: {{\"type\":\"content_block_stop\",\"index\":{index}}}\n\n"
73            ));
74        }
75        if !terminated {
76            out.push_str("event: message_delta\n");
77            out.push_str(
78                "data: {\"type\":\"message_delta\",\"delta\":{\"stop_reason\":\"max_tokens\"}}\n\n",
79            );
80            out.push_str("event: message_stop\n");
81            out.push_str("data: {\"type\":\"message_stop\"}\n\n");
82        }
83        out.into_bytes()
84    }
85}
86
87#[cfg(test)]
88mod tests {
89    use super::*;
90    use crate::extractor::DeltaExtractor;
91    use crate::target::{TargetKind, Targets};
92
93    #[test]
94    fn extracts_input_json_delta() {
95        let ext = Anthropic;
96        let mut t = Targets::new();
97        ext.on_event(
98            br#"{"type":"message_start","message":{"id":"msg_1","model":"claude-3"}}"#,
99            &mut t,
100        );
101        ext.on_event(
102            br#"{"type":"content_block_start","index":0,"content_block":{"type":"tool_use"}}"#,
103            &mut t,
104        );
105        ext.on_event(br#"{"type":"content_block_delta","index":0,"delta":{"type":"input_json_delta","partial_json":"{\"x\":1"}}"#, &mut t);
106        assert_eq!(t.id.as_deref(), Some("msg_1"));
107        assert_eq!(t.model.as_deref(), Some("claude-3"));
108        let state = t.iter().next().unwrap();
109        assert_eq!(state.kind, TargetKind::Block { index: 0 });
110        let r = state.repair();
111        assert!(r.consistent && r.safe);
112        assert_eq!(r.append, b"}");
113    }
114
115    #[test]
116    fn plain_text_delta_not_repaired() {
117        let ext = Anthropic;
118        let mut t = Targets::new();
119        ext.on_event(br#"{"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"Hello"}}"#, &mut t);
120        assert!(!t.iter().next().unwrap().repairable());
121    }
122
123    #[test]
124    fn message_stop_is_terminator() {
125        let ext = Anthropic;
126        assert!(ext.is_terminator(br#"{"type":"message_stop"}"#));
127        assert!(!ext.is_terminator(br#"{"type":"content_block_delta"}"#));
128    }
129}
suture_sse/anthropic.rs

suture_sse/
anthropic.rs
