suture_sse/
gemini.rs

1use crate::extractor::{json_escape, DeltaExtractor, Repair};
2use crate::target::{TargetKind, Targets};
3use serde_json::Value;
4
5/// Vertex AI Gemini SSE extractor (`streamGenerateContent?alt=sse`). Repairs
6/// reassembled `candidates[i].content.parts[].text` when it is JSON-looking
7/// (JSON-mode output). `functionCall` parts arrive whole and are ignored.
8pub struct Gemini;
9
10impl DeltaExtractor for Gemini {
11    fn on_event(&self, data: &[u8], targets: &mut Targets) {
12        let v: Value = match serde_json::from_slice(data) {
13            Ok(v) => v,
14            Err(_) => return,
15        };
16        let Some(cands) = v.get("candidates").and_then(Value::as_array) else {
17            return;
18        };
19        for cand in cands {
20            let idx = cand.get("index").and_then(Value::as_u64).unwrap_or(0) as usize;
21            let Some(parts) = cand
22                .get("content")
23                .and_then(|c| c.get("parts"))
24                .and_then(Value::as_array)
25            else {
26                continue;
27            };
28            for part in parts {
29                if let Some(text) = part.get("text").and_then(Value::as_str) {
30                    targets.feed(TargetKind::Content { choice: idx }, false, text.as_bytes());
31                }
32            }
33        }
34    }
35
36    fn is_terminator(&self, _data: &[u8]) -> bool {
37        false
38    }
39
40    fn synthesize(&self, repairs: &[Repair], _targets: &Targets, _terminated: bool) -> Vec<u8> {
41        let mut out = String::new();
42        for r in repairs {
43            let choice = if let TargetKind::Content { choice } = &r.kind {
44                *choice
45            } else {
46                continue;
47            };
48            let esc = json_escape(&r.append);
49            out.push_str(&format!(
50                "data: {{\"candidates\":[{{\"index\":{choice},\"content\":{{\"parts\":[{{\"text\":\"{esc}\"}}]}},\"finishReason\":\"length\"}}]}}\n\n"
51            ));
52        }
53        out.into_bytes()
54    }
55}
56
57#[cfg(test)]
58mod tests {
59    use super::*;
60    use crate::extractor::DeltaExtractor;
61    use crate::target::{TargetKind, Targets};
62
63    #[test]
64    fn extracts_json_text_parts() {
65        let ext = Gemini;
66        let mut t = Targets::new();
67        ext.on_event(
68            br#"{"candidates":[{"index":0,"content":{"role":"model","parts":[{"text":"{\"city\":"}]}}]}"#,
69            &mut t,
70        );
71        ext.on_event(
72            br#"{"candidates":[{"index":0,"content":{"parts":[{"text":"\"Par"}]}}]}"#,
73            &mut t,
74        );
75        let state = t.iter().next().expect("one target");
76        assert_eq!(state.kind, TargetKind::Content { choice: 0 });
77        assert!(state.repairable(), "json-looking text must be repairable");
78        let r = state.repair();
79        assert!(r.consistent && r.safe);
80        assert_eq!(r.append, b"\"}");
81    }
82
83    #[test]
84    fn plain_text_not_repaired() {
85        let ext = Gemini;
86        let mut t = Targets::new();
87        ext.on_event(
88            br#"{"candidates":[{"index":0,"content":{"parts":[{"text":"Hello there"}]}}]}"#,
89            &mut t,
90        );
91        assert!(!t.iter().next().unwrap().repairable());
92    }
93
94    #[test]
95    fn function_call_part_is_ignored() {
96        let ext = Gemini;
97        let mut t = Targets::new();
98        ext.on_event(
99            br#"{"candidates":[{"index":0,"content":{"parts":[{"functionCall":{"name":"f","args":{"x":1}}}]}}]}"#,
100            &mut t,
101        );
102        assert!(
103            t.iter().next().is_none(),
104            "functionCall parts create no target"
105        );
106    }
107
108    #[test]
109    fn never_terminator() {
110        let ext = Gemini;
111        assert!(!ext.is_terminator(br#"{"candidates":[{"finishReason":"STOP"}]}"#));
112    }
113
114    #[test]
115    fn multi_candidate_indexing() {
116        let ext = Gemini;
117        let mut t = Targets::new();
118        ext.on_event(
119            br#"{"candidates":[{"index":0,"content":{"parts":[{"text":"["}]}},{"index":1,"content":{"parts":[{"text":"{"}]}}]}"#,
120            &mut t,
121        );
122        let kinds: Vec<_> = t.iter().map(|s| s.kind.clone()).collect();
123        assert!(kinds.contains(&TargetKind::Content { choice: 0 }));
124        assert!(kinds.contains(&TargetKind::Content { choice: 1 }));
125    }
126}
suture_sse/gemini.rs

suture_sse/
gemini.rs
