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

rig_core/providers/gemini/streaming/
document.rs
