rig_core/providers/gemini/streaming/
document.rs

1//! The GenerateContent reassembler: a `streamGenerateContent` reply's chunks
2//! rebuilt into the `generateContent` body the same turn has unary. The
3//! REST, Vertex AI and gRPC wires share it, each over its chunks' REST JSON.
4//!
5//! ```
6//! use rig_core::providers::gemini::streaming::document::GenerateContentResponse;
7//! use serde_json::json;
8//!
9//! let mut document = GenerateContentResponse::default();
10//! for text in ["Hel", "lo"] {
11//!     let chunk = json!({"candidates": [{"content": {"parts": [{"text": text}], "role": "model"}, "index": 0}]});
12//!     document.chunk(chunk.as_object().cloned().unwrap_or_default());
13//! }
14//! assert_eq!(document.document()["candidates"][0]["content"]["parts"], json!([{"text": "Hello"}]));
15//! ```
16
17use serde_json::{Map, Value};
18
19use super::{GenerateContentChunk, GenerateContentDecoder, merge_part};
20use crate::json_utils::Lenient;
21use crate::wire::document::Reassemble;
22use crate::wire::{Decoder, WireEvent, WireFrame};
23
24/// The `generateContent` body a stream of chunks adds up to.
25///
26/// Candidates are kept by their `index`. Their parts append, and a text
27/// part continues the text part before it while `thought` stays the same,
28/// as the unary body states a run of text as one part: its text appends,
29/// a signature joins it unless both carry one, and its other fields
30/// replace. An empty text part with no signature carries nothing. A part
31/// that holds only a signature joins the part before it.
32/// `citationMetadata.citationSources` append. `promptFeedback` keeps its
33/// first value; every other field keeps its last non-null one, and a
34/// `null` only fills an absent field. A reply with no chunk has no
35/// document.
36#[derive(Debug, Default)]
37pub struct GenerateContentResponse {
38    /// The decoder's classifier.
39    classifier: GenerateContentDecoder,
40    document: Map<String, Value>,
41    /// The `index` of each candidate held, in order.
42    indices: Vec<u64>,
43    /// Signatures that arrived alone before any part of their candidate,
44    /// which the candidate's next part takes.
45    signatures: Vec<(u64, String)>,
46}
47
48impl GenerateContentResponse {
49    /// Absorb one chunk, in REST JSON.
50    ///
51    /// For rig's own crates; not covered by semver.
52    #[doc(hidden)]
53    pub fn chunk(&mut self, chunk: Map<String, Value>) {
54        for (key, value) in chunk {
55            match (key.as_str(), value) {
56                ("candidates", Value::Array(candidates)) => {
57                    for (position, candidate) in candidates.into_iter().enumerate() {
58                        self.candidate(position, candidate);
59                    }
60                }
61                // Candidates of another type fail the decoder; the document
62                // keeps the ones it holds.
63                ("candidates", _) => {}
64                ("promptFeedback", value) => {
65                    if self.document.get(&key).is_none_or(Value::is_null) {
66                        self.document.insert(key, value);
67                    }
68                }
69                (_, value) => set(&mut self.document, key, value),
70            }
71        }
72    }
73
74    /// The document the chunks add up to; `Null` when there were none.
75    ///
76    /// For rig's own crates; not covered by semver.
77    #[doc(hidden)]
78    pub fn document(self) -> Value {
79        if self.document.is_empty() {
80            Value::Null
81        } else {
82            Value::Object(self.document)
83        }
84    }
85
86    /// Fold one streamed candidate into the candidate of its `index`.
87    fn candidate(&mut self, position: usize, candidate: Value) {
88        let Value::Object(candidate) = candidate else {
89            return;
90        };
91        let index = candidate
92            .get("index")
93            .and_then(Value::as_u64)
94            .unwrap_or(position as u64);
95        let candidates = slot(&mut self.document, "candidates", Value::Array);
96        let Value::Array(candidates) = candidates else {
97            return;
98        };
99        let at = match self.indices.iter().position(|held| *held == index) {
100            Some(at) => at,
101            None => {
102                self.indices.push(index);
103                candidates.push(Value::Object(Map::new()));
104                candidates.len() - 1
105            }
106        };
107        let Some(Value::Object(held)) = candidates.get_mut(at) else {
108            return;
109        };
110        for (key, value) in candidate {
111            match (key.as_str(), value) {
112                ("content", Value::Object(content)) => {
113                    if let Value::Object(held) = slot(held, "content", Value::Object) {
114                        content_into(held, content, index, &mut self.signatures);
115                    }
116                }
117                ("citationMetadata", Value::Object(citations)) => {
118                    if let Value::Object(held) = slot(held, "citationMetadata", Value::Object) {
119                        for (key, value) in citations {
120                            match (key.as_str(), held.get_mut(&key), value) {
121                                (
122                                    "citationSources",
123                                    Some(Value::Array(sources)),
124                                    Value::Array(more),
125                                ) => {
126                                    sources.extend(more);
127                                }
128                                (_, _, value) => set(held, key, value),
129                            }
130                        }
131                    }
132                }
133                (_, value) => set(held, key, value),
134            }
135        }
136    }
137}
138
139/// Fold a streamed candidate's `content` into the content held so far.
140fn content_into(
141    held: &mut Map<String, Value>,
142    content: Map<String, Value>,
143    index: u64,
144    signatures: &mut Vec<(u64, String)>,
145) {
146    for (key, value) in content {
147        let Value::Array(more) = value else {
148            set(held, key, value);
149            continue;
150        };
151        if key != "parts" {
152            set(held, key, Value::Array(more));
153            continue;
154        }
155        let Value::Array(parts) = slot(held, "parts", Value::Array) else {
156            return;
157        };
158        for part in more {
159            part_into(parts, part, index, signatures);
160        }
161    }
162}
163
164/// Append one streamed part to a candidate's parts, joining the part before
165/// it where the unary body states them as one.
166fn part_into(
167    parts: &mut Vec<Value>,
168    mut part: Value,
169    index: u64,
170    signatures: &mut Vec<(u64, String)>,
171) {
172    let signature = part
173        .str("thoughtSignature")
174        .filter(|signature| !signature.is_empty())
175        .map(str::to_owned);
176    if let Some(text) = part.str("text") {
177        if text.is_empty() && signature.is_none() {
178            return;
179        }
180        let thought = part.bool("thought") == Some(true);
181        if let Some(last) = parts.last_mut()
182            && last.str("text").is_some()
183            && (last.bool("thought") == Some(true)) == thought
184            && !(signed(last) && signature.is_some())
185        {
186            merge_part(last, &part);
187            return;
188        }
189    } else if let Value::Object(fields) = &part {
190        let bare = ["thought", "thoughtSignature", "partMetadata"];
191        let data = fields.keys().any(|key| !bare.contains(&key.as_str()));
192        if !data && let Some(signature) = signature {
193            match parts.last_mut().and_then(Value::as_object_mut) {
194                Some(last) => {
195                    last.insert("thoughtSignature".to_owned(), Value::String(signature));
196                }
197                None => signatures.push((index, signature)),
198            }
199            return;
200        }
201    }
202    if let Some(at) = signatures.iter().position(|(held, _)| *held == index)
203        && let Some(fields) = part.as_object_mut()
204    {
205        let (_, signature) = signatures.remove(at);
206        fields
207            .entry("thoughtSignature")
208            .or_insert(Value::String(signature));
209    }
210    parts.push(part);
211}
212
213/// Whether `part` holds a non-empty signature.
214fn signed(part: &Value) -> bool {
215    part.str("thoughtSignature")
216        .is_some_and(|signature| !signature.is_empty())
217}
218
219/// Set `key` to `value`: a non-null value replaces, and `null` only fills
220/// an absent key.
221fn set(map: &mut Map<String, Value>, key: String, value: Value) {
222    if !value.is_null() || !map.contains_key(&key) {
223        map.insert(key, value);
224    }
225}
226
227/// The value at `key`, made an empty `kind` (a list or an object) when it is
228/// absent or of another type.
229fn slot<'a, T: Default>(
230    map: &'a mut Map<String, Value>,
231    key: &str,
232    kind: fn(T) -> Value,
233) -> &'a mut Value {
234    let slot = map.entry(key).or_insert(Value::Null);
235    if std::mem::discriminant(slot) != std::mem::discriminant(&kind(T::default())) {
236        *slot = kind(T::default());
237    }
238    slot
239}
240
241impl crate::wire::document::Serves<crate::operation::Completion> for GenerateContentResponse {}
242
243impl Reassemble<WireFrame> for GenerateContentResponse {
244    fn absorb(&mut self, frame: &WireFrame) {
245        // The driver reports what does not classify.
246        if let WireEvent::Known(GenerateContentChunk(chunk)) =
247            self.classifier.classify(frame.clone())
248        {
249            self.chunk(chunk);
250        }
251    }
252
253    fn finish(self) -> Value {
254        self.document()
255    }
256}
257
258#[cfg(test)]
259mod tests;
rig_core/providers/gemini/streaming/document.rs

rig_core/providers/gemini/streaming/
document.rs
